Coverage for pesummary/gw/conversions/snr.py: 61.2%

1# Licensed under an MIT style license -- see LICENSE.md 

2 

3import numpy as np 

4from pesummary.utils.decorators import array_input 

5from pesummary.utils.utils import logger 

6from pesummary.gw.pycbc import optimal_snr, compute_the_overlap 

7from pesummary.gw.conversions.angles import _dphi, _dpsi 

8 

9__author__ = [ 

10 "Stephen Fairhurst <stephen.fairhurst@ligo.org>", 

11 "Rhys Green <rhys.green@ligo.org>", 

12 "Charlie Hoy <charlie.hoy@ligo.org>", 

13 "Cameron Mills <joseph.mills@ligo.org>" 

14] 

15 

16 

17@array_input() 

18def _ifo_snr(IFO_abs_snr, IFO_snr_angle): 

19 """Return the matched filter SNR for a given IFO given samples for the 

20 absolute SNR and the angle 

21 """ 

22 return IFO_abs_snr * np.cos(IFO_snr_angle) 

23 

24 

25@array_input() 

26def _ifo_snr_from_real_and_imaginary(IFO_real_snr, IFO_imag_snr): 

27 """Return the matched filter SNR for a given IFO given samples for the 

28 real and imaginary SNR 

29 """ 

30 _complex = IFO_real_snr + IFO_imag_snr * 1j 

31 _abs = np.abs(_complex) 

32 return _ifo_snr(_abs, np.angle(_complex)) 

33 

34 

35@array_input() 

36def network_snr(snrs): 

37 """Return the network SNR for N IFOs 

38 

39 Parameters 

40 ---------- 

41 snrs: list 

42 list of numpy.array objects containing the snrs samples for a particular 

43 IFO 

44 """ 

45 squares = np.square(snrs) 

46 network_snr = np.sqrt(np.sum(squares, axis=0)) 

47 return network_snr 

48 

49 

50@array_input() 

51def network_matched_filter_snr(IFO_matched_filter_snrs, IFO_optimal_snrs): 

52 """Return the network matched filter SNR for a given detector network. Code 

53 adapted from Christopher Berry's python notebook 

54 

55 Parameters 

56 ---------- 

57 IFO_matched_filter_snrs: list 

58 list of matched filter SNRs for each IFO in the network 

59 IFO_optimal_snrs: list 

60 list of optimal SNRs 

61 """ 

62 for num, det_snr in enumerate(IFO_matched_filter_snrs): 

63 complex_snr = np.iscomplex(det_snr) 

64 convert_mf_snr = False 

65 try: 

66 if complex_snr: 

67 convert_mf_snr = True 

68 except ValueError: 

69 if any(_complex for _complex in complex_snr): 

70 convert_mf_snr = True 

71 if convert_mf_snr: 

72 IFO_matched_filter_snrs[num] = np.real(det_snr) 

73 network_optimal_snr = network_snr(IFO_optimal_snrs) 

74 network_matched_filter_snr = np.sum( 

75 [ 

76 mf_snr * opt_snr / network_optimal_snr for mf_snr, opt_snr in zip( 

77 IFO_matched_filter_snrs, IFO_optimal_snrs 

78 ) 

79 ], axis=0 

80 ) 

81 return network_matched_filter_snr 

82 

83 

84def _setup_frequencies(f_low, f_final, f_ref, psd): 

85 """Setup the final frequency and reference frequency to be compatible with 

86 the provided PSD. 

87 

88 Parameters 

89 ---------- 

90 f_low: float 

91 starting frequency you wish to use 

92 f_final: float 

93 final frequency you wish to use 

94 f_ref: float 

95 reference frequency you wish to use 

96 psd: dict 

97 dictionary of pycbc.frequencyseries.FrequencySeries objects containing 

98 the psd. Keys must be the detector identifier, e.g. H1, L1, ... 

99 """ 

100 detectors = list(psd.keys()) 

101 _f_final = psd[detectors[0]].sample_frequencies[-1] 

102 if f_final is None: 

103 f_final = _f_final 

104 elif f_final > _f_final: 

105 logger.warning( 

106 "The provided final frequency: {} does not match the final " 

107 "frequency in the PSD: {}. Using the final frequency stored in the " 

108 "PSD to prevent interpolation errors".format(f_final, _f_final) 

109 ) 

110 f_final = _f_final 

111 if f_ref is None: 

112 logger.warning("No reference frequency provided. Using f_low as default") 

113 f_ref = f_low 

