Coverage for pygeodyn/shear/blackBoxFormula2temp.py: 50%

1import numpy as np

2from scipy.special import lpmv

3import math

6Rc = 3485.0 * 10**0 # CMB radius (km)

7Ra = 6371.2 * 10**0 # Earth's reference radius (km)

10# =============================================

11# calculation of the spectrum

12# =============================================

15def calculate_spectrum(coef, L):

16 spectrum = np.zeros((L,))

17 iter = 0

18 for l in range(1, L + 1):

19 temp = 0

20 for m in range(l + 1):

21 if m == 0:

22 temp += coef[iter] ** 2

23 iter += 1

24 else:

25 temp += coef[iter] ** 2

26 temp += coef[iter + 1] ** 2

27 iter += 2

28 spectrum[l - 1] += temp * (l + 1)

29 return spectrum

32# =============================================

33# calculation Gauss points and weights

34# =============================================

37def gaussLeg(lower_bound, upper_bound, N):

38 M = (N + 1) // 2

39 XM = 0.5 * (upper_bound + lower_bound)

40 XL = 0.5 * (upper_bound - lower_bound)

42 gauss_pts = np.zeros(N)

43 gauss_wgths = np.zeros(N)

45 pp = 0

46 ppp = 0

48 EPS = 10**-14

50 for i in range(1, M + 1):

51 converged = False

52 z = np.cos(np.pi * (i - 0.25) / (N + 0.5))

53 while not converged:

54 p1 = 1.0

55 p2 = 0.0

57 for j in range(1, N + 1):

58 p3 = p2

59 p2 = p1

60 p1 = ((2.0 * j - 1.0) * z * p2 - (j - 1.0) * p3) / j

62 pp = N * (z * p1 - p2) / ((z - 1.0) * (z + 1.0))

63 ppp = N * (z * p1 - p2)

64 z1 = z

65 z = z1 - p1 / pp

67 if np.abs(z - z1) < EPS:

68 converged = not converged

70 gauss_pts[i - 1] = XM - XL * z

71 gauss_pts[N + 1 - i - 1] = XM + XL * z

72 gauss_wgths[i - 1] = 2.0 * XL * (1.0 - z) * (1.0 + z) / (ppp * ppp)

73 gauss_wgths[N + 1 - i - 1] = gauss_wgths[i - 1]

75 return gauss_pts, gauss_wgths

78# =============================================

79# calculation Gauss theta points and weights

80# =============================================

83def gaussPoints(thetaMin, thetaMax, N):

84 gauss_pts, gauss_wgths = gaussLeg(np.cos(thetaMax), np.cos(thetaMin), N)

85 gauss_theta = np.arccos(gauss_pts)

87 return gauss_theta, gauss_wgths

90# =============================================

91# calculation norm coefficient for QN Shmidt

92# =============================================

95def norm_coef(l, m, *args):

96 if m == 0:

97 return 1

98 else:

99 return np.sqrt((2) * ((math.factorial(l - m)) / (math.factorial(l + m))))

100

101

102# =============================================

103# calculation norm coefficient full norm

104# =============================================

105

106

107def full_norm(l, m, *args):

108 return np.sqrt(2 * l + 1) * np.sqrt(

109 np.math.factorial(l - m) / np.math.factorial(l + m)

110 )

111

112

113# =============================================

114# calculation of glm is needed for tlm

115# =============================================

116

117

118def calculate_glm_hat(l, m):

119 return np.sqrt((l + m) * (l - m))

120

121

122# =============================================

123# Associated Legendre functions Pn,m(x)

124# =============================================

125

126

127def assocLegendrePoly(

128 l,

129 m,

130 theta,

131):

132 Plm = lpmv(m, l, np.cos(theta)) * ((-1) ** m)

133 return Plm

134

135

136# =============================================

137# Derivative of Associated Legendre functions

138# dPn,m(x) / dtheta

139# =============================================

140

141

142def derivativeThetaAssocLegendrePoly(l, m, theta):

143 Plm = (l + 1) * (-np.cos(theta) / np.sin(theta)) * lpmv(m, l, np.cos(theta)) + (

144 l - m + 1

145 ) * (1 / np.sin(theta)) * lpmv(m, l + 1, np.cos(theta))

146 return Plm * ((-1) ** m)

147

148

149

150# =============================================

151# Building matrix corresponding qlm = M clm

152# =============================================

153

154

155def build_clm_to_qlm_matrix(Lmax):

156

157 Nu = Lmax * (Lmax + 2)

158 clm_to_qlm_matrix = np.zeros((Nu, Nu))

159

160 nu = 0

161 for lu in range(1, Lmax + 1):

162 for mu in range(0, lu + 1):

163 temp_coef = -2 * mu / (lu * (lu + 1))

164 if mu == 0:

165 nu += 1

166 else:

167 clm_to_qlm_matrix[nu, nu + 1] = -1 / temp_coef

168 clm_to_qlm_matrix[nu + 1, nu] = 1 / temp_coef

169 nu += 2

170

171 return clm_to_qlm_matrix

172

173

174# =============================================

175# Building matrix corresponding clm = M qlm

176# =============================================

177

178

179def build_qlm_to_clm_matrix(Lmax):

180

181 Nu = Lmax * (Lmax + 2)

182 qlm_to_clm_matrix = np.zeros((Nu, Nu))

183

184 nu = 0

185 for lu in range(1, Lmax + 1):

186 for mu in range(0, lu + 1):

187 temp_coef = -2 * mu / (lu * (lu + 1))

188 if mu == 0:

189 nu += 1

190 else:

191 qlm_to_clm_matrix[nu, nu + 1] = -temp_coef

192 qlm_to_clm_matrix[nu + 1, nu] = temp_coef

193 nu += 2

194

195 return qlm_to_clm_matrix

196

197

198# =============================================

199# Building matrix corresponding tlm = T qlm

200# =============================================

201

202

203def build_clm_to_tlm_matrix(Lmax):

204

205 Nu = Lmax * (Lmax + 2)

206 clm_to_tlm_matrix = np.zeros((Nu, Nu))

207

208 lumu_dict = {}

209

210 nu = 0

211 for lu in range(1, Lmax + 1):

212 for mu in range(0, lu + 1):

213 lumu_dict[(lu, mu, "cos")] = nu

214 nu += 1

215 if mu != 0:

216 lumu_dict[(lu, mu, "sin")] = nu

217 nu += 1

218

219 nu = 0

220 for lu in range(1, Lmax + 1):

221 for mu in range(0, lu + 1):

222 if lu <= 1:

223 pass

224 else:

225 left_coef = 1 / (2 * lu + 3)

226 rigth_coef = 1 / (2 * lu - 1)

227 if lu < Lmax:

228

229 if mu == 0:

230 pass

231 else:

232 temp_coef = -1 / (2 * mu) * (lu - 1) * (lu + 2)

233 nu = lumu_dict[(lu, mu, "cos")]

234 nu_plus = lumu_dict[(lu + 1, mu, "sin")]

235 if (lu - 1, mu, "sin") in lumu_dict:

236 nu_minus = lumu_dict[(lu - 1, mu, "sin")]

237

238 glm_plus = calculate_glm_hat(lu + 1, mu)

239 glm_minus = calculate_glm_hat(lu, mu)

240

241 if (lu - 1, mu, "sin") not in lumu_dict:

242 clm_to_tlm_matrix[nu, nu_plus] = (

243 -glm_plus * left_coef * temp_coef

244 )

245 else:

246 clm_to_tlm_matrix[nu, nu_plus] = (

247 -glm_plus * left_coef * temp_coef

248 )

249 clm_to_tlm_matrix[nu, nu_minus] = (

250 -glm_minus * rigth_coef * temp_coef

251 )

252

253 nu = lumu_dict[(lu, mu, "sin")]

254 nu_plus = lumu_dict[(lu + 1, mu, "cos")]

255 if (lu - 1, mu, "cos") in lumu_dict:

256 nu_minus = lumu_dict[(lu - 1, mu, "cos")]

257

258 if (lu - 1, mu, "cos") not in lumu_dict:

259 clm_to_tlm_matrix[nu, nu_plus] = (

260 glm_plus * left_coef * temp_coef

261 )

262 else:

263 clm_to_tlm_matrix[nu, nu_plus] = (

264 glm_plus * left_coef * temp_coef

265 )

266 clm_to_tlm_matrix[nu, nu_minus] = (

267 glm_minus * rigth_coef * temp_coef

268 )

269

270 else:

271 if mu == 0:

272 pass

273 else:

274 temp_coef = -1 / (2 * mu) * (lu - 1) * (lu + 2)

275

276 nu = lumu_dict[(lu, mu, "cos")]

277 if (lu - 1, mu, "sin") in lumu_dict:

278 nu_minus = lumu_dict[(lu - 1, mu, "sin")]

279

280 glm_minus = calculate_glm_hat(lu, mu)

281

282 if (lu - 1, mu, "sin") not in lumu_dict:

283 pass

284 else:

285 clm_to_tlm_matrix[nu, nu_minus] = (

286 -glm_minus * rigth_coef * temp_coef

287 )

288

289 nu = lumu_dict[(lu, mu, "sin")]

290 if mu <= lu - 1:

291 nu_minus = lumu_dict[(lu - 1, mu, "cos")]

292

293 if (lu - 1, mu, "cos") not in lumu_dict:

294 pass

295 else:

296 clm_to_tlm_matrix[nu, nu_minus] = (

297 glm_minus * rigth_coef * temp_coef

298 )

299

300 return clm_to_tlm_matrix

301

302

303# =============================================

304# Building matrix corresponding tlm = T qlm

305# for zonal part ql0 to tl0

306# =============================================

307

308

309def calculation_ql0_to_tl0_zonal_coefs(n):

310 def calc_al(l):

311 return (-1) * (2 * l + 1) * (2 * l + 2) / (4 * l + 3)

312

313 def calc_bl(l):

314 return (2 * l + 2) - (((2 * l + 2) ** 2) / (4 * l + 3))

315

316 leftPartPlus = calc_al(n)

317 if n > 0:

318 leftPartMinus = calc_bl(n - 1)

319 else:

320 leftPartMinus = 0

321

322 def calc_cl(l):

323 return ((l * (l + 1) * (l - 1)) / ((2 * l + 1) * (2 * l - 1))) - (

324 (3 * l * (l - 1)) / ((2 * l + 1) * (2 * l - 1))

325 )

326

327 def calc_dl(l):

328 return (

329 (((l + 1) ** 3) / ((2 * l + 1) * (2 * l + 3)))

330 + (((l**2) * (l + 1)) / ((2 * l + 1) * (2 * l - 1)))

331 - (((l + 1) ** 2) / (2 * l + 3))

332 + 3

333 - ((3 * ((l + 1) ** 2)) / ((2 * l + 1) * (2 * l + 3)))

334 - ((3 * (l**2)) / ((2 * l + 1) * (2 * l - 1)))

335 )

336

337 def calc_el(l):

338 return (

339 ((((l + 1) ** 2) * (l + 2)) / ((2 * l + 1) * (2 * l + 3)))

340 - (((l + 1) * (l + 2)) / (2 * l + 3))

341 - ((3 * (l + 1) * (l + 2)) / ((2 * l + 1) * (2 * l + 3)))

342 )

343

344 rightPartMin2 = calc_cl(2 * n + 2)

345 rightPartZero = calc_dl(2 * n)

346 if n > 0:

347 rightPartPlus2 = calc_el(2 * n - 2)

348 else:

349 rightPartPlus2 = 0

350

351 return leftPartPlus, leftPartMinus, rightPartMin2, rightPartZero, rightPartPlus2

352

353

354def build_qlm_to_tlm_zonal_matrix(Lmax_q):

355 Nq = Lmax_q * (Lmax_q + 2) + 1

356 Nu = (Lmax_q + 1) * (Lmax_q + 1 + 2)

357 matrix = np.zeros((Nu, Nq))

358 # gauss_theta, gauss_weights = gaussPoints(0, np.pi, 64)

359

360 u_dict = {}

361 nu = 0

362 for lu in range(1, (Lmax_q + 1) + 1):

363 for mu in range(0, lu + 1):

364 u_dict[lu, mu, "cos"] = nu

365 nu += 1

366 if mu != 0:

367 u_dict[lu, mu, "sin"] = nu

368 nu += 1

369

370 q_dict = {}

371 nq = 0

372 for lq in range(0, Lmax_q + 1):

373 for mq in range(0, lq + 1):

374 q_dict[lq, mq, "cos"] = nq

375 nq += 1

376 if mq != 0:

377 q_dict[lq, mq, "sin"] = nq

378 nq += 1

379

380 tprev = 0

