rg_transforms.py 22.7 KB
Newer Older
1
2
3
4
import warnings

import numpy as np
from d2o import distributed_data_object, STRATEGIES
5
from nifty.config import dependency_injector as gdi
6
import nifty.nifty_utilities as utilities
7

8
from keepers import Loggable
9
10
11
12

pyfftw = gdi.get('pyfftw')


13
class Transform(Loggable, object):
Jait Dixit's avatar
Jait Dixit committed
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
    """
        A generic fft object without any implementation.
    """

    def __init__(self, domain, codomain):
        pass

    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

39

Jait Dixit's avatar
Jait Dixit committed
40
class FFTW(Transform):
41
42
43
44
45
    """
        The pyfftw pendant of a fft object.
    """

    def __init__(self, domain, codomain):
Jait Dixit's avatar
Jait Dixit committed
46
47
        self.domain = domain
        self.codomain = codomain
48
49
50
51
52
53
54

        if 'pyfftw' not in gdi:
            raise ImportError("The module pyfftw is needed but not available.")

        # Enable caching for pyfftw.interfaces
        pyfftw.interfaces.cache.enable()

Jait Dixit's avatar
Jait Dixit committed
55
56
57
58
59
60
61
62
63
        # The plan_dict stores the FFTWTransformInfo objects which correspond
        # to a certain set of (field_val, domain, codomain) sets.
        self.info_dict = {}

        # 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,
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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
                           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)
        offset[0] = int(offset_input)
        # 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
125
        if temp_id not in self.centering_mask_dict:
126
127
128
129
            # 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
130
131
132
133
134
                              (np.tensordot(to_center,
                                            args +
                                            offset.reshape(offset.shape +
                                                           (1,) *
                                                           (np.array(
135
                                                              args).ndim - 1)),
Jait Dixit's avatar
Jait Dixit committed
136
                                            1)),
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
                (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]
Jait Dixit's avatar
Jait Dixit committed
164
165
            self.centering_mask_dict[temp_id] = centering_mask
        return self.centering_mask_dict[temp_id]
166

167
    def _get_transform_info(self, domain, codomain, axes, local_shape,
168
169
170
171
172
                            local_offset_Q, is_local, transform_shape=None,
                            **kwargs):
        # generate a id-tuple which identifies the domain-codomain setting
        temp_id = (domain.__hash__() ^
                   (101 * codomain.__hash__()) ^
Jait Dixit's avatar
Jait Dixit committed
173
174
                   (211 * transform_shape.__hash__()) ^
                   (131 * is_local.__hash__())
Theo Steininger's avatar
Theo Steininger committed
175
                   )
176
177

        # generate the plan_and_info object if not already there
Jait Dixit's avatar
Jait Dixit committed
178
        if temp_id not in self.info_dict:
179
            if is_local:
Jait Dixit's avatar
Jait Dixit committed
180
                self.info_dict[temp_id] = FFTWLocalTransformInfo(
181
                    domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
182
                    local_offset_Q, self, **kwargs
183
184
                )
            else:
Jait Dixit's avatar
Jait Dixit committed
185
                self.info_dict[temp_id] = FFTWMPITransfromInfo(
186
                    domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
187
                    local_offset_Q, self, transform_shape, **kwargs
188
189
                )

Jait Dixit's avatar
Jait Dixit committed
190
        return self.info_dict[temp_id]
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215

    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
216
                 for x, y in enumerate(val.shape)]
217
218
219
220
221
222
            )

        return val * mask

    def _atomic_mpi_transform(self, val, info, axes):
        # Apply codomain centering mask
223
        if reduce(lambda x, y: x + y, self.codomain.zerocenter):
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
            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:
            p.input_array[:] = val
        # Execute the plan
        p()

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

        # Apply domain centering mask
240
        if reduce(lambda x, y: x + y, self.domain.zerocenter):
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
            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
257

258
259
        current_info = self._get_transform_info(self.domain,
                                                self.codomain,
260
                                                axes,
261
                                                local_shape=local_val.shape,
Jait Dixit's avatar
Jait Dixit committed
262
                                                local_offset_Q=False,
263
264
265
266
                                                is_local=True,
                                                **kwargs)

        # Apply codomain centering mask
