power_indices.py 15.6 KB
Newer Older
1
2
3
4
5
6
# -*- coding: utf-8 -*-

import numpy as np
from d2o import distributed_data_object,\
                STRATEGIES as DISTRIBUTION_STRATEGIES

7
from nifty.config import dependency_injector as gdi
8

9
MPI = gdi['MPI']
10
11
12


class PowerIndices(object):
13
14
    def __init__(self, domain, distribution_strategy,
                 log=False, nbin=None, binbounds=None):
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
        """
            Returns an instance of the PowerIndices class. Given the shape and
            the density of a underlying rectangular grid it provides the user
            with the pindex, kindex, rho and pundex. The indices are bined
            according to the supplied parameter scheme. If wanted, computed
            results are stored for future reuse.

            Parameters
            ----------
            shape : tuple, list, ndarray
                Array-like object which specifies the shape of the underlying
                rectangular grid
            dgrid : tuple, list, ndarray
                Array-like object which specifies the step-width of the
                underlying grid
            zerocenter : boolean, tuple/list/ndarray of boolean *optional*
                Specifies which dimensions are zerocentered. (default:False)
            log : bool *optional*
                Flag specifying if the binning of the default indices is
                performed on logarithmic scale.
            nbin : integer *optional*
                Number of used bins for the binning of the default indices.
            binbounds : {list, array}
                Array-like inner boundaries of the used bins of the default
                indices.
        """
41
        self.domain = domain
42
43
        self.distribution_strategy = distribution_strategy

theos's avatar
theos committed
44
        # Compute the global k_array
45
        self.k_array = self.domain.get_distance_array(distribution_strategy)
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
        # Initialize the dictonary which stores all individual index-dicts
        self.global_dict = {}
        # Set self.default_parameters
        self.set_default(config_dict={'log': log,
                                      'nbin': nbin,
                                      'binbounds': binbounds})

    # Redirect the direct calls approaching a power_index instance to the
    # default_indices dict
    @property
    def default_indices(self):
        return self.get_index_dict(**self.default_parameters)

    def __getitem__(self, x):
        return self.default_indices.get(x)

    def __contains__(self, x):
        return self.default_indices.__contains__(x)

    def __iter__(self):
        return self.default_indices.__iter__()

    def __getattr__(self, x):
        return self.default_indices.__getattribute__(x)

    def set_default(self, **kwargs):
        """
            Sets the index-set which is specified by the parameters as the
            default for the power_index instance.

            Parameters
            ----------
            log : bool
                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
            -------
            None
        """
        parsed_kwargs = self._cast_config(**kwargs)
        self.default_parameters = parsed_kwargs

    def _cast_config(self, **kwargs):
        """
            internal helper function which casts the various combinations of
            possible parameters into a properly defaulted dictionary
        """
        temp_config_dict = kwargs.get('config_dict', None)
        if temp_config_dict is not None:
            return self._cast_config_helper(**temp_config_dict)
        else:
            defaults = self.default_parameters
            temp_log = kwargs.get("log", defaults['log'])
            temp_nbin = kwargs.get("nbin", defaults['nbin'])
            temp_binbounds = kwargs.get("binbounds", defaults['binbounds'])

            return self._cast_config_helper(log=temp_log,
                                            nbin=temp_nbin,
                                            binbounds=temp_binbounds)

    def _cast_config_helper(self, log, nbin, binbounds):
        """
            internal helper function which sets the defaults for the
            _cast_config function
        """

        try:
            temp_log = bool(log)
        except(TypeError):
            temp_log = False

        try:
            temp_nbin = int(nbin)
        except(TypeError):
            temp_nbin = None

        try:
            temp_binbounds = tuple(np.array(binbounds))
        except(TypeError):
            temp_binbounds = None

        temp_dict = {"log": temp_log,
                     "nbin": temp_nbin,
                     "binbounds": temp_binbounds}
        return temp_dict

    def get_index_dict(self, **kwargs):
        """
            Returns a dictionary containing the pindex, kindex, rho and pundex
            binned according to the supplied parameter scheme and a
            configuration dict containing this scheme.

            Parameters
            ----------
            store : bool
                Flag specifying if  the calculated index dictionary should be
                stored in the global_dict for future use.
            log : bool
                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
            -------
            index_dict : dict
                Contains the keys: 'config', 'pindex', 'kindex', 'rho' and
                'pundex'
        """
        # Cast the input arguments
        temp_config_dict = self._cast_config(**kwargs)
        # Compute a hashable identifier from the config which will be used
        # as dict key
        temp_key = self._freeze_config(temp_config_dict)
        # Check if the result should be stored for future use.
        storeQ = kwargs.get("store", True)
        # Try to find the requested index dict in the global_dict
        try:
            return self.global_dict[temp_key]
        except(KeyError):
            # If it is not found, calculate it.
            temp_index_dict = self._compute_index_dict(temp_config_dict)
            # Store it, if required
            if storeQ:
                self.global_dict[temp_key] = temp_index_dict
                # Important: If the result is stored, return a reference to
                # the dictionary entry, not anly a plain copy. Otherwise,
                # set_default breaks!
                return self.global_dict[temp_key]
            else:
                # Return the plain result.
                return temp_index_dict

    def _freeze_config(self, config_dict):
        """
            a helper function which forms a hashable identifying object from
            a config dictionary which can be used as key of a dict
        """
        return frozenset(config_dict.items())

    def _compute_index_dict(self, config_dict):
        """
            Internal helper function which takes a config_dict, asks for the
            pindex/kindex/rho/pundex set, and bins them according to the config
        """
        # if no binning is requested, compute the indices, build the dict,
        # and return it straight.
        if not config_dict["log"] and config_dict["nbin"] is None and \
                config_dict["binbounds"] is None:
            (temp_pindex, temp_kindex, temp_rho, temp_pundex) =\
