Coverage for pesummary/gw/conversions/tgr.py: 75.3%

166 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-12-09 22:34 +0000

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

2 

3from pesummary.utils.utils import logger 

4import numpy as np 

5from scipy.stats import gaussian_kde 

6import multiprocessing 

7from pesummary.utils.probability_dict import ProbabilityDict2D 

8from pesummary.utils.bounded_interp import RectBivariateSpline 

9 

10__author__ = [ 

11 "Aditya Vijaykumar <aditya.vijaykumar@ligo.org>", 

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

13] 

14 

15 

16def _wrapper_for_multiprocessing_kde(kde, *args): 

17 """Wrapper to evaluate a KDE on multiple cpus 

18 

19 Parameters 

20 ---------- 

21 kde: func 

22 KDE you wish to evaluate 

23 *args: tuple 

24 all args are passed to the KDE 

25 """ 

26 _reshape = (len(args[0]), len(args[1])) 

27 yy, xx = np.meshgrid(args[1], args[0]) 

28 _args = [xx.ravel(), yy.ravel()] 

29 return kde(_args).reshape(_reshape) 

30 

31 

32def _wrapper_for_multiprocessing_interp(interp, *args): 

33 """Wrapper to evaluate an interpolant on multiple cpus 

34 

35 Parameters 

36 ---------- 

37 interp: func 

38 interpolant you wish to use 

39 *args: tuple 

40 all args are passed to the interpolant 

41 """ 

42 return interp(*args).T 

43 

44 

45def _imrct_deviation_parameters_integrand_vectorized( 

46 final_mass, 

47 final_spin, 

48 v1, 

49 v2, 

50 P_final_mass_final_spin_i_interp_object, 

51 P_final_mass_final_spin_r_interp_object, 

52 multi_process=1, 

53 wrapper_function_for_multiprocess=None, 

54): 

55 """Compute the integrand of P(delta_final_mass/final_mass_bar, 

56 delta_final_spin/final_spin_bar). 

57 

58 Parameters 

59 ---------- 

60 final_mass: np.ndarray 

61 vector of values of final mass 

62 final_spin: np.ndarray 

63 vector of values of final spin 

64 v1: np.ndarray 

65 array of delta_final_mass/final_mass_bar values 

66 v2: np.ndarray 

67 array of delta_final_spin/final_spin_bar values 

68 P_final_mass_final_spin_i_interp_object: 

69 interpolated P_i(final_mass, final_spin) 

70 P_final_mass_final_spin_r_interp_object: 

71 interpolated P_r(final_mass, final_spin) 

72 multi_process: int 

73 Number of parallel processes. Default: 1 

74 wrapper_function_for_multiprocess: method 

75 Wrapper function for the multiprocessing. Default: None 

76 

77 Returns 

78 ------- 

79 np.array 

80 integrand of P(delta_final_mass/final_mass_bar, 

81 delta_final_spin/final_spin_bar) 

82 """ 

83 final_mass_mat, final_spin_mat = np.meshgrid(final_mass, final_spin) 

84 _abs = np.abs(final_mass_mat * final_spin_mat) 

85 _reshape = (len(v1), len(v1)) 

86 _v2, _v1 = np.meshgrid(v2, v1) 

87 v1, v2 = _v1.ravel(), _v2.ravel() 

88 v1, v2 = v1.reshape(len(v1), 1), v2.reshape(len(v2), 1) 

89 

90 # Create delta_final_mass and delta_final_spin vectors corresponding 

91 # to the given v1 and v2. 

92 

93 # These vectors have to be monotonically increasing in order to 

94 # evaluate the interpolated prob densities. Hence, for v1, v2 < 0, 

95 # flip them, evaluate the prob density (in column or row) and flip it back 

96 

97 # The definition of the delta_* parameters is taken from eq A1 of 

98 # Ghosh et al 2018, arXiv:1704.06784. 

99 

100 delta_final_mass_i = ((1.0 + v1 / 2.0)) * final_mass 

101 delta_final_spin_i = ((1.0 + v2 / 2.0)) * final_spin 

102 delta_final_mass_r = ((1.0 - v1 / 2.0)) * final_mass 

