powerspectrum.py 44.6 KB
Newer Older
Marco Selig's avatar
Marco Selig committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
## NIFTY (Numerical Information Field Theory) has been developed at the
## Max-Planck-Institute for Astrophysics.
##
## Copyright (C) 2013 Max-Planck-Society
##
## Author: Marco Selig
## Project homepage: <http://www.mpa-garching.mpg.de/ift/nifty/>
##
## 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/>.


from __future__ import division
import numpy as np
25
26
from mpi4py import MPI

Ultimanet's avatar
Ultimanet committed
27
from nifty.nifty_about import about
28
29
30
31
32
33
from nifty.nifty_mpi_data import distributed_data_object

class power_indices(object):
    def __init__(self, shape, dgrid, zerocentered=False, log=False, nbin=None, 
                 binbounds=None, comm=MPI.COMM_WORLD):
        """
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
            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
60
        self.comm = comm
Ultimanet's avatar
Ultimanet committed
61
62
        self.shape = np.array(shape, dtype = int)
        self.dgrid = np.abs(np.array(dgrid))
63
        if self.shape.shape != self.dgrid.shape:
64
65
            raise ValueError(about._errors.cstring("ERROR: The supplied shape\
                and dgrid have not the same dimensionality"))         
66
        self.zerocentered = self.__cast_zerocentered__(zerocentered)
Ultimanet's avatar
Ultimanet committed
67
68
69
70

        ## Compute the global kdict
        self.kdict = self.compute_kdict()
        
71
        
72
73
        ## Initialize the dictonary which stores all individual index-dicts
        self.global_dict={}
Ultimanet's avatar
Ultimanet committed
74
        
75
76
77
78
79
        ## Calculate the default dictonory according to the kwargs and set it 
        ## as default
        self.get_index_dict(log=log, nbin=nbin, binbounds=binbounds, 
                            store=True)
        self.set_default(log=log, nbin=nbin, binbounds=binbounds)
80
        
81
82
83
84
85
86
    ## Redirect the direct calls approaching a power_index instance to the 
    ## default_indices dict
    def __getitem__(self, x):
        return self.default_indices.get(x)
    def __getattr__(self, x):
        return self.default_indices.__getattribute__(x)
87
88
        
    def __cast_zerocentered__(self, zerocentered=False):
89
90
91
92
        """        
            internal helper function which brings the zerocentered input in 
            the form of a boolean-tuple
        """
93
94
95
96
97
98
99
100
101
        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 __cast_config__(self, *args, **kwargs):
102
103
104
105
        """
            internal helper function which casts the various combinations of 
            possible parameters into a properly defaulted dictionary
        """
106
107
108
109
110
111
112
113
114
115
116
117
118
        temp_config_dict = kwargs.get('config_dict', None)        
        if temp_config_dict != None:
            return self.__cast_config_helper__(**temp_config_dict)
        else:
            temp_log = kwargs.get("log", None)
            temp_nbin = kwargs.get("nbin", None)
            temp_binbounds = kwargs.get("binbounds", None)            
            
            return self.__cast_config_helper__(log=temp_log, 
                                               nbin=temp_nbin, 
                                               binbounds=temp_binbounds)
    
    def __cast_config_helper__(self, log, nbin, binbounds):
119
120
121
122
123
        """
            internal helper function which sets the defaults for the 
            __cast_config__ function
        """
        
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
        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 __freeze_config__(self, config_dict):
145
146
147
148
        """
            a helper function which forms a hashable identifying object from 
            a config dictionary which can be used as key of a dict
        """        
149
150
151
152
        return frozenset(config_dict.items())
        
    def set_default(self, *args, **kwargs):
        """
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
            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    
        """ 
        ## This shortcut relies on the fact, that get_index_dict returns a
        ## reference on the default dict and not a copy!!
        self.default_indices = self.get_index_dict(*args, **kwargs)         
173
174
175
176
        
    
    def get_index_dict(self, *args, **kwargs):
        """
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
            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        
201
        temp_config_dict = self.__cast_config__(*args, **kwargs)
202
203
        ## Compute a hashable identifier from the config which will be used 
        ## as dict key
204
        temp_key = self.__freeze_config__(temp_config_dict)
205
        ## Check if the result should be stored for future use.
206
        storeQ = kwargs.get("store", True)
207
        ## Try to find the requested index dict in the global_dict
208
209
210
        try:
            return self.global_dict[temp_key]
        except(KeyError):
211
212
213
            ## If it is not found, calculate it.
            temp_index_dict = self.__compute_index_dict__(temp_config_dict)
            ## Store it, if required
214
215
            if storeQ == True:
                self.global_dict[temp_key] = temp_index_dict
