diff --git a/bilby/gw/conversion.py b/bilby/gw/conversion.py index 00d957f98fea3b0eeba1459bf25f16cc47c39ab7..4be5634eb2977fac17db7b7c22fe30f2d6591916 100644 --- a/bilby/gw/conversion.py +++ b/bilby/gw/conversion.py @@ -2254,17 +2254,15 @@ def compute_snrs(sample, likelihood, npool=1): new_samples = [_compute_snrs(xx) for xx in tqdm(fill_args, file=sys.stdout)] for ii, ifo in enumerate(likelihood.interferometers): - matched_filter_snrs = list() - optimal_snrs = list() - mf_key = '{}_matched_filter_snr'.format(ifo.name) - optimal_key = '{}_optimal_snr'.format(ifo.name) + snr_updates = dict() + for key in new_samples[0][ii].snrs_as_sample.keys(): + snr_updates[f"{ifo.name}_{key}"] = list() for new_sample in new_samples: - matched_filter_snrs.append(new_sample[ii].complex_matched_filter_snr) - optimal_snrs.append(new_sample[ii].optimal_snr_squared.real ** 0.5) - - sample[mf_key] = matched_filter_snrs - sample[optimal_key] = optimal_snrs - + snr_update = new_sample[ii].snrs_as_sample + for key, val in snr_update.items(): + snr_updates[f"{ifo.name}_{key}"].append(val) + for k, v in snr_updates.items(): + sample[k] = v else: logger.debug('Not computing SNRs.') diff --git a/bilby/gw/likelihood/base.py b/bilby/gw/likelihood/base.py index 256237bfb9f8395a67a7152adb66fbd18892b9d1..e0a09c1e934eb8886e27831839602307d6aba532 100644 --- a/bilby/gw/likelihood/base.py +++ b/bilby/gw/likelihood/base.py @@ -128,6 +128,20 @@ class GravitationalWaveTransient(Likelihood): setattr(self, key, other) return self + @property + def snrs_as_sample(self) -> dict: + """Get the SNRs of this object as a sample dictionary + + Returns + ======= + dict + The dictionary of SNRs labelled accordingly + """ + return { + "matched_filter_snr" : self.complex_matched_filter_snr, + "optimal_snr" : self.optimal_snr_squared.real ** 0.5 + } + def __init__( self, interferometers, waveform_generator, time_marginalization=False, distance_marginalization=False, phase_marginalization=False, calibration_marginalization=False, priors=None,