103 delta_final_spin_r = ((1.0 - v2 / 2.0)) * final_spin 

104 

105 for num in range(len(v1)): 

106 if (1.0 + v1[num] / 2.0) < 0.0: 

107 delta_final_mass_i[num] = np.flipud(delta_final_mass_i[num]) 

108 if (1.0 + v2[num] / 2.0) < 0.0: 

109 delta_final_spin_i[num] = np.flipud(delta_final_spin_i[num]) 

110 if (1.0 - v1[num] / 2.0) < 0.0: 

111 delta_final_mass_r[num] = np.flipud(delta_final_mass_r[num]) 

112 if (1.0 - v2[num] / 2.0) < 0.0: 

113 delta_final_spin_r[num] = np.flipud(delta_final_spin_r[num]) 

114 

115 if multi_process > 1: 

116 with multiprocessing.Pool(multi_process) as pool: 

117 _length = len(delta_final_mass_i) 

118 _P_i = pool.starmap( 

119 wrapper_function_for_multiprocess, 

120 zip( 

121 [P_final_mass_final_spin_i_interp_object] * _length, 

122 delta_final_mass_i, 

123 delta_final_spin_i, 

124 ), 

125 ) 

126 _P_r = pool.starmap( 

127 wrapper_function_for_multiprocess, 

128 zip( 

129 [P_final_mass_final_spin_r_interp_object] * _length, 

130 delta_final_mass_r, 

131 delta_final_spin_r, 

132 ), 

133 ) 

134 P_i = np.array([i for i in _P_i]) 

135 P_r = np.array([i for i in _P_r]) 

136 else: 

137 P_i, P_r = [], [] 

138 for num in range(len(delta_final_mass_i)): 

139 P_i.append( 

140 wrapper_function_for_multiprocess( 

141 P_final_mass_final_spin_i_interp_object, 

142 delta_final_mass_i[num], 

143 delta_final_spin_i[num], 

144 ) 

145 ) 

146 P_r.append( 

147 wrapper_function_for_multiprocess( 

148 P_final_mass_final_spin_r_interp_object, 

149 delta_final_mass_r[num], 

150 delta_final_spin_r[num], 

151 ) 

152 ) 

153 

154 for num in range(len(v1)): 

155 if (1.0 + v1[num] / 2.0) < 0.0: 

156 P_i[num] = np.fliplr(P_i[num]) 

157 if (1.0 + v2[num] / 2.0) < 0.0: 

158 P_i[num] = np.flipud(P_i[num]) 

159 if (1.0 - v1[num] / 2.0) < 0.0: 

160 P_r[num] = np.fliplr(P_r[num]) 

161 if (1.0 - v2[num] / 2.0) < 0.0: 

162 P_r[num] = np.flipud(P_r[num]) 

163 

164 # The integration is performed according to eq A2 of Ghosh et al, 

165 # arXiv:1704.06784 

166 

167 _prod = np.array( 

168 [np.sum(_P_i * _P_r * _abs) for _P_i, _P_r in zip(P_i, P_r)] 

169 ).reshape(_reshape) 

170 return _prod 

171 

172 

173def _apply_args_and_kwargs(function, args, kwargs): 

174 """Apply a tuple of args and a dictionary of kwargs to a function 

175 

176 Parameters 

177 ---------- 

178 function: func 

179 function you wish to use 

180 args: tuple 

181 all args passed to function 

182 kwargs: dict 

183 all kwargs passed to function 

184 """ 

185 return function(*args, **kwargs) 

186 

187 

188def _imrct_deviation_parameters_integrand_series( 

189 final_mass, 

190 final_spin, 

191 v1, 

192 v2, 

193 P_final_mass_final_spin_i_interp_object, 

194 P_final_mass_final_spin_r_interp_object, 

195 multi_process=1, 

196 **kwargs 

197): 