216
217
218
                ## Important: If the result is stored, return a reference to 
                ## the dictionary entry, not anly a plain copy. Otherwise, 
                ## set_default breaks!
219
220
                return self.global_dict[temp_key]
            else:
221
                ## Return the plain result.
222
223
224
                return temp_index_dict
        
    
Ultimanet's avatar
Ultimanet committed
225
    def compute_kdict(self):
226
        """
227
228
229
230
231
232
233
234
235
236
            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
237
238
        """
        
239
        
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
        ##if(fourier):
        ##   dk = dgrid
        ##else:
        ##    dk = np.array([1/dgrid[i]/axes[i] for i in range(len(axes))])
        
        dk = self.dgrid
        shape = self.shape
        
        ## prepare the distributed_data_object        
        nkdict = distributed_data_object(global_shape=shape, 
                                         dtype=np.float128, 
                                         distribution_strategy="fftw")
        ## get the node's individual slice of the first dimension 
        slice_of_first_dimension = slice(*nkdict.distributor.local_slice[0:2])
        
        inds = []
        for a in shape:
            inds += [slice(0,a)]
        
        cords = np.ogrid[inds]
Ultimanet's avatar
Ultimanet committed
260

261
262
263
264
265
266
267
268
269
270
271
272
273
274
        dists = ((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)
        nkdict.set_local_data(dists)
        return nkdict
Ultimanet's avatar
Ultimanet committed
275
276
277
278
279
280
281
282
283
284
    
#    def compute_klength(self, kdict):
#        local_klength = np.sort(list(set(kdict.get_local_data().flatten())))
#        
#        global_klength = kdict.distributor._allgather(local_klength)
#        global_klength = np.array(global_klength).flatten()
#        global_klength = np.sort(list(set(global_klength)))
#
#        return global_klength

285

286
287
288
289
290
    def __compute_indices__(self, nkdict):
        """
        Internal helper function which computes pindex, kindex, rho and pundex
        from a given nkdict
        """
291
292
293
294
295
296
297
298
299
300
        ##########
        # kindex #        
        ##########
        ## compute the local kindex array        
        local_kindex = np.unique(nkdict.get_local_data())
        ## unify the local_kindex arrays
        global_kindex = self.comm.allgather(local_kindex)
        ## flatten the gathered lists        
        global_kindex = np.hstack(global_kindex)
        ## remove duplicates        
Ultimanet's avatar
Ultimanet committed
301
302
        global_kindex = np.unique(global_kindex)        

303
304
305
306
307
308
309
310
311
312
313
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
        ##########
        # 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.type,
                                                distribution_strategy='fftw')  
        ## store the local pindex data in the global_pindex d2o
        global_pindex.set_local_data(local_pindex)
        
        #######
        # rho #        
        #######
        ## Prepare the local pindex data in order to conut the degeneracy 
        ## factors
        temp = local_pindex.flatten()
        ## Remember: np.array.sort is an inplace function        
        temp.sort()        
        ## In local_count we save how many of the indvidual values in 
        ## local_value occured. Therefore we use np.unique to calculate the 
        ## offset...
        local_value, local_count = np.unique(temp, return_index=True)
        ## ...and then build the offset differences
        if local_count.shape != (0,):
            local_count = np.append(local_count[1:]-local_count[:-1],
                                    [temp.shape[0]-local_count[-1]])
        ## Prepare the global rho array, and store the individual counts in it
        ## rho has the same length as the kindex array
        local_rho = np.zeros(shape=global_kindex.shape, dtype=np.int)
        global_rho = np.empty_like(local_rho)
        ## Store the individual counts
        local_rho[local_value] = local_count
        ## Use Allreduce to spread the information
        self.comm.Allreduce(local_rho , global_rho, op=MPI.SUM)
        ##########
        # pundex #        
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
        ##########  
        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):
        """
        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 #        
        ##########
        ## Prepare the local data
        local_pindex = global_pindex.get_local_data()
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
        ## 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)
377
        return global_pundex
Ultimanet's avatar
Ultimanet committed
378
379
380
381
382
    
    def __compute_kdict_from_pindex_kindex__(self, pindex, kindex):
        tempindex = pindex.copy(dtype=kindex.dtype.type)        
        return tempindex.apply_scalar_function(lambda x: kindex[x])

383
384
385
386
387
388
389
390
391
    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 config_dict["log"]==False and config_dict["nbin"]==None and \
          config_dict["binbounds"]==None:
Ultimanet's avatar
Ultimanet committed
392
393
394
            (temp_pindex, temp_kindex, temp_rho, temp_pundex) =\
                                        self.__compute_indices__(self.kdict)
            temp_kdict = self.kdict
395
396
397
398
            
        ## if binning is required, make a recursive call to get the unbinned
        ## indices, bin them, compute the pundex and then return everything.
        else:
Ultimanet's avatar
Ultimanet committed
399
400
401
            ## Get the unbinned indices 
            temp_unbinned_indices = self.get_index_dict(store=False)            
            ## Bin them            
402
403
            (temp_pindex, temp_kindex, temp_rho, temp_pundex) = \
                self.__bin_power_indices__(temp_unbinned_indices, **config_dict)
Ultimanet's avatar
Ultimanet committed
404
405
406
407
            ## Make a binned version of kdict
            temp_kdict = self.__compute_kdict_from_pindex_kindex__(temp_pindex, 
                                                                   temp_kindex)
            
408
409
410
411
        temp_index_dict = {"config": config_dict, 
                               "pindex": temp_pindex,
                               "kindex": temp_kindex,
                               "rho": temp_rho,
Ultimanet's avatar
Ultimanet committed
412
413
                               "pundex": temp_pundex,
                               "kdict": temp_kdict}
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
        return temp_index_dict

    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]
        
        pindex_ = pindex.copy_empty()
        pindex_.set_local_data(reorder[pindex.get_local_data()])
        
        pundex_ = self.__compute_pundex__(pindex_, kindex_)     
        return pindex_, kindex_, rho_, pundex_
491
492
493



Marco Selig's avatar
Marco Selig committed
494
495


496
def draw_vector_nd(axes,dgrid,ps,symtype=0,fourier=False,zerocentered=False,kpack=None):
Marco Selig's avatar
Marco Selig committed
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534

    """
        Draws a n-dimensional field on a regular grid from a given power
        spectrum. The grid parameters need to be specified, together with a
        couple of global options explained below. The dimensionality of the
        field is determined automatically.

        Parameters
        ----------
        axes : ndarray
            An array with the length of each axis.

        dgrid : ndarray
            An array with the pixel length of each axis.

        ps : ndarray
            The power spectrum as a function of Fourier modes.

        symtype : int {0,1,2} : *optional*
            Whether the output should be real valued (0), complex-hermitian (1)
            or complex without symmetry (2). (default=0)

        fourier : bool : *optional*
            Whether the output should be in Fourier space or not
            (default=False).

        zerocentered : bool : *optional*
            Whether the output array should be zerocentered, i.e. starting with
            negative Fourier modes going over the zero mode to positive modes,
            or not zerocentered, where zero, positive and negative modes are
            simpy ordered consecutively.

        Returns
        -------
        field : ndarray
            The drawn random field.

    """
535
    if(kpack is None):
Marco Selig's avatar
Marco Selig committed
536
        kdict = np.fft.fftshift(nkdict_fast(axes,dgrid,fourier))
537
538
        klength = nklength(kdict)
    else:
539
        kdict = kpack[1][np.fft.ifftshift(kpack[0],axes=shiftaxes(zerocentered,st_to_zero_mode=False))]
540
        klength = kpack[1]
Marco Selig's avatar
Marco Selig committed
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

    #output is in position space
    if(not fourier):

        #output is real-valued
        if(symtype==0):
            vector = drawherm(klength,kdict,ps)
            if(np.any(zerocentered==True)):
                return np.real(np.fft.fftshift(np.fft.ifftn(vector),axes=shiftaxes(zerocentered)))
            else:
                return np.real(np.fft.ifftn(vector))

        #output is complex with hermitian symmetry
        elif(symtype==1):
            vector = drawwild(klength,kdict,ps,real_corr=2)
            if(np.any(zerocentered==True)):
                return np.fft.fftshift(np.fft.ifftn(np.real(vector)),axes=shiftaxes(zerocentered))
            else:
                return np.fft.ifftn(np.real(vector))

        #output is complex without symmetry
        else:
            vector = drawwild(klength,kdict,ps)
            if(np.any(zerocentered==True)):
                return np.fft.fftshift(np.fft.ifftn(vector),axes=shiftaxes(zerocentered))
            else:
                return np.fft.ifftn(vector)

    #output is in fourier space
    else:

        #output is real-valued
        if(symtype==0):
            vector = drawwild(klength,kdict,ps,real_corr=2)
            if np.any(zerocentered == True):
                return np.real(np.fft.fftshift(vector,axes=shiftaxes(zerocentered)))
            else:
                return np.real(vector)

        #output is complex with hermitian symmetry
        elif(symtype==1):
            vector = drawherm(klength,kdict,ps)
            if(np.any(zerocentered==True)):
                return np.fft.fftshift(vector,axes=shiftaxes(zerocentered))
            else:
                return vector

        #output is complex without symmetry
        else:
            vector = drawwild(klength,kdict,ps)
            if(np.any(zerocentered==True)):
                return np.fft.fftshift(vector,axes=shiftaxes(zerocentered))
            else:
                return vector


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
#def calc_ps(field,axes,dgrid,zerocentered=False,fourier=False):
#
#    """
#        Calculates the power spectrum of a given field assuming that the field
#        is statistically homogenous and isotropic.
#
#        Parameters
#        ----------
#        field : ndarray
#            The input field from which the power spectrum should be determined.
#
#        axes : ndarray
#            An array with the length of each axis.
#
#        dgrid : ndarray
#            An array with the pixel length of each axis.
#
#        zerocentered : bool : *optional*
#            Whether the output array should be zerocentered, i.e. starting with
#            negative Fourier modes going over the zero mode to positive modes,
#            or not zerocentered, where zero, positive and negative modes are
#            simpy ordered consecutively.
#
#        fourier : bool : *optional*
#            Whether the output should be in Fourier space or not
#            (default=False).
#
#    """
#
#    ## field absolutes
#    if(not fourier):
#        foufield = np.fft.fftshift(np.fft.fftn(field))
#    elif(np.any(zerocentered==False)):
#        foufield = np.fft.fftshift(field, axes=shiftaxes(zerocentered,st_to_zero_mode=True))
#    else:
#        foufield = field
#    fieldabs = np.abs(foufield)**2
#
Marco Selig's avatar
Marco Selig committed
635
#    kdict = nkdict_fast(axes,dgrid,fourier)
636
637
638
639
640
641
642
643
644
645
646
647
648
#    klength = nklength(kdict)
#
#    ## power spectrum
#    ps = np.zeros(klength.size)
#    rho = np.zeros(klength.size)
#    for ii in np.ndindex(kdict.shape):
#        position = np.searchsorted(klength,kdict[ii])
#        rho[position] += 1
#        ps[position] += fieldabs[ii]
#    ps = np.divide(ps,rho)
#    return ps