114 elif isinstance(f_ref, (list, np.ndarray)): 

115 f_ref = f_ref[0] 

116 return f_final, f_ref 

117 

118 

119def _setup_psd(psd, psd_default, df=1. / 256, **psd_default_kwargs): 

120 """Setup the PSD dictionary. If the provided PSD is empty, construct a 

121 PSD based on the default, this could either be analytic or not. If the 

122 provided PSD is not empty, simply return the provided PSD unchanged. 

123 

124 Parameters 

125 ---------- 

126 psd: dict 

127 dictionary containing the psd. Keys are the IFO and items are a 

128 pycbc.frequencyseries.FrequencySeries object 

129 psd_default: str/dict 

130 The default PSD to use. This can either be a string describing the 

131 analytic PSD or a dictionary with keyes showing the IFO and items 

132 either a pesummary.gw.file.psd.PSD or a 

133 pycbc.frequencyseries.FrequencySeries object 

134 psd_default_kwargs: dict, optional 

135 kwargs to pass to pesummary.gw.file.psd.pycbc_default_psd when 

136 a PSD is constructured based on the analytic default 

137 """ 

138 from pesummary.gw.file.psd import PSD 

139 ANALYTIC = False 

140 condition = isinstance(psd_default, dict) and len(psd_default) and ( 

141 all(isinstance(value, PSD) for ifo, value in psd_default.items()) 

142 ) 

143 f_low = psd_default_kwargs.get("f_low", None) 

144 f_final = psd_default_kwargs.get("f_final", None) 

145 psd_default_kwargs.update({"df": df}) 

146 frequency_cond = any(param is None for param in [f_low, f_final]) 

147 if psd == {} and condition: 

148 if frequency_cond: 

149 raise ValueError( 

150 "Please provide f_low and f_final as keyword arguments" 

151 ) 

152 for ifo, data in psd_default.items(): 

153 psd[ifo] = data.to_pycbc( 

154 f_low, f_high=f_final, f_high_override=True 

155 ) 

156 elif psd == {} and isinstance(psd_default, dict) and len(psd_default): 

157 for ifo, data in psd_default.items(): 

158 psd[ifo] = data 

159 elif psd == {}: 

160 from pesummary.gw.pycbc import pycbc_default_psd 

161 if isinstance(psd_default, dict): 

162 from pesummary import conf 

163 psd_default = conf.psd 

164 ANALYTIC = True 

165 psd_default_kwargs.update({"psd": psd_default}) 

166 psd = pycbc_default_psd(**psd_default_kwargs) 

167 

168 detectors = list(psd.keys()) 

169 if df != psd[detectors[0]].delta_f: 

170 from pesummary.gw.file.psd import PSDDict 

171 logger.warning( 

172 "Provided PSD has df={} and {} has been specified. Interpolation " 

173 "will be used".format(psd[detectors[0]].delta_f, df) 

174 ) 

175 if isinstance(psd, PSDDict): 

176 psd = psd.interpolate(f_low, df) 

177 else: 

178 from pesummary.gw.pycbc import interpolate_psd 

179 psd = { 

180 ifo: interpolate_psd(psd[ifo], f_low, df) for ifo in 

181 psd.keys() 

182 } 

183 return psd, ANALYTIC 

184 

185 

186def _make_waveform( 

187 approx, theta_jn, phi_jl, phase, psi_J, mass_1, mass_2, tilt_1, tilt_2, 

188 phi_12, a_1, a_2, beta, distance, apply_detector_response=True, **kwargs 

189): 

190 """Generate a frequency domain waveform 

191 

192 Parameters 

193 ---------- 

194 approx: str 

195 Name of the approximant you wish to use when generating the waveform 

196 theta_jn: float 

197 Angle between the total angular momentum and the line of sight 

198 phi_jl: float 

199 Azimuthal angle of the total orbital angular momentum around the 

200 total angular momentum 

201 phase: float 

202 The phase of the binary at coaelescence 

203 mass_1: float 

204 Primary mass of the binary 

205 mass_2: float 

206 Secondary mass of the binary 

207 tilt_1: float 

208 The angle between the total orbital angular momentum and the primary 

209 spin 

210 tilt_2: float 

211 The angle between the total orbital angular momentum and the primary 

212 spin 

213 phi_12: float 

214 The angle between the primary spin and the secondary spin 

215 a_1: float 

216 The spin magnitude on the larger object 

217 a_2: float 

218 The spin magnitude on the secondary object 

219 beta: float 

220 The opening angle of the system. Defined as the angle between the 

221 orbital angular momentum, L, and the total angular momentum J. 

222 apply_detector_response: Bool, optional 

223 if True apply an effective detector response and return 

224 fp * hp + fc * hc else return hp, hc. Default True 

225 **kwargs: dict 

226 All additional kwargs are passed to the 

227 pesummary.gw.waveform.fd_waveform function 

228 """ 

