operators.py 19.4 KB
Newer Older
Pierre Augier's avatar
Pierre Augier committed
1
2
"""Operators 3d (:mod:`fluidfft.fft3d.operators`)
=================================================
Pierre Augier's avatar
Pierre Augier committed
3

Pierre Augier's avatar
Pierre Augier committed
4
5
6
7
8
.. autoclass:: OperatorsPseudoSpectral3D
   :members:
   :undoc-members:

"""
Pierre Augier's avatar
Pierre Augier committed
9
10
11
12
from __future__ import print_function

from math import pi

13
from past.builtins import basestring
Pierre Augier's avatar
Pierre Augier committed
14

15
import numpy as np
16

Pierre Augier's avatar
Pierre Augier committed
17
18
from fluiddyn.util import mpi

19
20
from fluidfft import create_fft_object, empty_aligned
from fluidfft.util import _rescale_random
Pierre Augier's avatar
Pierre Augier committed
21
from fluidfft.fft2d.operators import _make_str_length
Pierre Augier's avatar
Pierre Augier committed
22

23
from .util_pythran import (
Pierre Augier's avatar
Pierre Augier committed
24
25
26
    project_perpk3d,
    divfft_from_vecfft,
    rotfft_from_vecfft,
Pierre Augier's avatar
Pierre Augier committed
27
    rotfft_from_vecfft_outin,
Pierre Augier's avatar
Pierre Augier committed
28
29
30
    vector_product,
    loop_spectra3d,
)
Pierre Augier's avatar
Pierre Augier committed
31

32
# from .dream_pythran import _vgradv_from_v2
Pierre Augier's avatar
Pierre Augier committed
33

Pierre Augier's avatar
Pierre Augier committed
34
35
36
37
if mpi.nb_proc > 1:
    MPI = mpi.MPI


Pierre Augier's avatar
Pierre Augier committed
38
39
40
__all__ = ["vector_product", "OperatorsPseudoSpectral3D"]


Pierre Augier's avatar
Pierre Augier committed
41
class OperatorsPseudoSpectral3D(object):
Pierre Augier's avatar
Pierre Augier committed
42
    """Perform 2D FFT and operations on data.
Pierre Augier's avatar
Pierre Augier committed
43

Pierre Augier's avatar
Pierre Augier committed
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
    Parameters
    ----------

    nx : int

      Global dimension over the x-axis (third dimension for the real arrays).

    ny : int

      Global dimension over the y-axis (second dimension for the real arrays).

    nz : int

      Global dimension over the y-axis (first dimension for the real arrays).

    lx : float

      Length of the domain along the x-axis.

    ly : float

      Length of the domain along the y-axis.

    lz : float

      Length of the domain along the z-axis.

    fft : str or FFT classes

      Name of module or string characterizing a method. It has to correspond to
Pierre Augier's avatar
Pierre Augier committed
74
      a module of fluidfft. The first part "fluidfft." of the module "path" can
Pierre Augier's avatar
Pierre Augier committed
75
76
77
78
79
      be omitted.

    coef_dealiasing : float

    """
Pierre Augier's avatar
Pierre Augier committed
80
81

    def __init__(self, nx, ny, nz, lx, ly, lz, fft=None, coef_dealiasing=1.):
Pierre Augier's avatar
Pierre Augier committed
82
83
84
85
        self.nx = self.nx_seq = nx
        self.ny = self.ny_seq = ny
        self.nz = self.nz_seq = nz

Pierre Augier's avatar
Pierre Augier committed
86
87
        if fft is None:
            if mpi.nb_proc == 1:
Pierre Augier's avatar
Pierre Augier committed
88
                fft = "fft3d.with_pyfftw"
Pierre Augier's avatar
Pierre Augier committed
89
            else:
Pierre Augier's avatar
Pierre Augier committed
90
                fft = "fft3d.mpi_with_fftwmpi3d"
Pierre Augier's avatar
Pierre Augier committed
91

Pierre Augier's avatar
Pierre Augier committed
92
        if isinstance(fft, basestring):
Pierre Augier's avatar
Pierre Augier committed
93
94
95
            if fft.lower() in ("sequential", "fftwpy"):
                fft = "fft3d.with_pyfftw"
            if any([fft.startswith(s) for s in ["fluidfft.", "fft3d."]]):
Pierre Augier's avatar
Pierre Augier committed
96
97
                op_fft = create_fft_object(fft, nz, ny, nx)
            else:
98
                raise ValueError(
Pierre Augier's avatar
Pierre Augier committed
99
100
101
102
103
                    "Cannot instantiate {}.".format(fft)
                    + " Expected something like 'fftwpy'"
                    " or 'fluidfft.fft3d.<method>' or 'fft3d.<method>'"
                )

Pierre Augier's avatar
Pierre Augier committed
104
105
106
107
108
109
110
111
        elif isinstance(fft, type):
            op_fft = fft(nz, ny, nx)
        else:
            op_fft = fft

        self._op_fft = op_fft
        self.type_fft = op_fft.__class__.__module__

Pierre Augier's avatar
Pierre Augier committed
112
113
114
115
116
        try:
            self.dim_first_fft = op_fft.get_dim_first_fft()
        except AttributeError:
            self.dim_first_fft = 2

Pierre Augier's avatar
Pierre Augier committed
117
118
119
        self.shapeX_seq = op_fft.get_shapeX_seq()
        self.shapeX_loc = op_fft.get_shapeX_loc()

Pierre Augier's avatar
Pierre Augier committed
120
        self._is_mpi_lib = self.shapeX_seq != self.shapeX_loc and mpi.nb_proc > 1
Pierre Augier's avatar
Pierre Augier committed
121

Pierre Augier's avatar
Pierre Augier committed
122
123
124
125
126
127
128
129
        Lx = self.Lx = float(lx)
        Ly = self.Ly = float(ly)
        Lz = self.Lz = float(lz)

        self.deltax = Lx / nx
        self.deltay = Ly / ny
        self.deltaz = Lz / nz

Pierre Augier's avatar
Pierre Augier committed
130
131
132
        self.x_seq = self.deltax * np.arange(nx)
        self.y_seq = self.deltay * np.arange(ny)
        self.z_seq = self.deltaz * np.arange(nz)
Pierre Augier's avatar
Pierre Augier committed
133

Pierre Augier's avatar
Pierre Augier committed
134
135
136
        self.deltakx = deltakx = 2 * pi / Lx
        self.deltaky = deltaky = 2 * pi / Ly
        self.deltakz = deltakz = 2 * pi / Lz
Pierre Augier's avatar
Pierre Augier committed
137

Pierre Augier's avatar
Pierre Augier committed
138
139
140
141
142
143
        self.ifft = self.ifft3d = op_fft.ifft
        self.fft = self.fft3d = op_fft.fft

        self.ifft_as_arg = op_fft.ifft_as_arg
        self.fft_as_arg = op_fft.fft_as_arg

144
145
146
147
148
        # try:
        # faster version which destroy the input
        self.ifft_as_arg_destroy = op_fft.ifft_as_arg_destroy
        # except AttributeError:
        #     self.ifft_as_arg_destroy = self.ifft_as_arg
149

Pierre Augier's avatar
Pierre Augier committed
150
        self.sum_wavenumbers = op_fft.sum_wavenumbers
151
152
        self.compute_energy_from_X = op_fft.compute_energy_from_X
        self.compute_energy_from_K = op_fft.compute_energy_from_K
Pierre Augier's avatar
Pierre Augier committed
153

Pierre Augier's avatar
Pierre Augier committed
154
        self.shapeK = self.shapeK_loc = op_fft.get_shapeK_loc()
Pierre Augier's avatar
Pierre Augier committed
155
        self.shapeK_seq = op_fft.get_shapeK_seq()
Pierre Augier's avatar
Pierre Augier committed
156
157
158
159
160
161
162
163
164
        self.nk0, self.nk1, self.nk2 = self.shapeK_loc

        order = op_fft.get_dimX_K()
        if order == (0, 1, 2):
            self.deltaks = deltakz, deltaky, deltakx
        elif order == (1, 0, 2):
            self.deltaks = deltaky, deltakz, deltakx
        elif order == (2, 1, 0):
            self.deltaks = deltakx, deltaky, deltakz
165
166
        elif order == (1, 2, 0):
            self.deltaks = deltaky, deltakx, deltakz
Pierre Augier's avatar
Pierre Augier committed
167
        else:
Pierre Augier's avatar
Pierre Augier committed
168
            print("order =", order)
Pierre Augier's avatar
Pierre Augier committed
169
170
            raise NotImplementedError

171
172
173
174
        for self.dimK_first_fft in range(3):
            if order[self.dimK_first_fft] == self.dim_first_fft:
                break