def calc_ps_fast(field,axes,dgrid,zerocentered=False,fourier=False,pindex=None,kindex=None,rho=None):
Marco Selig's avatar
Marco Selig committed
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

    """
        Calculates the power spectrum of a given field faster assuming that the
        field is statistically homogenous and isotropic.

        Parameters
        ----------
        field : ndarray
            The input field from which the power spectrum should be determined.

        axes : ndarray
            An array with the length of each axis.

        dgrid : ndarray
            An array with the pixel length of each axis.

        zerocentered : bool : *optional*
            Whether the output array should be zerocentered, i.e. starting with
            negative Fourier modes going over the zero mode to positive modes,
            or not zerocentered, where zero, positive and negative modes are
            simpy ordered consecutively.

        fourier : bool : *optional*
            Whether the output should be in Fourier space or not
            (default=False).

675
676
677
678
679
680
681
682
683
684
        pindex : ndarray
            Index of the Fourier grid points in a numpy.ndarray ordered
            following the zerocentered flag (default=None).

        kindex : ndarray
            Array of all k-vector lengths (default=None).

        rho : ndarray
            Degeneracy of the Fourier grid, indicating how many k-vectors in
            Fourier space have the same length (default=None).
Marco Selig's avatar
Marco Selig committed
685
686
687
688

    """
    ## field absolutes
    if(not fourier):
