foscat 3.6.1__py3-none-any.whl → 3.7.1__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.
Files changed (16) hide show
  1. foscat/CircSpline.py +19 -22
  2. foscat/FoCUS.py +61 -51
  3. foscat/Spline1D.py +12 -13
  4. foscat/Synthesis.py +16 -13
  5. foscat/alm.py +718 -580
  6. foscat/backend.py +141 -37
  7. foscat/scat_cov.py +2311 -430
  8. foscat/scat_cov2D.py +102 -2
  9. foscat/scat_cov_map.py +15 -2
  10. {foscat-3.6.1.dist-info → foscat-3.7.1.dist-info}/METADATA +3 -3
  11. foscat-3.7.1.dist-info/RECORD +26 -0
  12. {foscat-3.6.1.dist-info → foscat-3.7.1.dist-info}/WHEEL +1 -1
  13. foscat/alm_tools.py +0 -11
  14. foscat-3.6.1.dist-info/RECORD +0 -27
  15. /foscat-3.6.1.dist-info/LICENCE → /foscat-3.7.1.dist-info/LICENSE +0 -0
  16. {foscat-3.6.1.dist-info → foscat-3.7.1.dist-info}/top_level.txt +0 -0
foscat/alm.py CHANGED
@@ -1,178 +1,191 @@
1
+ import time
2
+
1
3
  import healpy as hp
2
4
  import numpy as np
3
- import time
4
5
 
5
- class alm():
6
6
 
7
- def __init__(self,backend=None,lmax=24,
8
- nside=None,limit_range=1E10):
9
-
7
+ class alm:
8
+
9
+ def __init__(self, backend=None, lmax=24, nside=None, limit_range=1e10):
10
+
10
11
  if backend is None:
11
12
  import foscat.scat_cov as sc
12
- self.sc=sc.funct()
13
- self.backend=self.sc.backend
13
+
14
+ self.sc = sc.funct()
15
+ self.backend = self.sc.backend
14
16
  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={}
17
+ self.backend = backend.backend
18
+
19
+ self._logtab = {}
20
+ self.lth = {}
21
+ self.lph = {}
22
+ self.matrix_shift_ph = {}
23
+ self.ratio_mm = {}
24
+ self.P_mm = {}
25
+ self.A = {}
26
+ self.B = {}
25
27
  if nside is not None:
26
- self.maxlog=6*nside+1
27
- self.lmax=3*nside
28
+ self.maxlog = 6 * nside + 1
29
+ self.lmax = 3 * nside
28
30
  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
31
+ self.lmax = lmax
32
+ self.maxlog = 2 * lmax + 1
33
+
34
+ for k in range(1, self.maxlog):
35
+ self._logtab[k] = self.backend.bk_log(self.backend.bk_cast(k))
36
+ self._logtab[0] = 0.0
35
37
 
36
38
  if nside is not None:
37
39
  self.ring_th(nside)
38
40
  self.ring_ph(nside)
39
41
  self.shift_ph(nside)
40
-
41
- self._limit_range=1/limit_range
42
- self._log_limit_range=np.log(limit_range)
43
-
44
42
 
45
- self.Yp={}
46
- self.Ym={}
43
+ self._limit_range = 1 / limit_range
44
+ self._log_limit_range = np.log(limit_range)
47
45
 
48
- def ring_th(self,nside):
46
+ self.Yp = {}
47
+ self.Ym = {}
48
+
49
+ def ring_th(self, nside):
49
50
  if nside not in self.lth:
50
- n=0
51
- ith=[]
52
- for k in range(nside-1):
53
- N=4*(k+1)
51
+ n = 0
52
+ ith = []
53
+ for k in range(nside - 1):
54
+ N = 4 * (k + 1)
54
55
  ith.append(n)
55
- n+=N
56
-
57
- for k in range(2*nside+1):
58
- N=4*nside
56
+ n += N
57
+
58
+ for k in range(2 * nside + 1):
59
+ N = 4 * nside
59
60
  ith.append(n)
60
- n+=N
61
- for k in range(nside-1):
62
- N=4*(nside-1-k)
61
+ n += N
62
+ for k in range(nside - 1):
63
+ N = 4 * (nside - 1 - k)
63
64
  ith.append(n)
64
- n+=N
65
-
66
- th,ph=hp.pix2ang(nside,ith)
67
-
68
- self.lth[nside]=th
65
+ n += N
66
+
67
+ th, ph = hp.pix2ang(nside, ith)
68
+
69
+ self.lth[nside] = th
69
70
  return self.lth[nside]
70
-
71
- def ring_ph(self,nside):
71
+
72
+ def ring_ph(self, nside):
72
73
  if nside not in self.lph:
73
- n=0
74
- iph=[]
75
- for k in range(nside-1):
76
- N=4*(k+1)
74
+ n = 0
75
+ iph = []
76
+ for k in range(nside - 1):
77
+ N = 4 * (k + 1)
77
78
  iph.append(n)
78
- n+=N
79
-
80
- for k in range(2*nside+1):
81
- N=4*nside
79
+ n += N
80
+
81
+ for k in range(2 * nside + 1):
82
+ N = 4 * nside
82
83
  iph.append(n)
83
- n+=N
84
- for k in range(nside-1):
85
- N=4*(nside-1-k)
84
+ n += N
85
+ for k in range(nside - 1):
86
+ N = 4 * (nside - 1 - k)
86
87
  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
-
88
+ n += N
89
+
90
+ th, ph = hp.pix2ang(nside, iph)
91
+
92
+ self.lph[nside] = ph
93
+
94
+ def shift_ph(self, nside):
95
+
95
96
  if nside not in self.matrix_shift_ph:
96
97
  self.ring_th(nside)
97
98
  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)))
99
+ x = (-1j * np.arange(3 * nside)).reshape(1, 3 * nside)
100
+ self.matrix_shift_ph[nside] = self.backend.bk_cast(
101
+ self.backend.bk_exp(x * self.lph[nside].reshape(4 * nside - 1, 1))
102
+ )
103
+
104
+ self.lmax = 3 * nside - 1
105
+
106
+ ratio_mm = {}
107
+
108
+ for m in range(3 * nside):
109
+ val = np.zeros([self.lmax - m + 1])
110
+ aval = np.zeros([self.lmax - m + 1])
111
+ bval = np.zeros([self.lmax - m + 1])
112
+
113
+ if m > 0:
114
+ val[0] = self.double_factorial_log(2 * m - 1) - 0.5 * np.sum(
115
+ np.log(1 + np.arange(2 * m))
116
+ )
112
117
  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))
118
+ val[0] = self.double_factorial_log(2 * m - 1)
119
+ if m < self.lmax:
120
+ aval[1] = 2 * m + 1
121
+ val[1] = val[0] - 0.5 * self.log(2 * m + 1)
122
+
123
+ for l in range(m + 2, self.lmax + 1):
124
+ aval[l - m] = (2 * l - 1) / (l - m)
125
+ bval[l - m] = (l + m - 1) / (l - m)
126
+ val[l - m] = (
127
+ val[l - m - 1]
128
+ + 0.5 * self.log(l - m)
129
+ - 0.5 * self.log(l + m)
130
+ )
131
+
132
+ self.A[nside, m] = self.backend.constant((aval))
133
+ self.B[nside, m] = self.backend.constant((bval))
134
+ self.ratio_mm[nside, m] = self.backend.constant(
135
+ np.sqrt(4 * np.pi) * np.expand_dims(np.exp(val), 1)
136
+ )
126
137
  # Calcul de P_{mm}(x)
127
- P_mm=np.ones([3*nside,4*nside-1])
128
- x=np.cos(self.lth[nside])
138
+ P_mm = np.ones([3 * nside, 4 * nside - 1])
139
+ x = np.cos(self.lth[nside])
129
140
  if m == 0:
130
141
  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):
142
+ for m in range(3 * nside - 1):
143
+ P_mm[m] = (0.5 - m % 2) * 2 * (1 - x**2) ** (m / 2)
144
+ self.P_mm[nside] = self.backend.constant(P_mm)
145
+
146
+ def init_Ys(self, s, nside):
136
147
 
137
- if (s,nside) not in self.Yp:
148
+ if (s, nside) not in self.Yp:
138
149
  import quaternionic
139
150
  import spherical
140
151
 
141
- ell_max = 3*nside-1 # Use the largest ℓ value you expect to need
152
+ ell_max = 3 * nside - 1 # Use the largest ℓ value you expect to need
142
153
  wigner = spherical.Wigner(ell_max)
143
154
 
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):
155
+ # th,ph=hp.pix2ang(nside,np.arange(12*nside*nside))
156
+
157
+ lth = self.ring_th(nside)
158
+
159
+ R = quaternionic.array.from_spherical_coordinates(lth, 0 * lth)
160
+ self.Yp[s, nside] = {}
161
+ self.Ym[s, nside] = {}
162
+ iplus = (wigner.sYlm(s, R) * (4 * np.pi / (12 * nside**2))).T.real
163
+ imoins = (wigner.sYlm(-s, R) * (4 * np.pi / (12 * nside**2))).T.real
164
+
165
+ for m in range(ell_max + 1):
166
+ idx = np.array([wigner.Yindex(k, m) for k in range(m, ell_max + 1)])
167
+ vnorm = 1 / np.expand_dims(
168
+ np.sqrt(2 * (np.arange(ell_max - m + 1) + m) + 1), 1
169
+ )
170
+ self.Yp[s, nside][m] = iplus[idx] * vnorm
171
+ self.Ym[s, nside][m] = imoins[idx] * vnorm
172
+
173
+ del iplus
174
+ del imoins
175
+ del wigner
176
+
177
+ def log(self, v):
165
178
  return np.log(v)
166
- if isinstance(v,np.ndarray):
179
+ if isinstance(v, np.ndarray):
167
180
  return np.array([self.backend.bk_log(self.backend.bk_cast(k)) for k in v])
168
- if v<self.maxlog:
181
+ if v < self.maxlog:
169
182
  return self._logtab[v]
170
183
  else:
171
- self._logtab[v]=self.backend.bk_log(self.backend.bk_cast(v))
184
+ self._logtab[v] = self.backend.bk_log(self.backend.bk_cast(v))
172
185
  return self._logtab[v]
173
186
 
174
187
  # Fonction pour calculer la double factorielle
175
- def double_factorial_log(self,n):
188
+ def double_factorial_log(self, n):
176
189
  if n <= 0:
177
190
  return 0.0
178
191
  result = 0.0
@@ -180,7 +193,7 @@ class alm():
180
193
  result += np.log(i)
181
194
  return result
182
195
 
183
- def recurrence_fn(self,states, inputs):
196
+ def recurrence_fn(self, states, inputs):
184
197
  """
185
198
  Fonction de récurrence pour tf.scan.
186
199
  states: un tuple (U_{n-1}, U_{n-2}) de forme [m]
@@ -190,62 +203,73 @@ class alm():
190
203
  a_n, b_n = inputs # a_n est de forme [m], b_n est un scalaire
191
204
  U_n = a_n * U_prev - b_n * U_prev2
192
205
  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
-
206
+
207
+ # Calcul des P_{lm}(x) pour tout l inclus dans [m,lmax]
208
+ def compute_legendre_m(self, x, m, lmax, nside):
209
+ result = np.zeros([lmax - m + 1, x.shape[0]])
210
+ ratio = np.zeros([lmax - m + 1, 1])
211
+
212
+ ratio[0, 0] = self.double_factorial_log(2 * m - 1) - 0.5 * np.sum(
213
+ self.log(1 + np.arange(2 * m))
214
+ )
215
+
200
216
  # Étape 1 : Calcul de P_{mm}(x)
201
217
  if m == 0:
202
218
  Pmm = 1.0
203
219
  else:
204
- #Pmm = (-1)**m * (1 - x**2)**(m/2)
205
- Pmm = (0.5-m%2)*2 * (1 - x**2)**(m/2)
206
-
207
-
220
+ # Pmm = (-1)**m * (1 - x**2)**(m/2)
221
+ Pmm = (0.5 - m % 2) * 2 * (1 - x**2) ** (m / 2)
222
+
208
223
  # Si l == m, c'est directement P_{mm}
209
- result[0] = Pmm
210
-
224
+ result[0] = Pmm
225
+
211
226
  if m == lmax:
212
- return result*np.exp(ratio)*np.sqrt(4*np.pi)
213
-
227
+ return result * np.exp(ratio) * np.sqrt(4 * np.pi)
228
+
214
229
  # Étape 2 : Calcul de P_{l+1, m}(x)
215
- result[1] = x * (2*m + 1) * result[0]
230
+ result[1] = x * (2 * m + 1) * result[0]
231
+
232
+ ratio[1, 0] = ratio[0, 0] - 0.5 * self.log(2 * m + 1)
216
233
 
217
- ratio[1,0] = ratio[0,0]-0.5*self.log(2*m+1)
218
-
219
234
  # É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)
235
+ for l in range(m + 2, lmax + 1):
236
+ result[l - m] = (
237
+ (2 * l - 1) * x * result[l - m - 1] - (l + m - 1) * result[l - m - 2]
238
+ ) / (l - m)
239
+ ratio[l - m, 0] = (
240
+ 0.5 * self.log(l - m) - 0.5 * self.log(l + m) + ratio[l - m - 1, 0]
241
+ )
242
+ if np.max(abs(result[l - m])) > self._limit_range:
243
+ result[l - m - 1] *= self._limit_range
244
+ result[l - m] *= self._limit_range
245
+ ratio[l - m - 1, 0] += self._log_limit_range
246
+ ratio[l - m, 0] += self._log_limit_range
247
+
248
+ return result * np.exp(ratio) * np.sqrt(4 * np.pi)
230
249
 
231
250
  # 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
-
251
+ def compute_legendre_m_old2(self, x, m, lmax, nside):
252
+
253
+ result = {}
254
+
236
255
  # Si l == m, c'est directement P_{mm}
237
- result[0] = self.P_mm[nside][m]
238
-
256
+ result[0] = self.P_mm[nside][m]
257
+
239
258
  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
-
259
+ v = self.backend.bk_reshape(
260
+ result[0] * self.ratio_mm[nside, m][0], [1, 4 * nside - 1]
261
+ )
262
+ return self.backend.bk_complex(v, 0 * v)
263
+
243
264
  # Étape 2 : Calcul de P_{l+1, m}(x)
244
- result[1] = x * self.A[nside,m][1] * result[0]
245
-
265
+ result[1] = x * self.A[nside, m][1] * result[0]
266
+
246
267
  # É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]
268
+ for l in range(m + 2, lmax + 1):
269
+ result[l - m] = (
270
+ self.A[nside, m][l - m] * x * result[l - m - 1]
271
+ - self.B[nside, m][l - m] * result[l - m - 2]
272
+ )
249
273
  """
250
274
  if np.max(abs(result[l-m]))>self._limit_range:
251
275
  result[l-m-1]*= self._limit_range
@@ -253,45 +277,62 @@ class alm():
253
277
  ratio[l-m-1]+= self._log_limit_range
254
278
  ratio[l-m]+= self._log_limit_range