Pierre Augier's avatar
Pierre Augier committed
175
176
177
178
179
180
181
182
183
        k0_adim_loc, k1_adim_loc, k2_adim_loc = op_fft.get_k_adim_loc()

        self.k0 = self.deltaks[0] * k0_adim_loc
        self.k1 = self.deltaks[1] * k1_adim_loc
        self.k2 = self.deltaks[2] * k2_adim_loc

        # oh that's strange!
        K1, K0, K2 = np.meshgrid(self.k1, self.k0, self.k2, copy=False)

184
185
186
187
        K0 = np.ascontiguousarray(K0)
        K1 = np.ascontiguousarray(K1)
        K2 = np.ascontiguousarray(K2)

188
        assert K0.shape == self.shapeK_loc, (K0.shape, self.shapeK_loc)
Pierre Augier's avatar
Pierre Augier committed
189

Pierre Augier's avatar
Pierre Augier committed
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
        if order == (0, 1, 2):
            self.Kz = K0
            self.Ky = K1
            self.Kx = K2
        elif order == (1, 0, 2):
            self.Ky = K0
            self.Kz = K1
            self.Kx = K2
        elif order == (2, 1, 0):
            self.Kx = K0
            self.Ky = K1
            self.Kz = K2
        elif order == (1, 2, 0):
            self.Ky = K0
            self.Kx = K1
            self.Kz = K2
        else:
Pierre Augier's avatar
Pierre Augier committed
207
            print("order =", order)
Pierre Augier's avatar
Pierre Augier committed
208
            raise NotImplementedError
Pierre Augier's avatar
Pierre Augier committed
209

Pierre Augier's avatar
Pierre Augier committed
210
211
        self.K2 = K0 ** 2 + K1 ** 2 + K2 ** 2
        self.K8 = self.K2 ** 4
Pierre Augier's avatar
Pierre Augier committed
212

213
        self.seq_indices_first_K = op_fft.get_seq_indices_first_K()
214
        self.seq_indices_first_X = op_fft.get_seq_indices_first_X()
Pierre Augier's avatar
Pierre Augier committed
215

Pierre Augier's avatar
Pierre Augier committed
216
        K_square_nozero = self.K2.copy()
Pierre Augier's avatar
Pierre Augier committed
217

218
        if all(index == 0 for index in self.seq_indices_first_K):
Pierre Augier's avatar
Pierre Augier committed
219
220
            K_square_nozero[0, 0, 0] = 1e-14

Pierre Augier's avatar
Pierre Augier committed
221
        self.inv_K_square_nozero = 1. / K_square_nozero
Pierre Augier's avatar
Pierre Augier committed
222
223
224

        self.coef_dealiasing = coef_dealiasing

225
226
227
        CONDKX = abs(self.Kx) >= self.coef_dealiasing * self.k2.max()
        CONDKY = abs(self.Ky) >= self.coef_dealiasing * self.k1.max()
        CONDKZ = abs(self.Kz) >= self.coef_dealiasing * self.k0.max()
Pierre Augier's avatar
Pierre Augier committed
228
        where_dealiased = np.logical_or(CONDKX, CONDKY, CONDKZ)
229
        self.where_dealiased = np.array(where_dealiased, dtype=np.uint8)
230
231
232
233

        self.gather_Xspace = op_fft.gather_Xspace
        self.scatter_Xspace = op_fft.scatter_Xspace

Pierre Augier's avatar
Pierre Augier committed
234
        if mpi.nb_proc > 1:
Pierre Augier's avatar
Pierre Augier committed
235
            self.comm = mpi.comm
Pierre Augier's avatar
Pierre Augier committed
236

Pierre Augier's avatar
Pierre Augier committed
237
238
        self.rank = mpi.rank

Pierre Augier's avatar
Pierre Augier committed
239
        # initialization spectra
Pierre Augier's avatar
Pierre Augier committed
240
241
242
        self.nkx_spectra = nx // 2 + 1
        self.nky_spectra = ny // 2 + 1
        self.nkz_spectra = nz // 2 + 1
Pierre Augier's avatar
Pierre Augier committed
243
244
245
246
247

        self.kxmax_spectra = self.deltakx * self.nkx_spectra
        self.kymax_spectra = self.deltaky * self.nky_spectra
        self.kzmax_spectra = self.deltakz * self.nkz_spectra