Marco Selig's avatar
Marco Selig committed
689
690
691
        foufield = np.fft.fftshift(np.fft.fftn(field))
    elif(np.any(zerocentered==False)):
        foufield = np.fft.fftshift(field, axes=shiftaxes(zerocentered,st_to_zero_mode=True))
Marco Selig's avatar
Marco Selig committed
692
693
694
    else:
        foufield = field
    fieldabs = np.abs(foufield)**2
695
696
697
698

    if(rho is None):
        if(pindex is None):
            ## kdict
Marco Selig's avatar
Marco Selig committed
699
            kdict = nkdict_fast(axes,dgrid,fourier)
700
701
702
703
704
705
706
707
708
709
710
711
712
            ## klength
            if(kindex is None):
                klength = nklength(kdict)
            else:
                klength = kindex
            ## power spectrum
            ps = np.zeros(klength.size)
            rho = np.zeros(klength.size)
            for ii in np.ndindex(kdict.shape):
                position = np.searchsorted(klength,kdict[ii])
                ps[position] += fieldabs[ii]
                rho[position] += 1
        else:
Marco Selig's avatar
Marco Selig committed
713
714
715
            ## zerocenter pindex
            if(np.any(zerocentered==False)):
                pindex = np.fft.fftshift(pindex, axes=shiftaxes(zerocentered,st_to_zero_mode=True))