229 from pesummary.gw.waveform import fd_waveform 

230 _samples = { 

231 "theta_jn": [theta_jn], "phi_jl": [phi_jl], "phase": [phase], 

232 "mass_1": [mass_1], "mass_2": [mass_2], "tilt_1": [tilt_1], 

233 "tilt_2": [tilt_2], "phi_12": [phi_12], "a_1": [a_1], 

234 "a_2": [a_2], "luminosity_distance": [distance] 

235 } 

236 waveforms = fd_waveform( 

237 _samples, approx, kwargs.get("df", 1. / 256), 

238 kwargs.get("f_low", 20.), kwargs.get("f_final", 1024.), 

239 f_ref=kwargs.get("f_ref", 20.), ind=0, pycbc=True, 

240 mode_array=kwargs.get("mode_array", None) 

241 ) 

242 hp, hc = waveforms["h_plus"], waveforms["h_cross"] 

243 if kwargs.get("flen", None) is not None: 

244 flen = kwargs.get("flen") 

245 hp.resize(flen) 

246 hc.resize(flen) 

247 if not apply_detector_response: 

248 return hp, hc 

249 dpsi = _dpsi(theta_jn, phi_jl, beta) 

250 fp = np.cos(2 * (psi_J - dpsi)) 

251 fc = -1. * np.sin(2 * (psi_J - dpsi)) 

252 h = (fp * hp + fc * hc) 

253 h *= np.exp(2j * _dphi(theta_jn, phi_jl, beta)) 

254 return h 

255 

256 

257def _calculate_precessing_harmonics( 

258 mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, phi_12, beta, distance, 

259 harmonics=[0, 1], approx="IMRPhenomPv2", **kwargs 

260): 

261 """Decompose a precessing waveform into a series of harmonics as defined 

262 in Fairhurst et al. arXiv:1908.05707 

263 

264 Parameters 

265 ---------- 

266 mass_1: np.ndarray 

267 Primary mass of the bianry 

268 mass_2: np.ndarray 

269 Secondary mass of the binary 

270 a_1: np.ndarray 

271 The spin magnitude on the larger object 

272 a_2: np.ndarray 

273 The spin magnitude on the secondary object 

274 tilt_1: np.ndarray 

275 The angle between the total orbital angular momentum and the primary 

276 spin 

277 tilt_2: np.ndarray 

278 The angle between the total orbital angular momentum and the secondary 

279 spin 

280 phi_12: np.ndarray 

281 The angle between the primary spin and the secondary spin 

282 beta: np.ndarray 

283 The angle between the total angular momentum and the total orbital 

284 angular momentum 

285 harmonics: list, optional 

286 List of harmonics which you wish to calculate. Default [0, 1] 

287 approximant: str, optional 

288 Approximant to use for the decomposition. Default IMRPhenomPv2 

289 """ 

290 harm = {} 

291 if (0 in harmonics) or (4 in harmonics): 

292 h0 = _make_waveform( 

293 approx, 0., 0., 0., 0., mass_1, mass_2, tilt_1, tilt_2, 

294 phi_12, a_1, a_2, beta, distance, **kwargs 

295 ) 

296 hpi4 = _make_waveform( 

297 approx, 0., 0., np.pi / 4, np.pi / 4, mass_1, mass_2, 

298 tilt_1, tilt_2, phi_12, a_1, a_2, beta, distance, **kwargs 

299 ) 

300 if (0 in harmonics): 

301 harm[0] = (h0 - hpi4) / 2 

302 if (4 in harmonics): 

303 harm[4] = (h0 + hpi4) / 2 

304 if (1 in harmonics) or (3 in harmonics): 

305 h0 = _make_waveform( 

306 approx, np.pi / 2, 0., np.pi / 4, np.pi / 4, mass_1, mass_2, 

307 tilt_1, tilt_2, phi_12, a_1, a_2, beta, distance, 

308 **kwargs 

309 ) 

310 hpi2 = _make_waveform( 

311 approx, np.pi / 2, np.pi / 2, 0., np.pi / 4, mass_1, mass_2, 

312 tilt_1, tilt_2, phi_12, a_1, a_2, beta, distance, 

313 **kwargs 

314 ) 

