ml4gw 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of ml4gw might be problematic. Click here for more details.

@@ -0,0 +1,779 @@
1
+ from typing import Dict, Tuple
2
+
3
+ import torch
4
+ from torchtyping import TensorType
5
+
6
+ from ..constants import MPC_SEC, MTSUN_SI, PI
7
+ from .phenom_d import IMRPhenomD
8
+
9
+
10
+ class IMRPhenomPv2(IMRPhenomD):
11
+ def __init__(self):
12
+ super().__init__()
13
+
14
+ def forward(
15
+ self,
16
+ fs: TensorType,
17
+ chirp_mass: TensorType,
18
+ mass_ratio: TensorType,
19
+ s1x: TensorType,
20
+ s1y: TensorType,
21
+ s1z: TensorType,
22
+ s2x: TensorType,
23
+ s2y: TensorType,
24
+ s2z: TensorType,
25
+ dist_mpc: TensorType,
26
+ tc: TensorType,
27
+ phiRef: TensorType,
28
+ incl: TensorType,
29
+ f_ref: float,
30
+ ):
31
+ """
32
+ IMRPhenomPv2 waveform
33
+
34
+ Args:
35
+ fs :
36
+ Frequency series in Hz.
37
+ chirp_mass :
38
+ Chirp mass in solar masses.
39
+ mass_ratio :
40
+ Mass ratio m1/m2.
41
+ s1x :
42
+ Spin component x of the first BH.
43
+ s1y :
44
+ Spin component y of the first BH.
45
+ s1z :
46
+ Spin component z of the first BH.
47
+ s2x :
48
+ Spin component x of the second BH.
49
+ s2y :
50
+ Spin component y of the second BH.
51
+ s2z :
52
+ Spin component z of the second BH.
53
+ dist_mpc :
54
+ Luminosity distance in Mpc.
55
+ tc :
56
+ Coalescence time.
57
+ phiRef :
58
+ Reference phase.
59
+ incl :
60
+ Inclination angle.
61
+ f_ref :
62
+ Reference frequency in Hz.
63
+
64
+ Returns:
65
+ hc, hp: Tuple[torch.Tensor, torch.Tensor]
66
+ Cross and plus polarizations
67
+
68
+ Note: m1 must be larger than m2.
69
+ """
70
+
71
+ m2 = chirp_mass * (1.0 + mass_ratio) ** 0.2 / mass_ratio**0.6
72
+ m1 = m2 * mass_ratio
73
+
74
+ # # flip m1 m2. For some reason LAL uses this convention for PhenomPv2
75
+ m1, m2 = m2, m1
76
+ s1x, s2x = s2x, s1x
77
+ s1y, s2y = s2y, s1y
78
+ s1z, s2z = s2z, s1z
79
+
80
+ (
81
+ chi1_l,
82
+ chi2_l,
83
+ chip,
84
+ thetaJN,
85
+ alpha0,
86
+ phi_aligned,
87
+ zeta_polariz,
88
+ ) = self.convert_spins(
89
+ m1, m2, f_ref, phiRef, incl, s1x, s1y, s1z, s2x, s2y, s2z
90
+ )
91
+
92
+ phic = 2 * phi_aligned
93
+ q = m2 / m1 # q>=1
94
+ M = m1 + m2
95
+ chi_eff = (m1 * chi1_l + m2 * chi2_l) / M
96
+ chil = (1.0 + q) / q * chi_eff
97
+ eta = m1 * m2 / (M * M)
98
+ eta2 = eta * eta
99
+ Seta = torch.sqrt(1.0 - 4.0 * eta)
100
+ chi = self.chiPN(Seta, eta, chi2_l, chi1_l)
101
+ chi22 = chi2_l * chi2_l
102
+ chi12 = chi1_l * chi1_l
103
+ xi = -1.0 + chi
104
+ m_sec = M * MTSUN_SI
105
+ piM = PI * m_sec
106
+
107
+ omega_ref = piM * f_ref
108
+ logomega_ref = torch.log(omega_ref)
109
+ omega_ref_cbrt = (piM * f_ref) ** (1 / 3) # == v0
110
+ omega_ref_cbrt2 = omega_ref_cbrt * omega_ref_cbrt
111
+
112
+ angcoeffs = self.ComputeNNLOanglecoeffs(q, chil, chip)
113
+
114
+ alphaNNLOoffset = (
115
+ angcoeffs["alphacoeff1"] / omega_ref
116
+ + angcoeffs["alphacoeff2"] / omega_ref_cbrt2
117
+ + angcoeffs["alphacoeff3"] / omega_ref_cbrt
118
+ + angcoeffs["alphacoeff4"] * logomega_ref
119
+ + angcoeffs["alphacoeff5"] * omega_ref_cbrt
120
+ )
121
+
122
+ epsilonNNLOoffset = (
123
+ angcoeffs["epsiloncoeff1"] / omega_ref
124
+ + angcoeffs["epsiloncoeff2"] / omega_ref_cbrt2
125
+ + angcoeffs["epsiloncoeff3"] / omega_ref_cbrt
126
+ + angcoeffs["epsiloncoeff4"] * logomega_ref
127
+ + angcoeffs["epsiloncoeff5"] * omega_ref_cbrt
128
+ )
129
+
130
+ Y2m2 = self.SpinWeightedY(thetaJN, 0, -2, 2, -2)
131
+ Y2m1 = self.SpinWeightedY(thetaJN, 0, -2, 2, -1)
132
+ Y20 = self.SpinWeightedY(thetaJN, 0, -2, 2, -0)
133
+ Y21 = self.SpinWeightedY(thetaJN, 0, -2, 2, 1)
134
+ Y22 = self.SpinWeightedY(thetaJN, 0, -2, 2, 2)
135
+ Y2 = torch.stack((Y2m2, Y2m1, Y20, Y21, Y22))
136
+
137
+ hPhenomDs, diffRDphase = self.PhenomPOneFrequency(
138
+ fs,
139
+ m2,
140
+ m1,
141
+ eta,
142
+ eta2,
143
+ Seta,
144
+ chi2_l,
145
+ chi1_l,
146
+ chi12,
147
+ chi22,
148
+ chip,
149
+ phic,
150
+ M,
151
+ xi,
152
+ dist_mpc,
153
+ )
154
+
155
+ hp, hc = self.PhenomPCoreTwistUp(
156
+ fs,
157
+ hPhenomDs,
158
+ eta,
159
+ chi1_l,
160
+ chi2_l,
161
+ chip,
162
+ M,
163
+ angcoeffs,
164
+ Y2,
165
+ alphaNNLOoffset - alpha0,
166
+ epsilonNNLOoffset,
167
+ )
168
+ t0 = (diffRDphase.unsqueeze(1)) / (2 * PI)
169
+ phase_corr = torch.cos(2 * PI * fs * (t0)) - 1j * torch.sin(
170
+ 2 * PI * fs * (t0)
171
+ )
172
+ M_s = (m1 + m2) * MTSUN_SI
173
+ phase_corr_tc = torch.exp(
174
+ -1j * fs * M_s.unsqueeze(1) * tc.unsqueeze(1)
175
+ )
176
+ hp *= phase_corr * phase_corr_tc
177
+ hc *= phase_corr * phase_corr_tc
178
+
179
+ c2z = torch.cos(2 * zeta_polariz).unsqueeze(1)
180
+ s2z = torch.sin(2 * zeta_polariz).unsqueeze(1)
181
+ hplus = c2z * hp + s2z * hc
182
+ hcross = c2z * hc - s2z * hp
183
+ return hcross, hplus
184
+
185
+ def PhenomPCoreTwistUp(
186
+ self,
187
+ fHz: TensorType,
188
+ hPhenom: TensorType,
189
+ eta: TensorType,
190
+ chi1_l: TensorType,
191
+ chi2_l: TensorType,
192
+ chip: TensorType,
193
+ M: TensorType,
194
+ angcoeffs: Dict[str, TensorType],
195
+ Y2m: TensorType,
196
+ alphaoffset: TensorType,
197
+ epsilonoffset: TensorType,
198
+ ) -> Tuple[TensorType, TensorType]:
199
+ assert angcoeffs is not None
200
+ assert Y2m is not None
201
+ f = fHz * MTSUN_SI * M.unsqueeze(1) # Frequency in geometric units
202
+ q = (1.0 + torch.sqrt(1.0 - 4.0 * eta) - 2.0 * eta) / (2.0 * eta)
203
+ m1 = 1.0 / (1.0 + q) # Mass of the smaller BH for unit total mass M=1.
204
+ m2 = q / (1.0 + q) # Mass of the larger BH for unit total mass M=1.
205
+ Sperp = chip * (
206
+ m2 * m2
207
+ ) # Dimensionfull spin component in the orbital plane.
208
+ # S_perp = S_2_perp chi_eff = m1 * chi1_l + m2 * chi2_l
209
+ # effective spin for M=1
210
+
211
+ SL = chi1_l * m1 * m1 + chi2_l * m2 * m2 # Dimensionfull aligned spin.
212
+
213
+ omega = PI * f
214
+ logomega = torch.log(omega)
215
+ omega_cbrt = (omega) ** (1 / 3)
216
+ omega_cbrt2 = omega_cbrt * omega_cbrt
217
+ alpha = (
218
+ (
219
+ angcoeffs["alphacoeff1"] / omega.mT
220
+ + angcoeffs["alphacoeff2"] / omega_cbrt2.mT
221
+ + angcoeffs["alphacoeff3"] / omega_cbrt.mT
222
+ + angcoeffs["alphacoeff4"] * logomega.mT
223
+ + angcoeffs["alphacoeff5"] * omega_cbrt.mT
224
+ )
225
+ - alphaoffset
226
+ ).mT
227
+
228
+ epsilon = (
229
+ (
230
+ angcoeffs["epsiloncoeff1"] / omega.mT
231
+ + angcoeffs["epsiloncoeff2"] / omega_cbrt2.mT
232
+ + angcoeffs["epsiloncoeff3"] / omega_cbrt.mT
233
+ + angcoeffs["epsiloncoeff4"] * logomega.mT
234
+ + angcoeffs["epsiloncoeff5"] * omega_cbrt.mT
235
+ )
236
+ - epsilonoffset
237
+ ).mT
238
+
239
+ cBetah, sBetah = self.WignerdCoefficients(
240
+ omega_cbrt.mT, SL, eta, Sperp
241
+ )
242
+
243
+ cBetah2 = cBetah * cBetah
244
+ cBetah3 = cBetah2 * cBetah
245
+ cBetah4 = cBetah3 * cBetah
246
+ sBetah2 = sBetah * sBetah
247
+ sBetah3 = sBetah2 * sBetah
248
+ sBetah4 = sBetah3 * sBetah
249
+
250
+ hp_sum = 0
251
+ hc_sum = 0
252
+
253
+ cexp_i_alpha = torch.exp(1j * alpha)
254
+ cexp_2i_alpha = cexp_i_alpha * cexp_i_alpha
255
+ cexp_mi_alpha = 1.0 / cexp_i_alpha
256
+ cexp_m2i_alpha = cexp_mi_alpha * cexp_mi_alpha
257
+ T2m = (
258
+ cexp_2i_alpha.mT * cBetah4.mT * Y2m[0]
259
+ - cexp_i_alpha.mT * 2 * cBetah3.mT * sBetah.mT * Y2m[1]
260
+ + 1
261
+ * torch.sqrt(torch.tensor(6))
262
+ * sBetah2.mT
263
+ * cBetah2.mT
264
+ * Y2m[2]
265
+ - cexp_mi_alpha.mT * 2 * cBetah.mT * sBetah3.mT * Y2m[3]
266
+ + cexp_m2i_alpha.mT * sBetah4.mT * Y2m[4]
267
+ ).mT
268
+ Tm2m = (
269
+ cexp_m2i_alpha.mT * sBetah4.mT * torch.conj(Y2m[0])
270
+ + cexp_mi_alpha.mT
271
+ * 2
272
+ * cBetah.mT
273
+ * sBetah3.mT
274
+ * torch.conj(Y2m[1])
275
+ + 1
276
+ * torch.sqrt(torch.tensor(6))
277
+ * sBetah2.mT
278
+ * cBetah2.mT
279
+ * torch.conj(Y2m[2])
280
+ + cexp_i_alpha.mT * 2 * cBetah3.mT * sBetah.mT * torch.conj(Y2m[3])
281
+ + cexp_2i_alpha.mT * cBetah4.mT * torch.conj(Y2m[4])
282
+ ).mT
283
+ hp_sum = T2m + Tm2m
284
+ hc_sum = 1j * (T2m - Tm2m)
285
+
286
+ eps_phase_hP = torch.exp(-2j * epsilon) * hPhenom / 2.0
287
+
288
+ hp = eps_phase_hP * hp_sum
289
+ hc = eps_phase_hP * hc_sum
290
+
291
+ return hp, hc
292
+
293
+ def PhenomPOneFrequency(
294
+ self,
295
+ fs,
296
+ m1,
297
+ m2,
298
+ eta,
299
+ eta2,
300
+ Seta,
301
+ chi1,
302
+ chi2,
303
+ chi12,
304
+ chi22,
305
+ chip,
306
+ phic,
307
+ M,
308
+ xi,
309
+ dist_mpc,
310
+ ):
311
+ """
312
+ m1, m2: in solar masses
313
+ phic: Orbital phase at peak of the underlying non precessing model
314
+ M: Total mass (Solar masses)
315
+ """
316
+
317
+ M_s = M * MTSUN_SI
318
+ Mf = torch.outer(M_s, fs)
319
+ fRD, _ = self.phP_get_fRD_fdamp(m1, m2, chi1, chi2, chip)
320
+
321
+ phase, _ = self.phenom_d_phase(Mf, m1, m2, eta, eta2, chi1, chi2, xi)
322
+ phase = (phase.mT - (phic + PI / 4.0)).mT
323
+ Amp = self.phenom_d_amp(
324
+ Mf, m1, m2, eta, eta2, Seta, chi1, chi2, chi12, chi22, xi, dist_mpc
325
+ )[0]
326
+ Amp0 = self.get_Amp0(Mf, eta)
327
+ dist_s = dist_mpc * MPC_SEC
328
+ Amp = ((Amp0 * Amp).mT * (M_s**2.0) / dist_s).mT
329
+ # phase -= 2. * phic; # line 1316 ???
330
+ hPhenom = Amp * (torch.exp(-1j * phase))
331
+
332
+ fRDs = torch.outer(
333
+ fRD, torch.linspace(0.5, 1.5, 101, device=fRD.device)
334
+ )
335
+ delta_fRds = torch.median(torch.diff(fRDs, axis=1), axis=1)[0]
336
+ MfRDs = torch.zeros_like(fRDs)
337
+ for i in range(fRD.shape[0]):
338
+ MfRDs[i, :] = torch.outer(M_s, fRDs[i, :])[i, :]
339
+ RD_phase = self.phenom_d_phase(
340
+ MfRDs, m1, m2, eta, eta2, chi1, chi2, xi
341
+ )[0]
342
+ diff = torch.diff(RD_phase, axis=1)
343
+ diffRDphase = (diff[:, 1:] + diff[:, :-1]) / (
344
+ 2 * delta_fRds.unsqueeze(1)
345
+ )
346
+ diffRDphase = -diffRDphase[:, 50]
347
+ # MfRD = torch.outer(M_s, fRD)
348
+ # Dphase = torch.diag(
349
+ # -self.phenom_d_phase(
350
+ # MfRD, m1, m2, eta, eta2, chi1, chi2, xi)[1] * M_s
351
+ # ).view(-1, 1)
352
+ return hPhenom, diffRDphase
353
+
354
+ # Utility functions
355
+
356
+ def interpolate(
357
+ self, x: TensorType, xp: TensorType, fp: TensorType
358
+ ) -> TensorType:
359
+ """One-dimensional linear interpolation for monotonically
360
+ increasing sample points.
361
+
362
+ Returns the one-dimensional piecewise linear interpolant to a function
363
+ with given data points :math:`(xp, fp)`, evaluated at :math:`x`
364
+
365
+ Args:
366
+ x: the :math:`x`-coordinates at which to evaluate the interpolated
367
+ values.
368
+ xp: the :math:`x`-coordinates of data points, must be increasing.
369
+ fp: the :math:`y`-coordinates of data points, same length as `xp`.
370
+
371
+ Returns:
372
+ the interpolated values, same size as `x`.
373
+ """
374
+ original_shape = x.shape
375
+ x = x.flatten()
376
+ xp = xp.flatten()
377
+ fp = fp.flatten()
378
+
379
+ m = (fp[1:] - fp[:-1]) / (xp[1:] - xp[:-1]) # slope
380
+ b = fp[:-1] - (m * xp[:-1])
381
+
382
+ indices = torch.searchsorted(xp, x, right=False) - 1
383
+
384
+ interpolated = m[indices] * x + b[indices]
385
+
386
+ return interpolated.reshape(original_shape)
387
+
388
+ def ROTATEZ(self, angle: TensorType, x, y, z):
389
+ tmp_x = x * torch.cos(angle) - y * torch.sin(angle)
390
+ tmp_y = x * torch.sin(angle) + y * torch.cos(angle)
391
+ return tmp_x, tmp_y, z
392
+
393
+ def ROTATEY(self, angle, x, y, z):
394
+ tmp_x = x * torch.cos(angle) + z * torch.sin(angle)
395
+ tmp_z = -x * torch.sin(angle) + z * torch.cos(angle)
396
+ return tmp_x, y, tmp_z
397
+
398
+ def L2PNR(self, v: TensorType, eta: TensorType) -> TensorType:
399
+ eta2 = eta**2
400
+ x = v**2
401
+ x2 = x**2
402
+ tmp = (
403
+ eta
404
+ * (
405
+ 1.0
406
+ + (1.5 + eta / 6.0) * x
407
+ + (3.375 - (19.0 * eta) / 8.0 - eta2 / 24.0) * x2
408
+ )
409
+ ) / x**0.5
410
+
411
+ return tmp
412
+
413
+ def convert_spins(
414
+ self,
415
+ m1: TensorType,
416
+ m2: TensorType,
417
+ f_ref: float,
418
+ phiRef: TensorType,
419
+ incl: TensorType,
420
+ s1x: TensorType,
421
+ s1y: TensorType,
422
+ s1z: TensorType,
423
+ s2x: TensorType,
424
+ s2y: TensorType,
425
+ s2z: TensorType,
426
+ ) -> Tuple[
427
+ TensorType,
428
+ TensorType,
429
+ TensorType,
430
+ TensorType,
431
+ TensorType,
432
+ TensorType,
433
+ TensorType,
434
+ ]:
435
+ M = m1 + m2
436
+ m1_2 = m1 * m1
437
+ m2_2 = m2 * m2
438
+ eta = m1 * m2 / (M * M) # Symmetric mass-ratio
439
+
440
+ # From the components in the source frame, we can easily determine
441
+ # chi1_l, chi2_l, chip and phi_aligned, which we need to return.
442
+ # We also compute the spherical angles of J,
443
+ # which we need to transform to the J frame
444
+
445
+ # Aligned spins
446
+ chi1_l = s1z # Dimensionless aligned spin on BH 1
447
+ chi2_l = s2z # Dimensionless aligned spin on BH 2
448
+
449
+ # Magnitude of the spin projections in the orbital plane
450
+ S1_perp = m1_2 * torch.sqrt(s1x**2 + s1y**2)
451
+ S2_perp = m2_2 * torch.sqrt(s2x**2 + s2y**2)
452
+
453
+ A1 = 2 + (3 * m2) / (2 * m1)
454
+ A2 = 2 + (3 * m1) / (2 * m2)
455
+ ASp1 = A1 * S1_perp
456
+ ASp2 = A2 * S2_perp
457
+ num = torch.maximum(ASp1, ASp2)
458
+ den = A2 * m2_2 # warning: this assumes m2 > m1
459
+ chip = num / den
460
+
461
+ m_sec = M * MTSUN_SI
462
+ piM = PI * m_sec
463
+ v_ref = (piM * f_ref) ** (1 / 3)
464
+ L0 = M * M * self.L2PNR(v_ref, eta)
465
+ J0x_sf = m1_2 * s1x + m2_2 * s2x
466
+ J0y_sf = m1_2 * s1y + m2_2 * s2y
467
+ J0z_sf = L0 + m1_2 * s1z + m2_2 * s2z
468
+ J0 = torch.sqrt(J0x_sf * J0x_sf + J0y_sf * J0y_sf + J0z_sf * J0z_sf)
469
+
470
+ thetaJ_sf = torch.arccos(J0z_sf / J0)
471
+
472
+ phiJ_sf = torch.arctan2(J0y_sf, J0x_sf)
473
+
474
+ phi_aligned = -phiJ_sf
475
+
476
+ # First we determine kappa
477
+ # in the source frame, the components of N are given in
478
+ # Eq (35c) of T1500606-v6
479
+ Nx_sf = torch.sin(incl) * torch.cos(PI / 2.0 - phiRef)
480
+ Ny_sf = torch.sin(incl) * torch.sin(PI / 2.0 - phiRef)
481
+ Nz_sf = torch.cos(incl)
482
+
483
+ tmp_x = Nx_sf
484
+ tmp_y = Ny_sf
485
+ tmp_z = Nz_sf
486
+
487
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(-phiJ_sf, tmp_x, tmp_y, tmp_z)
488
+ tmp_x, tmp_y, tmp_z = self.ROTATEY(-thetaJ_sf, tmp_x, tmp_y, tmp_z)
489
+
490
+ kappa = -torch.arctan2(tmp_y, tmp_x)
491
+
492
+ # Then we determine alpha0, by rotating LN
493
+ tmp_x, tmp_y, tmp_z = 0, 0, 1
494
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(-phiJ_sf, tmp_x, tmp_y, tmp_z)
495
+ tmp_x, tmp_y, tmp_z = self.ROTATEY(-thetaJ_sf, tmp_x, tmp_y, tmp_z)
496
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(kappa, tmp_x, tmp_y, tmp_z)
497
+
498
+ alpha0 = torch.arctan2(tmp_y, tmp_x)
499
+
500
+ # Finally we determine thetaJ, by rotating N
501
+ tmp_x, tmp_y, tmp_z = Nx_sf, Ny_sf, Nz_sf
502
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(-phiJ_sf, tmp_x, tmp_y, tmp_z)
503
+ tmp_x, tmp_y, tmp_z = self.ROTATEY(-thetaJ_sf, tmp_x, tmp_y, tmp_z)
504
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(kappa, tmp_x, tmp_y, tmp_z)
505
+ Nx_Jf, Nz_Jf = tmp_x, tmp_z
506
+ thetaJN = torch.arccos(Nz_Jf)
507
+
508
+ # Finally, we need to redefine the polarizations:
509
+ # PhenomP's polarizations are defined following Arun et al
510
+ # (arXiv:0810.5336)
511
+ # i.e. projecting the metric onto the P,Q,N triad defined with
512
+ # P=NxJ/|NxJ| (see (2.6) in there).
513
+ # By contrast, the triad X,Y,N used in LAL
514
+ # ("waveframe" in the nomenclature of T1500606-v6)
515
+ # is defined in e.g. eq (35) of this document
516
+ # (via its components in the source frame;
517
+ # note we use the default Omega=Pi/2).
518
+ # Both triads differ from each other by a rotation around N by an angle
519
+ # \zeta and we need to rotate the polarizations accordingly by 2\zeta
520
+
521
+ Xx_sf = -torch.cos(incl) * torch.sin(phiRef)
522
+ Xy_sf = -torch.cos(incl) * torch.cos(phiRef)
523
+ Xz_sf = torch.sin(incl)
524
+ tmp_x, tmp_y, tmp_z = Xx_sf, Xy_sf, Xz_sf
525
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(-phiJ_sf, tmp_x, tmp_y, tmp_z)
526
+ tmp_x, tmp_y, tmp_z = self.ROTATEY(-thetaJ_sf, tmp_x, tmp_y, tmp_z)
527
+ tmp_x, tmp_y, tmp_z = self.ROTATEZ(kappa, tmp_x, tmp_y, tmp_z)
528
+
529
+ # Now the tmp_a are the components of X in the J frame
530
+ # We need the polar angle of that vector in the P,Q basis of Arun et al
531
+ # P = NxJ/|NxJ| and since we put N in the (pos x)z half plane of the J
532
+ # frame
533
+ PArunx_Jf = 0.0
534
+ PAruny_Jf = -1.0
535
+ PArunz_Jf = 0.0
536
+
537
+ # Q = NxP
538
+ QArunx_Jf = Nz_Jf
539
+ QAruny_Jf = 0.0
540
+ QArunz_Jf = -Nx_Jf
541
+
542
+ # Calculate the dot products XdotPArun and XdotQArun
543
+ XdotPArun = tmp_x * PArunx_Jf + tmp_y * PAruny_Jf + tmp_z * PArunz_Jf
544
+ XdotQArun = tmp_x * QArunx_Jf + tmp_y * QAruny_Jf + tmp_z * QArunz_Jf
545
+
546
+ zeta_polariz = torch.arctan2(XdotQArun, XdotPArun)
547
+ return chi1_l, chi2_l, chip, thetaJN, alpha0, phi_aligned, zeta_polariz
548
+
549
+ # TODO: add input and output types
550
+ def SpinWeightedY(self, theta, phi, s, l, m): # noqa: E741
551
+ "copied from SphericalHarmonics.c in LAL"
552
+ if s == -2:
553
+ if l == 2: # noqa: E741
554
+ if m == -2:
555
+ fac = (
556
+ torch.sqrt(torch.tensor(5.0 / (64.0 * PI)))
557
+ * (1.0 - torch.cos(theta))
558
+ * (1.0 - torch.cos(theta))
559
+ )
560
+ elif m == -1:
561
+ fac = (
562
+ torch.sqrt(torch.tensor(5.0 / (16.0 * PI)))
563
+ * torch.sin(theta)
564
+ * (1.0 - torch.cos(theta))
565
+ )
566
+ elif m == 0:
567
+ fac = (
568
+ torch.sqrt(torch.tensor(15.0 / (32.0 * PI)))
569
+ * torch.sin(theta)
570
+ * torch.sin(theta)
571
+ )
572
+ elif m == 1:
573
+ fac = (
574
+ torch.sqrt(torch.tensor(5.0 / (16.0 * PI)))
575
+ * torch.sin(theta)
576
+ * (1.0 + torch.cos(theta))
577
+ )
578
+ elif m == 2:
579
+ fac = (
580
+ torch.sqrt(torch.tensor(5.0 / (64.0 * PI)))
581
+ * (1.0 + torch.cos(theta))
582
+ * (1.0 + torch.cos(theta))
583
+ )
584
+ else:
585
+ raise ValueError(
586
+ f"Invalid mode s={s}, l={l}, m={m} - require |m| <= l"
587
+ )
588
+ return fac * torch.complex(
589
+ torch.cos(torch.tensor(m * phi)),
590
+ torch.sin(torch.tensor(m * phi)),
591
+ )
592
+
593
+ def WignerdCoefficients(
594
+ self, v: TensorType, SL: TensorType, eta: TensorType, Sp: TensorType
595
+ ) -> Tuple[TensorType, TensorType]:
596
+ # We define the shorthand s := Sp / (L + SL)
597
+ L = self.L2PNR(v, eta)
598
+ s = (Sp / (L + SL)).mT
599
+ s2 = s**2
600
+ cos_beta = 1.0 / (1.0 + s2) ** 0.5
601
+ cos_beta_half = ((1.0 + cos_beta) / 2.0) ** 0.5 # cos(beta/2)
602
+ sin_beta_half = ((1.0 - cos_beta) / 2.0) ** 0.5 # sin(beta/2)
603
+
604
+ return cos_beta_half, sin_beta_half
605
+
606
+ def ComputeNNLOanglecoeffs(
607
+ self, q: TensorType, chil: TensorType, chip: TensorType
608
+ ) -> Dict[str, TensorType]:
609
+ m2 = q / (1.0 + q)
610
+ m1 = 1.0 / (1.0 + q)
611
+ dm = m1 - m2
612
+ mtot = 1.0
613
+ eta = m1 * m2 # mtot = 1
614
+ eta2 = eta * eta
615
+ eta3 = eta2 * eta
616
+ eta4 = eta3 * eta
617
+ mtot2 = mtot * mtot
618
+ mtot4 = mtot2 * mtot2
619
+ mtot6 = mtot4 * mtot2
620
+ mtot8 = mtot6 * mtot2
621
+ chil2 = chil * chil
622
+ chip2 = chip * chip
623
+ chip4 = chip2 * chip2
624
+ dm2 = dm * dm
625
+ dm3 = dm2 * dm
626
+ m2_2 = m2 * m2
627
+ m2_3 = m2_2 * m2
628
+ m2_4 = m2_3 * m2
629
+ m2_5 = m2_4 * m2
630
+ m2_6 = m2_5 * m2
631
+ m2_7 = m2_6 * m2
632
+ m2_8 = m2_7 * m2
633
+
634
+ angcoeffs = {}
635
+ angcoeffs["alphacoeff1"] = -0.18229166666666666 - (5 * dm) / (
636
+ 64.0 * m2
637
+ )
638
+
639
+ angcoeffs["alphacoeff2"] = (-15 * dm * m2 * chil) / (
640
+ 128.0 * mtot2 * eta
641
+ ) - (35 * m2_2 * chil) / (128.0 * mtot2 * eta)
642
+
643
+ angcoeffs["alphacoeff3"] = (
644
+ -1.7952473958333333
645
+ - (4555 * dm) / (7168.0 * m2)
646
+ - (15 * chip2 * dm * m2_3) / (128.0 * mtot4 * eta2)
647
+ - (35 * chip2 * m2_4) / (128.0 * mtot4 * eta2)
648
+ - (515 * eta) / 384.0
649
+ - (15 * dm2 * eta) / (256.0 * m2_2)
650
+ - (175 * dm * eta) / (256.0 * m2)
651
+ )
652
+
653
+ angcoeffs["alphacoeff4"] = (
654
+ -(35 * PI) / 48.0
655
+ - (5 * dm * PI) / (16.0 * m2)
656
+ + (5 * dm2 * chil) / (16.0 * mtot2)
657
+ + (5 * dm * m2 * chil) / (3.0 * mtot2)
658
+ + (2545 * m2_2 * chil) / (1152.0 * mtot2)
659
+ - (5 * chip2 * dm * m2_5 * chil) / (128.0 * mtot6 * eta3)
660
+ - (35 * chip2 * m2_6 * chil) / (384.0 * mtot6 * eta3)
661
+ + (2035 * dm * m2 * chil) / (21504.0 * mtot2 * eta)
662
+ + (2995 * m2_2 * chil) / (9216.0 * mtot2 * eta)
663
+ )
664
+
665
+ angcoeffs["alphacoeff5"] = (
666
+ 4.318908476114694
667
+ + (27895885 * dm) / (2.1676032e7 * m2)
668
+ - (15 * chip4 * dm * m2_7) / (512.0 * mtot8 * eta4)
669
+ - (35 * chip4 * m2_8) / (512.0 * mtot8 * eta4)
670
+ - (485 * chip2 * dm * m2_3) / (14336.0 * mtot4 * eta2)
671
+ + (475 * chip2 * m2_4) / (6144.0 * mtot4 * eta2)
672
+ + (15 * chip2 * dm2 * m2_2) / (256.0 * mtot4 * eta)
673
+ + (145 * chip2 * dm * m2_3) / (512.0 * mtot4 * eta)
674
+ + (575 * chip2 * m2_4) / (1536.0 * mtot4 * eta)
675
+ + (39695 * eta) / 86016.0
676
+ + (1615 * dm2 * eta) / (28672.0 * m2_2)
677
+ - (265 * dm * eta) / (14336.0 * m2)
678
+ + (955 * eta2) / 576.0
679
+ + (15 * dm3 * eta2) / (1024.0 * m2_3)
680
+ + (35 * dm2 * eta2) / (256.0 * m2_2)
681
+ + (2725 * dm * eta2) / (3072.0 * m2)
682
+ - (15 * dm * m2 * PI * chil) / (16.0 * mtot2 * eta)
683
+ - (35 * m2_2 * PI * chil) / (16.0 * mtot2 * eta)
684
+ + (15 * chip2 * dm * m2_7 * chil2) / (128.0 * mtot8 * eta4)
685
+ + (35 * chip2 * m2_8 * chil2) / (128.0 * mtot8 * eta4)
686
+ + (375 * dm2 * m2_2 * chil2) / (256.0 * mtot4 * eta)
687
+ + (1815 * dm * m2_3 * chil2) / (256.0 * mtot4 * eta)
688
+ + (1645 * m2_4 * chil2) / (192.0 * mtot4 * eta)
689
+ )
690
+
691
+ angcoeffs["epsiloncoeff1"] = -0.18229166666666666 - (5 * dm) / (
692
+ 64.0 * m2
693
+ )
694
+ angcoeffs["epsiloncoeff2"] = (-15 * dm * m2 * chil) / (
695
+ 128.0 * mtot2 * eta
696
+ ) - (35 * m2_2 * chil) / (128.0 * mtot2 * eta)
697
+ angcoeffs["epsiloncoeff3"] = (
698
+ -1.7952473958333333
699
+ - (4555 * dm) / (7168.0 * m2)
700
+ - (515 * eta) / 384.0
701
+ - (15 * dm2 * eta) / (256.0 * m2_2)
702
+ - (175 * dm * eta) / (256.0 * m2)
703
+ )
704
+ angcoeffs["epsiloncoeff4"] = (
705
+ -(35 * PI) / 48.0
706
+ - (5 * dm * PI) / (16.0 * m2)
707
+ + (5 * dm2 * chil) / (16.0 * mtot2)
708
+ + (5 * dm * m2 * chil) / (3.0 * mtot2)
709
+ + (2545 * m2_2 * chil) / (1152.0 * mtot2)
710
+ + (2035 * dm * m2 * chil) / (21504.0 * mtot2 * eta)
711
+ + (2995 * m2_2 * chil) / (9216.0 * mtot2 * eta)
712
+ )
713
+ angcoeffs["epsiloncoeff5"] = (
714
+ 4.318908476114694
715
+ + (27895885 * dm) / (2.1676032e7 * m2)
716
+ + (39695 * eta) / 86016.0
717
+ + (1615 * dm2 * eta) / (28672.0 * m2_2)
718
+ - (265 * dm * eta) / (14336.0 * m2)
719
+ + (955 * eta2) / 576.0
720
+ + (15 * dm3 * eta2) / (1024.0 * m2_3)
721
+ + (35 * dm2 * eta2) / (256.0 * m2_2)
722
+ + (2725 * dm * eta2) / (3072.0 * m2)
723
+ - (15 * dm * m2 * PI * chil) / (16.0 * mtot2 * eta)
724
+ - (35 * m2_2 * PI * chil) / (16.0 * mtot2 * eta)
725
+ + (375 * dm2 * m2_2 * chil2) / (256.0 * mtot4 * eta)
726
+ + (1815 * dm * m2_3 * chil2) / (256.0 * mtot4 * eta)
727
+ + (1645 * m2_4 * chil2) / (192.0 * mtot4 * eta)
728
+ )
729
+ return angcoeffs
730
+
731
+ def FinalSpin_inplane(
732
+ self,
733
+ m1: TensorType,
734
+ m2: TensorType,
735
+ chi1_l: TensorType,
736
+ chi2_l: TensorType,
737
+ chip: TensorType,
738
+ ) -> TensorType:
739
+ M = m1 + m2
740
+ eta = m1 * m2 / (M * M)
741
+ eta2 = eta * eta
742
+ # m1 > m2, the convention used in phenomD
743
+ # (not the convention of internal phenomP)
744
+ mass_ratio = m1 / m2
745
+ af_parallel = self.FinalSpin0815(eta, eta2, chi1_l, chi2_l)
746
+ Sperp = chip * mass_ratio * mass_ratio
747
+ af = torch.copysign(
748
+ torch.ones_like(af_parallel), af_parallel
749
+ ) * torch.sqrt(Sperp * Sperp + af_parallel * af_parallel)
750
+ return af
751
+
752
+ def phP_get_fRD_fdamp(
753
+ self, m1, m2, chi1_l, chi2_l, chip
754
+ ) -> Tuple[TensorType, TensorType]:
755
+ # m1 > m2 should hold here
756
+ finspin = self.FinalSpin_inplane(m1, m2, chi1_l, chi2_l, chip)
757
+ m1_s = m1 * MTSUN_SI
758
+ m2_s = m2 * MTSUN_SI
759
+ M_s = m1_s + m2_s
760
+ eta_s = m1_s * m2_s / (M_s**2.0)
761
+ eta_s2 = eta_s * eta_s
762
+ Erad = self.PhenomInternal_EradRational0815(
763
+ eta_s, eta_s2, chi1_l, chi2_l
764
+ )
765
+ fRD = self.interpolate(finspin, self.qnmdata_a, self.qnmdata_fring) / (
766
+ 1.0 - Erad
767
+ )
768
+ fdamp = self.interpolate(
769
+ finspin, self.qnmdata_a, self.qnmdata_fdamp
770
+ ) / (1.0 - Erad)
771
+ return fRD / M_s, fdamp / M_s
772
+
773
+ def get_Amp0(self, fM_s: TensorType, eta: TensorType) -> TensorType:
774
+ Amp0 = (
775
+ (2.0 / 3.0 * eta.unsqueeze(1)) ** (1.0 / 2.0)
776
+ * (fM_s) ** (-7.0 / 6.0)
777
+ * PI ** (-1.0 / 6.0)
778
+ )
779
+ return Amp0