381 for n in range(0, Lmax_q // 2 + 1):

382 a, b, c, d, e = calculation_ql0_to_tl0_zonal_coefs(n)

383 if n == 0:

384 t2np = u_dict[2 * n + 1, 0, "cos"]

385 q2np = q_dict[2 * n + 2, 0, "cos"]

386 q2n = q_dict[2 * n, 0, "cos"]

387 matrix[t2np, q2np] = c / a

388 matrix[t2np, q2n] = d / a

389 tprev = t2np

390 else:

391 t2np = u_dict[2 * n + 1, 0, "cos"]

392 if 2 * n + 2 <= Lmax_q:

393 q2np = q_dict[2 * n + 2, 0, "cos"]

394 q2n = q_dict[2 * n, 0, "cos"]

395 q2nm = q_dict[2 * n - 2, 0, "cos"]

396 if 2 * n + 2 <= Lmax_q:

397 matrix[t2np, q2np] = c / a

398 matrix[t2np, q2n] = d / a

399 matrix[t2np, q2nm] = e / a

400 matrix[t2np, :] -= matrix[tprev, :] * b / a

401 tprev = t2np

403 return matrix

407# =============================================

408# calculation of full normalisation coefficient

409# =============================================

410

411

412def norm_coef_Sph(l, m):

413 # return np.sqrt(2 * l + 1) * np.sqrt(np.math.factorial(l - m) / np.math.factorial(l + m))

414 return (

415 np.sqrt(2 * l + 1)

416 * np.sqrt(np.math.factorial(l - m) / np.math.factorial(l + m))

417 * ((-1) ** m)

418 )

419

420

421# =============================================

422# calculation of fully normalised legendre poly

423# Plm hat as in notes

424# =============================================

425

426

427def Plm_hat(l, m, theta):

428 return norm_coef_Sph(l, m) * assocLegendrePoly(l, m, theta)

429

430

431# =============================================

432# calculation of magnetic field data at thetas

433# and phis using generalised

434# spherical harmonics

435# =============================================

436

437

438def calculate_Br_from_ReIm(LBr, thetas, phis, Br_glm_r_i, **kwargs):

439 if "normFunc" in kwargs:

440 normFunc = kwargs["normFunc"]

441 else:

442

443 def normFunc(*args):

444 return 1

445

446 Br_data_matrix = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

447 Br_coef = []

448 for lbr in range(1, LBr + 1):

449 for mbr in range(-lbr, lbr + 1):

450 Br_coef.append(normFunc(lbr, mbr))

451 Br_coef.append(Br_coef[-1])

452

453 Br_coef = np.array(Br_coef)

454

455 Br_data_temp_theta = []

456 for theta_num in thetas:

457 Br_temp_theta = []

458 nbr = 0

459 for lbr in range(1, LBr + 1):

460 for mbr in range(-lbr, lbr + 1):

461 Br_temp_theta.append(Plm_hat(lbr, mbr, theta_num))

462 nbr += 1

463 Br_temp_theta.append(Br_temp_theta[-1])

464 nbr += 1

465

466 Br_data_temp_theta.append(np.array(Br_temp_theta))

467

468 Br_data_temp_phi = []

469 for phi_num in phis:

470 Br_temp_phi = []

471 nbr = 0

472 for lbr in range(1, LBr + 1):

473 for mbr in range(-lbr, lbr + 1):

474 Br_temp_phi.append(np.exp(1j * mbr * phi_num))

475 nbr += 1

476 Br_temp_phi.append(1j * np.exp(1j * mbr * phi_num))

477 nbr += 1

478

479 Br_data_temp_phi.append(np.array(Br_temp_phi))

480

481 for tn in range(len(thetas)):

482 for pn in range(len(phis)):

483 Br_data_matrix[tn, pn] = (

484 Br_coef

485 * Br_data_temp_theta[tn]

486 * Br_data_temp_phi[pn]

487 * Br_glm_r_i.reshape(

488 -1,

489 )

490 ).sum()

491

492 return Br_data_matrix

493

494

495def calculate_rdBr_dr_from_ReIm(LBr, thetas, phis, Br_glm_r_i, **kwargs):

496 if "normFunc" in kwargs:

497 normFunc = kwargs["normFunc"]

498 else:

499

500 def normFunc(*args):

501 return 1

502

503 Br_data_matrix = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

504 Br_coef = []

505 for lbr in range(1, LBr + 1):

506 for mbr in range(-lbr, lbr + 1):

507 Br_coef.append(normFunc(lbr, mbr))

508 Br_coef.append(Br_coef[-1])

509

510 Br_coef = np.array(Br_coef)

511

512 Br_data_temp_theta = []

513 for theta_num in thetas:

514 Br_temp_theta = []

515 nbr = 0

516 for lbr in range(1, LBr + 1):

517 for mbr in range(-lbr, lbr + 1):

518 Br_temp_theta.append(Plm_hat(lbr, mbr, theta_num))

519 nbr += 1

520 Br_temp_theta.append(Br_temp_theta[-1])

521 nbr += 1

522

523 Br_data_temp_theta.append(np.array(Br_temp_theta))

524

525 Br_data_temp_phi = []

526 for phi_num in phis:

527 Br_temp_phi = []

528 nbr = 0

529 for lbr in range(1, LBr + 1):

530 for mbr in range(-lbr, lbr + 1):

531 Br_temp_phi.append(np.exp(1j * mbr * phi_num))

532 nbr += 1

533 Br_temp_phi.append(1j * np.exp(1j * mbr * phi_num))

534 nbr += 1

535

536 Br_data_temp_phi.append(np.array(Br_temp_phi))

537

538 for tn in range(len(thetas)):

539 for pn in range(len(phis)):

540 Br_data_matrix[tn, pn] = (

541 Br_coef

542 * Br_data_temp_theta[tn]

543 * Br_data_temp_phi[pn]

544 * Br_glm_r_i.reshape(

545 -1,

546 )

547 ).sum()

548

549 return Br_data_matrix

550

551

552def calculate_b_plus_from_ReIm(LBr, thetas, phis, b_coeff, **kwargs):

553 if "normFunc" in kwargs:

554 normFunc = kwargs["normFunc"]

555 else:

556

557 def normFunc(*args):

558 return 1

559

560 Br_data_matrix = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

561 Br_coef = []

562 for lbr in range(1, LBr + 1):

563 for mbr in range(-lbr, lbr + 1):

564 Br_coef.append(normFunc(lbr, mbr))

565 Br_coef.append(Br_coef[-1])

566

567 Br_coef = np.array(Br_coef)

568

569 Br_data_temp_theta = []

570 for theta_num in thetas:

571 Br_temp_theta = []

572 nbr = 0

573 for lbr in range(1, LBr + 1):

574 for mbr in range(-lbr, lbr + 1):

575 Br_temp_theta.append(

576 generalised_Plm_plus(l=lbr, m=mbr, theta=theta_num)

577 )

578 nbr += 1

579 Br_temp_theta.append(Br_temp_theta[-1])

580 nbr += 1

581

582 Br_data_temp_theta.append(np.array(Br_temp_theta))

583

584 Br_data_temp_phi = []

585 for phi_num in phis:

586 Br_temp_phi = []

587 nbr = 0

588 for lbr in range(1, LBr + 1):

589 for mbr in range(-lbr, lbr + 1):

590 Br_temp_phi.append(np.exp(1j * mbr * phi_num))

591 nbr += 1

592 Br_temp_phi.append(1j * np.exp(1j * mbr * phi_num))

593 nbr += 1

594

595 Br_data_temp_phi.append(np.array(Br_temp_phi))

596

597 for tn in range(len(thetas)):

598 for pn in range(len(phis)):

599 Br_data_matrix[tn, pn] = (

600 Br_coef

601 * Br_data_temp_theta[tn]

602 * Br_data_temp_phi[pn]

603 * b_coeff.reshape(

604 -1,

605 )

606 ).sum()

607

608 return Br_data_matrix

609

610

611def calculate_b_minus_from_ReIm(LBr, thetas, phis, b_coeff, **kwargs):

612 if "normFunc" in kwargs:

613 normFunc = kwargs["normFunc"]

614 else:

615

616 def normFunc(*args):

617 return 1

618

619 Br_data_matrix = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

620 Br_coef = []

621 for lbr in range(1, LBr + 1):

622 for mbr in range(-lbr, lbr + 1):

623 Br_coef.append(normFunc(lbr, mbr))

624 Br_coef.append(Br_coef[-1])

625

626 Br_coef = np.array(Br_coef)

627

628 Br_data_temp_theta = []

629 for theta_num in thetas:

630 Br_temp_theta = []

631 nbr = 0

632 for lbr in range(1, LBr + 1):

633 for mbr in range(-lbr, lbr + 1):

634 Br_temp_theta.append(

635 generalised_Plm_minus(l=lbr, m=mbr, theta=theta_num)

636 )

637 nbr += 1

638 Br_temp_theta.append(Br_temp_theta[-1])

639 nbr += 1

640

641 Br_data_temp_theta.append(np.array(Br_temp_theta))

642

643 Br_data_temp_phi = []

644 for phi_num in phis:

645 Br_temp_phi = []

646 nbr = 0

647 for lbr in range(1, LBr + 1):

648 for mbr in range(-lbr, lbr + 1):

649 Br_temp_phi.append(np.exp(1j * mbr * phi_num))

650 nbr += 1

651 Br_temp_phi.append(1j * np.exp(1j * mbr * phi_num))

652 nbr += 1

653

654 Br_data_temp_phi.append(np.array(Br_temp_phi))

655

656 for tn in range(len(thetas)):

657 for pn in range(len(phis)):

658 Br_data_matrix[tn, pn] = (

659 Br_coef

660 * Br_data_temp_theta[tn]

661 * Br_data_temp_phi[pn]

662 * b_coeff.reshape(

663 -1,

664 )

665 ).sum()

666

667 return Br_data_matrix

668

669

670# =============================================

671# calculation of additional information for

672# cos sin and Re Im representation:

673# - number of elements

674# - position of each l and m element

675# =============================================

676

677

678def build_real_imag_and_cossin_dict_pos(L, q=False):

679 x_in_real_imag_form_dict = {}

680 Nx_real_imag_form = 0

681 start_l = 1

682 if q:

683 start_l = 0

684 for lx in range(start_l, L + 1):

685 for mx in range(-lx, lx + 1):

686 x_in_real_imag_form_dict[Nx_real_imag_form] = (lx, mx, "r")

687 Nx_real_imag_form += 1

688 x_in_real_imag_form_dict[Nx_real_imag_form] = (lx, mx, "i")

689 Nx_real_imag_form += 1

690

691 x_cos_sin_dict = {}

692 Nx_cossin_form = 0

693 for lx in range(start_l, L + 1):

694 for mx in range(0, lx + 1):

695 x_cos_sin_dict[Nx_cossin_form] = (lx, mx, "cos")

696 Nx_cossin_form += 1

697 if mx != 0:

698 x_cos_sin_dict[Nx_cossin_form] = (lx, mx, "sin")

699 Nx_cossin_form += 1

700

701 return x_in_real_imag_form_dict, Nx_real_imag_form, x_cos_sin_dict, Nx_cossin_form

702

703

704# =============================================

705# building transformation matrix to calculate

706# Re Im form from cos sin

707# =============================================

708

709

710def building_transformation_matrix_from_cossin_to_R_I_form_B(

711 Nx_real_imag_form, Nx_cossin_form, x_in_real_imag_form_dict, x_cos_sin_dict

712):

713 x_Matrix_from_cossin_to_R_I_form = np.zeros((Nx_real_imag_form, Nx_cossin_form))

714

715 for nx_r_i in range(Nx_real_imag_form):

716 lx, mx, r_i = x_in_real_imag_form_dict[nx_r_i]

717 nx_cs = 0

718

719 if mx == 0:

720 if r_i == "r":

721 for pos, deg in x_cos_sin_dict.items():

722 if (lx, mx, "cos") == deg:

723 nx_cs = pos

724 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (lx + 1) / np.sqrt(

725 2 * lx + 1

726 )

727

728 if mx > 0:

729 if r_i == "r":

730 for pos, deg in x_cos_sin_dict.items():

731 if (lx, mx, "cos") == deg:

732 nx_cs = pos

733 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

734 np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1) / 2

735 )

736 if r_i == "i":

737 for pos, deg in x_cos_sin_dict.items():

738 if (lx, mx, "sin") == deg:

739 nx_cs = pos

740 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

741 -np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1) / 2

742 )

743

744 if mx < 0:

745 if r_i == "r":

746 for pos, deg in x_cos_sin_dict.items():

747 if (lx, -mx, "cos") == deg:

748 nx_cs = pos

749 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

750 np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1) * ((-1) ** mx) / 2

751 )

752 if r_i == "i":

753 for pos, deg in x_cos_sin_dict.items():

754 if (lx, -mx, "sin") == deg:

755 nx_cs = pos

756 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

757 np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1) * ((-1) ** mx) / 2

758 )

759

760 return x_Matrix_from_cossin_to_R_I_form

761

762

763def building_transformation_matrix_from_cossin_to_R_I_form(

764 Nx_real_imag_form, Nx_cossin_form, x_in_real_imag_form_dict, x_cos_sin_dict

765):

766 x_Matrix_from_cossin_to_R_I_form = np.zeros((Nx_real_imag_form, Nx_cossin_form))

767

768 for nx_r_i in range(Nx_real_imag_form):

769 lx, mx, r_i = x_in_real_imag_form_dict[nx_r_i]

770 nx_cs = 0

771

772 if mx == 0:

773 if r_i == "r":

774 for pos, deg in x_cos_sin_dict.items():

775 if (lx, mx, "cos") == deg:

776 nx_cs = pos

777 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = 1 / np.sqrt(

778 2 * lx + 1

779 )

780

781 if mx > 0:

782 if r_i == "r":

783 for pos, deg in x_cos_sin_dict.items():

784 if (lx, mx, "cos") == deg:

785 nx_cs = pos

786 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

787 np.sqrt(2) / np.sqrt(2 * lx + 1) / 2

788 )

789 if r_i == "i":

790 for pos, deg in x_cos_sin_dict.items():

791 if (lx, mx, "sin") == deg:

792 nx_cs = pos

793 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

794 -np.sqrt(2) / np.sqrt(2 * lx + 1) / 2

795 )

796

797 if mx < 0:

798 if r_i == "r":

799 for pos, deg in x_cos_sin_dict.items():

800 if (lx, -mx, "cos") == deg:

801 nx_cs = pos

802 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

803 np.sqrt(2) / np.sqrt(2 * lx + 1) * ((-1) ** mx) / 2

804 )

805 if r_i == "i":

806 for pos, deg in x_cos_sin_dict.items():

807 if (lx, -mx, "sin") == deg:

808 nx_cs = pos

809 x_Matrix_from_cossin_to_R_I_form[nx_r_i, nx_cs] = (

810 np.sqrt(2) / np.sqrt(2 * lx + 1) * ((-1) ** mx) / 2

811 )

812

813 return x_Matrix_from_cossin_to_R_I_form

814

815

816# =============================================

817# building transformation matrix to calculate

818# cos sin form from Re Im

819# =============================================

820

821

822def building_transformation_matrix_from_R_I_to_cossin_form_B(

823 Nx_real_imag_form, Nx_cossin_form, x_in_real_imag_form_dict, x_cos_sin_dict

824):

