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

import numpy as np
from nifty.keepers import about,\
                          global_configuration as gc,\
                          global_dependency_injector as gdi
theos's avatar
theos committed
7
8
from nifty.d2o import distributed_data_object,\
                      STRATEGIES as DISTRIBUTION_STRATEGIES
9
10

MPI = gdi[gc['mpi_module']]
Ultima's avatar
Ultima committed
11
hp = gdi.get('healpy')
12
13
14
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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55


class power_indices(object):
    def __init__(self, datamodel, allowed_distribution_strategies, comm,
                 log=False, nbin=None, binbounds=None):
        """
            Returns an instance of the power_indices 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
            zerocentered : 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.
        """
        if comm is None:
            self.comm = getattr(gdi[gc['mpi_module']], gc['default_comm'])
        else:
            self.comm = comm

        self.datamodel = datamodel
        self.allowed_distribution_strategies = allowed_distribution_strategies

        # Compute the global kdict
        self.kdict = self.compute_kdict()
        # Initialize the dictonary which stores all individual index-dicts
        self.global_dict = {}
        # Set self.default_parameters
Ultima's avatar
Ultima committed
56
57
58
        self.set_default(config_dict={'log': log,
                                      'nbin': nbin,
                                      'binbounds': binbounds})
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

    # 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:
Ultima's avatar
Ultima committed
109
110
111
112
            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'])
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
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221

            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 compute_kdict(self):
        raise NotImplementedError(
            about._errors.cstring(
                "ERROR: No generic compute_kdict method implemented."))

    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) =\
                self._compute_indices(self.kdict)
            temp_kdict = self.kdict

        # 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
Ultima's avatar
Ultima committed
222
223
224
225
            temp_unbinned_indices = self.get_index_dict(nbin=None,
                                                        binbounds=None,
                                                        log=False,
                                                        store=False)
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
            # Bin them
            (temp_pindex, temp_kindex, temp_rho, temp_pundex) = \
                self._bin_power_indices(
                    temp_unbinned_indices, **config_dict)
            # Make a binned version of kdict
            temp_kdict = self._compute_kdict_from_pindex_kindex(temp_pindex,
                                                                temp_kindex)

        temp_index_dict = {"config": config_dict,
                           "pindex": temp_pindex,
                           "kindex": temp_kindex,
                           "rho": temp_rho,
                           "pundex": temp_pundex,
                           "kdict": temp_kdict}
        return temp_index_dict

    def _compute_kdict_from_pindex_kindex(self, pindex, kindex):
        if isinstance(pindex, distributed_data_object):
            tempindex = pindex.copy(dtype=kindex.dtype)
245
246
            result = tempindex.apply_scalar_function(
                                lambda x: kindex[x.astype(np.dtype('int'))])
247
248
249
250
251
        else:
            result = kindex[pindex].astype(dtype=kindex.dtype)
        return result

    def _compute_indices(self, nkdict):
csongor's avatar
csongor committed
252
        if self.datamodel in ['np']:
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
            return self._compute_indices_np(nkdict)
        elif self.datamodel in self.allowed_distribution_strategies:
            return self._compute_indices_d2o(nkdict)
        else:
            raise ValueError(about._errors.cstring(
                'ERROR: Datamodel is not supported.'))

    def _compute_indices_d2o(self, nkdict):
        """
        Internal helper function which computes pindex, kindex, rho and pundex
        from a given nkdict
        """
        raise NotImplementedError(
            about._errors.cstring(
                "ERROR: No generic _compute_indices_d2o method implemented."))

    def _compute_pundex_d2o(self, global_pindex, global_kindex):
        """
        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.
        """
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
303
304
305
306
307
308
309
310
311
312
313
        if self.datamodel 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
            self.comm.Allreduce(local_pundex, global_pundex, op=MPI.MIN)
            return global_pundex

        elif self.datamodel in DISTRIBUTION_STRATEGIES['not']:
            ##########
            # pundex #
            ##########
            pundex = np.unique(global_pindex.get_local_data(),
                               return_index=True)[1]
            return pundex
        else:
            raise NotImplementedError(about._errors.cstring(
                "ERROR: _compute_pundex_d2o for given datamodel not " +
                "implemented."))
