Read about our upcoming Code of Conduct on this issue

This instance will be upgraded to Heptapod 0.28.1 on 2022-01-26 at 16:00 UTC+1 (a few minutes of down time)

spatiotemporal_spectra.py 44.7 KB
Newer Older
1
"""
2
3
Spatiotemporal Spectra (:mod:`fluidsim.base.output.spatiotemporal_spectra`)
===========================================================================
4
5
6

Provides:

Pierre Augier's avatar
Pierre Augier committed
7
8
9
10
11
.. autoclass:: SpatioTemporalSpectra3D
   :members:
   :private-members:

.. autoclass:: SpatioTemporalSpectra2D
12
13
14
15
16
17
18
19
   :members:
   :private-members:

"""

from pathlib import Path
from logging import warn
from math import pi
20

21
22
23
import numpy as np
from scipy import signal
import h5py
Jason Reneuve's avatar
Jason Reneuve committed
24
from rich.progress import Progress
25
26

from fluiddyn.util import mpi
27
from fluidsim.util.util import open_patient
28
29
from fluidsim.base.output.base import SpecificOutput

30
from transonic import boost, Array, Type
31

32
33
Uf32f64 = Type(np.float32, np.float64)
A = Array[Uf32f64, "1d"]
34

35
36
37

@boost
def find_index_first_geq(arr: A, value: Uf32f64):
38
39
40
41
42
43
44
    """find the first index such that `arr[index] >= value`"""
    for i, v in enumerate(arr):
        if v >= value:
            return i
    raise ValueError("No index such that `arr[index] >= value`")


45
46
@boost
def find_index_first_g(arr: A, value: Uf32f64):
47
48
49
50
51
52
53
    """find the first index such that `arr[index] > value`"""
    for i, v in enumerate(arr):
        if v > value:
            return i
    raise ValueError("No index such that `arr[index] > value`")


54
55
@boost
def find_index_first_l(arr: A, value: Uf32f64):
56
57
58
59
60
61
62
    """find the first index such that `arr[index] < value`"""
    for i, v in enumerate(arr):
        if v < value:
            return i
    raise ValueError("No index such that `arr[index] < value`")


63
def filter_tmins_paths(tmin, tmins, paths):
64
65
    if tmins.size == 1:
        return tmins, paths
Jason Reneuve's avatar
Jason Reneuve committed
66
    delta_tmin = np.diff(tmins).min()
67
    start = find_index_first_l(tmin - tmins, delta_tmin)
68
69
70
    return tmins[start:], paths[start:]


71
72
@boost
def get_arange_minmax(times: A, tmin: Uf32f64, tmax: Uf32f64):
73
74
75
76
77
78
    """get a range of index for which `tmin <= times[i] <= tmax`

    This assumes that `times` is sorted.

    """

79
    if tmin <= times[0]:
80
81
        start = 0
    else:
82
        start = find_index_first_geq(times, tmin)
83

84
    if tmax >= times[-1]:
85
86
        stop = len(times)
    else:
87
        stop = find_index_first_g(times, tmax)
88
89
90
91

    return np.arange(start, stop)


Pierre Augier's avatar
Pierre Augier committed
92
class SpatioTemporalSpectra3D(SpecificOutput):
93
94
95
96
97
    """
    Computes the spatiotemporal spectra.
    """

    _tag = "spatiotemporal_spectra"
Pierre Augier's avatar
Pierre Augier committed
98
    nb_dim = 3
99

Pierre Augier's avatar
Pierre Augier committed
100
101
102
    @classmethod
    def _complete_params_with_default(cls, params):
        params.output.periods_save._set_attrib(cls._tag, 0)
103
        params.output._set_child(
Pierre Augier's avatar
Pierre Augier committed
104
            cls._tag,
105
            attribs={
106
                "probes_region": None,
107
                "file_max_size": 10.0,  # MB
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
108
                "SAVE_AS_COMPLEX64": True,
109
110
111
            },
        )

112
        params.output.spatiotemporal_spectra._set_doc(
113
            """
114
            probes_region: int tuple (default:None)
115
116
117

                Boundaries of the region to record in the spectral domain.

Jason Reneuve's avatar
Jason Reneuve committed
118
                probes_region = (ikxmax, ikymax, ikzmax), in nondimensional units (mode indices).
Jason Reneuve's avatar
tests    
Jason Reneuve committed
119
120
121

                The resulting ranges for each wavevectors are :

Jason Reneuve's avatar
Jason Reneuve committed
122
                    ikx in [0 , ikxmax]
Jason Reneuve's avatar
tests    
Jason Reneuve committed
123

Jason Reneuve's avatar
Jason Reneuve committed
124
                    iky in [-ikymax+1 , ikymax]
Jason Reneuve's avatar
tests    
Jason Reneuve committed
125

Jason Reneuve's avatar
Jason Reneuve committed
126
                    ikz in [-ikzmax+1 , ikzmax]
127

Jason Reneuve's avatar
Jason Reneuve committed
128
                If None, set all ikmax = 4.
129
130
131
132
133

            file_max_size: float (default: 10.0)

                Maximum size of one time series file, in megabytes.

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
134
            SAVE_AS_COMPLEX64: bool (default: True)
135

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
136
137
138
                If set to False, probes data is saved as complex128.

                Warning : saving as complex128 reduces digital noise at high frequency, but doubles the size of the output!
139
140
141
142
143
144
145

            """
        )

    def __init__(self, output):
        params = output.sim.params
        try:
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
146
            params_st_spec = params.output.spatiotemporal_spectra
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
        except AttributeError:
            warn(
                "Cannot initialize spatiotemporal spectra output because "
                "`params` does not contain parameters for this class."
            )
            return

        super().__init__(
            output,
            period_save=params.output.periods_save.spatiotemporal_spectra,
        )

        oper = self.sim.oper

        # Parameters
        self.period_save = params.output.periods_save.spatiotemporal_spectra

        self.path_dir = Path(self.sim.output.path_run) / "spatiotemporal"
        self.keys_fields = self.sim.info_solver.classes.State.keys_state_phys

        if not output._has_to_save:
            self.period_save = 0.0
        if self.period_save == 0.0:
            return

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
172
        if params_st_spec.probes_region is not None:
Pierre Augier's avatar
Pierre Augier committed
173
174
175
176
            if self.nb_dim == 3:
                ikxmax, ikymax, ikzmax = params_st_spec.probes_region
            else:
                ikxmax, ikymax = params_st_spec.probes_region
177
        else:
Pierre Augier's avatar
Pierre Augier committed
178
179
180
181
            ikxmax = ikymax = 4
            if self.nb_dim == 3:
                ikzmax = 4