315 if (1 in harmonics): 

316 harm[1] = -1. * (h0 + hpi2) / 4 

317 if (3 in harmonics): 

318 harm[3] = -1. * (h0 - hpi2) / 4 

319 if (2 in harmonics): 

320 h0 = _make_waveform( 

321 approx, np.pi / 2, 0., 0., 0., mass_1, mass_2, 

322 tilt_1, tilt_2, phi_12, a_1, a_2, beta, 

323 distance, **kwargs 

324 ) 

325 hpi2 = _make_waveform( 

326 approx, np.pi / 2, np.pi / 2, 0., 0., mass_1, mass_2, 

327 tilt_1, tilt_2, phi_12, a_1, a_2, beta, distance, 

328 **kwargs 

329 ) 

330 harm[2] = (h0 + hpi2) / 6 

331 return harm 

332 

333 

334def _make_waveform_from_precessing_harmonics( 

335 harmonic_dict, theta_jn, phi_jl, phase, f_plus_j, f_cross_j 

336): 

337 """Generate waveform for a binary merger with given precessing harmonics and 

338 orientation 

339 

340 Parameters 

341 ---------- 

342 harmonic_dict: dict 

343 harmonics to include 

344 theta_jn: np.ndarray 

345 the angle between total angular momentum and line of sight 

346 phi_jl: np.ndarray 

347 the initial polarization phase 

348 phase: np.ndarray 

349 the initial orbital phase 

350 psi_J: np.ndarray 

351 the polarization angle in the J-aligned frame 

352 f_plus_j: np.ndarray 

353 The Detector plus response function as defined using the J-aligned frame 

354 f_cross_j: np.ndarray 

355 The Detector cross response function as defined using the J-aligned 

356 frame 

357 """ 

358 A = _harmonic_amplitudes( 

359 theta_jn, phi_jl, f_plus_j, f_cross_j, harmonic_dict 

360 ) 

361 h_app = 0 

362 for k, harm in harmonic_dict.items(): 

363 if h_app: 

364 h_app += A[k] * harm 

365 else: 

366 h_app = A[k] * harm 

367 h_app *= np.exp(2j * phase + 2j * phi_jl) 

368 return h_app 

369 

370 

371def _harmonic_amplitudes( 

372 theta_jn, phi_jl, f_plus_j, f_cross_j, harmonics=[0, 1] 

373): 

374 """Calculate the amplitudes of the precessing harmonics as a function of 

375 orientation 

376 

377 Parameters 

378 ---------- 

379 theta_jn: np.ndarray 

380 the angle between J and line of sight 

381 phi_jl: np.ndarray 

382 the precession phase 

383 f_plus_j: np.ndarray 

384 The Detector plus response function as defined using the J-aligned frame 

385 f_cross_j: np.ndarray 

386 The Detector cross response function as defined using the J-aligned 

387 frame 

388 harmonics: list, optional 

389 The list of harmonics you wish to return. Default is [0, 1] 

390 """ 

391 amp = {} 

392 if 0 in harmonics: 

393 amp[0] = ( 

394 (1 + np.cos(theta_jn)**2) / 2 * f_plus_j 

395 - 1j * np.cos(theta_jn) * f_cross_j 

396 ) 

397 if 1 in harmonics: 

398 amp[1] = 2 * np.exp(-1j * phi_jl) * ( 

399 np.sin(theta_jn) * np.cos(theta_jn) * f_plus_j 

400 - 1j * np.sin(theta_jn) * f_cross_j 

401 ) 

402 if 2 in harmonics: 

403 amp[2] = 3 * np.exp(-2j * phi_jl) * (np.sin(theta_jn)**2) * f_plus_j 

404 if 3 in harmonics: 

405 amp[3] = 2 * np.exp(-3j * phi_jl) * ( 

406 -np.sin(theta_jn) * np.cos(theta_jn) * f_plus_j 

407 - 1j * np.sin(theta_jn) * f_cross_j 

408 ) 

409 if 4 in harmonics: 

410 amp[4] = np.exp(-4j * phi_jl) * ( 

411 (1 + np.cos(theta_jn)**2) / 2 * f_plus_j 

412 + 1j * np.cos(theta_jn) * f_cross_j 

413 ) 

414 return amp 

415 

416 

417def _prec_ratio(theta_jn, phi_jl, psi_J, b_bar, ra, dec, time, detector): 

