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

theos's avatar
theos committed
19
20
import numpy as np

21
22
import d2o

23
24
from power_index_factory import PowerIndexFactory

25
from nifty.spaces.space import Space
26
from nifty.spaces.rg_space import RGSpace
theos's avatar
theos committed
27
28


Theo Steininger's avatar
Theo Steininger committed
29
class PowerSpace(Space):
Theo Steininger's avatar
Theo Steininger committed
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
    """ NIFTY class for spaces of power spectra.

    Parameters
    ----------
    harmonic_partner : Space
        The harmonic Space of which this is the power space.
    distribution_strategy : str *optional*
        The distribution strategy used for the distributed_data_objects
        derived from this PowerSpace, e.g. the pindex.
        (default : 'not')
    logarithmic : bool *optional*
        True if logarithmic binning should be used (default : False).
    nbin : {int, None} *optional*
        The number of bins that should be used for power spectrum binning
        (default : None).
        if nbin == None, then nbin is set to the length of kindex.
    binbounds :  {list, array-like} *optional*
        Array-like inner boundaries of the used bins of the default
        indices.
        (default : None)
        if binbounds == None :
            Calculates the bounds from the kindex while applying the
            logarithmic and nbin keywords.

    Attributes
    ----------
    pindex : distributed_data_object
        TODO add description
    kindex : numpy.ndarray
        TODO add description
    pundex : numpy.ndarray
        TODO add description
    rho : numpy.ndarray
        The amount of k-modes that get mapped to one power bin is given by
        rho.


    Notes
    -----
    A power space is the result of a projection of a harmonic space where
    k-modes of equal length get mapped to one power index.

    """
73

74
75
    # ---Overwritten properties and methods---

76
    def __init__(self, harmonic_partner=RGSpace((1,)),
77
                 distribution_strategy='not',
78
                 logarithmic=False, nbin=None, binbounds=None):
Martin Reinecke's avatar
Martin Reinecke committed
79
        super(PowerSpace, self).__init__()
80
81
        self._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
                                  '_k_array']
82

83
        if not isinstance(harmonic_partner, Space):
84
            raise ValueError(
85
86
                "harmonic_partner must be a Space.")
        if not harmonic_partner.harmonic:
87
            raise ValueError(
88
89
                "harmonic_partner must be a harmonic space.")
        self._harmonic_partner = harmonic_partner
90

Jait Dixit's avatar
Jait Dixit committed
91
        power_index = PowerIndexFactory.get_power_index(
92
                        domain=self.harmonic_partner,
Jait Dixit's avatar
Jait Dixit committed
93
                        distribution_strategy=distribution_strategy,
94
                        logarithmic=logarithmic,
Jait Dixit's avatar
Jait Dixit committed
95
96
                        nbin=nbin,
                        binbounds=binbounds)
97
98

        config = power_index['config']
99
        self._logarithmic = config['logarithmic']
100
101
102
103
104
105
        self._nbin = config['nbin']
        self._binbounds = config['binbounds']

        self._pindex = power_index['pindex']
        self._kindex = power_index['kindex']
        self._rho = power_index['rho']
106
107
        self._pundex = power_index['pundex']
        self._k_array = power_index['k_array']
108

Theo Steininger's avatar
Theo Steininger committed
109
    def pre_cast(self, x, axes):
Theo Steininger's avatar
Theo Steininger committed
110
111
112
113
114
115
        """ Casts power spectrum functions to discretized power spectra.

        This function takes an array or a function. If it is an array it does
        nothing, otherwise it interpretes the function as power spectrum and
        evaluates it at every k-mode.

116
117
118
        Parameters
        ----------
        x : {array-like, function array-like -> array-like}
Theo Steininger's avatar
Theo Steininger committed
119
120
121
122
123
124
            power spectrum given either in discretized form or implicitly as a
            function
        axes : tuple of ints
            Specifies the axes of x which correspond to this space. For
            explicifying the power spectrum function, this is ignored.

125
126
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
127
128
129
        array-like
            discretized power spectrum

130
        """
Theo Steininger's avatar
Theo Steininger committed
131

132
133
134
135
136
        if callable(x):
            return x(self.kindex)
        else:
            return x

137
138
139
140
141
    # ---Mandatory properties and methods---

    @property
    def harmonic(self):
        return True
142

143
144
    @property
    def shape(self):
145
        return self.kindex.shape
146

147
148
149
150
151
152
153
    @property
    def dim(self):
        return self.shape[0]

    @property
    def total_volume(self):
        # every power-pixel has a volume of 1
Jait Dixit's avatar
Jait Dixit committed
154
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
155
156

    def copy(self):
157
        distribution_strategy = self.pindex.distribution_strategy
158
        return self.__class__(harmonic_partner=self.harmonic_partner,
159
                              distribution_strategy=distribution_strategy,
160
                              logarithmic=self.logarithmic,
161
                              nbin=self.nbin,
Martin Reinecke's avatar
Martin Reinecke committed
162
                              binbounds=self.binbounds)
163

164
    def weight(self, x, power=1, axes=None, inplace=False):
Jait Dixit's avatar
Jait Dixit committed
165
166
        reshaper = [1, ] * len(x.shape)
        # we know len(axes) is always 1
167
168
        reshaper[axes[0]] = self.shape[0]

169
        weight = self.rho.reshape(reshaper)