Pierre Augier's avatar
Pierre Augier committed
248
        self.deltak = self.deltak_spectra3d = max(
Pierre Augier's avatar
Pierre Augier committed
249
250
            self.deltakx, self.deltaky, self.deltakz
        )
Pierre Augier's avatar
Pierre Augier committed
251
        self.kmax_spectra3d = min(
Pierre Augier's avatar
Pierre Augier committed
252
253
            self.kxmax_spectra, self.kymax_spectra, self.kzmax_spectra
        )
Pierre Augier's avatar
Pierre Augier committed
254
        self.nk_spectra3d = max(
Pierre Augier's avatar
Pierre Augier committed
255
256
            2, int(self.kmax_spectra3d / self.deltak_spectra3d)
        )
Pierre Augier's avatar
Pierre Augier committed
257
258
        self.k_spectra3d = self.deltak_spectra3d * np.arange(self.nk_spectra3d)

Pierre Augier's avatar
Pierre Augier committed
259
    # self.tmp_fields_fft = tuple(self.create_arrayK() for n in range(6))
260

Pierre Augier's avatar
Pierre Augier committed
261
    def produce_str_describing_grid(self):
Pierre Augier's avatar
Pierre Augier committed
262
        """Produce a short string describing the grid."""
Pierre Augier's avatar
Pierre Augier committed
263
        return "{}x{}x{}".format(self.nx_seq, self.ny_seq, self.nz_seq)
264

Pierre Augier's avatar
Pierre Augier committed
265
    def produce_str_describing_oper(self):
Pierre Augier's avatar
Pierre Augier committed
266
        """Produce a short string describing the operator."""
Pierre Augier's avatar
Pierre Augier committed
267
268
269
270
        str_Lx = _make_str_length(self.Lx)
        str_Ly = _make_str_length(self.Ly)
        str_Lz = _make_str_length(self.Lz)

Pierre Augier's avatar
Pierre Augier committed
271
272
273
        return ("{}x{}x{}_V" + str_Lx + "x" + str_Ly + "x" + str_Lz).format(
            self.nx_seq, self.ny_seq, self.nz_seq
        )
Pierre Augier's avatar
Pierre Augier committed
274
275
276
277
278
279
280
281
282

    def produce_long_str_describing_oper(self):
        """Produce a string describing the operator."""

        str_Lx = _make_str_length(self.Lx)
        str_Ly = _make_str_length(self.Ly)
        str_Lz = _make_str_length(self.Lz)

        return (
Pierre Augier's avatar
Pierre Augier committed
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
            "type fft: "
            + self.type_fft
            + "\n"
            + "nx = {:6d} ; ny = {:6d} ; nz = {:6d}\n".format(
                self.nx_seq, self.ny_seq, self.nz_seq
            )
            + "Lx = "
            + str_Lx
            + " ; Ly = "
            + str_Ly
            + " ; Lz = "
            + str_Lz
            + "\n"
        )

    def _get_shapeX(self, shape="loc"):
        if shape.lower() == "loc":
Pierre Augier's avatar
Pierre Augier committed
300
            return self.shapeX_loc
Pierre Augier's avatar
Pierre Augier committed
301
302

        elif shape.lower() == "seq":
Pierre Augier's avatar
Pierre Augier committed
303
            return self.shapeX_seq
Pierre Augier's avatar
Pierre Augier committed
304

Pierre Augier's avatar
Pierre Augier committed
305
        else:
Pierre Augier's avatar
Pierre Augier committed
306
            raise ValueError('shape should be "loc" or "seq"')
Pierre Augier's avatar
Pierre Augier committed
307

Pierre Augier's avatar
Pierre Augier committed
308
309
    def _get_shapeK(self, shape="loc"):
        if shape.lower() == "loc":
310
            return self.shapeK_loc
Pierre Augier's avatar
Pierre Augier committed
311
312

        elif shape.lower() == "seq":
313
            return self.shapeK_seq
Pierre Augier's avatar
Pierre Augier committed
314

315
316
317
        else:
            raise ValueError('shape should be "loc" or "seq"')

Pierre Augier's avatar
Pierre Augier committed
318
    def create_arrayX(self, value=None, shape="loc"):
Pierre Augier's avatar
Pierre Augier committed
319
320
        """Return a constant array in real space."""
        shapeX = self._get_shapeX(shape)
321
322
323
        field = empty_aligned(shapeX)
        if value is not None:
            field.fill(value)
