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
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
10__author__ = [
11 "Aditya Vijaykumar <aditya.vijaykumar@ligo.org>",
12 "Charlie Hoy <charlie.hoy@ligo.org>"
13]
16def _wrapper_for_multiprocessing_kde(kde, *args):
17 """Wrapper to evaluate a KDE on multiple cpus
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)
32def _wrapper_for_multiprocessing_interp(interp, *args):
33 """Wrapper to evaluate an interpolant on multiple cpus
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
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).
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
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)
90 # Create delta_final_mass and delta_final_spin vectors corresponding
91 # to the given v1 and v2.
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
97 # The definition of the delta_* parameters is taken from eq A1 of
98 # Ghosh et al 2018, arXiv:1704.06784.
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
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])
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 )
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])
164 # The integration is performed according to eq A2 of Ghosh et al,
165 # arXiv:1704.06784
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
173def _apply_args_and_kwargs(function, args, kwargs):
174 """Apply a tuple of args and a dictionary of kwargs to a function
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)
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.
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
258def imrct_deviation_parameters_integrand(*args, vectorize=False, **kwargs):
259 """Compute the final mass and final spin deviation parameters
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)
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
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
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 )
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
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 )
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]
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 )
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
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
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
460 Returns
461 -------
462 imrct_deviations: ProbabilityDict2d
463 2d pdf of the IMRCT deviation parameters
464 data: dict
465 Metadata
466 """
467 import time
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
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"
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]