rg_transforms.py 25 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
Theo Steininger's avatar
Theo Steininger committed
13
14
15
16
17
#
# Copyright(C) 2013-2017 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik
# and financially supported by the Studienstiftung des deutschen Volkes.
18

19
20
21
22
import warnings

import numpy as np
from d2o import distributed_data_object, STRATEGIES
23
from nifty.config import dependency_injector as gdi
24
import nifty.nifty_utilities as utilities
25

26
from keepers import Loggable
27

Theo Steininger's avatar
Theo Steininger committed
28
fftw = gdi.get('fftw')
29
30


31
class Transform(Loggable, object):
Jait Dixit's avatar
Jait Dixit committed
32
33
34
35
    """
        A generic fft object without any implementation.
    """

36
    def __init__(self, domain, codomain):
Jait Dixit's avatar
Jait Dixit committed
37
38
        self.domain = domain
        self.codomain = codomain
39

Jait Dixit's avatar
Jait Dixit committed
40
41
42
43
44
        # initialize the dictionary which stores the values from
        # get_centering_mask
        self.centering_mask_dict = {}

    def get_centering_mask(self, to_center_input, dimensions_input,
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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
                           offset_input=False):
        """
            Computes the mask, used to (de-)zerocenter domain and target
            fields.

            Parameters
            ----------
            to_center_input : tuple, list, numpy.ndarray
                A tuple of booleans which dimensions should be
                zero-centered.

            dimensions_input : tuple, list, numpy.ndarray
                A tuple containing the mask's desired shape.

            offset_input : int, boolean
                Specifies whether the zero-th dimension starts with an odd
                or and even index, i.e. if it is shifted.

            Returns
            -------
            result : np.ndarray
                A 1/-1-alternating mask.
        """
        # cast input
        to_center = np.array(to_center_input)
        dimensions = np.array(dimensions_input)

        # if none of the dimensions are zero centered, return a 1
        if np.all(to_center == 0):
            return 1

        if np.all(dimensions == np.array(1)) or \
                np.all(dimensions == np.array([1])):
            return dimensions
        # The dimensions of size 1 must be sorted out for computing the
        # centering_mask. The depth of the array will be restored in the
        # end.
        size_one_dimensions = []
        temp_dimensions = []
        temp_to_center = []
        for i in range(len(dimensions)):
            if dimensions[i] == 1:
                size_one_dimensions += [True]
            else:
                size_one_dimensions += [False]
                temp_dimensions += [dimensions[i]]
                temp_to_center += [to_center[i]]
        dimensions = np.array(temp_dimensions)
        to_center = np.array(temp_to_center)
        # cast the offset_input into the shape of to_center
        offset = np.zeros(to_center.shape, dtype=int)
Theo Steininger's avatar
Theo Steininger committed
96
97
98
99
        # if the first dimension has length 1 and has an offset, restore the
        # global minus by hand
        if not size_one_dimensions[0]:
            offset[0] = int(offset_input)
100
101
102
103
104
105
106
107
108
        # check for dimension match
        if to_center.size != dimensions.size:
            raise TypeError(
                'The length of the supplied lists does not match.')

        # build up the value memory
        # compute an identifier for the parameter set
        temp_id = tuple(
            (tuple(to_center), tuple(dimensions), tuple(offset)))
Jait Dixit's avatar
Jait Dixit committed
109
        if temp_id not in self.centering_mask_dict:
