power_space.py 5.52 KB
Newer Older
theos's avatar
theos committed
1
2
# -*- coding: utf-8 -*-

theos's avatar
theos committed
3
4
import numpy as np

5
6
from keepers import Versionable

7
8
import d2o

9
10
from power_index_factory import PowerIndexFactory

11
from nifty.spaces.space import Space
12
from nifty.spaces.rg_space import RGSpace
13
from nifty.nifty_utilities import cast_axis_to_tuple
theos's avatar
theos committed
14
15


16
17
class PowerSpace(Versionable, Space):

18
19
    # ---Overwritten properties and methods---

20
21
    def __init__(self, harmonic_domain=RGSpace((1,)),
                 distribution_strategy='not',
Jait Dixit's avatar
Jait Dixit committed
22
                 log=False, nbin=None, binbounds=None,
23
24
25
                 dtype=np.dtype('float')):

        super(PowerSpace, self).__init__(dtype)
26
27
        self._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
                                  '_k_array']
28
29

        if not isinstance(harmonic_domain, Space):
30
31
            raise ValueError(
                "harmonic_domain must be a Space.")
32
        if not harmonic_domain.harmonic:
33
34
            raise ValueError(
                "harmonic_domain must be a harmonic space.")
35
36
        self._harmonic_domain = harmonic_domain

Jait Dixit's avatar
Jait Dixit committed
37
38
39
40
41
42
        power_index = PowerIndexFactory.get_power_index(
                        domain=self.harmonic_domain,
                        distribution_strategy=distribution_strategy,
                        log=log,
                        nbin=nbin,
                        binbounds=binbounds)
43
44
45
46
47
48
49
50
51

        config = power_index['config']
        self._log = config['log']
        self._nbin = config['nbin']
        self._binbounds = config['binbounds']

        self._pindex = power_index['pindex']
        self._kindex = power_index['kindex']
        self._rho = power_index['rho']
52
53
        self._pundex = power_index['pundex']
        self._k_array = power_index['k_array']
54

55
56
57
58
59
60
    def pre_cast(self, x, axes=None):
        if callable(x):
            return x(self.kindex)
        else:
            return x

61
62
63
64
65
    # ---Mandatory properties and methods---

    @property
    def harmonic(self):
        return True
66

67
68
    @property
    def shape(self):
69
        return self.kindex.shape
70

71
72
73
74
75
76
77
    @property
    def dim(self):
        return self.shape[0]

    @property
    def total_volume(self):
        # every power-pixel has a volume of 1
78
79
80
        return reduce(lambda x, y: x*y, self.pindex.shape)

    def copy(self):
81
        distribution_strategy = self.pindex.distribution_strategy
82
        return self.__class__(harmonic_domain=self.harmonic_domain,
83
                              distribution_strategy=distribution_strategy,
84
85
86
87
                              log=self.log,
                              nbin=self.nbin,
                              binbounds=self.binbounds,
                              dtype=self.dtype)
88

89
    def weight(self, x, power=1, axes=None, inplace=False):
90
91
92
93
        total_shape = x.shape

        axes = cast_axis_to_tuple(axes, len(total_shape))
        if len(axes) != 1:
94
95
            raise ValueError(
                "axes must be of length 1.")
96
97
98
99

        reshaper = [1, ] * len(total_shape)
        reshaper[axes[0]] = self.shape[0]

100
        weight = self.rho.reshape(reshaper)
101
102
        if power != 1:
            weight = weight ** power
103
104
105
106
107
108

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
109
110
111

        return result_x

112
    def get_distance_array(self, distribution_strategy):
113
114
115
116
        result = d2o.distributed_data_object(
                                self.kindex,
                                distribution_strategy=distribution_strategy)
        return result
theos's avatar
theos committed
117

118
    def get_fft_smoothing_kernel_function(self, sigma):
119
        raise NotImplementedError(
120
            "There is no fft smoothing function for PowerSpace.")
theos's avatar
theos committed
121

122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
    # ---Added properties and methods---

    @property
    def harmonic_domain(self):
        return self._harmonic_domain

    @property
    def log(self):
        return self._log

    @property
    def nbin(self):
        return self._nbin

    @property
    def binbounds(self):
        return self._binbounds

    @property
    def pindex(self):
        return self._pindex

    @property
    def kindex(self):
        return self._kindex

    @property
    def rho(self):
        return self._rho
151

152
153
154
155
156
157
158
    @property
    def pundex(self):
        return self._pundex

    @property
    def k_array(self):
        return self._k_array
159
160
161
162

    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Jait Dixit's avatar
Jait Dixit committed
163
164
165
        hdf5_group['kindex'] = self.kindex
        hdf5_group['rho'] = self.rho
        hdf5_group['pundex'] = self.pundex
166
167
168
169
        hdf5_group['dtype'] = self.dtype.name
        hdf5_group['log'] = self.log
        hdf5_group['nbin'] = str(self.nbin)
        hdf5_group['binbounds'] = str(self.binbounds)
170
171
172
173
174
175
176
177
178

        return {
            'harmonic_domain': self.harmonic_domain,
            'pindex': self.pindex,
            'k_array': self.k_array
        }

    @classmethod
    def _from_hdf5(cls, hdf5_group, loopback_get):
Jait Dixit's avatar
Jait Dixit committed
179
180
181
182
183
        # make an empty PowerSpace object
        new_ps = EmptyPowerSpace()
        # reset class
        new_ps.__class__ = cls
        # set all values
184
        new_ps.dtype = np.dtype(hdf5_group['dtype'][()])
Jait Dixit's avatar
Jait Dixit committed
185
        new_ps._harmonic_domain = loopback_get('harmonic_domain')
186
187
188
        new_ps._log = hdf5_group['log'][()]
        exec('new_ps._nbin = ' + hdf5_group['nbin'][()])
        exec('new_ps._binbounds = ' + hdf5_group['binbounds'][()])
Jait Dixit's avatar
Jait Dixit committed
189
190
191
192
193
194
195
196
197
198
199
200
201

        new_ps._pindex = loopback_get('pindex')
        new_ps._kindex = hdf5_group['kindex'][:]
        new_ps._rho = hdf5_group['rho'][:]
        new_ps._pundex = hdf5_group['pundex'][:]
        new_ps._k_array = loopback_get('k_array')

        return new_ps


class EmptyPowerSpace(PowerSpace):
    def __init__(self):
        pass