198 """ 

199 Creates the over the deviation parameter space. 

200 

201 Parameters 

202 ---------- 

203 final_mass: np.ndarray 

204 vector of values of final mass 

205 final_spin: np.ndarray 

206 vector of values of final spin 

207 v1: np.ndarray 

208 array of delta_final_mass/final_mass_bar values 

209 v2: np.ndarray 

210 array of delta_final_spin/final_spin_bar values 

211 P_final_mass_final_spin_i_interp_object: 

212 interpolated P_i(final_mass, final_spin) 

213 P_final_massfinal_spin_r_interp_object: 

214 interpolated P_r(final_mass, final_spin) 

215 """ 

216 P = np.zeros(shape=(len(v1), len(v2))) 

217 if multi_process == 1: 

218 logger.debug("Performing calculation on a single cpu. This may take some time") 

219 for i, _v2 in enumerate(v2): 

220 for j, _v1 in enumerate(v1): 

221 P[i, j] = _imrct_deviation_parameters_integrand_vectorized( 

222 final_mass, 

223 final_spin, 

224 [_v1], 

225 [_v2], 

226 P_final_mass_final_spin_i_interp_object, 

227 P_final_mass_final_spin_r_interp_object, 

228 multi_process=1, 

229 **kwargs 

230 ) 

231 else: 

232 logger.debug("Splitting the calculation across {} cpus".format(multi_process)) 

233 _v1, _v2 = np.meshgrid(v1, v2) 

234 _v1, _v2 = _v1.ravel(), _v2.ravel() 

235 with multiprocessing.Pool(multi_process) as pool: 

236 args = [ 

237 [final_mass] * len(_v1), 

238 [final_spin] * len(_v1), 

239 np.atleast_2d(_v1).T.tolist(), 

240 np.atleast_2d(_v2).T.tolist(), 

241 [P_final_mass_final_spin_i_interp_object] * len(_v1), 

242 [P_final_mass_final_spin_r_interp_object] * len(_v1), 

243 ] 

244 kwargs["multi_process"] = 1 

245 _args = np.array(args, dtype=object).T 

246 _P = pool.starmap( 

247 _apply_args_and_kwargs, 

248 zip( 

249 [_imrct_deviation_parameters_integrand_vectorized] * len(_args), 

250 _args, 

251 [kwargs] * len(_args), 

252 ), 

253 ) 

254 P = np.array([i[0] for i in _P]).reshape(len(v1), len(v2)) 

255 return P 

256 

257 

258def imrct_deviation_parameters_integrand(*args, vectorize=False, **kwargs): 

259 """Compute the final mass and final spin deviation parameters 

260 

261 Parameters 

262 ---------- 

263 *args: tuple 

264 all args passed to either 

265 _imrct_deviation_parameters_integrand_vectorized or 

266 _imrct_deviation_parameters_integrand_series 

267 vectorize: bool 

268 Vectorize the calculation. Note that vectorize=True uses up a lot 

269 of memory 

270 kwargs: dict, optional 

271 all kwargs passed to either 

272 _imrct_deviation_parameters_integrand_vectorized or 

273 _imrct_deviation_parameters_integrand_series 

274 """ 

275 if vectorize: 

276 return _imrct_deviation_parameters_integrand_vectorized(*args, **kwargs) 

277 return _imrct_deviation_parameters_integrand_series(*args, **kwargs) 

278 

279 

280def imrct_deviation_parameters_from_final_mass_final_spin( 

281 final_mass_inspiral, 

282 final_spin_inspiral, 

283 final_mass_postinspiral, 

284 final_spin_postinspiral, 

285 N_bins=101, 

286 final_mass_deviation_lim=1, 

287 final_spin_deviation_lim=1, 

288 multi_process=1, 

289 use_kde=False, 

290 kde=gaussian_kde, 

291 kde_kwargs=dict(), 

292 interp_method=RectBivariateSpline, 

293 interp_kwargs=dict(fill_value=0.0, bounds_error=False, kx=1, ky=1), 

294 vectorize=False, 

295): 