418 """Calculate the ratio between the leading and first precession terms: Zeta 

419 as defined in Fairhurst et al. arXiv:1908.05707 

420 

421 Parameters 

422 ---------- 

423 theta_jn: float 

424 The angle between the total angular momentum and the line of sight 

425 phi_jl: np.ndarray 

426 the precession phase 

427 psi_J: float 

428 The polarization of the binary defined with respect to the total 

429 angular momentum 

430 b_bar: float 

431 Tangent of the average angle between the total angular momentum and the 

432 total orbital angular momentum during inspiral 

433 ra: float 

434 The right ascension of the binary 

435 dec: float 

436 The declinartion of the binary 

437 time: float 

438 The merger time of the binary 

439 detector: str 

440 The name of the detector you wish to calculate the ratio for 

441 """ 

442 from pesummary.gw.waveform import antenna_response 

443 samples = {"ra": [ra], "dec": [dec], "psi": [psi_J], "geocent_time": [time]} 

444 f_plus, f_cross = antenna_response(samples, detector) 

445 ratio = _prec_ratio_plus_cross(theta_jn, phi_jl, f_plus, f_cross, b_bar) 

446 return ratio 

447 

448 

449def _prec_ratio_plus_cross(theta_jn, phi_jl, f_plus, f_cross, b_bar): 

450 """Calculate the ratio between the leading and first precession harmonics 

451 given an antenna response pattern. Zeta as defined in Fairhurst et al. 

452 arXiv:1908.05707 

453 

454 Parameters 

455 ---------- 

456 theta_jn: float/np.ndarray 

457 The angle between the total angular momentum and the line of sight 

458 phi_jl: np.ndarray 

459 the precession phase 

460 f_plus: float/np.ndarray 

461 The plus polarization factor for a given sky location / orientation 

462 f_cross: float/np.ndarray 

463 The cross polarization factor for a given sky location / orientation 

464 b_bar: float/np.ndarray 

465 Tangent of the average angle between the total angular momentum and the 

466 total orbital angular momentum during inspiral 

467 """ 

468 amplitudes = _harmonic_amplitudes( 

469 theta_jn, phi_jl, f_plus, f_cross, harmonics=[0, 1] 

470 ) 

471 A0 = amplitudes[0] 

472 A1 = amplitudes[1] 

473 return b_bar * A1 / A0 

474 

475 

476@array_input( 

477 ignore_kwargs=[ 

478 "f_low", "f_final", "psd", "approx", "f_ref", "return_data_used", 

479 "multi_process", "df", "multipole" 

480 ] 

481) 

482def multipole_snr( 

483 mass_1, mass_2, spin_1z, spin_2z, psi, iota, ra, dec, time, distance, phase, 

484 f_low=20., f_final=1024., psd={}, approx="IMRPhenomXHM", f_ref=None, 

485 return_data_used=False, multi_process=1, df=1. / 256, 

486 multipole=[21, 33, 44], psd_default="aLIGOZeroDetHighPower" 

487): 

488 """Calculate the multipole snr as defined in Mills et al. 

489 arXiv: 

490 

491 Parameters 

492 ---------- 

493 mass_1: float/np.ndarray 

494 Primary mass of the bianry 

495 mass_2: float/np.ndarray 

496 Secondary mass of the binary 

497 spin_1z: float/np.ndarray 

498 The primary spin aligned with the total orbital angular momentum 

499 spin_2z: float/np.ndarray 

500 The secondary spin aligned with the total orbital angular momentum 

501 psi: float/np.ndarray 

502 The polarization angle of the binary 

503 iota: float/np.ndarray 

504 The angle between the total orbital angular momentum and the line of 

505 sight 

506 ra: float/np.ndarray 

507 The right ascension of the source 

508 dec: float/np.ndarray 

509 The declination of the source 

510 time: float/np.ndarray 

511 The merger time of the binary 

512 distance: float/np.ndarray 

513 The luminosity distance of the source 

514 phase: float/np.ndarray 

515 The phase of the source 

516 f_low: float, optional 

517 Low frequency to use for integration. Default 20Hz 

518 f_final: float, optional 

519 Final frequency to use for integration. Default 1024Hz 

520 psd: dict, optional 

521 Dictionary of pycbc.types.frequencyseries.FrequencySeries objects, one 

522 for each detector. Default is to use the aLIGOZeroDetHighPower PSD 

523 approx: str, optional 

524 The aligned spin higher order multipole approximant you wish to use. 

525 Default IMRPhenomXHM 

526 f_ref: float, optional 

527 Reference frequency where the spins are defined. Default is f_low 

528 return_data_used: Bool, optional 

529 if True, return a dictionary containing information about what data was 

530 used. Default False 

531 multi_process: int, optional 

532 The number of cpus to use when computing the precessing_snr. Default 1 

533 df: float, optional 

534 Frequency spacing between frequency samples. Default 1./256 

535 multipole: list, optional 

536 List of multipoles to calculate the SNR for. Default [21, 33, 44] 

537 """ 

