fft_operator_support.py 6.6 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
# 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/>.
#
# 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.

from __future__ import division
import numpy as np
Martin Reinecke's avatar
Martin Reinecke committed
21
22
from .. import nifty_utilities as utilities
from ..low_level_library import hartley
Martin Reinecke's avatar
Martin Reinecke committed
23
from .. import dobj
Martin Reinecke's avatar
Martin Reinecke committed
24
25
from ..field import Field
from ..spaces.gl_space import GLSpace
Martin Reinecke's avatar
Martin Reinecke committed
26

Martin Reinecke's avatar
Martin Reinecke committed
27

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
28
class Transformation(object):
Martin Reinecke's avatar
Martin Reinecke committed
29
30
31
32
    def __init__(self, pdom, hdom, space):
        self.pdom = pdom
        self.hdom = hdom
        self.space = space
Martin Reinecke's avatar
Martin Reinecke committed
33
34
35


class RGRGTransformation(Transformation):
Martin Reinecke's avatar
Martin Reinecke committed
36
    def __init__(self, pdom, hdom, space):
Martin Reinecke's avatar
Martin Reinecke committed
37
        import pyfftw
Martin Reinecke's avatar
Martin Reinecke committed
38
        super(RGRGTransformation, self).__init__(pdom, hdom, space)
Martin Reinecke's avatar
Martin Reinecke committed
39
        pyfftw.interfaces.cache.enable()
Martin Reinecke's avatar
Martin Reinecke committed
40
41
42
43
44
45
46
        # correct for forward/inverse 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.
        self.fct_p2h = pdom[space].scalar_dvol()
        self.fct_h2p = 1./(pdom[space].scalar_dvol()*hdom[space].dim)
Martin Reinecke's avatar
Martin Reinecke committed
47
48
49
50
51

    @property
    def unitary(self):
        return True

Martin Reinecke's avatar
Martin Reinecke committed
52
    def transform(self, x):
Martin Reinecke's avatar
Martin Reinecke committed
53
54
55
56
57
        """
        RG -> RG transform method.

        Parameters
        ----------
Martin Reinecke's avatar
Martin Reinecke committed
58
59
        x : Field
            The field to be transformed
Martin Reinecke's avatar
Martin Reinecke committed
60
        """
Martin Reinecke's avatar
Martin Reinecke committed
61
        from pyfftw.interfaces.numpy_fft import fftn
Martin Reinecke's avatar
Martin Reinecke committed
62
63
        axes = x.domain.axes[self.space]
        p2h = x.domain == self.pdom
Martin Reinecke's avatar
Martin Reinecke committed
64
        tdom = self.hdom if p2h else self.pdom
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
65
66
        if dobj.distaxis(x.val) in axes:
            tmpax = (dobj.distaxis(x.val),)
Martin Reinecke's avatar
Martin Reinecke committed
67
68
            tmp = dobj.redistribute(x.val, nodist=tmpax)
            ldat = dobj.local_data(tmp)
Martin Reinecke's avatar
Martin Reinecke committed
69
            if len(axes) == 1:  # only one transform needed
Martin Reinecke's avatar
Martin Reinecke committed
70
                ldat = hartley(ldat, axes=tmpax)
Martin Reinecke's avatar
Martin Reinecke committed
71
                tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=dobj.distaxis(tmp))
Martin Reinecke's avatar
Martin Reinecke committed
72
                tmp = dobj.redistribute(tmp, dist=tmpax[0])
Martin Reinecke's avatar
Martin Reinecke committed
73
74
            else:  # two separate transforms
                ldat = fftn(ldat, axes=tmpax)
Martin Reinecke's avatar
Martin Reinecke committed
75
                tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=dobj.distaxis(tmp))
Martin Reinecke's avatar
Martin Reinecke committed
76
                tmp = dobj.redistribute(tmp, dist=tmpax[0])
Martin Reinecke's avatar
Martin Reinecke committed
77
                tmpax = tuple(i for i in axes if i not in tmpax)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
78
                ldat = dobj.local_data(tmp)
Martin Reinecke's avatar
Martin Reinecke committed
79
80
                ldat = fftn(ldat, axes=tmpax)
                ldat = ldat.real+ldat.imag
Martin Reinecke's avatar
Martin Reinecke committed
81
                tmp = dobj.from_local_data(tmp.shape, ldat, distaxis=dobj.distaxis(tmp))
Martin Reinecke's avatar
Martin Reinecke committed
82
            Tval = Field(tdom, tmp)
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
83
        else:
Martin Reinecke's avatar
Martin Reinecke committed
84
            ldat = dobj.local_data(x.val)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
85
86
            # these two alternatives are equivalent, with the second being faster
            if False:
Martin Reinecke's avatar
Martin Reinecke committed
87
                ldat = fftn(ldat, axes=axes)
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
88
89
                ldat = ldat.real+ldat.imag
            else:
Martin Reinecke's avatar
Martin Reinecke committed
90
                ldat = hartley(ldat, axes=axes)
Martin Reinecke's avatar
Martin Reinecke committed
91
            tmp = dobj.from_local_data(x.val.shape, ldat, distaxis=dobj.distaxis(x.val))
Martin Reinecke's avatar
Martin Reinecke committed
92
            Tval = Field(tdom, tmp)
Martin Reinecke's avatar
Martin Reinecke committed
93
94
95
        fct = self.fct_p2h if p2h else self.fct_h2p
        if fct != 1:
            Tval *= fct