182
183
        ikxmax = min(int(round(ikxmax)), params.oper.nx // 2)
        ikymax = min(int(round(ikymax)), params.oper.ny // 2)
Pierre Augier's avatar
Pierre Augier committed
184
185

        if self.nb_dim == 3:
186
            ikzmax = min(int(round(ikzmax)), params.oper.nz // 2)
Jason Reneuve's avatar
Jason Reneuve committed
187
            self.probes_region = ikxmax, ikymax, ikzmax
Pierre Augier's avatar
Pierre Augier committed
188
189
        else:
            self.probes_region = ikxmax, ikymax
190

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
191
192
        self.file_max_size = params_st_spec.file_max_size
        self.SAVE_AS_COMPLEX64 = params_st_spec.SAVE_AS_COMPLEX64
193

194
195
196
        # region must be int tuple
        ikxmax = int(ikxmax)
        ikymax = int(ikymax)
Pierre Augier's avatar
Pierre Augier committed
197
198
        if self.nb_dim == 3:
            ikzmax = int(ikzmax)
199

200
        # dimensions order in Fourier space
Pierre Augier's avatar
Pierre Augier committed
201
202
203
204
205
206
        if self.nb_dim == 3:
            self.dims_order = np.array(oper.oper_fft.get_dimX_K())
        else:
            self.dims_order = np.arange(2)
            if oper.oper_fft.get_is_transposed():
                self.dims_order = self.dims_order[::-1]
207
208
209
210
211
212
213

        # data directory
        if mpi.rank == 0:
            self.path_dir.mkdir(exist_ok=True)
        if mpi.nb_proc > 1:
            mpi.comm.barrier()

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
214
215
216
217
218
219
220
221
        # data type and size
        if self.SAVE_AS_COMPLEX64:
            self.datatype = np.complex64
            size_1_number = 8e-6
        else:
            self.datatype = np.complex128
            size_1_number = 16e-6

222
223
224
225
        # check for existing files
        paths = sorted(self.path_dir.glob("rank*.h5"))
        if paths:
            # check values in files
226
            with open_patient(paths[0], "r") as file:
227
228
                if file.attrs["nb_proc"] != mpi.nb_proc:
                    raise ValueError("process number is different from files")
Jason Reneuve's avatar
Jason Reneuve committed
229
                if (file.attrs["dims_order"] != self.dims_order).any():
230
                    raise ValueError("dimensions order is different from files")
Jason Reneuve's avatar
Jason Reneuve committed
231
232
                if (file.attrs["probes_region"] != self.probes_region).any():
                    raise ValueError("probes region is different from files")
233
            # init from files
Jason Reneuve's avatar
Jason Reneuve committed
234
            INIT_FROM_PARAMS = False
235
236
237
            paths = [p for p in paths if p.name.startswith(f"rank{mpi.rank:05}")]
            if paths:
                self.path_file = paths[-1]
238
                with open_patient(self.path_file, "r") as file:
239
                    self.index_file = file.attrs["index_file"]
240
                    self.probes_k0adim_loc = file["probes_k0adim_loc"][:]
241
                    self.probes_ik0_loc = file["probes_ik0_loc"][:]
Pierre Augier's avatar
Pierre Augier committed
242
                    self.probes_k1adim_loc = file["probes_k1adim_loc"][:]
243
                    self.probes_ik1_loc = file["probes_ik1_loc"][:]
Pierre Augier's avatar
Pierre Augier committed
244
245
246
247
248

                    if self.nb_dim == 3:
                        self.probes_k2adim_loc = file["probes_k2adim_loc"][:]
                        self.probes_ik2_loc = file["probes_ik2_loc"][:]

Jason Reneuve's avatar
Jason Reneuve committed
249
                    self.probes_nb_loc = self.probes_ik0_loc.size
250
251
252
253
254
255
256
257
258
259
260
261
262
263
                    self.number_times_in_file = file["times"].size
                    self.t_last_save = file["times"][-1]
            else:
                # no probes in proc
                self.path_file = None
                self.index_file = 0
                self.number_times_in_file = 0
                self.probes_nb_loc = 0
                self.probes_ik0_loc = []
                self.probes_ik1_loc = []
                self.probes_ik2_loc = []

        else:
            # no files were found : initialize from params
Jason Reneuve's avatar
Jason Reneuve committed
264
            INIT_FROM_PARAMS = True
Pierre Augier's avatar
Pierre Augier committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287

            if self.nb_dim == 3:
                # pair kx,ky,kz with k0,k1,k2
                iksmax = np.array([ikzmax, ikymax, ikxmax])
                iksmin = np.array([1 - ikzmax, 1 - ikymax, 0])
                ik0max, ik1max, ik2max = iksmax[self.dims_order]
                ik0min, ik1min, ik2min = iksmin[self.dims_order]

                # local probes indices
                (
                    k0_adim_loc,
                    k1_adim_loc,
                    k2_adim_loc,
                ) = oper.oper_fft.get_k_adim_loc()
                K0_adim, K1_adim, K2_adim = np.meshgrid(
                    k0_adim_loc, k1_adim_loc, k2_adim_loc, indexing="ij"
                )
            else:
                iksmax = np.array([ikymax, ikxmax])
                iksmin = np.array([1 - ikymax, 0])
                ik0max, ik1max = iksmax[self.dims_order]
                ik0min, ik1min = iksmin[self.dims_order]

Pierre Augier's avatar
Pierre Augier committed
288
289
290
291
292
293
                kx_adim_loc = np.array(
                    (oper.kx_loc / oper.deltakx).round(), dtype=int
                )
                ky_adim_loc = np.array(
                    (oper.ky_loc / oper.deltaky).round(), dtype=int
                )
Pierre Augier's avatar
Pierre Augier committed
294
                if oper.oper_fft.get_is_transposed():
Pierre Augier's avatar
Pierre Augier committed
295
296
                    k0_adim_loc = kx_adim_loc
                    k1_adim_loc = ky_adim_loc
Pierre Augier's avatar
Pierre Augier committed
297
                else:
Pierre Augier's avatar
Pierre Augier committed
298
299
                    k0_adim_loc = ky_adim_loc
                    k1_adim_loc = kx_adim_loc
Pierre Augier's avatar
Pierre Augier committed
300
301
302
303
304

                K0_adim, K1_adim = np.meshgrid(
                    k0_adim_loc, k1_adim_loc, indexing="ij"
                )

305
            cond_region = (
Jason Reneuve's avatar
Jason Reneuve committed
306
307
308
309
                (K0_adim >= ik0min)
                & (K0_adim <= ik0max)
                & (K1_adim >= ik1min)
                & (K1_adim <= ik1max)
310
            )
Pierre Augier's avatar
Pierre Augier committed
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327

            if self.nb_dim == 3:
                cond_region = (
                    cond_region & (K2_adim >= ik2min) & (K2_adim <= ik2max)
                )

            if self.nb_dim == 3:
                (
                    self.probes_ik0_loc,
                    self.probes_ik1_loc,
                    self.probes_ik2_loc,
                ) = np.where(cond_region)
            else:
                (
                    self.probes_ik0_loc,
                    self.probes_ik1_loc,
                ) = np.where(cond_region)
328

Jason Reneuve's avatar
Jason Reneuve committed
329
            self.probes_nb_loc = self.probes_ik0_loc.size
330

331
            # local probes wavenumbers (nondimensional)
Pierre Augier's avatar
Pierre Augier committed
332
333
334
335
336
            self.probes_k0adim_loc = self._get_data_probe_from_field(K0_adim)
            self.probes_k1adim_loc = self._get_data_probe_from_field(K1_adim)

            if self.nb_dim == 3:
                self.probes_k2adim_loc = self._get_data_probe_from_field(K2_adim)
337

Jason Reneuve's avatar
Jason Reneuve committed
338
            # initialize files info
339
340
341
342
343
344
345
346
347
348
349
350
            self.index_file = 0
            self.number_times_in_file = 0
            self.t_last_save = -self.period_save

        # size of a single write: nb_fields * probes_nb_loc + time
        probes_write_size = (
            len(self.keys_fields) * self.probes_nb_loc + 1
        ) * size_1_number
        self.max_number_times_in_file = int(
            self.file_max_size / probes_write_size
        )

Jason Reneuve's avatar
Jason Reneuve committed
351
352
353
354
        # initialize files
        if INIT_FROM_PARAMS and self.probes_nb_loc > 0:
            self._init_new_file(tmin_file=self.sim.time_stepping.t)

355
356
357
358
    def _init_files(self, arrays_1st_time=None):
        # we don't want to do anything when this function is called.
        pass

Jason Reneuve's avatar
Jason Reneuve committed
359
    def _init_new_file(self, tmin_file=None):
360
        """Initializes a new file"""
Jason Reneuve's avatar
Jason Reneuve committed
361
362
363
364
365
366
367
368
369
        if tmin_file is not None:
            # max number of digits = int(log10(t_end)) + 1
            # add .3f precision = 4 additional characters
            # +2 by anticipation of potential restarts
            str_width = int(np.log10(self.sim.params.time_stepping.t_end)) + 7
            ind_str = f"tmin{tmin_file:0{str_width}.3f}"
        else:
            ind_str = f"file{self.index_file:04}"
        self.path_file = self.path_dir / f"rank{mpi.rank:05}_{ind_str}.h5"
370
        with open_patient(self.path_file, "w") as file:
371
372
373
            file.attrs["nb_proc"] = mpi.nb_proc
            file.attrs["dims_order"] = self.dims_order
            file.attrs["index_file"] = self.index_file
Jason Reneuve's avatar
Jason Reneuve committed
374
            file.attrs["probes_region"] = self.probes_region
Jason Reneuve's avatar
Jason Reneuve committed
375
376
            file.attrs["period_save"] = self.period_save
            file.attrs["max_number_times_in_file"] = self.max_number_times_in_file
377
            create_ds = file.create_dataset
378
            create_ds("probes_k0adim_loc", data=self.probes_k0adim_loc)
379
            create_ds("probes_ik0_loc", data=self.probes_ik0_loc)
Pierre Augier's avatar
Pierre Augier committed
380
            create_ds("probes_k1adim_loc", data=self.probes_k1adim_loc)
381
            create_ds("probes_ik1_loc", data=self.probes_ik1_loc)
Pierre Augier's avatar
Pierre Augier committed
382
383
384
            if self.nb_dim == 3:
                create_ds("probes_k2adim_loc", data=self.probes_k2adim_loc)
                create_ds("probes_ik2_loc", data=self.probes_ik2_loc)
385
386
            for key in self.keys_fields:
                create_ds(
387
                    key + "_Fourier_loc",
388
389
                    (self.probes_nb_loc, 1),
                    maxshape=(self.probes_nb_loc, None),
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
390
                    dtype=self.datatype,
391
392
393
394
395
                )
            create_ds("times", (1,), maxshape=(None,))

    def _write_to_file(self, data):
        """Writes a file with the temporal data"""
396
        with open_patient(self.path_file, "a") as file:
397
398
399
400
401
402
403
404
405
            for k, v in data.items():
                dset = file[k]
                if k.startswith("times"):
                    dset.resize((self.number_times_in_file,))
                    dset[-1] = v
                else:
                    dset.resize((self.probes_nb_loc, self.number_times_in_file))
                    dset[:, -1] = v

406
    def _add_probes_data_to_dict(self, data, key):
407
        """Probes fields in Fourier space and append data to a dict object"""
Pierre Augier's avatar
Pierre Augier committed
408
409
410
        data[key + "_Fourier_loc"] = self._get_data_probe_from_field(
            self.sim.state.get_var(f"{key}_fft")
        )
411
412
413

    def _online_save(self):
        """Prepares data and writes to file"""
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
414
415
416
417
418
419
420
421
        tsim = self.sim.time_stepping.t
        if (
            tsim + 1e-15
        ) // self.period_save > self.t_last_save // self.period_save:
            # if max write number is reached, init new file
            if self.number_times_in_file >= self.max_number_times_in_file:
                self.index_file += 1
                self.number_times_in_file = 0
Jason Reneuve's avatar
Jason Reneuve committed
422
                self._init_new_file(tmin_file=self.sim.time_stepping.t)
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
423
424
425
426
427
428
            # get data from probes
            data = {"times": self.sim.time_stepping.t}
            for key in self.keys_fields:
                self._add_probes_data_to_dict(data, key)
            # write to file
            self.number_times_in_file += 1
429
430
            if self.probes_nb_loc > 0:
                self._write_to_file(data)
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
431
            self.t_last_save = tsim
432

433
    def load_time_series(self, keys=None, tmin=0, tmax=None, dtype=None):
434
        """load time series from files"""
435
436
437
438
439
440

        if mpi.nb_proc > 1:
            raise RuntimeError(
                "This postprocessing function should not be called with MPI."
            )

441
442
        if keys is None:
            keys = self.keys_fields
443
444
445
446
447
448
449
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        # get ranks
        paths = sorted(self.path_dir.glob("rank*.h5"))
        ranks = sorted({int(p.name[4:9]) for p in paths})

450
451
452
453
454
455
        # get times and dimensions order from the files of first rank
        print(f"load times series...")
        paths_1st_rank = [
            p for p in paths if p.name.startswith(f"rank{ranks[0]:05}")
        ]

456
        with open_patient(paths_1st_rank[0], "r") as file:
Pierre Augier's avatar
Pierre Augier committed
457
            dims_order = file.attrs["dims_order"]
458
            region = file.attrs["probes_region"]
459
            if dtype is None:
460
                dtype = file[keys[0] + "_Fourier_loc"].dtype
461

Jason Reneuve's avatar
Jason Reneuve committed
462
463
        # get list of useful files, from tmin
        tmins_files = np.array([float(p.name[14:-3]) for p in paths_1st_rank])
464
465
466
        tmins_files, paths_1st_rank = filter_tmins_paths(
            tmin, tmins_files, paths_1st_rank
        )
Jason Reneuve's avatar
Jason Reneuve committed
467
468
469
470
471
472
473
474
475

        with Progress() as progress:
            npaths = len(paths_1st_rank)
            task_files = progress.add_task(
                "Getting times from rank 0...", total=npaths
            )

            times = []
            for ip, path in enumerate(paths_1st_rank):
476
                with open_patient(path, "r") as file:
Jason Reneuve's avatar
Jason Reneuve committed
477
478
479
480
481
482
483
484
                    if tmins_files[ip] > tmax:
                        progress.update(task_files, completed=npaths)
                        break
                    times_file = file["times"][:]
                    cond_times = (times_file >= tmin) & (times_file <= tmax)
                    times.append(times_file[cond_times])
                    progress.update(task_files, advance=1)

485
486
        times = np.concatenate(times)

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
487
488
489
        tmin = times.min()
        tmax = times.max()
        print(f"tmin={tmin:8.6g}, tmax={tmax:8.6g}, nit={times.size}")
490
491

        # get sequential shape of Fourier space
Pierre Augier's avatar
Pierre Augier committed
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
        if self.nb_dim == 3:
            ikxmax, ikymax, ikzmax = region
            iksmax = np.array([ikzmax, ikymax, ikxmax])
            iksmin = np.array([1 - ikzmax, 1 - ikymax, 0])
            ik0max, ik1max, ik2max = iksmax[dims_order]
            ik0min, ik1min, ik2min = iksmin[dims_order]
            shape_series = (
                ik0max + 1 - ik0min,
                ik1max + 1 - ik1min,
                ik2max + 1 - ik2min,
                times.size,
            )
        else:
            ikxmax, ikymax = region
            iksmax = np.array([ikymax, ikxmax])
            iksmin = np.array([1 - ikymax, 0])
            ik0max, ik1max = iksmax[dims_order]
            ik0min, ik1min = iksmin[dims_order]
            shape_series = (
                ik0max + 1 - ik0min,
                ik1max + 1 - ik1min,
                times.size,
            )
515
516

        # load series, rebuild as state_spect arrays + time
517
518
519
        series = {
            f"{k}_Fourier": np.empty(shape_series, dtype=dtype) for k in keys
        }
Jason Reneuve's avatar
Jason Reneuve committed
520
521
        with Progress() as progress:
            task_ranks = progress.add_task("Rearranging...", total=len(ranks))
Jason Reneuve's avatar
Jason Reneuve committed
522
            task_files = progress.add_task("Rank 00000...", total=npaths)
Jason Reneuve's avatar
Jason Reneuve committed
523
524
525
526
527
            # loop on all files, rank by rank
            for rank in ranks:
                paths_rank = [
                    p for p in paths if p.name.startswith(f"rank{rank:05}")
                ]
Jason Reneuve's avatar
Jason Reneuve committed
528
529
530

                # get list of useful files, from tmin
                tmins_files = np.array([float(p.name[14:-3]) for p in paths_rank])
531
532
533
                tmins_files, paths_rank = filter_tmins_paths(
                    tmin, tmins_files, paths_rank
                )
Jason Reneuve's avatar
Jason Reneuve committed
534

Jason Reneuve's avatar
Jason Reneuve committed
535
536
537
538
539
540
541
542
                npaths = len(paths_rank)
                progress.update(
                    task_files,
                    description=f"Rank {rank:05}...",
                    total=npaths,
                    completed=0,
                )

Jason Reneuve's avatar
Jason Reneuve committed
543
544
545
546
547
548
549
                # for a given rank, paths are sorted by time
                for ip, path_file in enumerate(paths_rank):
                    # break after the last useful file
                    if tmins_files[ip] > tmax:
                        progress.update(task_files, completed=npaths)
                        break

550
                    with open_patient(path_file, "r") as file:
Jason Reneuve's avatar
Jason Reneuve committed
551
552
                        # time indices
                        times_file = file["times"][:]
Jason Reneuve's avatar
Jason Reneuve committed
553
554
555
556
                        its_file = get_arange_minmax(times_file, tmin, tmax)
                        tmin_keep = times_file[its_file[0]]
                        tmax_keep = times_file[its_file[-1]]
                        its = get_arange_minmax(times, tmin_keep, tmax_keep)
Jason Reneuve's avatar
Jason Reneuve committed
557
558
559
560

                        # k_adim_loc = global probes indices!
                        ik0 = file["probes_k0adim_loc"][:]
                        ik1 = file["probes_k1adim_loc"][:]
Pierre Augier's avatar
Pierre Augier committed
561
562
                        if self.nb_dim == 3:
                            ik2 = file["probes_k2adim_loc"][:]
Jason Reneuve's avatar
Jason Reneuve committed
563
564

                        # load data at desired times for all keys_fields
565
                        for key in keys:
566
                            skey = key + "_Fourier"
Jason Reneuve's avatar
Jason Reneuve committed
567
                            data = file[skey + "_loc"][:, its_file]
Pierre Augier's avatar
Pierre Augier committed
568
569
570
571
572
573
574

                            if self.nb_dim == 3:
                                for i, it in enumerate(its):
                                    series[skey][ik0, ik1, ik2, it] = data[:, i]
                            else:
                                for i, it in enumerate(its):
                                    series[skey][ik0, ik1, it] = data[:, i]
Jason Reneuve's avatar
Jason Reneuve committed
575
576
577
578
579
580

                    # update rich task
                    progress.update(task_files, advance=1)

                # update rich task
                progress.update(task_ranks, advance=1)
581

Jason Reneuve's avatar
Jason Reneuve committed
582
583
584
        # add Ki_adim arrays, times and dims order
        k0_adim = np.r_[0 : ik0max + 1, ik0min:0]
        k1_adim = np.r_[0 : ik1max + 1, ik1min:0]
Pierre Augier's avatar
Pierre Augier committed
585
586
587
588
589
590
591
592
593

        if self.nb_dim == 3:
            k2_adim = np.r_[0 : ik2max + 1, ik2min:0]
            K0_adim, K1_adim, K2_adim = np.meshgrid(
                k0_adim, k1_adim, k2_adim, indexing="ij"
            )
        else:
            K0_adim, K1_adim = np.meshgrid(k0_adim, k1_adim, indexing="ij")

Jason Reneuve's avatar
Jason Reneuve committed
594
595
596
597
598
        series.update(
            {
                "K0_adim": K0_adim,
                "K1_adim": K1_adim,
                "times": times,
Pierre Augier's avatar
Pierre Augier committed
599
                "dims_order": dims_order,
Jason Reneuve's avatar
Jason Reneuve committed
600
601
            }
        )
Pierre Augier's avatar
Pierre Augier committed
602
603
        if self.nb_dim == 3:
            series["K2_adim"] = K2_adim
Jason Reneuve's avatar
Jason Reneuve committed
604

605
        return series
Jason Reneuve's avatar
Jason Reneuve committed
606

607
608
    def _compute_spectrum(self, data):
        if not hasattr(self, "f_sample"):
609
610
611
            paths = sorted(self.path_dir.glob("rank*.h5"))
            with h5py.File(paths[0], "r") as file:
                self.f_sample = 1.0 / file.attrs["period_save"]
612
613
614
615
616
            self.domega = 2 * pi * self.f_sample / data.shape[-1]

        # TODO: I'm not sure if detrend=False is good in prod, but it's much
        # better for testing
        freq, spectrum = signal.periodogram(
617
618
619
620
            data,
            fs=self.f_sample,
            scaling="spectrum",
            detrend=False,
621
            return_onesided=False,
622
623
624
        )
        return freq, spectrum / self.domega

625
    def compute_spectra(self, tmin=0, tmax=None, dtype=None):
Jason Reneuve's avatar
Jason Reneuve committed
626
627
628
629
630
        """compute spatiotemporal spectra from files"""
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        # load time series as state_spect arrays + times
Jason Reneuve's avatar
fixes    
Jason Reneuve committed
631
        series = self.load_time_series(tmin=tmin, tmax=tmax, dtype=dtype)
Jason Reneuve's avatar
Jason Reneuve committed
632
633

        # compute spectra
634
        print("performing time fft...")
Jason Reneuve's avatar
Jason Reneuve committed
635

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
636
        spectra = {k: v for k, v in series.items() if k.startswith("K")}
Jason Reneuve's avatar
Jason Reneuve committed
637
638

        for key, data in series.items():
639
            if "_Fourier" not in key:
Jason Reneuve's avatar
Jason Reneuve committed
640
                continue
641
            key_spectrum = "spectrum_" + key.split("_Fourier")[0]
642
            freq, spectrum = self._compute_spectrum(data)
643
            spectra[key_spectrum] = spectrum
Jason Reneuve's avatar
Jason Reneuve committed
644

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
645
646
        spectra["omegas"] = 2 * pi * freq
        spectra["dims_order"] = series["dims_order"]
Jason Reneuve's avatar
Jason Reneuve committed
647

Jason Reneuve's avatar
fixes    
Jason Reneuve committed
648
        return spectra
Pierre Augier's avatar
Pierre Augier committed
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676

    def _get_data_probe_from_field(self, field):
        return field[
            self.probes_ik0_loc,
            self.probes_ik1_loc,
            self.probes_ik2_loc,
        ]


class SpatioTemporalSpectra2D(SpatioTemporalSpectra3D):
    nb_dim = 2

    def _get_data_probe_from_field(self, field):
        return field[
            self.probes_ik0_loc,
            self.probes_ik1_loc,
        ]


class SpatioTemporalSpectraNS:
    def _get_path_saved_spectra(self, tmin, tmax, dtype, save_urud):
        base = f"periodogram_{tmin}_{tmax}"
        if dtype is not None:
            base += f"_{dtype}"
        if save_urud:
            base += "_urud"
        return self.path_dir / (base + ".h5")

Jason Reneuve's avatar
Jason Reneuve committed
677
678
679
680
681
682
683
684
    def _get_path_saved_tspectra(self, tmin, tmax, dtype, save_urud):
        base = f"periodogram_temporal_{tmin}_{tmax}"
        if dtype is not None:
            base += f"_{dtype}"
        if save_urud:
            base += "_urud"
        return self.path_dir / (base + ".h5")

Pierre Augier's avatar
Pierre Augier committed
685
686
687
    def save_spectra_kzkhomega(
        self, tmin=0, tmax=None, dtype=None, save_urud=False
    ):
Jason Reneuve's avatar
Jason Reneuve committed
688
689
690
691
692
        """
        save:
            - the spatiotemporal spectra, with a cylindrical average in k-space
            - the temporal spectra, with an average on the whole k-space
        """
Pierre Augier's avatar
Pierre Augier committed
693
694
695
696
697
698
699
700
701
702
703
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        # compute spectra
        print("Computing spectra...")
        spectra = self.compute_spectra(tmin=tmin, tmax=tmax, dtype=dtype)

        # get kz, kh
        params_oper = self.sim.params.oper
        deltakx = 2 * pi / params_oper.Lx
        order = spectra["dims_order"]
704
        KX = deltakx * spectra[f"K{order[-1]}_adim"]
Jason Reneuve's avatar
Jason Reneuve committed
705
        kx_max = self.sim.params.oper.nx // 2 * deltakx
706
707
708
709
710
711
712

        if self.nb_dim == 3:
            deltaky = 2 * pi / params_oper.Ly
            deltakz = 2 * pi / params_oper.Lz
            KY = deltaky * spectra[f"K{order[1]}_adim"]
            KH = np.sqrt(KX ** 2 + KY ** 2)
            deltakh = max(deltakx, deltaky)
713
            khmax_spectra = max(KX.max(), KY.max())
714
715
716
717
718
719
720
721
            del KY
        else:
            # in 2d, vertical (here "z") is y
            deltakz = 2 * pi / params_oper.Ly
            KH = abs(KX)
            deltakh = deltakx
            khmax_spectra = KX.max()

Pierre Augier's avatar
Pierre Augier committed
722
723
724
725
726
727
728
729
730
        KZ = deltakz * spectra[f"K{order[0]}_adim"]

        kz_spectra = np.arange(0, KZ.max() + 1e-15, deltakz)

        nkh_spectra = max(2, int(khmax_spectra / deltakh))
        kh_spectra = deltakh * np.arange(nkh_spectra)

        # get one-sided frequencies
        omegas = spectra["omegas"]
731
        nomegas = (omegas.size + 1) // 2
Pierre Augier's avatar
Pierre Augier committed
732
733
        omegas_onesided = abs(omegas[:nomegas])

Jason Reneuve's avatar
Jason Reneuve committed
734
735
        # kzkhomega : perform cylindrical average
        # temporal spectra : average on Fourier space
Pierre Augier's avatar
Pierre Augier committed
736
737
738
739
740
        spectra_kzkhomega = {
            "kz_spectra": kz_spectra,
            "kh_spectra": kh_spectra,
            "omegas": omegas_onesided,
        }
Jason Reneuve's avatar
Jason Reneuve committed
741
        tspectra = {"omegas": omegas_onesided}
Pierre Augier's avatar
Pierre Augier committed
742
743
744
745
746
747
        for key, data in spectra.items():
            if not key.startswith("spectrum_"):
                continue
            spectra_kzkhomega[key] = self.compute_spectrum_kzkhomega(
                np.ascontiguousarray(data), kh_spectra, kz_spectra, KX, KZ, KH
            )
Jason Reneuve's avatar
Jason Reneuve committed
748
749
750
751
752
753
754
755
            tspectrum = self._sum_wavenumber(data, KX, kx_max)
            # one-sided frequencies
            tspectrum_onesided = np.zeros(nomegas)
            tspectrum_onesided[0] = tspectrum[0]
            tspectrum_onesided[1:] = (
                tspectrum[1:nomegas] + tspectrum[-1:-nomegas:-1]
            )
            tspectra[key] = tspectrum_onesided
Pierre Augier's avatar
Pierre Augier committed
756
757
758
759

        del spectra

        # total kinetic energy
760
761
762
763
764
765
        if self.nb_dim == 3:
            spectra_kzkhomega["spectrum_K"] = 0.5 * (
                spectra_kzkhomega["spectrum_vx"]
                + spectra_kzkhomega["spectrum_vy"]
                + spectra_kzkhomega["spectrum_vz"]
            )
Jason Reneuve's avatar
Jason Reneuve committed
766
767
768
769
770
            tspectra["spectrum_K"] = 0.5 * (
                tspectra["spectrum_vx"]
                + tspectra["spectrum_vy"]
                + tspectra["spectrum_vz"]
            )
771
772
773
774
775
        else:
            spectra_kzkhomega["spectrum_K"] = 0.5 * (
                spectra_kzkhomega["spectrum_ux"]
                + spectra_kzkhomega["spectrum_uy"]
            )
Jason Reneuve's avatar
Jason Reneuve committed
776
777
778
            tspectra["spectrum_K"] = 0.5 * (
                tspectra["spectrum_ux"] + tspectra["spectrum_uy"]
            )
Pierre Augier's avatar
Pierre Augier committed
779
780
781
782
783
784
785

        # potential energy
        try:
            N = self.sim.params.N
            spectra_kzkhomega["spectrum_A"] = (
                0.5 / N ** 2 * spectra_kzkhomega["spectrum_b"]
            )
Jason Reneuve's avatar
Jason Reneuve committed
786
            tspectra["spectrum_A"] = 0.5 / N ** 2 * tspectra["spectrum_b"]
Pierre Augier's avatar
Pierre Augier committed
787
788
789
        except AttributeError:
            pass

Jason Reneuve's avatar
Jason Reneuve committed
790
        # save to files
Pierre Augier's avatar
Pierre Augier committed
791
792
793
794
795
796
797
        path_file = self._get_path_saved_spectra(tmin, tmax, dtype, save_urud)
        with h5py.File(path_file, "w") as file:
            file.attrs["tmin"] = tmin
            file.attrs["tmax"] = tmax
            for key, val in spectra_kzkhomega.items():
                file.create_dataset(key, data=val)

Jason Reneuve's avatar
Jason Reneuve committed
798
799
800
801
802
803
804
        path_file = self._get_path_saved_tspectra(tmin, tmax, dtype, save_urud)
        with h5py.File(path_file, "w") as file:
            file.attrs["tmin"] = tmin
            file.attrs["tmax"] = tmax
            for key, val in tspectra.items():
                file.create_dataset(key, data=val)

Pierre Augier's avatar
Pierre Augier committed
805
806
807
808
        # toroidal/poloidal decomposition
        if save_urud:
            print("Computing ur, ud spectra...")
            spectra_urud_kzkhomega = {}
Jason Reneuve's avatar
Jason Reneuve committed
809
            tspectra_urud = {}
Pierre Augier's avatar
Pierre Augier committed
810
811
812
813
814
815
816
817
818
            spectra = self.compute_spectra_urud(tmin=tmin, tmax=tmax, dtype=dtype)

            for key, data in spectra.items():
                if not key.startswith("spectrum_"):
                    continue
                spectra_urud_kzkhomega[key] = self.compute_spectrum_kzkhomega(
                    np.ascontiguousarray(data), kh_spectra, kz_spectra, KX, KZ, KH
                )
                spectra_kzkhomega[key] = spectra_urud_kzkhomega[key]
Jason Reneuve's avatar
Jason Reneuve committed
819
820
821
822
823
824
825
826
827
                tspectrum = self._sum_wavenumber(data, KX, kx_max)
                # one-sided frequencies
                tspectrum_onesided = np.zeros(nomegas)
                tspectrum_onesided[0] = tspectrum[0]
                tspectrum_onesided[1:] = (
                    tspectrum[1:nomegas] + tspectrum[-1:-nomegas:-1]
                )
                tspectra_urud[key] = tspectrum_onesided
                tspectra[key] = tspectra_urud[key]
Pierre Augier's avatar
Pierre Augier committed
828

Jason Reneuve's avatar
Jason Reneuve committed
829
            path_file = self._get_path_saved_spectra(tmin, tmax, dtype, save_urud)
Pierre Augier's avatar
Pierre Augier committed
830
831
832
833
            with h5py.File(path_file, "a") as file:
                for key, val in spectra_urud_kzkhomega.items():
                    file.create_dataset(key, data=val)

Jason Reneuve's avatar
Jason Reneuve committed
834
835
836
837
838
839
840
841
            path_file = self._get_path_saved_tspectra(
                tmin, tmax, dtype, save_urud
            )
            with h5py.File(path_file, "a") as file:
                for key, val in tspectra_urud.items():
                    file.create_dataset(key, data=val)

        return spectra_kzkhomega, tspectra
Pierre Augier's avatar
Pierre Augier committed
842

Jason Reneuve's avatar
Jason Reneuve committed
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
    def load_spectra_kzkhomega(
        self, tmin=0, tmax=None, dtype=None, save_urud=False
    ):
        """load kzkhomega spectra from file"""
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        spectra = {}

        path_file = self._get_path_saved_spectra(tmin, tmax, dtype, save_urud)
        with h5py.File(path_file, "r") as file:
            for key in file.keys():
                spectra[key] = file[key][...]

        return spectra

Pierre Augier's avatar
Pierre Augier committed
859
860
861
862
863
864
865
866
867
868
869
870
871
    def plot_kzkhomega(
        self,
        key_field=None,
        tmin=0,
        tmax=None,
        dtype=None,
        equation=None,
        xmax=None,
        ymax=None,
        cmap="viridis",
        vmin=None,
        vmax=None,
    ):
872
873
874
875
876
        """plot the spatiotemporal spectra, with a cylindrical average in k-space

        equation must start with 'omega=', 'kh=', 'kz=', 'ikh=' or 'ikz='.

        """
Jason Reneuve's avatar
Jason Reneuve committed
877
        keys_plot = self.keys_fields.copy()
878
879
880
        if self.nb_dim == 3:
            keys_plot.extend(["Khd", "Khr", "Kp"])

Pierre Augier's avatar
Pierre Augier committed
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
        if key_field is None:
            key_field = keys_plot[0]
        if key_field not in keys_plot:
            raise KeyError(f"possible keys are {keys_plot}")
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        key_spect = "spectrum_" + key_field
        if key_field.startswith("Kh") or key_field.startswith("Kp"):
            save_urud = True
        else:
            save_urud = False

        path_file = self._get_path_saved_spectra(tmin, tmax, dtype, save_urud)
        path_urud = self._get_path_saved_spectra(tmin, tmax, dtype, True)
        if path_urud.exists() and not path_file.exists():
            path_file = path_urud

        spectra_kzkhomega = {}

        # compute and save spectra if needed
        if not path_file.exists():
903
904
905
906
907
908
            if self.nb_dim == 3:
                self.save_spectra_kzkhomega(
                    tmin=tmin, tmax=tmax, dtype=dtype, save_urud=save_urud
                )
            else:
                self.save_spectra_kzkhomega(tmin=tmin, tmax=tmax, dtype=dtype)
Pierre Augier's avatar
Pierre Augier committed
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929

        # load spectrum
        with h5py.File(path_file, "r") as file:
            if key_spect.startswith("spectrum_Kp"):
                spectrum = file["spectrum_Khd"][:] + 0.5 * file["spectrum_vz"][:]
            else:
                spectrum = file[key_spect][...]
            if dtype == "complex64":
                float_dtype = "float32"
            elif dtype == "complex128":
                float_dtype = "float64"
            if dtype:
                spectrum = spectrum.astype(float_dtype)
            spectra_kzkhomega[key_spect] = spectrum
            spectra_kzkhomega["kh_spectra"] = file["kh_spectra"][...]
            spectra_kzkhomega["kz_spectra"] = file["kz_spectra"][...]
            spectra_kzkhomega["omegas"] = file["omegas"][...]

        # slice along equation
        if equation is None:
            equation = f"omega=0"
930
931
932
933
934
935
936
937
938
939
940
        elif equation.startswith("kh="):
            kh = eval(equation[len("kh=") :])
            kh_spectra = spectra_kzkhomega["kh_spectra"]
            ikh = abs(kh_spectra - kh).argmin()
            equation = f"ikh={ikh}"
        elif equation.startswith("kz="):
            kz = eval(equation[len("kz=") :])
            kz_spectra = spectra_kzkhomega["kz_spectra"]
            ikz = abs(kz_spectra - kz).argmin()
            equation = f"ikz={ikz}"

Pierre Augier's avatar
Pierre Augier committed
941
        if equation.startswith("omega="):
942
943
944
945
946
            try:
                variables = dict(N=self.sim.params.N)
            except AttributeError:
                variables = dict()
            omega = eval(equation[len("omega=") :], variables)
Pierre Augier's avatar
Pierre Augier committed
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
            omegas = spectra_kzkhomega["omegas"]
            iomega = abs(omegas - omega).argmin()
            spect = spectra_kzkhomega[key_spect][:, :, iomega]
            xaxis = np.arange(spectra_kzkhomega["kh_spectra"].size)
            yaxis = np.arange(spectra_kzkhomega["kz_spectra"].size)
            xlabel = r"$k_h/\delta k_h$"
            ylabel = r"$k_z/\delta k_z$"
            omega = omegas[iomega]
            equation = r"$\omega=$" + f"{omega:.2g}"
            # use reduced frequency for stratified fluids
            try:
                N = self.sim.params.N
                equation = r"$\omega/N=$" + f"{omega/N:.2g}"
            except AttributeError:
                pass
        elif equation.startswith("ikh="):
            ikh = eval(equation[len("ikh=") :])
            kh_spectra = spectra_kzkhomega["kh_spectra"]
            spect = spectra_kzkhomega[key_spect][:, ikh, :].transpose()

            xaxis = np.arange(spectra_kzkhomega["kz_spectra"].size)
            yaxis = spectra_kzkhomega["omegas"]
            # use reduced frequency for stratified fluids
            try:
                N = self.sim.params.N
                yaxis /= N
            except AttributeError:
                pass

            xlabel = r"$k_z/\delta k_z$"
            ylabel = r"$\omega/N$"
            kh = kh_spectra[ikh]
            equation = f"$k_h = {ikh}\\delta k_h = {kh:.2g}$"
        elif equation.startswith("ikz="):
            ikz = eval(equation[len("ikz=") :])
            kz_spectra = spectra_kzkhomega["kz_spectra"]
            spect = spectra_kzkhomega[key_spect][ikz, :, :].transpose()

            xaxis = np.arange(spectra_kzkhomega["kh_spectra"].size)
            yaxis = spectra_kzkhomega["omegas"]
            # use reduced frequency for stratified fluids
            try:
                N = self.sim.params.N
                yaxis /= N
            except AttributeError:
                pass

            xlabel = r"$k_h/\delta k_h$"
            ylabel = r"$\omega/N$"
            kz = kz_spectra[ikz]
            equation = f"$k_z = {ikz}\\delta k_z = {kz:.2g}$"
        else:
            raise NotImplementedError(
                "equation must start with 'omega=', 'kh=', 'kz=', 'ikh=' or 'ikz='"
            )

        # plot
        fig, ax = self.output.figure_axe()
        ax.set_xlabel(xlabel)
        ax.set_ylabel(ylabel)

        # no log(0)
        spect += 1e-15

        im = ax.pcolormesh(
            xaxis,
            yaxis,
            np.log10(spect),
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            shading="nearest",
        )
        fig.colorbar(im)

        ax.set_title(
            f"{key_field} spatiotemporal spectra {equation}\n"
            f"tmin={tmin:.3f}, tmax={tmax:.3f}\n" + self.output.summary_simul
        )

        # add dispersion relation : omega = N * kh / sqrt(kh ** 2 + kz ** 2)
        try:
            N = self.sim.params.N
        except AttributeError:
            return
        dkh_over_dkz = (
            spectra_kzkhomega["kz_spectra"][1]
            / spectra_kzkhomega["kh_spectra"][1]
        )
        if equation.startswith(r"$\omega"):
1037
1038
1039
            if omega > 0:
                ikz_disp = np.sqrt(N ** 2 / omega ** 2 - 1) / dkh_over_dkz * xaxis
                ax.plot(xaxis, ikz_disp, "k+", linewidth=2)
Pierre Augier's avatar
Pierre Augier committed
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
        elif equation.startswith(r"$k_h"):
            omega_disp = ikh / np.sqrt(ikh ** 2 + dkh_over_dkz ** 2 * xaxis ** 2)
            ax.plot(xaxis, omega_disp, "k+", linewidth=2)
        elif equation.startswith(r"$k_z"):
            omega_disp = xaxis / np.sqrt(
                xaxis ** 2 + dkh_over_dkz ** 2 * ikz ** 2
            )
            ax.plot(xaxis, omega_disp, "k+", linewidth=2)
        else:
            raise ValueError("wrong equation for dispersion relation")

        # set axis limits after plotting dispersion relation
        if xmax is None:
            xmax = xaxis.max()
        if ymax is None:
            ymax = yaxis.max()
        ax.set_xlim((0, xmax))
        ax.set_ylim((0, ymax))

    def compute_spectra_urud(self, tmin=0, tmax=None, dtype=None):
        raise NotImplementedError

    def compute_temporal_spectra(
        self, tmin=0, tmax=None, dtype=None, compute_urud=False
    ):
        """compute the temporal spectra by averaging over Fourier space"""
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        tspectra = {}

        # compute kxkykzomega spectra
        spectra = self.compute_spectra(tmin=tmin, tmax=tmax, dtype=dtype)
        if compute_urud:
            spectra.update(
                self.compute_spectra_urud(tmin=tmin, tmax=tmax, dtype=dtype)
            )

Jason Reneuve's avatar
Jason Reneuve committed
1078
        # one-sided frequencies
1079
        nomegas = (spectra["omegas"].size + 1) // 2
Jason Reneuve's avatar
Jason Reneuve committed
1080
1081
        tspectra["omegas"] = spectra["omegas"][:nomegas]

Pierre Augier's avatar
Pierre Augier committed
1082
        order = spectra["dims_order"]
Pierre Augier's avatar
Pierre Augier committed
1083
        KX = spectra[f"K{order[-1]}_adim"]
Pierre Augier's avatar
Pierre Augier committed
1084
1085
1086
1087
1088
1089
1090
        deltakx = 2 * pi / self.sim.params.oper.Lx
        kx_max = self.sim.params.oper.nx // 2 * deltakx

        # average over Fourier space (kx,ky,kz)
        for key, spectrum in spectra.items():
            if not key.startswith("spectrum_"):
                continue
Jason Reneuve's avatar
Jason Reneuve committed
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
            tspectrum = self._sum_wavenumber(spectrum, KX, kx_max)
            # one-sided frequencies
            tspectrum_onesided = np.zeros(nomegas)
            tspectrum_onesided[0] = tspectrum[0]
            tspectrum_onesided[1:] = (
                tspectrum[1:nomegas] + tspectrum[-1:-nomegas:-1]
            )
            tspectra[key] = tspectrum_onesided

        # total kinetic energy
        if self.nb_dim == 3:
            tspectra["spectrum_K"] = 0.5 * (
                tspectra["spectrum_vx"]
                + tspectra["spectrum_vy"]
                + tspectra["spectrum_vz"]
            )
        else:
            tspectra["spectrum_K"] = 0.5 * (
                tspectra["spectrum_ux"] + tspectra["spectrum_uy"]
            )
Pierre Augier's avatar
Pierre Augier committed
1111

Jason Reneuve's avatar
Jason Reneuve committed
1112
1113
1114
1115
1116
1117
        # potential energy
        try:
            N = self.sim.params.N
            tspectra["spectrum_A"] = 0.5 / N ** 2 * tspectra["spectrum_b"]
        except AttributeError:
            pass
Pierre Augier's avatar
Pierre Augier committed
1118
1119

        return tspectra
Jason Reneuve's avatar
Jason Reneuve committed
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138

    def plot_temporal_spectra(
        self,
        key_field=None,
        tmin=0,
        tmax=None,
        dtype=None,
    ):
        """plot the temporal spectra computed from the 4d spectra"""
        keys_plot = self.keys_fields.copy()
        if self.nb_dim == 3:
            keys_plot.extend(["Khd", "Khr", "Kp"])
        if key_field is None:
            key_field = keys_plot[0]
        if key_field not in keys_plot:
            raise KeyError(f"possible keys are {keys_plot}")
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

1139
        if self.nb_dim == 3:
1140
            # much simpler for 3d
1141
1142
1143
            save_urud = True
        else:
            save_urud = False
Jason Reneuve's avatar
Jason Reneuve committed
1144
1145
1146
1147
1148
1149
1150
1151
1152
        path_file = self._get_path_saved_tspectra(tmin, tmax, dtype, save_urud)
        if path_file.exists():
            tspectra = self.load_temporal_spectra(
                tmin=tmin, tmax=tmax, dtype=dtype, save_urud=save_urud
            )
        else:
            tspectra = self.save_temporal_spectra(
                tmin=tmin, tmax=tmax, dtype=dtype, save_urud=save_urud
            )
Jason Reneuve's avatar
Jason Reneuve committed
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176

        # plot
        fig, ax = self.output.figure_axe()
        ax.set_xlabel(r"$\omega$")
        ax.set_ylabel("spectrum")
        ax.set_title(
            f"{key_field} temporal spectrum (tmin={tmin:.3f}, tmax={tmax:.3f})\n"
            + self.output.summary_simul
        )
        ax.set_xscale("log")
        ax.set_yscale("log")

        # specific to strat
        try:
            N = self.sim.params.N
        except AttributeError:
            ax.plot(
                tspectra["omegas"],
                tspectra["spectrum_" + key_field],
                "k",
                linewidth=2,
            )
        else:
            omegas = tspectra["omegas"] / N
1177
1178
1179
1180
            if self.nb_dim == 3:
                # polo/toro/potential decomposition
                EKp = tspectra["spectrum_Khd"] + 0.5 * tspectra["spectrum_vz"]
                EKhr = tspectra["spectrum_Khr"]
1181
1182
1183
1184
                EK = EKhr + EKp
                ax.plot(omegas, EK, "r", linewidth=2, label=r"$E_K$")
                ax.plot(omegas, EKhr, "r--", linewidth=1, label=r"$E_{K,toro}$")
                ax.plot(omegas, EKp, "r-.", linewidth=1, label=r"$E_{K,polo}$")
1185
1186
1187
1188
            else:
                # kinetic energy
                EK = tspectra["spectrum_K"]
                ax.plot(omegas, EK, "r", linewidth=2, label=r"$E_K$")
1189
            EK_N = EK[abs(omegas - 1).argmin()]  # value at N
1190
            EA = tspectra["spectrum_A"]
Jason Reneuve's avatar
Jason Reneuve committed
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209

            ax.plot(omegas, EA, "b", linewidth=2, label=r"$E_A$")
            ax.set_title(
                f"kinetic/potential energy spectrum (tmin={tmin:.3f}, tmax={tmax:.3f})\n"
                + self.output.summary_simul
            )

            # resonant modes
            if self.nb_dim == 3:
                aspect_ratio = self.sim.oper.Lx / self.sim.oper.Lz
            else:
                aspect_ratio = self.sim.oper.Lx / self.sim.oper.Ly

            def modes(nx, nz):
                return np.sqrt(nx ** 2 / (nx ** 2 + aspect_ratio ** 2 * nz ** 2))

            nxs = np.arange(1, 11)
            modes_nz1 = modes(nxs, 1)
            modes_nz2 = modes(nxs, 2)
1210
            modes_y = np.full_like(modes_nz1, fill_value=10 * EK_N)
Jason Reneuve's avatar
Jason Reneuve committed
1211
1212
1213
1214
1215
1216

            ax.plot(modes_nz1, modes_y, "o", label="modes $n_z=1$")
            ax.plot(modes_nz2, modes_y * 3, "o", label="modes $n_z=2$")

            # omega^-2 scaling
            omegas_scaling = np.arange(0.4, 1 + 1e-15, 0.01)
1217
            scaling_y = EK_N * omegas_scaling ** -2
Jason Reneuve's avatar
Jason Reneuve committed
1218
1219
1220

            ax.plot(omegas_scaling, scaling_y, "k--")

1221
            # eye guide at N
1222
            ax.axvline(1, linestyle="dotted")
Jason Reneuve's avatar
Jason Reneuve committed
1223

1224
1225
1226
1227
1228
1229
1230
1231
            # eye guide at omega_f (specific to some forcing types)
            forcing_type = self.sim.params.forcing.type
            if forcing_type in ["watu_coriolis", "milestone"]:
                if forcing_type == "watu_coriolis":
                    omega_f = self.sim.params.forcing.watu_coriolis.omega_f
                elif forcing_type == "milestone":
                    period = self.sim.forcing.get_info()["period"]
                    omega_f = 2 * pi / period
1232
                ax.axvline(omega_f / N, linestyle="dotted")
Jason Reneuve's avatar
Jason Reneuve committed
1233
1234
1235
1236

            ax.set_xlabel(r"$\omega/N$")

            ax.legend()
Jason Reneuve's avatar
Jason Reneuve committed
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272

    def load_temporal_spectra(
        self, tmin=0, tmax=None, dtype=None, save_urud=False
    ):
        """load temporal spectra from file"""
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        tspectra = {}

        path_file = self._get_path_saved_tspectra(tmin, tmax, dtype, save_urud)
        with h5py.File(path_file, "r") as file:
            for key in file.keys():
                tspectra[key] = file[key][...]

        return tspectra

    def save_temporal_spectra(
        self, tmin=0, tmax=None, dtype=None, save_urud=False
    ):
        """compute temporal spectra from files"""
        if tmax is None:
            tmax = self.sim.params.time_stepping.t_end

        tspectra = self.compute_temporal_spectra(
            tmin=tmin, tmax=tmax, dtype=dtype, compute_urud=save_urud
        )

        path_file = self._get_path_saved_tspectra(tmin, tmax, dtype, save_urud)
        with h5py.File(path_file, "w") as file:
            file.attrs["tmin"] = tmin
            file.attrs["tmax"] = tmax
            for key, val in tspectra.items():
                file.create_dataset(key, data=val)

        return tspectra