296 """Compute the IMR Consistency Test deviation parameters. 

297 Code borrows from the implementation in lalsuite: 

298 https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalinference/python/lalinference/imrtgr/imrtgrutils.py 

299 and 

300 https://git.ligo.org/lscsoft/lalsuite/-/blob/master/lalinference/bin/imrtgr_imr_consistency_test.py 

301 

302 Parameters 

303 ---------- 

304 final_mass_inspiral: np.ndarray 

305 values of final mass calculated from the inspiral part 

306 final_spin_inspiral: np.ndarray 

307 values of final spin calculated from the inspiral part 

308 final_mass_postinspiral: np.ndarray 

309 values of final mass calculated from the post-inspiral part 

310 final_spin_postinspiral: np.ndarray 

311 values of final spin calculated from the post-inspiral part 

312 final_mass_deviation_lim: float 

313 Maximum absolute value of the final mass deviation parameter. Default 1. 

314 final_spin_deviation_lim: float 

315 Maximum absolute value of the final spin deviation parameter. Default 1. 

316 N_bins: int, optional 

317 Number of equally spaced bins between [-final_mass_deviation_lim, 

318 final_mass_deviation_lim] and [-final_spin_deviation_lim, 

319 final_spin_deviation_lim]. Default is 101. 

320 multi_process: int 

321 Number of parallel processes. Default is 1. 

322 use_kde: bool 

323 If `True`, uses kde instead of interpolation. Default is False. 

324 kde: method 

325 KDE method to use. Default is scipy.stats.gaussian_kde 

326 kde_kwargs: dict 

327 Arguments to be passed to the KDE method 

328 interp_method: method 

329 Interpolation method to use. Default is scipy.interpolate.interp2d 

330 interp_kwargs: dict, optional 

331 Arguments to be passed to the interpolation method 

332 Default is `dict(fill_value=0.0, bounds_error=False)` 

333 vectorize: bool 

334 if True, use vectorized imrct_deviation_parameters_integrand 

335 function. This is quicker but does consume more memory. Default: False 

336 

337 Returns 

338 ------- 

339 imrct_deviations: ProbabilityDict2d 

340 Contains the 2d pdf of the IMRCT deviation parameters 

341 """ 

342 # Find the maximum values 

343 final_mass_lim = np.max(np.append(final_mass_inspiral, final_mass_postinspiral)) 

344 final_spin_lim = np.max( 

345 np.abs(np.append(final_spin_inspiral, final_spin_postinspiral)) 

346 ) 

347 

348 # bin the data 

349 final_mass_bins = np.linspace(-final_mass_lim, final_mass_lim, N_bins) 

350 diff_final_mass = final_mass_bins[1] - final_mass_bins[0] 

351 final_spin_bins = np.linspace(-final_spin_lim, final_spin_lim, N_bins) 

352 diff_final_spin = final_spin_bins[1] - final_spin_bins[0] 

353 final_mass_intp = (final_mass_bins[:-1] + final_mass_bins[1:]) * 0.5 

354 final_spin_intp = (final_spin_bins[:-1] + final_spin_bins[1:]) * 0.5 

355 if use_kde: 

356 logger.debug("Using KDE to interpolate data") 

357 # kde the samples for final mass and final spin 

358 final_mass_intp = np.append( 

359 final_mass_intp, final_mass_bins[-1] + diff_final_mass 

360 ) 

361 final_spin_intp = np.append( 

362 final_spin_intp, final_spin_bins[-1] + diff_final_spin 

363 ) 

364 inspiral_interp = kde( 

365 np.array([final_mass_inspiral, final_spin_inspiral]), **kde_kwargs 

366 ) 

367 postinspiral_interp = kde( 

368 np.array([final_mass_postinspiral, final_spin_postinspiral]), **kde_kwargs 

369 ) 

370 _wrapper_function = _wrapper_for_multiprocessing_kde 

371 else: 

372 logger.debug("Interpolating 2d histogram data") 

373 _inspiral_2d_histogram, _, _ = np.histogram2d( 

374 final_mass_inspiral, 

375 final_spin_inspiral, 

376 bins=(final_mass_bins, final_spin_bins), 

377 density=True, 

378 ) 

379 _postinspiral_2d_histogram, _, _ = np.histogram2d( 

380 final_mass_postinspiral, 

381 final_spin_postinspiral, 

382 bins=(final_mass_bins, final_spin_bins), 

383 density=True, 

384 ) 

