Coverage for pesummary/gw/conversions/tgr.py: 75.6%
168 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-05-02 08:42 +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
6from scipy.interpolate import interp2d
7import multiprocessing
8from pesummary.utils.probability_dict import ProbabilityDict2D
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)
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=interp2d,
293 interp_kwargs=dict(fill_value=0.0, bounds_error=False),
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 # transpose density to go from (X,Y) indexing returned by
386 # np.histogram2d() to array (i,j) indexing for further computations.
387 # From now onwards, different rows (i) correspond to different values
388 # of final mass and different columns (j) correspond to different
389 # values of final_spin
390 _inspiral_2d_histogram = _inspiral_2d_histogram.T
391 _postinspiral_2d_histogram = _postinspiral_2d_histogram.T
392 inspiral_interp = interp_method(
393 final_mass_intp, final_spin_intp, _inspiral_2d_histogram, **interp_kwargs
394 )
395 postinspiral_interp = interp_method(
396 final_mass_intp, final_spin_intp, _postinspiral_2d_histogram, **interp_kwargs
397 )
398 _wrapper_function = _wrapper_for_multiprocessing_interp
400 final_mass_deviation_vec = np.linspace(
401 -final_mass_deviation_lim, final_mass_deviation_lim, N_bins
402 )
403 final_spin_deviation_vec = np.linspace(
404 -final_spin_deviation_lim, final_spin_deviation_lim, N_bins
405 )
407 diff_final_mass_deviation = final_mass_deviation_vec[1] - final_mass_deviation_vec[0]
408 diff_final_spin_deviation = final_spin_deviation_vec[1] - final_spin_deviation_vec[0]
410 P_final_mass_deviation_final_spin_deviation = imrct_deviation_parameters_integrand(
411 final_mass_intp,
412 final_spin_intp,
413 final_mass_deviation_vec,
414 final_spin_deviation_vec,
415 inspiral_interp,
416 postinspiral_interp,
417 multi_process=multi_process,
418 vectorize=vectorize,
419 wrapper_function_for_multiprocess=_wrapper_function,
420 )
422 imrct_deviations = ProbabilityDict2D(
423 {
424 "final_mass_final_spin_deviations": [
425 final_mass_deviation_vec,
426 final_spin_deviation_vec,
427 P_final_mass_deviation_final_spin_deviation
428 / np.sum(P_final_mass_deviation_final_spin_deviation),
429 ]
430 }
431 )
432 return imrct_deviations
435def generate_imrct_deviation_parameters(
436 samples, evolve_spins_forward=True, inspiral_string="inspiral",
437 postinspiral_string="postinspiral", approximant=None, f_low=None,
438 return_samples_used=False, **kwargs
439):
440 """Generate deviation parameter pdfs for the IMR Consistency Test
442 Parameters
443 ----------
444 samples: MultiAnalysisSamplesDict
445 Dictionary containing inspiral and postinspiral samples
446 evolve_spins_forward: bool
447 If `True`, evolve spins to the ISCO frequency. Default: True.
448 inspiral_string: string
449 Identifier for the inspiral samples
450 postinspiral_string: string
451 Identifier for the post-inspiral samples
452 approximant: dict, optional
453 The approximant used for the inspiral and postinspiral analyses.
454 Keys of the dictionary must be the same as the inspiral_string and
455 postinspiral_string. Default None
456 f_low: dict, optional
457 The low frequency cut-off used for the inspiral and postinspiral
458 analyses. Keys of the dictionary must be the same as the inspiral_string
459 and postinspiral_string. Default None
460 return_samples_used: Bool, optional
461 if True, return the samples which were used to generate the IMRCT deviation
462 parameters. These samples will match the input but may include remnant
463 samples if they were not previously present
464 kwargs: dict, optional
465 Keywords to be passed to imrct_deviation_parameters_from_final_mass_final_spin
467 Returns
468 -------
469 imrct_deviations: ProbabilityDict2d
470 2d pdf of the IMRCT deviation parameters
471 data: dict
472 Metadata
473 """
474 import time
476 remnant_condition = lambda _dictionary, _suffix: all(
477 "{}{}".format(param, _suffix) not in _dictionary.keys() for
478 param in ["final_mass", "final_spin"]
479 )
480 evolved = np.ones(2, dtype=bool)
481 suffix = [""]
482 evolve_spins = ["ISCO"]
483 if not evolve_spins_forward:
484 suffix = ["_non_evolved"] + suffix
485 evolve_spins = [False, False]
486 fits_data = {}
487 for idx, (key, sample) in enumerate(samples.items()):
488 zipped = zip(suffix, evolve_spins)
489 for num, (_suffix, _evolve_spins) in enumerate(zipped):
490 cond = remnant_condition(sample, _suffix)
491 _found_msg = (
492 "Found {} remnant properties in the posterior table "
493 "for {}. Using these for calculation."
494 )
495 if not cond:
496 logger.info(
497 _found_msg.format(
498 "evolved" if not len(_suffix) else "non-evolved",
499 key
500 )
501 )
502 if len(_suffix):
503 evolved[idx] = False
504 break
505 elif not remnant_condition(sample, ""):
506 logger.info(_found_msg.format("evolved", key))
507 evolved[idx] = True
508 break
509 else:
510 logger.warning(
511 "{} remnant properties not found in the posterior "
512 "table for {}. Trying to calculate them.".format(
513 "Evolved" if not len(_suffix) else "Non-evolved",
514 key
515 )
516 )
517 returned_extra_kwargs = sample.generate_all_posterior_samples(
518 evolve_spins=_evolve_spins, return_kwargs=True,
519 approximant=(
520 approximant[key] if approximant is not None else None
521 ), f_low=f_low[key] if f_low is not None else None
522 )
523 _cond = remnant_condition(sample, _suffix)
524 if not _cond:
525 logger.info(
526 "{} remnant properties generated. Using these "
527 "samples for calculation".format(
528 "Evolved" if not len(_suffix) else "Non-evolved"
529 )
530 )
531 for fit in ["final_mass_NR_fits", "final_spin_NR_fits"]:
532 fits_data["{} {}".format(key, fit)] = returned_extra_kwargs[
533 "meta_data"
534 ][fit]
535 if len(_suffix):
536 evolved[idx] = False
537 break
539 if num == 1:
540 raise ValueError(
541 "Unable to compute the remnant properties"
542 )
543 if not all(_evolved == evolved[0] for _evolved in evolved):
544 keys = list(samples.keys())
545 _inspiral_index = keys.index(inspiral_string)
546 _postinspiral_index = keys.index(postinspiral_string)
547 raise ValueError(
548 "Using {} remnant properties for the inspiral and {} remnant "
549 "properties for the postinspiral. This must be the same for "
550 "the calculation".format(
551 "evolved" if evolved[_inspiral_index] else "non-evolved",
552 "non-evolved" if not evolved[_postinspiral_index] else "evolved"
553 )
554 )
555 samples_string = "final_{}"
556 if not evolved[0]:
557 samples_string += "_non_evolved"
559 logger.info("Calculating IMRCT deviation parameters and GR Quantile")
560 t0 = time.time()
561 imrct_deviations = imrct_deviation_parameters_from_final_mass_final_spin(
562 samples[inspiral_string][samples_string.format("mass")],
563 samples[inspiral_string][samples_string.format("spin")],
564 samples[postinspiral_string][samples_string.format("mass")],
565 samples[postinspiral_string][samples_string.format("spin")],
566 **kwargs,
567 )
568 gr_quantile = (
569 imrct_deviations[
570 "final_mass_final_spin_deviations"
571 ].minimum_encompassing_contour_level(0.0, 0.0)
572 * 100
573 )
574 t1 = time.time()
575 data = kwargs.copy()
576 data["evolve_spins"] = evolved
577 data["Time (seconds)"] = round(t1 - t0, 2)
578 data["GR Quantile (%)"] = gr_quantile[0]
579 data.update(fits_data)
580 logger.info(
581 "Calculation Finished in {} seconds. GR Quantile is {} %.".format(
582 data["Time (seconds)"], round(data["GR Quantile (%)"], 2)
583 )
584 )
585 if return_samples_used:
586 return imrct_deviations, data, evolved[0], samples
587 return imrct_deviations, data, evolved[0]