power_space.py 10.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.
Theo Steininger's avatar
Theo Steininger committed
18

19
20
import numpy as np

Martin Reinecke's avatar
Martin Reinecke committed
21
from d2o import distributed_data_object
22

23
from nifty.spaces.space import Space
Theo Steininger's avatar
Theo Steininger committed
24

Martin Reinecke's avatar
Martin Reinecke committed
25
26
_PSCache = {}

Theo Steininger's avatar
Theo Steininger committed
27

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

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

    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.

    """
86

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

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

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

Martin Reinecke's avatar
Martin Reinecke committed
100
101
102
103
        # 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")
104

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

Martin Reinecke's avatar
Martin Reinecke committed
108
109
110
111
112
113
        key = (harmonic_partner, distribution_strategy, logarithmic, nbin,
               binbounds)
        if _PSCache.get(key) is not None:
            (self._binbounds, self._pindex, self._kindex, self._rho) \
                = _PSCache[key]
            return
114

Martin Reinecke's avatar
Martin Reinecke committed
115
116
117
118
119
120
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
151
152
153
154
155
156
157
158
159
160
161
162
163
        self._binbounds = None
        if logarithmic is None and nbin is None and binbounds is None:
            bb = self._harmonic_partner.get_natural_binbounds()
        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
                kindex = self._harmonic_partner.get_unique_distances()
                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)
            self._binbounds = tuple(bb)

        dists = self._harmonic_partner.get_distance_array(
                distribution_strategy)
        # Compute pindex, kindex and rho according to bb
        self._pindex = distributed_data_object(
            global_shape=dists.shape,
            dtype=np.int,
            distribution_strategy=distribution_strategy)

        self._pindex.set_local_data(np.searchsorted(
            bb, dists.get_local_data()))  # also expensive!
        self._rho = self._pindex.bincount().get_full_data()
        self._kindex = self._pindex.bincount(
            weights=dists).get_full_data()/self._rho

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

Theo Steininger's avatar
Theo Steininger committed
165
    def pre_cast(self, x, axes):
Theo Steininger's avatar
Theo Steininger committed
166
167
168
169
170
171
        """ 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.

172
173
174
        Parameters
        ----------
        x : {array-like, function array-like -> array-like}
Theo Steininger's avatar
Theo Steininger committed
175
176
177
178
179
180
            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.

181
182
        Returns
        -------
Theo Steininger's avatar
Theo Steininger committed
183
184
185
        array-like
            discretized power spectrum

186
        """
Theo Steininger's avatar
Theo Steininger committed
187

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

190
191
    # ---Mandatory properties and methods---

192
193
    def __repr__(self):
        return ("PowerSpace(harmonic_partner=%r, distribution_strategy=%r, "
Martin Reinecke's avatar
Martin Reinecke committed
194
                "binbounds=%r)"
195
                % (self.harmonic_partner, self.pindex.distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
196
                   self._binbounds))
197

198
199
200
    @property
    def harmonic(self):
        return True
201

202
203
    @property
    def shape(self):
204
        return self.kindex.shape
205

206
207
208
209
210
211
212
    @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
213
        return float(reduce(lambda x, y: x*y, self.pindex.shape))
214
215

    def copy(self):
216
        distribution_strategy = self.pindex.distribution_strategy
217
        return self.__class__(harmonic_partner=self.harmonic_partner,
218
                              distribution_strategy=distribution_strategy,
Martin Reinecke's avatar
Martin Reinecke committed
219
                              binbounds=self._binbounds)
220

Martin Reinecke's avatar
Martin Reinecke committed
221
    def weight(self, x, power, axes, inplace=False):
Jait Dixit's avatar
Jait Dixit committed
222
223
        reshaper = [1, ] * len(x.shape)
        # we know len(axes) is always 1
224
225
        reshaper[axes[0]] = self.shape[0]

226
        weight = self.rho.reshape(reshaper)
227
        if power != 1:
228
            weight = weight ** np.float(power)
229
230
231
232
233
234

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
235
236
237

        return result_x

238
    def get_distance_array(self, distribution_strategy):
Martin Reinecke's avatar
Martin Reinecke committed
239
        return distributed_data_object(
Martin Reinecke's avatar
Martin Reinecke committed
240
                                self.kindex, dtype=np.float64,
241
                                distribution_strategy=distribution_strategy)
242

243
    def get_fft_smoothing_kernel_function(self, sigma):
244
        raise NotImplementedError(
245
            "There is no fft smoothing function for PowerSpace.")
246

247
248
249
    # ---Added properties and methods---

    @property
250
    def harmonic_partner(self):
Theo Steininger's avatar
Theo Steininger committed
251
        """ Returns the Space of which this is the power space.
252
253
        """
        return self._harmonic_partner
254
255

    @property
Martin Reinecke's avatar
Martin Reinecke committed
256
257
    def binbounds(self):
        return self._binbounds
258
259
260

    @property
    def pindex(self):
261
        """ A distributed_data_object having the shape of the harmonic partner
Theo Steininger's avatar
Theo Steininger committed
262
263
        space containing the indices of the power bin a pixel belongs to.
        """
264
265
266
267
        return self._pindex

    @property
    def kindex(self):
Theo Steininger's avatar
Theo Steininger committed
268
269
        """ Sorted array of all k-modes.
        """
270
271
272
273
        return self._kindex

    @property
    def rho(self):
Theo Steininger's avatar
Theo Steininger committed
274
275
        """Degeneracy factor of the individual k-vectors.
        """
276
        return self._rho
277

278
279
280
    # ---Serialization---

    def _to_hdf5(self, hdf5_group):
Martin Reinecke's avatar
Martin Reinecke committed
281
282
283
284
        hdf5_group.attrs['binbounds'] = str(self._binbounds)
        hdf5_group.attrs['distribution_strategy'] = \
            self._pindex.distribution_strategy

285
        return {
286
            'harmonic_partner': self.harmonic_partner,
287
288
289
        }

    @classmethod
Theo Steininger's avatar
Theo Steininger committed
290
    def _from_hdf5(cls, hdf5_group, repository):
Martin Reinecke's avatar
Martin Reinecke committed
291
292
293
294
        hp = repository.get('harmonic_partner', hdf5_group)
        exec("bb = " + hdf5_group.attrs['binbounds'])
        ds = hdf5_group.attrs['distribution_strategy']
        return PowerSpace(hp, ds, binbounds=bb)