255
279
  """
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])
280
+ result = self.backend.bk_reshape(
281
+ self.backend.bk_concat([result[k] for k in range(lmax + 1 - m)], axis=0),
282
+ [lmax + 1 - m, 4 * nside - 1],
283
+ )
257
284
 
258
- return self.backend.bk_complex(result*self.ratio_mm[nside,m],0*result)
285
+ return self.backend.bk_complex(result * self.ratio_mm[nside, m], 0 * result)
286
+
287
+ def compute_legendre_m_old(self, x, m, lmax, nside):
259
288
 
260
-
261
- def compute_legendre_m_old(self,x,m,lmax,nside):
262
-
263
289
  import tensorflow as tf
264
- result={}
265
-
290
+
291
+ result = {}
292
+
266
293
  # Si l == m, c'est directement P_{mm}
267
- U_0 = self.P_mm[nside][m]
268
-
294
+ U_0 = self.P_mm[nside][m]
295
+
269
296
  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)
297
+ v = self.backend.bk_reshape(
298
+ U_0 * self.ratio_mm[nside, m][0], [1, 4 * nside - 1]
299
+ )
300
+ return self.backend.bk_complex(v, 0 * v)
279
301
 
280
- a_values = self.backend.bk_expand_dims(self.A[nside,m],1)*self.backend.bk_expand_dims(x,0)
302
+ # Étape 2 : Calcul de P_{l+1, m}(x)
303
+ U_1 = x * self.A[nside, m][1] * U_0
304
+ if m == lmax - 1:
305
+ result = tf.concat(
306
+ [
307
+ self.backend.bk_expand_dims(U_0, 0),
308
+ self.backend.bk_expand_dims(U_1, 0),
309
+ ],
310
+ 0,
311
+ )
312
+ return self.backend.bk_complex(result * self.ratio_mm[nside, m], 0 * result)
313
+
314
+ a_values = self.backend.bk_expand_dims(
315
+ self.A[nside, m], 1
316
+ ) * self.backend.bk_expand_dims(x, 0)
281
317
  # Initialiser les états avec (U_1, U_0) pour chaque m
282
318
  initial_states = (U_1, U_0)
283
- inputs = (a_values[2:], self.B[nside,m][2:])
319
+ inputs = (a_values[2:], self.B[nside, m][2:])
284
320
  # Appliquer tf.scan
285
321
  result = tf.scan(self.recurrence_fn, inputs, initializer=initial_states)
286
322
  # 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)
323
+ result = tf.concat(
324
+ [
325
+ self.backend.bk_expand_dims(U_0, 0),
326
+ self.backend.bk_expand_dims(U_1, 0),
327
+ result[0],
328
+ ],
329
+ axis=0,
330
+ )
290
331
  """
291
332
  # Étape 3 : Récurence pour l > m + 1
292
333
  for l in range(m + 2, lmax+1):
293
334
  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
-
335
+
295
336
  if np.max(abs(result[l-m]))>self._limit_range:
296
337
  result[l-m-1]*= self._limit_range
297
338
  result[l-m]*= self._limit_range
@@ -299,501 +340,598 @@ class alm():
299
340
  ratio[l-m]+= self._log_limit_range
300
341
  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
