fft_operator_support.py 4.42 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
24
from ..field import Field
from ..spaces.gl_space import GLSpace
Martin Reinecke's avatar
Martin Reinecke committed
25
26


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


class RGRGTransformation(Transformation):
Martin Reinecke's avatar
Martin Reinecke committed
35
    def __init__(self, pdom, hdom, space):
Martin Reinecke's avatar
Martin Reinecke committed
36
        import pyfftw
Martin Reinecke's avatar
Martin Reinecke committed
37
        super(RGRGTransformation, self).__init__(pdom, hdom, space)
Martin Reinecke's avatar
Martin Reinecke committed
38
        pyfftw.interfaces.cache.enable()
Martin Reinecke's avatar
Martin Reinecke committed
39
40
41
42
43
44
45
        # 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
46
47
48
49
50

    @property
    def unitary(self):
        return True

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

        Parameters
        ----------
Martin Reinecke's avatar
Martin Reinecke committed
57
58
        x : Field
            The field to be transformed
Martin Reinecke's avatar
Martin Reinecke committed
59
        """
Martin Reinecke's avatar
Martin Reinecke committed
60
61
62
63
        axes = x.domain.axes[self.space]
        p2h = x.domain == self.pdom
        if p2h:
            Tval = Field(self.hdom, hartley(x.val, axes))
Martin Reinecke's avatar
Martin Reinecke committed
64
        else:
Martin Reinecke's avatar
Martin Reinecke committed
65
66
67
68
            Tval = Field(self.pdom, hartley(x.val, axes))
        fct = self.fct_p2h if p2h else self.fct_h2p
        if fct != 1:
            Tval *= fct
Martin Reinecke's avatar
Martin Reinecke committed
69

Martin Reinecke's avatar
Martin Reinecke committed
70
        return Tval
Martin Reinecke's avatar
Martin Reinecke committed
71
72
73


class SlicingTransformation(Transformation):
Martin Reinecke's avatar
Martin Reinecke committed
74
75
76
77
78
79
80
81
82
83
    def transform(self, x):
        p2h = x.domain == self.pdom
        if p2h:
            res = Field(self.hdom, dtype=x.dtype)

            for slice in utilities.get_slice_list(x.shape,
                                                  x.domain.axes[self.space]):
                res.val[slice] = self._slice_p2h(x.val[slice])
        else:
            res = Field(self.pdom, dtype=x.dtype)
Martin Reinecke's avatar
Martin Reinecke committed
84

Martin Reinecke's avatar
Martin Reinecke committed
85
86
87
88
            for slice in utilities.get_slice_list(x.shape,
                                                  x.domain.axes[self.space]):
                res.val[slice] = self._slice_h2p(x.val[slice])
        return res
Martin Reinecke's avatar
Martin Reinecke committed
89
90
91


def buildLm(nr, lmax):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
92
    res = np.empty((len(nr)+lmax+1)//2, dtype=(nr[0]*1j).dtype)
Martin Reinecke's avatar
Martin Reinecke committed
93
94
95
96
97
98
    res[0:lmax+1] = nr[0:lmax+1]
    res[lmax+1:] = np.sqrt(0.5)*(nr[lmax+1::2] + 1j*nr[lmax+2::2])
    return res


def buildIdx(nr, lmax):
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
99
100
101
102
103
    res = np.empty((lmax+1)*(lmax+1), dtype=nr[0].real.dtype)
    res[0:lmax+1] = nr[0:lmax+1].real
    res[lmax+1::2] = np.sqrt(2)*nr[lmax+1:].real
    res[lmax+2::2] = np.sqrt(2)*nr[lmax+1:].imag
    return res
Martin Reinecke's avatar
Martin Reinecke committed
104
105


Martin Reinecke's avatar
Martin Reinecke committed
106
107
108
class SphericalTransformation(SlicingTransformation):
    def __init__(self, pdom, hdom, space):
        super(SphericalTransformation, self).__init__(pdom, hdom, space)
Martin Reinecke's avatar
Martin Reinecke committed
109
110
        from pyHealpix import sharpjob_d

Martin Reinecke's avatar
Martin Reinecke committed
111
112
113
114
115
116
117
        self.lmax = self.hdom[self.space].lmax
        self.sjob = sharpjob_d()
        self.sjob.set_triangular_alm_info(self.hdom[self.space].lmax,
                                          self.hdom[self.space].mmax)
        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
118
        else:
Martin Reinecke's avatar
Martin Reinecke committed
119
            self.sjob.set_Healpix_geometry(self.pdom[self.space].nside)
Martin Reinecke's avatar
Martin Reinecke committed
120
121
122
123
124

    @property
    def unitary(self):
        return False

Martin Reinecke's avatar
Martin Reinecke committed
125
126
127
    def _slice_p2h(self, inp):
        rr = self.sjob.map2alm(inp)
        return buildIdx(rr, lmax=self.lmax)
Martin Reinecke's avatar
Martin Reinecke committed
128

Martin Reinecke's avatar
Martin Reinecke committed
129
130
131
    def _slice_h2p(self, inp):
        result = buildLm(inp, lmax=self.lmax)
        return self.sjob.alm2map(result)