power_space.py 11.8 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

23
24
from d2o import distributed_data_object,\
    STRATEGIES as DISTRIBUTION_STRATEGIES
25

Martin Reinecke's avatar
Martin Reinecke committed
26
from ...spaces.space import Space
Martin Reinecke's avatar
Martin Reinecke committed
27
from functools import reduce
Martin Reinecke's avatar
Martin Reinecke committed
28
from ...config import nifty_configuration as gc
theos's avatar
theos committed
29
30


Theo Steininger's avatar
Theo Steininger committed
31
class PowerSpace(Space):
Theo Steininger's avatar
Theo Steininger committed
32
33
34
35
36
37
38
39
40
41
42
    """ 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
43
        True if logarithmic binning should be used (default : None).
Theo Steininger's avatar
Theo Steininger committed
44
45
46
47
48
    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
49
50
51
        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
52
        (default : None)
Martin Reinecke's avatar
Martin Reinecke committed
53
        if binbounds == None:
Theo Steininger's avatar
Theo Steininger committed
54
55
            Calculates the bounds from the kindex while applying the
            logarithmic and nbin keywords.
Martin Reinecke's avatar
Martin Reinecke committed
56
57
58
59
60
    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
61
62
63
64

    Attributes
    ----------
    pindex : distributed_data_object
65
66
        This holds the information which pixel of the harmonic partner gets
        mapped to which power bin
Theo Steininger's avatar
Theo Steininger committed
67
    kindex : numpy.ndarray
68
        Sorted array of all k-modes.
Theo Steininger's avatar
Theo Steininger committed
69
70
71
    rho : numpy.ndarray
        The amount of k-modes that get mapped to one power bin is given by
        rho.
72
73
74
75
76
77
78
79
    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
80
81
    binbounds :  tuple or None
        Boundaries between the power spectrum bins
Theo Steininger's avatar
Theo Steininger committed
82
83
84
85
86
87
88

    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.

    """
89

90
91
    _powerIndexCache = {}

92
93
    # ---Overwritten properties and methods---

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

99
100
101
102
103
104
105
        if distribution_strategy is None:
            distribution_strategy = gc['default_distribution_strategy']
        elif distribution_strategy not in DISTRIBUTION_STRATEGIES['global']:
            raise ValueError(
                    "distribution_strategy must be a global-type "
                    "strategy.")

Martin Reinecke's avatar
Martin Reinecke committed
106
107
108
        if not (isinstance(harmonic_partner, Space) and
                harmonic_partner.harmonic):
            raise ValueError("harmonic_partner must be a harmonic space.")
109
        self._harmonic_partner = harmonic_partner
110

Martin Reinecke's avatar
Martin Reinecke committed
111
112
113
114
        # 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")
115

Martin Reinecke's avatar
Martin Reinecke committed
116
117
        if binbounds is not None:
            binbounds = tuple(binbounds)
118

Martin Reinecke's avatar
Martin Reinecke committed
119
120
        key = (harmonic_partner, distribution_strategy, logarithmic, nbin,
               binbounds)
121
122
123
124
125
126
127
128
129
130
        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(
131
                                harmonic_partner=self.harmonic_partner,
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
                                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]

147
148
    @staticmethod
    def _compute_binbounds(harmonic_partner, distribution_strategy,
149
                           logarithmic, nbin, binbounds):
150

Martin Reinecke's avatar
Martin Reinecke committed
151
        if logarithmic is None and nbin is None and binbounds is None:
152
            result = None
Martin Reinecke's avatar
Martin Reinecke committed
153
154
155
156
157
158
159
160
161
162
163
        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
164
                kindex = harmonic_partner.get_unique_distances()
Martin Reinecke's avatar
Martin Reinecke committed
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
                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)
181
182
            result = tuple(bb)
        return result
183

184
185
    @staticmethod
    def _compute_pindex(harmonic_partner, distance_array, binbounds,
186
                        distribution_strategy):
187

Martin Reinecke's avatar
Martin Reinecke committed
188
        # Compute pindex, kindex and rho according to bb
189
190
191
192
193
        pindex = distributed_data_object(
                                global_shape=distance_array.shape,
                                dtype=np.int,
                                distribution_strategy=distribution_strategy)
        if binbounds is None:
194
            binbounds = harmonic_partner.get_natural_binbounds()
195
196
197
        pindex.set_local_data(
                np.searchsorted(binbounds, distance_array.get_local_data()))
        return pindex