Martin Reinecke's avatar
Martin Reinecke committed
96

Martin Reinecke's avatar
Martin Reinecke committed
97
        return Tval
Martin Reinecke's avatar
Martin Reinecke committed
98
99


100
class SphericalTransformation(Transformation):
Martin Reinecke's avatar
Martin Reinecke committed
101
102
    def __init__(self, pdom, hdom, space):
        super(SphericalTransformation, self).__init__(pdom, hdom, space)
Martin Reinecke's avatar
Martin Reinecke committed
103
104
        from pyHealpix import sharpjob_d

Martin Reinecke's avatar
Martin Reinecke committed
105
        self.lmax = self.hdom[self.space].lmax
106
        self.mmax = self.hdom[self.space].mmax
Martin Reinecke's avatar
Martin Reinecke committed
107
        self.sjob = sharpjob_d()
108
        self.sjob.set_triangular_alm_info(self.lmax, self.mmax)
Martin Reinecke's avatar
Martin Reinecke committed
109
110
111
        if isinstance(self.pdom[self.space], GLSpace):
            self.sjob.set_Gauss_geometry(self.pdom[self.space].nlat,
                                         self.pdom[self.space].nlon)
Martin Reinecke's avatar
Martin Reinecke committed
112
        else:
Martin Reinecke's avatar
Martin Reinecke committed
113
            self.sjob.set_Healpix_geometry(self.pdom[self.space].nside)
Martin Reinecke's avatar
Martin Reinecke committed
114
115
116
117
118

    @property
    def unitary(self):
        return False

Martin Reinecke's avatar
Martin Reinecke committed
119
120
    def _slice_p2h(self, inp):
        rr = self.sjob.map2alm(inp)
121
122
123
124
125
126
127
        assert len(rr) == ((self.mmax+1)*(self.mmax+2))//2 + \
                          (self.mmax+1)*(self.lmax-self.mmax)
        res = np.empty(2*len(rr)-self.lmax-1, dtype=rr[0].real.dtype)
        res[0:self.lmax+1] = rr[0:self.lmax+1].real
        res[self.lmax+1::2] = np.sqrt(2)*rr[self.lmax+1:].real
        res[self.lmax+2::2] = np.sqrt(2)*rr[self.lmax+1:].imag
        return res
Martin Reinecke's avatar
Martin Reinecke committed
128

Martin Reinecke's avatar
Martin Reinecke committed
129
    def _slice_h2p(self, inp):
130
131
132
133
134
135
136
137
138
139
        res = np.empty((len(inp)+self.lmax+1)//2, dtype=(inp[0]*1j).dtype)
        assert len(res) == ((self.mmax+1)*(self.mmax+2))//2 + \
                           (self.mmax+1)*(self.lmax-self.mmax)
        res[0:self.lmax+1] = inp[0:self.lmax+1]
        res[self.lmax+1:] = np.sqrt(0.5)*(inp[self.lmax+1::2] +
                                          1j*inp[self.lmax+2::2])
        return self.sjob.alm2map(res)

    def transform(self, x):
        axes = x.domain.axes[self.space]
Martin Reinecke's avatar
Martin Reinecke committed
140
141
        axis = axes[0]
        tval = x.val
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
142
        if dobj.distaxis(tval) == axis:
Martin Reinecke's avatar
Martin Reinecke committed
143
            tval = dobj.redistribute(tval, nodist=(axis,))
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
144
        distaxis = dobj.distaxis(tval)
Martin Reinecke's avatar
Martin Reinecke committed
145

Martin Reinecke's avatar
Martin Reinecke committed
146
        p2h = x.domain == self.pdom
Martin Reinecke's avatar
Martin Reinecke committed
147
        idat = dobj.local_data(tval)
148
        if p2h:
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
149
            odat = np.empty(dobj.local_shape(self.hdom.shape, distaxis=distaxis), dtype=x.dtype)
Martin Reinecke's avatar
Martin Reinecke committed
150
151
152

            for slice in utilities.get_slice_list(idat.shape, axes):
                odat[slice] = self._slice_p2h(idat[slice])
Martin Reinecke's avatar
Martin Reinecke committed
153
            odat = dobj.from_local_data(self.hdom.shape, odat, distaxis)
Martin Reinecke's avatar
Martin Reinecke committed
154
            if distaxis != dobj.distaxis(x.val):
Martin Reinecke's avatar
Martin Reinecke committed
155
                odat = dobj.redistribute(odat, dist=dobj.distaxis(x.val))
Martin Reinecke's avatar
Martin Reinecke committed
156
            return Field(self.hdom, odat)
157
        else:
Martin Reinecke's avatar
Martin Reinecke committed
158
            odat = np.empty(dobj.local_shape(self.pdom.shape, distaxis=distaxis), dtype=x.dtype)
Martin Reinecke's avatar
Martin Reinecke committed
159
160
            for slice in utilities.get_slice_list(idat.shape, axes):
                odat[slice] = self._slice_h2p(idat[slice])
Martin Reinecke's avatar
Martin Reinecke committed
161
            odat = dobj.from_local_data(self.pdom.shape, odat, distaxis)
Martin Reinecke's avatar
Martin Reinecke committed
162
            if distaxis != dobj.distaxis(x.val):
Martin Reinecke's avatar
Martin Reinecke committed
163
164
                odat = dobj.redistribute(odat, dist=dobj.distaxis(x.val))
            return Field(self.pdom, odat)