Pierre Augier's avatar
Pierre Augier committed
324
325
        return field

Pierre Augier's avatar
Pierre Augier committed
326
    def create_arrayK(self, value=None, shape="loc"):
327
328
        """Return a constant array in real space."""
        shapeK = self._get_shapeK(shape)
329
330
331
        field = empty_aligned(shapeK, dtype=np.complex128)
        if value is not None:
            field.fill(value)
332
333
        return field

Pierre Augier's avatar
Pierre Augier committed
334
    def create_arrayX_random(self, shape="loc", min_val=None, max_val=None):
Pierre Augier's avatar
Pierre Augier committed
335
        """Return a random array in real space."""
336
        shape = self._get_shapeX(shape)
337
338
        values = np.random.random(shape)
        return _rescale_random(values, min_val, max_val)
339

Pierre Augier's avatar
Pierre Augier committed
340
    def create_arrayK_random(self, shape="loc", min_val=None, max_val=None):
341
342
        """Return a random array in real space."""
        shape = self._get_shapeK(shape)
Pierre Augier's avatar
Pierre Augier committed
343
        values = (np.random.random(shape) + 1j * np.random.random(shape))
344
        return _rescale_random(values, min_val, max_val)
Pierre Augier's avatar
Pierre Augier committed
345

346
347
348
349
350
351
352
    def sum_wavenumbers_versatile(self, field_fft):
        """Compute the sum over all wavenumbers (versatile version).

        This function should return the same result than
        :func:`sum_wavenumbers`.

        It is here mainly to check that the classes are well implemented.
Pierre Augier's avatar
Pierre Augier committed
353

354
        """
Pierre Augier's avatar
Pierre Augier committed
355
356
357
358
359
360
361
362
363
364
365
        spectrum3d_loc = self._compute_spectrum3d_loc(field_fft)
        result = spectrum3d_loc.sum()

        if self._is_mpi_lib:
            result = mpi.comm.allreduce(result, op=mpi.MPI.SUM)

        return result

    def _compute_spectrum3d_loc(self, field_fft):
        """"""

366
        dimK_first_fft = self.dimK_first_fft
Pierre Augier's avatar
Pierre Augier committed
367
368

        nx_seq = self.shapeX_seq[self.dim_first_fft]
369
        # nk_seq = self.shapeK_seq[dimK_first_fft]
Pierre Augier's avatar
Pierre Augier committed
370
371
372
373
374
        nk_loc = self.shapeK_loc[dimK_first_fft]
        ik_start = self.seq_indices_first_K[dimK_first_fft]
        ik_stop = ik_start + nk_loc

        # the copy is important: no *= !
Pierre Augier's avatar
Pierre Augier committed
375
        field_fft = 2 * field_fft
Pierre Augier's avatar
Pierre Augier committed
376
377
378
379
380
381

        if ik_start == 0:
            if dimK_first_fft == 2:
                slice0 = np.s_[:, :, 0]
            elif dimK_first_fft == 0:
                slice0 = np.s_[0, :, :]
Pierre Augier's avatar
Pierre Augier committed
382
383
            elif dimK_first_fft == 1:
                slice0 = np.s_[:, 0, :]
Pierre Augier's avatar
Pierre Augier committed
384
385
            else:
                raise NotImplementedError
Pierre Augier's avatar
Pierre Augier committed
386

Pierre Augier's avatar
Pierre Augier committed
387
            field_fft[slice0] /= 2
Pierre Augier's avatar
Pierre Augier committed
388
        if ik_stop == nx_seq // 2 + 1 and nx_seq % 2 == 0:
Pierre Augier's avatar
Pierre Augier committed
389
390
391
392
            if dimK_first_fft == 2:
                slice_last = np.s_[:, :, -1]
            elif dimK_first_fft == 0:
                slice_last = np.s_[-1, :, :]
Pierre Augier's avatar
Pierre Augier committed
393
394
            elif dimK_first_fft == 1:
                slice_last = np.s_[:, -1, :]
Pierre Augier's avatar
Pierre Augier committed
395
396
            else:
                raise NotImplementedError
Pierre Augier's avatar
Pierre Augier committed
397

Pierre Augier's avatar
Pierre Augier committed
398
399
400
            field_fft[slice_last] /= 2

        return field_fft
401

Pierre Augier's avatar
Pierre Augier committed
402
    def project_perpk3d(self, vx_fft, vy_fft, vz_fft):
