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

Martin Reinecke's avatar
Martin Reinecke committed
19
#from builtins import str
Martin Reinecke's avatar
Martin Reinecke committed
20
import ast
theos's avatar
theos committed
21
22
import numpy as np

Martin Reinecke's avatar
Martin Reinecke committed
23
from d2o import distributed_data_object
24

Martin Reinecke's avatar
Martin Reinecke committed
25
from ..space import Space
Martin Reinecke's avatar
Martin Reinecke committed
26
from functools import reduce
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
    """ 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*
Martin Reinecke's avatar
Martin Reinecke committed
41
        True if logarithmic binning should be used (default : None).
Theo Steininger's avatar
Theo Steininger committed
42
43
44
45
46
    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*
Martin Reinecke's avatar
Martin Reinecke committed
47
48
49
        Boundaries between the power spectrum bins.
        (If binbounds has n entries, there will be n+1 bins, the first bin
        starting at -inf, the last bin ending at +inf.)
Theo Steininger's avatar
Theo Steininger committed
50
        (default : None)
Martin Reinecke's avatar
Martin Reinecke committed
51
        if binbounds == None:
Theo Steininger's avatar
Theo Steininger committed
52
53
            Calculates the bounds from the kindex while applying the
            logarithmic and nbin keywords.
Martin Reinecke's avatar
Martin Reinecke committed
54
55
56
57
58
    Note: if "bindounds" is not None, both "logarithmic" and "nbin" must be
        None!
    Note: if "binbounds", "logarithmic", and "nbin" are all None, then
        "natural" binning is performed, i.e. there will be one bin for every
        distinct k-vector length.
Theo Steininger's avatar
Theo Steininger committed
59
60
61
62

    Attributes
    ----------
    pindex : distributed_data_object
63
64
        This holds the information which pixel of the harmonic partner gets
        mapped to which power bin
Theo Steininger's avatar
Theo Steininger committed
65
    kindex : numpy.ndarray
66
        Sorted array of all k-modes.
Theo Steininger's avatar
Theo Steininger committed
67
68
69
    rho : numpy.ndarray
        The amount of k-modes that get mapped to one power bin is given by
        rho.
70
71
72
73
74
75
76
77
    dim : np.int
        Total number of dimensionality, i.e. the number of pixels.
    harmonic : bool
        Specifies whether the space is a signal or harmonic space.
    total_volume : np.float
        The total volume of the space.
    shape : tuple of np.ints
        The shape of the space's data array.
Martin Reinecke's avatar
Martin Reinecke committed
78
79
    binbounds :  tuple or None
        Boundaries between the power spectrum bins
Theo Steininger's avatar
Theo Steininger committed
80
81
82
83
84
85
86

    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.

    """
87

88
89
    _powerIndexCache = {}

90
91
    # ---Overwritten properties and methods---

92
    def __init__(self, harmonic_partner, distribution_strategy='not',
Martin Reinecke's avatar
Martin Reinecke committed
93
                 logarithmic=None, nbin=None, binbounds=None):
Martin Reinecke's avatar
Martin Reinecke committed
94
        super(PowerSpace, self).__init__()
Martin Reinecke's avatar
Martin Reinecke committed
95
        self._ignore_for_hash += ['_pindex', '_kindex', '_rho']
96

Martin Reinecke's avatar
Martin Reinecke committed
97
98
99
        if not (isinstance(harmonic_partner, Space) and
                harmonic_partner.harmonic):
            raise ValueError("harmonic_partner must be a harmonic space.")
100
        self._harmonic_partner = harmonic_partner
101

Martin Reinecke's avatar
Martin Reinecke committed
102
103
104
105
        # sanity check
        if binbounds is not None and not(nbin is None and logarithmic is None):
            raise ValueError(
                "if binbounds is defined, nbin and logarithmic must be None")
106

Martin Reinecke's avatar
Martin Reinecke committed
107
108
        if binbounds is not None:
            binbounds = tuple(binbounds)
109