825 x_Matrix_from_R_I_to_cossin_form = np.zeros((Nx_cossin_form, Nx_real_imag_form))

826

827 for nx_r_i in range(Nx_real_imag_form):

828 lx, mx, r_i = x_in_real_imag_form_dict[nx_r_i]

829 nx_cs = 0

830

831 if mx == 0:

832 if r_i == "r":

833 for pos, deg in x_cos_sin_dict.items():

834 if (lx, mx, "cos") == deg:

835 nx_cs = pos

836 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

837 (lx + 1) / np.sqrt(2 * lx + 1)

838 )

839

840 if mx > 0:

841 if r_i == "r":

842 for pos, deg in x_cos_sin_dict.items():

843 if (lx, mx, "cos") == deg:

844 nx_cs = pos

845 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

846 np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1)

847 )

848 if r_i == "i":

849 for pos, deg in x_cos_sin_dict.items():

850 if (lx, mx, "sin") == deg:

851 nx_cs = pos

852 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

853 -np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1)

854 )

855

856 if mx < 0:

857 if r_i == "r":

858 for pos, deg in x_cos_sin_dict.items():

859 if (lx, -mx, "cos") == deg:

860 nx_cs = pos

861 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

862 np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1) * ((-1) ** mx)

863 )

864 if r_i == "i":

865 for pos, deg in x_cos_sin_dict.items():

866 if (lx, -mx, "sin") == deg:

867 nx_cs = pos

868 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

869 np.sqrt(2) * (lx + 1) / np.sqrt(2 * lx + 1) * ((-1) ** mx)

870 )

871

872 return x_Matrix_from_R_I_to_cossin_form

873

874

875def building_transformation_matrix_from_R_I_to_cossin_form(

876 Nx_real_imag_form, Nx_cossin_form, x_in_real_imag_form_dict, x_cos_sin_dict

877):

878 x_Matrix_from_R_I_to_cossin_form = np.zeros((Nx_cossin_form, Nx_real_imag_form))

879

880 for nx_r_i in range(Nx_real_imag_form):

881 lx, mx, r_i = x_in_real_imag_form_dict[nx_r_i]

882 nx_cs = 0

883

884 if mx == 0:

885 if r_i == "r":

886 for pos, deg in x_cos_sin_dict.items():

887 if (lx, mx, "cos") == deg:

888 nx_cs = pos

889 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

890 1 / np.sqrt(2 * lx + 1)

891 )

892

893 if mx > 0:

894 if r_i == "r":

895 for pos, deg in x_cos_sin_dict.items():

896 if (lx, mx, "cos") == deg:

897 nx_cs = pos

898 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

899 np.sqrt(2) / np.sqrt(2 * lx + 1)

900 )

901 if r_i == "i":

902 for pos, deg in x_cos_sin_dict.items():

903 if (lx, mx, "sin") == deg:

904 nx_cs = pos

905 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

906 -np.sqrt(2) / np.sqrt(2 * lx + 1)

907 )

908

909 if mx < 0:

910 if r_i == "r":

911 for pos, deg in x_cos_sin_dict.items():

912 if (lx, -mx, "cos") == deg:

913 nx_cs = pos

914 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

915 np.sqrt(2) / np.sqrt(2 * lx + 1) * ((-1) ** mx)

916 )

917 if r_i == "i":

918 for pos, deg in x_cos_sin_dict.items():

919 if (lx, -mx, "sin") == deg:

920 nx_cs = pos

921 x_Matrix_from_R_I_to_cossin_form[nx_cs, nx_r_i] = 1 / (

922 np.sqrt(2) / np.sqrt(2 * lx + 1) * ((-1) ** mx)

923 )

924

925 return x_Matrix_from_R_I_to_cossin_form

926

927

928# =============================================

929# Precalculate Polinimials

930# =============================================

931

932

933def precalculatePoly(L, thetas, PolyFunc, **kwargs):

934 if "normFunc" in kwargs:

935 normFunc = kwargs["normFunc"]

936 else:

937

938 def normFunc(*args):

939 return 1

940

941 res = {}

942 for l in range(1, L + 1):

943 for m in range(-l, l + 1):

944 for tn in range(len(thetas)):

945 res[(l, m, thetas[tn])] = PolyFunc(

946 l=l, m=m, theta=thetas[tn]

947 ) * normFunc(l, m)

948

949 return res

950

951

952# =============================================

953# Projecting on spherical harmonics for

954# Re Im form of coefficients

955# =============================================

956

957

958def projectionFFTCoreReIm(

959 gauss_theta, gauss_weights, input_matrix, Lmax, tmax, **kwargs

960):

961 if "precalculatedPolyDict" in kwargs:

962 precalculatedPolyDict = kwargs["precalculatedPolyDict"]

963 else:

964 if "PolyFunc" in kwargs:

965 PolyFunc = kwargs["PolyFunc"]

966

967 if "normFunc" in kwargs:

968 normFunc = kwargs["normFunc"]

969 else:

970

971 def normFunc(*args):

972 return 1

973

974 precalculatedPolyDict = precalculatePoly(

975 Lmax, gauss_theta, PolyFunc, normFunc

976 )

977 else:

978 raise Exception("No Polynomial")

979

980 pmax = 2 * tmax

981

982 lsvmsv_dict = {}

983

984 nsv = 0

985 for lsv in range(1, Lmax + 1):

986 for msv in range(-lsv, lsv + 1):

987 lsvmsv_dict[nsv] = (lsv, msv, "r")

988 nsv += 1

989 lsvmsv_dict[nsv] = (lsv, msv, "i")

990 nsv += 1

991

992 N = nsv

993

994 return_matrix = np.zeros((N,))

995

996 for tn in range(len(gauss_theta)):

997 theta = gauss_theta[tn]

998 weight = gauss_weights[tn]

999

1000 ftf = np.fft.fft(input_matrix[tn, :]) / input_matrix.shape[1]

1001

1002 LM = 0

1003 for l in range(1, Lmax + 1):

1004

1005 for m in range(-l, l + 1):

1006 return_matrix[LM] += (

1007 precalculatedPolyDict[(l, m, gauss_theta[tn])]

1008 * ftf[m].real

1009 * weight

1010 )

1011

1012 LM += 1

1013

1014 return_matrix[LM] += (

1015 precalculatedPolyDict[(l, m, gauss_theta[tn])]

1016 * ftf[m].imag

1017 * weight

1018 )

1019

1020 LM += 1

1021

1022 return_matrix = return_matrix / 2

1023

1024 return return_matrix

1025

1026

1027def projectionFFTCoreToEarthReIm(

1028 gauss_theta, gauss_weights, input_matrix, Lmax, tmax, **kwargs

1029):

1030 if "precalculatedPolyDict" in kwargs:

1031 precalculatedPolyDict = kwargs["precalculatedPolyDict"]

1032 else:

1033 if "PolyFunc" in kwargs:

1034 PolyFunc = kwargs["PolyFunc"]

1035

1036 if "normFunc" in kwargs:

1037 normFunc = kwargs["normFunc"]

1038 else:

1039

1040 def normFunc(*args):

1041 return 1

1042

1043 precalculatedPolyDict = precalculatePoly(

1044 Lmax, gauss_theta, PolyFunc, normFunc

1045 )

1046 else:

1047 raise Exception("No Polynomial")

1048

1049 pmax = 2 * tmax

1050

1051 lsvmsv_dict = {}

1052

1053 nsv = 0

1054 for lsv in range(1, Lmax + 1):

1055 for msv in range(-lsv, lsv + 1):

1056 lsvmsv_dict[nsv] = (lsv, msv, "r")

1057 nsv += 1

1058 lsvmsv_dict[nsv] = (lsv, msv, "i")

1059 nsv += 1

1060

1061 N = nsv

1062

1063 return_matrix = np.zeros((N,))

1064

1065 for tn in range(len(gauss_theta)):

1066 theta = gauss_theta[tn]

1067 weight = gauss_weights[tn]

1068

1069 ftf = np.fft.fft(input_matrix[tn, :]) / input_matrix.shape[1]

1070

1071 LM = 0

1072 for l in range(1, Lmax + 1):

1073

1074 for m in range(-l, l + 1):

1075 temp = precalculatedPolyDict[(l, m, theta)] * weight

1076

1077 return_matrix[LM] += temp * ftf[m].real

1078

1079 LM += 1

1080

1081 return_matrix[LM] += temp * ftf[m].imag

1082

1083 LM += 1

1084

1085 n = 0

1086 for l in range(1, Lmax + 1):

1087 for m in range(-l, l + 1):

1088 return_matrix[n] = return_matrix[n] / ((Ra / Rc) ** (l + 2))

1089 n += 1

1090 return_matrix[n] = return_matrix[n] / ((Ra / Rc) ** (l + 2))

1091 n += 1

1092

1093 return_matrix = return_matrix / 2

1094

1095 return return_matrix

1096

1097

1098# =============================================

1099# generalised spherical harmonics plus

1100# =============================================

1101

1102

1103def generalised_Plm_plus(theta, l, m):

1104 ##############################################

1105 # Generalised l=0 m=0 not exists at the moment

1106 ##############################################

1107 if l != 0:

1108 norm = -np.sqrt(l * (l + 1))

1109 hat_norm = norm_coef_Sph(l, m)

1110 temp = (derivativeThetaAssocLegendrePoly(l, m, theta) * hat_norm) + (

1111 (m / np.sin(theta)) * Plm_hat(l, m, theta)

1112 )

1113 return temp / norm

1114 else:

1115 return 0

1116

1117

1118# =============================================

1119# generalised spherical harmonics minus

1120# =============================================

1121

1122

1123def generalised_Plm_minus(theta, l, m):

1124 ##############################################

1125 # Generalised l=0 m=0 not exists at the moment

1126 ##############################################

1127 if l != 0:

1128 norm = np.sqrt(l * (l + 1))

1129 hat_norm = norm_coef_Sph(l, m)

1130 temp = (derivativeThetaAssocLegendrePoly(l, m, theta) * hat_norm) - (

1131 (m / np.sin(theta)) * Plm_hat(l, m, theta)

1132 )

1133 return temp / norm

1134 else:

1135 return 0

1136

1137

1138# =============================================

1139# matrix for multiplication Re Im form by 1i

1140# =============================================

1141

1142

1143def building_transformation_matrix_multiply_by_i(

1144 Nx_real_imag_form, x_in_real_imag_form_dict

1145):

1146 matrix_multiply_by_i = np.zeros((Nx_real_imag_form, Nx_real_imag_form))

1147

1148 for nx_r_i in range(Nx_real_imag_form):

1149 lx, mx, r_i = x_in_real_imag_form_dict[nx_r_i]

1150

1151 if r_i == "r":

1152 matrix_multiply_by_i[nx_r_i + 1, nx_r_i] = 1

1153 if r_i == "i":

1154 matrix_multiply_by_i[nx_r_i - 1, nx_r_i] = -1

1155

1156 return matrix_multiply_by_i

1157

1158

1159# =============================================

1160# matrix for calculating u from Re Im form

1161# =============================================

1162

1163

1164def calculate_u_for_Re_Im(Lu, thetas, phis, u_r_i, Plm_func):

1165 u_Re_Im = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

1166 u_data_temp_theta = []

1167 nt = 0

1168 for theta_num in thetas:

1169 u_temp_theta = []

1170 nu = 0

1171 for lu in range(1, Lu + 1):

1172 for mu in range(-lu, lu + 1):

1173 # u_temp_theta.append(Plm_func(l=lu, m=mu, theta=theta_num) * gauss_weights[nt])

1174 u_temp_theta.append(Plm_func(l=lu, m=mu, theta=theta_num))

1175 u_temp_theta.append(u_temp_theta[-1])

1176 nu += 1

1177 nt += 1

1178

1179 u_data_temp_theta.append(np.array(u_temp_theta))

1180

1181 u_data_temp_phi = []

1182 for phi_num in phis:

1183 u_temp_phi = []

1184 nbr = 0

1185 for lu in range(1, Lu + 1):

1186 for mu in range(-lu, lu + 1):

1187 u_temp_phi.append(np.exp(1j * mu * phi_num))

1188 nbr += 1

1189 u_temp_phi.append(1j * np.exp(1j * mu * phi_num))

1190 nbr += 1

1191

1192 u_data_temp_phi.append(np.array(u_temp_phi))

1193

1194 for tn in range(len(thetas)):

1195 for pn in range(len(phis)):

1196 u_Re_Im[tn, pn] = (

1197 u_data_temp_theta[tn]

1198 * u_data_temp_phi[pn]

1199 * u_r_i.reshape(

1200 -1,

1201 )

1202 ).sum()

1203

1204 return u_Re_Im

1205

1206

1207#########################################

1208# Calculating ulm_plus and ulm_minus

1209# forward matrices

1210#########################################

1211

1212

1213def calculating_ulm_plus_minus_from_qlm(

1214 Lu,

1215 Nu_real_imag_form,

1216 QtC_matrix_real_imag_form,

1217 QtT_matrix_real_imag_form,

1218 Matrix_multiply_by_i,

1219):

1220 #########################################

1221 # Building norm matrix

1222 #########################################

1223

1224 func_l = []

1225 for lu in range(1, Lu + 1):

1226 for mu in range(-lu, lu + 1):

1227 func_l.append(np.sqrt(lu * (lu + 1)) / np.sqrt(2))

1228 func_l.append(np.sqrt(lu * (lu + 1)) / np.sqrt(2))

1229 func_l = np.array(func_l)

1230

1231 #########################################

1232 # Building forward matrices that

1233 # calculate ulm+ and ulm- from qlm

1234 #########################################

1235

1236 ulm_plus_matrix = (

1237 QtC_matrix_real_imag_form - Matrix_multiply_by_i.dot(QtT_matrix_real_imag_form)

1238 ) * func_l.reshape(Nu_real_imag_form, 1)

1239 ulm_minus_matrix = (

1240 QtC_matrix_real_imag_form + Matrix_multiply_by_i.dot(QtT_matrix_real_imag_form)

1241 ) * func_l.reshape(Nu_real_imag_form, 1)