267
        if reduce(lambda x, y: x + y, self.codomain.zerocenter):
268
269
270
271
272
273
274
275
276
277
278
            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
279
        if reduce(lambda x, y: x + y, self.domain.zerocenter):
280
281
282
283
284
285
286
287
288
            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
Theo Steininger's avatar
Theo Steininger committed
289
            result_dtype = np.result_type(np.complex, self.codomain.dtype)
290
            return_val = val.copy_empty(global_shape=val.shape,
Theo Steininger's avatar
Theo Steininger committed
291
                                        dtype=result_dtype)
292
293
294
295
296
297
298
299
            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')
300
        self.logger.info("Repacking d2o to fftw distribution strategy")
301
302
303
304
305
306
307
308
309
310
311
312
313
        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
314
315
316
317
        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)
318

Theo Steininger's avatar
Theo Steininger committed
319
        result_dtype = np.result_type(np.complex, self.codomain.dtype)
320
        return_val = val.copy_empty(global_shape=val.shape,
Theo Steininger's avatar
Theo Steininger committed
321
                                    dtype=result_dtype)
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339

        # 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
340
341
                    temp_val = np.empty_like(
                        local_val,
Theo Steininger's avatar
Theo Steininger committed
342
                        dtype=result_dtype
Jait Dixit's avatar
Jait Dixit committed
343
                    )
344
345
346
347
348
349
350
351
352
353
354
                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
355
                axes = (0, )
356
357

            if current_info is None:
358
359
360
                transform_shape = list(inp.shape)
                transform_shape[0] = val.shape[0]

361
362
363
                current_info = self._get_transform_info(
                    self.domain,
                    self.codomain,
364
                    axes,
365
366
367
                    local_shape=val.local_shape,
                    local_offset_Q=local_offset_Q,
                    is_local=False,
368
                    transform_shape=tuple(transform_shape),
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
                    **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
391
    def transform(self, val, axes, **kwargs):
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
        """
            The pyfftw transform function.

            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):
415
            raise ValueError("Provided axes does not match array shape")
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448

        # 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):
449
    def __init__(self, domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
450
                 local_offset_Q, fftw_context, **kwargs):
451
452
453
        if pyfftw is None:
            raise ImportError("The module pyfftw is needed but not available.")

Theo Steininger's avatar
Theo Steininger committed
454
455
456
457
458
459
        shape = (local_shape if axes is None else
                 [y for x, y in enumerate(local_shape) if x in axes])

        self.cmask_domain = fftw_context.get_centering_mask(domain.zerocenter,
                                                            shape,
                                                            local_offset_Q)
460

Jait Dixit's avatar
Jait Dixit committed
461
        self.cmask_codomain = fftw_context.get_centering_mask(
Theo Steininger's avatar
Theo Steininger committed
462
463
464
                                                        codomain.zerocenter,
                                                        shape,
                                                        local_offset_Q)
465
466
467

        # If both domain and codomain are zero-centered the result,
        # will get a global minus. Store the sign to correct it.
468
469
        self.sign = (-1) ** np.sum(np.array(domain.zerocenter) *
                                   np.array(codomain.zerocenter) *
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
                                   (np.array(domain.shape) // 2 % 2))

    @property
    def cmask_domain(self):
        return self._domain_centering_mask

    @cmask_domain.setter
    def cmask_domain(self, cmask):
        self._domain_centering_mask = cmask

    @property
    def cmask_codomain(self):
        return self._codomain_centering_mask

    @cmask_codomain.setter
    def cmask_codomain(self, cmask):
        self._codomain_centering_mask = cmask

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

    @sign.setter
    def sign(self, sign):
        self._sign = sign


class FFTWLocalTransformInfo(FFTWTransformInfo):
498
    def __init__(self, domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
499
                 local_offset_Q, fftw_context, **kwargs):
500
501
        super(FFTWLocalTransformInfo, self).__init__(domain,
                                                     codomain,
502
                                                     axes,
503
504
                                                     local_shape,
                                                     local_offset_Q,
Jait Dixit's avatar
Jait Dixit committed
505
                                                     fftw_context,
506
507
508
509
510
511
512
513
514
515
516
517
                                                     **kwargs)
        if codomain.harmonic:
            self._fftw_interface = pyfftw.interfaces.numpy_fft.fftn
        else:
            self._fftw_interface = pyfftw.interfaces.numpy_fft.ifftn

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


class FFTWMPITransfromInfo(FFTWTransformInfo):
518
    def __init__(self, domain, codomain, axes, local_shape,
Jait Dixit's avatar
Jait Dixit committed
519
                 local_offset_Q, fftw_context, transform_shape, **kwargs):
520
521
        super(FFTWMPITransfromInfo, self).__init__(domain,
                                                   codomain,
522
                                                   axes,
523
524
                                                   local_shape,
                                                   local_offset_Q,
Jait Dixit's avatar
Jait Dixit committed
525
                                                   fftw_context,
526
527
528
529
530
531
532
533
534
535
536
537
538
539
                                                   **kwargs)
        self._plan = pyfftw.create_mpi_plan(
            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
540
541
542
543
544
545
546
547
548
549
550
551

class GFFT(Transform):
    """
        The gfft pendant of a fft object.

        Parameters
        ----------
        fft_module_name : String
            Switch between the gfft module used: 'gfft' and 'gfft_dummy'

    """

552
    def __init__(self, domain, codomain, fft_module=None):
Jait Dixit's avatar
Jait Dixit committed
553
        if fft_module is None:
554
            fft_module = gdi['gfft_dummy']
Jait Dixit's avatar
Jait Dixit committed
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583

        self.domain = domain
        self.codomain = codomain
        self.fft_machine = fft_module

    def transform(self, val, axes, **kwargs):
        """
            The gfft transform function.

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

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

            **kwargs : *optional*
                Further kwargs are not processed.

            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):