716
            ## power spectrum
717
            ps = np.zeros(np.max(pindex)+1)
718
719
720
721
722
723
            rho = np.zeros(ps.size)
            for ii in np.ndindex(pindex.shape):
                ps[pindex[ii]] += fieldabs[ii]
                rho[pindex[ii]] += 1
    elif(pindex is None):
        ## kdict
Marco Selig's avatar
Marco Selig committed
724
        kdict = nkdict_fast(axes,dgrid,fourier)
725
726
727
728
729
730
731
732
733
734
735
        ## klength
        if(kindex is None):
            klength = nklength(kdict)
        else:
            klength = kindex
        ## power spectrum
        ps = np.zeros(klength.size)
        for ii in np.ndindex(kdict.shape):
            position = np.searchsorted(klength,kdict[ii])
            ps[position] += fieldabs[ii]
    else:
Marco Selig's avatar
Marco Selig committed
736
737
738
        ## zerocenter pindex
        if(np.any(zerocentered==False)):
            pindex = np.fft.fftshift(pindex, axes=shiftaxes(zerocentered,st_to_zero_mode=True))
739
740
741
742
743
744
        ## power spectrum
        ps = np.zeros(rho.size)
        for ii in np.ndindex(pindex.shape):
            ps[pindex[ii]] += fieldabs[ii]

    ps = np.divide(ps,rho)
Marco Selig's avatar
Marco Selig committed
745
746
    return ps