1242

1243 return ulm_plus_matrix, ulm_minus_matrix

1244

1245

1246#########################################

1247# calculating u+ and u-

1248#########################################

1249

1250

1251def calculate_u_data_matrices(

1252 Nu_real_imag_form, u_in_real_imag_form_dict, thetas, phis, debug=False, **kwargs

1253):

1254 if "u_plus_from_qlm_matrix" in kwargs:

1255 u_plus_from_qlm_matrix = kwargs["u_plus_from_qlm_matrix"]

1256 u_minus_from_qlm_matrix = kwargs["u_minus_from_qlm_matrix"]

1257 Nu = u_minus_from_qlm_matrix.shape[1]

1258 else:

1259 u_plus_from_qlm_matrix = np.eye(Nu_real_imag_form)

1260 u_minus_from_qlm_matrix = np.eye(Nu_real_imag_form)

1261 Nu = Nu_real_imag_form

1262

1263 u_plus_data_matrices = []

1264 u_minus_data_matrices = []

1265

1266 u_plus_forward_matrix = np.zeros(

1267 (len(thetas) * len(phis), Nu_real_imag_form), dtype=np.complex128

1268 )

1269 u_minus_forward_matrix = np.zeros(

1270 (len(thetas) * len(phis), Nu_real_imag_form), dtype=np.complex128

1271 )

1272 ntp = 0

1273

1274 if debug:

1275 if "u_plus_forward_matrix" in kwargs:

1276 forward_type = kwargs["forward_type"]

1277 if forward_type == "complex":

1278 u_plus_forward_matrix_complex = kwargs["u_plus_forward_matrix"]

1279 u_minus_forward_matrix_complex = kwargs["u_minus_forward_matrix"]

1280

1281 for i in range(u_plus_forward_matrix_complex.shape[1]):

1282 u_plus_forward_matrix[:, 2 * i] = u_plus_forward_matrix_complex[

1283 :, i

1284 ]

1285 u_plus_forward_matrix[:, 2 * i + 1] = (

1286 1j * u_plus_forward_matrix_complex[:, i]

1287 )

1288

1289 u_minus_forward_matrix[:, 2 * i] = u_minus_forward_matrix_complex[

1290 :, i

1291 ]

1292 u_minus_forward_matrix[:, 2 * i + 1] = (

1293 1j * u_minus_forward_matrix_complex[:, i]

1294 )

1295

1296 else:

1297 u_plus_forward_matrix = kwargs["u_plus_forward_matrix"]

1298 u_minus_forward_matrix = kwargs["u_minus_forward_matrix"]

1299 else:

1300 for tn in range(len(thetas)):

1301 for pn in range(len(phis)):

1302 for nu in range(int(Nu_real_imag_form / 2)):

1303 lu, mu, r_i_u = u_in_real_imag_form_dict[2 * nu]

1304 u_plus_forward_matrix[ntp, 2 * nu] = generalised_Plm_plus(

1305 l=lu, m=mu, theta=thetas[tn]

1306 ) * np.exp(1j * mu * phis[pn])

1307 u_minus_forward_matrix[ntp, 2 * nu] = generalised_Plm_minus(

1308 l=lu, m=mu, theta=thetas[tn]

1309 ) * np.exp(1j * mu * phis[pn])

1310 u_plus_forward_matrix[ntp, 2 * nu + 1] = (

1311 1j * u_plus_forward_matrix[ntp, 2 * nu]

1312 )

1313 u_minus_forward_matrix[ntp, 2 * nu + 1] = (

1314 1j * u_minus_forward_matrix[ntp, 2 * nu]

1315 )

1316 ntp += 1

1317

1318 for nu in range(Nu):

1319 vector = np.zeros((Nu,))

1320 vector[nu] = 1

1321 u_plus_data_matrices.append(

1322 u_plus_forward_matrix.dot(u_plus_from_qlm_matrix.dot(vector)).reshape(

1323 (len(thetas), len(phis))

1324 )

1325 )

1326 u_minus_data_matrices.append(

1327 u_minus_forward_matrix.dot(u_minus_from_qlm_matrix.dot(vector)).reshape(

1328 (len(thetas), len(phis))

1329 )

1330 )

1331 else:

1332 if "u_plus_forward_matrix" in kwargs:

1333 forward_type = kwargs["forward_type"]

1334 if forward_type == "complex":

1335 u_plus_forward_matrix_complex = kwargs["u_plus_forward_matrix"]

1336 u_minus_forward_matrix_complex = kwargs["u_minus_forward_matrix"]

1337

1338 for i in range(u_plus_forward_matrix_complex.shape[1]):

1339 u_plus_forward_matrix[:, 2 * i] = u_plus_forward_matrix_complex[

1340 :, i

1341 ]

1342 u_plus_forward_matrix[:, 2 * i + 1] = (

1343 1j * u_plus_forward_matrix_complex[:, i]

1344 )

1345

1346 u_minus_forward_matrix[:, 2 * i] = u_minus_forward_matrix_complex[

1347 :, i

1348 ]

1349 u_minus_forward_matrix[:, 2 * i + 1] = (

1350 1j * u_minus_forward_matrix_complex[:, i]

1351 )

1352

1353 else:

1354 u_plus_forward_matrix = kwargs["u_plus_forward_matrix"]

1355 u_minus_forward_matrix = kwargs["u_minus_forward_matrix"]

1356 else:

1357 for tn in range(len(thetas)):

1358 for pn in range(len(phis)):

1359 for nu in range(int(Nu_real_imag_form / 2)):

1360 lu, mu, r_i_u = u_in_real_imag_form_dict[2 * nu]

1361 u_plus_forward_matrix[ntp, 2 * nu] = generalised_Plm_plus(

1362 l=lu, m=mu, theta=thetas[tn]

1363 ) * np.exp(1j * mu * phis[pn])

1364 u_minus_forward_matrix[ntp, 2 * nu] = generalised_Plm_minus(

1365 l=lu, m=mu, theta=thetas[tn]

1366 ) * np.exp(1j * mu * phis[pn])

1367 u_plus_forward_matrix[ntp, 2 * nu + 1] = (

1368 1j * u_plus_forward_matrix[ntp, 2 * nu]

1369 )

1370 u_minus_forward_matrix[ntp, 2 * nu + 1] = (

1371 1j * u_minus_forward_matrix[ntp, 2 * nu]

1372 )

1373 ntp += 1

1374

1375 for nu in range(Nu):

1376 vector = np.zeros((Nu,))

1377 vector[nu] = 1

1378 u_plus_data_matrices.append(

1379 u_plus_forward_matrix.dot(u_plus_from_qlm_matrix.dot(vector)).reshape(

1380 (len(thetas), len(phis))

1381 )

1382 )

1383 u_minus_data_matrices.append(

1384 u_minus_forward_matrix.dot(u_minus_from_qlm_matrix.dot(vector)).reshape(

1385 (len(thetas), len(phis))

1386 )

1387 )

1388

1389 return u_plus_data_matrices, u_minus_data_matrices

1390

1391

1392def calculate_u_zero_data_matrices(

1393 Nu_real_imag_form, u_in_real_imag_form_dict, thetas, phis, debug=False

1394):

1395 u_zero_data_matrices = []

1396

1397 if debug:

1398 for nu in range(Nu_real_imag_form):

1399 u_zero_data_matrix = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

1400

1401 lu, mu, r_i_u = u_in_real_imag_form_dict[nu]

1402

1403 norm = lu * (lu + 1)

1404

1405 u_zero_temp_theta = []

1406 for theta_num in thetas:

1407 u_zero_temp_theta.append(Plm_hat(l=lu, m=mu, theta=theta_num))

1408

1409 u_zero_temp_phi = []

1410 for phi_num in phis:

1411 if r_i_u == "r":

1412 u_zero_temp_phi.append(np.exp(1j * mu * phi_num))

1413 else:

1414 u_zero_temp_phi.append(1j * np.exp(1j * mu * phi_num))

1415

1416 for tn in range(len(thetas)):

1417 for pn in range(len(phis)):

1418 u_zero_data_matrix[tn, pn] = (

1419 u_zero_temp_theta[tn] * u_zero_temp_phi[pn] * norm

1420 ).sum()

1421

1422 u_zero_data_matrices.append(u_zero_data_matrix)

1423 else:

1424 for nu in range(Nu_real_imag_form):

1425 u_zero_data_matrix = np.zeros((len(thetas), len(phis)), dtype=np.complex128)

1426

1427 lu, mu, r_i_u = u_in_real_imag_form_dict[nu]

1428

1429 norm = lu * (lu + 1)

1430

1431 u_zero_temp_theta = []

1432 for theta_num in thetas:

1433 u_zero_temp_theta.append(Plm_hat(l=lu, m=mu, theta=theta_num))

1434

1435 u_zero_temp_phi = []

1436 for phi_num in phis:

1437 if r_i_u == "r":

1438 u_zero_temp_phi.append(np.exp(1j * mu * phi_num))

1439 else:

1440 u_zero_temp_phi.append(1j * np.exp(1j * mu * phi_num))

1441

1442 for tn in range(len(thetas)):

1443 for pn in range(len(phis)):

1444 u_zero_data_matrix[tn, pn] = (

1445 u_zero_temp_theta[tn] * u_zero_temp_phi[pn] * norm

1446 ).sum()

1447

1448 u_zero_data_matrices.append(u_zero_data_matrix)

1449

1450 return u_zero_data_matrices

1451

1452

1453#########################################

1454# calculating v+ and v-

1455# from u+ and u- via projecting

1456# using fft

1457#########################################

1458

1459

1460def calculate_v_plus_minus(

1461 LSV,

1462 Nu_real_imag_form,

1463 Br_data_matrix,

1464 u_plus_data_matrices,

1465 u_minus_data_matrices,

1466 gauss_thetas,

1467 gauss_weights,

1468 tmax,

1469 debug=False,

1470):

1471 (

1472 SV_in_real_imag_form_dict,

1473 Nsv_real_imag_form,

1474 SV_cos_sin_dict,

1475 Nsv_cossin_form,

1476 ) = build_real_imag_and_cossin_dict_pos(LSV)

1477

1478 SV_Matrix_from_cossin_to_R_I_form = (

1479 building_transformation_matrix_from_cossin_to_R_I_form_B(

1480 Nsv_real_imag_form,

1481 Nsv_cossin_form,

1482 SV_in_real_imag_form_dict,

1483 SV_cos_sin_dict,

1484 )

1485 )

1486 SV_Matrix_from_R_I_to_cossin_form = (

1487 building_transformation_matrix_from_R_I_to_cossin_form_B(

1488 Nsv_real_imag_form,

1489 Nsv_cossin_form,

1490 SV_in_real_imag_form_dict,

1491 SV_cos_sin_dict,

1492 )

1493 )

1494

1495 v_plus_matrices = np.zeros((Nsv_real_imag_form, Nu_real_imag_form))

1496 v_minus_matrices = np.zeros((Nsv_real_imag_form, Nu_real_imag_form))

1497

1498 def one_norm(*args):

1499 return 1

1500

1501 precalculatedPolyDictMinus = precalculatePoly(

1502 LSV, gauss_thetas, generalised_Plm_minus

1503 )

1504 precalculatedPolyDictPlus = precalculatePoly(

1505 LSV, gauss_thetas, generalised_Plm_plus

1506 )

1507

1508 if debug:

1509 for nu in range(Nu_real_imag_form):

1510 temp_plus = 1j * Br_data_matrix * u_plus_data_matrices[nu]

1511 temp_minus = -1j * Br_data_matrix * u_minus_data_matrices[nu]

1512

1513 projected_plus = projectionFFTCoreToEarthReIm(

1514 gauss_thetas,

1515 gauss_weights,

1516 temp_plus,

1517 LSV,

1518 tmax,

1519 precalculatedPolyDict=precalculatedPolyDictPlus,

1520 )

1521 projected_minus = projectionFFTCoreToEarthReIm(

1522 gauss_thetas,

1523 gauss_weights,

1524 temp_minus,

1525 LSV,

1526 tmax,

1527 precalculatedPolyDict=precalculatedPolyDictMinus,

1528 )

1529

1530 v_plus_matrices[:, nu] = projected_plus

1531 v_minus_matrices[:, nu] = projected_minus

1532 else:

1533 for nu in range(Nu_real_imag_form):

1534 temp_plus = 1j * Br_data_matrix * u_plus_data_matrices[nu]

1535 temp_minus = -1j * Br_data_matrix * u_minus_data_matrices[nu]

1536

1537 projected_plus = projectionFFTCoreToEarthReIm(

1538 gauss_thetas,

1539 gauss_weights,

1540 temp_plus,

1541 LSV,

1542 tmax,

1543 precalculatedPolyDict=precalculatedPolyDictPlus,

1544 )

1545 projected_minus = projectionFFTCoreToEarthReIm(

1546 gauss_thetas,

1547 gauss_weights,

1548 temp_minus,

1549 LSV,

1550 tmax,

1551 precalculatedPolyDict=precalculatedPolyDictMinus,

1552 )

1553

1554 v_plus_matrices[:, nu] = projected_plus

1555 v_minus_matrices[:, nu] = projected_minus

1556

1557 return (

1558 v_plus_matrices,

1559 v_minus_matrices,

1560 )

1561

1562

1563#########################################

1564# calculating dBr/dt from v+ and v-

1565#########################################

1566

1567

1568def calculating_SV_from_v_matrix(

1569 LSV,

1570 Nsv_real_imag_form,

1571 SV_in_real_imag_form_dict,

1572 v_plus_matrices_from_qlm,

1573 v_minus_matrices_from_qlm,

1574):

1575 norm_SV_Re_Im = np.ones((Nsv_real_imag_form, 1))

1576

1577 r_norm = 1

1578 nsv = 0

1579 for lsv in range(1, LSV + 1):

1580 temp = np.sqrt(lsv * (lsv + 1)) / np.sqrt(2) / r_norm