385 inspiral_interp = interp_method( 

386 final_mass_intp, final_spin_intp, _inspiral_2d_histogram, **interp_kwargs 

387 ) 

388 postinspiral_interp = interp_method( 

389 final_mass_intp, final_spin_intp, _postinspiral_2d_histogram, **interp_kwargs 

390 ) 

391 _wrapper_function = _wrapper_for_multiprocessing_interp 

392 

393 final_mass_deviation_vec = np.linspace( 

394 -final_mass_deviation_lim, final_mass_deviation_lim, N_bins 

395 ) 

396 final_spin_deviation_vec = np.linspace( 

397 -final_spin_deviation_lim, final_spin_deviation_lim, N_bins 

398 ) 

399 

400 diff_final_mass_deviation = final_mass_deviation_vec[1] - final_mass_deviation_vec[0] 

401 diff_final_spin_deviation = final_spin_deviation_vec[1] - final_spin_deviation_vec[0] 

402 

403 P_final_mass_deviation_final_spin_deviation = imrct_deviation_parameters_integrand( 

404 final_mass_intp, 

405 final_spin_intp, 

406 final_mass_deviation_vec, 

407 final_spin_deviation_vec, 

408 inspiral_interp, 

409 postinspiral_interp, 

410 multi_process=multi_process, 

411 vectorize=vectorize, 

412 wrapper_function_for_multiprocess=_wrapper_function, 

413 ) 

414 

415 imrct_deviations = ProbabilityDict2D( 

416 { 

417 "final_mass_final_spin_deviations": [ 

418 final_mass_deviation_vec, 

419 final_spin_deviation_vec, 

420 P_final_mass_deviation_final_spin_deviation 

421 / np.sum(P_final_mass_deviation_final_spin_deviation), 

422 ] 

423 } 

424 ) 

425 return imrct_deviations 

426 

427 

428def generate_imrct_deviation_parameters( 

429 samples, evolve_spins_forward=True, inspiral_string="inspiral", 

430 postinspiral_string="postinspiral", approximant=None, f_low=None, 

431 return_samples_used=False, **kwargs 

432): 

433 """Generate deviation parameter pdfs for the IMR Consistency Test 

434 

435 Parameters 

436 ---------- 

437 samples: MultiAnalysisSamplesDict 

438 Dictionary containing inspiral and postinspiral samples 

439 evolve_spins_forward: bool 

440 If `True`, evolve spins to the ISCO frequency. Default: True. 

441 inspiral_string: string 

442 Identifier for the inspiral samples 

443 postinspiral_string: string 

444 Identifier for the post-inspiral samples 

445 approximant: dict, optional 

446 The approximant used for the inspiral and postinspiral analyses. 

447 Keys of the dictionary must be the same as the inspiral_string and 

448 postinspiral_string. Default None 

449 f_low: dict, optional 

450 The low frequency cut-off used for the inspiral and postinspiral 

451 analyses. Keys of the dictionary must be the same as the inspiral_string 

452 and postinspiral_string. Default None 

453 return_samples_used: Bool, optional 

454 if True, return the samples which were used to generate the IMRCT deviation 

455 parameters. These samples will match the input but may include remnant 

456 samples if they were not previously present 

457 kwargs: dict, optional 

458 Keywords to be passed to imrct_deviation_parameters_from_final_mass_final_spin 

459 

460 Returns 

461 ------- 

462 imrct_deviations: ProbabilityDict2d 

463 2d pdf of the IMRCT deviation parameters 

464 data: dict 

465 Metadata 

466 """ 

467 import time 

468 

469 remnant_condition = lambda _dictionary, _suffix: all( 

470 "{}{}".format(param, _suffix) not in _dictionary.keys() for 

471 param in ["final_mass", "final_spin"] 

472 ) 

473 evolved = np.ones(2, dtype=bool) 

474 suffix = [""] 

475 evolve_spins = ["ISCO"] 

476 if not evolve_spins_forward: 

477 suffix = ["_non_evolved"] + suffix 

478 evolve_spins = [False, False] 

479 fits_data = {} 

480 for idx, (key, sample) in enumerate(samples.items()): 

