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

19
import ast
20 21
import numpy as np

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

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


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

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

    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.

    """
85

86 87
    _powerIndexCache = {}

88 89
    # ---Overwritten properties and methods---

90
    def __init__(self, harmonic_partner, 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
        key = (harmonic_partner, distribution_strategy, logarithmic, nbin,
               binbounds)
110 111 112 113 114 115 116 117 118 119
        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(
120
                                harmonic_partner=self.harmonic_partner,
121 122 123 124 125 126 127 128 129 130 131 132 133 134 135
                                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]

136 137
    @staticmethod
    def _compute_binbounds(harmonic_partner, distribution_strategy,
138
                           logarithmic, nbin, binbounds):
139

Martin Reinecke's avatar
Martin Reinecke committed
140
        if logarithmic is None and nbin is None and binbounds is None:
141
            result = None
Martin Reinecke's avatar
Martin Reinecke committed
142 143 144 145 146 147 148 149 150 151 152
        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
153
                kindex = harmonic_partner.get_unique_distances()
Martin Reinecke's avatar
Martin Reinecke committed
154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169
                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)
170 171 172
            result = tuple(bb)
        return result

173 174
    @staticmethod
    def _compute_pindex(harmonic_partner, distance_array, binbounds,
175
                        distribution_strategy):
Martin Reinecke's avatar
Martin Reinecke committed
176 177

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

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

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

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

209
        """
Theo Steininger's avatar
Theo Steininger committed
210

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

213 214
    # ---Mandatory properties and methods---

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

221 222 223
    @property
    def harmonic(self):
        return True
224

225 226
    @property
    def shape(self):
227
        return self.kindex.shape
228

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

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

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

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

        if inplace:
            x *= weight
            result_x = x
        else:
            result_x = x*weight
258 259 260

        return result_x

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

266
    def get_fft_smoothing_kernel_function(self, sigma):
267
        raise NotImplementedError(
268
            "There is no fft smoothing function for PowerSpace.")
269

270 271 272
    # ---Added properties and methods---

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

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

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

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

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

301 302 303
    # ---Serialization---

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

308
        return {
309
            'harmonic_partner': self.harmonic_partner,
310 311 312
        }

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