Ultimanet's avatar
Ultimanet committed
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
#
#def get_power_index(axes,dgrid,zerocentered,irred=False,fourier=True):
#
#    """
#        Returns the index of the Fourier grid points in a numpy
#        array, ordered following the zerocentered flag.
#
#        Parameters
#        ----------
#        axes : ndarray
#            An array with the length of each axis.
#
#        dgrid : ndarray
#            An array with the pixel length of each axis.
#
#        zerocentered : bool
#            Whether the output array should be zerocentered, i.e. starting with
#            negative Fourier modes going over the zero mode to positive modes,
#            or not zerocentered, where zero, positive and negative modes are
#            simpy ordered consecutively.
#
#        irred : bool : *optional*
#            If True, the function returns an array of all k-vector lengths and
#            their degeneracy factors. If False, just the power index array is
#            returned.
#
#        fourier : bool : *optional*
#            Whether the output should be in Fourier space or not
#            (default=False).
#
#        Returns
#        -------
#            index or {klength, rho} : scalar or list
#                Returns either an array of all k-vector lengths and
#                their degeneracy factors or just the power index array
#                depending on the flag irred.
#
#    """
#
#    ## kdict, klength
#    if(np.any(zerocentered==False)):
#        kdict = np.fft.fftshift(nkdict_fast(axes,dgrid,fourier),axes=shiftaxes(zerocentered,st_to_zero_mode=True))
#    else:
#        kdict = nkdict_fast(axes,dgrid,fourier)
#    klength = nklength(kdict)
#    ## output
#    if(irred):
#        rho = np.zeros(klength.shape,dtype=np.int)
#        for ii in np.ndindex(kdict.shape):
#            rho[np.searchsorted(klength,kdict[ii])] += 1
#        return klength,rho
#    else:
#        ind = np.empty(axes,dtype=np.int)
#        for ii in np.ndindex(kdict.shape):
#            ind[ii] = np.searchsorted(klength,kdict[ii])
#        return ind
#
#
#def get_power_indices(axes,dgrid,zerocentered,fourier=True):
#    """
#        Returns the index of the Fourier grid points in a numpy
#        array, ordered following the zerocentered flag.
#
#        Parameters
#        ----------
#        axes : ndarray
#            An array with the length of each axis.
#
#        dgrid : ndarray
#            An array with the pixel length of each axis.
#
#        zerocentered : bool
#            Whether the output array should be zerocentered, i.e. starting with
#            negative Fourier modes going over the zero mode to positive modes,
#            or not zerocentered, where zero, positive and negative modes are
#            simpy ordered consecutively.
#
#        irred : bool : *optional*
#            If True, the function returns an array of all k-vector lengths and
#            their degeneracy factors. If False, just the power index array is
#            returned.
#
#        fourier : bool : *optional*
#            Whether the output should be in Fourier space or not
#            (default=False).
#
#        Returns
#        -------
#        index, klength, rho : ndarrays
#            Returns the power index array, an array of all k-vector lengths and
#            their degeneracy factors.
#
#    """
#
#    ## kdict, klength
#    if(np.any(zerocentered==False)):
#        kdict = np.fft.fftshift(nkdict_fast(axes,dgrid,fourier),axes=shiftaxes(zerocentered,st_to_zero_mode=True))
#    else:
#        kdict = nkdict_fast(axes,dgrid,fourier)
#    klength = nklength(kdict)
#    ## output
#    ind = np.empty(axes,dtype=np.int)
#    rho = np.zeros(klength.shape,dtype=np.int)
#    for ii in np.ndindex(kdict.shape):
#        ind[ii] = np.searchsorted(klength,kdict[ii])
#        rho[ind[ii]] += 1
#    return ind,klength,rho
#
#
#def get_power_indices2(axes,dgrid,zerocentered,fourier=True):
#    """
#        Returns the index of the Fourier grid points in a numpy
#        array, ordered following the zerocentered flag.
#
#        Parameters
#        ----------
#        axes : ndarray
#            An array with the length of each axis.
#
#        dgrid : ndarray
#            An array with the pixel length of each axis.
#
#        zerocentered : bool
#            Whether the output array should be zerocentered, i.e. starting with
#            negative Fourier modes going over the zero mode to positive modes,
#            or not zerocentered, where zero, positive and negative modes are
#            simpy ordered consecutively.
#
#        irred : bool : *optional*
#            If True, the function returns an array of all k-vector lengths and
#            their degeneracy factors. If False, just the power index array is
#            returned.
#
#        fourier : bool : *optional*
#            Whether the output should be in Fourier space or not
#            (default=False).
#
#        Returns
#        -------
#        index, klength, rho : ndarrays
#            Returns the power index array, an array of all k-vector lengths and
#            their degeneracy factors.
#
#    """
#
#    ## kdict, klength
#    if(np.any(zerocentered==False)):
#        kdict = np.fft.fftshift(nkdict_fast2(axes,dgrid,fourier),axes=shiftaxes(zerocentered,st_to_zero_mode=True))
#    else:
#        kdict = nkdict_fast2(axes,dgrid,fourier)
#
#    klength,rho,ind = nkdict_to_indices(kdict)
#
#    return ind,klength,rho
#
#def nkdict_to_indices(kdict):
#
#    kindex,pindex = np.unique(kdict,return_inverse=True)
#    pindex = pindex.reshape(kdict.shape)
#
#    rho = pindex.flatten()
#    rho.sort()
#    rho = np.unique(rho,return_index=True,return_inverse=False)[1]
#    rho = np.append(rho[1:]-rho[:-1],[np.prod(pindex.shape)-rho[-1]])
#
#    return kindex,rho,pindex
#
#
#
#def bin_power_indices(pindex,kindex,rho,log=False,nbin=None,binbounds=None):
#    """
#        Returns the (re)binned power indices associated with the Fourier grid.
#
#        Parameters
#        ----------
#        pindex : ndarray
#            Index of the Fourier grid points in a numpy.ndarray ordered
#            following the zerocentered flag (default=None).
#        kindex : ndarray
#            Array of all k-vector lengths (default=None).
#        rho : ndarray
#            Degeneracy of the Fourier grid, indicating how many k-vectors in
#            Fourier space have the same length (default=None).
#        log : bool
#            Flag specifying if the binning is performed on logarithmic scale
#            (default: False).
#        nbin : integer
#            Number of used bins (default: None).
#        binbounds : {list, array}
#            Array-like inner boundaries of the used bins (default: None).
#
#        Returns
#        -------
#        pindex, kindex, rho : ndarrays
#            The (re)binned power indices.
#
#    """
#    ## 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]
#
#    return reorder[pindex],kindex_,rho_
#
978

979

Marco Selig's avatar
Marco Selig committed
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
def nhermitianize(field,zerocentered):

    """
        Hermitianizes an arbitrary n-dimensional field. Becomes relatively slow
        for large n.

        Parameters
        ----------
        field : ndarray
            The input field that should be hermitianized.

        zerocentered : bool
            Whether the output array should be zerocentered, i.e. starting with
            negative Fourier modes going over the zero mode to positive modes,
            or not zerocentered, where zero, positive and negative modes are
            simpy ordered consecutively.

        Returns
        -------
        hermfield : ndarray
            The hermitianized field.

    """
    ## shift zerocentered axes
    if(np.any(zerocentered==True)):
        field = np.fft.fftshift(field, axes=shiftaxes(zerocentered))