Pierre Augier's avatar
Pierre Augier committed
403
404
405
406
407
        """Project (inplace) a vector perpendicular to the wavevector.

        The resulting vector is divergence-free.

        """
Pierre Augier's avatar
Pierre Augier committed
408
409
410
411
412
413
414
415
416
        project_perpk3d(
            vx_fft,
            vy_fft,
            vz_fft,
            self.Kx,
            self.Ky,
            self.Kz,
            self.inv_K_square_nozero,
        )
Pierre Augier's avatar
Pierre Augier committed
417

Pierre Augier's avatar
Pierre Augier committed
418
    def divfft_from_vecfft(self, vx_fft, vy_fft, vz_fft):
Pierre Augier's avatar
Pierre Augier committed
419
        """Return the divergence of a vector in spectral space."""
Pierre Augier's avatar
Pierre Augier committed
420
421
422
423
424
425
426
        # float64[][][]
        Kx = self.Kx
        Ky = self.Ky
        Kz = self.Kz

        return divfft_from_vecfft(vx_fft, vy_fft, vz_fft, Kx, Ky, Kz)

427
428
429
430
431
432
433
434
    def rotfft_from_vecfft(self, vx_fft, vy_fft, vz_fft):
        """Return the curl of a vector in spectral space."""
        # float64[][][]
        Kx = self.Kx
        Ky = self.Ky
        Kz = self.Kz

        return rotfft_from_vecfft(vx_fft, vy_fft, vz_fft, Kx, Ky, Kz)
Pierre Augier's avatar
Pierre Augier committed
435

436
437
438
    def rotfft_from_vecfft_outin(
        self, vx_fft, vy_fft, vz_fft, rotxfft, rotyfft, rotzfft
    ):
Pierre Augier's avatar
Pierre Augier committed
439
440
441
442
443
444
        """Return the curl of a vector in spectral space."""
        # float64[][][]
        Kx = self.Kx
        Ky = self.Ky
        Kz = self.Kz

445
446
447
        rotfft_from_vecfft_outin(
            vx_fft, vy_fft, vz_fft, Kx, Ky, Kz, rotxfft, rotyfft, rotzfft
        )
Pierre Augier's avatar
Pierre Augier committed
448

Pierre Augier's avatar
Pierre Augier committed
449
450
    def div_vb_fft_from_vb(self, vx, vy, vz, b):
        r"""Compute :math:`\nabla \cdot (\boldsymbol{v} b)` in spectral space.
Pierre Augier's avatar
Pierre Augier committed
451

Pierre Augier's avatar
Pierre Augier committed
452
453
        """
        fft3d = self.fft3d
Pierre Augier's avatar
Pierre Augier committed
454

Pierre Augier's avatar
Pierre Augier committed
455
456
457
        vxbfft = fft3d(vx * b)
        vybfft = fft3d(vy * b)
        vzbfft = fft3d(vz * b)
Pierre Augier's avatar
Pierre Augier committed
458

Pierre Augier's avatar
Pierre Augier committed
459
460
461
        return divfft_from_vecfft(
            vxbfft, vybfft, vzbfft, self.Kx, self.Ky, self.Kz
        )
Pierre Augier's avatar
Pierre Augier committed
462

Pierre Augier's avatar
Pierre Augier committed
463
    def rotzfft_from_vxvyfft(self, vx_fft, vy_fft):
Pierre Augier's avatar
Pierre Augier committed
464
        """Compute the z component of the curl in spectral space."""
Pierre Augier's avatar
Pierre Augier committed
465
466
467
468
469
470
471
        return 1j * (self.Kx * vy_fft - self.Ky * vx_fft)

    def get_XYZ_loc(self):
        """Compute the local 3d arrays with the x, y, and y values.

        """

472
        if self.shapeX_seq != self.shapeX_loc:
Pierre Augier's avatar
Pierre Augier committed
473
            i0_seq_start, i1_seq_start, i2_seq_start = self.seq_indices_first_X
Pierre Augier's avatar
Pierre Augier committed
474
            if self.shapeX_seq[1:] != self.shapeX_loc[1:]:
Pierre Augier's avatar
Pierre Augier committed
475
                # general solution
476
477
478
479
                # mpi.print_sorted(
                #     'in get_XYZ_loc:',
                #     '(i0_seq_start, i1_seq_start, i2_seq_start):',
                #     (i0_seq_start, i1_seq_start, i2_seq_start))