170
171
        if power != 1:
            weight = weight ** power
172
173
174
175
176
177

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
178
179
180

        return result_x

181
    def get_distance_array(self, distribution_strategy):
182
        result = d2o.distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
183
                                self.kindex, dtype=np.float64,
184
185
                                distribution_strategy=distribution_strategy)
        return result
theos's avatar
theos committed
186

187
    def get_fft_smoothing_kernel_function(self, sigma):
188
        raise NotImplementedError(
189
            "There is no fft smoothing function for PowerSpace.")
theos's avatar
theos committed
190

191
192
193
    # ---Added properties and methods---

    @property
194
    def harmonic_partner(self):
Theo Steininger's avatar
Theo Steininger committed
195
        """ Returns the Space of which this is the power space.
196
197
        """
        return self._harmonic_partner
198
199

    @property
200
    def logarithmic(self):
Theo Steininger's avatar
Theo Steininger committed
201
        """ Returns True if logarithmic binning is used.
202
203
        """
        return self._logarithmic
204
205
206

    @property
    def nbin(self):
Theo Steininger's avatar
Theo Steininger committed
207
        """ Returns the number of power bins if specfied during initialization.
208
        """
209
210
211
212
        return self._nbin

    @property
    def binbounds(self):
Theo Steininger's avatar
Theo Steininger committed
213
        """ Inner boundaries of the used bins if specfied during initialization.
214
        """
215
216
217
218
        return self._binbounds

    @property
    def pindex(self):
Theo Steininger's avatar
Theo Steininger committed
219
220
221
        """ A distributed_data_objects having the shape of the harmonic partner
        space containing the indices of the power bin a pixel belongs to.
        """
222
223
224
225
        return self._pindex

    @property
    def kindex(self):
Theo Steininger's avatar
Theo Steininger committed
226
227
        """ Sorted array of all k-modes.
        """
228
229
230
231
        return self._kindex

    @property
    def rho(self):
Theo Steininger's avatar
Theo Steininger committed
232
233
        """Degeneracy factor of the individual k-vectors.
        """
234
        return self._rho
235

236
237
    @property
    def pundex(self):
Theo Steininger's avatar
Theo Steininger committed
238
239
240
241
        """ An array for which the n-th entry gives the flat index of the
        first occurence of a k-vector with length==kindex[n] in the
        k_array.
        """
242
243
244
245
        return self._pundex

    @property
    def k_array(self):
Theo Steininger's avatar
Theo Steininger committed
246
247
248
        """ An array containing distances to the grid center (i.e. zero-mode)
        for every k-mode in the grid of the harmonic partner space.
        """
249
        return self._k_array
250
251
252
253

    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Jait Dixit's avatar
Jait Dixit committed
254
255
256
        hdf5_group['kindex'] = self.kindex
        hdf5_group['rho'] = self.rho
        hdf5_group['pundex'] = self.pundex
257
        hdf5_group['logarithmic'] = self.logarithmic
Theo Steininger's avatar
Theo Steininger committed
258
259
260
        # Store nbin as string, since it can be None
        hdf5_group.attrs['nbin'] = str(self.nbin)
        hdf5_group.attrs['binbounds'] = str(self.binbounds)
261
262

        return {
263
            'harmonic_partner': self.harmonic_partner,
264
265
266
267
268
            'pindex': self.pindex,
            'k_array': self.k_array
        }

    @classmethod
Theo Steininger's avatar
Theo Steininger committed
269
    def _from_hdf5(cls, hdf5_group, repository):
Jait Dixit's avatar
Jait Dixit committed
270
271
272
273
        # make an empty PowerSpace object
        new_ps = EmptyPowerSpace()
        # reset class
        new_ps.__class__ = cls
Jait Dixit's avatar
Jait Dixit committed
274
        # call instructor so that classes are properly setup
Martin Reinecke's avatar
Martin Reinecke committed
275
        super(PowerSpace, new_ps).__init__()
Jait Dixit's avatar
Jait Dixit committed
276
        # set all values
Theo Steininger's avatar
Theo Steininger committed
277
278
        new_ps._harmonic_partner = repository.get('harmonic_partner',
                                                  hdf5_group)
279
        new_ps._logarithmic = hdf5_group['logarithmic'][()]
Theo Steininger's avatar
Theo Steininger committed
280
281
        exec('new_ps._nbin = ' + hdf5_group.attrs['nbin'])
        exec('new_ps._binbounds = ' + hdf5_group.attrs['binbounds'])
Jait Dixit's avatar
Jait Dixit committed
282

Theo Steininger's avatar
Theo Steininger committed
283
        new_ps._pindex = repository.get('pindex', hdf5_group)
Jait Dixit's avatar
Jait Dixit committed
284
285
286
        new_ps._kindex = hdf5_group['kindex'][:]
        new_ps._rho = hdf5_group['rho'][:]
        new_ps._pundex = hdf5_group['pundex'][:]
Theo Steininger's avatar
Theo Steininger committed
287
        new_ps._k_array = repository.get('k_array', hdf5_group)
Jait Dixit's avatar
Jait Dixit committed
288
        new_ps._ignore_for_hash += ['_pindex', '_kindex', '_rho', '_pundex',
289
                                    '_k_array']
Jait Dixit's avatar
Jait Dixit committed
290
291
292
293
294
295
296

        return new_ps


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