foscat 3.0.9__py3-none-any.whl → 3.6.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
foscat/alm.py ADDED
@@ -0,0 +1,690 @@
1
+ import healpy as hp
2
+ import numpy as np
3
+ import time
4
+
5
+ class alm():
6
+
7
+ def __init__(self,backend=None,lmax=24,
8
+ nside=None,limit_range=1E10):
9
+
10
+ if backend is None:
11
+ import foscat.scat_cov as sc
12
+ self.sc=sc.funct()
13
+ self.backend=self.sc.backend
14
+ else:
15
+ self.backend=backend.backend
16
+
17
+ self._logtab={}
18
+ self.lth={}
19
+ self.lph={}
20
+ self.matrix_shift_ph={}
21
+ self.ratio_mm={}
22
+ self.P_mm={}
23
+ self.A={}
24
+ self.B={}
25
+ if nside is not None:
26
+ self.maxlog=6*nside+1
27
+ self.lmax=3*nside
28
+ else:
29
+ self.lmax=lmax
30
+ self.maxlog=2*lmax+1
31
+
32
+ for k in range(1,self.maxlog):
33
+ self._logtab[k]=self.backend.bk_log(self.backend.bk_cast(k))
34
+ self._logtab[0]=0.0
35
+
36
+ if nside is not None:
37
+ self.ring_th(nside)
38
+ self.ring_ph(nside)
39
+ self.shift_ph(nside)
40
+
41
+ self._limit_range=1/limit_range
42
+ self._log_limit_range=np.log(limit_range)
43
+
44
+
45
+ self.Yp={}
46
+ self.Ym={}
47
+
48
+ def ring_th(self,nside):
49
+ if nside not in self.lth:
50
+ n=0
51
+ ith=[]
52
+ for k in range(nside-1):
53
+ N=4*(k+1)
54
+ ith.append(n)
55
+ n+=N
56
+
57
+ for k in range(2*nside+1):
58
+ N=4*nside
59
+ ith.append(n)
60
+ n+=N
61
+ for k in range(nside-1):
62
+ N=4*(nside-1-k)
63
+ ith.append(n)
64
+ n+=N
65
+
66
+ th,ph=hp.pix2ang(nside,ith)
67
+
68
+ self.lth[nside]=th
69
+ return self.lth[nside]
70
+
71
+ def ring_ph(self,nside):
72
+ if nside not in self.lph:
73
+ n=0
74
+ iph=[]
75
+ for k in range(nside-1):
76
+ N=4*(k+1)
77
+ iph.append(n)
78
+ n+=N
79
+
80
+ for k in range(2*nside+1):
81
+ N=4*nside
82
+ iph.append(n)
83
+ n+=N
84
+ for k in range(nside-1):
85
+ N=4*(nside-1-k)
86
+ iph.append(n)
87
+ n+=N
88
+
89
+ th,ph=hp.pix2ang(nside,iph)
90
+
91
+ self.lph[nside]=ph
92
+
93
+ def shift_ph(self,nside):
94
+
95
+ if nside not in self.matrix_shift_ph:
96
+ self.ring_th(nside)
97
+ self.ring_ph(nside)
98
+ x=(-1J*np.arange(3*nside)).reshape(1,3*nside)
99
+ self.matrix_shift_ph[nside]=self.backend.bk_cast(self.backend.bk_exp(x*self.lph[nside].reshape(4*nside-1,1)))
100
+
101
+ self.lmax=3*nside-1
102
+
103
+ ratio_mm={}
104
+
105
+ for m in range(3*nside):
106
+ val=np.zeros([self.lmax-m+1])
107
+ aval=np.zeros([self.lmax-m+1])
108
+ bval=np.zeros([self.lmax-m+1])
109
+
110
+ if m>0:
111
+ val[0]=self.double_factorial_log(2*m - 1)-0.5*np.sum(np.log(1+np.arange(2*m)))
112
+ else:
113
+ val[0]=self.double_factorial_log(2*m - 1)
114
+ if m<self.lmax:
115
+ aval[1]=(2*m + 1)
116
+ val[1] = val[0]-0.5*self.log(2*m+1)
117
+
118
+ for l in range(m + 2, self.lmax+1):
119
+ aval[l-m]=(2*l - 1)/ (l - m)
120
+ bval[l-m]=(l + m - 1)/ (l - m)
121
+ val[l-m] = val[l-m-1] + 0.5*self.log(l-m) - 0.5*self.log(l+m)
122
+
123
+ self.A[nside,m]=self.backend.constant((aval))
124
+ self.B[nside,m]=self.backend.constant((bval))
125
+ self.ratio_mm[nside,m]=self.backend.constant(np.sqrt(4*np.pi)*np.expand_dims(np.exp(val),1))
126
+ # Calcul de P_{mm}(x)
127
+ P_mm=np.ones([3*nside,4*nside-1])
128
+ x=np.cos(self.lth[nside])
129
+ if m == 0:
130
+ P_mm[m] = 1.0
131
+ for m in range(3*nside-1):
132
+ P_mm[m] = (0.5-m%2)*2 * (1 - x**2)**(m/2)
133
+ self.P_mm[nside]=self.backend.constant(P_mm)
134
+
135
+ def init_Ys(self,s,nside):
136
+
137
+ if (s,nside) not in self.Yp:
138
+ import quaternionic
139
+ import spherical
140
+
141
+ ell_max = 3*nside-1 # Use the largest ℓ value you expect to need
142
+ wigner = spherical.Wigner(ell_max)
143
+
144
+ #th,ph=hp.pix2ang(nside,np.arange(12*nside*nside))
145
+
146
+ lth=self.ring_th(nside)
147
+
148
+ R = quaternionic.array.from_spherical_coordinates(lth, 0*lth)
149
+ self.Yp[s,nside] = {}
150
+ self.Ym[s,nside] = {}
151
+ iplus = (wigner.sYlm( s, R)*(4*np.pi/(12*nside**2))).T.real
152
+ imoins = (wigner.sYlm(-s, R)*(4*np.pi/(12*nside**2))).T.real
153
+
154
+ for m in range(ell_max+1):
155
+ idx=np.array([wigner.Yindex(k, m) for k in range(m,ell_max+1)])
156
+ vnorm=1/np.expand_dims(np.sqrt(2*(np.arange(ell_max-m+1)+m)+1),1)
157
+ self.Yp[s,nside][m] = iplus[idx]*vnorm
158
+ self.Ym[s,nside][m] = imoins[idx]*vnorm
159
+
160
+ del(iplus)
161
+ del(imoins)
162
+ del(wigner)
163
+
164
+ def log(self,v):
165
+ return np.log(v)
166
+ if isinstance(v,np.ndarray):
167
+ return np.array([self.backend.bk_log(self.backend.bk_cast(k)) for k in v])
168
+ if v<self.maxlog:
169
+ return self._logtab[v]
170
+ else:
171
+ self._logtab[v]=self.backend.bk_log(self.backend.bk_cast(v))
172
+ return self._logtab[v]
173
+
174
+ # Fonction pour calculer la double factorielle
175
+ def double_factorial_log(self,n):
176
+ if n <= 0:
177
+ return 0.0
178
+ result = 0.0
179
+ for i in range(n, 0, -2):
180
+ result += np.log(i)
181
+ return result
182
+
183
+ def recurrence_fn(self,states, inputs):
184
+ """
185
+ Fonction de récurrence pour tf.scan.
186
+ states: un tuple (U_{n-1}, U_{n-2}) de forme [m]
187
+ inputs: un tuple (a_n(x), b_n) où a_n(x) est de forme [m]
188
+ """
189
+ U_prev, U_prev2 = states
190
+ a_n, b_n = inputs # a_n est de forme [m], b_n est un scalaire
191
+ U_n = a_n * U_prev - b_n * U_prev2
192
+ return (U_n, U_prev) # Avancer les états
193
+ # Calcul des P_{lm}(x) pour tout l inclus dans [m,lmax]
194
+ def compute_legendre_m(self,x,m,lmax,nside):
195
+ result=np.zeros([lmax-m+1,x.shape[0]])
196
+ ratio=np.zeros([lmax-m+1,1])
197
+
198
+ ratio[0,0] = self.double_factorial_log(2*m - 1)-0.5*np.sum(self.log(1+np.arange(2*m)))
199
+
200
+ # Étape 1 : Calcul de P_{mm}(x)
201
+ if m == 0:
202
+ Pmm = 1.0
203
+ else:
204
+ #Pmm = (-1)**m * (1 - x**2)**(m/2)
205
+ Pmm = (0.5-m%2)*2 * (1 - x**2)**(m/2)
206
+
207
+
208
+ # Si l == m, c'est directement P_{mm}
209
+ result[0] = Pmm
210
+
211
+ if m == lmax:
212
+ return result*np.exp(ratio)*np.sqrt(4*np.pi)
213
+
214
+ # Étape 2 : Calcul de P_{l+1, m}(x)
215
+ result[1] = x * (2*m + 1) * result[0]
216
+
217
+ ratio[1,0] = ratio[0,0]-0.5*self.log(2*m+1)
218
+
219
+ # Étape 3 : Récurence pour l > m + 1
220
+ for l in range(m + 2, lmax+1):
221
+ result[l-m] = ((2*l - 1) * x * result[l-m-1] - (l + m - 1) * result[l-m-2]) / (l - m)
222
+ ratio[l-m,0] = 0.5*self.log(l-m)-0.5*self.log(l+m)+ratio[l-m-1,0]
223
+ if np.max(abs(result[l-m]))>self._limit_range:
224
+ result[l-m-1]*= self._limit_range
225
+ result[l-m]*= self._limit_range
226
+ ratio[l-m-1,0]+= self._log_limit_range
227
+ ratio[l-m,0]+= self._log_limit_range
228
+
229
+ return result*np.exp(ratio)*np.sqrt(4*np.pi)
230
+
231
+ # Calcul des P_{lm}(x) pour tout l inclus dans [m,lmax]
232
+ def compute_legendre_m_old2(self,x,m,lmax,nside):
233
+
234
+ result={}
235
+
236
+ # Si l == m, c'est directement P_{mm}
237
+ result[0] = self.P_mm[nside][m]
238
+
239
+ if m == lmax:
240
+ v=self.backend.bk_reshape(result[0]*self.ratio_mm[nside,m][0],[1,4*nside-1])
241
+ return self.backend.bk_complex(v,0*v)
242
+
243
+ # Étape 2 : Calcul de P_{l+1, m}(x)
244
+ result[1] = x * self.A[nside,m][1] * result[0]
245
+
246
+ # Étape 3 : Récurence pour l > m + 1
247
+ for l in range(m + 2, lmax+1):
248
+ result[l-m] = self.A[nside,m][l-m] * x * result[l-m-1] - self.B[nside,m][l-m] * result[l-m-2]
249
+ """
250
+ if np.max(abs(result[l-m]))>self._limit_range:
251
+ result[l-m-1]*= self._limit_range
252
+ result[l-m]*= self._limit_range
253
+ ratio[l-m-1]+= self._log_limit_range
254
+ ratio[l-m]+= self._log_limit_range
255
+ """
256
+ result=self.backend.bk_reshape(self.backend.bk_concat([result[k] for k in range(lmax+1-m)],axis=0),[lmax+1-m,4*nside-1])
257
+
258
+ return self.backend.bk_complex(result*self.ratio_mm[nside,m],0*result)
259
+
260
+
261
+ def compute_legendre_m_old(self,x,m,lmax,nside):
262
+
263
+ import tensorflow as tf
264
+ result={}
265
+
266
+ # Si l == m, c'est directement P_{mm}
267
+ U_0 = self.P_mm[nside][m]
268
+
269
+ if m == lmax:
270
+ v=self.backend.bk_reshape(U_0*self.ratio_mm[nside,m][0],[1,4*nside-1])
271
+ return self.backend.bk_complex(v,0*v)
272
+
273
+ # Étape 2 : Calcul de P_{l+1, m}(x)
274
+ U_1 = x * self.A[nside,m][1] * U_0
275
+ if m == lmax-1:
276
+ result = tf.concat([self.backend.bk_expand_dims(U_0,0),
277
+ self.backend.bk_expand_dims(U_1,0)],0)
278
+ return self.backend.bk_complex(result*self.ratio_mm[nside,m],0*result)
279
+
280
+ a_values = self.backend.bk_expand_dims(self.A[nside,m],1)*self.backend.bk_expand_dims(x,0)
281
+ # Initialiser les états avec (U_1, U_0) pour chaque m
282
+ initial_states = (U_1, U_0)
283
+ inputs = (a_values[2:], self.B[nside,m][2:])
284
+ # Appliquer tf.scan
285
+ result = tf.scan(self.recurrence_fn, inputs, initializer=initial_states)
286
+ # Le premier élément de result contient les U[n]
287
+ result = tf.concat([self.backend.bk_expand_dims(U_0,0),
288
+ self.backend.bk_expand_dims(U_1,0),
289
+ result[0]], axis=0)
290
+ """
291
+ # Étape 3 : Récurence pour l > m + 1
292
+ for l in range(m + 2, lmax+1):
293
+ result[l-m] = self.A[nside,m][l-m] * x * result[l-m-1] - self.B[nside,m][l-m] * result[l-m-2]
294
+
295
+ if np.max(abs(result[l-m]))>self._limit_range:
296
+ result[l-m-1]*= self._limit_range
297
+ result[l-m]*= self._limit_range
298
+ ratio[l-m-1]+= self._log_limit_range
299
+ ratio[l-m]+= self._log_limit_range
300
+ result=self.backend.bk_reshape(self.backend.bk_concat([result[k] for k in range(lmax+1-m)],axis=0),[lmax+1-m,4*nside-1])
301
+ """
302
+
303
+ return self.backend.bk_complex(result*self.ratio_mm[nside,m],0*result)
304
+
305
+
306
+ # Calcul des s_P_{lm}(x) pour tout l inclus dans [m,lmax]
307
+ def compute_legendre_spin2_m(self,co_th,si_th,m,lmax):
308
+ result=np.zeros([lmax-m+2,co_th.shape[0]])
309
+ ratio =np.zeros([lmax-m+2,1])
310
+
311
+ ratio[1,0] = self.double_factorial_log(2*m - 1)-0.5*np.sum(self.log(1+np.arange(2*m)))
312
+ # Étape 1 : Calcul de P_{mm}(x)
313
+ if m == 0:
314
+ Pmm = 1.0
315
+ else:
316
+ #Pmm = (-1)**m * (1 - x**2)**(m/2)
317
+ Pmm = (0.5-m%2)*2 * (1 - co_th**2)**(m/2)
318
+
319
+
320
+ # Si l == m, c'est directement P_{mm}
321
+ result[1] = Pmm
322
+
323
+ if m == lmax:
324
+ ylm=result*np.exp(ratio)
325
+ ylm[1:]*=(np.sqrt(4*np.pi*(2*(np.arange(lmax-m+1)+m)+1))).reshape(lmax+1-m,1)
326
+
327
+ else:
328
+ # Étape 2 : Calcul de P_{l+1, m}(x)
329
+ result[2] = co_th * (2*m + 1) * result[0]
330
+
331
+ ratio[2,0] = ratio[1,0]-self.log(2*m+1)/2
332
+
333
+ # Étape 3 : Récurence pour l > m + 1
334
+ for l in range(m + 2, lmax+1):
335
+ result[l-m+1] = ((2*l - 1) * co_th * result[l-m] - (l + m - 1) * result[l-m-1]) / (l - m)
336
+ ratio[l-m+1,0] = (self.log(l-m)-self.log(l+m))/2+ratio[l-m,0]
337
+ if np.max(abs(result[l-m+1]))>self._limit_range:
338
+ result[l-m]*= self._limit_range
339
+ result[l-m+1]*= self._limit_range
340
+ ratio[l-m,0]+= self._log_limit_range
341
+ ratio[l-m+1,0]+= self._log_limit_range
342
+
343
+ ylm=result*np.exp(ratio)
344
+ ylm[1:]*=(np.sqrt(4*np.pi*(2*(np.arange(lmax-m+1)+m)+1))).reshape(lmax+1-m,1)
345
+
346
+ ell=(np.arange(lmax+1-m)+m).reshape(lmax+1-m,1)
347
+
348
+ cot_th=co_th/si_th
349
+ si2_th=si_th*si_th
350
+
351
+ a = (2*m**2-ell*(ell+1))/(si2_th.reshape(1,si2_th.shape[0]))+ell*(ell-1)*cot_th*cot_th
352
+ b = 2*m*(ell-1)*cot_th/si_th
353
+ w=np.zeros([lmax+1-m,1])
354
+ l=ell[ell>1]
355
+ w[ell>1]=np.sqrt(1/((l+2)*(l+1)*(l)*(l-1)))
356
+ w=w.reshape(lmax+1-m,1)
357
+
358
+ alpha_plus=w*(a+b)
359
+ alpha_moins=w*(a-b)
360
+
361
+ a=2*np.sqrt((2*ell+1)/(2*ell-1)*(ell*ell-m*m))
362
+ b=m/si2_th
363
+
364
+ beta_plus=w*a*(cot_th/si_th+b)
365
+ beta_moins=w*a*(cot_th/si_th-b)
366
+
367
+ ylm_plus = alpha_plus*ylm[1:]+ beta_plus*ylm[:-1]
368
+ ylm_moins = alpha_moins*ylm[1:] + beta_moins*ylm[:-1]
369
+
370
+ return ylm_plus,ylm_moins
371
+
372
+ def rfft2fft(self,val,axis=0):
373
+ r=self.backend.bk_rfft(val)
374
+ if axis==0:
375
+ r_inv=self.backend.bk_reverse(self.backend.bk_conjugate(r[1:-1]),axis=axis)
376
+ else:
377
+ r_inv=self.backend.bk_reverse(self.backend.bk_conjugate(r[:,1:-1]),axis=axis)
378
+ return self.backend.bk_concat([r,r_inv],axis=axis)
379
+
380
+ def comp_tf(self,im,nside,realfft=False):
381
+
382
+ n=0
383
+
384
+ ft_im=[]
385
+ for k in range(nside-1):
386
+ N=4*(k+1)
387
+
388
+ if realfft:
389
+ tmp=self.rfft2fft(im[n:n+N])
390
+ else:
391
+ tmp=self.backend.bk_fft(im[n:n+N])
392
+
393
+ l_n=tmp.shape[0]
394
+
395
+ if l_n<3*nside+1:
396
+ repeat_n=3*nside//l_n+1
397
+ tmp=self.backend.bk_tile(tmp,repeat_n,axis=0)
398
+
399
+ ft_im.append(tmp[0:3*nside])
400
+
401
+ n+=N
402
+ if nside>1:
403
+ result=self.backend.bk_reshape(self.backend.bk_concat(ft_im,axis=0),[nside-1,3*nside])
404
+
405
+ N=4*nside*(2*nside+1)
406
+ v=self.backend.bk_reshape(im[n:n+N],[2*nside+1,4*nside])
407
+ if realfft:
408
+ v_fft=self.rfft2fft(v,axis=1)[:,:3*nside]
409
+ else:
410
+ v_fft=self.backend.bk_fft(v)[:,:3*nside]
411
+
412
+ n+=N
413
+ if nside>1:
414
+ result=self.backend.bk_concat([result,v_fft],axis=0)
415
+ else:
416
+ result=v_fft
417
+
418
+ if nside>1:
419
+ ft_im=[]
420
+ for k in range(nside-1):
421
+ N=4*(nside-1-k)
422
+
423
+ if realfft:
424
+ tmp=self.rfft2fft(im[n:n+N])[0:l_n]
425
+ else:
426
+ tmp=self.backend.bk_fft(im[n:n+N])[0:l_n]
427
+
428
+ l_n=tmp.shape[0]
429
+
430
+ if l_n<3*nside+1:
431
+ repeat_n=3*nside//l_n+1
432
+ tmp=self.backend.bk_tile(tmp,repeat_n,axis=0)
433
+
434
+ ft_im.append(tmp[0:3*nside])
435
+ n+=N
436
+
437
+ lastresult=self.backend.bk_reshape(self.backend.bk_concat(ft_im,axis=0),[nside-1,3*nside])
438
+ return self.backend.bk_concat([result,lastresult],axis=0)*self.matrix_shift_ph[nside]
439
+ else:
440
+ return result*self.matrix_shift_ph[nside]
441
+
442
+ def anafast(self,im,map2=None,nest=False,spin=2):
443
+
444
+ """The `anafast` function computes the L1 and L2 norm power spectra.
445
+
446
+ Currently, it is not optimized for single-pass computation due to the relatively inefficient computation of \(Y_{lm}\).
447
+ Nonetheless, it utilizes TensorFlow and can be integrated into gradient computations.
448
+
449
+ Input:
450
+ - `im`: a vector of size \([12 \times \text{Nside}^2]\) for scalar data, or of size \([2, 12 \times \text{Nside}^2]\) for Q,U polar data,
451
+ or of size \([3, 12 \times \text{Nside}^2]\) for I,Q,U polar data.
452
+ - `map2` (optional): a vector of size \([12 \times \text{Nside}^2]\) for scalar data, or of size
453
+ \([3, 12 \times \text{Nside}^2]\) for polar data. If provided, cross power spectra will be computed.
454
+ - `nest=True`: alters the ordering of the input maps.
455
+ - `spin=2` for 1/2 spin data as Q and U. Spin=1 for seep fields
456
+
457
+ Output:
458
+ -A tensor of size \([l_{\text{max}} \times (l_{\text{max}}-1)]\) formatted as \([6, \ldots]\),
459
+ ordered as TT, EE, BB, TE, EB.TBanafast function computes L1 and L2 norm powerspctra.
460
+
461
+ """
462
+ i_im=self.backend.bk_cast(im)
463
+ if map2 is not None:
464
+ i_map2=self.backend.bk_cast(map2)
465
+
466
+ doT=True
467
+ if len(i_im.shape)==1: # nopol
468
+ nside=int(np.sqrt(i_im.shape[0]//12))
469
+ else:
470
+ if i_im.shape[0]==2:
471
+ doT=False
472
+ nside=int(np.sqrt(i_im.shape[1]//12))
473
+
474
+ self.shift_ph(nside)
475
+
476
+ if doT: # nopol
477
+ if len(i_im.shape)==2: # pol
478
+ l_im=i_im[0]
479
+ if map2 is not None:
480
+ l_map2=i_map2[0]
481
+ else:
482
+ l_im=i_im
483
+ if map2 is not None:
484
+ l_map2=i_map2
485
+
486
+ if nest:
487
+ idx=hp.ring2nest(nside,np.arange(12*nside**2))
488
+ if len(i_im.shape)==1: # nopol
489
+ ft_im=self.comp_tf(self.backend.bk_gather(l_im,idx),nside,realfft=True)
490
+ if map2 is not None:
491
+ ft_im2=self.comp_tf(self.backend.bk_gather(l_map2,idx),nside,realfft=True)
492
+ else:
493
+ ft_im=self.comp_tf(l_im,nside,realfft=True)
494
+ if map2 is not None:
495
+ ft_im2=self.comp_tf(l_map2,nside,realfft=True)
496
+
497
+ lth=self.ring_th(nside)
498
+
499
+ co_th=np.cos(lth)
500
+
501
+ lmax=3*nside-1
502
+
503
+ cl2=None
504
+ cl2_L1=None
505
+ dt2=0
506
+ dt3=0
507
+ dt4=0
508
+ if len(i_im.shape)==2: # nopol
509
+
510
+ self.init_Ys(spin,nside)
511
+
512
+ if nest:
513
+ idx=hp.ring2nest(nside,np.arange(12*nside**2))
514
+ l_Q=self.backend.bk_gather(i_im[int(doT)],idx)
515
+ l_U=self.backend.bk_gather(i_im[1+int(doT)],idx)
516
+ ft_im_Pp=self.comp_tf(self.backend.bk_complex(l_Q,l_U),nside)
517
+ ft_im_Pm=self.comp_tf(self.backend.bk_complex(l_Q,-l_U),nside)
518
+ if map2 is not None:
519
+ l_Q=self.backend.bk_gather(i_map2[int(doT)],idx)
520
+ l_U=self.backend.bk_gather(i_map2[1+int(doT)],idx)
521
+ ft_im2_Pp=self.comp_tf(self.backend.bk_complex(l_Q,l_U),nside)
522
+ ft_im2_Pm=self.comp_tf(self.backend.bk_complex(l_Q,-l_U),nside)
523
+ else:
524
+ ft_im_Pp=self.comp_tf(self.backend.bk_complex(i_im[int(doT)],i_im[1+int(doT)]),nside)
525
+ ft_im_Pm=self.comp_tf(self.backend.bk_complex(i_im[int(doT)],-i_im[1+int(doT)]),nside)
526
+ if map2 is not None:
527
+ ft_im2_Pp=self.comp_tf(self.backend.bk_complex(i_map2[int(doT)],i_map2[1+int(doT)]),nside)
528
+ ft_im2_Pm=self.comp_tf(self.backend.bk_complex(i_map2[int(doT)],-i_map2[1+int(doT)]),nside)
529
+
530
+ for m in range(lmax+1):
531
+
532
+ plm=self.compute_legendre_m(co_th,m,3*nside-1,nside)/(12*nside**2)
533
+
534
+ if doT:
535
+ tmp=self.backend.bk_reduce_sum(plm*ft_im[:,m],1)
536
+
537
+ if map2 is not None:
538
+ tmp2=self.backend.bk_reduce_sum(plm*ft_im2[:,m],1)
539
+ else:
540
+ tmp2=tmp
541
+
542
+ if len(i_im.shape)==2: # pol
543
+ plmp=self.Yp[spin,nside][m]
544
+ plmm=self.Ym[spin,nside][m]
545
+
546
+ tmpp=self.backend.bk_reduce_sum(plmp*ft_im_Pp[:,m],1)
547
+ tmpm=self.backend.bk_reduce_sum(plmm*ft_im_Pm[:,m],1)
548
+
549
+ almE=-(tmpp+tmpm)/2.0
550
+ almB=(tmpp-tmpm)/(2J)
551
+
552
+ if map2 is not None:
553
+ tmpp2=self.backend.bk_reduce_sum(plmp*ft_im2_Pp[:,m],1)
554
+ tmpm2=self.backend.bk_reduce_sum(plmm*ft_im2_Pm[:,m],1)
555
+
556
+ almE2=-(tmpp2+tmpm2)/2.0
557
+ almB2=(tmpp2-tmpm2)/(2J)
558
+ else:
559
+ almE2=almE
560
+ almB2=almB
561
+
562
+ if doT:
563
+ tmpTT=self.backend.bk_real((tmp*self.backend.bk_conjugate(tmp2)))
564
+ tmpTE=self.backend.bk_real((tmp*self.backend.bk_conjugate(almE2)))
565
+ tmpTB=-self.backend.bk_real((tmp*self.backend.bk_conjugate(almB2)))
566
+
567
+ tmpEE=self.backend.bk_real((almE*self.backend.bk_conjugate(almE2)))
568
+ tmpBB=self.backend.bk_real((almB*self.backend.bk_conjugate(almB2)))
569
+ tmpEB=-self.backend.bk_real((almE*self.backend.bk_conjugate(almB2)))
570
+
571
+ if map2 is not None:
572
+ tmpEB=(tmpEB-self.backend.bk_real((almE2*self.backend.bk_conjugate(almB))))/2
573
+
574
+ if doT:
575
+ tmpTE=(tmpTE+self.backend.bk_real((tmp2*self.backend.bk_conjugate(almE))))/2
576
+ tmpTB=(tmpTB-self.backend.bk_real((tmp2*self.backend.bk_conjugate(almB))))/2
577
+
578
+
579
+ if m==0:
580
+ if doT:
581
+ l_cl=self.backend.bk_concat([tmpTT,tmpEE,tmpBB,tmpTE,tmpEB,tmpTB],0)
582
+ else:
583
+ l_cl=self.backend.bk_concat([tmpEE,tmpBB,tmpEB],0)
584
+ else:
585
+ offset_tensor=self.backend.bk_zeros((m),dtype=self.backend.all_bk_type)
586
+ if doT:
587
+ l_cl=self.backend.bk_concat([self.backend.bk_concat([offset_tensor,tmpTT],axis=0),
588
+ self.backend.bk_concat([offset_tensor,tmpEE],axis=0),
589
+ self.backend.bk_concat([offset_tensor,tmpBB],axis=0),
590
+ self.backend.bk_concat([offset_tensor,tmpTE],axis=0),
591
+ self.backend.bk_concat([offset_tensor,tmpEB],axis=0),
592
+ self.backend.bk_concat([offset_tensor,tmpTB],axis=0)],axis=0)
593
+ else:
594
+ l_cl=self.backend.bk_concat([self.backend.bk_concat([offset_tensor,tmpEE],axis=0),
595
+ self.backend.bk_concat([offset_tensor,tmpBB],axis=0),
596
+ self.backend.bk_concat([offset_tensor,tmpEB],axis=0)],axis=0)
597
+
598
+ if doT:
599
+ l_cl=self.backend.bk_reshape(l_cl,[6,lmax+1])
600
+ else:
601
+ l_cl=self.backend.bk_reshape(l_cl,[3,lmax+1])
602
+ else:
603
+ tmp=self.backend.bk_real((tmp*self.backend.bk_conjugate(tmp2)))
604
+ if m==0:
605
+ l_cl=tmp
606
+ else:
607
+ offset_tensor=self.backend.bk_zeros((m),dtype=self.backend.all_bk_type)
608
+ l_cl=self.backend.bk_concat([offset_tensor,tmp],axis=0)
609
+
610
+ if cl2 is None:
611
+ cl2=l_cl
612
+ else:
613
+ cl2+=2*l_cl
614
+
615
+ #cl2=cl2*(4*np.pi) #self.backend.bk_sqrt(self.backend.bk_cast(4*np.pi)) #(2*np.arange(cl2.shape[0])+1)))
616
+
617
+ cl2_l1=self.backend.bk_L1(cl2)
618
+
619
+ return cl2,cl2_l1
620
+
621
+ def map2alm(self,im,nest=False):
622
+ nside=int(np.sqrt(im.shape[0]//12))
623
+
624
+ if nest:
625
+ idx=hp.ring2nest(nside,np.arange(12*nside**2))
626
+ ft_im=self.comp_tf(self.backend.bk_cast(self.backend.bk_gather(im,idx)),nside,realfft=True)
627
+ else:
628
+ ft_im=self.comp_tf(self.backend.bk_cast(im),nside,realfft=True)
629
+
630
+ lth=self.ring_th(nside)
631
+
632
+ co_th=np.cos(lth)
633
+
634
+ lmax=3*nside-1
635
+
636
+ alm=None
637
+ for m in range(lmax+1):
638
+ plm=self.compute_legendre_m(co_th,m,3*nside-1,nside)/(12*nside**2)
639
+
640
+ tmp=self.backend.bk_reduce_sum(plm*ft_im[:,m],1)
641
+ if m==0:
642
+ alm=tmp
643
+ else:
644
+ alm=self.backend.bk_concat([alm,tmp],axis=0)
645
+
646
+ return alm
647
+
648
+ def map2alm_spin(self,im_Q,im_U,spin=2,nest=False):
649
+
650
+ if spin==0:
651
+ return self.map2alm(im_Q,nest=nest),self.map2alm(im_U,nest=nest)
652
+
653
+ nside=int(np.sqrt(im_Q.shape[0]//12))
654
+
655
+ lth=self.ring_th(nside)
656
+
657
+ co_th=np.cos(lth)
658
+
659
+ if nest:
660
+ idx=hp.ring2nest(nside,np.arange(12*nside**2))
661
+ l_Q=self.backend.bk_gather(im_Q,idx)
662
+ l_U=self.backend.bk_gather(im_U,idx)
663
+ ft_im_1=self.comp_tf(self.backend.bk_complex(l_Q,l_U),nside)
664
+ ft_im_2=self.comp_tf(self.backend.bk_complex(l_Q,-l_U),nside)
665
+ else:
666
+ ft_im_1=self.comp_tf(self.backend.bk_complex(im_Q,im_U),nside)
667
+ ft_im_2=self.comp_tf(self.backend.bk_complex(im_Q,-im_U),nside)
668
+
669
+ lmax=3*nside-1
670
+
671
+ alm=None
672
+ for m in range(lmax+1):
673
+ #not yet debug use spherical
674
+ #plmp1,plmm1=self.compute_legendre_spin2_m(co_th,si_th,m,3*nside-1)
675
+ #plmp1/=(12*nside**2)
676
+ #plmm1/=(12*nside**2)
677
+
678
+ plmp=self.Yp[spin,nside][m]
679
+ plmm=self.Ym[spin,nside][m]
680
+
681
+ tmpp=self.backend.bk_reduce_sum(plmp*ft_im_1[:,m],1)
682
+ tmpm=self.backend.bk_reduce_sum(plmm*ft_im_2[:,m],1)
683
+ if m==0:
684
+ almE=-(tmpp+tmpm)/2.0
685
+ almB=(tmpp-tmpm)/(2J)
686
+ else:
687
+ almE=self.backend.bk_concat([almE,-(tmpp+tmpm)/2],axis=0)
688
+ almB=self.backend.bk_concat([almB,(tmpp-tmpm)/(2J)],axis=0)
689
+
690
+ return almE,almB
foscat/alm_tools.py ADDED
@@ -0,0 +1,11 @@
1
+ import numpy as np
2
+
3
+ #====================================================================================================================
4
+ # This class is an automatic traduction of the fortran healpix software
5
+ #====================================================================================================================
6
+
7
+
8
+ class alm_tools():
9
+ def __init__(self):
10
+ pass
11
+