584
            raise ValueError("Provided axes does not match array shape")
Jait Dixit's avatar
Jait Dixit committed
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611

        # GFFT doesn't accept d2o objects as input. Consolidate data from
        # all nodes into numpy.ndarray before proceeding.
        if isinstance(val, distributed_data_object):
            temp_inp = val.get_full_data()
        else:
            temp_inp = val

        # Array for storing the result
        return_val = None

        for slice_list in utilities.get_slice_list(temp_inp.shape, axes):

            # don't copy the whole data array
            if slice_list == [slice(None, None)]:
                inp = temp_inp
            else:
                # initialize the return_val object if needed
                if return_val is None:
                    return_val = np.empty_like(temp_inp)
                inp = temp_inp[slice_list]

            inp = self.fft_machine.gfft(
                inp,
                in_ax=[],
                out_ax=[],
                ftmachine='fft' if self.codomain.harmonic else 'ifft',
612
613
614
615
                in_zero_center=map(bool, self.domain.zerocenter),
                out_zero_center=map(bool, self.codomain.zerocenter),
                # enforce_hermitian_symmetry=bool(self.codomain.complexity),
                enforce_hermitian_symmetry=False,
Jait Dixit's avatar
Jait Dixit committed
616
617
618
619
620
621
622
623
624
                W=-1,
                alpha=-1,
                verbose=False
            )
            if slice_list == [slice(None, None)]:
                return_val = inp
            else:
                return_val[slice_list] = inp

Theo Steininger's avatar
Theo Steininger committed
625
        result_dtype = np.result_type(np.complex, self.codomain.dtype)
Jait Dixit's avatar
Jait Dixit committed
626
        if isinstance(val, distributed_data_object):
Theo Steininger's avatar
Theo Steininger committed
627
            new_val = val.copy_empty(dtype=result_dtype)
Jait Dixit's avatar
Jait Dixit committed
628
629
630
            new_val.set_full_data(return_val, copy=False)
            return_val = new_val
        else:
Theo Steininger's avatar
Theo Steininger committed
631
            return_val = return_val.astype(result_type, copy=False)
Jait Dixit's avatar
Jait Dixit committed
632
633

        return return_val