110
111
112
113
            # use np.tile in order to stack the core alternation scheme
            # until the desired format is constructed.
            core = np.fromfunction(
                lambda *args: (-1) **
Jait Dixit's avatar
Jait Dixit committed
114
115
116
117
118
                              (np.tensordot(to_center,
                                            args +
                                            offset.reshape(offset.shape +
                                                           (1,) *
                                                           (np.array(
119
                                                              args).ndim - 1)),
Jait Dixit's avatar
Jait Dixit committed
120
                                            1)),
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
                (2,) * to_center.size)
            # Cast the core to the smallest integers we can get
            core = core.astype(np.int8)

            centering_mask = np.tile(core, dimensions // 2)
            # for the dimensions of odd size corresponding slices must be
            # added
            for i in range(centering_mask.ndim):
                # check if the size of the certain dimension is odd or even
                if (dimensions % 2)[i] == 0:
                    continue
                # prepare the slice object
                temp_slice = (slice(None),) * i + (slice(-2, -1, 1),) + \
                             (slice(None),) * (centering_mask.ndim - 1 - i)
                # append the slice to the centering_mask
                centering_mask = np.append(centering_mask,
                                           centering_mask[temp_slice],
                                           axis=i)
            # Add depth to the centering_mask where the length of a
            # dimension was one
            temp_slice = ()
            for i in range(len(size_one_dimensions)):
                if size_one_dimensions[i]:
                    temp_slice += (None,)
                else:
                    temp_slice += (slice(None),)
            centering_mask = centering_mask[temp_slice]
Theo Steininger's avatar
Theo Steininger committed
148
149
150
151
152
            # if the first dimension has length 1 and has an offset, restore
            # the global minus by hand
            if size_one_dimensions[0] and offset_input:
                centering_mask *= -1

Jait Dixit's avatar
Jait Dixit committed
153
154
            self.centering_mask_dict[temp_id] = centering_mask
        return self.centering_mask_dict[temp_id]
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179

    def _apply_mask(self, val, mask, axes):
        """
            Apply centering mask to an array.

            Parameters
            ----------
            val: distributed_data_object or numpy.ndarray
                The value-array on which the mask should be applied.

            mask: numpy.ndarray
                The mask to be applied.

            axes: tuple
                The axes which are to be transformed.

            Returns
            -------
            distributed_data_object or np.nd_array
                Mask input array by multiplying it with the mask.
        """
        # reshape mask if necessary
        if axes:
            mask = mask.reshape(
                [y if x in axes else 1
Jait Dixit's avatar
Jait Dixit committed
180
                 for x, y in enumerate(val.shape)]
181
182
183
            )
        return val * mask

Theo Steininger's avatar
Theo Steininger committed
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
    def transform(self, val, axes, **kwargs):
        """
            A generic ff-transform function.

            Parameters
            ----------
            field_val : distributed_data_object
                The value-array of the field which is supposed to
                be transformed.

            domain : nifty.rg.nifty_rg.rg_space
                The domain of the space which should be transformed.

            codomain : nifty.rg.nifty_rg.rg_space
                The taget into which the field should be transformed.
        """
        raise NotImplementedError


Martin Reinecke's avatar
Martin Reinecke committed
203
class MPIFFT(Transform):
Theo Steininger's avatar
Theo Steininger committed
204
    """
Martin Reinecke's avatar
Martin Reinecke committed
205
        The MPI-parallel FFTW pendant of a fft object.
Theo Steininger's avatar
Theo Steininger committed
206
207
208
209
    """

    def __init__(self, domain, codomain):

Theo Steininger's avatar
Theo Steininger committed
210
        if not hasattr(fftw, 'FFTW_MPI'):
Martin Reinecke's avatar
Martin Reinecke committed
211
212
            raise ImportError(
                "The MPI FFTW module is needed but not available.")
Theo Steininger's avatar
Theo Steininger committed
213

Martin Reinecke's avatar
Martin Reinecke committed
214
        super(MPIFFT, self).__init__(domain, codomain)
Theo Steininger's avatar
Theo Steininger committed
215

Martin Reinecke's avatar
Martin Reinecke committed
216
        # Enable caching
Theo Steininger's avatar
Theo Steininger committed
217
        fftw.interfaces.cache.enable()
Theo Steininger's avatar
Theo Steininger committed
218
219
220
221
222
223
224
225
226

        # The plan_dict stores the FFTWTransformInfo objects which correspond
        # to a certain set of (field_val, domain, codomain) sets.
        self.info_dict = {}

    def _get_transform_info(self, domain, codomain, axes, local_shape,
                            local_offset_Q, is_local, transform_shape=None,
                            **kwargs):
        # generate a id-tuple which identifies the domain-codomain setting
Theo Steininger's avatar
Theo Steininger committed
227
        temp_id = (domain, codomain, transform_shape, is_local)
Theo Steininger's avatar
Theo Steininger committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243

        # generate the plan_and_info object if not already there
        if temp_id not in self.info_dict:
            if is_local:
                self.info_dict[temp_id] = FFTWLocalTransformInfo(
                    domain, codomain, axes, local_shape,
                    local_offset_Q, self, **kwargs
                )
            else:
                self.info_dict[temp_id] = FFTWMPITransfromInfo(
                    domain, codomain, axes, local_shape,
                    local_offset_Q, self, transform_shape, **kwargs
                )

        return self.info_dict[temp_id]

244
245
    def _atomic_mpi_transform(self, val, info, axes):
        # Apply codomain centering mask
246
        if reduce(lambda x, y: x + y, self.codomain.zerocenter):
247
248
249
250
251
252
            temp_val = np.copy(val)
            val = self._apply_mask(temp_val, info.cmask_codomain, axes)

        p = info.plan
        # Load the value into the plan
        if p.has_input:
253
            p.input_array[None] = val
254
255
256
257
258
259
260
261
262
        # Execute the plan
        p()

        if p.has_output:
            result = p.output_array
        else:
            return None

        # Apply domain centering mask
263
        if reduce(lambda x, y: x + y, self.domain.zerocenter):
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
            result = self._apply_mask(result, info.cmask_domain, axes)

        # Correct the sign if needed
        result *= info.sign

        return result

    def _local_transform(self, val, axes, **kwargs):
        ####
        # val must be numpy array or d2o with slicing distributor
        ###

        try:
            local_val = val.get_local_data(copy=False)
        except(AttributeError):
            local_val = val
Jait Dixit's avatar
Jait Dixit committed
280

281
282
        current_info = self._get_transform_info(self.domain,
                                                self.codomain,
283
                                                axes,
284
                                                local_shape=local_val.shape,
Jait Dixit's avatar
Jait Dixit committed
285
                                                local_offset_Q=False,
286
287
288
289
                                                is_local=True,
                                                **kwargs)

        # Apply codomain centering mask
290
        if reduce(lambda x, y: x + y, self.codomain.zerocenter):
291
292
293
294
295
296
297
298
299
300
301
            temp_val = np.copy(local_val)
            local_val = self._apply_mask(temp_val,
                                         current_info.cmask_codomain, axes)

        local_result = current_info.fftw_interface(
            local_val,
            axes=axes,
            planner_effort='FFTW_ESTIMATE'
        )

        # Apply domain centering mask
302
        if reduce(lambda x, y: x + y, self.domain.zerocenter):
303
304
305
306
307
308
309
310
311
312
            local_result = self._apply_mask(local_result,
                                            current_info.cmask_domain, axes)

        # Correct the sign if needed
        if current_info.sign != 1:
            local_result *= current_info.sign

        try:
            # Create return object and insert results inplace
            return_val = val.copy_empty(global_shape=val.shape,
313
                                        dtype=np.complex)
314
315
316
317
318
319
320
321
            return_val.set_local_data(data=local_result, copy=False)
        except(AttributeError):
            return_val = local_result

        return return_val

    def _repack_to_fftw_and_transform(self, val, axes, **kwargs):
        temp_val = val.copy_empty(distribution_strategy='fftw')
322
        self.logger.info("Repacking d2o to fftw distribution strategy")
323
324
325
326
327
328
329
330
331
332
333
334
335
        temp_val.set_full_data(val, copy=False)

        # Recursive call to transform
        result = self.transform(temp_val, axes, **kwargs)

        return_val = result.copy_empty(
            distribution_strategy=val.distribution_strategy)
        return_val.set_full_data(data=result, copy=False)

        return return_val

    def _mpi_transform(self, val, axes, **kwargs):

Jait Dixit's avatar
Jait Dixit committed
336
337
338
339
        local_offset_list = np.cumsum(
            np.concatenate([[0, ], val.distributor.all_local_slices[:, 2]])
        )
        local_offset_Q = bool(local_offset_list[val.distributor.comm.rank] % 2)
340
        return_val = val.copy_empty(global_shape=val.shape,
341
                                    dtype=np.complex)
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359

        # Extract local data
        local_val = val.get_local_data(copy=False)

        # Create temporary storage for slices
        temp_val = None

        # If axes tuple includes all axes, set it to None
        if axes is not None:
            if set(axes) == set(range(len(val.shape))):
                axes = None

        current_info = None
        for slice_list in utilities.get_slice_list(local_val.shape, axes):
            if slice_list == [slice(None, None)]:
                inp = local_val
            else:
                if temp_val is None:
Jait Dixit's avatar
Jait Dixit committed
360
361
                    temp_val = np.empty_like(
                        local_val,
362
                        dtype=np.complex
Jait Dixit's avatar
Jait Dixit committed
363
                    )
364
365
366
367
368
369
370
371
372
373
374
                inp = local_val[slice_list]

            # This is in order to make FFTW behave properly when slicing input
            # over MPI ranks when the input is 1-dimensional. The default
            # behaviour is to optimize to take advantage of byte-alignment,
            # which doesn't match the slicing strategy for multi-dimensional
            # data.
            original_shape = None
            if len(inp.shape) == 1:
                original_shape = inp.shape
                inp = inp.reshape(inp.shape[0], 1)
Theo Steininger's avatar
Theo Steininger committed
375
                axes = (0, )
376
377
378
379
                if original_shape[0]%2!=0:
                    raise AttributeError("MPI-FFTs of onedimensional arrays "
                        "with odd length are currently not supported due to a "
                        "bug in FFTW. Please use a grid with even length.")
380
381

            if current_info is None:
382
383
384
                transform_shape = list(inp.shape)
                transform_shape[0] = val.shape[0]

385
386
387
                current_info = self._get_transform_info(
                    self.domain,
                    self.codomain,
388
                    axes,
389
390
391
                    local_shape=val.local_shape,
                    local_offset_Q=local_offset_Q,
                    is_local=False,
392
                    transform_shape=tuple(transform_shape),
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
                    **kwargs
                )

            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                result = self._atomic_mpi_transform(inp, current_info, axes)

            if result is None:
                temp_val = np.empty_like(local_val)
            elif slice_list == [slice(None, None)]:
                temp_val = result
            else:
                # Reverting to the original shape i.e. before the input was
                # augmented with 1 to make FFTW behave properly.
                if original_shape is not None:
                    result = result.reshape(original_shape)
                temp_val[slice_list] = result

        return_val.set_local_data(data=temp_val, copy=False)

        return return_val

Jait Dixit's avatar
Jait Dixit committed
415
    def transform(self, val, axes, **kwargs):
416
        """
Martin Reinecke's avatar
Martin Reinecke committed
417
            The MPI-parallel FFTW transform function.
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438

            Parameters
            ----------
            val : distributed_data_object or numpy.ndarray
                The value-array of the field which is supposed to
                be transformed.

            axes: tuple, None
                The axes which should be transformed.

            **kwargs : *optional*
                Further kwargs are passed to the create_mpi_plan routine.

            Returns
            -------
            result : np.ndarray or distributed_data_object
                Fourier-transformed pendant of the input field.
        """
        # Check if the axes provided are valid given the shape
        if axes is not None and \
                not all(axis in range(len(val.shape)) for axis in axes):
439
            raise ValueError("Provided axes does not match array shape")
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472

        # If the input is a numpy array we transform it locally
        if not isinstance(val, distributed_data_object):
            # Cast to a np.ndarray
            temp_val = np.asarray(val)

            # local transform doesn't apply transforms inplace
            return_val = self._local_transform(temp_val, axes)
        else:
            if val.distribution_strategy in STRATEGIES['slicing']:
                if axes is None or 0 in axes:
                    if val.distribution_strategy != 'fftw':
                        return_val = \
                            self._repack_to_fftw_and_transform(
                                val, axes, **kwargs
                            )
                    else:
                        return_val = self._mpi_transform(
                            val, axes, **kwargs
                        )
                else:
                    return_val = self._local_transform(
                        val, axes, **kwargs
                    )
            else:
                return_val = self._repack_to_fftw_and_transform(
                    val, axes, **kwargs
                )

        return return_val


class FFTWTransformInfo(object):
473
    def __init__(self, domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
474
                 local_offset_Q, fftw_context, **kwargs):
Theo Steininger's avatar
Theo Steininger committed
475
        if not hasattr(fftw, 'FFTW_MPI'):
Martin Reinecke's avatar
Martin Reinecke committed
476
477
            raise ImportError(
                "The MPI FFTW module is needed but not available.")
478

Theo Steininger's avatar
Theo Steininger committed
479
480
481
        shape = (local_shape if axes is None else
                 [y for x, y in enumerate(local_shape) if x in axes])

Theo Steininger's avatar
Theo Steininger committed
482
483
484
        self._cmask_domain = fftw_context.get_centering_mask(domain.zerocenter,
                                                             shape,
                                                             local_offset_Q)
485

Theo Steininger's avatar
Theo Steininger committed
486
487
488
489
        self._cmask_codomain = fftw_context.get_centering_mask(
                                                         codomain.zerocenter,
                                                         shape,
                                                         local_offset_Q)
490
491
492

        # If both domain and codomain are zero-centered the result,
        # will get a global minus. Store the sign to correct it.
Theo Steininger's avatar
Theo Steininger committed
493
494
495
        self._sign = (-1) ** np.sum(np.array(domain.zerocenter) *
                                    np.array(codomain.zerocenter) *
                                    (np.array(domain.shape) // 2 % 2))
496
497
498

    @property
    def cmask_domain(self):
Theo Steininger's avatar
Theo Steininger committed
499
        return self._cmask_domain
500
501
502

    @property
    def cmask_codomain(self):
Theo Steininger's avatar
Theo Steininger committed
503
        return self._cmask_codomain
504
505
506
507
508
509
510

    @property
    def sign(self):
        return self._sign


class FFTWLocalTransformInfo(FFTWTransformInfo):
511
    def __init__(self, domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
512
                 local_offset_Q, fftw_context, **kwargs):
513
514
        super(FFTWLocalTransformInfo, self).__init__(domain,
                                                     codomain,
515
                                                     axes,
516
517
                                                     local_shape,
                                                     local_offset_Q,
Jait Dixit's avatar
Jait Dixit committed
518
                                                     fftw_context,
519
520
                                                     **kwargs)
        if codomain.harmonic:
Theo Steininger's avatar
Theo Steininger committed
521
            self._fftw_interface = fftw.interfaces.numpy_fft.fftn
522
        else:
Theo Steininger's avatar
Theo Steininger committed
523
            self._fftw_interface = fftw.interfaces.numpy_fft.ifftn
524
525
526
527
528
529
530

    @property
    def fftw_interface(self):
        return self._fftw_interface


class FFTWMPITransfromInfo(FFTWTransformInfo):
531
    def __init__(self, domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
532
                 local_offset_Q, fftw_context, transform_shape, **kwargs):
533
534
        super(FFTWMPITransfromInfo, self).__init__(domain,
                                                   codomain,
535
                                                   axes,
536
537
                                                   local_shape,
                                                   local_offset_Q,
Jait Dixit's avatar
Jait Dixit committed
538
                                                   fftw_context,
539
                                                   **kwargs)
Theo Steininger's avatar
Theo Steininger committed
540
        self._plan = fftw.create_mpi_plan(
541
542
543
544
545
546
547
548
549
550
551
552
            input_shape=transform_shape,
            input_dtype='complex128',
            output_dtype='complex128',
            direction='FFTW_FORWARD' if codomain.harmonic else 'FFTW_BACKWARD',
            flags=["FFTW_ESTIMATE"],
            **kwargs
        )

    @property
    def plan(self):
        return self._plan

Jait Dixit's avatar
Jait Dixit committed
553

Theo Steininger's avatar
Theo Steininger committed
554
class SerialFFT(Transform):
Jait Dixit's avatar
Jait Dixit committed
555
    """
Theo Steininger's avatar
Theo Steininger committed
556
        The numpy fft pendant of a fft object.
Jait Dixit's avatar
Jait Dixit committed
557
558

    """
Theo Steininger's avatar
Theo Steininger committed
559
560
    def __init__(self, domain, codomain, use_fftw):
        super(SerialFFT, self).__init__(domain, codomain)
561

Theo Steininger's avatar
Theo Steininger committed
562
        if use_fftw and (fftw is None):
563
            raise ImportError(
Theo Steininger's avatar
Theo Steininger committed
564
                "The serial FFTW module is needed but not available.")
565

Theo Steininger's avatar
Theo Steininger committed
566
        self._use_fftw = use_fftw
567
        # Enable caching
Theo Steininger's avatar
Theo Steininger committed
568
569
        if self._use_fftw:
            fftw.interfaces.cache.enable()
Jait Dixit's avatar
Jait Dixit committed
570
571
572

    def transform(self, val, axes, **kwargs):
        """
Martin Reinecke's avatar
Martin Reinecke committed
573
            The scalar FFT transform function.
Jait Dixit's avatar
Jait Dixit committed
574
575
576

            Parameters
            ----------
Theo Steininger's avatar
Theo Steininger committed
577
            val : distributed_data_object or numpy.ndarray
Jait Dixit's avatar
Jait Dixit committed
578
579
580
                The value-array of the field which is supposed to
                be transformed.

Theo Steininger's avatar
Theo Steininger committed
581
            axes: tuple, None
Jait Dixit's avatar
Jait Dixit committed
582
583
584
                The axes which should be transformed.

            **kwargs : *optional*
Theo Steininger's avatar
Theo Steininger committed
585
                Further kwargs are passed to the create_mpi_plan routine.
Jait Dixit's avatar
Jait Dixit committed
586
587
588
589
590
591

            Returns
            -------
            result : np.ndarray or distributed_data_object
                Fourier-transformed pendant of the input field.
        """
592

Jait Dixit's avatar
Jait Dixit committed
593
594
595
        # Check if the axes provided are valid given the shape
        if axes is not None and \
                not all(axis in range(len(val.shape)) for axis in axes):
596
            raise ValueError("Provided axes does not match array shape")
Jait Dixit's avatar
Jait Dixit committed
597

Theo Steininger's avatar
Theo Steininger committed
598
        return_val = val.copy_empty(global_shape=val.shape,
599
                                    dtype=np.complex)
Jait Dixit's avatar
Jait Dixit committed
600

Theo Steininger's avatar
Theo Steininger committed
601
602
        if (axes is None) or (0 in axes) or \
           (val.distribution_strategy not in STRATEGIES['slicing']):
Jait Dixit's avatar
Jait Dixit committed
603

Theo Steininger's avatar
Theo Steininger committed
604
605
            if val.distribution_strategy == 'not':
                local_val = val.get_local_data(copy=False)
Jait Dixit's avatar
Jait Dixit committed
606
            else:
Theo Steininger's avatar
Theo Steininger committed
607
608
609
610
611
612
                local_val = val.get_full_data()

            result_data = self._atomic_transform(local_val=local_val,
                                                 axes=axes,
                                                 local_offset_Q=False)
            return_val.set_full_data(result_data, copy=False)
Jait Dixit's avatar
Jait Dixit committed
613
614

        else:
Theo Steininger's avatar
Theo Steininger committed
615
616
617
618
619
620
621
622
623
624
625
626
            local_offset_list = np.cumsum(
                    np.concatenate([[0, ],
                                    val.distributor.all_local_slices[:, 2]]))
            local_offset_Q = \
                bool(local_offset_list[val.distributor.comm.rank] % 2)

            local_val = val.get_local_data()
            result_data = self._atomic_transform(local_val=local_val,
                                                 axes=axes,
                                                 local_offset_Q=local_offset_Q)

            return_val.set_local_data(result_data, copy=False)
Jait Dixit's avatar
Jait Dixit committed
627
628

        return return_val
Theo Steininger's avatar
Theo Steininger committed
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645

    def _atomic_transform(self, local_val, axes, local_offset_Q):

        # some auxiliaries for the mask computation
        local_shape = local_val.shape
        shape = (local_shape if axes is None else
                 [y for x, y in enumerate(local_shape) if x in axes])

        # Apply codomain centering mask
        if reduce(lambda x, y: x + y, self.codomain.zerocenter):
            temp_val = np.copy(local_val)
            mask = self.get_centering_mask(self.codomain.zerocenter,
                                           shape,
                                           local_offset_Q)
            local_val = self._apply_mask(temp_val, mask, axes)

        # perform the transformation
Theo Steininger's avatar
Theo Steininger committed
646
        if self._use_fftw:
647
            if self.codomain.harmonic:
Theo Steininger's avatar
Theo Steininger committed
648
                result_val = fftw.interfaces.numpy_fft.fftn(
Martin Reinecke's avatar
Martin Reinecke committed
649
                             local_val, axes=axes)
650
            else:
Theo Steininger's avatar
Theo Steininger committed
651
                result_val = fftw.interfaces.numpy_fft.ifftn(
Martin Reinecke's avatar
Martin Reinecke committed
652
                             local_val, axes=axes)
653
        else:
654
655
656
657
            if self.codomain.harmonic:
                result_val = np.fft.fftn(local_val, axes=axes)
            else:
                result_val = np.fft.ifftn(local_val, axes=axes)
Theo Steininger's avatar
Theo Steininger committed
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674

        # Apply domain centering mask
        if reduce(lambda x, y: x + y, self.domain.zerocenter):
            mask = self.get_centering_mask(self.domain.zerocenter,
                                           shape,
                                           local_offset_Q)
            result_val = self._apply_mask(result_val, mask, axes)

        # If both domain and codomain are zero-centered the result,
        # will get a global minus. Store the sign to correct it.
        sign = (-1) ** np.sum(np.array(self.domain.zerocenter) *
                              np.array(self.codomain.zerocenter) *
                              (np.array(self.domain.shape) // 2 % 2))
        if sign != 1:
            result_val *= sign

        return result_val