314
315
316
317
318
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
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510

    def _compute_indices_np(self, nkdict):
        """
        Internal helper function which computes pindex, kindex, rho and pundex
        from a given nkdict
        """
        raise NotImplementedError(
            about._errors.cstring(
                "ERROR: No generic _compute_indices_np method implemented."))

    def _compute_pundex_np(self, pindex, kindex):
        """
        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.
        """
        ##########
        # pundex #
        ##########
        pundex = np.unique(pindex, return_index=True)[1]
        return pundex

    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]

        if self.datamodel == 'np':
            pindex_ = reorder[pindex]
            pundex_ = self._compute_pundex_np(pindex_, kindex_)
        elif self.datamodel in self.allowed_distribution_strategies:
            pindex_ = pindex.copy_empty()
            pindex_.set_local_data(reorder[pindex.get_local_data()])
            pundex_ = self._compute_pundex_d2o(pindex_, kindex_)
        else:
            raise ValueError(about._errors.cstring(
                'ERROR: Datamodel is not supported.'))

        return pindex_, kindex_, rho_, pundex_


class rg_power_indices(power_indices):

    def __init__(self, shape, dgrid, datamodel,
                 allowed_distribution_strategies,
                 zerocentered=False, log=False, nbin=None,
                 binbounds=None, comm=None):
        """
            Returns an instance of the power_indices 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
            zerocentered : 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.
        """
        # Basic inits and consistency checks
        self.shape = np.array(shape, dtype=int)
        self.dgrid = np.abs(np.array(dgrid))
        if self.shape.shape != self.dgrid.shape:
            raise ValueError(about._errors.cstring("ERROR: The supplied shape\
                and dgrid have not the same dimensionality"))
        self.zerocentered = self._cast_zerocentered(zerocentered)

        super(rg_power_indices, self).__init__(
            datamodel=datamodel,
            allowed_distribution_strategies=allowed_distribution_strategies,
            comm=comm,
            log=log,
            nbin=nbin,
            binbounds=binbounds)

    def _cast_zerocentered(self, zerocentered=False):
        """
            internal helper function which brings the zerocentered input in
            the form of a boolean-tuple
        """
        zc = np.array(zerocentered).astype(bool)
        if zc.shape == self.shape.shape:
            return tuple(zc)
        else:
            temp = np.empty(shape=self.shape.shape, dtype=bool)
            temp[:] = zc
            return tuple(temp)

    def compute_kdict(self):
        """
            Calculates an n-dimensional array with its entries being the
            lengths of the k-vectors from the zero point of the grid.

            Parameters
            ----------
            None : All information is taken from the parent object.

            Returns
            -------
            nkdict : distributed_data_object
        """
        shape = self.shape
        if self.datamodel == 'np':
            slice_of_first_dimension = slice(0, shape[0])
            nkdict = self._compute_kdict_helper(slice_of_first_dimension)

        elif self.datamodel in self.allowed_distribution_strategies:
            # prepare the distributed_data_object
            nkdict = distributed_data_object(
                                global_shape=shape,
                                dtype=np.float128,
theos's avatar
theos committed
511
512
                                distribution_strategy=self.datamodel,
                                comm=self.comm)
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
            if self.datamodel in DISTRIBUTION_STRATEGIES['slicing']:
                # get the node's individual slice of the first dimension
                slice_of_first_dimension = slice(
                                        *nkdict.distributor.local_slice[0:2])
            elif self.datamodel in DISTRIBUTION_STRATEGIES['not']:
                slice_of_first_dimension = slice(0, shape[0])
            else:
                raise ValueError(about._errors.cstring(
                    "ERROR: Unsupported distribution strategy"))
            dists = self._compute_kdict_helper(slice_of_first_dimension)
            nkdict.set_local_data(dists)

        else:
            raise ValueError(about._errors.cstring(

                "ERROR: Unsupported datamodel"))
        return nkdict

    def _compute_kdict_helper(self, slice_of_first_dimension):
        dk = self.dgrid
        shape = self.shape

        inds = []
        for a in shape:
            inds += [slice(0, a)]

        cords = np.ogrid[inds]

        dists = ((np.float128(0) + cords[0] - shape[0] // 2) * dk[0])**2
        # apply zerocenteredQ shift
        if self.zerocentered[0] == False:
            dists = np.fft.fftshift(dists)
        # only save the individual slice
        dists = dists[slice_of_first_dimension]
        for ii in range(1, len(shape)):
            temp = ((cords[ii] - shape[ii] // 2) * dk[ii])**2
            if self.zerocentered[ii] == False:
                temp = np.fft.fftshift(temp)
            dists = dists + temp
        dists = np.sqrt(dists)
        return dists

    def _compute_indices_d2o(self, nkdict):
        """
        Internal helper function which computes pindex, kindex, rho and pundex
        from a given nkdict
        """
        ##########
        # kindex #
        ##########
        global_kindex = nkdict.unique()

        ##########
        # pindex #
        ##########
        # compute the local pindex slice on basis of the local nkdict data
        local_pindex = np.searchsorted(global_kindex, nkdict.get_local_data())
        # prepare the distributed_data_object
        global_pindex = distributed_data_object(
                                        global_shape=nkdict.shape,
                                        dtype=local_pindex.dtype,
                                        distribution_strategy=self.datamodel,
                                        comm=self.comm)
        # store the local pindex data in the global_pindex d2o
        global_pindex.set_local_data(local_pindex)

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

        ##########
        # pundex #
        ##########
        global_pundex = self._compute_pundex_d2o(global_pindex,
                                                 global_kindex)

        return global_pindex, global_kindex, global_rho, global_pundex

    def _compute_indices_np(self, nkdict):
        """
        Internal helper function which computes pindex, kindex, rho and pundex
        from a given nkdict
        """
        ##########
        # kindex #
        ##########
        kindex = np.unique(nkdict)

        ##########
        # pindex #
        ##########
        pindex = np.searchsorted(kindex, nkdict)

        #######
        # rho #
        #######
        rho = np.bincount(pindex.flatten())

        ##########
        # pundex #
        ##########
        pundex = self._compute_pundex_np(pindex, kindex)

        return pindex, kindex, rho, pundex


class lm_power_indices(power_indices):

    def __init__(self, lmax, dim, datamodel,
                 allowed_distribution_strategies,
                 zerocentered=False, log=False, nbin=None,
                 binbounds=None, comm=None):
        """
            Returns an instance of the power_indices 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
            zerocentered : 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.
        """
        # Basic inits and consistency checks
        self.lmax = np.uint(lmax)
        self.dim = np.uint(dim)
        super(lm_power_indices, self).__init__(
            datamodel=datamodel,
            allowed_distribution_strategies=allowed_distribution_strategies,
            comm=comm,
            log=log,
            nbin=nbin,
            binbounds=binbounds)

    def compute_kdict(self):
        """
            Calculates an n-dimensional array with its entries being the
            lengths of the k-vectors from the zero point of the grid.

            Parameters
            ----------
            None : All information is taken from the parent object.

            Returns
            -------
            nkdict : distributed_data_object
        """
theos's avatar
theos committed
676
677
678
679
680
681
682
683

        if self.datamodel != 'not':
            about.warnings.cprint(
                "WARNING: full kdict is temporarily stored on every node " +
                "altough disribution strategy != 'not'!")

        if self.datamodel in self.allowed_distribution_strategies:
            if 'healpy' in gdi:
684
685
686
                nkdict = hp.Alm.getlm(self.lmax, i=None)[0]
            else:
                nkdict = self._getlm()[0]
theos's avatar
theos committed
687
688
689
690
            nkdict = distributed_data_object(
                         nkdict,
                         distribution_strategy=self.datamodel,
                         comm=self.comm)
691
692
693
694
695
696
697
698
699
700
701
702
703
        else:
            raise ValueError(about._errors.cstring(
                "ERROR: Unsupported datamodel"))
        return nkdict

    def _getlm(self):  # > compute all (l,m)
        index = np.arange(self.dim)
        n = 2 * self.lmax + 1
        m = np.ceil((n - np.sqrt(n**2 - 8 * (index - self.lmax))) / 2
                    ).astype(np.int)
        l = index - self.lmax * m + m * (m - 1) // 2
        return l, m

csongor's avatar
csongor committed
704
    def _compute_indices(self, nkdict):
theos's avatar
theos committed
705
        if self.datamodel in self.allowed_distribution_strategies:
csongor's avatar
csongor committed
706
707
708
709
710
            return self._compute_indices_d2o(nkdict)
        else:
            raise ValueError(about._errors.cstring(
                'ERROR: Datamodel is not supported.'))

711
712
713
714
715
716
717
718
719
720
721
722
723
    def _compute_indices_d2o(self, nkdict):
        """
        Internal helper function which computes pindex, kindex, rho and pundex
        from a given nkdict
        """
        ##########
        # kindex #
        ##########
        kindex = np.arange(self.lmax + 1, dtype=np.float)

        ##########
        # pindex #
        ##########
theos's avatar
theos committed
724
        pindex = nkdict.copy(dtype=np.int)
725
726
727
728
729
730
731
732
733

        #######
        # rho #
        #######
        rho = (2 * kindex + 1).astype(np.int)

        ##########
        # pundex #
        ##########
theos's avatar
theos committed
734
        pundex = self._compute_pundex_d2o(pindex, kindex)
735
736

        return pindex, kindex, rho, pundex