538 from pesummary.gw.waveform import antenna_response 

539 from pesummary.utils.utils import iterator 

540 import multiprocessing 

541 

542 if isinstance(f_low, (list, np.ndarray)): 

543 f_low = f_low[0] 

544 if any(mm not in [21, 33, 44] for mm in multipole): 

545 raise ValueError( 

546 "Currently, only the SNR in the (l, m) = [(2, 1), (3, 3), (4, 4)] " 

547 "multipoles can be calculated. Please provide any multipole within " 

548 "multipole=[21, 33, 44]" 

549 ) 

550 multipole += [22] 

551 psd, ANALYTIC = _setup_psd( 

552 psd, psd_default, mass_1=mass_1, mass_2=mass_2, spin_1z=spin_1z, 

553 spin_2z=spin_2z, f_low=f_low, detectors=["H1", "L1"], f_final=f_final, 

554 df=df 

555 ) 

556 f_final, f_ref = _setup_frequencies(f_low, f_final, f_ref, psd) 

557 flen = int(f_final / df) + 1 

558 _samples = {"ra": ra, "dec": dec, "psi": psi, "geocent_time": time} 

559 detectors = list(psd.keys()) 

560 antenna = { 

561 detector: antenna_response(_samples, detector) for detector in detectors 

562 } 

563 _f_plus = {key: value[0] for key, value in antenna.items()} 

564 _f_cross = {key: value[1] for key, value in antenna.items()} 

565 with multiprocessing.Pool(multi_process) as pool: 

566 f_plus = np.array( 

567 [dict(zip(_f_plus, item)) for item in zip(*_f_plus.values())] 

568 ) 

569 f_cross = np.array( 

570 [dict(zip(_f_cross, item)) for item in zip(*_f_cross.values())] 

571 ) 

572 args = np.array([ 

573 mass_1, mass_2, spin_1z, spin_2z, psi, iota, ra, dec, time, distance, 

574 phase, [f_low] * len(mass_1), [f_final] * len(mass_1), 

575 [psd] * len(mass_1), [approx] * len(mass_1), [f_ref] * len(mass_1), 

576 [df] * len(mass_1), [flen] * len(mass_1), [multipole] * len(mass_1), 

577 f_plus, f_cross 

578 ], dtype=object).T 

579 rho_hm = np.array( 

580 list( 

581 iterator( 

582 pool.imap(_wrapper_for_multipole_snr, args), tqdm=True, 

583 desc="Calculating rho_hm", logger=logger, total=len(mass_1) 

584 ) 

585 ), dtype=object 

586 ) 

587 rho_hm = np.asarray(np.nan_to_num(rho_hm, 0).T, dtype=np.float64) 

588 rho_hm = np.sqrt(rho_hm) 

589 if return_data_used: 

590 psd_used = "stored" if not ANALYTIC else list(psd.values())[0].__name__ 

591 data_used = {"psd": psd_used, "approximant": approx, "f_final": f_final} 

592 return rho_hm, data_used 

593 return rho_hm 

594 

595 

596@array_input( 

597 ignore_kwargs=[ 

598 "f_low", "psd", "approx", "f_final", "f_ref", "return_data_used", 

599 "multi_process", "duration", "df", "psd_default", "debug" 

600 ] 

601) 

602def precessing_snr( 

603 mass_1, mass_2, beta, psi_J, a_1, a_2, tilt_1, tilt_2, phi_12, theta_jn, 

604 ra, dec, time, phi_jl, distance, phase, f_low=20., psd={}, spin_1z=None, 

605 spin_2z=None, chi_eff=None, approx="IMRPhenomPv2", f_final=1024., 

606 f_ref=None, return_data_used=False, multi_process=1, duration=None, 

607 df=1. / 256, psd_default="aLIGOZeroDetHighPower", debug=True 

608): 