Martin Reinecke's avatar
Martin Reinecke committed
110
111
        key = (harmonic_partner, distribution_strategy, logarithmic, nbin,
               binbounds)
112
113
114
115
116
117
118
119
120
121
        if self._powerIndexCache.get(key) is None:
            distance_array = \
                self.harmonic_partner.get_distance_array(distribution_strategy)
            temp_binbounds = self._compute_binbounds(
                                  harmonic_partner=self.harmonic_partner,
                                  distribution_strategy=distribution_strategy,
                                  logarithmic=logarithmic,
                                  nbin=nbin,
                                  binbounds=binbounds)
            temp_pindex = self._compute_pindex(
122
                                harmonic_partner=self.harmonic_partner,
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
                                distance_array=distance_array,
                                binbounds=temp_binbounds,
                                distribution_strategy=distribution_strategy)
            temp_rho = temp_pindex.bincount().get_full_data()
            temp_kindex = \
                (temp_pindex.bincount(weights=distance_array).get_full_data() /
                 temp_rho)
            self._powerIndexCache[key] = (temp_binbounds,
                                          temp_pindex,
                                          temp_kindex,
                                          temp_rho)

        (self._binbounds, self._pindex, self._kindex, self._rho) = \
            self._powerIndexCache[key]

138
139
    @staticmethod
    def _compute_binbounds(harmonic_partner, distribution_strategy,
140
                           logarithmic, nbin, binbounds):
141

Martin Reinecke's avatar
Martin Reinecke committed
142
        if logarithmic is None and nbin is None and binbounds is None:
143
            result = None
Martin Reinecke's avatar
Martin Reinecke committed
144
145
146
147
148
149
150
151
152
153
154
        else:
            if binbounds is not None:
                bb = np.sort(np.array(binbounds))
            else:
                if logarithmic is not None:
                    logarithmic = bool(logarithmic)
                if nbin is not None:
                    nbin = int(nbin)

                # equidistant binning (linear or log)
                # MR FIXME: this needs to improve
155
                kindex = harmonic_partner.get_unique_distances()
Martin Reinecke's avatar
Martin Reinecke committed
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
                if (logarithmic):
                    k = np.r_[0, np.log(kindex[1:])]
                else:
                    k = kindex
                dk = np.max(k[2:] - k[1:-1])  # minimum dk to avoid empty bins
                if(nbin is None):
                    nbin = int((k[-1] - 0.5 * (k[2] + k[1])) /
                               dk - 0.5)  # maximal nbin
                else:
                    nbin = min(int(nbin), int(
                        (k[-1] - 0.5 * (k[2] + k[1])) / dk + 2.5))
                    dk = (k[-1] - 0.5 * (k[2] + k[1])) / (nbin - 2.5)
                bb = np.r_[0.5 * (3 * k[1] - k[2]),
                           0.5 * (k[1] + k[2]) + dk * np.arange(nbin-2)]
                if(logarithmic):
                    bb = np.exp(bb)
172
173
            result = tuple(bb)
        return result
174

175
176
    @staticmethod
    def _compute_pindex(harmonic_partner, distance_array, binbounds,
177
                        distribution_strategy):
178

Martin Reinecke's avatar
Martin Reinecke committed
179
        # Compute pindex, kindex and rho according to bb
180
181
182
183
184
        pindex = distributed_data_object(
                                global_shape=distance_array.shape,
                                dtype=np.int,
                                distribution_strategy=distribution_strategy)
        if binbounds is None:
185
            binbounds = harmonic_partner.get_natural_binbounds()
186
187
188
        pindex.set_local_data(
                np.searchsorted(binbounds, distance_array.get_local_data()))
        return pindex
189

Theo Steininger's avatar
Theo Steininger committed
190
    def pre_cast(self, x, axes):
Theo Steininger's avatar
Theo Steininger committed
191
192
193
194
195
196
        """ 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.

197
198
199
        Parameters
        ----------
        x : {array-like, function array-like -> array-like}
Theo Steininger's avatar
Theo Steininger committed
200
201
202
203
204
205
            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.

206
207
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
208
209
210
        array-like
            discretized power spectrum

211
        """