198

Theo Steininger's avatar
Theo Steininger committed
199
    def pre_cast(self, x, axes):
Theo Steininger's avatar
Theo Steininger committed
200
201
202
203
204
205
        """ 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.

206
207
208
        Parameters
        ----------
        x : {array-like, function array-like -> array-like}
Theo Steininger's avatar
Theo Steininger committed
209
210
211
212
213
214
            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.

215
216
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
217
218
219
        array-like
            discretized power spectrum

220
        """
Theo Steininger's avatar
Theo Steininger committed
221

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

224
225
    # ---Mandatory properties and methods---

226
227
    def __repr__(self):
        return ("PowerSpace(harmonic_partner=%r, distribution_strategy=%r, "
Martin Reinecke's avatar
Martin Reinecke committed
228
                "binbounds=%r)"
229
                % (self.harmonic_partner, self.pindex.distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
230
                   self._binbounds))
231

232
233
234
    @property
    def harmonic(self):
        return True
235

236
237
    @property
    def shape(self):
238
        return self.kindex.shape
239

240
241
242
243
244
245
246
    @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
247
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
248
249

    def copy(self):
250
        distribution_strategy = self.pindex.distribution_strategy
251
        return self.__class__(harmonic_partner=self.harmonic_partner,
252
                              distribution_strategy=distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
253
                              binbounds=self._binbounds)
254

Martin Reinecke's avatar
Martin Reinecke committed
255
    def weight(self, x, power, axes, inplace=False):
Jait Dixit's avatar
Jait Dixit committed
256
257
        reshaper = [1, ] * len(x.shape)
        # we know len(axes) is always 1
258
259
        reshaper[axes[0]] = self.shape[0]

260
        weight = self.rho.reshape(reshaper)
261
        if power != 1:
262
            weight = weight ** np.float(power)
263
264
265
266
267
268

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
269
270
271

        return result_x

272
    def get_distance_array(self, distribution_strategy):
Martin Reinecke's avatar
Martin Reinecke committed
273
        return distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
274
                                self.kindex, dtype=np.float64,
275
                                distribution_strategy=distribution_strategy)
theos's avatar
theos committed
276

277
    def get_fft_smoothing_kernel_function(self, sigma):
278
        raise NotImplementedError(
279
            "There is no fft smoothing function for PowerSpace.")
theos's avatar
theos committed
280

281
282
283
    # ---Added properties and methods---

    @property
284
    def harmonic_partner(self):
Theo Steininger's avatar
Theo Steininger committed
285
        """ Returns the Space of which this is the power space.
286
287
        """
        return self._harmonic_partner
288
289

    @property
Martin Reinecke's avatar
Martin Reinecke committed
290
291
    def binbounds(self):
        return self._binbounds
292
293
294

    @property
    def pindex(self):
295
        """ A distributed_data_object having the shape of the harmonic partner
Theo Steininger's avatar
Theo Steininger committed
296
297
        space containing the indices of the power bin a pixel belongs to.
        """
298
299
300
301
        return self._pindex

    @property
    def kindex(self):
Theo Steininger's avatar
Theo Steininger committed
302
303
        """ Sorted array of all k-modes.
        """
304
305
306
307
        return self._kindex

    @property
    def rho(self):
Theo Steininger's avatar
Theo Steininger committed
308
309
        """Degeneracy factor of the individual k-vectors.
        """
310
        return self._rho
311

312
313
314
    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Martin Reinecke's avatar
Martin Reinecke committed
315
316
317
318
        hdf5_group.attrs['binbounds'] = str(self._binbounds)
        hdf5_group.attrs['distribution_strategy'] = \
            self._pindex.distribution_strategy

319
        return {
320
            'harmonic_partner': self.harmonic_partner,
321
322
323
        }

    @classmethod
Theo Steininger's avatar
Theo Steininger committed
324
    def _from_hdf5(cls, hdf5_group, repository):
Martin Reinecke's avatar
Martin Reinecke committed
325
        hp = repository.get('harmonic_partner', hdf5_group)
Martin Reinecke's avatar
Martin Reinecke committed
326
        bb = ast.literal_eval(hdf5_group.attrs['binbounds'])
Martin Reinecke's avatar
Martin Reinecke committed
327
328
        ds = hdf5_group.attrs['distribution_strategy']
        return PowerSpace(hp, ds, binbounds=bb)