609 """Calculate the precessing snr as defined in Fairhurst et al. 

610 arXiv:1908.05707 

611 

612 Parameters 

613 ---------- 

614 mass_1: np.ndarray 

615 Primary mass of the bianry 

616 mass_2: np.ndarray 

617 Secondary mass of the binary 

618 beta: np.ndarray 

619 The angle between the total angular momentum and the total orbital 

620 angular momentum 

621 psi_J: np.ndarray 

622 The polarization angle as defined with respect to the total angular 

623 momentum 

624 a_1: np.ndarray 

625 The spin magnitude on the larger object 

626 a_2: np.ndarray 

627 The spin magnitude on the secondary object 

628 tilt_1: np.ndarray 

629 The angle between the total orbital angular momentum and the primary 

630 spin 

631 tilt_2: np.ndarray 

632 The angle between the total orbital angular momentum and the secondary 

633 spin 

634 phi_12: np.ndarray 

635 The angle between the primary spin and the secondary spin 

636 theta_jn: np.ndarray 

637 The angle between the total orbital angular momentum and the line of 

638 sight 

639 ra: np.ndarray 

640 The right ascension of the source 

641 dec: np.ndarray 

642 The declination of the source 

643 time: np.ndarray 

644 The merger time of the binary 

645 phi_jl: np.ndarray 

646 the precession phase 

647 f_low: float, optional 

648 The low frequency cut-off to use for integration. Default is 20Hz 

649 psd: dict, optional 

650 Dictionary of pycbc.types.frequencyseries.FrequencySeries objects, one 

651 for each detector. Default is to use the aLIGOZeroDetHighPower PSD 

652 spin_1z: np.ndarray, optional 

653 The primary spin aligned with the total orbital angular momentum 

654 spin_2z: np.ndarray, optional 

655 The secondary spin aligned with the total orbital angular momentum 

656 chi_eff: np.ndarray, optional 

657 Effective spin of the binary 

658 approx: str, optional 

659 The approximant you wish to use. Default IMRPhenomPv2 

660 f_final: float, optional 

661 Final frequency to use for integration. Default 1024Hz 

662 f_ref: float, optional 

663 Reference frequency where the spins are defined. Default is f_low 

664 return_data_used: Bool, optional 

665 if True, return a dictionary containing information about what data was 

666 used. Default False 

667 multi_process: int, optional 

668 The number of cpus to use when computing the precessing_snr. Default 1 

669 duration: float, optional 

670 maximum IMR duration to use to estimate delta_f when PSD is not 

671 provided. 

672 debug: Bool, optional 

673 if True, return posteriors for b_bar and the overlap between the 0th 

674 and 1st harmonics. These are useful for debugging. 

675 """ 

676 from pesummary.gw.waveform import antenna_response 

677 from pesummary.utils.utils import iterator 

678 import multiprocessing 

679 

680 if isinstance(f_low, (list, np.ndarray)): 

681 f_low = f_low[0] 

682 psd, ANALYTIC = _setup_psd( 

683 psd, psd_default, mass_1=mass_1, mass_2=mass_2, spin_1z=spin_1z, 

684 spin_2z=spin_2z, chi_eff=chi_eff, f_low=f_low, duration=duration, 

685 detectors=["H1", "L1"], f_final=f_final, df=df 

686 ) 

687 f_final, f_ref = _setup_frequencies(f_low, f_final, f_ref, psd) 

688 flen = int(f_final / df) + 1 

689 _samples = {"ra": ra, "dec": dec, "psi": psi_J, "geocent_time": time} 

690 detectors = list(psd.keys()) 

691 antenna = { 

692 detector: antenna_response(_samples, detector) for detector in detectors 

693 } 

694 _f_plus_j = {key: value[0] for key, value in antenna.items()} 

695 _f_cross_j = {key: value[1] for key, value in antenna.items()} 

696 with multiprocessing.Pool(multi_process) as pool: 

697 f_plus_j = np.array( 

698 [dict(zip(_f_plus_j, item)) for item in zip(*_f_plus_j.values())] 

699 ) 

700 f_cross_j = np.array( 

701 [dict(zip(_f_cross_j, item)) for item in zip(*_f_cross_j.values())] 

702 ) 

703 dphi = _dphi(theta_jn, phi_jl, beta) 

704 args = np.array([ 

705 mass_1, mass_2, a_1, a_2, tilt_1, tilt_2, phi_12, 

706 theta_jn, beta, psi_J, ra, dec, time, [approx] * len(mass_1), 

707 [psd] * len(mass_1), [detectors] * len(mass_1), phi_jl, distance, 

708 phase - dphi, f_plus_j, f_cross_j, [f_low] * len(mass_1), 

709 [df] * len(mass_1), [f_final] * len(mass_1), [flen] * len(mass_1), 

710 [f_ref] * len(mass_1), [debug] * len(mass_1) 

711 ], dtype=object).T 