Pierre Augier's avatar
Pierre Augier committed
480

Pierre Augier's avatar
Pierre Augier committed
481
482
483
                z_loc = self.z_seq[i0_seq_start:i0_seq_start + self.shapeX_loc[0]]
                y_loc = self.y_seq[i1_seq_start:i1_seq_start + self.shapeX_loc[1]]
                x_loc = self.x_seq[i2_seq_start:i2_seq_start + self.shapeX_loc[2]]
Pierre Augier's avatar
Pierre Augier committed
484

Pierre Augier's avatar
Pierre Augier committed
485
486
487
            # mpi.print_sorted('z_loc', z_loc)
            # mpi.print_sorted('y_loc', y_loc)
            # mpi.print_sorted('x_loc', x_loc)
Pierre Augier's avatar
Pierre Augier committed
488

Pierre Augier's avatar
Pierre Augier committed
489
490
491
492
            else:
                # 1d decomposition
                x_loc = self.x_seq
                y_loc = self.y_seq
Pierre Augier's avatar
Pierre Augier committed
493
                z_loc = self.z_seq[i0_seq_start:i0_seq_start + self.shapeX_loc[0]]
Pierre Augier's avatar
Pierre Augier committed
494
495
496
497
498
499
500
501
502
503
        else:
            x_loc = self.x_seq
            y_loc = self.y_seq
            z_loc = self.z_seq

        Y, Z, X = np.meshgrid(y_loc, z_loc, x_loc, copy=False)

        assert X.shape == Y.shape == Z.shape == self.shapeX_loc

        return X, Y, Z
504
505
506
507
508
509
510

    def compute_1dspectra(self, energy_fft):
        """Compute the 1D spectra.

        Returns
        -------

511
        spectrum_kx
512

513
        spectrum_ky
514

515
        spectrum_kz
516
517

        """
518
        # nk0, nk1, nk2 = self.shapeK_loc
519
520
        spectrum_k0k1k2 = self._compute_spectrum3d_loc(energy_fft)
        dimX_K = self._op_fft.get_dimX_K()
521

522
        if self._is_mpi_lib:
Pierre Augier's avatar
Pierre Augier committed
523

524
525
            def compute_spectrum_ki(dimXi):
                ni = self.shapeX_seq[dimXi]
Pierre Augier's avatar
Pierre Augier committed
526
                nk_spectra = ni // 2 + 1
527
                dimK = dimX_K.index(dimXi)
528
                dims_for_sum = tuple(dim for dim in range(3) if dim != dimK)
529
530
531
532
                spectrum_tmp_loc = spectrum_k0k1k2.sum(axis=dims_for_sum)
                istart = self.seq_indices_first_K[dimK]
                nk_loc = self.shapeK_loc[dimK]

533
                if self.dimK_first_fft != dimK:
534
                    spectrum_tmp_seq = np.zeros(ni)
Pierre Augier's avatar
Pierre Augier committed
535
                    spectrum_tmp_seq[istart:istart + nk_loc] = spectrum_tmp_loc
536
                    spectrum_ki = spectrum_tmp_seq[:nk_spectra]
Pierre Augier's avatar
Pierre Augier committed
537
                    nk1 = (ni + 1) // 2
538
539
540
                    spectrum_ki[1:nk1] += spectrum_tmp_seq[nk_spectra:][::-1]
                else:
                    spectrum_tmp_seq = np.zeros(nk_spectra)
Pierre Augier's avatar
Pierre Augier committed
541
                    spectrum_tmp_seq[istart:istart + nk_loc] = spectrum_tmp_loc
542
543
544
                    spectrum_ki = spectrum_tmp_seq

                spectrum_ki = mpi.comm.allreduce(spectrum_ki, op=mpi.MPI.SUM)
545

546
547
548
549
550
                return spectrum_ki

            spectrum_kx = compute_spectrum_ki(dimXi=2)
            spectrum_ky = compute_spectrum_ki(dimXi=1)
            spectrum_kz = compute_spectrum_ki(dimXi=0)
551

552
        else:
Pierre Augier's avatar
Pierre Augier committed
553

554
555
            def compute_spectrum_ki(dimXi):
                ni = self.shapeX_seq[dimXi]
Pierre Augier's avatar
Pierre Augier committed
556
                nk_spectra = ni // 2 + 1
557
                dimK = dimX_K.index(dimXi)
558
                dims_for_sum = tuple(dim for dim in range(3) if dim != dimK)
