power_indices.py 7.19 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.
18
19

import numpy as np
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
20
from d2o import distributed_data_object
21
22

class PowerIndices(object):
23
    """Computes helpful quantities to deal with power spectra.
24

Martin Reinecke's avatar
step 4    
Martin Reinecke committed
25
26
    Given the shape and the density of a underlying grid this class provides
    the user with the pindex, kindex and rho. The indices are binned
Martin Reinecke's avatar
step 3    
Martin Reinecke committed
27
    according to the supplied parameter scheme.
28

29
    Parameters
30
    ----------
31
32
33
34
35
    domain : NIFTy harmonic space
        The space for which the power indices get computed
    distribution_strategy : str
        The distribution_strategy that will be used for the k_array and pindex
        distributed_data_object.
36
    """
Martin Reinecke's avatar
step 2    
Martin Reinecke committed
37
    def __init__(self, domain, distribution_strategy):
38
        self.domain = domain
39
40
        self.distribution_strategy = distribution_strategy

theos's avatar
theos committed
41
        # Compute the global k_array
42
        self.k_array = self.domain.get_distance_array(distribution_strategy)
43
44


Martin Reinecke's avatar
step 3    
Martin Reinecke committed
45
    def get_index_dict(self, logarithmic, nbin, binbounds):
46
        """
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
47
            Returns a dictionary containing the pindex, kindex and rho
48
49
50
51
52
            binned according to the supplied parameter scheme and a
            configuration dict containing this scheme.

            Parameters
            ----------
53
            logarithmic : bool
54
55
56
57
58
59
60
61
62
63
64
                Flag specifying if the binning is performed on logarithmic
                scale.
            nbin : integer
                Number of used bins.
            binbounds : {list, array}
                Array-like inner boundaries of the used bins.

        """

        # if no binning is requested, compute the indices, build the dict,
        # and return it straight.
Martin Reinecke's avatar
step 3    
Martin Reinecke committed
65
        if not logarithmic and nbin is None and binbounds is None:
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
66
            (temp_pindex, temp_kindex, temp_rho) =\
theos's avatar
theos committed
67
68
                self._compute_indices(self.k_array)
            temp_k_array = self.k_array
69
70

        # if binning is required, make a recursive call to get the unbinned
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
71
        # indices, bin them, and then return everything.
72
73
        else:
            # Get the unbinned indices
Martin Reinecke's avatar
step 3    
Martin Reinecke committed
74
            pindex, kindex, rho, dummy = self.get_index_dict(nbin=None,
75
                                                        binbounds=None,
Martin Reinecke's avatar
step 3    
Martin Reinecke committed
76
                                                        logarithmic=False)
77
            # Bin them
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
78
            (temp_pindex, temp_kindex, temp_rho) = \
79
                self._bin_power_indices(
Martin Reinecke's avatar
step 3    
Martin Reinecke committed
80
                    pindex, kindex, rho, logarithmic, nbin, binbounds)
theos's avatar
theos committed
81
            # Make a binned version of k_array
Martin Reinecke's avatar
step 4    
Martin Reinecke committed
82
83
84
            tempindex = temp_pindex.copy(dtype=temp_kindex.dtype)
            temp_k_array = tempindex.apply_scalar_function(
                            lambda x: temp_kindex[x.astype(np.dtype('int'))])
85

Martin Reinecke's avatar
step 3    
Martin Reinecke committed
86
        return temp_pindex, temp_kindex, temp_rho, temp_k_array
87

theos's avatar
theos committed
88
    def _compute_indices(self, k_array):
89
        """
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
90
        Internal helper function which computes pindex, kindex and rho
theos's avatar
theos committed
91
        from a given k_array
92
        """
93
94
95
        ##########
        # kindex #
        ##########
theos's avatar
theos committed
96
        global_kindex = k_array.unique()
97
98
99
100

        ##########
        # pindex #
        ##########
theos's avatar
theos committed
101
102
        # compute the local pindex slice on basis of the local k_array data
        local_pindex = np.searchsorted(global_kindex, k_array.get_local_data())
103
104
        # prepare the distributed_data_object
        global_pindex = distributed_data_object(
theos's avatar
theos committed
105
                            global_shape=k_array.shape,
106
107
108
109
110
111
112
113
114
                            dtype=local_pindex.dtype,
                            distribution_strategy=self.distribution_strategy)
        # store the local pindex data in the global_pindex d2o
        global_pindex.set_local_data(local_pindex)

        #######
        # rho #
        #######
        global_rho = global_pindex.bincount().get_full_data()
115

Martin Reinecke's avatar
step 1    
Martin Reinecke committed
116
        return global_pindex, global_kindex, global_rho
117

Martin Reinecke's avatar
step 3    
Martin Reinecke committed
118
    def _bin_power_indices(self, pindex, kindex, rho, logarithmic, nbin, binbounds):
119
120
121
122
123
124
125
126
127
128
129
        """
            Returns the binned power indices associated with the Fourier grid.

            Parameters
            ----------
            pindex : distributed_data_object
                Index of the Fourier grid points in a distributed_data_object.
            kindex : ndarray
                Array of all k-vector lengths.
            rho : ndarray
                Degeneracy factor of the individual k-vectors.
130
            logarithmic : bool
131
132
133
134
135
136
137
138
139
140
                Flag specifying if the binning is performed on logarithmic
                scale.
            nbin : integer
                Number of used bins.
            binbounds : {list, array}
                Array-like inner boundaries of the used bins.

            Returns
            -------
            pindex : distributed_data_object
Martin Reinecke's avatar
step 1    
Martin Reinecke committed
141
            kindex, rho : ndarrays
142
143
144
145
146
147
148
149
150
                The (re)binned power indices.

        """

        # boundaries
        if(binbounds is not None):
            binbounds = np.sort(binbounds)
        # equal binning
        else:
151
152
153
            if(logarithmic is None):
                logarithmic = False
            if(logarithmic):
154
155
156
157
158
159
160
161
162
163
164
165
166
                k = np.r_[0, np.log(kindex[1:])]
            else:
                k = kindex
            dk = np.max(k[2:] - k[1:-1])  # minimal dk
            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)
            binbounds = np.r_[0.5 * (3 * k[1] - k[2]),
                              0.5 * (k[1] + k[2]) + dk * np.arange(nbin - 2)]
167
            if(logarithmic):
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
                binbounds = np.exp(binbounds)
        # reordering
        reorder = np.searchsorted(binbounds, kindex)
        rho_ = np.zeros(len(binbounds) + 1, dtype=rho.dtype)
        kindex_ = np.empty(len(binbounds) + 1, dtype=kindex.dtype)
        for ii in range(len(reorder)):
            if(rho_[reorder[ii]] == 0):
                kindex_[reorder[ii]] = kindex[ii]
                rho_[reorder[ii]] += rho[ii]
            else:
                kindex_[reorder[ii]] = ((kindex_[reorder[ii]] *
                                         rho_[reorder[ii]] +
                                         kindex[ii] * rho[ii]) /
                                        (rho_[reorder[ii]] + rho[ii]))
                rho_[reorder[ii]] += rho[ii]

184
185
        pindex_ = pindex.copy_empty()
        pindex_.set_local_data(reorder[pindex.get_local_data()])
186

Martin Reinecke's avatar
step 1    
Martin Reinecke committed
187
        return pindex_, kindex_, rho_