1581 for m in range(-lsv, lsv + 1):

1582 norm_SV_Re_Im[nsv] = temp

1583 norm_SV_Re_Im[nsv + 1] = temp

1584 nsv += 2

1585

1586 dBr_dt_matrix = v_plus_matrices_from_qlm - v_minus_matrices_from_qlm

1587

1588 dBr_dt_matrix = dBr_dt_matrix * norm_SV_Re_Im * (-1)

1589

1590 Matrix_multiply_by_i_SV = building_transformation_matrix_multiply_by_i(

1591 Nsv_real_imag_form, SV_in_real_imag_form_dict

1592 )

1593

1594 dBr_dt_matrix_from_qlm = Matrix_multiply_by_i_SV.dot(dBr_dt_matrix)

1595

1596 return dBr_dt_matrix_from_qlm

1597

1598

1599def calculating_SV_from_v_matrix_ulm(

1600 LSV,

1601 Nsv_real_imag_form,

1602 SV_in_real_imag_form_dict,

1603 v_plus_matrices_from_ulm,

1604 v_minus_matrices_from_ulm,

1605):

1606 norm_SV_Re_Im = np.ones((Nsv_real_imag_form, 1))

1607

1608 r_norm = 1

1609 nsv = 0

1610 for lsv in range(1, LSV + 1):

1611 temp = np.sqrt(lsv * (lsv + 1)) / np.sqrt(2) / r_norm

1612 for m in range(-lsv, lsv + 1):

1613 norm_SV_Re_Im[nsv] = temp

1614 norm_SV_Re_Im[nsv + 1] = temp

1615 nsv += 2

1616

1617 Nq = v_plus_matrices_from_ulm.shape[1]

1618 dBr_dt_matrix = np.zeros((Nsv_real_imag_form, Nq * 2))

1619 dBr_dt_matrix[:, :Nq] = -v_minus_matrices_from_ulm

1620 dBr_dt_matrix[:, Nq:] = v_plus_matrices_from_ulm

1621

1622 dBr_dt_matrix = dBr_dt_matrix * norm_SV_Re_Im * (-1)

1623

1624 Matrix_multiply_by_i_SV = building_transformation_matrix_multiply_by_i(

1625 Nsv_real_imag_form, SV_in_real_imag_form_dict

1626 )

1627

1628 dBr_dt_matrix_from_ulm = Matrix_multiply_by_i_SV.dot(dBr_dt_matrix)

1629

1630 return dBr_dt_matrix_from_ulm

1631

1632

1633##########################################################################################################

1634##########################################################################################################

1635

1636

1637def calculate_t_from_u(

1638 LSV,

1639 LBr,

1640 Nu_real_imag_form,

1641 gauss_thetas,

1642 gauss_weights,

1643 phis,

1644 Br_glm_re_im,

1645 u_plus_data_matrices,

1646 u_minus_data_matrices,

1647 tmax=64,

1648 debug=True,

1649):

1650 rdBr_glm_r_i = np.array(Br_glm_re_im)

1651

1652 n = 0

1653 for l_ in range(1, LBr + 1):

1654 for m_ in range(-l_, l_ + 1):

1655 rdBr_glm_r_i[n, 0] = rdBr_glm_r_i[n, 0] * (-1) * (l_ + 1)

1656 n += 1

1657 rdBr_glm_r_i[n, 0] = rdBr_glm_r_i[n, 0] * (-1) * (l_ + 1)

1658 n += 1

1659

1660 rdBr_dt_data_matrix = calculate_rdBr_dr_from_ReIm(

1661 LBr, gauss_thetas, phis, rdBr_glm_r_i

1662 )

1663

1664 (

1665 SV_in_real_imag_form_dict,

1666 Nsv_real_imag_form,

1667 SV_cos_sin_dict,

1668 Nsv_cossin_form,

1669 ) = build_real_imag_and_cossin_dict_pos(LSV)

1670

1671 t_plus_matrix = np.zeros((Nsv_real_imag_form, Nu_real_imag_form))

1672 t_minus_matrix = np.zeros((Nsv_real_imag_form, Nu_real_imag_form))

1673

1674 def one_norm(*args):

1675 return 1

1676

1677 precalculatedPolyDictMinus = precalculatePoly(

1678 LSV, gauss_thetas, generalised_Plm_minus

1679 )

1680 precalculatedPolyDictPlus = precalculatePoly(

1681 LSV, gauss_thetas, generalised_Plm_plus

1682 )

1683

1684 if debug:

1685 for nu in range(Nu_real_imag_form):

1686 temp_plus = 1j * rdBr_dt_data_matrix * u_plus_data_matrices[nu]

1687 temp_minus = -1j * rdBr_dt_data_matrix * u_minus_data_matrices[nu]

1688

1689

1690 projected_plus = projectionFFTCoreReIm(

1691 gauss_thetas,

1692 gauss_weights,

1693 temp_plus,

1694 LSV,

1695 tmax,

1696 precalculatedPolyDict=precalculatedPolyDictPlus,

1697 )

1698 projected_minus = projectionFFTCoreReIm(

1699 gauss_thetas,

1700 gauss_weights,

1701 temp_minus,

1702 LSV,

1703 tmax,

1704 precalculatedPolyDict=precalculatedPolyDictMinus,

1705 )

1706

1707 t_plus_matrix[:, nu] = projected_plus

1708 t_minus_matrix[:, nu] = projected_minus

1709 else:

1710 for nu in range(Nu_real_imag_form):

1711 temp_plus = 1j * rdBr_dt_data_matrix * u_plus_data_matrices[nu]

1712 temp_minus = -1j * rdBr_dt_data_matrix * u_minus_data_matrices[nu]

1713

1714 projected_plus = projectionFFTCoreReIm(

1715 gauss_thetas,

1716 gauss_weights,

1717 temp_plus,

1718 LSV,

1719 tmax,

1720 precalculatedPolyDict=precalculatedPolyDictPlus,

1721 )

1722 projected_minus = projectionFFTCoreReIm(

1723 gauss_thetas,

1724 gauss_weights,

1725 temp_minus,

1726 LSV,

1727 tmax,

1728 precalculatedPolyDict=precalculatedPolyDictMinus,

1729 )

1730

1731 t_plus_matrix[:, nu] = projected_plus

1732 t_minus_matrix[:, nu] = projected_minus

1733

1734 return t_plus_matrix, t_minus_matrix

1735

1736

1737def calculate_t_from_qlm(

1738 LSV,

1739 LBr,

1740 Nu_real_imag_form,

1741 gauss_thetas,

1742 gauss_weights,

1743 phis,

1744 Br_glm_re_im,

1745 u_plus_data_matrices,

1746 u_minus_data_matrices,

1747 ulm_plus_matrix_from_qlm,

1748 ulm_minus_matrix_from_qlm,

1749 tmax=64,

1750 debug=True,

1751):

1752 t_plus_matrix, t_minus_matrix = calculate_t_from_u(

1753 LSV,

1754 LBr,

1755 Nu_real_imag_form,

1756 gauss_thetas,

1757 gauss_weights,

1758 phis,

1759 Br_glm_re_im,

1760 u_plus_data_matrices,

1761 u_minus_data_matrices,

1762 tmax,

1763 debug,

1764 )

1765

1766 return t_plus_matrix.dot(ulm_plus_matrix_from_qlm), t_minus_matrix.dot(

1767 ulm_minus_matrix_from_qlm

1768 )

1769

1770

1771def calculate_v0_from_u_minus_plus(

1772 LSV,

1773 LBr,

1774 Nsv_real_imag_form,

1775 gauss_thetas,

1776 gauss_weights,

1777 phis,

1778 blm_re_im,

1779 u_plus_data_matrices,

1780 u_minus_data_matrices,

1781 tmax=64,

1782 debug=True,

1783):

1784 nb = 0

1785 for lb in range(1, LBr + 1):

1786 nrm = -np.sqrt(lb / (2 * (lb + 1)))

1787 for mb in range(-lb, lb + 1):

1788 blm_re_im[nb] = blm_re_im[nb] * nrm

1789 blm_re_im[nb + 1] = blm_re_im[nb + 1] * nrm

1790 nb += 2

1791

1792 b_plus_matrix = calculate_b_plus_from_ReIm(LBr, gauss_thetas, phis, blm_re_im)

1793 b_minus_matrix = calculate_b_minus_from_ReIm(LBr, gauss_thetas, phis, blm_re_im)

1794

1795 v_zero_forward_matrix = np.zeros(

1796 (Nsv_real_imag_form, len(u_plus_data_matrices) * 2)

1797 )

1798

1799 def one_norm(*args):

1800 return 1

1801

1802 precalculatedPolyDictHat = precalculatePoly(LSV, gauss_thetas, Plm_hat)

1803

1804 if debug:

1805 for nu in range(len(u_plus_data_matrices)):

1806 temp_1 = 1j * b_minus_matrix * u_plus_data_matrices[nu]

1807 temp_2 = -1j * b_plus_matrix * u_minus_data_matrices[nu]

1808

1809 projected_1 = projectionFFTCoreReIm(

1810 gauss_thetas,

1811 gauss_weights,

1812 temp_1,

1813 LSV,

1814 tmax,

1815 precalculatedPolyDict=precalculatedPolyDictHat,

1816 )

1817 projected_2 = projectionFFTCoreReIm(

1818 gauss_thetas,

1819 gauss_weights,

1820 temp_2,

1821 LSV,

1822 tmax,

1823 precalculatedPolyDict=precalculatedPolyDictHat,

1824 )

1825

1826 v_zero_forward_matrix[:, nu] = projected_2

1827 v_zero_forward_matrix[:, nu + len(u_plus_data_matrices)] = projected_1

1828 else:

1829 for nu in range(len(u_plus_data_matrices)):

1830 temp_1 = 1j * b_minus_matrix * u_plus_data_matrices[nu]

1831 temp_2 = -1j * b_plus_matrix * u_minus_data_matrices[nu]

1832

1833 projected_1 = projectionFFTCoreReIm(

1834 gauss_thetas,

1835 gauss_weights,

1836 temp_1,

1837 LSV,

1838 tmax,

1839 precalculatedPolyDict=precalculatedPolyDictHat,

1840 )

1841 projected_2 = projectionFFTCoreReIm(

1842 gauss_thetas,

1843 gauss_weights,

1844 temp_2,

1845 LSV,

1846 tmax,

1847 precalculatedPolyDict=precalculatedPolyDictHat,

1848 )

1849

1850 v_zero_forward_matrix[:, nu] = projected_2

1851 v_zero_forward_matrix[:, nu + len(u_plus_data_matrices)] = projected_1

1852

1853 return v_zero_forward_matrix

1854

1855

1856def calculate_v0_from_qlm(

1857 LSV,

1858 LBr,

1859 Nsv_real_imag_form,

1860 gauss_thetas,

1861 gauss_weights,

1862 phis,

1863 blm_re_im,

1864 u_plus_data_matrices,

1865 u_minus_data_matrices,

1866 u_minus_plus_from_qlm_matrix,

1867 tmax=64,

1868 debug=True,

1869):

1870 v_zero_forward_matrix = calculate_v0_from_u_minus_plus(

1871 LSV,

1872 LBr,

1873 Nsv_real_imag_form,

1874 gauss_thetas,

1875 gauss_weights,

1876 phis,

1877 blm_re_im,

1878 u_plus_data_matrices,

1879 u_minus_data_matrices,

1880 tmax,

1881 debug,

1882 )

1883 return v_zero_forward_matrix.dot(u_minus_plus_from_qlm_matrix)

1884

1885

1886def calculate_p_from_clm(

1887 LSV,

1888 LBr,

1889 Nsv_real_imag_form,

1890 Nu_real_imag_form,

1891 Br_glm_re_im,

1892 gauss_thetas,

1893 gauss_weights,

1894 phis,

1895 u_zero_data_matrices,

1896 tmax=64,

1897 debug=True,

1898):

1899 blm_re_im = np.array(Br_glm_re_im)

1900

1901 nb = 0

1902 for lb in range(1, LBr + 1):

1903 nrm = -np.sqrt(lb / (2 * (lb + 1)))

1904 for mb in range(-lb, lb + 1):

1905 blm_re_im[nb] = blm_re_im[nb] * nrm

1906 blm_re_im[nb + 1] = blm_re_im[nb + 1] * nrm

1907 nb += 2

1908

1909 b_plus_matrix = calculate_b_plus_from_ReIm(LBr, gauss_thetas, phis, blm_re_im)

1910 b_minus_matrix = calculate_b_minus_from_ReIm(LBr, gauss_thetas, phis, blm_re_im)

1911

1912 p_plus_matrix = np.zeros((Nsv_real_imag_form, Nu_real_imag_form))

1913 p_minus_matrix = np.zeros((Nsv_real_imag_form, Nu_real_imag_form))

1914

1915 temp_pluses = []

1916 temp_minuses = []

1917

1918 precalculatedPolyDictMinus = precalculatePoly(

1919 LSV, gauss_thetas, generalised_Plm_minus

1920 )

1921 precalculatedPolyDictPlus = precalculatePoly(

1922 LSV, gauss_thetas, generalised_Plm_plus

1923 )

1924

1925 if debug:

1926 for nu in range(Nu_real_imag_form):

1927 temp_plus = -1j * (b_plus_matrix) * u_zero_data_matrices[nu]

1928 temp_minus = 1j * (b_minus_matrix) * u_zero_data_matrices[nu]

1929

1930 temp_pluses.append(temp_plus)

1931 temp_minuses.append(temp_minus)

1932

1933 projected_plus = projectionFFTCoreReIm(

1934 gauss_thetas,

1935 gauss_weights,

1936 temp_plus,

1937 LSV,

1938 tmax,

1939 precalculatedPolyDict=precalculatedPolyDictPlus,

1940 )

1941 projected_minus = projectionFFTCoreReIm(

1942 gauss_thetas,

1943 gauss_weights,

1944 temp_minus,

1945 LSV,

1946 tmax,

1947 precalculatedPolyDict=precalculatedPolyDictMinus,

1948 )