Theo Steininger's avatar
Theo Steininger committed
212

Martin Reinecke's avatar
Martin Reinecke committed
213
        return x(self.kindex) if callable(x) else x
214

215
216
    # ---Mandatory properties and methods---

217
218
    def __repr__(self):
        return ("PowerSpace(harmonic_partner=%r, distribution_strategy=%r, "
Martin Reinecke's avatar
Martin Reinecke committed
219
                "binbounds=%r)"
220
                % (self.harmonic_partner, self.pindex.distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
221
                   self._binbounds))
222

223
224
225
    @property
    def harmonic(self):
        return True
226

227
228
    @property
    def shape(self):
229
        return self.kindex.shape
230

231
232
233
234
235
236
237
    @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
238
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
239
240

    def copy(self):
241
        distribution_strategy = self.pindex.distribution_strategy
242
        return self.__class__(harmonic_partner=self.harmonic_partner,
243
                              distribution_strategy=distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
244
                              binbounds=self._binbounds)
245

Martin Reinecke's avatar
Martin Reinecke committed
246
    def weight(self, x, power, axes, inplace=False):
Jait Dixit's avatar
Jait Dixit committed
247
248
        reshaper = [1, ] * len(x.shape)
        # we know len(axes) is always 1
249
250
        reshaper[axes[0]] = self.shape[0]

251
        weight = self.rho.reshape(reshaper)
252
        if power != 1:
253
            weight = weight ** np.float(power)
254
255
256
257
258
259

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
260
261
262

        return result_x

263
    def get_distance_array(self, distribution_strategy):
Martin Reinecke's avatar
Martin Reinecke committed
264
        return distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
265
                                self.kindex, dtype=np.float64,
266
                                distribution_strategy=distribution_strategy)
theos's avatar
theos committed
267

268
    def get_fft_smoothing_kernel_function(self, sigma):
269
        raise NotImplementedError(
270
            "There is no fft smoothing function for PowerSpace.")
theos's avatar
theos committed
271

272
273
274
    # ---Added properties and methods---

    @property
275
    def harmonic_partner(self):
Theo Steininger's avatar
Theo Steininger committed
276
        """ Returns the Space of which this is the power space.
277
278
        """
        return self._harmonic_partner
279
280

    @property
Martin Reinecke's avatar
Martin Reinecke committed
281
282
    def binbounds(self):
        return self._binbounds
283
284
285

    @property
    def pindex(self):
286
        """ A distributed_data_object having the shape of the harmonic partner
Theo Steininger's avatar
Theo Steininger committed
287
288
        space containing the indices of the power bin a pixel belongs to.
        """
289
290
291
292
        return self._pindex

    @property
    def kindex(self):
Theo Steininger's avatar
Theo Steininger committed
293
294
        """ Sorted array of all k-modes.
        """
295
296
297
298
        return self._kindex

    @property
    def rho(self):
Theo Steininger's avatar
Theo Steininger committed
299
300
        """Degeneracy factor of the individual k-vectors.
        """
301
        return self._rho
302

303
304
305
    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Martin Reinecke's avatar
Martin Reinecke committed
306
307
308
309
        hdf5_group.attrs['binbounds'] = str(self._binbounds)
        hdf5_group.attrs['distribution_strategy'] = \
            self._pindex.distribution_strategy

310
        return {
311
            'harmonic_partner': self.harmonic_partner,
312
313
314
        }

    @classmethod
Theo Steininger's avatar
Theo Steininger committed
315
    def _from_hdf5(cls, hdf5_group, repository):
Martin Reinecke's avatar
Martin Reinecke committed
316
        hp = repository.get('harmonic_partner', hdf5_group)
Martin Reinecke's avatar
Martin Reinecke committed
317
        bb = ast.literal_eval(hdf5_group.attrs['binbounds'])
Martin Reinecke's avatar
Martin Reinecke committed
318
319
        ds = hdf5_group.attrs['distribution_strategy']
        return PowerSpace(hp, ds, binbounds=bb)