#    for index in np.ndenumerate(field):
#        negind = tuple(-np.array(index[0]))
#        field[negind] = np.conjugate(index[1])
#        if(field[negind]==field[index[0]]):
#            field[index[0]] = np.abs(index[1])*(np.sign(index[1].real)+(np.sign(index[1].real)==0)*np.sign(index[1].imag)).astype(np.int)
    subshape = np.array(field.shape,dtype=np.int) ## == axes
    maxindex = subshape//2
    subshape[np.argmax(subshape)] = subshape[np.argmax(subshape)]//2+1 ## ~half larges axis
    for ii in np.ndindex(tuple(subshape)):
        negii = tuple(-np.array(ii))
        field[negii] = np.conjugate(field[ii])
    for ii in np.ndindex((2,)*maxindex.size):
        index = tuple(ii*maxindex)
        field[index] = np.abs(field[index])*(np.sign(field[index].real)+(np.sign(field[index].real)==0)*-np.sign(field[index].imag)).astype(np.int) ## minus since overwritten before
    ## reshift zerocentered axes
    if(np.any(zerocentered==True)):
        field = np.fft.fftshift(field,axes=shiftaxes(zerocentered))
    return field

def nhermitianize_fast(field,zerocentered,special=False):

    """
        Hermitianizes an arbitrary n-dimensional field faster.
        Still becomes comparably slow for large n.

        Parameters
        ----------
        field : ndarray
            The input field that should be hermitianized.

        zerocentered : bool
            Whether the output array should be zerocentered, i.e. starting with
            negative Fourier modes going over the zero mode to positive modes,
            or not zerocentered, where zero, positive and negative modes are
            simpy ordered consecutively.

        special : bool, *optional*
            Must be True for random fields drawn from Gaussian or pm1
            distributions.

        Returns
        -------
        hermfield : ndarray
            The hermitianized field.

    """
    ## shift zerocentered axes
    if(np.any(zerocentered==True)):
        field = np.fft.fftshift(field, axes=shiftaxes(zerocentered))
    dummy = np.conjugate(field)
    ## mirror conjugate field
    for ii in range(field.ndim):
        dummy = np.swapaxes(dummy,0,ii)
        dummy = np.flipud(dummy)
        dummy = np.roll(dummy,1,axis=0)
        dummy = np.swapaxes(dummy,0,ii)
    if(special): ## special normalisation for certain random fields
        field = np.sqrt(0.5)*(field+dummy)
        maxindex = np.array(field.shape,dtype=np.int)//2
Ultimanet's avatar
Ultimanet committed
1065
        print ('maxindex: ', maxindex)
Marco Selig's avatar
Marco Selig committed
1066
        for ii in np.ndindex((2,)*maxindex.size):
Ultimanet's avatar
Ultimanet committed
1067
            print ('ii: ', ii)
Marco Selig's avatar
Marco Selig committed
1068
            index = tuple(ii*maxindex)
Ultimanet's avatar
Ultimanet committed
1069
            print ('index: ', index)
Marco Selig's avatar
Marco Selig committed
1070
1071
            field[index] *= np.sqrt(0.5)
    else: ## regular case
Ultimanet's avatar
Ultimanet committed
1072
1073
        field = 0.5*(field+dummy)
        #field = dummy
Marco Selig's avatar
Marco Selig committed
1074
1075
1076
1077
    ## reshift zerocentered axes
    if(np.any(zerocentered==True)):
        field = np.fft.fftshift(field,axes=shiftaxes(zerocentered))
    return field
1078
1079
    
    
Marco Selig's avatar
Marco Selig committed
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
def random_hermitian_pm1(datatype,zerocentered,shape):

    """
        Draws a set of hermitianized random, complex pm1 numbers.

    """

    field = np.random.randint(4,high=None,size=np.prod(shape,axis=0,dtype=np.int,out=None)).reshape(shape,order='C')
    dummy = np.copy(field)
    ## mirror field
    for ii in range(field.ndim):
        dummy = np.swapaxes(dummy,0,ii)
        dummy = np.flipud(dummy)
        dummy = np.roll(dummy,1,axis=0)
        dummy = np.swapaxes(dummy,0,ii)
    field = (field+dummy+2*(field>dummy)*((field+dummy)%2))%4 ## wicked magic
    x = np.array([1+0j,0+1j,-1+0j,0-1j],dtype=datatype)[field]
    ## (re)shift zerocentered axes
    if(np.any(zerocentered==True)):
        field = np.fft.fftshift(field,axes=shiftaxes(zerocentered))
    return x


#-----------------------------------------------------------------------------
# Auxiliary functions
#-----------------------------------------------------------------------------

def shiftaxes(zerocentered,st_to_zero_mode=False):

    """
        Shifts the axes in a special way needed for some functions
    """

    axes = []
    for ii in range(len(zerocentered)):
        if(st_to_zero_mode==False)and(zerocentered[ii]):
            axes += [ii]
        if(st_to_zero_mode==True)and(not zerocentered[ii]):
            axes += [ii]
    return axes


