rgrgtransformation.py 5.56 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
# NIFTy
# Copyright (C) 2017  Theo Steininger
#
# Author: Theo Steininger
#
# 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/>.

Jait Dixit's avatar
Jait Dixit committed
19
20
import numpy as np
from transformation import Transformation
Theo Steininger's avatar
Theo Steininger committed
21
from rg_transforms import FFTW, NUMPYFFT
Jait Dixit's avatar
Jait Dixit committed
22
23
24
25
from nifty import RGSpace, nifty_configuration


class RGRGTransformation(Transformation):
Jait Dixit's avatar
Jait Dixit committed
26
    def __init__(self, domain, codomain=None, module=None):
27
28
        super(RGRGTransformation, self).__init__(domain, codomain,
                                                 module=module)
Jait Dixit's avatar
Jait Dixit committed
29
30

        if module is None:
Theo Steininger's avatar
Theo Steininger committed
31
            if nifty_configuration['fft_module'] == 'fftw':
Jait Dixit's avatar
Jait Dixit committed
32
                self._transform = FFTW(self.domain, self.codomain)
Theo Steininger's avatar
Theo Steininger committed
33
34
            elif nifty_configuration['fft_module'] == 'numpy':
                self._transform = NUMPYFFT(self.domain, self.codomain)
Jait Dixit's avatar
Jait Dixit committed
35
            else:
Jait Dixit's avatar
Jait Dixit committed
36
37
                raise ValueError('ERROR: unknow default FFT module:' +
                                 nifty_configuration['fft_module'])
Jait Dixit's avatar
Jait Dixit committed
38
        else:
Theo Steininger's avatar
Theo Steininger committed
39
            if module == 'fftw':
Jait Dixit's avatar
Jait Dixit committed
40
                self._transform = FFTW(self.domain, self.codomain)
Theo Steininger's avatar
Theo Steininger committed
41
42
            elif module == 'numpy':
                self._transform = NUMPYFFT(self.domain, self.codomain)
Jait Dixit's avatar
Jait Dixit committed
43
44
            else:
                raise ValueError('ERROR: unknow FFT module:' + module)
Jait Dixit's avatar
Jait Dixit committed
45

46
    @classmethod
Martin Reinecke's avatar
Martin Reinecke committed
47
    def get_codomain(cls, domain, zerocenter=None):
Jait Dixit's avatar
Jait Dixit committed
48
49
50
51
52
53
54
55
56
        """
            Generates a compatible codomain to which transformations are
            reasonable, i.e.\  either a shifted grid or a Fourier conjugate
            grid.

            Parameters
            ----------
            domain: RGSpace
                Space for which a codomain is to be generated
Martin Reinecke's avatar
Martin Reinecke committed
57
            zerocenter : {bool, numpy.ndarray}, *optional*
Jait Dixit's avatar
Jait Dixit committed
58
59
60
61
62
63
64
65
66
                Whether or not the grid is zerocentered for each axis or not
                (default: None).

            Returns
            -------
            codomain : nifty.rg_space
                A compatible codomain.
        """
        if not isinstance(domain, RGSpace):
67
            raise TypeError("domain needs to be a RGSpace")
Jait Dixit's avatar
Jait Dixit committed
68

Martin Reinecke's avatar
Martin Reinecke committed
69
        # parse the zerocenter input
70
        if zerocenter is None:
71
            zerocenter = domain.zerocenter
Jait Dixit's avatar
Jait Dixit committed
72
73
        # if the input is something scalar, cast it to a boolean
        else:
74
            temp = np.empty_like(domain.zerocenter)
75
            temp[:] = zerocenter
76
            zerocenter = temp
Jait Dixit's avatar
Jait Dixit committed
77
78

        # calculate the initialization parameters
79
80
        distances = 1 / (np.array(domain.shape) *
                         np.array(domain.distances))
Jait Dixit's avatar
Jait Dixit committed
81

82
        new_space = RGSpace(domain.shape,
83
84
                            zerocenter=zerocenter,
                            distances=distances,
Martin Reinecke's avatar
Martin Reinecke committed
85
                            harmonic=(not domain.harmonic))
86
87

        # better safe than sorry
88
        cls.check_codomain(domain, new_space)
Jait Dixit's avatar
Jait Dixit committed
89
90
        return new_space

91
92
    @classmethod
    def check_codomain(cls, domain, codomain):
Jait Dixit's avatar
Jait Dixit committed
93
        if not isinstance(domain, RGSpace):
94
            raise TypeError("domain is not a RGSpace")
Jait Dixit's avatar
Jait Dixit committed
95
96

        if not isinstance(codomain, RGSpace):
97
            raise TypeError("domain is not a RGSpace")
Jait Dixit's avatar
Jait Dixit committed
98

99
100
        if not np.all(np.array(domain.shape) ==
                      np.array(codomain.shape)):
101
102
            raise AttributeError("The shapes of domain and codomain must be "
                                 "identical.")
Jait Dixit's avatar
Jait Dixit committed
103
104

        if domain.harmonic == codomain.harmonic:
105
106
            raise AttributeError("domain.harmonic and codomain.harmonic must "
                                 "not be the same.")
107

Jait Dixit's avatar
Jait Dixit committed
108
109
        # Check if the distances match, i.e. dist' = 1 / (num * dist)
        if not np.all(
110
111
112
            np.absolute(np.array(domain.shape) *
                        np.array(domain.distances) *
                        np.array(codomain.distances) - 1) <
113
                10**-7):
114
115
            raise AttributeError("The grid-distances of domain and codomain "
                                 "do not match.")
Jait Dixit's avatar
Jait Dixit committed
116

117
        super(RGRGTransformation, cls).check_codomain(domain, codomain)
Jait Dixit's avatar
Jait Dixit committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

    def transform(self, val, axes=None, **kwargs):
        """
        RG -> RG transform method.

        Parameters
        ----------
        val : np.ndarray or distributed_data_object
            The value array which is to be transformed

        axes : None or tuple
            The axes along which the transformation should take place

        """
        if self._transform.codomain.harmonic:
133
134
135
136
137
            # correct for forward fft.
            # naively one would set power to 0.5 here in order to
            # apply effectively a factor of 1/sqrt(N) to the field.
            # BUT: the pixel volumes of the domain and codomain are different.
            # Hence, in order to produce the same scalar product, power===1.
Jait Dixit's avatar
Jait Dixit committed
138
            val = self._transform.domain.weight(val, power=1, axes=axes)
Jait Dixit's avatar
Jait Dixit committed
139
140
141
142
143

        # Perform the transformation
        Tval = self._transform.transform(val, axes, **kwargs)

        if not self._transform.codomain.harmonic:
144
145
            # correct for inverse fft.
            # See discussion above.
Jait Dixit's avatar
Jait Dixit committed
146
            Tval = self._transform.codomain.weight(Tval, power=-1, axes=axes)
Jait Dixit's avatar
Jait Dixit committed
147
148

        return Tval