1949

1950 p_plus_matrix[:, nu] = projected_plus

1951 p_minus_matrix[:, nu] = projected_minus

1952 else:

1953 for nu in range(Nu_real_imag_form):

1954 temp_plus = -1j * (b_plus_matrix) * u_zero_data_matrices[nu]

1955 temp_minus = 1j * (b_minus_matrix) * u_zero_data_matrices[nu]

1956

1957 temp_pluses.append(temp_plus)

1958 temp_minuses.append(temp_minus)

1959

1960 projected_plus = projectionFFTCoreReIm(

1961 gauss_thetas,

1962 gauss_weights,

1963 temp_plus,

1964 LSV,

1965 tmax,

1966 precalculatedPolyDict=precalculatedPolyDictPlus,

1967 )

1968 projected_minus = projectionFFTCoreReIm(

1969 gauss_thetas,

1970 gauss_weights,

1971 temp_minus,

1972 LSV,

1973 tmax,

1974 precalculatedPolyDict=precalculatedPolyDictMinus,

1975 )

1976

1977 p_plus_matrix[:, nu] = projected_plus

1978 p_minus_matrix[:, nu] = projected_minus

1979

1980 return p_plus_matrix, p_minus_matrix

1981

1982

1983def calculate_p_from_qlm(

1984 LSV,

1985 LBr,

1986 Nsv_real_imag_form,

1987 Nu_real_imag_form,

1988 Br_glm_re_im,

1989 gauss_thetas,

1990 gauss_weights,

1991 phis,

1992 u_zero_data_matrices,

1993 QtC_matrix_real_imag_form,

1994 tmax=64,

1995 debug=True,

1996):

1997 p_plus_matrix, p_minus_matrix = calculate_p_from_clm(

1998 LSV,

1999 LBr,

2000 Nsv_real_imag_form,

2001 Nu_real_imag_form,

2002 Br_glm_re_im,

2003 gauss_thetas,

2004 gauss_weights,

2005 phis,

2006 u_zero_data_matrices,

2007 tmax,

2008 debug,

2009 )

2010

2011 return p_plus_matrix.dot(QtC_matrix_real_imag_form), p_minus_matrix.dot(

2012 QtC_matrix_real_imag_form

2013 )

2014

2015

2016def calculate_u_plus_minus_forward_temp_for_s(Lq, gauss_thetas, phis, debug=False):

2017 (

2018 Q_real_imag_dict,

2019 Nq_real_imag_form,

2020 Q_cos_sin_dict,

2021 Nq_cossin_form,

2022 ) = build_real_imag_and_cossin_dict_pos(Lq, q=True)

2023

2024 u_plus_forward_matrix = np.zeros(

2025 (len(gauss_thetas) * len(phis), Nq_real_imag_form), dtype=np.complex128

2026 )

2027 u_minus_forward_matrix = np.zeros(

2028 (len(gauss_thetas) * len(phis), Nq_real_imag_form), dtype=np.complex128

2029 )

2030 ntp = 0

2031

2032 if debug:

2033 for tn in range(len(gauss_thetas)):

2034 for pn in range(len(phis)):

2035 for nq in range(int(Nq_real_imag_form / 2)):

2036 lq, mq, r_i_q = Q_real_imag_dict[2 * nq]

2037 u_plus_forward_matrix[ntp, 2 * nq] = generalised_Plm_plus(

2038 l=lq, m=mq, theta=gauss_thetas[tn]

2039 ) * np.exp(1j * mq * phis[pn])

2040 u_minus_forward_matrix[ntp, 2 * nq] = generalised_Plm_minus(

2041 l=lq, m=mq, theta=gauss_thetas[tn]

2042 ) * np.exp(1j * mq * phis[pn])

2043 u_plus_forward_matrix[ntp, 2 * nq + 1] = (

2044 1j * u_plus_forward_matrix[ntp, 2 * nq]

2045 )

2046 u_minus_forward_matrix[ntp, 2 * nq + 1] = (

2047 1j * u_minus_forward_matrix[ntp, 2 * nq]

2048 )

2049 ntp += 1

2050 else:

2051 for tn in range(len(gauss_thetas)):

2052 for pn in range(len(phis)):

2053 for nq in range(int(Nq_real_imag_form / 2)):

2054 lq, mq, r_i_q = Q_real_imag_dict[2 * nq]

2055 u_plus_forward_matrix[ntp, 2 * nq] = generalised_Plm_plus(

2056 l=lq, m=mq, theta=gauss_thetas[tn]

2057 ) * np.exp(1j * mq * phis[pn])

2058 u_minus_forward_matrix[ntp, 2 * nq] = generalised_Plm_minus(

2059 l=lq, m=mq, theta=gauss_thetas[tn]

2060 ) * np.exp(1j * mq * phis[pn])

2061 u_plus_forward_matrix[ntp, 2 * nq + 1] = (

2062 1j * u_plus_forward_matrix[ntp, 2 * nq]

2063 )

2064 u_minus_forward_matrix[ntp, 2 * nq + 1] = (

2065 1j * u_minus_forward_matrix[ntp, 2 * nq]

2066 )

2067 ntp += 1

2068

2069 return u_plus_forward_matrix, u_minus_forward_matrix

2070

2071

2072def calculate_s_from_delta(

2073 LSV,

2074 LBr,

2075 Lq,

2076 Br_glm_re_im,

2077 gauss_thetas,

2078 gauss_weights,

2079 phis,

2080 u_plus_forward_matrix,

2081 u_minus_forward_matrix,

2082 tmax=64,

2083 debug=False,

2084):

2085 Br_data_matrix = calculate_Br_from_ReIm(LBr, gauss_thetas, phis, Br_glm_re_im)

2086

2087 (

2088 SV_in_real_imag_form_dict,

2089 Nsv_real_imag_form,

2090 SV_cos_sin_dict,

2091 Nsv_cossin_form,

2092 ) = build_real_imag_and_cossin_dict_pos(LSV)

2093

2094 (

2095 Q_real_imag_dict,

2096 Nq_real_imag_form,

2097 Q_cos_sin_dict,

2098 Nq_cossin_form,

2099 ) = build_real_imag_and_cossin_dict_pos(Lq, q=True)

2100

2101 s_plus_from_delta_plus_matrix = np.zeros((Nsv_real_imag_form, Nq_real_imag_form))

2102 s_minus_from_delta_minus_matrix = np.zeros((Nsv_real_imag_form, Nq_real_imag_form))

2103

2104 def one_norm(*args):

2105 return 1

2106

2107 precalculatedPolyDictMinus = precalculatePoly(

2108 LSV, gauss_thetas, generalised_Plm_minus

2109 )

2110 precalculatedPolyDictPlus = precalculatePoly(

2111 LSV, gauss_thetas, generalised_Plm_plus

2112 )

2113

2114 if debug:

2115 for nq in range(Nq_real_imag_form):

2116 array = np.zeros((Nq_real_imag_form,))

2117 array[nq] = 1

2118

2119 temp_plus = (

2120 1j

2121 * Br_data_matrix

2122 * u_plus_forward_matrix.dot(array).reshape(tmax, tmax * 2)

2123 )

2124 temp_minus = (

2125 -1j

2126 * Br_data_matrix

2127 * u_minus_forward_matrix.dot(array).reshape(tmax, tmax * 2)

2128 )

2129

2130 projected_plus = projectionFFTCoreReIm(

2131 gauss_thetas,

2132 gauss_weights,

2133 temp_plus,

2134 LSV,

2135 tmax,

2136 precalculatedPolyDict=precalculatedPolyDictPlus,

2137 )

2138 projected_minus = projectionFFTCoreReIm(

2139 gauss_thetas,

2140 gauss_weights,

2141 temp_minus,

2142 LSV,

2143 tmax,

2144 precalculatedPolyDict=precalculatedPolyDictMinus,

2145 )

2146

2147 s_plus_from_delta_plus_matrix[:, nq] = projected_plus

2148 s_minus_from_delta_minus_matrix[:, nq] = projected_minus

2149 else:

2150 for nq in range(Nq_real_imag_form):

2151 array = np.zeros((Nq_real_imag_form,))

2152 array[nq] = 1

2153

2154 temp_plus = (

2155 1j

2156 * Br_data_matrix

2157 * u_plus_forward_matrix.dot(array).reshape(tmax, tmax * 2)

2158 )

2159 temp_minus = (

2160 -1j

2161 * Br_data_matrix

2162 * u_minus_forward_matrix.dot(array).reshape(tmax, tmax * 2)

2163 )

2164

2165 projected_plus = projectionFFTCoreReIm(

2166 gauss_thetas,

2167 gauss_weights,

2168 temp_plus,

2169 LSV,

2170 tmax,

2171 precalculatedPolyDict=precalculatedPolyDictPlus,

2172 )

2173 projected_minus = projectionFFTCoreReIm(

2174 gauss_thetas,

2175 gauss_weights,

2176 temp_minus,

2177 LSV,

2178 tmax,

2179 precalculatedPolyDict=precalculatedPolyDictMinus,

2180 )

2181

2182 s_plus_from_delta_plus_matrix[:, nq] = projected_plus

2183 s_minus_from_delta_minus_matrix[:, nq] = projected_minus

2184

2185 return s_plus_from_delta_plus_matrix, s_minus_from_delta_minus_matrix

2186

2187

2188def calculate_gamma(l, m):

2189 return np.sqrt(((l - m) * (l + m)) / ((2 * l + 1) * (2 * l - 1)))

2190

2191

2192def calculate_psi1_to_psi2(L):

2193 saved_data = {}

2194 saved_indeces = []

2195

2196 for m in range(0, L + 1):

2197 psi_coefs_pos = {(2 * n + m, m): n for n in range(0, int((L - m) / 2) + 1)}

2198 y_coefs_pos = {(2 * n + m + 1, m): n for n in range(0, int((L - m) / 2) + 1)}

2199 temp_matrix_sin = np.zeros((int((L - m) / 2) + 1, int((L - m) / 2) + 1))

2200 temp_matrix_cos = np.zeros((int((L - m) / 2) + 1, int((L - m) / 2) + 1))

2201 saved_indeces = saved_indeces + list(psi_coefs_pos.keys())

2202 for n in range(0, int((L - m) / 2) + 1):

2203 l_2 = 2 * n + m + 1

2204 l = l_2 - 1

2205 gamma = calculate_gamma(l_2, m)

2206 temp_sin = gamma * (l_2 - 1) * 4

2207 temp_cos = gamma

2208 ind_psi = psi_coefs_pos[(l, m)]

2209 ind_y = y_coefs_pos[(l_2, m)]

2210 temp_matrix_sin[ind_y, ind_psi] = temp_sin

2211 temp_matrix_cos[ind_y, ind_psi] = temp_cos

2212

2213 if l_2 + 1 <= L:

2214 l = l_2 + 1

2215 gamma = calculate_gamma(l_2 + 1, m)

2216 temp_sin = gamma * (l_2 + 2) * (-1) * 4

2217 # temp_sin = gamma * (l_2 + 2) * (-1)

2218 temp_cos = gamma

2219 ind_psi = psi_coefs_pos[(l, m)]

2220 ind_y = y_coefs_pos[(l_2, m)]

2221 temp_matrix_sin[ind_y, ind_psi] = temp_sin

2222 temp_matrix_cos[ind_y, ind_psi] = temp_cos

2223

2224 temp_matrix = np.linalg.inv(temp_matrix_cos).dot(temp_matrix_sin)

2225

2226 saved_data[m] = {

2227 "psi_pos": psi_coefs_pos,

2228 "y_pos": y_coefs_pos,

2229 "matrix": temp_matrix,

2230 }

2231

2232 return {"ind_pairs": saved_indeces, "m_fixed_data": saved_data}

2233

2234

2235def calculate_psi1_to_psi2_matrix(Lq, data):

2236 (

2237 Q_real_imag_dict,

2238 Nq_real_imag_form,

2239 Q_cos_sin_dict,

2240 Nq_cossin_form,

2241 ) = build_real_imag_and_cossin_dict_pos(Lq, q=True)

2242

2243 matrix = np.zeros((Nq_real_imag_form, Nq_real_imag_form))

2244

2245 ind_pairs = data["ind_pairs"]

2246 data_matrices = data["m_fixed_data"]

2247

2248 for nq_r_i in range(Nq_real_imag_form):

2249 lq, mq, r_i = Q_real_imag_dict[nq_r_i]

2250 if (lq, mq) in ind_pairs:

2251 ind_psi1 = data_matrices[mq]["psi_pos"][(lq, mq)]

2252 for psi_2_pair in data_matrices[mq]["psi_pos"].keys():

2253 ind_psi2 = data_matrices[mq]["psi_pos"][psi_2_pair]

2254 r_i_2 = r_i

2255 forw_matrix = data_matrices[mq]["matrix"]

2256

2257 npsi2_r_i_pos = list(Q_real_imag_dict.keys())[

2258 list(Q_real_imag_dict.values()).index(

2259 (psi_2_pair[0], psi_2_pair[1], r_i_2)

2260 )

2261 ]

2262 matrix[npsi2_r_i_pos, nq_r_i] = forw_matrix[ind_psi2, ind_psi1]

2263

2264 npsi2_r_i_neg = list(Q_real_imag_dict.keys())[

2265 list(Q_real_imag_dict.values()).index(

2266 (psi_2_pair[0], -psi_2_pair[1], r_i_2)

2267 )

2268 ]

2269 ind_psi1_neg = list(Q_real_imag_dict.keys())[

2270 list(Q_real_imag_dict.values()).index((lq, -mq, r_i))

2271 ]

2272 matrix[npsi2_r_i_neg, ind_psi1_neg] = forw_matrix[ind_psi2, ind_psi1]

2273

2274 return matrix

2275

2276

2277def calculate_s_tilda_from_qlm_tilda(Lq, rc=1):

2278 (

2279 Q_in_real_imag_form_dict,

2280 Nq_real_imag_form,

2281 Q_cos_sin_dict,

2282 Nq_cossin_form,

2283 ) = build_real_imag_and_cossin_dict_pos(Lq, q=True)