342
  """
302
-
303
- return self.backend.bk_complex(result*self.ratio_mm[nside,m],0*result)
304
343
 
305
-
344
+ return self.backend.bk_complex(result * self.ratio_mm[nside, m], 0 * result)
345
+
306
346
  # 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)))
347
+ def compute_legendre_spin2_m(self, co_th, si_th, m, lmax):
348
+ result = np.zeros([lmax - m + 2, co_th.shape[0]])
349
+ ratio = np.zeros([lmax - m + 2, 1])
350
+
351
+ ratio[1, 0] = self.double_factorial_log(2 * m - 1) - 0.5 * np.sum(
352
+ self.log(1 + np.arange(2 * m))
353
+ )
312
354
  # Étape 1 : Calcul de P_{mm}(x)
313
355
  if m == 0:
314
356
  Pmm = 1.0
315
357
  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
-
358
+ # Pmm = (-1)**m * (1 - x**2)**(m/2)
359
+ Pmm = (0.5 - m % 2) * 2 * (1 - co_th**2) ** (m / 2)
360
+
320
361
  # Si l == m, c'est directement P_{mm}
321
- result[1] = Pmm
322
-
362
+ result[1] = Pmm
363
+
323
364
  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)
365
+ ylm = result * np.exp(ratio)
366
+ ylm[1:] *= (
367
+ np.sqrt(4 * np.pi * (2 * (np.arange(lmax - m + 1) + m) + 1))
368
+ ).reshape(lmax + 1 - m, 1)
326
369
 
327
370
  else:
328
371
  # Étape 2 : Calcul de P_{l+1, m}(x)
329
- result[2] = co_th * (2*m + 1) * result[0]
372
+ result[2] = co_th * (2 * m + 1) * result[0]
373
+
374
+ ratio[2, 0] = ratio[1, 0] - self.log(2 * m + 1) / 2
330
375
 
331
- ratio[2,0] = ratio[1,0]-self.log(2*m+1)/2
332
-
333
376
  # É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)
377
+ for l in range(m + 2, lmax + 1):
378
+ result[l - m + 1] = (
379
+ (2 * l - 1) * co_th * result[l - m]
380
+ - (l + m - 1) * result[l - m - 1]
381
+ ) / (l - m)
382
+ ratio[l - m + 1, 0] = (self.log(l - m) - self.log(l + m)) / 2 + ratio[
383
+ l - m, 0
384
+ ]
385
+ if np.max(abs(result[l - m + 1])) > self._limit_range:
386
+ result[l - m] *= self._limit_range
387
+ result[l - m + 1] *= self._limit_range
388
+ ratio[l - m, 0] += self._log_limit_range
389
+ ratio[l - m + 1, 0] += self._log_limit_range
390
+
391
+ ylm = result * np.exp(ratio)
392
+ ylm[1:] *= (
393
+ np.sqrt(4 * np.pi * (2 * (np.arange(lmax - m + 1) + m) + 1))
394
+ ).reshape(lmax + 1 - m, 1)
395
+
396
+ ell = (np.arange(lmax + 1 - m) + m).reshape(lmax + 1 - m, 1)
397
+
398
+ cot_th = co_th / si_th
399
+ si2_th = si_th * si_th
400
+
401
+ a = (2 * m**2 - ell * (ell + 1)) / (
402
+ si2_th.reshape(1, si2_th.shape[0])
403
+ ) + ell * (ell - 1) * cot_th * cot_th
404
+ b = 2 * m * (ell - 1) * cot_th / si_th
405
+ w = np.zeros([lmax + 1 - m, 1])
406
+ l = ell[ell > 1]
407
+ w[ell > 1] = np.sqrt(1 / ((l + 2) * (l + 1) * (l) * (l - 1)))
408
+ w = w.reshape(lmax + 1 - m, 1)
409
+
410
+ alpha_plus = w * (a + b)
411
+ alpha_moins = w * (a - b)
412
+
413
+ a = 2 * np.sqrt((2 * ell + 1) / (2 * ell - 1) * (ell * ell - m * m))
414
+ b = m / si2_th
415
+
416
+ beta_plus = w * a * (cot_th / si_th + b)
417
+ beta_moins = w * a * (cot_th / si_th - b)
418
+
419
+ ylm_plus = alpha_plus * ylm[1:] + beta_plus * ylm[:-1]
420
+ ylm_moins = alpha_moins * ylm[1:] + beta_moins * ylm[:-1]
421
+
422
+ return ylm_plus, ylm_moins
423
+
424
+ def rfft2fft(self, val, axis=0):
425
+ r = self.backend.bk_rfft(val)
426
+ if axis == 0:
427
+ r_inv = self.backend.bk_reverse(
428
+ self.backend.bk_conjugate(r[1:-1]), axis=axis
429
+ )
376
430
  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 irfft2fft(self,val,N,axis=0):
381
- if axis==0:
382
- return self.backend.bk_irfft(val[0:N//2+1])
431
+ r_inv = self.backend.bk_reverse(
432
+ self.backend.bk_conjugate(r[:, 1:-1]), axis=axis
433
+ )
434
+ return self.backend.bk_concat([r, r_inv], axis=axis)
435
+
436
+ def irfft2fft(self, val, N, axis=0):
437
+ if axis == 0:
438
+ return self.backend.bk_irfft(val[0 : N // 2 + 1])
383
439
  else:
384
- return self.backend.bk_irfft(val[:,0:N//2+1])
385
-
386
- def comp_tf(self,im,nside,realfft=False):
387
-
440
+ return self.backend.bk_irfft(val[:, 0 : N // 2 + 1])
441
+
442
+ def comp_tf(self, im, nside, realfft=False):
443
+
388
444
  self.shift_ph(nside)
389
- n=0
390
-
391
- ft_im=[]
392
- for k in range(nside-1):
393
- N=4*(k+1)
394
-
445
+ n = 0
446
+
447
+ ft_im = []
448
+ for k in range(nside - 1):
449
+ N = 4 * (k + 1)
450
+
395
451
  if realfft:
396
- tmp=self.rfft2fft(im[n:n+N])
452
+ tmp = self.rfft2fft(im[n : n + N])
397
453
  else:
398
- tmp=self.backend.bk_fft(im[n:n+N])
399
-
400
- l_n=tmp.shape[0]
401
-
402
- if l_n<3*nside+1:
403
- repeat_n=3*nside//l_n+1
404
- tmp=self.backend.bk_tile(tmp,repeat_n,axis=0)
405
-
406
- ft_im.append(tmp[0:3*nside])
407
-
408
- n+=N
409
- if nside>1:
410
- result=self.backend.bk_reshape(self.backend.bk_concat(ft_im,axis=0),[nside-1,3*nside])
411
-
412
- N=4*nside*(2*nside+1)
413
- v=self.backend.bk_reshape(im[n:n+N],[2*nside+1,4*nside])
454
+ tmp = self.backend.bk_fft(im[n : n + N])
455
+
456
+ l_n = tmp.shape[0]
457
+
458
+ if l_n < 3 * nside + 1:
459
+ repeat_n = 3 * nside // l_n + 1
460
+ tmp = self.backend.bk_tile(tmp, repeat_n, axis=0)
461
+
462
+ ft_im.append(tmp[0 : 3 * nside])
463
+
464
+ n += N
465
+ if nside > 1:
466
+ result = self.backend.bk_reshape(
467
+ self.backend.bk_concat(ft_im, axis=0), [nside - 1, 3 * nside]
468
+ )
469
+
470
+ N = 4 * nside * (2 * nside + 1)
471
+ v = self.backend.bk_reshape(im[n : n + N], [2 * nside + 1, 4 * nside])
414
472
  if realfft:
415
- v_fft=self.rfft2fft(v,axis=1)[:,:3*nside]
473
+ v_fft = self.rfft2fft(v, axis=1)[:, : 3 * nside]
416
474
  else:
417
- v_fft=self.backend.bk_fft(v)[:,:3*nside]
475
+ v_fft = self.backend.bk_fft(v)[:, : 3 * nside]
418
476
 
419
- n+=N
420
- if nside>1:
421
- result=self.backend.bk_concat([result,v_fft],axis=0)
477
+ n += N
478
+ if nside > 1:
479
+ result = self.backend.bk_concat([result, v_fft], axis=0)
422
480
  else:
423
- result=v_fft
424
-
425
- if nside>1:
426
- ft_im=[]
427
- for k in range(nside-1):
428
- N=4*(nside-1-k)
481
+ result = v_fft
482
+
483
+ if nside > 1:
484
+ ft_im = []
485
+ for k in range(nside - 1):
486
+ N = 4 * (nside - 1 - k)
429
487
 
430
488
  if realfft:
431
- tmp=self.rfft2fft(im[n:n+N])[0:l_n]
489
+ tmp = self.rfft2fft(im[n : n + N])[0:l_n]
432
490
  else:
433
- tmp=self.backend.bk_fft(im[n:n+N])[0:l_n]
491
+ tmp = self.backend.bk_fft(im[n : n + N])[0:l_n]
434
492
 
435
- l_n=tmp.shape[0]
493
+ l_n = tmp.shape[0]
436
494
 
437
- if l_n<3*nside+1:
438
- repeat_n=3*nside//l_n+1
439
- tmp=self.backend.bk_tile(tmp,repeat_n,axis=0)
495
+ if l_n < 3 * nside + 1:
496
+ repeat_n = 3 * nside // l_n + 1
497
+ tmp = self.backend.bk_tile(tmp, repeat_n, axis=0)
440
498
 
441
- ft_im.append(tmp[0:3*nside])
442
- n+=N
499
+ ft_im.append(tmp[0 : 3 * nside])
500
+ n += N
443
501
 
444
- lastresult=self.backend.bk_reshape(self.backend.bk_concat(ft_im,axis=0),[nside-1,3*nside])
445
- return self.backend.bk_concat([result,lastresult],axis=0)*self.matrix_shift_ph[nside]
502
+ lastresult = self.backend.bk_reshape(
503
+ self.backend.bk_concat(ft_im, axis=0), [nside - 1, 3 * nside]
504
+ )
505
+ return (
506
+ self.backend.bk_concat([result, lastresult], axis=0)
507
+ * self.matrix_shift_ph[nside]
508
+ )
446
509
  else:
447
- return result*self.matrix_shift_ph[nside]
510
+ return result * self.matrix_shift_ph[nside]
511
+
512
+ def icomp_tf(self, i_im, nside, realfft=False):
448
513
 
449
-
450
- def icomp_tf(self,i_im,nside,realfft=False):
451
-
452
514
  self.shift_ph(nside)
453
-
454
- n=0
455
- im=[]
456
- ft_im=i_im*self.backend.bk_conjugate(self.matrix_shift_ph[nside])
457
-
458
- for k in range(nside-1):
459
- N=4*(k+1)
460
-
515
+
516
+ n = 0
517
+ im = []
518
+ ft_im = i_im * self.backend.bk_conjugate(self.matrix_shift_ph[nside])
519
+
520
+ for k in range(nside - 1):
521
+ N = 4 * (k + 1)
522
+
461
523
  if realfft:
462
- tmp=self.irfft2fft(ft_im[k],N)
524
+ tmp = self.irfft2fft(ft_im[k], N)
463
525
  else:
464
- tmp=self.backend.bk_ifft(im[k],N)
465
-
526
+ tmp = self.backend.bk_ifft(im[k], N)
527
+
466
528
  im.append(tmp[0:N])
467
-
468
- n+=N
469
-
470
- if nside>1:
471
- result=self.backend.bk_concat(im,axis=0)
472
-
473
- N=4*nside*(2*nside+1)
474
- v=ft_im[nside-1:3*nside,0:2*nside+1]
529
+
530
+ n += N
531
+
532
+ if nside > 1:
533
+ result = self.backend.bk_concat(im, axis=0)
534
+
535
+ N = 4 * nside * (2 * nside + 1)
536
+ v = ft_im[nside - 1 : 3 * nside, 0 : 2 * nside + 1]
475
537
  if realfft:
476
- v_fft=self.backend.bk_reshape(self.irfft2fft(v,N,axis=1),[4*nside*(2*nside+1)])
538
+ v_fft = self.backend.bk_reshape(
539
+ self.irfft2fft(v, N, axis=1), [4 * nside * (2 * nside + 1)]
540
+ )
477
541
  else:
478
- v_fft=self.backend.bk_ifft(v)
542
+ v_fft = self.backend.bk_ifft(v)
479
543
 
480
- n+=N
481
- if nside>1:
482
- result=self.backend.bk_concat([result,v_fft],axis=0)
544
+ n += N
545
+ if nside > 1:
546
+ result = self.backend.bk_concat([result, v_fft], axis=0)
483
547
  else:
484
- result=v_fft
485
-
486
- if nside>1:
487
- im=[]
488
- for k in range(nside-1):
489
- N=4*(nside-1-k)
490
-
548
+ result = v_fft
549
+
550
+ if nside > 1:
551
+ im = []
552
+ for k in range(nside - 1):
553
+ N = 4 * (nside - 1 - k)
554
+
491
555
  if realfft:
492
- tmp=self.irfft2fft(ft_im[k+3*nside],N)
556
+ tmp = self.irfft2fft(ft_im[k + 3 * nside], N)
493
557
  else:
494
- tmp=self.backend.bk_ifft(im[k+3*nside],N)
495
-
558
+ tmp = self.backend.bk_ifft(im[k + 3 * nside], N)
559
+
496
560
  im.append(tmp[0:N])
497
-
498
- n+=N
499
561
 
500
- return self.backend.bk_concat([result]+im,axis=0)
562
+ n += N
563
+
564
+ return self.backend.bk_concat([result] + im, axis=0)
501
565
  else:
502
566
  return result
503
-
504
- def anafast(self,im,map2=None,nest=False,spin=2):
505
-
506
- """The `anafast` function computes the L1 and L2 norm power spectra.
507
567
 