712 

713 rho_p = np.array( 

714 list( 

715 iterator( 

716 pool.imap(_wrapper_for_precessing_snr, args), tqdm=True, 

717 desc="Calculating rho_p", logger=logger, total=len(mass_1) 

718 ) 

719 ), dtype=object 

720 ) 

721 

722 if debug: 

723 rho_ps, b_bars, overlaps, snrs = {}, {}, {}, {} 

724 for num, _dict in enumerate([rho_ps, b_bars, overlaps, snrs]): 

725 for key in rho_p[0][0]: 

726 _dict[key] = np.nan_to_num( 

727 [dictionary[num][key] for dictionary in rho_p], 0 

728 ) 

729 _return = [ 

730 np.sqrt(np.sum([_dict[i] for i in detectors], axis=0)) if num == 0 

731 or num == 3 else np.mean([_dict[i] for i in detectors], axis=0) for 

732 num, _dict in enumerate([rho_ps, b_bars, overlaps, snrs]) 

733 ] 

734 else: 

735 rho_p = { 

736 key: np.nan_to_num([dictionary[key] for dictionary in rho_p], 0) for 

737 key in rho_p[0] 

738 } 

739 _return = np.sqrt(np.sum([rho_p[i] for i in detectors], axis=0)) 

740 if return_data_used: 

741 psd_used = "stored" if not ANALYTIC else list(psd.values())[0].__name__ 

742 data_used = {"psd": psd_used, "approximant": approx, "f_final": f_final} 

743 return _return, data_used 

744 return _return 

745 

746 

747def _wrapper_for_multipole_snr(args): 

748 """Wrapper function for _multipole_snr for a pool of workers 

749 

750 Parameters 

751 ---------- 

752 args: tuple 

753 All args passed to _precessing_snr 

754 """ 

755 return _multipole_snr(*args) 

756 

757 

758def _wrapper_for_precessing_snr(args): 

759 """Wrapper function for _precessing_snr for a pool of workers 

760 

761 Parameters 

762 ---------- 

763 args: tuple 

764 All args passed to _precessing_snr 

765 """ 

766 return _precessing_snr(*args) 

767 

768 

769def _calculate_b_bar( 

770 harmonic_dict, psd_dict, low_frequency_cutoff=20., 

771 high_frequency_cutoff=1024., return_snrs=False 

772): 

773 """Calculate the tangent of half the opening angle as defined in 

774 Fairhurst et al. arXiv:1908.05707 

775 

776 Parameters 

777 ---------- 

778 harmonic_dict: dict 

779 dictionary of precessing harmonics. Key is the harmonic number 

780 (0, 1, 2...) and item is the 

781 `pycbc.types.frequencyseries.FrequencySeries` object 

782 psd_dict: dict 

783 dictionary of psds. Key is the IFO and item is the 

784 `pycbc.types.frequencyseries.FrequencySeries` object 

785 low_frequency_cutoff: float, optional 

786 Low frequency-cutoff to use for integration. Default 20Hz 

787 high_frequency_cutoff: float, optional 

788 The final frequency to use for integration. Default 1024Hz 

789 return_snrs: Bool, optional 

790 if True, return the snrs of each harmonic 

791 """ 

792 _optimal_snr = lambda harmonic, psd: optimal_snr( 

793 harmonic, psd, low_frequency_cutoff=low_frequency_cutoff, 

794 high_frequency_cutoff=high_frequency_cutoff 

795 ) 

796 rhos = { 

797 detector: { 

798 0: _optimal_snr(harmonic_dict[0], psd_dict[detector]), 

799 1: _optimal_snr(harmonic_dict[1], psd_dict[detector]) 

800 } for detector in psd_dict.keys() 

801 } 

802 b_bar = {detector: snr[1] / snr[0] for detector, snr in rhos.items()} 

803 if return_snrs: 

804 return b_bar, rhos 

805 return b_bar 

806 

807 

808def _mode_array_map(mode_key, approx): 

809 """Return the mode_array in pycbc format for requested mode 

810 

811 Parameters 

812 ---------- 

813 mode_key: int/str 

814 Mode key e.g. 22/'22'. 

815 approx: str 

816 Waveform approximant. 

817 

818 Returns 

819 ------- 

820 mode_array: list 

821 pesummary.gw.waveform.fd_waveform appropriate list of lm modes. 

822 """ 

823 mode_array_dict = { 

824 "22": [[2, 2], [2, -2]],