hplmtransformation.py 4.6 KB
Newer Older
1
import numpy as np
Jait Dixit's avatar
Jait Dixit committed
2
from transformation import Transformation
3
4
5
from d2o import distributed_data_object
from nifty.config import dependency_injector as gdi
import nifty.nifty_utilities as utilities
Jait Dixit's avatar
Jait Dixit committed
6
from nifty import HPSpace, LMSpace
7

csongor's avatar
csongor committed
8
9
import lm_transformation_factory as ltf

10
11
12
hp = gdi.get('healpy')


Jait Dixit's avatar
Jait Dixit committed
13
class HPLMTransformation(Transformation):
Jait Dixit's avatar
Jait Dixit committed
14
    def __init__(self, domain, codomain=None, module=None):
15
16
17
        if 'healpy' not in gdi:
            raise ImportError("The module healpy is needed but not available")

Jait Dixit's avatar
Jait Dixit committed
18
19
20
21
        if codomain is None:
            self.domain = domain
            self.codomain = self.get_codomain(domain)
        elif self.check_codomain(domain, codomain):
Jait Dixit's avatar
Jait Dixit committed
22
23
24
25
26
            self.domain = domain
            self.codomain = codomain
        else:
            raise ValueError("ERROR: Incompatible codomain!")

Jait Dixit's avatar
Jait Dixit committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    @staticmethod
    def get_codomain(domain):
        """
            Generates a compatible codomain to which transformations are
            reasonable, i.e.\  an instance of the :py:class:`lm_space` class.

            Parameters
            ----------
            domain: HPSpace
                Space for which a codomain is to be generated

            Returns
            -------
            codomain : LMSpace
                A compatible codomain.
        """
        if domain is None:
            raise ValueError('ERROR: cannot generate codomain for None')

        if not isinstance(domain, HPSpace):
            raise TypeError('ERROR: domain needs to be a HPSpace')

49
        lmax = 3 * domain.nside - 1
Jait Dixit's avatar
Jait Dixit committed
50
51
52
        mmax = lmax
        return LMSpace(lmax=lmax, mmax=mmax, dtype=np.dtype('complex128'))

Jait Dixit's avatar
Jait Dixit committed
53
54
55
56
57
58
59
60
61
62
63
    @staticmethod
    def check_codomain(domain, codomain):
        if not isinstance(domain, HPSpace):
            raise TypeError('ERROR: domain is not a HPSpace')

        if codomain is None:
            return False

        if not isinstance(codomain, LMSpace):
            raise TypeError('ERROR: codomain must be a LMSpace.')

64
65
66
        nside = domain.nside
        lmax = codomain.lmax
        mmax = codomain.mmax
Jait Dixit's avatar
Jait Dixit committed
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85

        if (3 * nside - 1 != lmax) or (lmax != mmax):
            return False

        return True

    def transform(self, val, axes=None, **kwargs):
        """
        HP -> LM 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

        """
86
87
88
89
        # get by number of iterations from kwargs
        niter = kwargs['niter'] if 'niter' in kwargs else 0

        if self.domain.discrete:
Jait Dixit's avatar
Jait Dixit committed
90
            val = self.domain.weight(val, power=-0.5, axes=axes)
91
92

        # shorthands for transform parameters
93
94
        lmax = self.codomain.lmax
        mmax = self.codomain.mmax
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

        if isinstance(val, distributed_data_object):
            temp_val = val.get_full_data()
        else:
            temp_val = val

        return_val = None

        for slice_list in utilities.get_slice_list(temp_val.shape, axes):
            if slice_list == [slice(None, None)]:
                inp = temp_val
            else:
                if return_val is None:
                    return_val = np.empty_like(temp_val)
                inp = temp_val[slice_list]

csongor's avatar
csongor committed
111
112
113
114
115
116
117
118
119
            if self.domain.dtype == np.dtype('complex128'):
                inpReal = hp.map2alm(
                    np.real(inp).astype(np.float64, copy=False),
                    lmax=lmax, mmax=mmax, iter=niter, pol=True,
                    use_weights=False, datapath=None)
                inpImg = hp.map2alm(
                    np.imag(inp).astype(np.float64, copy=False),
                    lmax=lmax, mmax=mmax, iter=niter, pol=True,
                    use_weights=False, datapath=None)
120
121
                inpReal = ltf.buildIdx(inpReal, lmax=lmax)
                inpImg = ltf.buildIdx(inpImg, lmax=lmax)
csongor's avatar
csongor committed
122
123
124
125
126
                inp = inpReal + inpImg * 1j
            else:
                inp = hp.map2alm(inp.astype(np.float64, copy=False),
                                 lmax=lmax, mmax=mmax, iter=niter, pol=True,
                                 use_weights=False, datapath=None)
127
                inp = ltf.buildIdx(inp, lmax=lmax)
128
129
130
131
132
133
134
135
136
137
138
            if slice_list == [slice(None, None)]:
                return_val = inp
            else:
                return_val[slice_list] = inp

        if isinstance(val, distributed_data_object):
            new_val = val.copy_empty(dtype=self.codomain.dtype)
            new_val.set_full_data(return_val, copy=False)
        else:
            return_val = return_val.astype(self.codomain.dtype, copy=False)

Jait Dixit's avatar
Jait Dixit committed
139
        return return_val