power_space.py 11.1 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
    """ 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')
Martin Reinecke's avatar
Martin Reinecke committed
42
43
44
45
46
47
48
49
50
51
52
53
    binbounds: None, or tuple/array/list of float
        if None:
            There will be as many bins as there are distinct k-vector lengths
            in the harmonic partner space.
            The "binbounds" property of the PowerSpace will also be None.

        else:
            the bin bounds requested for this PowerSpace. The array
            must be sorted and strictly ascending. The first entry is the right
            boundary of the first bin, and the last entry is the left boundary
            of the last bin, i.e. thee will be len(binbounds)+1 bins in total,
            with the first and last bins reaching to -+infinity, respectively.
Theo Steininger's avatar
Theo Steininger committed
54
55
56
57
58
        (default : None)

    Attributes
    ----------
    pindex : distributed_data_object
59
60
        This holds the information which pixel of the harmonic partner gets
        mapped to which power bin
Theo Steininger's avatar
Theo Steininger committed
61
    kindex : numpy.ndarray
62
        Sorted array of all k-modes.
Theo Steininger's avatar
Theo Steininger committed
63
64
65
    rho : numpy.ndarray
        The amount of k-modes that get mapped to one power bin is given by
        rho.
66
67
68
    dim : np.int
        Total number of dimensionality, i.e. the number of pixels.
    harmonic : bool
Martin Reinecke's avatar
Martin Reinecke committed
69
        Always True for this space.
70
71
72
73
    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
74
75
76
    binbounds : tuple or None
        Boundaries between the power spectrum bins; None is used to indicate
        natural binning
Theo Steininger's avatar
Theo Steininger committed
77
78
79
80
81
82
83

    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.

    """
84

85
86
    _powerIndexCache = {}

87
88
    # ---Overwritten properties and methods---

Martin Reinecke's avatar
Martin Reinecke committed
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
    @staticmethod
    def linear_binbounds (nbin, first_bound, last_bound):
        """
        nbin: integer
            the number of bins
        first_bound, last_bound: float
            the k values for the right boundary of the first bin and the left
            boundary of the last bin, respectively. They are given in length
            units of the harmonic partner space.
        This will produce a binbounds array with nbin-1 entries with
        binbounds[0]=first_bound and binbounds[-1]=last_bound and the remaining
        values equidistantly spaced (in linear scale) between these two.
        """
        nbin = int(nbin)
        assert nbin>=3, "nbin must be at least 3"
Martin Reinecke's avatar
Martin Reinecke committed
104
        return np.linspace(float(first_bound),float(last_bound),nbin-1)
Martin Reinecke's avatar
Martin Reinecke committed
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

    @staticmethod
    def logarithmic_binbounds (nbin, first_bound, last_bound):
        """
        nbin: integer
            the number of bins
        first_bound, last_bound: float
            the k values for the right boundary of the first bin and the left
            boundary of the last bin, respectively. They are given in length
            units of the harmonic partner space.
        This will produce a binbounds array with nbin-1 entries with
        binbounds[0]=first_bound and binbounds[-1]=last_bound and the remaining
        values equidistantly spaced (in natural logarithmic scale)
        between these two.
        """
Martin Reinecke's avatar
Martin Reinecke committed
120
121
122
123
124
        nbin = int(nbin)
        assert nbin>=3, "nbin must be at least 3"
        return np.logspace(np.log(float(first_bound)),
                           np.log(float(last_bound)),
                           nbin-1, base=np.e)
Martin Reinecke's avatar
Martin Reinecke committed
125

126
    def __init__(self, harmonic_partner, distribution_strategy=None,
Martin Reinecke's avatar
Martin Reinecke committed
127
                 binbounds=None):
Martin Reinecke's avatar
Martin Reinecke committed
128
        super(PowerSpace, self).__init__()
Martin Reinecke's avatar
Martin Reinecke committed
129
        self._ignore_for_hash += ['_pindex', '_kindex', '_rho']
130

131
132
133
134
135
136
137
        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
138
139
140
        if not (isinstance(harmonic_partner, Space) and
                harmonic_partner.harmonic):
            raise ValueError("harmonic_partner must be a harmonic space.")
141
        self._harmonic_partner = harmonic_partner
142

Martin Reinecke's avatar
Martin Reinecke committed
143
144
        if binbounds is not None:
            binbounds = tuple(binbounds)
145

Martin Reinecke's avatar
Martin Reinecke committed
146
        key = (harmonic_partner, distribution_strategy, binbounds)
147
148
149
150
        if self._powerIndexCache.get(key) is None:
            distance_array = \
                self.harmonic_partner.get_distance_array(distribution_strategy)
            temp_pindex = self._compute_pindex(
151
                                harmonic_partner=self.harmonic_partner,
152
                                distance_array=distance_array,
Martin Reinecke's avatar
Martin Reinecke committed
153
                                binbounds=binbounds,
154
155
                                distribution_strategy=distribution_strategy)
            temp_rho = temp_pindex.bincount().get_full_data()
Martin Reinecke's avatar
Martin Reinecke committed
156
            assert not np.any(temp_rho==0), "empty bins detected"
157
158
159
            temp_kindex = \
                (temp_pindex.bincount(weights=distance_array).get_full_data() /
                 temp_rho)