559
560
561
                spectrum_tmp = spectrum_k0k1k2.sum(axis=dims_for_sum)
                if self.dimK_first_fft != dimK:
                    spectrum_ki = spectrum_tmp[:nk_spectra]
Pierre Augier's avatar
Pierre Augier committed
562
                    nk1 = (ni + 1) // 2
563
564
565
566
567
568
569
570
571
                    spectrum_ki[1:nk1] += spectrum_tmp[nk_spectra:][::-1]
                else:
                    spectrum_ki = spectrum_tmp
                return spectrum_ki

            spectrum_kx = compute_spectrum_ki(dimXi=2)
            spectrum_ky = compute_spectrum_ki(dimXi=1)
            spectrum_kz = compute_spectrum_ki(dimXi=0)

Pierre Augier's avatar
Pierre Augier committed
572
573
574
575
576
        return (
            spectrum_kx / self.deltakx,
            spectrum_ky / self.deltaky,
            spectrum_kz / self.deltakz,
        )
577

578
579
580
    def compute_3dspectrum(self, energy_fft):
        """Compute the 3D spectrum.

Pierre Augier's avatar
Pierre Augier committed
581
        The corresponding wavenumber array is ``self.k_spectra3d``.
582
583

        """
Pierre Augier's avatar
Pierre Augier committed
584
585
586
587
588
589
        K2 = self.K2
        ks = self.k_spectra3d
        spectrum_k0k1k2 = self._compute_spectrum3d_loc(energy_fft)
        spectrum3d = loop_spectra3d(spectrum_k0k1k2, ks, K2)
        if self._is_mpi_lib:
            spectrum3d = mpi.comm.allreduce(spectrum3d, op=mpi.MPI.SUM)
Pierre Augier's avatar
Pierre Augier committed
590
        return spectrum3d / self.deltak_spectra3d
591
592
593
594

    def compute_spectra_2vars(self, energy_fft):
        """Compute spectra vs 2 variables.

Pierre Augier's avatar
Pierre Augier committed
595
596
597
598
599
600
601
        .. warning::

           Not implemented!

        .. todo::

           Implement the method :func:`compute_spectra_2vars`.
602
603
604
605
606
607
608
609
610
611
612
613

        Returns
        -------

        E_kx_kyz

        E_ky_kzx

        E_kz_kxy

        """
        raise NotImplementedError
Pierre Augier's avatar
Pierre Augier committed
614
615


Pierre Augier's avatar
Pierre Augier committed
616
617
618
619
620
# This one is actually not so useful!
# def get_cross_section(self, equation='x=0', to_process=0):
#     """Get a 2d cross section.

#     .. warning::
Pierre Augier's avatar
Pierre Augier committed
621

Pierre Augier's avatar
Pierre Augier committed
622
#        Not implemented!
Pierre Augier's avatar
Pierre Augier committed
623

Pierre Augier's avatar
Pierre Augier committed
624
#     .. todo::
Pierre Augier's avatar
Pierre Augier committed
625

Pierre Augier's avatar
Pierre Augier committed
626
627
628
#        Implement the method :func:`get_cross_section`.  We need a
#        not-implemented method :func:`get_seq_indices_first_X` in the C++
#        classes...
Pierre Augier's avatar
Pierre Augier committed
629

Pierre Augier's avatar
Pierre Augier committed
630
631
#        We first have to implement the very simple cases for which
#        ``equation`` is equal to:
Pierre Augier's avatar
Pierre Augier committed
632

Pierre Augier's avatar
Pierre Augier committed
633
634
635
636
637
638
#        - x = 2.
#        - y = 2.
#        - z = 2.
#        - ix = 10
#        - iy = 10
#        - iz = 10
Pierre Augier's avatar
Pierre Augier committed
639

Pierre Augier's avatar
Pierre Augier committed
640
641
#     Parameters
#     ----------
Pierre Augier's avatar
Pierre Augier committed
642

Pierre Augier's avatar
Pierre Augier committed
643
#     equation: str
Pierre Augier's avatar
Pierre Augier committed
644

Pierre Augier's avatar
Pierre Augier committed
645
646
#       Equation defining the cross-section. We should be able to use the
#       variables x, y, z, ix, iy and iz.
Pierre Augier's avatar
Pierre Augier committed
647

Pierre Augier's avatar
Pierre Augier committed
648
649
#     """
#     raise NotImplementedError