From f729525f8064bf09f38df1aaef7ecfb8aab0f616 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 20 Aug 2024 11:06:38 +0100 Subject: [PATCH 1/4] Updated the profiling tests. --- profiling/cuda/cuda_likelihood.py | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/profiling/cuda/cuda_likelihood.py b/profiling/cuda/cuda_likelihood.py index d2194f3..20b460e 100644 --- a/profiling/cuda/cuda_likelihood.py +++ b/profiling/cuda/cuda_likelihood.py @@ -3,12 +3,12 @@ from heron.models.lalnoise import AdvancedLIGO from heron.injection import make_injection_zero_noise from heron.detector import AdvancedLIGOHanford -from heron.likelihood import TimeDomainLikelihoodModelUncertaintyPyTorch +from heron.likelihood import TimeDomainLikelihoodPyTorch from heron.models.lalsimulation import IMRPhenomPv2, IMRPhenomPv2_FakeUncertainty -def profile_likelihood(): - waveform = IMRPhenomPv2_FakeUncertainty() +def profile_likelihood_pytorch_nouncert(): + waveform = IMRPhenomPv2() psd_model = AdvancedLIGO() injections = make_injection_zero_noise(waveform=IMRPhenomPv2, @@ -22,22 +22,21 @@ def profile_likelihood(): data = injections['H1'] - likelihood = TimeDomainLikelihoodModelUncertaintyPyTorch(data, psd=psd_model) + likelihood = TimeDomainLikelihoodPyTorch(data, psd=psd_model) print(likelihood.device) - for m1 in np.linspace(20, 50, 100): - test_waveform = waveform.time_domain(parameters={"m1": m1*u.solMass, - "m2": 30*u.solMass, - "gpstime": 4000, - "distance": 410 * u.megaparsec}, times=data.times) + test_waveform = waveform.time_domain(parameters={"m1": 30*u.solMass, + "m2": 30*u.solMass, + "gpstime": 4000, + "distance": 410 * u.megaparsec}, times=data.times) - projected_waveform = test_waveform.project(AdvancedLIGOHanford(), - ra=0, dec=0, - phi_0=0, psi=0, - iota=0) + projected_waveform = test_waveform.project(AdvancedLIGOHanford(), + ra=0, dec=0, + phi_0=0, psi=0, + iota=0) - log_like = likelihood.log_likelihood(projected_waveform) + log_like = likelihood.log_likelihood(projected_waveform) from torch.profiler import profile, record_function, ProfilerActivity @@ -45,7 +44,7 @@ def profile_likelihood(): with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: with record_function("model_likelihood"): - profile_likelihood() + profile_likelihood_pytorch_nouncert() print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=10)) print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10)) From eefa5ef2072f37c536fb21bd799352ce314bd02f Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Thu, 22 Aug 2024 16:17:29 +0100 Subject: [PATCH 2/4] Further inference improvements. --- heron/likelihood.py | 44 ++++++++++++++++++++++++----------- heron/models/lalsimulation.py | 2 +- tests/test_inference.py | 5 ++-- 3 files changed, 34 insertions(+), 17 deletions(-) diff --git a/heron/likelihood.py b/heron/likelihood.py index e7e6fee..44999f0 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -126,6 +126,9 @@ class TimeDomainLikelihoodModelUncertainty(TimeDomainLikelihood): def __init__(self, data, psd, waveform=None, detector=None): super().__init__(data, psd, waveform, detector) + self.norm_factor_2 = np.max(self.C) + self.norm_factor = np.sqrt(self.norm_factor_2) + def _normalisation(self, K, S): norm = ( -1.5 * K.shape[0] * self.log(2 * self.pi) @@ -135,12 +138,13 @@ def _normalisation(self, K, S): ) return norm - def _weighted_data(self): + def _weighted_data(self, indices): """Return the weighted data component""" # TODO This can all be pre-computed + (a, b) = indices if not hasattr(self, "weighted_data_CACHE"): dw = self.weighted_data_CACHE = ( - -0.5 * self.data.T @ self.solve(self.C, self.data) + -0.5 * (np.array(self.data)/np.sqrt(self.norm_factor))[a[0]:a[1]].T @ self.solve((self.C/self.norm_factor_2)[a[0]:a[1], a[0]:a[1]], self.data[a[0]:a[1]]) ) else: dw = self.weighted_data_CACHE @@ -150,22 +154,34 @@ def _weighted_model(self, mu, K): """Return the inner product of the GPR mean""" return -0.5 * np.array(mu).T @ self.solve(K, mu) - def _weighted_cross(self, mu, K): - a = self.solve(self.C, self.data) + self.solve(K, mu) - b = self.inverse_C + self.inverse(K) - return 0.5 * a.T @ self.solve(b, a) + def _weighted_cross(self, mu, K, indices): + # NB the first part of this is repeated elsewhere + (a,b) = indices + C = (self.C/self.norm_factor_2)[a[0]:a[1],a[0]:a[1]] + data = (self.data/self.norm_factor)[a[0]:a[1]] + + A = (self.solve(C, data) - self.solve(K, mu)) + B = (self.inverse_C*self.norm_factor_2)[a[0]:a[1],a[0]:a[1]] + self.inverse(K) + return 0.5 * A.T @ self.solve(B, A) - def log_likelihood(self, waveform): - like = self._weighted_cross(waveform.data, waveform.covariance) + def log_likelihood(self, waveform, norm=True): + a, b = self.timeseries.determine_overlap(self, waveform) + + wf = np.array(waveform.data)[b[0]:b[1]] + wc = waveform.covariance[b[0]:b[1],b[0]:b[1]] + wc /= self.norm_factor_2 #np.max(wc) + wf /= self.norm_factor #np.sqrt(np.max(wc)) + + like = - self._weighted_cross(wf, wc, indices=(a,b)) # print("cross term", like) - A = self._weighted_data() + A = self._weighted_data((a, b)) # print("data term", A) - B = self._weighted_model(waveform.data, waveform.covariance) + B = self._weighted_model(wf, wc) # print("model term", B) - like = like + A + B - norm = self._normalisation(waveform.covariance, self.C) + like = like - A - B + #N = self._normalisation(waveform.covariance/np.sqrt(normalise), self.C/normalise) # print("normalisation", norm) - like += norm + #like += (N if norm else 0) return like @@ -337,6 +353,8 @@ def _weighted_cross(self, mu, K): return 0.5 * a.T @ self.solve(b, a) def log_likelihood(self, waveform): + a, b = self.timeseries.determine_overlap(self, waveform) + waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) waveform_c = torch.tensor( waveform.covariance, device=self.device, dtype=torch.double diff --git a/heron/models/lalsimulation.py b/heron/models/lalsimulation.py index 4086358..559b90d 100644 --- a/heron/models/lalsimulation.py +++ b/heron/models/lalsimulation.py @@ -198,7 +198,7 @@ def __init__(self, covariance=1e-24): def time_domain(self, parameters, times=None): waveform_dict = super().time_domain(parameters, times) - covariance = torch.eye(len(waveform_dict["plus"].times)) * self.covariance + covariance = np.eye((len(waveform_dict["plus"].times))) * self.covariance**2 for wave in waveform_dict.waveforms.values(): # Artificially add a covariance function to each of these wave.covariance = covariance diff --git a/tests/test_inference.py b/tests/test_inference.py index e968bdd..9c7875f 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -189,7 +189,8 @@ def test_likelihood_numpy_equivalent(self): self.assertTrue(mass_ratios[np.argmax(log_likes)] == 0.6) self.assertTrue(np.all((np.array(log_likes) - np.array(log_likes_n)) < 0.001)) - + + class Test_Filter(unittest.TestCase): """Test that filters can be applied correctly to data.""" @@ -225,8 +226,6 @@ def test_snr(self): phi_0=0, psi=0, iota=0) - f = projected_waveform.plot() - f.savefig("projected_waveform.png") snr = likelihood.snr(projected_waveform) self.assertTrue(snr > 40 and snr < 45) From 69110f42550faa2588daed3c19c945c1aeedc110 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Thu, 22 Aug 2024 18:07:46 +0100 Subject: [PATCH 3/4] Added a new test for the numpy uncertainty likelihood. --- heron/likelihood.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/heron/likelihood.py b/heron/likelihood.py index 44999f0..2a05f61 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -179,9 +179,9 @@ def log_likelihood(self, waveform, norm=True): B = self._weighted_model(wf, wc) # print("model term", B) like = like - A - B - #N = self._normalisation(waveform.covariance/np.sqrt(normalise), self.C/normalise) + N = self._normalisation(waveform.covariance/self.norm_factor, self.C/self.norm_factor_2) # print("normalisation", norm) - #like += (N if norm else 0) + like += (N if norm else 0) return like From 8023c1ec749a4e851c28915691e7c9d6823bded6 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Thu, 22 Aug 2024 18:08:32 +0100 Subject: [PATCH 4/4] Actually add the tests. --- tests/test_inference_uncertain.py | 63 +++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 tests/test_inference_uncertain.py diff --git a/tests/test_inference_uncertain.py b/tests/test_inference_uncertain.py new file mode 100644 index 0000000..cc95fbd --- /dev/null +++ b/tests/test_inference_uncertain.py @@ -0,0 +1,63 @@ +import unittest + +import numpy as np +import astropy.units as u +import bilby.gw.prior + +from heron.models.lalsimulation import SEOBNRv3, IMRPhenomPv2, IMRPhenomPv2_FakeUncertainty +from heron.models.lalnoise import AdvancedLIGO +from heron.injection import make_injection, make_injection_zero_noise +from heron.detector import Detector, AdvancedLIGOHanford, AdvancedLIGOLivingston, AdvancedVirgo +from heron.likelihood import MultiDetector, TimeDomainLikelihood, TimeDomainLikelihoodModelUncertainty, TimeDomainLikelihoodPyTorch +#, TimeDomainLikelihoodModelUncertaintyPyTorch + +from heron.inference import heron_inference, parse_dict, load_yaml + +from torch.cuda import is_available + +CUDA_NOT_AVAILABLE = not is_available() + + +class Test_Likelihood_ZeroNoise_With_Uncertainty(unittest.TestCase): + """ + Test likelihoods on a zero noise injection. + """ + + def setUp(self): + self.waveform = IMRPhenomPv2_FakeUncertainty() + self.psd_model = AdvancedLIGO() + + self.injections = make_injection_zero_noise(waveform=IMRPhenomPv2, + injection_parameters={"distance": 1000*u.megaparsec, + "mass_ratio": 0.6, + "gpstime": 0, + "total_mass": 60 * u.solMass}, + detectors={"AdvancedLIGOHanford": "AdvancedLIGO", + "AdvancedLIGOLivingston": "AdvancedLIGO"} + ) + + + + def test_likelihood_maximum_at_true_value_mass_ratio(self): + + data = self.injections['H1'] + + likelihood = TimeDomainLikelihoodModelUncertainty(data, psd=self.psd_model) + mass_ratios = np.linspace(0.1, 1.0, 100) + + log_likes = [] + for mass_ratio in mass_ratios: + + test_waveform = self.waveform.time_domain(parameters={"distance": 1000*u.megaparsec, + "mass_ratio": mass_ratio, + "gpstime": 0, + "total_mass": 60 * u.solMass}, times=likelihood.times) + projected_waveform = test_waveform.project(AdvancedLIGOHanford(), + ra=0, dec=0, + gpstime=0, + phi_0=0, psi=0, + iota=0) + + log_likes.append(likelihood.log_likelihood(projected_waveform, norm=False)) + + self.assertTrue(np.abs(mass_ratios[np.argmax(log_likes)] - 0.6) < 0.1)