2284

2285 Matrix_multiply_by_i_Q = building_transformation_matrix_multiply_by_i(

2286 Nq_real_imag_form, Q_in_real_imag_form_dict

2287 )

2288

2289 matrix_plus = np.zeros((Nq_real_imag_form, Nq_real_imag_form))

2290 matrix_minus = np.zeros((Nq_real_imag_form, Nq_real_imag_form))

2291

2292 for ls in range(Lq + 1):

2293 for ms in range(ls + 1):

2294 for r_i in ["r", "i"]:

2295 r_i_2 = r_i

2296 coef = 0

2297 ###################################

2298 # l = 0 is under the question

2299 ###################################

2300 if ls > 0:

2301 coef = np.sqrt(1 / (2 * ls * (ls + 1) * (rc**2)))

2302

2303 # q_{l-1}^m

2304 if ls - 1 >= ms:

2305 gamma = calculate_gamma(ls, ms)

2306 ns_r_i = list(Q_in_real_imag_form_dict.keys())[

2307 list(Q_in_real_imag_form_dict.values()).index((ls, ms, r_i))

2308 ]

2309 nq_r_i = list(Q_in_real_imag_form_dict.keys())[

2310 list(Q_in_real_imag_form_dict.values()).index(

2311 (ls - 1, ms, r_i_2)

2312 )

2313 ]

2314 temp = (ls + 1) * gamma

2315 matrix_plus[ns_r_i, nq_r_i] = temp * coef

2316 matrix_minus[ns_r_i, nq_r_i] = -temp * coef

2317

2318 if ms != 0:

2319 ns_r_i = list(Q_in_real_imag_form_dict.keys())[

2320 list(Q_in_real_imag_form_dict.values()).index(

2321 (ls, -ms, r_i)

2322 )

2323 ]

2324 nq_r_i = list(Q_in_real_imag_form_dict.keys())[

2325 list(Q_in_real_imag_form_dict.values()).index(

2326 (ls - 1, -ms, r_i_2)

2327 )

2328 ]

2329 matrix_plus[ns_r_i, nq_r_i] = temp * coef

2330 matrix_minus[ns_r_i, nq_r_i] = -temp * coef

2331

2332 # q_{l}^m

2333 ns_r_i = list(Q_in_real_imag_form_dict.keys())[

2334 list(Q_in_real_imag_form_dict.values()).index((ls, ms, r_i))

2335 ]

2336 nq_r_i = list(Q_in_real_imag_form_dict.keys())[

2337 list(Q_in_real_imag_form_dict.values()).index((ls, ms, r_i_2))

2338 ]

2339 temp = -ms

2340 matrix_plus[ns_r_i, nq_r_i] = temp * coef

2341 matrix_minus[ns_r_i, nq_r_i] = temp * coef

2342

2343 if ms != 0:

2344 ns_r_i = list(Q_in_real_imag_form_dict.keys())[

2345 list(Q_in_real_imag_form_dict.values()).index((ls, -ms, r_i))

2346 ]

2347 nq_r_i = list(Q_in_real_imag_form_dict.keys())[

2348 list(Q_in_real_imag_form_dict.values()).index((ls, -ms, r_i_2))

2349 ]

2350 matrix_plus[ns_r_i, nq_r_i] = -temp * coef

2351 matrix_minus[ns_r_i, nq_r_i] = -temp * coef

2352 # matrix_plus[ns_r_i, nq_r_i] = temp * coef

2353 # matrix_minus[ns_r_i, nq_r_i] = temp * coef

2354

2355 # q_{l+1}^m

2356 if ls + 1 <= Lq:

2357 gamma = calculate_gamma(ls + 1, ms)

2358 r_i_2 = r_i

2359 ns_r_i = list(Q_in_real_imag_form_dict.keys())[

2360 list(Q_in_real_imag_form_dict.values()).index((ls, ms, r_i))

2361 ]

2362 nq_r_i = list(Q_in_real_imag_form_dict.keys())[

2363 list(Q_in_real_imag_form_dict.values()).index(

2364 (ls + 1, ms, r_i_2)

2365 )

2366 ]

2367 temp = ls * gamma

2368 matrix_plus[ns_r_i, nq_r_i] = -temp * coef

2369 matrix_minus[ns_r_i, nq_r_i] = temp * coef

2370

2371 if ms != 0:

2372 ns_r_i = list(Q_in_real_imag_form_dict.keys())[

2373 list(Q_in_real_imag_form_dict.values()).index(

2374 (ls, -ms, r_i)

2375 )

2376 ]

2377 nq_r_i = list(Q_in_real_imag_form_dict.keys())[

2378 list(Q_in_real_imag_form_dict.values()).index(

2379 (ls + 1, -ms, r_i_2)

2380 )

2381 ]

2382 matrix_plus[ns_r_i, nq_r_i] = -temp * coef

2383 matrix_minus[ns_r_i, nq_r_i] = temp * coef

2384

2385 matrix_plus = Matrix_multiply_by_i_Q.dot(matrix_plus)

2386 matrix_minus = Matrix_multiply_by_i_Q.dot(matrix_minus)

2387

2388 return matrix_plus, matrix_minus

2389

2390

2391def calculate_s_hat_from_qlm(Lq, rc=1):

2392 (

2393 Q_in_real_imag_form_dict,

2394 Nq_real_imag_form,

2395 Q_cos_sin_dict,

2396 Nq_cossin_form,

2397 ) = build_real_imag_and_cossin_dict_pos(Lq, q=True)

2398

2399 Matrix_multiply_by_i = building_transformation_matrix_multiply_by_i(

2400 Nq_real_imag_form, Q_in_real_imag_form_dict

2401 )

2402

2403 matrix_plus = np.zeros((Nq_real_imag_form, Nq_real_imag_form))

2404 matrix_minus = np.zeros((Nq_real_imag_form, Nq_real_imag_form))

2405

2406 for m in range(0, Lq + 1):

2407 for n in range(0, int((Lq - m) / 2) + 1):

2408 l_2 = 2 * n + m + 1

2409 if l_2 <= Lq:

2410 # print(l_2, m)

2411 coef = np.sqrt((2 * n + m + 1) * (2 * n + m + 2)) / (

2412 np.sqrt(2) * (rc**2)

2413 )

2414 for r_i_2 in ["r", "i"]:

2415 n2_r_i_pos = list(Q_in_real_imag_form_dict.keys())[

2416 list(Q_in_real_imag_form_dict.values()).index((l_2, m, r_i_2))

2417 ]

2418 n2_r_i_neg = list(Q_in_real_imag_form_dict.keys())[

2419 list(Q_in_real_imag_form_dict.values()).index((l_2, -m, r_i_2))

2420 ]

2421 r_i_q = r_i_2

2422 l = 2 * n + m

2423 nq_r_i_pos = list(Q_in_real_imag_form_dict.keys())[

2424 list(Q_in_real_imag_form_dict.values()).index((l, m, r_i_q))

2425 ]

2426 nq_r_i_neg = list(Q_in_real_imag_form_dict.keys())[

2427 list(Q_in_real_imag_form_dict.values()).index((l, -m, r_i_q))

2428 ]

2429 temp = coef * (2 * n + m) * calculate_gamma(2 * n + m + 1, m)

2430 matrix_plus[n2_r_i_pos, nq_r_i_pos] = temp

2431 matrix_minus[n2_r_i_pos, nq_r_i_pos] = -temp

2432

2433 # temp = coef * (2 * n + (-m)) * calculate_gamma(2 * n + (-m) + 1, (-m))

2434 matrix_plus[n2_r_i_neg, nq_r_i_neg] = temp

2435 matrix_minus[n2_r_i_neg, nq_r_i_neg] = -temp

2436

2437 if (2 * (n + 1) + m) <= Lq:

2438 l = 2 * (n + 1) + m

2439 nq_r_i_pos = list(Q_in_real_imag_form_dict.keys())[

2440 list(Q_in_real_imag_form_dict.values()).index((l, m, r_i_q))

2441 ]

2442 nq_r_i_neg = list(Q_in_real_imag_form_dict.keys())[

2443 list(Q_in_real_imag_form_dict.values()).index(

2444 (l, -m, r_i_q)

2445 )

2446 ]

2447 temp = (

2448 -coef * (2 * n + m + 3) * calculate_gamma(2 * n + m + 2, m)

2449 )

2450 matrix_plus[n2_r_i_pos, nq_r_i_pos] = temp

2451 matrix_minus[n2_r_i_pos, nq_r_i_pos] = -temp

2452

2453 matrix_plus[n2_r_i_neg, nq_r_i_neg] = temp

2454 matrix_minus[n2_r_i_neg, nq_r_i_neg] = -temp

2455

2456 matrix_plus = Matrix_multiply_by_i.dot(matrix_plus)

2457 matrix_minus = Matrix_multiply_by_i.dot(matrix_minus)

2458

2459 return matrix_plus, matrix_minus

2460

2461

2462def calculate_u_forward(Lu, thetas, phis, Plm_func):

2463 (

2464 u_in_real_imag_form_dict,

2465 Nu_real_imag_form,

2466 u_cos_sin_dict,

2467 Nu_cossin_form,

2468 ) = build_real_imag_and_cossin_dict_pos(Lu)

2469

2470 u_forward_matrix = np.zeros(

2471 (len(thetas) * len(phis), Nu_real_imag_form), dtype=np.complex128

2472 )

2473 ntp = 0

2474 for tn in range(len(thetas)):

2475 for pn in range(len(phis)):

2476 for nu in range(int(Nu_real_imag_form / 2)):

2477 lu, mu, r_i_u = u_in_real_imag_form_dict[2 * nu]

2478 u_forward_matrix[ntp, 2 * nu] = Plm_func(

2479 l=lu, m=mu, theta=thetas[tn]

2480 ) * np.exp(1j * mu * phis[pn])

2481 u_forward_matrix[ntp, 2 * nu + 1] = 1j * u_forward_matrix[ntp, 2 * nu]

2482 ntp += 1

2483

2484 return u_forward_matrix

2485

2486

2487def from_re_im_to_complex(x_in_re_im):

2488 nu_re_im = x_in_re_im.reshape(

2489 -1,

2490 ).shape[0]

2491 x_in_re_im = x_in_re_im.reshape(-1, 1)

2492 x_in_complex = np.zeros((int(nu_re_im / 2), 1), dtype=np.complex128)

2493 for i in range(int(nu_re_im / 2)):

2494 x_in_complex[i, 0] = x_in_re_im[2 * i, 0] + 1j * x_in_re_im[2 * i + 1, 0]

2495 return x_in_complex

2496

2497

2498def from_complex_to_re_im(x_in_complex):

2499 nu_complex = x_in_complex.reshape(

2500 -1,

2501 ).shape[0]

2502 x_in_complex = x_in_complex.reshape(-1, 1)

2503 x_in_re_im = np.zeros((int(nu_complex * 2), 1))

2504 for i in range(int(nu_complex)):

2505 x_in_re_im[2 * i, 0] = x_in_complex[i, 0].real

2506 x_in_re_im[2 * i + 1, 0] = x_in_complex[i, 0].imag

2507 return x_in_re_im

2508

2509

2510def calculate_u_forward_complex_old(Lu, thetas, phis, Plm_func, debug=False):

2511 (

2512 u_in_complex_form_dict,

2513 Nu_complex_form,

2514 u_cos_sin_dict,

2515 Nu_cossin_form,

2516 ) = build_complex_and_cossin_dict_pos(Lu)

2517

2518 u_forward_matrix = np.zeros(

2519 (len(thetas) * len(phis), Nu_complex_form), dtype=np.complex128

2520 )

2521 ntp = 0

2522 if debug == False:

2523 for tn in range(len(thetas)):

2524 for pn in range(len(phis)):

2525 for nu in range(Nu_complex_form):

2526 lu, mu = u_in_complex_form_dict[nu]

2527 u_forward_matrix[ntp, nu] = Plm_func(

2528 l=lu, m=mu, theta=thetas[tn]

2529 ) * np.exp(1j * mu * phis[pn])

2530 ntp += 1

2531 else:

2532 for tn in range(len(thetas)):

2533 for pn in range(len(phis)):

2534 for nu in range(Nu_complex_form):

2535 lu, mu = u_in_complex_form_dict[nu]

2536 u_forward_matrix[ntp, nu] = Plm_func(

2537 l=lu, m=mu, theta=thetas[tn]

2538 ) * np.exp(1j * mu * phis[pn])

2539 ntp += 1

2540

2541 return u_forward_matrix

2542

2543

2544def calculate_u_forward_complex(Lu, thetas, phis, Plm_func, debug=False):

2545 (

2546 u_in_complex_form_dict,

2547 Nu_complex_form,

2548 u_cos_sin_dict,

2549 Nu_cossin_form,

2550 ) = build_complex_and_cossin_dict_pos(Lu)

2551

2552 u_forward_matrix = np.zeros(

2553 (len(thetas) * len(phis), Nu_complex_form), dtype=np.complex128

2554 )

2555 ntp = 0

2556 if debug == False:

2557 for nu in range(Nu_complex_form):

2558 ntp = 0

2559 for tn in range(len(thetas)):

2560 lu, mu = u_in_complex_form_dict[nu]

2561 Plm_temp = Plm_func(

2562 l=lu, m=mu, theta=thetas[tn]

2563 )

2564 for pn in range(len(phis)):

2565 u_forward_matrix[ntp, nu] = Plm_temp * np.exp(1j * mu * phis[pn])

2566 ntp += 1

2567 else:

2568 for nu in range(Nu_complex_form):

2569 ntp = 0

2570 for tn in range(len(thetas)):

2571 lu, mu = u_in_complex_form_dict[nu]

2572 Plm_temp = Plm_func(

2573 l=lu, m=mu, theta=thetas[tn]

2574 )

2575 for pn in range(len(phis)):

2576 u_forward_matrix[ntp, nu] = Plm_temp * np.exp(1j * mu * phis[pn])