481 zipped = zip(suffix, evolve_spins) 

482 for num, (_suffix, _evolve_spins) in enumerate(zipped): 

483 cond = remnant_condition(sample, _suffix) 

484 _found_msg = ( 

485 "Found {} remnant properties in the posterior table " 

486 "for {}. Using these for calculation." 

487 ) 

488 if not cond: 

489 logger.info( 

490 _found_msg.format( 

491 "evolved" if not len(_suffix) else "non-evolved", 

492 key 

493 ) 

494 ) 

495 if len(_suffix): 

496 evolved[idx] = False 

497 break 

498 elif not remnant_condition(sample, ""): 

499 logger.info(_found_msg.format("evolved", key)) 

500 evolved[idx] = True 

501 break 

502 else: 

503 logger.warning( 

504 "{} remnant properties not found in the posterior " 

505 "table for {}. Trying to calculate them.".format( 

506 "Evolved" if not len(_suffix) else "Non-evolved", 

507 key 

508 ) 

509 ) 

510 returned_extra_kwargs = sample.generate_all_posterior_samples( 

511 evolve_spins=_evolve_spins, return_kwargs=True, 

512 approximant=( 

513 approximant[key] if approximant is not None else None 

514 ), f_low=f_low[key] if f_low is not None else None 

515 ) 

516 _cond = remnant_condition(sample, _suffix) 

517 if not _cond: 

518 logger.info( 

519 "{} remnant properties generated. Using these " 

520 "samples for calculation".format( 

521 "Evolved" if not len(_suffix) else "Non-evolved" 

522 ) 

523 ) 

524 for fit in ["final_mass_NR_fits", "final_spin_NR_fits"]: 

525 fits_data["{} {}".format(key, fit)] = returned_extra_kwargs[ 

526 "meta_data" 

527 ][fit] 

528 if len(_suffix): 

529 evolved[idx] = False 

530 break 

531 

532 if num == 1: 

533 raise ValueError( 

534 "Unable to compute the remnant properties" 

535 ) 

536 if not all(_evolved == evolved[0] for _evolved in evolved): 

537 keys = list(samples.keys()) 

538 _inspiral_index = keys.index(inspiral_string) 

539 _postinspiral_index = keys.index(postinspiral_string) 

540 raise ValueError( 

541 "Using {} remnant properties for the inspiral and {} remnant " 

542 "properties for the postinspiral. This must be the same for " 

543 "the calculation".format( 

544 "evolved" if evolved[_inspiral_index] else "non-evolved", 

545 "non-evolved" if not evolved[_postinspiral_index] else "evolved" 

546 ) 

547 ) 

548 samples_string = "final_{}" 

549 if not evolved[0]: 

550 samples_string += "_non_evolved" 

551 

552 logger.info("Calculating IMRCT deviation parameters and GR Quantile") 

553 t0 = time.time() 

554 imrct_deviations = imrct_deviation_parameters_from_final_mass_final_spin( 

555 samples[inspiral_string][samples_string.format("mass")], 

556 samples[inspiral_string][samples_string.format("spin")], 

557 samples[postinspiral_string][samples_string.format("mass")], 

558 samples[postinspiral_string][samples_string.format("spin")], 

559 **kwargs, 

560 ) 

561 gr_quantile = ( 

562 imrct_deviations[ 

563 "final_mass_final_spin_deviations" 

564 ].minimum_encompassing_contour_level(0.0, 0.0) 

565 * 100 

566 ) 

567 t1 = time.time() 

568 data = kwargs.copy() 

569 data["evolve_spins"] = evolved 

570 data["Time (seconds)"] = round(t1 - t0, 2) 

571 data["GR Quantile (%)"] = gr_quantile[0] 

572 data.update(fits_data) 

573 logger.info( 

574 "Calculation Finished in {} seconds. GR Quantile is {} %.".format( 

575 data["Time (seconds)"], round(data["GR Quantile (%)"], 2) 

576 ) 

577 ) 

578 if return_samples_used: 

579 return imrct_deviations, data, evolved[0], samples 

580 return imrct_deviations, data, evolved[0]