Martin Reinecke's avatar
Martin Reinecke committed
160
            self._powerIndexCache[key] = (binbounds,
161
162
163
164
165
166
167
                                          temp_pindex,
                                          temp_kindex,
                                          temp_rho)

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

168
169
    @staticmethod
    def _compute_pindex(harmonic_partner, distance_array, binbounds,
170
                        distribution_strategy):
171

Martin Reinecke's avatar
Martin Reinecke committed
172
        # Compute pindex, kindex and rho according to bb
173
174
175
176
177
        pindex = distributed_data_object(
                                global_shape=distance_array.shape,
                                dtype=np.int,
                                distribution_strategy=distribution_strategy)
        if binbounds is None:
178
            binbounds = harmonic_partner.get_natural_binbounds()
179
180
181
        pindex.set_local_data(
                np.searchsorted(binbounds, distance_array.get_local_data()))
        return pindex
182

Theo Steininger's avatar
Theo Steininger committed
183
    def pre_cast(self, x, axes):
Theo Steininger's avatar
Theo Steininger committed
184
185
186
187
188
189
        """ 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.

190
191
192
        Parameters
        ----------
        x : {array-like, function array-like -> array-like}
Theo Steininger's avatar
Theo Steininger committed
193
194
195
196
197
198
            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.

199
200
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
201
202
203
        array-like
            discretized power spectrum

204
        """
Theo Steininger's avatar
Theo Steininger committed
205

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

208
209
    # ---Mandatory properties and methods---

210
211
    def __repr__(self):
        return ("PowerSpace(harmonic_partner=%r, distribution_strategy=%r, "
Martin Reinecke's avatar
Martin Reinecke committed
212
                "binbounds=%r)"
213
                % (self.harmonic_partner, self.pindex.distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
214
                   self._binbounds))
215

216
217
218
    @property
    def harmonic(self):
        return True
219

220
221
    @property
    def shape(self):
222
        return self.kindex.shape
223

224
225
226
227
228
229
230
    @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
231
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
232
233

    def copy(self):
234
        distribution_strategy = self.pindex.distribution_strategy
235
        return self.__class__(harmonic_partner=self.harmonic_partner,
236
                              distribution_strategy=distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
237
                              binbounds=self._binbounds)
238

Martin Reinecke's avatar
Martin Reinecke committed
239
    def weight(self, x, power, axes, inplace=False):
Jait Dixit's avatar
Jait Dixit committed
240
241
        reshaper = [1, ] * len(x.shape)
        # we know len(axes) is always 1
242
243
        reshaper[axes[0]] = self.shape[0]

244
        weight = self.rho.reshape(reshaper)
245
        if power != 1:
246
            weight = weight ** np.float(power)
247
248
249
250
251
252

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
253
254
255

        return result_x

256
    def get_distance_array(self, distribution_strategy):
Martin Reinecke's avatar
Martin Reinecke committed
257
        return distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
258
                                self.kindex, dtype=np.float64,
259
                                distribution_strategy=distribution_strategy)
theos's avatar
theos committed
260

261
    def get_fft_smoothing_kernel_function(self, sigma):
262
        raise NotImplementedError(
263
            "There is no fft smoothing function for PowerSpace.")
theos's avatar
theos committed
264

265
266
267
    # ---Added properties and methods---

    @property
268
    def harmonic_partner(self):
Theo Steininger's avatar
Theo Steininger committed
269
        """ Returns the Space of which this is the power space.
270
271
        """
        return self._harmonic_partner
272
273

    @property
Martin Reinecke's avatar
Martin Reinecke committed
274
275
    def binbounds(self):
        return self._binbounds
276
277
278

    @property
    def pindex(self):
279
        """ A distributed_data_object having the shape of the harmonic partner
Theo Steininger's avatar
Theo Steininger committed
280
281
        space containing the indices of the power bin a pixel belongs to.
        """
282
283
284
285
        return self._pindex

    @property
    def kindex(self):
Theo Steininger's avatar
Theo Steininger committed
286
287
        """ Sorted array of all k-modes.
        """
288
289
290
291
        return self._kindex

    @property
    def rho(self):
Theo Steininger's avatar
Theo Steininger committed
292
293
        """Degeneracy factor of the individual k-vectors.
        """
294
        return self._rho
295

296
297
298
    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Martin Reinecke's avatar
Martin Reinecke committed
299
300
301
302
        hdf5_group.attrs['binbounds'] = str(self._binbounds)
        hdf5_group.attrs['distribution_strategy'] = \
            self._pindex.distribution_strategy

303
        return {
304
            'harmonic_partner': self.harmonic_partner,
305
306
307
        }

    @classmethod
Theo Steininger's avatar
Theo Steininger committed
308
    def _from_hdf5(cls, hdf5_group, repository):
Martin Reinecke's avatar
Martin Reinecke committed
309
        hp = repository.get('harmonic_partner', hdf5_group)
Martin Reinecke's avatar
Martin Reinecke committed
310
        bb = ast.literal_eval(hdf5_group.attrs['binbounds'])
Martin Reinecke's avatar
Martin Reinecke committed
311
312
        ds = hdf5_group.attrs['distribution_strategy']
        return PowerSpace(hp, ds, binbounds=bb)