2577 ntp += 1

2578

2579 return u_forward_matrix

2580

2581

2582def build_complex_and_cossin_dict_pos(L):

2583 x_in_comlex_form_dict = {}

2584 Nx_comlex_form = 0

2585 for lx in range(1, L + 1):

2586 for mx in range(-lx, lx + 1):

2587 x_in_comlex_form_dict[Nx_comlex_form] = (lx, mx)

2588 Nx_comlex_form += 1

2589

2590 x_cos_sin_dict = {}

2591 Nx_cossin_form = 0

2592 for lx in range(1, L + 1):

2593 for mx in range(0, lx + 1):

2594 x_cos_sin_dict[Nx_cossin_form] = (lx, mx, "cos")

2595 Nx_cossin_form += 1

2596 if mx != 0:

2597 x_cos_sin_dict[Nx_cossin_form] = (lx, mx, "sin")

2598 Nx_cossin_form += 1

2599

2600 return x_in_comlex_form_dict, Nx_comlex_form, x_cos_sin_dict, Nx_cossin_form

2601

2602

2603# =============================================

2604# Building matrix corresponding clm = M qlm

2605# Re Im form

2606# =============================================

2607

2608

2609def build_qlm_to_clm_matrix_Re_Im(Lmax):

2610 Nu = Lmax * (Lmax + 2) * 2

2611 qlm_to_clm_matrix = np.zeros((Nu, Nu))

2612

2613 nu = 0

2614 for lu in range(1, Lmax + 1):

2615 for mu in range(-lu, lu + 1):

2616 temp_coef = -2 * mu / (lu * (lu + 1))

2617 if mu == 0:

2618 nu += 2

2619 else:

2620 qlm_to_clm_matrix[nu, nu + 1] = -temp_coef

2621 qlm_to_clm_matrix[nu + 1, nu] = temp_coef

2622 nu += 2

2623

2624 return qlm_to_clm_matrix

2625

2626

2627# =============================================

2628# Building matrix corresponding tlm = T qlm

2629# Re Im

2630# =============================================

2631

2632

2633def build_qlm_to_tlm_matrix_Re_Im(Lmax):

2634 def calculate_gamma_lm(l, m):

2635 return np.sqrt((l - m) * (l + m) / (2 * l + 1) / (2 * l - 1))

2636

2637 Nq = Lmax * (Lmax + 2) * 2

2638 Lu = Lmax + 1

2639 Nu = Lu * (Lu + 2) * 2

2640 qlm_to_tlm_matrix = np.zeros((Nu, Nq))

2641

2642 lumu_dict = {}

2643

2644 nu = 0

2645 for lu in range(1, (Lmax + 1) + 1):

2646 for mu in range(-lu, lu + 1):

2647 lumu_dict[(lu, mu, "r")] = nu

2648 nu += 1

2649 lumu_dict[(lu, mu, "i")] = nu

2650 nu += 1

2651

2652 nu = 0

2653 for lu in range(1, Lu + 1):

2654 for mu in range(-lu, lu + 1):

2655 if lu <= 1:

2656 pass

2657 else:

2658 left_coef = lu * (lu - 1)

2659 rigth_coef = (lu + 1) * (lu + 2)

2660 if lu < Lmax:

2661 if mu == 0:

2662 pass

2663 else:

2664 temp_coef = -1 / (lu * (lu + 1))

2665 nu = lumu_dict[(lu, mu, "r")]

2666 nu_plus = lumu_dict[(lu + 1, mu, "r")]

2667 if (lu - 1, mu, "r") in lumu_dict:

2668 nu_minus = lumu_dict[(lu - 1, mu, "r")]

2669

2670 gamma_l_plus = calculate_gamma_lm(lu + 1, mu)

2671 gamma_l = calculate_gamma_lm(lu, mu)

2672

2673 if (lu - 1, mu, "r") not in lumu_dict:

2674 qlm_to_tlm_matrix[nu, nu_plus] = (

2675 gamma_l_plus * left_coef * temp_coef

2676 )

2677 else:

2678 qlm_to_tlm_matrix[nu, nu_plus] = (

2679 gamma_l_plus * left_coef * temp_coef

2680 )

2681 qlm_to_tlm_matrix[nu, nu_minus] = (

2682 gamma_l * rigth_coef * temp_coef

2683 )

2684

2685 nu = lumu_dict[(lu, mu, "i")]

2686 nu_plus = lumu_dict[(lu + 1, mu, "i")]

2687 if (lu - 1, mu, "i") in lumu_dict:

2688 nu_minus = lumu_dict[(lu - 1, mu, "i")]

2689

2690 if (lu - 1, mu, "i") not in lumu_dict:

2691 qlm_to_tlm_matrix[nu, nu_plus] = (

2692 gamma_l_plus * left_coef * temp_coef

2693 )

2694 else:

2695 qlm_to_tlm_matrix[nu, nu_plus] = (

2696 gamma_l_plus * left_coef * temp_coef

2697 )

2698 qlm_to_tlm_matrix[nu, nu_minus] = (

2699 gamma_l * rigth_coef * temp_coef

2700 )

2701

2702 else:

2703 if mu == 0:

2704 pass

2705 else:

2706 temp_coef = -1 / (lu * (lu + 1))

2707

2708 nu = lumu_dict[(lu, mu, "r")]

2709 if (lu - 1, mu, "r") in lumu_dict:

2710 nu_minus = lumu_dict[(lu - 1, mu, "r")]

2711

2712 gamma_l = calculate_gamma_lm(lu, mu)

2713

2714 if (lu - 1, mu, "r") not in lumu_dict:

2715 pass

2716 else:

2717 qlm_to_tlm_matrix[nu, nu_minus] = (

2718 gamma_l * rigth_coef * temp_coef

2719 )

2720

2721 nu = lumu_dict[(lu, mu, "i")]

2722 if (lu - 1, mu, "i") in lumu_dict:

2723 nu_minus = lumu_dict[(lu - 1, mu, "i")]

2724

2725 if (lu - 1, mu, "i") not in lumu_dict:

2726 pass

2727 else:

2728 qlm_to_tlm_matrix[nu, nu_minus] = (

2729 gamma_l * rigth_coef * temp_coef

2730 )

2731

2732 return qlm_to_tlm_matrix

2733

2734

2735# =============================================

2736# Building matrix corresponding tlm = T qlm

2737# for zonal part ql0 to tl0 Re Im form

2738# =============================================

2739

2740

2741def calculation_ql0_to_tl0_zonal_coefs(n):

2742 def calc_al(l):

2743 return (

2744 (-1) * (2 * l + 1) * (2 * l + 2) / np.sqrt(4 * l + 3) / np.sqrt(4 * l + 1)

2745 )

2746

2747 def calc_bl(l):

2748 return (2 * l + 1) * (2 * l + 2) / np.sqrt(4 * l + 3) / np.sqrt(4 * l + 5)

2749

2750 leftPartPlus = calc_al(n)

2751 if n > 0:

2752 leftPartMinus = calc_bl(n - 1)

2753 else:

2754 leftPartMinus = 0

2755

2756 def calc_cl(l):

2757 return (1 / np.sqrt(2 * l + 1)) * ((l * (l + 1)) / (2 * l - 1)) * (

2758 (l - 1) / np.sqrt(2 * l - 3)

2759 ) - (3 / np.sqrt(2 * l + 1)) * (l / (2 * l - 1)) * (

2760 (l - 1) / np.sqrt(2 * l - 3)

2761 )

2762

2763 def calc_dl(l):

2764 return (

2765 (-1) * (1 / (2 * l + 1)) * (l * (l + 1) * (l + 1) / (2 * l + 3))

2766 + (l / (2 * l + 1)) * (l * (l + 1) / (2 * l - 1))

2767 + 3

2768 - (3 * (l + 1) * (l + 1) / (2 * l + 1) / (2 * l + 3))

2769 - (3 * l * l / (2 * l + 1) / (2 * l - 1))

2770 )

2771

2772 def calc_el(l):

2773 return (-1) * (1 / np.sqrt(2 * l + 1)) * (l * (l + 1) / (2 * l + 1)) * (

2774 (l + 2) / np.sqrt(2 * l + 5)

2775 ) - (3 / np.sqrt(2 * l + 1)) * ((l + 1) / (2 * l + 3)) * (

2776 (l + 2) / np.sqrt(2 * l + 5)

2777 )

2778

2779 rightPartMin2 = calc_cl(2 * n + 2)

2780 rightPartZero = calc_dl(2 * n)

2781 if n > 0:

2782 rightPartPlus2 = calc_el(2 * n - 2)

2783 else:

2784 rightPartPlus2 = 0

2785

2786 return leftPartPlus, leftPartMinus, rightPartMin2, rightPartZero, rightPartPlus2

2787

2788

2789def build_qlm_to_tlm_zonal_matrix_Re_Im(Lmax_q):

2790 Lu_max = Lmax_q + 1

2791 Nq = Lmax_q * (Lmax_q + 2) * 2 + 2

2792 Nu = Lu_max * (Lu_max + 2) * 2

2793 matrix = np.zeros((Nu, Nq))

2794 # gauss_theta, gauss_weights = gaussPoints(0, np.pi, 64)

2795

2796 u_dict = {}

2797 nu = 0

2798 for lu in range(1, Lu_max + 1):

2799 for mu in range(-lu, lu + 1):

2800 u_dict[lu, mu, "r"] = nu

2801 nu += 1

2802 u_dict[lu, mu, "i"] = nu

2803 nu += 1

2804

2805 q_dict = {}

2806 nq = 0

2807 for lq in range(0, Lmax_q + 1):

2808 for mq in range(-lq, lq + 1):

2809 q_dict[lq, mq, "r"] = nq

2810 nq += 1

2811 q_dict[lq, mq, "i"] = nq

2812 nq += 1

2813

2814 tprev = 0

2815 for n in range(0, Lmax_q // 2 + 1):

2816 a, b, c, d, e = calculation_ql0_to_tl0_zonal_coefs(n)

2817 if n == 0:

2818 t2np = u_dict[2 * n + 1, 0, "r"]

2819 q2np = q_dict[2 * n + 2, 0, "r"]

2820 q2n = q_dict[2 * n, 0, "r"]

2821 matrix[t2np, q2np] = c / a

2822 matrix[t2np, q2n] = d / a

2823 tprev = t2np

2824 else:

2825 a0, _, _, _, _ = calculation_ql0_to_tl0_zonal_coefs(n - 1)

2826 t2np = u_dict[2 * n + 1, 0, "r"]

2827 if 2 * n + 2 <= Lmax_q:

2828 q2np = q_dict[2 * n + 2, 0, "r"]

2829 q2n = q_dict[2 * n, 0, "r"]

2830 q2nm = q_dict[2 * n - 2, 0, "r"]

2831 if 2 * n + 2 <= Lmax_q:

2832 matrix[t2np, q2np] = c / a

2833 matrix[t2np, q2n] = d / a

2834 matrix[t2np, q2nm] = e / a

2835 matrix[t2np, :] -= matrix[tprev, :] * b / a0

2836 tprev = t2np

2837

2838 return matrix

2839

2840

2841def build_qlm_to_tlm_zonal_matrix_Re_Im_v2(Lmax_q):

2842 Lu_max = Lmax_q + 1

2843 Nq = Lmax_q * (Lmax_q + 2) * 2 + 2

2844 Nu = Lu_max * (Lu_max + 2) * 2

2845 matrix = np.zeros((Nu, Nq))

2846

2847 u_in_re_im_form_dict, Nu_re_im_form, _, _ = build_real_imag_and_cossin_dict_pos(

2848 Lu_max, q=False

2849 )

2850 q_in_re_im_form_dict, Nq_re_im_form, _, _ = build_real_imag_and_cossin_dict_pos(

2851 Lmax_q, q=True

2852 )

2853

2854 nu = 0

2855 for lu in range(1, Lu_max + 1):

2856 for mu in range(-lu, lu + 1):

2857 if mu == 0:

2858 lq_m1 = lu - 1

2859 lq_p1 = lu + 1

2860

2861 if lq_m1 >= 0:

2862 gamma = calculate_gamma(lu, 0)

2863 nq = list(q_in_re_im_form_dict.keys())[

2864 list(q_in_re_im_form_dict.values()).index((lq_m1, 0, "r"))

2865 ]

2866 matrix[nu, nq] = gamma * (-(lu + 2) / lu)

2867 matrix[nu + 1, nq + 1] = gamma * (-(lu + 2) / lu)

2868

2869 if lq_p1 <= Lmax_q:

2870 gamma = calculate_gamma(lu + 1, 0)

2871 nq = list(q_in_re_im_form_dict.keys())[

2872 list(q_in_re_im_form_dict.values()).index((lq_p1, 0, "r"))

2873 ]

2874 matrix[nu, nq] = gamma * (-(lu - 1) / (lu + 1))

2875 matrix[nu + 1, nq + 1] = gamma * (-(lu - 1) / (lu + 1))

2877 nu += 2

2879 return matrix

2882def compare_plus_minus_coefficients(plus_coefficients, minus_coefficients, re_im_dict):

2883 for i in range(len(plus_coefficients.reshape((-1,)))):

2884 l, m, reim = re_im_dict[i]

2885 if m < 0:

2886 m_ = -m

2887 pos = list(re_im_dict.keys())[

2888 list(re_im_dict.values()).index((l, m_, reim))

2889 ]

2890 print(re_im_dict[i])

2891 if reim == "r":

2892 print(

2893 "plus", plus_coefficients[i], ((-1) ** m_) * minus_coefficients[pos]

2894 )

2895 print(

2896 "minus",

2897 minus_coefficients[i],

2898 ((-1) ** m_) * plus_coefficients[pos],

2899 )

2900 else:

2901 print(

2902 "plus",

2903 plus_coefficients[i],

2904 -((-1) ** m_) * minus_coefficients[pos],

2905 )

2906 print(

2907 "minus",

2908 minus_coefficients[i],

2909 -((-1) ** m_) * plus_coefficients[pos],

2910 )