Ultimanet's avatar
Ultimanet committed
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
#def nkdict(axes,dgrid,fourier=True):
#    """
#        Calculates an n-dimensional array with its entries being the lengths of
#        the k-vectors from the zero point of the Fourier grid.
#
#    """
#    if(fourier):
#        dk = dgrid
#    else:
#        dk = np.array([1/axes[i]/dgrid[i] for i in range(len(axes))])
#
#    kdict = np.empty(axes)
#    for ii in np.ndindex(kdict.shape):
#        kdict[ii] = np.sqrt(np.sum(((ii-axes//2)*dk)**2))
#    return kdict
#
#
Marco Selig's avatar
Marco Selig committed
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
def nkdict_fast(axes,dgrid,fourier=True):
    """
        Calculates an n-dimensional array with its entries being the lengths of
        the k-vectors from the zero point of the Fourier grid.

    """
    if(fourier):
        dk = dgrid
    else:
        dk = np.array([1/dgrid[i]/axes[i] for i in range(len(axes))])

    temp_vecs = np.array(np.where(np.ones(axes)),dtype='float').reshape(np.append(len(axes),axes))
    temp_vecs = np.rollaxis(temp_vecs,0,len(temp_vecs.shape))
    temp_vecs -= axes//2
    temp_vecs *= dk
    temp_vecs *= temp_vecs
    return np.sqrt(np.sum((temp_vecs),axis=-1))


Ultimanet's avatar
Ultimanet committed
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
#def nkdict_fast2(axes,dgrid,fourier=True):
#    """
#        Calculates an n-dimensional array with its entries being the lengths of
#        the k-vectors from the zero point of the grid.
#
#    """
#    if(fourier):
#        dk = dgrid
#    else:
#        dk = np.array([1/dgrid[i]/axes[i] for i in range(len(axes))])
#
#    inds = []
#    for a in axes:
#        inds += [slice(0,a)]
#    cords = np.ogrid[inds]
#
#    dists = ((cords[0]-axes[0]//2)*dk[0])**2
#    for ii in range(1,len(axes)):
#        dists = dists + ((cords[ii]-axes[ii]//2)*dk[ii])**2
#    dists = np.sqrt(dists)
#
#    return dists
#
#
Marco Selig's avatar
Marco Selig committed
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
def nklength(kdict):
    return np.sort(list(set(kdict.flatten())))


def drawherm(klength,kdict,ps):

    """
        Draws a hermitian random field from a Gaussian distribution.

    """

#    vector = np.zeros(kdict.shape,dtype='complex')
#    for ii in np.ndindex(vector.shape):
#        if(vector[ii]==np.complex(0.,0.)):
#            vector[ii] = np.sqrt(0.5*ps[np.searchsorted(klength,kdict[ii])])*np.complex(np.random.normal(0.,1.),np.random.normal(0.,1.))
#            negii = tuple(-np.array(ii))
#            vector[negii] = np.conjugate(vector[ii])
#            if(vector[negii]==vector[ii]):
#                vector[ii] = np.float(np.sqrt(ps[np.searchsorted(klength,kdict[ii])]))*np.random.normal(0.,1.)
#    return vector
    vec = np.random.normal(loc=0,scale=1,size=kdict.size).reshape(kdict.shape)
    vec = np.fft.fftn(vec)/np.sqrt(np.prod(kdict.shape))
    for ii in np.ndindex(kdict.shape):
        vec[ii] *= np.sqrt(ps[np.searchsorted(klength,kdict[ii])])
    return vec


#def drawwild(vector,klength,kdict,ps,real_corr=1): ## vector = np.zeros(kdict.shape,dtype=np.complex)
#    for ii in np.ndindex(vector.shape):
#        vector[ii] = np.sqrt(real_corr*0.5*ps[klength==kdict[ii]])*np.complex(np.random.normal(0.,1.),np.random.normal(0.,1.))
#    return vector

def drawwild(klength,kdict,ps,real_corr=1):

    """
        Draws a field of arbitrary symmetry from a Gaussian distribution.

    """

    vec = np.empty(kdict.size,dtype=np.complex)
    vec.real = np.random.normal(loc=0,scale=np.sqrt(real_corr*0.5),size=kdict.size)
    vec.imag = np.random.normal(loc=0,scale=np.sqrt(real_corr*0.5),size=kdict.size)
    vec = vec.reshape(kdict.shape)
    for ii in np.ndindex(kdict.shape):
        vec[ii] *= np.sqrt(ps[np.searchsorted(klength,kdict[ii])])
    return vec