508
- Currently, it is not optimized for single-pass computation due to the relatively inefficient computation of \(Y_{lm}\).
568
+ def anafast(self, im, map2=None, nest=False, spin=2):
569
+ """The `anafast` function computes the L1 and L2 norm power spectra.
570
+
571
+ Currently, it is not optimized for single-pass computation due to the relatively inefficient computation of \(Y_{lm}\).
509
572
  Nonetheless, it utilizes TensorFlow and can be integrated into gradient computations.
510
573
 
511
574
  Input:
512
- - `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,
575
+ - `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,
513
576
  or of size \([3, 12 \times \text{Nside}^2]\) for I,Q,U polar data.
514
- - `map2` (optional): a vector of size \([12 \times \text{Nside}^2]\) for scalar data, or of size
577
+ - `map2` (optional): a vector of size \([12 \times \text{Nside}^2]\) for scalar data, or of size
515
578
  \([3, 12 \times \text{Nside}^2]\) for polar data. If provided, cross power spectra will be computed.
516
579
  - `nest=True`: alters the ordering of the input maps.
517
580
  - `spin=2` for 1/2 spin data as Q and U. Spin=1 for seep fields
518
581
 
519
582
  Output:
520
- -A tensor of size \([l_{\text{max}} \times (l_{\text{max}}-1)]\) formatted as \([6, \ldots]\),
583
+ -A tensor of size \([l_{\text{max}} \times (l_{\text{max}}-1)]\) formatted as \([6, \ldots]\),
521
584
  ordered as TT, EE, BB, TE, EB.TBanafast function computes L1 and L2 norm powerspctra.
522
585
 
523
586
  """
524
- i_im=self.backend.bk_cast(im)
587
+ i_im = self.backend.bk_cast(im)
525
588
  if map2 is not None:
526
- i_map2=self.backend.bk_cast(map2)
527
-
528
- doT=True
529
- if len(i_im.shape)==1: # nopol
530
- nside=int(np.sqrt(i_im.shape[0]//12))
589
+ i_map2 = self.backend.bk_cast(map2)
590
+
591
+ doT = True
592
+ if len(i_im.shape) == 1: # nopol
593
+ nside = int(np.sqrt(i_im.shape[0] // 12))
531
594
  else:
532
- if i_im.shape[0]==2:
533
- doT=False
534
- nside=int(np.sqrt(i_im.shape[1]//12))
535
-
595
+ if i_im.shape[0] == 2:
596
+ doT = False
597
+ nside = int(np.sqrt(i_im.shape[1] // 12))
598
+
536
599
  self.shift_ph(nside)
537
-
538
- if doT: # nopol
539
- if len(i_im.shape)==2: # pol
540
- l_im=i_im[0]
600
+
601
+ if doT: # nopol
602
+ if len(i_im.shape) == 2: # pol
603
+ l_im = i_im[0]
541
604
  if map2 is not None:
542
- l_map2=i_map2[0]
605
+ l_map2 = i_map2[0]
543
606
  else:
544
- l_im=i_im
607
+ l_im = i_im
545
608
  if map2 is not None:
546
- l_map2=i_map2
547
-
609
+ l_map2 = i_map2
610
+
548
611
  if nest:
549
- idx=hp.ring2nest(nside,np.arange(12*nside**2))
550
- if len(i_im.shape)==1: # nopol
551
- ft_im=self.comp_tf(self.backend.bk_gather(l_im,idx),nside,realfft=True)
612
+ idx = hp.ring2nest(nside, np.arange(12 * nside**2))
613
+ if len(i_im.shape) == 1: # nopol
614
+ ft_im = self.comp_tf(
615
+ self.backend.bk_gather(l_im, idx), nside, realfft=True
616
+ )
552
617
  if map2 is not None:
553
- ft_im2=self.comp_tf(self.backend.bk_gather(l_map2,idx),nside,realfft=True)
618
+ ft_im2 = self.comp_tf(
619
+ self.backend.bk_gather(l_map2, idx), nside, realfft=True
620
+ )
554
621
  else:
555
- ft_im=self.comp_tf(l_im,nside,realfft=True)
622
+ ft_im = self.comp_tf(l_im, nside, realfft=True)
556
623
  if map2 is not None:
557
- ft_im2=self.comp_tf(l_map2,nside,realfft=True)
558
-
559
- lth=self.ring_th(nside)
560
-
561
- co_th=np.cos(lth)
562
-
563
- lmax=3*nside-1
564
-
565
- cl2=None
566
- cl2_L1=None
567
- dt2=0
568
- dt3=0
569
- dt4=0
570
- if len(i_im.shape)==2: # nopol
571
-
572
- self.init_Ys(spin,nside)
573
-
624
+ ft_im2 = self.comp_tf(l_map2, nside, realfft=True)
625
+
626
+ lth = self.ring_th(nside)
627
+
628
+ co_th = np.cos(lth)
629
+
630
+ lmax = 3 * nside - 1
631
+
632
+ cl2 = None
633
+ cl2_L1 = None
634
+ dt2 = 0
635
+ dt3 = 0
636
+ dt4 = 0
637
+ if len(i_im.shape) == 2: # nopol
638
+
639
+ self.init_Ys(spin, nside)
640
+
574
641
  if nest:
575
- idx=hp.ring2nest(nside,np.arange(12*nside**2))
576
- l_Q=self.backend.bk_gather(i_im[int(doT)],idx)
577
- l_U=self.backend.bk_gather(i_im[1+int(doT)],idx)
578
- ft_im_Pp=self.comp_tf(self.backend.bk_complex(l_Q,l_U),nside)
579
- ft_im_Pm=self.comp_tf(self.backend.bk_complex(l_Q,-l_U),nside)
642
+ idx = hp.ring2nest(nside, np.arange(12 * nside**2))
643
+ l_Q = self.backend.bk_gather(i_im[int(doT)], idx)
644
+ l_U = self.backend.bk_gather(i_im[1 + int(doT)], idx)
645
+ ft_im_Pp = self.comp_tf(self.backend.bk_complex(l_Q, l_U), nside)
646
+ ft_im_Pm = self.comp_tf(self.backend.bk_complex(l_Q, -l_U), nside)
580
647
  if map2 is not None:
581
- l_Q=self.backend.bk_gather(i_map2[int(doT)],idx)
582
- l_U=self.backend.bk_gather(i_map2[1+int(doT)],idx)
583
- ft_im2_Pp=self.comp_tf(self.backend.bk_complex(l_Q,l_U),nside)
584
- ft_im2_Pm=self.comp_tf(self.backend.bk_complex(l_Q,-l_U),nside)
648
+ l_Q = self.backend.bk_gather(i_map2[int(doT)], idx)
649
+ l_U = self.backend.bk_gather(i_map2[1 + int(doT)], idx)
650
+ ft_im2_Pp = self.comp_tf(self.backend.bk_complex(l_Q, l_U), nside)
651
+ ft_im2_Pm = self.comp_tf(self.backend.bk_complex(l_Q, -l_U), nside)
585
652
  else:
586
- ft_im_Pp=self.comp_tf(self.backend.bk_complex(i_im[int(doT)],i_im[1+int(doT)]),nside)
587
- ft_im_Pm=self.comp_tf(self.backend.bk_complex(i_im[int(doT)],-i_im[1+int(doT)]),nside)
653
+ ft_im_Pp = self.comp_tf(
654
+ self.backend.bk_complex(i_im[int(doT)], i_im[1 + int(doT)]), nside
655
+ )
656
+ ft_im_Pm = self.comp_tf(
657
+ self.backend.bk_complex(i_im[int(doT)], -i_im[1 + int(doT)]), nside
658
+ )
588
659
  if map2 is not None:
589
- ft_im2_Pp=self.comp_tf(self.backend.bk_complex(i_map2[int(doT)],i_map2[1+int(doT)]),nside)
590
- ft_im2_Pm=self.comp_tf(self.backend.bk_complex(i_map2[int(doT)],-i_map2[1+int(doT)]),nside)
660
+ ft_im2_Pp = self.comp_tf(
661
+ self.backend.bk_complex(i_map2[int(doT)], i_map2[1 + int(doT)]),
662
+ nside,
663
+ )
664
+ ft_im2_Pm = self.comp_tf(
665
+ self.backend.bk_complex(
666
+ i_map2[int(doT)], -i_map2[1 + int(doT)]
667
+ ),
668
+ nside,
669
+ )
670
+
671
+ for m in range(lmax + 1):
672
+
673
+ plm = self.compute_legendre_m(co_th, m, 3 * nside - 1, nside) / (
674
+ 12 * nside**2
675
+ )
591
676
 
592
- for m in range(lmax+1):
593
-
594
- plm=self.compute_legendre_m(co_th,m,3*nside-1,nside)/(12*nside**2)
595
-
596
677
  if doT:
597
- tmp=self.backend.bk_reduce_sum(plm*ft_im[:,m],1)
678
+ tmp = self.backend.bk_reduce_sum(plm * ft_im[:, m], 1)
598
679
 
599
680
  if map2 is not None:
600
- tmp2=self.backend.bk_reduce_sum(plm*ft_im2[:,m],1)
681
+ tmp2 = self.backend.bk_reduce_sum(plm * ft_im2[:, m], 1)
601
682
  else:
602
- tmp2=tmp
603
-
604
- if len(i_im.shape)==2: # pol
605
- plmp=self.Yp[spin,nside][m]
606
- plmm=self.Ym[spin,nside][m]
607
-
608
- tmpp=self.backend.bk_reduce_sum(plmp*ft_im_Pp[:,m],1)
609
- tmpm=self.backend.bk_reduce_sum(plmm*ft_im_Pm[:,m],1)
610
-
611
- almE=-(tmpp+tmpm)/2.0
612
- almB=(tmpp-tmpm)/(2J)
613
-
683
+ tmp2 = tmp
684
+
685
+ if len(i_im.shape) == 2: # pol
686
+ plmp = self.Yp[spin, nside][m]
687
+ plmm = self.Ym[spin, nside][m]
688
+
689
+ tmpp = self.backend.bk_reduce_sum(plmp * ft_im_Pp[:, m], 1)
690
+ tmpm = self.backend.bk_reduce_sum(plmm * ft_im_Pm[:, m], 1)
691
+
692
+ almE = -(tmpp + tmpm) / 2.0
693
+ almB = (tmpp - tmpm) / (2j)
694
+
614
695
  if map2 is not None:
615
- tmpp2=self.backend.bk_reduce_sum(plmp*ft_im2_Pp[:,m],1)
616
- tmpm2=self.backend.bk_reduce_sum(plmm*ft_im2_Pm[:,m],1)
617
-
618
- almE2=-(tmpp2+tmpm2)/2.0
619
- almB2=(tmpp2-tmpm2)/(2J)
696
+ tmpp2 = self.backend.bk_reduce_sum(plmp * ft_im2_Pp[:, m], 1)
697
+ tmpm2 = self.backend.bk_reduce_sum(plmm * ft_im2_Pm[:, m], 1)
698
+
699
+ almE2 = -(tmpp2 + tmpm2) / 2.0
700
+ almB2 = (tmpp2 - tmpm2) / (2j)
620
701
  else:
621
- almE2=almE
622
- almB2=almB
623
-
702
+ almE2 = almE
703
+ almB2 = almB
704
+
624
705
  if doT:
625
- tmpTT=self.backend.bk_real((tmp*self.backend.bk_conjugate(tmp2)))
626
- tmpTE=self.backend.bk_real((tmp*self.backend.bk_conjugate(almE2)))
627
- tmpTB=-self.backend.bk_real((tmp*self.backend.bk_conjugate(almB2)))
628
-
629
- tmpEE=self.backend.bk_real((almE*self.backend.bk_conjugate(almE2)))
630
- tmpBB=self.backend.bk_real((almB*self.backend.bk_conjugate(almB2)))
631
- tmpEB=-self.backend.bk_real((almE*self.backend.bk_conjugate(almB2)))
632
-
706
+ tmpTT = self.backend.bk_real(
707
+ (tmp * self.backend.bk_conjugate(tmp2))
708
+ )
709
+ tmpTE = self.backend.bk_real(
710
+ (tmp * self.backend.bk_conjugate(almE2))
711
+ )
712
+ tmpTB = -self.backend.bk_real(
713
+ (tmp * self.backend.bk_conjugate(almB2))
714
+ )
715
+
716
+ tmpEE = self.backend.bk_real((almE * self.backend.bk_conjugate(almE2)))
717
+ tmpBB = self.backend.bk_real((almB * self.backend.bk_conjugate(almB2)))
718
+ tmpEB = -self.backend.bk_real((almE * self.backend.bk_conjugate(almB2)))
719
+
633
720
  if map2 is not None:
634
- tmpEB=(tmpEB-self.backend.bk_real((almE2*self.backend.bk_conjugate(almB))))/2
635
-
636
- if doT:
637
- tmpTE=(tmpTE+self.backend.bk_real((tmp2*self.backend.bk_conjugate(almE))))/2
638
- tmpTB=(tmpTB-self.backend.bk_real((tmp2*self.backend.bk_conjugate(almB))))/2
639
-
721
+ tmpEB = (
722
+ tmpEB
723
+ - self.backend.bk_real(
724
+ (almE2 * self.backend.bk_conjugate(almB))
725
+ )
726
+ ) / 2
640
727
 
641
- if m==0:
642
728
  if doT:
643
- l_cl=self.backend.bk_concat([tmpTT,tmpEE,tmpBB,tmpTE,tmpEB,tmpTB],0)
729
+ tmpTE = (
730
+ tmpTE
731
+ + self.backend.bk_real(
732
+ (tmp2 * self.backend.bk_conjugate(almE))
733
+ )
734
+ ) / 2
735
+ tmpTB = (
736
+ tmpTB
737
+ - self.backend.bk_real(
738
+ (tmp2 * self.backend.bk_conjugate(almB))
739
+ )
740
+ ) / 2
741
+
742
+ if m == 0:
743
+ if doT:
744
+ l_cl = self.backend.bk_concat(
745
+ [tmpTT, tmpEE, tmpBB, tmpTE, tmpEB, tmpTB], 0
746
+ )
644
747
  else:
645
- l_cl=self.backend.bk_concat([tmpEE,tmpBB,tmpEB],0)
748
+ l_cl = self.backend.bk_concat([tmpEE, tmpBB, tmpEB], 0)
646
749
  else:
647
- offset_tensor=self.backend.bk_zeros((m),dtype=self.backend.all_bk_type)
750
+ offset_tensor = self.backend.bk_zeros(
751
+ (m), dtype=self.backend.all_bk_type
752
+ )
648
753
  if doT:
649
- l_cl=self.backend.bk_concat([self.backend.bk_concat([offset_tensor,tmpTT],axis=0),
650
- self.backend.bk_concat([offset_tensor,tmpEE],axis=0),
651
- self.backend.bk_concat([offset_tensor,tmpBB],axis=0),
652
- self.backend.bk_concat([offset_tensor,tmpTE],axis=0),
653
- self.backend.bk_concat([offset_tensor,tmpEB],axis=0),
654
- self.backend.bk_concat([offset_tensor,tmpTB],axis=0)],axis=0)
754
+ l_cl = self.backend.bk_concat(
755
+ [
756
+ self.backend.bk_concat([offset_tensor, tmpTT], axis=0),
757
+ self.backend.bk_concat([offset_tensor, tmpEE], axis=0),
758
+ self.backend.bk_concat([offset_tensor, tmpBB], axis=0),
759
+ self.backend.bk_concat([offset_tensor, tmpTE], axis=0),
760
+ self.backend.bk_concat([offset_tensor, tmpEB], axis=0),
761
+ self.backend.bk_concat([offset_tensor, tmpTB], axis=0),
762
+ ],
763
+ axis=0,
764
+ )
655
765
  else:
656
- l_cl=self.backend.bk_concat([self.backend.bk_concat([offset_tensor,tmpEE],axis=0),
657
- self.backend.bk_concat([offset_tensor,tmpBB],axis=0),
658
- self.backend.bk_concat([offset_tensor,tmpEB],axis=0)],axis=0)
659
-
766
+ l_cl = self.backend.bk_concat(
767
+ [
768
+ self.backend.bk_concat([offset_tensor, tmpEE], axis=0),
769
+ self.backend.bk_concat([offset_tensor, tmpBB], axis=0),
770
+ self.backend.bk_concat([offset_tensor, tmpEB], axis=0),
771
+ ],
772
+ axis=0,
773
+ )
774
+
660
775
  if doT:
661
- l_cl=self.backend.bk_reshape(l_cl,[6,lmax+1])
776
+ l_cl = self.backend.bk_reshape(l_cl, [6, lmax + 1])
662
777
  else:
663
- l_cl=self.backend.bk_reshape(l_cl,[3,lmax+1])
778
+ l_cl = self.backend.bk_reshape(l_cl, [3, lmax + 1])
664
779
  else:
665
- tmp=self.backend.bk_real((tmp*self.backend.bk_conjugate(tmp2)))
666
- if m==0:
667
- l_cl=tmp
780
+ tmp = self.backend.bk_real((tmp * self.backend.bk_conjugate(tmp2)))
781
+ if m == 0:
782
+ l_cl = tmp
668
783
  else:
669
- offset_tensor=self.backend.bk_zeros((m),dtype=self.backend.all_bk_type)
670
- l_cl=self.backend.bk_concat([offset_tensor,tmp],axis=0)
671
-
784
+ offset_tensor = self.backend.bk_zeros(
785
+ (m), dtype=self.backend.all_bk_type
786
+ )
787
+ l_cl = self.backend.bk_concat([offset_tensor, tmp], axis=0)
788
+
672
789
  if cl2 is None:
673
- cl2=l_cl
790
+ cl2 = l_cl
674
791
  else:
675
- cl2+=2*l_cl
676
-
677
- #cl2=cl2*(4*np.pi) #self.backend.bk_sqrt(self.backend.bk_cast(4*np.pi)) #(2*np.arange(cl2.shape[0])+1)))
678
-
679
- cl2_l1=self.backend.bk_L1(cl2)
680
-
681
- return cl2,cl2_l1
682
-
683
- def map2alm(self,im,nest=False):
684
- nside=int(np.sqrt(im.shape[0]//12))
685
-
686
- ph=self.shift_ph(nside)
687
-
792
+ cl2 += 2 * l_cl
793
+
794
+ # cl2=cl2*(4*np.pi) #self.backend.bk_sqrt(self.backend.bk_cast(4*np.pi)) #(2*np.arange(cl2.shape[0])+1)))
795
+
796
+ cl2_l1 = self.backend.bk_L1(cl2)
797
+
798
+ return cl2, cl2_l1
799
+
800
+ def map2alm(self, im, nest=False):
801
+ nside = int(np.sqrt(im.shape[0] // 12))
802
+
803
+ ph = self.shift_ph(nside)
804
+
688
805
  if nest:
689
- idx=hp.ring2nest(nside,np.arange(12*nside**2))
690
- ft_im=self.comp_tf(self.backend.bk_cast(self.backend.bk_gather(im,idx)),nside,realfft=True)
806
+ idx = hp.ring2nest(nside, np.arange(12 * nside**2))
807
+ ft_im = self.comp_tf(
808
+ self.backend.bk_cast(self.backend.bk_gather(im, idx)),
809
+ nside,
810
+ realfft=True,
811
+ )
691
812
  else:
692
- ft_im=self.comp_tf(self.backend.bk_cast(im),nside,realfft=True)
693
-
694
- lth=self.ring_th(nside)
695
-
696
- co_th=np.cos(lth)
697
-
698
- lmax=3*nside-1
699
-
700
- alm=None
701
- for m in range(lmax+1):
702
- plm=self.compute_legendre_m(co_th,m,3*nside-1,nside)/(12*nside**2)
703
-
704
- tmp=self.backend.bk_reduce_sum(plm*ft_im[:,m],1)
705
- if m==0:
706
- alm=tmp
813
+ ft_im = self.comp_tf(self.backend.bk_cast(im), nside, realfft=True)
814
+
815
+ lth = self.ring_th(nside)
816
+
817
+ co_th = np.cos(lth)
818
+
819
+ lmax = 3 * nside - 1
820
+
821
+ alm = None
822
+ for m in range(lmax + 1):
823
+ plm = self.compute_legendre_m(co_th, m, 3 * nside - 1, nside) / (
824
+ 12 * nside**2
825
+ )
826
+
827
+ tmp = self.backend.bk_reduce_sum(plm * ft_im[:, m], 1)
828
+ if m == 0:
829
+ alm = tmp
707
830
  else:
708
- alm=self.backend.bk_concat([alm,tmp],axis=0)
709
-
831
+ alm = self.backend.bk_concat([alm, tmp], axis=0)
832
+
710
833
  return alm
711
834
 
712
-
713
- def alm2map(self,nside,alm):
835
+ def alm2map(self, nside, alm):
836
+
837
+ lth = self.ring_th(nside)
838
+
839
+ co_th = np.cos(lth)
714
840
 
715
- lth=self.ring_th(nside)
841
+ ft_im = []
716
842
 
717
- co_th=np.cos(lth)
718
-
719
- ft_im=[]
843
+ n = 0
720
844
 
721
- n=0
845
+ lmax = 3 * nside - 1
722
846
 
723
- lmax=3*nside-1
724
-
725
- for m in range(lmax+1):
726
- plm=self.compute_legendre_m(co_th,m,3*nside-1,nside)/(12*nside**2)
847
+ for m in range(lmax + 1):
848
+ plm = self.compute_legendre_m(co_th, m, 3 * nside - 1, nside) / (
849
+ 12 * nside**2
850
+ )
727
851
 
728
- print(alm[n:n+lmax-m+1].shape,plm.shape)
729
- ft_im.append(self.backend.bk_reduce_sum(self.backend.bk_reshape(alm[n:n+lmax-m+1],[lmax-m+1,1])*plm,0))
852
+ print(alm[n : n + lmax - m + 1].shape, plm.shape)
853
+ ft_im.append(
854
+ self.backend.bk_reduce_sum(
855
+ self.backend.bk_reshape(
856
+ alm[n : n + lmax - m + 1], [lmax - m + 1, 1]
857
+ )
858
+ * plm,
859
+ 0,
860
+ )
861
+ )
862
+
863
+ n = n + lmax - m + 1
864
+
865
+ return self.backend.bk_reshape(
866
+ self.backend.bk_concat(ft_im, 0), [lmax + 1, 4 * nside - 1]
867
+ )
730
868
 
731
- n=n+lmax-m+1
732
-
733
- return self.backend.bk_reshape(self.backend.bk_concat(ft_im,0),[lmax+1,4*nside-1])
734
-
735
-
736
869
  if nest:
737
- idx=hp.ring2nest(nside,np.arange(12*nside**2))
738
- ft_im=self.comp_tf(self.backend.bk_cast(self.backend.bk_gather(im,idx)),nside,realfft=True)
870
+ idx = hp.ring2nest(nside, np.arange(12 * nside**2))
871
+ ft_im = self.comp_tf(
872
+ self.backend.bk_cast(self.backend.bk_gather(im, idx)),
873
+ nside,
874
+ realfft=True,
875
+ )
739
876
  else:
740
- ft_im=self.comp_tf(self.backend.bk_cast(im),nside,realfft=True)
741
-
742
-
743
- lmax=3*nside-1
744
-
745
- alm=None
746
- for m in range(lmax+1):
747
- plm=self.compute_legendre_m(co_th,m,3*nside-1,nside)/(12*nside**2)
748
-
749
- tmp=self.backend.bk_reduce_sum(plm*ft_im[:,m],1)
750
- if m==0:
751
- alm=tmp
877
+ ft_im = self.comp_tf(self.backend.bk_cast(im), nside, realfft=True)
878
+
879
+ lmax = 3 * nside - 1
880
+
881
+ alm = None
882
+ for m in range(lmax + 1):
883
+ plm = self.compute_legendre_m(co_th, m, 3 * nside - 1, nside) / (
884
+ 12 * nside**2
885
+ )
886
+
887
+ tmp = self.backend.bk_reduce_sum(plm * ft_im[:, m], 1)
888
+ if m == 0:
889
+ alm = tmp
752
890
  else:
753
- alm=self.backend.bk_concat([alm,tmp],axis=0)
754
-
891
+ alm = self.backend.bk_concat([alm, tmp], axis=0)
892
+
755
893
  return o_map
756
894
 
757
- def map2alm_spin(self,im_Q,im_U,spin=2,nest=False):
758
-
759
- if spin==0:
760
- return self.map2alm(im_Q,nest=nest),self.map2alm(im_U,nest=nest)
895
+ def map2alm_spin(self, im_Q, im_U, spin=2, nest=False):
896
+
897
+ if spin == 0:
898
+ return self.map2alm(im_Q, nest=nest), self.map2alm(im_U, nest=nest)
761
899
 
762
- nside=int(np.sqrt(im_Q.shape[0]//12))
900
+ nside = int(np.sqrt(im_Q.shape[0] // 12))
901
+
902
+ lth = self.ring_th(nside)
903
+
904
+ co_th = np.cos(lth)
763
905
 
764
- lth=self.ring_th(nside)
765
-
766
- co_th=np.cos(lth)
767
-
768
906
  if nest:
769
- idx=hp.ring2nest(nside,np.arange(12*nside**2))
770
- l_Q=self.backend.bk_gather(im_Q,idx)
771
- l_U=self.backend.bk_gather(im_U,idx)
772
- ft_im_1=self.comp_tf(self.backend.bk_complex(l_Q,l_U),nside)
773
- ft_im_2=self.comp_tf(self.backend.bk_complex(l_Q,-l_U),nside)
907
+ idx = hp.ring2nest(nside, np.arange(12 * nside**2))
908
+ l_Q = self.backend.bk_gather(im_Q, idx)
909
+ l_U = self.backend.bk_gather(im_U, idx)
910
+ ft_im_1 = self.comp_tf(self.backend.bk_complex(l_Q, l_U), nside)
911
+ ft_im_2 = self.comp_tf(self.backend.bk_complex(l_Q, -l_U), nside)
774
912
  else:
775
- ft_im_1=self.comp_tf(self.backend.bk_complex(im_Q,im_U),nside)
776
- ft_im_2=self.comp_tf(self.backend.bk_complex(im_Q,-im_U),nside)
777
-
778
- lmax=3*nside-1
779
-
780
- alm=None
781
- for m in range(lmax+1):
782
- #not yet debug use spherical
783
- #plmp1,plmm1=self.compute_legendre_spin2_m(co_th,si_th,m,3*nside-1)
784
- #plmp1/=(12*nside**2)
785
- #plmm1/=(12*nside**2)
786
-
787
- plmp=self.Yp[spin,nside][m]
788
- plmm=self.Ym[spin,nside][m]
789
-
790
- tmpp=self.backend.bk_reduce_sum(plmp*ft_im_1[:,m],1)
791
- tmpm=self.backend.bk_reduce_sum(plmm*ft_im_2[:,m],1)
792
- if m==0:
793
- almE=-(tmpp+tmpm)/2.0
794
- almB=(tmpp-tmpm)/(2J)
913
+ ft_im_1 = self.comp_tf(self.backend.bk_complex(im_Q, im_U), nside)
914
+ ft_im_2 = self.comp_tf(self.backend.bk_complex(im_Q, -im_U), nside)
915
+
916
+ lmax = 3 * nside - 1
917
+
918
+ alm = None
919
+ for m in range(lmax + 1):
920
+ # not yet debug use spherical
921
+ # plmp1,plmm1=self.compute_legendre_spin2_m(co_th,si_th,m,3*nside-1)
922
+ # plmp1/=(12*nside**2)
923
+ # plmm1/=(12*nside**2)
924
+
925
+ plmp = self.Yp[spin, nside][m]
926
+ plmm = self.Ym[spin, nside][m]
927
+
928
+ tmpp = self.backend.bk_reduce_sum(plmp * ft_im_1[:, m], 1)
929
+ tmpm = self.backend.bk_reduce_sum(plmm * ft_im_2[:, m], 1)
930
+ if m == 0:
931
+ almE = -(tmpp + tmpm) / 2.0
932
+ almB = (tmpp - tmpm) / (2j)
795
933
  else:
796
- almE=self.backend.bk_concat([almE,-(tmpp+tmpm)/2],axis=0)
797
- almB=self.backend.bk_concat([almB,(tmpp-tmpm)/(2J)],axis=0)
798
-
799
- return almE,almB
934
+ almE = self.backend.bk_concat([almE, -(tmpp + tmpm) / 2], axis=0)
935
+ almB = self.backend.bk_concat([almB, (tmpp - tmpm) / (2j)], axis=0)
936
+
937
+ return almE, almB