theos's avatar
theos committed
203
204
                self._compute_indices(self.k_array)
            temp_k_array = self.k_array
205
206
207
208
209
210
211
212
213
214
215
216
217

        # if binning is required, make a recursive call to get the unbinned
        # indices, bin them, compute the pundex and then return everything.
        else:
            # Get the unbinned indices
            temp_unbinned_indices = self.get_index_dict(nbin=None,
                                                        binbounds=None,
                                                        log=False,
                                                        store=False)
            # Bin them
            (temp_pindex, temp_kindex, temp_rho, temp_pundex) = \
                self._bin_power_indices(
                    temp_unbinned_indices, **config_dict)
theos's avatar
theos committed
218
219
220
            # Make a binned version of k_array
            temp_k_array = self._compute_k_array_from_pindex_kindex(
                               temp_pindex, temp_kindex)
221
222
223
224
225
226

        temp_index_dict = {"config": config_dict,
                           "pindex": temp_pindex,
                           "kindex": temp_kindex,
                           "rho": temp_rho,
                           "pundex": temp_pundex,
theos's avatar
theos committed
227
                           "k_array": temp_k_array}
228
229
        return temp_index_dict

theos's avatar
theos committed
230
    def _compute_k_array_from_pindex_kindex(self, pindex, kindex):
231
232
233
        tempindex = pindex.copy(dtype=kindex.dtype)
        result = tempindex.apply_scalar_function(
                            lambda x: kindex[x.astype(np.dtype('int'))])
234
235
        return result

theos's avatar
theos committed
236
    def _compute_indices(self, k_array):
237
238
        """
        Internal helper function which computes pindex, kindex, rho and pundex
theos's avatar
theos committed
239
        from a given k_array
240
        """
241
242
243
        ##########
        # kindex #
        ##########
theos's avatar
theos committed
244
        global_kindex = k_array.unique()
245
246
247
248

        ##########
        # pindex #
        ##########
theos's avatar
theos committed
249
250
        # compute the local pindex slice on basis of the local k_array data
        local_pindex = np.searchsorted(global_kindex, k_array.get_local_data())
251
252
        # prepare the distributed_data_object
        global_pindex = distributed_data_object(
theos's avatar
theos committed
253
                            global_shape=k_array.shape,
254
255
256
257
258
259
260
261
262
                            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()
263

264
265
266
267
268
269
270
271
272
        ##########
        # pundex #
        ##########
        global_pundex = self._compute_pundex(global_pindex,
                                             global_kindex)

        return global_pindex, global_kindex, global_rho, global_pundex

    def _compute_pundex(self, global_pindex, global_kindex):
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
        """
        Internal helper function which computes the pundex array from a
        pindex and a kindex array. This function is separated from the
        _compute_indices function as it is needed in _bin_power_indices,
        too.
        """
        if self.distribution_strategy in DISTRIBUTION_STRATEGIES['slicing']:
            ##########
            # pundex #
            ##########
            # Prepare the local data
            local_pindex = global_pindex.get_local_data()
            # Compute the local pundices for the local pindices
            (temp_uniqued_pindex, local_temp_pundex) = np.unique(
                                                            local_pindex,
                                                            return_index=True)
            # Shift the local pundices by the nodes' local_dim_offset
            local_temp_pundex += global_pindex.distributor.local_dim_offset

            # Prepare the pundex arrays used for the Allreduce operation
            # pundex has the same length as the kindex array
            local_pundex = np.zeros(shape=global_kindex.shape, dtype=np.int)
            # Set the default value higher than the maximal possible pundex
            # value so that MPI.MIN can sort out the default
            local_pundex += np.prod(global_pindex.shape) + 1
            # Set the default value higher than the length
            global_pundex = np.empty_like(local_pundex)
            # Store the individual pundices in the local_pundex array
            local_pundex[temp_uniqued_pindex] = local_temp_pundex
            # Use Allreduce to find the first occurences/smallest pundices
303
304
305
            global_pindex.comm.Allreduce(local_pundex,
                                         global_pundex,
                                         op=MPI.MIN)
306
307
308
309
310
311
312
313
314
315
            return global_pundex

        elif self.distribution_strategy in DISTRIBUTION_STRATEGIES['not']:
            ##########
            # pundex #
            ##########
            pundex = np.unique(global_pindex.get_local_data(),
                               return_index=True)[1]
            return pundex
        else:
316
317
318
            raise NotImplementedError(
                "_compute_pundex_d2o not implemented for given "
                "distribution_strategy")
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395

    def _bin_power_indices(self, index_dict, **kwargs):
        """
            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.
            log : bool
                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
            kindex, rho, pundex : ndarrays
                The (re)binned power indices.

        """
        # Cast the given config
        temp_config_dict = self._cast_config(**kwargs)
        log = temp_config_dict['log']
        nbin = temp_config_dict['nbin']
        binbounds = temp_config_dict['binbounds']

        # Extract the necessary indices from the supplied index dict
        pindex = index_dict["pindex"]
        kindex = index_dict["kindex"]
        rho = index_dict["rho"]

        # boundaries
        if(binbounds is not None):
            binbounds = np.sort(binbounds)
        # equal binning
        else:
            if(log is None):
                log = False
            if(log):
                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)]
            if(log):
                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]

396
397
398
        pindex_ = pindex.copy_empty()
        pindex_.set_local_data(reorder[pindex.get_local_data()])
        pundex_ = self._compute_pundex(pindex_, kindex_)
399
400

        return pindex_, kindex_, rho_, pundex_