From ca1392b83e209072ee7bc23fa7d5d7329876bedd Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 15:04:05 +0100 Subject: [PATCH 01/25] Fixed the lalnoise implementation. --- heron/models/lalnoise.py | 27 ++++++++------------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/heron/models/lalnoise.py b/heron/models/lalnoise.py index 528a569..ddd3783 100644 --- a/heron/models/lalnoise.py +++ b/heron/models/lalnoise.py @@ -38,7 +38,7 @@ def frequency_domain( ) self.psd_function(psd_data, flow=lower_frequency) psd_data = psd_data.data.data - #psd_data[frequencies < mask_below] = psd_data[frequencies > mask_below][0] + psd_data[frequencies < mask_below] = psd_data[frequencies > mask_below][0] psd = PSD(psd_data, frequencies=frequencies) return psd @@ -46,25 +46,14 @@ def covariance_matrix(self, times): """ Return a time-domain representation of this power spectral density. """ - dt = times[1] - times[0] N = len(times) T = times[-1] - times[0] df = 1 / T frequencies = torch.arange(len(times) // 2 + 1) * df.value - psd = np.array(self.frequency_domain(df=df, frequencies=frequencies).data) - psd[-1] = psd[-2] - # import matplotlib.pyplot as plt - # f, ax = plt.subplots(1,1) - # ax.plot(frequencies, psd) - # f.savefig("psd.png") - # Calculate the autocovariance from a one-sided PSD - acf = 0.5*np.real(np.fft.irfft(psd*df, n=(N)))*T - # The covariance is then the Toeplitz matrix formed from the acf - # f, ax = plt.subplots(1,1) - # ax.plot(acf) - # f.savefig("acf.png") - return scipy.linalg.toeplitz(acf) + psd = self.frequency_domain(df=df, frequencies=frequencies) + ts = np.fft.irfft(psd, n=(N)) # * (N*N/dt/dt/2), n=(N)) + return scipy.linalg.toeplitz(ts) def time_domain(self, times): return self.covariance_matrix(times) @@ -94,17 +83,17 @@ def time_series(self, times): frequencies = torch.arange(len(times) // 2 + 1) * df reals = np.random.randn(len(frequencies)) imags = np.random.randn(len(frequencies)) - psd = np.array(self.frequency_domain(df=df, frequencies=frequencies).data) - psd[-1] = psd[-2] - S = 0.5 * np.sqrt(psd / df) #* T inside sqrt # np.sqrt(N * N / 4 / (T) * psd.value) + psd = self.frequency_domain(frequencies=frequencies) + + S = 0.5 * np.sqrt(psd.value * T) # np.sqrt(N * N / 4 / (T) * psd.value) noise_r = S * (reals) noise_i = S * (imags) noise_f = noise_r + 1j * noise_i - return TimeSeries(data=np.fft.irfft(noise_f, n=(N))*df*N, times=times) + return TimeSeries(data=np.fft.irfft(noise_f, n=(N)), times=times) class AdvancedLIGODesignSensitivity2018(LALSimulationPSD): From 1e7dbd9106f70ed1254bb24ae74cbdb823c62c1d Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 15:45:27 +0100 Subject: [PATCH 02/25] Updates to inference tests. --- tests/test_inference.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_inference.py b/tests/test_inference.py index d91b895..4f525fc 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -19,7 +19,8 @@ from heron.models.lalnoise import AdvancedLIGO from heron.injection import make_injection from heron.detector import Detector, AdvancedLIGOHanford, AdvancedLIGOLivingston, AdvancedVirgo -from heron.likelihood import MultiDetector, TimeDomainLikelihood, TimeDomainLikelihoodModelUncertainty, TimeDomainLikelihoodPyTorch, TimeDomainLikelihoodModelUncertaintyPyTorch +from heron.likelihood import MultiDetector, TimeDomainLikelihood, TimeDomainLikelihoodModelUncertainty +# TimeDomainLikelihoodPyTorch, TimeDomainLikelihoodModelUncertaintyPyTorch from heron.inference import heron_inference, parse_dict, load_yaml @@ -55,6 +56,9 @@ def test_snr(self): test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, "m2": 30*u.solMass, "distance": 1000 * u.megaparsec}, times=likelihood.times) + + print("noise max", likelihood.C.max()) + print("waveform max", np.array(test_waveform['plus'].data).max()) snr = likelihood.snr(test_waveform.project(AdvancedLIGOHanford(), ra=0, dec=0, phi_0=0, psi=0, @@ -176,6 +180,8 @@ def test_parser_psds(self): # def test_sampler(self): # heron_inference("tests/test_inference_config.yaml") + +@unittest.skip("Skipping gpytorch tests until these are nearer being ready") class Test_PyTorch(unittest.TestCase): """Test that the pytorch likelihoods work.""" From fe12e3d705c45712a09a62956f96fccab7e41edb Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 16:20:46 +0100 Subject: [PATCH 03/25] Attempts to fix snr test. --- heron/injection.py | 49 ----------------------------------------- heron/likelihood.py | 9 ++++---- tests/test_inference.py | 26 ++++++++++++---------- 3 files changed, 19 insertions(+), 65 deletions(-) diff --git a/heron/injection.py b/heron/injection.py index e4178a8..665d253 100644 --- a/heron/injection.py +++ b/heron/injection.py @@ -63,52 +63,6 @@ def make_injection( return injections -def make_injection_zero_noise( - waveform=IMRPhenomPv2, - injection_parameters={}, - times=None, - detectors=None, - framefile=None, -): - - parameters = {"ra": 0, "dec": 0, "psi": 0, "theta_jn": 0, "phase": 0, 'gpstime': 4000} - parameters.update(injection_parameters) - - waveform = waveform() - - if times is None: - times = np.linspace(-0.5, 0.1, int(0.6 * 4096)) + parameters['gpstime'] - waveform = waveform.time_domain( - parameters, - times=times, - ) - - injections = {} - for detector, psd_model in detectors.items(): - detector = KNOWN_IFOS[detector]() - channel = f"{detector.abbreviation}:Injection" - logger.info(f"Making injection for {detector} in channel {channel}") - psd_model = KNOWN_PSDS[psd_model]() - #data = psd_model.time_series(times) - - # import matplotlib - # matplotlib.use("agg") - # from gwpy.plot import Plot - # f = Plot(data, waveform.project(detector), data+waveform.project(detector), separate=False) - # f.savefig(f"{detector.abbreviation}_injected_waveform.png") - - injection = waveform.project(detector) - injection.channel = channel - injections[detector.abbreviation] = injection - likelihood = TimeDomainLikelihood(injection, psd=psd_model) - snr = likelihood.snr(waveform.project(detector)) - logger.info(f"Optimal Filter SNR: {snr}") - - if framefile: - filename = f"{detector.abbreviation}_{framefile}.gwf" - logger.info(f"Saving framefile to {filename}") - injection.write(filename, format="gwf") - def injection_parameters_add_units(parameters): UNITS = {"luminosity_distance": u.megaparsec, "m1": u.solMass, "m2": u.solMass} @@ -121,9 +75,6 @@ def injection_parameters_add_units(parameters): @click.command() @click.option("--settings") def injection(settings): - - click.secho("Creating an injection with the heron injection engine") - settings = load_yaml(settings) if "logging" in settings: diff --git a/heron/likelihood.py b/heron/likelihood.py index c0e28eb..e0bc335 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -92,11 +92,12 @@ def snr(self, waveform): """ Calculate the signal to noise ratio for a given waveform. """ - factor = 1e30 + factor = 1e22 N = len(self.times) - h_h = ( - (np.array(waveform.data).T*factor @ self.solve(self.C*factor**2, np.array(waveform.data)*factor)) * self.dt - ) + dt = (self.times[1] - self.times[0]).value + T = (self.times[-1] - self.times[0]).value + print("dt", dt, "N", N) + h_h = (np.array(waveform.data).T @ self.solve(self.C, np.array(waveform.data))) * (dt * dt / N / 4) return np.sqrt(np.abs(h_h)) def snr_f(self, waveform): diff --git a/tests/test_inference.py b/tests/test_inference.py index 4f525fc..cb99d5a 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -55,7 +55,7 @@ def test_snr(self): test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, "m2": 30*u.solMass, - "distance": 1000 * u.megaparsec}, times=likelihood.times) + "distance": 410 * u.megaparsec}, times=data.times) print("noise max", likelihood.C.max()) print("waveform max", np.array(test_waveform['plus'].data).max()) @@ -63,22 +63,24 @@ def test_snr(self): ra=0, dec=0, phi_0=0, psi=0, iota=0)) + print("snr", snr) self.assertTrue(snr > 40 and snr < 45) - # def test_snr_f(self): - # data = self.injections['H1'] + def test_snr_f(self): + data = self.injections['H1'] - # likelihood = TimeDomainLikelihood(data, psd=self.psd_model) + likelihood = TimeDomainLikelihood(data, psd=self.psd_model) - # test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, - # "m2": 30*u.solMass, - # "distance": 410 * u.megaparsec}, times=data.times) + test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, + "m2": 30*u.solMass, + "distance": 410 * u.megaparsec}, times=data.times) - # snr = likelihood.snr_f(test_waveform.project(AdvancedLIGOHanford(), - # ra=0, dec=0, - # phi_0=0, psi=0, - # iota=0)) - # self.assertTrue(snr > 80 and snr < 90) + snr = likelihood.snr_f(test_waveform.project(AdvancedLIGOHanford(), + ra=0, dec=0, + phi_0=0, psi=0, + iota=0)) + print("f-domain snr", snr) + self.assertTrue(snr > 80 and snr < 90) From 58e6e14f76c9bc33078210809982c5e399561fa9 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 16:59:51 +0100 Subject: [PATCH 04/25] Reverted various changes which break the GPR construction again. --- heron/likelihood.py | 499 +++++++++++++++++---------------------- heron/models/__init__.py | 36 +-- heron/types.py | 28 ++- tests/test_inference.py | 39 +-- 4 files changed, 268 insertions(+), 334 deletions(-) diff --git a/heron/likelihood.py b/heron/likelihood.py index e0bc335..32371db 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -3,8 +3,6 @@ using the waveform models it supports. """ -from gwpy.timeseries import TimeSeries - import numpy as np import torch @@ -26,15 +24,8 @@ class LikelihoodBase: class Likelihood(LikelihoodBase): - def abs(self, A): - return np.abs(A) - - def einsum(self, *args, **kwargs): - return np.einsum(*args, **kwargs) - def logdet(self, K): - (sign, logabsdet) = np.linalg.slogdet(K) - return logabsdet + return np.linalg.slogdet(K).logabsdet def inverse(self, A): return np.linalg.inv(A) @@ -55,26 +46,24 @@ def pi(self): class TimeDomainLikelihood(Likelihood): - def __init__(self, data, psd, waveform=None, detector=None, fixed_parameters={}, timing_basis=None): + def __init__( + self, + data, + psd, + waveform=None, + detector=None, + fixed_parameters={}, + timing_basis=None, + ): self.psd = psd self.data = np.array(data.data) - self.data_ts = data self.times = data.times - - self.logger = logger = logging.getLogger( - "heron.likelihood.TimeDomainLikelihood" - ) - self.N = len(self.times) + self.C = self.psd.covariance_matrix(times=self.times) - factor = 1e30 - self.inverse_C = self.inverse(self.C*factor**2) - self.dt = self.abs((self.times[1] - self.times[0]).value) - - self.normalisation = - (self.N/2) * self.log(2*self.pi) + (self.logdet(self.C*1e30) - self.log(1e30)) *self.dt - #* (self.dt * self.dt / 4) / 4 - self.logger.info(f"Normalisation: {self.normalisation}") + self.inverse_C = np.linalg.inv(self.C) + self.dt = (self.times[1] - self.times[0]).value self.N = len(self.times) if waveform is not None: @@ -85,127 +74,96 @@ def __init__(self, data, psd, waveform=None, detector=None, fixed_parameters={}, self.fixed_parameters = fixed_parameters if timing_basis is not None: - self.fixed_parameters['reference_frame'] = timing_basis + self.fixed_parameters["reference_frame"] = timing_basis + self.logger = logger = logging.getLogger( + "heron.likelihood.TimeDomainLikelihood" + ) def snr(self, waveform): """ Calculate the signal to noise ratio for a given waveform. """ - factor = 1e22 - N = len(self.times) dt = (self.times[1] - self.times[0]).value - T = (self.times[-1] - self.times[0]).value - print("dt", dt, "N", N) - h_h = (np.array(waveform.data).T @ self.solve(self.C, np.array(waveform.data))) * (dt * dt / N / 4) + N = len(self.times) + h_h = ( + (np.array(waveform.data).T @ self.solve(self.C, np.array(waveform.data))) + * (dt * dt / N / 4) + / 4 + ) return np.sqrt(np.abs(h_h)) - def snr_f(self, waveform): - dt = (self.times[1] - self.times[0]).value - T = (self.times[-1] - self.times[0]).value - wf = np.fft.rfft(waveform.data * dt) - S = self.psd.frequency_domain(frequencies=np.arange(0, 0.5/dt, 1/T)) - A = (4 * (wf.conj()*wf)[:-1] / S.value[:-1] / T) - return np.sqrt(np.real(np.sum(A))) - - def log_likelihood(self, waveform, norm=True): - """ - Calculate the log likelihood of a given waveform and the data. - - Parameters - ---------- - waveform : `heron.types.Waveform` - The waveform to compare to the data. - - Returns - ------- - float - The log-likelihood for the waveform. - """ - factor = 1e30 - try: - assert(np.all(self.times == waveform.times)) - except AssertionError: - print(self.times, waveform.times) - raise Exception - residual = (self.data * factor) - (np.array(waveform.data) * factor) - weighted_residual = (residual.T @ self.inverse_C @ residual) * (self.dt) - # why is this negative using Toeplitz? - self.logger.info(f"residual: {residual}; chisq: {weighted_residual}") - out = - 0.5 * weighted_residual - if norm: - out += self.normalisation - return out + def log_likelihood(self, waveform): + residual = np.array(self.data.data) - np.array(waveform.data) + weighted_residual = ( + (residual) @ self.solve(self.C, residual) * (self.dt * self.dt / 4) / 4 + ) + normalisation = self.logdet(2 * np.pi * self.C) + return -0.5 * weighted_residual + 0.5 * normalisation def __call__(self, parameters): self.logger.info(parameters) keys = set(parameters.keys()) - extrinsic = {"phase", "psi", "ra", "dec", "theta_jn", "zenith", "azimuth", "gpstime"} - conversions = {"geocent_time", "mass_ratio", "total_mass", "luminosity_distance", "chirp_mass"} - bad_keys = keys - set(self.waveform.allowed_parameters) - extrinsic - conversions + extrinsic = {"phase", "psi", "ra", "dec", "theta_jn"} + conversions = {"mass_ratio", "total_mass", "luminosity_distance"} + bad_keys = keys - set(self.waveform._args.keys()) - extrinsic - conversions if len(bad_keys) > 0: print("The following keys were not recognised", bad_keys) - if self.fixed_parameters: - parameters.update(self.fixed_parameters) + parameters.update(self.fixed_parameters) test_waveform = self.waveform.time_domain( parameters=parameters, times=self.times ) projected_waveform = test_waveform.project(self.detector) - llike = self.log_likelihood(projected_waveform) - self.logger.info(f"log likelihood: {llike}") - return llike + return self.log_likelihood(projected_waveform) class TimeDomainLikelihoodModelUncertainty(TimeDomainLikelihood): - def __init__(self, data, psd, waveform=None, detector=None, fixed_parameters=None, timing_basis=None): - super().__init__(data, psd, waveform, detector, fixed_parameters, timing_basis) + def __init__(self, data, psd, waveform=None, detector=None): + super().__init__(data, psd, waveform, detector) + + def _normalisation(self, K, S): + norm = ( + -1.5 * K.shape[0] * self.log(2 * self.pi) + - 0.5 * self.logdet(K) + + 0.5 * self.logdet(self.C) + - 0.5 * self.logdet(self.solve(K, self.C) + self.eye(K.shape[0])) + ) + return norm def _weighted_data(self): """Return the weighted data component""" # TODO This can all be pre-computed - factor = 1e23 - factor_sq = factor**2 - if not hasattr(self, "weighted_data_CACHE"): dw = self.weighted_data_CACHE = ( - -0.5 * (self.data*factor) @ self.solve(self.C*factor_sq, self.data*factor) + -0.5 * self.data.T @ self.solve(self.C, self.data) ) else: dw = self.weighted_data_CACHE return dw - def log_likelihood(self, waveform, norm=True): - - waveform_d = np.array(waveform.data) - waveform_c = np.array(waveform.covariance) - K = waveform_c - mu = waveform_d - factor = 1/np.max(K) - factor_sq = factor**2 - factor_sqi = factor**-2 - - K = K*factor_sq - C = self.C * factor_sq - Ci = self.inverse(self.C * factor_sq) - Ki = self.inverse(K) - A = self.inverse(C + K) - mu = mu * factor - data = self.data*factor - - sigma = self.inverse(Ki+Ci)*factor_sqi - B = (self.einsum("ij,i", Ki, mu) + (self.einsum("ij,i", Ci, data)))*factor - - N = - (self.N / 2) * self.log(2*self.pi) + 0.5 * (self.logdet(sigma) - self.logdet(self.C) - self.logdet(K) - 3 * self.log(factor_sq)) * self.dt + 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 log_likelihood(self, waveform): + like = self._weighted_cross(waveform.data, waveform.covariance) + # print("cross term", like) + A = self._weighted_data() + # print("data term", A) + B = self._weighted_model(waveform.data, waveform.covariance) + # print("model term", B) + like = like + A + B + norm = self._normalisation(waveform.covariance, self.C) + # print("normalisation", norm) + like += norm - data_like = (self._weighted_data()) - model_like = -0.5 * self.einsum("i,ij,j", mu, Ki, mu) - shift = + 0.5 * self.einsum('i,ij,j', B, sigma, B) #* ((self.dt) / 4) - like = data_like + model_like + shift - - if norm: - like += N return like @@ -219,183 +177,164 @@ def __init__(self, *args): for detector in args: if isinstance(detector, LikelihoodBase): self._likelihoods.append(detector) - self.logger = logger = logging.getLogger( - "heron.likelihood.MultiDetector" - ) def __call__(self, parameters): out = 0 - self.logger.info(f"Calling likelihood at {parameters}") for detector in self._likelihoods: out += detector(parameters) return out -# class LikelihoodPyTorch(Likelihood): - -# def logdet(self, K): -# A = torch.slogdet(K) -# return A.logabsdet #* A.sign - -# def inverse(self, A): -# out, info = torch.linalg.inv_ex(A) -# if info == 0: -# return out -# else: -# raise ValueError(f"Matrix could not be inverted: {info}") - -# def solve(self, A, B): -# return torch.linalg.solve(A, B) - -# def abs(self, A): -# return torch.abs(A) - -# def eye(self, N, *args, **kwargs): -# return torch.eye(N, device=self.device, dtype=torch.double) - -# def log(self, A): -# return torch.log(A) - -# @property -# def pi(self): -# return torch.tensor(torch.pi, device=self.device) - -# def einsum(self, *args, **kwargs): -# return torch.einsum(*args, **kwargs) - - -# class TimeDomainLikelihoodPyTorch(LikelihoodPyTorch): - -# def __init__(self, data, psd, waveform=None, detector=None, fixed_parameters={}, timing_basis=None): -# self.logger = logger = logging.getLogger( -# "heron.likelihood.TimeDomainLikelihoodPyTorch" -# ) -# self.device = device -# self.logger.info(f"Using device {device}") -# self.psd = psd - -# self.data = torch.tensor(data.data, device=self.device, dtype=torch.double) -# self.times = data.times - -# self.C = self.psd.covariance_matrix(times=self.times) -# self.C = torch.tensor(self.C, device=self.device) -# self.inverse_C = self.inverse(self.C) - -# self.dt = (self.times[1] - self.times[0]).value -# self.N = len(self.times) - -# if waveform is not None: -# self.waveform = waveform - -# if detector is not None: -# self.detector = detector - -# self.fixed_parameters = fixed_parameters -# if timing_basis is not None: -# self.fixed_parameters['reference_frame'] = timing_basis - -# def snr(self, waveform): -# """ -# Calculate the signal to noise ratio for a given waveform. -# """ -# dt = (self.times[1] - self.times[0]).value -# self.logger.info(f"Contains {self.N} points with dt {dt}.") -# waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) -# h_h = (waveform_d @ self.solve(self.C, waveform_d)) * (dt / 4) -# return 2*torch.sqrt(torch.abs(h_h)) - -# def snr_f(self, waveform): -# dt = (self.times[1] - self.times[0]).value -# T = (self.times[-1] - self.times[0]).value -# wf = np.fft.rfft(waveform.data * dt) -# # Single-sided PSD -# S = self.psd.frequency_domain(frequencies=np.arange(0, 0.5/dt, 1/T)) -# A = (4 * (wf.conj()*wf)[:-1] / (S.value[:-1]) / T) -# return np.sqrt(np.real(np.sum(A))) - - -# def log_likelihood(self, waveform, norm=True): -# waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) -# residual = self.data - waveform_d -# weighted_residual = ( -# (residual) @ self.solve(self.C, residual) * (self.dt * self.dt / 4) / 4 -# ) -# normalisation = self.logdet(2 * np.pi * self.C) * self.N -# #print("normalisation", normalisation) -# like = -0.5 * weighted_residual -# if norm: -# like -= 0.5 * normalisation -# return like - -# def __call__(self, parameters): -# self.logger.info("Called with", parameters) - -# keys = set(parameters.keys()) -# extrinsic = {"phase", "psi", "ra", "dec", "theta_jn", "zenith", "azimuth"} -# conversions = {"mass_ratio", "total_mass", "luminosity_distance"} -# bad_keys = keys - set(self.waveform._args.keys()) - extrinsic - conversions -# if len(bad_keys) > 0: -# print("The following keys were not recognised", bad_keys) -# parameters.update(self.fixed_parameters) -# test_waveform = self.waveform.time_domain( -# parameters=parameters, times=self.times -# ) -# projected_waveform = test_waveform.project(self.detector) -# return self.log_likelihood(projected_waveform) - - -# class TimeDomainLikelihoodModelUncertaintyPyTorch(TimeDomainLikelihoodPyTorch): - -# def __init__(self, data, psd, waveform=None, detector=None, fixed_parameters=None, timing_basis=None): -# super().__init__(data, psd, waveform, detector, fixed_parameters, timing_basis) - -# def _weighted_data(self): -# """Return the weighted data component""" -# # TODO This can all be pre-computed -# factor = 1e23 -# factor_sq = factor**2 - -# if not hasattr(self, "weighted_data_CACHE"): -# dw = self.weighted_data_CACHE = ( -# -0.5 * (self.data*factor) @ self.solve(self.C*factor_sq, self.data*factor) # * (self.dt * self.dt / 4) / 4 -# ) -# else: -# dw = self.weighted_data_CACHE -# return dw - -# def log_likelihood(self, waveform, norm=True): -# print("size", waveform.covariance.size) -# waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) -# waveform_c = torch.tensor( -# waveform.covariance, device=self.device, dtype=torch.double -# ) -# K = waveform_c -# mu = waveform_d - -# #print(torch.sum(mu.cpu()-self.data.cpu())) - -# factor = 1/torch.max(K) -# factor_sq = factor**2 -# factor_sqi = factor**-2 - -# K = K*factor_sq -# # C = self.C * factor_sq -# # Ci = self.inverse(self.C * factor_sq) -# #Ki = self.inverse(K) -# # A = Ci + Ki #self.inverse(C + K) -# # mu = mu * factor -# # data = self.data*factor - -# # sigma = self.inverse(A)*factor_sqi -# # B = (self.einsum("ij,i", Ki, mu) + (self.einsum("ij,i", Ci, data)))*factor -# # N = - (self.N / 2) * self.log(2*self.pi) + 0.5 * (self.logdet(sigma) - self.logdet(self.C) - self.logdet(K) - 3 * self.log(factor_sq)) * self.dt -# data_like = (self._weighted_data()) -# model_like = -0.5 * mu @ self.solve(K, mu) #self.einsum("i,ij,j", mu, Ki, mu) -# # shift = + 0.5 * self.einsum('i,ij,j', B, sigma, B) #* ((self.dt) / 4) -# like = data_like + model_like # + shift - -# # if norm: -# # like += N - -# return like +class LikelihoodPyTorch(Likelihood): + + def logdet(self, K): + return torch.slogdet(K).logabsdet + + def inverse(self, A): + out, info = torch.linalg.inv_ex(A) + if info == 0: + return out + else: + raise ValueError(f"Matrix could not be inverted: {info}") + + def solve(self, A, B): + return torch.linalg.solve(A, B) + + def eye(self, N, *args, **kwargs): + return torch.eye(N, device=self.device, dtype=torch.double) + + def log(self, A): + return torch.log(A) + + @property + def pi(self): + return torch.tensor(torch.pi, device=self.device) + + +class TimeDomainLikelihoodPyTorch(LikelihoodPyTorch): + + def __init__( + self, + data, + psd, + waveform=None, + detector=None, + fixed_parameters={}, + timing_basis=None, + ): + self.logger = logger = logging.getLogger( + "heron.likelihood.TimeDomainLikelihoodPyTorch" + ) + self.device = device + self.logger.info(f"Using device {device}") + self.psd = psd + + self.data = torch.tensor(data.data, device=self.device, dtype=torch.double) + self.times = data.times + + self.C = self.psd.covariance_matrix(times=self.times) + self.C = torch.tensor(self.C, device=self.device) + self.inverse_C = torch.linalg.inv(self.C) + + self.dt = (self.times[1] - self.times[0]).value + self.N = len(self.times) + + if waveform is not None: + self.waveform = waveform + + if detector is not None: + self.detector = detector + + self.fixed_parameters = fixed_parameters + if timing_basis is not None: + self.fixed_parameters["reference_frame"] = timing_basis + + def snr(self, waveform): + """ + Calculate the signal to noise ratio for a given waveform. + """ + dt = (self.times[1] - self.times[0]).value + N = len(self.times) + waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) + h_h = (waveform_d.T @ self.solve(self.C, waveform_d)) * (dt * dt / N / 4) / 4 + return torch.sqrt(torch.abs(h_h)) + + def log_likelihood(self, waveform): + waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) + residual = self.data - waveform_d + weighted_residual = ( + (residual) @ self.solve(self.C, residual) * (self.dt * self.dt / 4) / 4 + ) + normalisation = self.logdet(2 * np.pi * self.C) + return -0.5 * weighted_residual + 0.5 * normalisation + + def __call__(self, parameters): + self.logger.info(parameters) + + keys = set(parameters.keys()) + extrinsic = {"phase", "psi", "ra", "dec", "theta_jn"} + conversions = {"mass_ratio", "total_mass", "luminosity_distance"} + bad_keys = keys - set(self.waveform._args.keys()) - extrinsic - conversions + if len(bad_keys) > 0: + print("The following keys were not recognised", bad_keys) + parameters.update(self.fixed_parameters) + test_waveform = self.waveform.time_domain( + parameters=parameters, times=self.times + ) + projected_waveform = test_waveform.project(self.detector) + return self.log_likelihood(projected_waveform) + + +class TimeDomainLikelihoodModelUncertaintyPyTorch(TimeDomainLikelihoodPyTorch): + + def __init__(self, data, psd, waveform=None, detector=None): + super().__init__(data, psd, waveform, detector) + + def _normalisation(self, K, S): + norm = ( + -1.5 * K.shape[0] * self.log(2 * self.pi) + - 0.5 * self.logdet(K) + + 0.5 * self.logdet(self.C) + - 0.5 * self.logdet(self.solve(K, self.C) + self.eye(K.shape[0])) + ) + return norm + + def _weighted_data(self): + """Return the weighted data component""" + # TODO This can all be pre-computed + if not hasattr(self, "weighted_data_CACHE"): + dw = self.weighted_data_CACHE = ( + -0.5 * self.data.T @ self.solve(self.C, self.data) + ) + else: + dw = self.weighted_data_CACHE + return dw + + def _weighted_model(self, mu, K): + """Return the inner product of the GPR mean""" + mu = torch.tensor(mu, device=self.device, dtype=torch.double) + return -0.5 * 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 log_likelihood(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 + ) + like = self._weighted_cross(waveform_d, waveform_c) + # print("cross term", like) + A = self._weighted_data() + # print("data term", A) + B = self._weighted_model(waveform_d, waveform_c) + # print("model term", B) + like = like + A + B + norm = self._normalisation(waveform_c, self.C) + # print("normalisation", norm) + like += norm + + return like diff --git a/heron/models/__init__.py b/heron/models/__init__.py index e535f84..f3ba0fd 100644 --- a/heron/models/__init__.py +++ b/heron/models/__init__.py @@ -1,5 +1,5 @@ import torch -from lal import antenna, MSUN_SI +from lal import antenna from astropy import units as u @@ -25,32 +25,14 @@ def _convert_luminosity_distance(self, args): args["distance"] = args.pop("luminosity_distance") return args - def _convert_mass_ratio_total_mass(self, args): - """ - Convert a mass ratio and a total mass into individual component masses. - If the masses have no units they are assumed to be in SI units. - - Parameters - ---------- - args['total_mass'] : float, `astropy.units.Quantity` - The total mass for the system. - args['mass_ratio'] : float - The mass ratio of the system using the convention m2/m1 - """ - args["m1"] = (args["total_mass"] / (1 + args["mass_ratio"])) - args["m2"] = (args["total_mass"] / (1 + (1 / args["mass_ratio"]))) - # Do these have units? - # If not then we can skip some relatively expensive operations and apply a heuristic. - if isinstance(args["m1"], u.Quantity): - args["m1"] = args["m1"].to_value(u.kilogram) - args["m2"] = args["m2"].to_value(u.kilogram) - if (not isinstance(args["m1"], u.Quantity)) and (args["m1"] < 1000): - # This appears to be in solar masses - args["m1"] *= MSUN_SI - if (not isinstance(args["m2"], u.Quantity)) and (args["m2"] < 1000): - # This appears to be in solar masses - args["m2"] *= MSUN_SI - + def _convert_mass_ratio_total_mass(self, args): + + args["m1"] = (args["total_mass"] / (1 + args["mass_ratio"])).to_value( + u.kilogram + ) + args["m2"] = (args["total_mass"] / (1 + 1 / args["mass_ratio"])).to_value( + u.kilogram + ) args.pop("total_mass") args.pop("mass_ratio") diff --git a/heron/types.py b/heron/types.py index 3a64760..14346e6 100644 --- a/heron/types.py +++ b/heron/types.py @@ -102,10 +102,9 @@ def project( The declination of the signal source. """ - if not time: time = self.waveforms["plus"].epoch.value - + if ((ra is None) and (dec is None)) and ( ("ra" in self._parameters) and ("dec" in self._parameters) ): @@ -114,21 +113,30 @@ def project( dt = detector.geocentre_delay(ra=ra, dec=dec, times=time) - elif ("azimuth" in self._parameters.keys()) and ("zenith" in self._parameters.keys()) and ("reference_frame" in self._parameters.keys()): + elif ( + ("azimuth" in self._parameters.keys()) + and ("zenith" in self._parameters.keys()) + and ("reference_frame" in self._parameters.keys()) + ): # Use horizontal coordinates. det1 = cached_detector_by_prefix[self._parameters["reference_frame"][0]] det2 = cached_detector_by_prefix[self._parameters["reference_frame"][1]] - tg, ra, dec = DetFrameToEquatorial( - det1, det2, time, self._parameters["azimuth"], self._parameters["zenith"] + ra, dec, dt = DetFrameToEquatorial( + det1, + det2, + time, + self._parameters["azimuth"], + self._parameters["zenith"], ) - dt = time - tg + elif (ra is None) and (dec is None): raise ValueError("Right ascension and declination must both be specified.") else: dt = detector.geocentre_delay(ra=ra, dec=dec, times=time) + if "plus" in self.waveforms and "cross" in self.waveforms: - + if not iota and "theta_jn" in self._parameters: iota = self._parameters["theta_jn"] elif isinstance(iota, type(None)): @@ -180,15 +188,15 @@ def project( else: projected_covariance = None - bins = dt / (self.waveforms["plus"].dt) - projected_waveform = Waveform( - data=array_library.roll(array_library.pad(projected_data, 5000), int(bins.value))[5000:-5000], + data=projected_data, variance=projected_variance, covariance=projected_covariance, times=self.waveforms["plus"].times, ) + projected_waveform.shift(dt) + return projected_waveform else: diff --git a/tests/test_inference.py b/tests/test_inference.py index cb99d5a..e56416f 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -55,32 +55,35 @@ def test_snr(self): test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, "m2": 30*u.solMass, - "distance": 410 * u.megaparsec}, times=data.times) + "distance": 1000 * u.megaparsec}, times=data.times) - print("noise max", likelihood.C.max()) - print("waveform max", np.array(test_waveform['plus'].data).max()) - snr = likelihood.snr(test_waveform.project(AdvancedLIGOHanford(), + projected_waveform = test_waveform.project(AdvancedLIGOHanford(), ra=0, dec=0, phi_0=0, psi=0, - iota=0)) + iota=0) + + f = projected_waveform.plot() + f.savefig("projected_waveform.png") + + snr = likelihood.snr(projected_waveform) print("snr", snr) self.assertTrue(snr > 40 and snr < 45) - def test_snr_f(self): - data = self.injections['H1'] + # def test_snr_f(self): + # data = self.injections['H1'] - likelihood = TimeDomainLikelihood(data, psd=self.psd_model) + # likelihood = TimeDomainLikelihood(data, psd=self.psd_model) - test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, - "m2": 30*u.solMass, - "distance": 410 * u.megaparsec}, times=data.times) + # test_waveform = self.waveform.time_domain(parameters={"m1": 35*u.solMass, + # "m2": 30*u.solMass, + # "distance": 410 * u.megaparsec}, times=data.times) - snr = likelihood.snr_f(test_waveform.project(AdvancedLIGOHanford(), - ra=0, dec=0, - phi_0=0, psi=0, - iota=0)) - print("f-domain snr", snr) - self.assertTrue(snr > 80 and snr < 90) + # snr = likelihood.snr_f(test_waveform.project(AdvancedLIGOHanford(), + # ra=0, dec=0, + # phi_0=0, psi=0, + # iota=0)) + # print("f-domain snr", snr) + # self.assertTrue(snr > 80 and snr < 90) @@ -117,6 +120,7 @@ def test_likelihood_with_uncertainty(self): log_like = likelihood.log_likelihood(projected_waveform) + @unittest.skip("The likelihood with uncertainty isn't working yet.") def test_sampling_with_uncertainty(self): waveform = IMRPhenomPv2_FakeUncertainty() likelihood = TimeDomainLikelihoodModelUncertainty(self.injections['H1'], @@ -137,6 +141,7 @@ def test_sampling_with_uncertainty(self): log_like = likelihood(parameters=parameters) self.assertTrue(-2400 < log_like < -2200) + @unittest.skip("The likelihood with uncertainty isn't working yet.") def test_sampling_with_uncertainty_multi(self): waveform = IMRPhenomPv2_FakeUncertainty() likelihood = MultiDetector(TimeDomainLikelihoodModelUncertainty(self.injections['H1'], From 5fb79d08aac779e9a9759c9eec7078cb67d7a0d1 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 17:10:25 +0100 Subject: [PATCH 05/25] Update the github actions. --- .github/workflows/documentation.yaml | 37 ++++++++++++++++++++++++ .github/workflows/review.yml | 43 +++++++--------------------- 2 files changed, 47 insertions(+), 33 deletions(-) create mode 100644 .github/workflows/documentation.yaml diff --git a/.github/workflows/documentation.yaml b/.github/workflows/documentation.yaml new file mode 100644 index 0000000..4e0281a --- /dev/null +++ b/.github/workflows/documentation.yaml @@ -0,0 +1,37 @@ +name: Documentation +on: + push: + branches: + - master + - v*-preview + +jobs: + docs: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + + - name: Set up Python 3.8 + uses: actions/setup-python@v2 + with: + python-version: 3.8 + + - name: Build HTML Docs + run: | + sudo apt install pandoc + pip install sphinx + pip install kentigern + pip install -r requirements.txt + pip install -r requirements_dev.txt + pip install . + cd docs + make multi + + - name: SCP Deploy HTML Docs + uses: horochx/deploy-via-scp@v1.0.1 + with: + local: docs/_build/html/* + remote: /home/danwilliams/code.daniel-williams.co.uk/heron/ + host: ${{ secrets.sshhost }} + user: ${{ secrets.sshuser }} + key: ${{ secrets.sshkey }} diff --git a/.github/workflows/review.yml b/.github/workflows/review.yml index 975c0bf..de2f11a 100644 --- a/.github/workflows/review.yml +++ b/.github/workflows/review.yml @@ -1,23 +1,16 @@ -# This is a basic workflow to help you get started with Actions - -name: CI - -# Controls when the action will run. Triggers the workflow on push or pull request -# events but only for the master branch +name: Tests on: push: - branches: [ review ] + branches: + - master + - v*-preview pull_request: - branches: [ review ] + branches: + - master -# A workflow run is made up of one or more jobs that can run sequentially or in parallel jobs: - # This workflow contains a single job called "build" - build: - # The type of runner that the job will run on + test: runs-on: ubuntu-latest - - # Steps represent a sequence of tasks that will be executed as part of the job steps: - uses: actions/checkout@v2 @@ -25,23 +18,7 @@ jobs: uses: actions/setup-python@v2 with: python-version: 3.8 - - - name: Build HTML Docs + + - name: Run Tests run: | - sudo apt install pandoc - pip install sphinx - pip install kentigern - pip install -r requirements.txt - pip install -r requirements_dev.txt - pip install . - cd docs - make html - - - name: SCP Deploy HTML Docs - uses: horochx/deploy-via-scp@v1.0.1 - with: - local: docs/_build/html/* - remote: /home/danwilliams/code.daniel-williams.co.uk/heron/review - host: ${{ secrets.sshhost }} - user: ${{ secrets.sshuser }} - key: ${{ secrets.sshkey }} + python -m unittest discover tests/ From 379bb880a511ad2bdc4247778e823e354fb18ad7 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 17:18:12 +0100 Subject: [PATCH 06/25] Add install requirements to test runner. --- .github/workflows/review.yml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.github/workflows/review.yml b/.github/workflows/review.yml index de2f11a..f7765f0 100644 --- a/.github/workflows/review.yml +++ b/.github/workflows/review.yml @@ -19,6 +19,11 @@ jobs: with: python-version: 3.8 + - name: Install dependencies + run: | + pip install -r requirements.txt + pip install -r requirements_test.txt + - name: Run Tests run: | python -m unittest discover tests/ From 918d53fa005f9e3c14a86151d704332afbaf225c Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 17:21:42 +0100 Subject: [PATCH 07/25] Modify the docs build. --- .github/workflows/documentation.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/documentation.yaml b/.github/workflows/documentation.yaml index 4e0281a..8e1c98d 100644 --- a/.github/workflows/documentation.yaml +++ b/.github/workflows/documentation.yaml @@ -19,19 +19,19 @@ jobs: - name: Build HTML Docs run: | sudo apt install pandoc - pip install sphinx + pip install sphinx sphinx-multiversion pip install kentigern pip install -r requirements.txt pip install -r requirements_dev.txt pip install . cd docs - make multi + make html - name: SCP Deploy HTML Docs uses: horochx/deploy-via-scp@v1.0.1 with: local: docs/_build/html/* - remote: /home/danwilliams/code.daniel-williams.co.uk/heron/ + remote: /home/danwilliams/code.daniel-williams.co.uk/heron/${{ github.ref_name }} host: ${{ secrets.sshhost }} user: ${{ secrets.sshuser }} key: ${{ secrets.sshkey }} From 9a07570b97aae2e5c18ebf6924b220ccc62c997e Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 16 Aug 2024 21:17:47 +0100 Subject: [PATCH 08/25] Fixed waveform alignment in likelihood. --- heron/injection.py | 48 +++++++++++++++++++++++++++++++ heron/likelihood.py | 14 ++++++---- heron/types.py | 62 +++++++++++++++++++++++++++++++++++------ tests/test_inference.py | 47 +++++++++++++++++++++++++++++-- 4 files changed, 155 insertions(+), 16 deletions(-) diff --git a/heron/injection.py b/heron/injection.py index 665d253..5e7556b 100644 --- a/heron/injection.py +++ b/heron/injection.py @@ -63,6 +63,54 @@ def make_injection( return injections +def make_injection_zero_noise( + waveform=IMRPhenomPv2, + injection_parameters={}, + times=None, + detectors=None, + framefile=None, +): + + parameters = {"ra": 0, "dec": 0, "psi": 0, "theta_jn": 0, "phase": 0, 'gpstime': 4000} + parameters.update(injection_parameters) + + waveform = waveform() + + if times is None: + times = np.linspace(-0.5, 0.1, int(0.6 * 4096)) + parameters['gpstime'] + waveform = waveform.time_domain( + parameters, + times=times, + ) + + injections = {} + for detector, psd_model in detectors.items(): + detector = KNOWN_IFOS[detector]() + channel = f"{detector.abbreviation}:Injection" + logger.info(f"Making injection for {detector} in channel {channel}") + psd_model = KNOWN_PSDS[psd_model]() + #data = psd_model.time_series(times) + + # import matplotlib + # matplotlib.use("agg") + # from gwpy.plot import Plot + # f = Plot(data, waveform.project(detector), data+waveform.project(detector), separate=False) + # f.savefig(f"{detector.abbreviation}_injected_waveform.png") + + injection = waveform.project(detector) + injection.channel = channel + injections[detector.abbreviation] = injection + likelihood = TimeDomainLikelihood(injection, psd=psd_model) + snr = likelihood.snr(waveform.project(detector)) + logger.info(f"Optimal Filter SNR: {snr}") + + if framefile: + filename = f"{detector.abbreviation}_{framefile}.gwf" + logger.info(f"Saving framefile to {filename}") + injection.write(filename, format="gwf") + + return injections + def injection_parameters_add_units(parameters): UNITS = {"luminosity_distance": u.megaparsec, "m1": u.solMass, "m2": u.solMass} diff --git a/heron/likelihood.py b/heron/likelihood.py index 32371db..9de6390 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -56,7 +56,7 @@ def __init__( timing_basis=None, ): self.psd = psd - + self.timeseries = data self.data = np.array(data.data) self.times = data.times @@ -86,19 +86,21 @@ def snr(self, waveform): """ dt = (self.times[1] - self.times[0]).value N = len(self.times) + w = np.array(waveform.data) h_h = ( - (np.array(waveform.data).T @ self.solve(self.C, np.array(waveform.data))) + (w.T @ self.solve(self.C, w)) * (dt * dt / N / 4) / 4 ) return np.sqrt(np.abs(h_h)) - def log_likelihood(self, waveform): - residual = np.array(self.data.data) - np.array(waveform.data) + def log_likelihood(self, waveform, norm=True): + a, b = self.timeseries.determine_overlap(self, waveform) + residual = np.array(self.data.data[a[0]:a[1]]) - np.array(waveform.data[b[0]:b[1]]) weighted_residual = ( - (residual) @ self.solve(self.C, residual) * (self.dt * self.dt / 4) / 4 + (residual) @ self.solve(self.C[a[0]:a[1],b[0]:b[1]], residual) * (self.dt * self.dt / 4) / 4 ) - normalisation = self.logdet(2 * np.pi * self.C) + normalisation = self.logdet(2 * np.pi * self.C[a[0]:a[1],b[0]:b[1]]) if norm else 0 return -0.5 * weighted_residual + 0.5 * normalisation def __call__(self, parameters): diff --git a/heron/types.py b/heron/types.py index 14346e6..42765d4 100644 --- a/heron/types.py +++ b/heron/types.py @@ -8,6 +8,7 @@ from lalinference import DetFrameToEquatorial import numpy as array_library +import numpy as np import matplotlib.pyplot as plt @@ -16,8 +17,58 @@ class TimeSeries(TimeSeries): Overload the GWPy timeseries so that additional methods can be defined upon it. """ - pass + def determine_overlap(self, timeseries_a, timeseries_b): + def is_in(time, timeseries): + diff = np.min(np.abs(timeseries - time)) + if diff < (timeseries[1] - timeseries[0]): + return True, diff + else: + return False, diff + + overlap = None + if ( + is_in(timeseries_a.times[-1], timeseries_b.times)[0] + and is_in(timeseries_b.times[0], timeseries_a.times)[0] + ): + overlap = timeseries_b.times[0], timeseries_a.times[-1] + elif ( + is_in(timeseries_a.times[0], timeseries_b.times)[0] + and is_in(timeseries_b.times[-1], timeseries_a.times)[0] + ): + overlap = timeseries_a.times[0], timeseries_b.times[-1] + elif ( + is_in(timeseries_b.times[0], timeseries_a.times)[0] + and is_in(timeseries_b.times[-1], timeseries_a.times)[0] + and not is_in(timeseries_a.times[-1], timeseries_b.times)[0] + ): + overlap = timeseries_b.times[0], timeseries_b.times[-1] + elif ( + is_in(timeseries_a.times[0], timeseries_b.times)[0] + and is_in(timeseries_a.times[-1], timeseries_b.times)[0] + and not is_in(timeseries_b.times[-1], timeseries_a.times)[0] + ): + overlap = timeseries_a.times[0], timeseries_a.times[-1] + else: + overlap = None + return None + start_a = np.argmin(np.abs(timeseries_a.times - overlap[0])) + finish_a = np.argmin(np.abs(timeseries_a.times - overlap[-1])) + + start_b = np.argmin(np.abs(timeseries_b.times - overlap[0])) + finish_b = np.argmin(np.abs(timeseries_b.times - overlap[-1])) + return (start_a, finish_a), (start_b, finish_b) + + def align(self, waveform_b): + """ + Align this waveform with another one by altering the phase. + """ + + indices = self.determine_overlap(self, waveform_b) + + return self[indices[0][0]:indices[0][1]], waveform_b[indices[1][0]: indices[1][1]] + + class PSD(FrequencySeries): def __init__(self, data, frequencies, *args, **kwargs): @@ -40,7 +91,7 @@ def __init__(self, variance=None, covariance=None, *args, **kwargs): def __new__(self, variance=None, covariance=None, *args, **kwargs): # if "covariance" in kwargs: # self.covariance = kwargs.pop("covariance") - waveform = super(Waveform, self).__new__(TimeSeriesBase, *args, **kwargs) + waveform = super(Waveform, self).__new__(TimeSeries, *args, **kwargs) waveform.covariance = covariance waveform.variance = variance @@ -50,11 +101,6 @@ def __new__(self, variance=None, covariance=None, *args, **kwargs): # def dt(self): # return self.waveform.times[1] - self.waveform.times[0] - def align(self, waveform_b): - """ - Align this waveform with another one by altering the phase. - """ - pass class WaveformDict: @@ -194,7 +240,7 @@ def project( covariance=projected_covariance, times=self.waveforms["plus"].times, ) - + projected_waveform.shift(dt) return projected_waveform diff --git a/tests/test_inference.py b/tests/test_inference.py index e56416f..130e7f0 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -17,7 +17,7 @@ from heron.models.lalsimulation import SEOBNRv3, IMRPhenomPv2, IMRPhenomPv2_FakeUncertainty from heron.models.lalnoise import AdvancedLIGO -from heron.injection import make_injection +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 @@ -29,6 +29,49 @@ CUDA_NOT_AVAILABLE = not is_available() +class Test_Likelihood_ZeroNoise(unittest.TestCase): + """ + Test likelihoods on a zero noise injection. + """ + + def setUp(self): + self.waveform = IMRPhenomPv2() + 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_no_norm(self): + data = self.injections['H1'] + + from gwpy.plot import Plot + + likelihood = TimeDomainLikelihood(data, psd=self.psd_model) + + test_waveform = self.waveform.time_domain(parameters={"distance": 1000*u.megaparsec, + "mass_ratio": 0.6, + "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) + + f = Plot(data, projected_waveform) + f.savefig("projected_waveform.png") + + log_like = likelihood.log_likelihood(projected_waveform, norm=False) + + self.assertTrue(log_like <= 1e-5) + + class Test_Filter(unittest.TestCase): """Test that filters can be applied correctly to data.""" @@ -44,6 +87,7 @@ def setUp(self): "AdvancedLIGOLivingston": "AdvancedLIGO"} ) + def test_timedomain_psd(self): noise = self.psd_model.time_domain(times=self.injections['H1'].times) #print(noise) @@ -66,7 +110,6 @@ def test_snr(self): f.savefig("projected_waveform.png") snr = likelihood.snr(projected_waveform) - print("snr", snr) self.assertTrue(snr > 40 and snr < 45) # def test_snr_f(self): From 7f20b36133505a774f4b416d615f154eda8bdb56 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Mon, 19 Aug 2024 11:43:34 +0100 Subject: [PATCH 09/25] Fix location of logabsdet. --- heron/likelihood.py | 3 ++- tests/test_inference.py | 27 ++++++++++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/heron/likelihood.py b/heron/likelihood.py index 9de6390..b540f03 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -25,7 +25,8 @@ class LikelihoodBase: class Likelihood(LikelihoodBase): def logdet(self, K): - return np.linalg.slogdet(K).logabsdet + (sign, logabsdet) = np.linalg.slogdet(K) + return logabsdet def inverse(self, A): return np.linalg.inv(A) diff --git a/tests/test_inference.py b/tests/test_inference.py index 130e7f0..0532545 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -70,7 +70,32 @@ def test_likelihood_no_norm(self): log_like = likelihood.log_likelihood(projected_waveform, norm=False) self.assertTrue(log_like <= 1e-5) - + + + def test_likelihood_no_norm(self): + data = self.injections['H1'] + + from gwpy.plot import Plot + + likelihood = TimeDomainLikelihood(data, psd=self.psd_model) + + test_waveform = self.waveform.time_domain(parameters={"distance": 1000*u.megaparsec, + "mass_ratio": 0.6, + "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) + + f = Plot(data, projected_waveform) + f.savefig("projected_waveform.png") + + log_like = likelihood.log_likelihood(projected_waveform) + + self.assertTrue(log_like <= 1e-5) + class Test_Filter(unittest.TestCase): """Test that filters can be applied correctly to data.""" From 97590b7d88b963a012db59390128e15b55cd0161 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Mon, 19 Aug 2024 14:56:45 +0100 Subject: [PATCH 10/25] Updated tests to ensure compatibility between torch and numpy --- heron/injection.py | 2 +- heron/likelihood.py | 24 +++-- profiling/{ => cprofile}/likelihood.py | 0 profiling/{ => cuda}/cuda_likelihood.py | 0 tests/test_inference.py | 116 +++++++++++++++++++++--- 5 files changed, 124 insertions(+), 18 deletions(-) rename profiling/{ => cprofile}/likelihood.py (100%) rename profiling/{ => cuda}/cuda_likelihood.py (100%) diff --git a/heron/injection.py b/heron/injection.py index 5e7556b..c6f2e6d 100644 --- a/heron/injection.py +++ b/heron/injection.py @@ -109,7 +109,7 @@ def make_injection_zero_noise( logger.info(f"Saving framefile to {filename}") injection.write(filename, format="gwf") - return injections + return injections def injection_parameters_add_units(parameters): UNITS = {"luminosity_distance": u.megaparsec, "m1": u.solMass, "m2": u.solMass} diff --git a/heron/likelihood.py b/heron/likelihood.py index b540f03..e7e6fee 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -102,7 +102,7 @@ def log_likelihood(self, waveform, norm=True): (residual) @ self.solve(self.C[a[0]:a[1],b[0]:b[1]], residual) * (self.dt * self.dt / 4) / 4 ) normalisation = self.logdet(2 * np.pi * self.C[a[0]:a[1],b[0]:b[1]]) if norm else 0 - return -0.5 * weighted_residual + 0.5 * normalisation + return 0.5 * weighted_residual + 0.5 * normalisation def __call__(self, parameters): self.logger.info(parameters) @@ -233,6 +233,8 @@ def __init__( self.logger.info(f"Using device {device}") self.psd = psd + self.timeseries = data + self.data = torch.tensor(data.data, device=self.device, dtype=torch.double) self.times = data.times @@ -263,14 +265,24 @@ def snr(self, waveform): h_h = (waveform_d.T @ self.solve(self.C, waveform_d)) * (dt * dt / N / 4) / 4 return torch.sqrt(torch.abs(h_h)) - def log_likelihood(self, waveform): + def log_likelihood(self, waveform, norm=True): + a, b = self.timeseries.determine_overlap(self, waveform) + residual = np.array(self.data.data[a[0]:a[1]]) - np.array(waveform.data[b[0]:b[1]]) + weighted_residual = ( + (residual) @ self.solve(self.C[a[0]:a[1],b[0]:b[1]], residual) * (self.dt * self.dt / 4) / 4 + ) + normalisation = self.logdet(2 * np.pi * self.C[a[0]:a[1],b[0]:b[1]]) if norm else 0 + return 0.5 * weighted_residual + 0.5 * normalisation + + def log_likelihood(self, waveform, norm=True): + a, b = self.timeseries.determine_overlap(self, waveform) waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) - residual = self.data - waveform_d + residual = self.data[a[0]:a[1]] - waveform_d[b[0]:b[1]] weighted_residual = ( - (residual) @ self.solve(self.C, residual) * (self.dt * self.dt / 4) / 4 + (residual) @ self.solve(self.C[a[0]:a[1],b[0]:b[1]], residual) * (self.dt * self.dt / 4) / 4 ) - normalisation = self.logdet(2 * np.pi * self.C) - return -0.5 * weighted_residual + 0.5 * normalisation + normalisation = self.logdet(2 * np.pi * self.C[a[0]:a[1],b[0]:b[1]]) if norm else 0 + return 0.5 * weighted_residual + 0.5 * normalisation def __call__(self, parameters): self.logger.info(parameters) diff --git a/profiling/likelihood.py b/profiling/cprofile/likelihood.py similarity index 100% rename from profiling/likelihood.py rename to profiling/cprofile/likelihood.py diff --git a/profiling/cuda_likelihood.py b/profiling/cuda/cuda_likelihood.py similarity index 100% rename from profiling/cuda_likelihood.py rename to profiling/cuda/cuda_likelihood.py diff --git a/tests/test_inference.py b/tests/test_inference.py index 0532545..e968bdd 100644 --- a/tests/test_inference.py +++ b/tests/test_inference.py @@ -19,8 +19,8 @@ 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.likelihood import MultiDetector, TimeDomainLikelihood, TimeDomainLikelihoodModelUncertainty, TimeDomainLikelihoodPyTorch +#, TimeDomainLikelihoodModelUncertaintyPyTorch from heron.inference import heron_inference, parse_dict, load_yaml @@ -50,7 +50,7 @@ def setUp(self): def test_likelihood_no_norm(self): data = self.injections['H1'] - from gwpy.plot import Plot + # from gwpy.plot import Plot likelihood = TimeDomainLikelihood(data, psd=self.psd_model) @@ -64,20 +64,63 @@ def test_likelihood_no_norm(self): phi_0=0, psi=0, iota=0) - f = Plot(data, projected_waveform) - f.savefig("projected_waveform.png") + # f = Plot(data, projected_waveform) + # f.savefig("projected_waveform.png") log_like = likelihood.log_likelihood(projected_waveform, norm=False) self.assertTrue(log_like <= 1e-5) + def test_likelihood_maximum_at_true_value_mass_ratio(self): + + data = self.injections['H1'] + + likelihood = TimeDomainLikelihood(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)) + + self.assertTrue(mass_ratios[np.argmax(log_likes)] == 0.6) + + +class Test_PyTorch_Likelihood_ZeroNoise(unittest.TestCase): + """ + Test likelihoods on a zero noise injection. + """ + + def setUp(self): + self.waveform = IMRPhenomPv2() + 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_no_norm(self): data = self.injections['H1'] - from gwpy.plot import Plot + # from gwpy.plot import Plot - likelihood = TimeDomainLikelihood(data, psd=self.psd_model) + likelihood = TimeDomainLikelihoodPyTorch(data, psd=self.psd_model) test_waveform = self.waveform.time_domain(parameters={"distance": 1000*u.megaparsec, "mass_ratio": 0.6, @@ -89,12 +132,63 @@ def test_likelihood_no_norm(self): phi_0=0, psi=0, iota=0) - f = Plot(data, projected_waveform) - f.savefig("projected_waveform.png") + log_like = likelihood.log_likelihood(projected_waveform, norm=False) - log_like = likelihood.log_likelihood(projected_waveform) + self.assertTrue(log_like.cpu().numpy() <= 1e-5) - self.assertTrue(log_like <= 1e-5) + + def test_likelihood_maximum_at_true_value_mass_ratio(self): + + data = self.injections['H1'] + + likelihood = TimeDomainLikelihoodPyTorch(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).cpu().numpy()) + + self.assertTrue(mass_ratios[np.argmax(log_likes)] == 0.6) + + + def test_likelihood_numpy_equivalent(self): + + data = self.injections['H1'] + + likelihood = TimeDomainLikelihoodPyTorch(data, psd=self.psd_model) + numpy_likelihood = TimeDomainLikelihood(data, psd=self.psd_model) + mass_ratios = np.linspace(0.1, 1.0, 100) + + log_likes = [] + log_likes_n = [] + 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).cpu().numpy()) + log_likes_n.append(numpy_likelihood.log_likelihood(projected_waveform)) + + 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): From f729525f8064bf09f38df1aaef7ecfb8aab0f616 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 20 Aug 2024 11:06:38 +0100 Subject: [PATCH 11/25] 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 6d46060213efcedef519fbd83f493bf8f01feaa3 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 20 Aug 2024 11:06:38 +0100 Subject: [PATCH 12/25] Updated the profiling tests. --- profiling/cuda_likelihood.py | 29 ++++++++++++++--------------- 1 file changed, 14 insertions(+), 15 deletions(-) diff --git a/profiling/cuda_likelihood.py b/profiling/cuda_likelihood.py index d2194f3..20b460e 100644 --- a/profiling/cuda_likelihood.py +++ b/profiling/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 967ac4e1d53e4d8f88ba044e17e8e3c21d543e26 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 20 Aug 2024 11:19:19 +0100 Subject: [PATCH 13/25] Add cuda likelihood test. --- profiling/cuda/cuda_likelihood.py | 50 +++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) create mode 100644 profiling/cuda/cuda_likelihood.py diff --git a/profiling/cuda/cuda_likelihood.py b/profiling/cuda/cuda_likelihood.py new file mode 100644 index 0000000..20b460e --- /dev/null +++ b/profiling/cuda/cuda_likelihood.py @@ -0,0 +1,50 @@ +import numpy as np +import astropy.units as u +from heron.models.lalnoise import AdvancedLIGO +from heron.injection import make_injection_zero_noise +from heron.detector import AdvancedLIGOHanford +from heron.likelihood import TimeDomainLikelihoodPyTorch +from heron.models.lalsimulation import IMRPhenomPv2, IMRPhenomPv2_FakeUncertainty + + +def profile_likelihood_pytorch_nouncert(): + waveform = IMRPhenomPv2() + psd_model = AdvancedLIGO() + + injections = make_injection_zero_noise(waveform=IMRPhenomPv2, + injection_parameters={"m1": 35*u.solMass, + "m2": 30*u.solMass, + "gpstime": 4000, + "distance": 410 * u.megaparsec}, + detectors={"AdvancedLIGOHanford": "AdvancedLIGO", + "AdvancedLIGOLivingston": "AdvancedLIGO"} + ) + + data = injections['H1'] + + likelihood = TimeDomainLikelihoodPyTorch(data, psd=psd_model) + print(likelihood.device) + + + 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) + + log_like = likelihood.log_likelihood(projected_waveform) + + +from torch.profiler import profile, record_function, ProfilerActivity + +with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: + with record_function("model_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 56eaba09e7ee56da9fb408f0b6505177dd4c6bc4 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 20 Aug 2024 14:06:06 +0100 Subject: [PATCH 14/25] Moved cuda likelihood. --- profiling/cuda_likelihood.py | 50 ------------------------------------ 1 file changed, 50 deletions(-) delete mode 100644 profiling/cuda_likelihood.py diff --git a/profiling/cuda_likelihood.py b/profiling/cuda_likelihood.py deleted file mode 100644 index 20b460e..0000000 --- a/profiling/cuda_likelihood.py +++ /dev/null @@ -1,50 +0,0 @@ -import numpy as np -import astropy.units as u -from heron.models.lalnoise import AdvancedLIGO -from heron.injection import make_injection_zero_noise -from heron.detector import AdvancedLIGOHanford -from heron.likelihood import TimeDomainLikelihoodPyTorch -from heron.models.lalsimulation import IMRPhenomPv2, IMRPhenomPv2_FakeUncertainty - - -def profile_likelihood_pytorch_nouncert(): - waveform = IMRPhenomPv2() - psd_model = AdvancedLIGO() - - injections = make_injection_zero_noise(waveform=IMRPhenomPv2, - injection_parameters={"m1": 35*u.solMass, - "m2": 30*u.solMass, - "gpstime": 4000, - "distance": 410 * u.megaparsec}, - detectors={"AdvancedLIGOHanford": "AdvancedLIGO", - "AdvancedLIGOLivingston": "AdvancedLIGO"} - ) - - data = injections['H1'] - - likelihood = TimeDomainLikelihoodPyTorch(data, psd=psd_model) - print(likelihood.device) - - - 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) - - log_like = likelihood.log_likelihood(projected_waveform) - - -from torch.profiler import profile, record_function, ProfilerActivity - -with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof: - with record_function("model_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 15/25] 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 16/25] 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 17/25] 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) From e0da2a47485646ecbaf92d7d559f774317f5f826 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 23 Aug 2024 12:29:37 +0100 Subject: [PATCH 18/25] Added first tests of the CUDA likelihood. --- heron/likelihood.py | 33 ++++++++++-------- tests/test_inference_uncertain.py | 57 +++++++++++++++++++++++++++++-- 2 files changed, 72 insertions(+), 18 deletions(-) diff --git a/heron/likelihood.py b/heron/likelihood.py index 2a05f61..2b304d8 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -322,52 +322,55 @@ class TimeDomainLikelihoodModelUncertaintyPyTorch(TimeDomainLikelihoodPyTorch): def __init__(self, data, psd, waveform=None, detector=None): super().__init__(data, psd, waveform, detector) - def _normalisation(self, K, S): + def _normalisation(self, K, S, indices): + (ind_a, ind_b) = indices norm = ( -1.5 * K.shape[0] * self.log(2 * self.pi) - 0.5 * self.logdet(K) - + 0.5 * self.logdet(self.C) - - 0.5 * self.logdet(self.solve(K, self.C) + self.eye(K.shape[0])) + + 0.5 * self.logdet(self.C[ind_a[0]:ind_a[1],ind_a[0]:ind_a[1]]) + - 0.5 * self.logdet(self.solve(K, self.C[ind_a[0]:ind_a[1],ind_a[0]:ind_a[1]]) + self.eye(K.shape[0])) ) return norm - def _weighted_data(self): + def _weighted_data(self, indices): """Return the weighted data component""" # TODO This can all be pre-computed + (ind_a, ind_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) ) else: dw = self.weighted_data_CACHE - return dw + return dw#[ind_a[0]:ind_a[1]] def _weighted_model(self, mu, K): """Return the inner product of the GPR mean""" mu = torch.tensor(mu, device=self.device, dtype=torch.double) return -0.5 * 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) + def _weighted_cross(self, mu, K, indices): + (ind_a, ind_b) = indices + a = self.solve(self.C[ind_a[0]:ind_a[1],ind_a[0]:ind_a[1]], self.data[ind_a[0]:ind_a[1]]) + self.solve(K, mu) + b = self.inverse_C[ind_a[0]:ind_a[1],ind_a[0]:ind_a[1]] + self.inverse(K) return 0.5 * a.T @ self.solve(b, a) - def log_likelihood(self, waveform): + def log_likelihood(self, waveform, norm=True): a, b = self.timeseries.determine_overlap(self, waveform) - waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double) + waveform_d = torch.tensor(waveform.data, device=self.device, dtype=torch.double)[b[0]:b[1]] waveform_c = torch.tensor( waveform.covariance, device=self.device, dtype=torch.double - ) - like = self._weighted_cross(waveform_d, waveform_c) + )[b[0]:b[1], b[0]:b[1]] + like = self._weighted_cross(waveform_d, waveform_c, indices=(a,b)) # print("cross term", like) - A = self._weighted_data() + A = self._weighted_data(indices=(a,b)) # print("data term", A) B = self._weighted_model(waveform_d, waveform_c) # print("model term", B) like = like + A + B - norm = self._normalisation(waveform_c, self.C) + normalisation = self._normalisation(waveform_c, self.C, indices=(a,b)) # print("normalisation", norm) - like += norm + like += normalisation if norm else 0 return like diff --git a/tests/test_inference_uncertain.py b/tests/test_inference_uncertain.py index cc95fbd..ed0c482 100644 --- a/tests/test_inference_uncertain.py +++ b/tests/test_inference_uncertain.py @@ -7,9 +7,15 @@ 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.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 @@ -61,3 +67,48 @@ def test_likelihood_maximum_at_true_value_mass_ratio(self): log_likes.append(likelihood.log_likelihood(projected_waveform, norm=False)) self.assertTrue(np.abs(mass_ratios[np.argmax(log_likes)] - 0.6) < 0.1) + + +class Test_Likelihood_ZeroNoise_With_Uncertainty_PyTorch(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 = TimeDomainLikelihoodModelUncertaintyPyTorch(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).cpu().numpy()) + + self.assertTrue(np.abs(mass_ratios[np.argmax(log_likes)] - 0.6) < 0.1) From 58ac10de4b5dd1a26f09e8fae9ab7b7831fa6c3a Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 23 Aug 2024 12:35:49 +0100 Subject: [PATCH 19/25] Update the tests so the cuda-requiring ones are skipped in CI --- tests/test_inference_uncertain.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_inference_uncertain.py b/tests/test_inference_uncertain.py index ed0c482..eff2b64 100644 --- a/tests/test_inference_uncertain.py +++ b/tests/test_inference_uncertain.py @@ -68,7 +68,7 @@ def test_likelihood_maximum_at_true_value_mass_ratio(self): self.assertTrue(np.abs(mass_ratios[np.argmax(log_likes)] - 0.6) < 0.1) - +@unittest.skipIf(CUDA_NOT_AVAILABLE, "CUDA is not installed on this system") class Test_Likelihood_ZeroNoise_With_Uncertainty_PyTorch(unittest.TestCase): """ Test likelihoods on a zero noise injection. From 55b32379df0dd94da9e8fdd6b7ba13b36692906c Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Mon, 26 Aug 2024 12:52:29 +0100 Subject: [PATCH 20/25] Update handling of mass ratio --- heron/likelihood.py | 2 +- heron/models/__init__.py | 22 ++++++++++++++-------- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/heron/likelihood.py b/heron/likelihood.py index 2b304d8..3c8fd05 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -108,7 +108,7 @@ def __call__(self, parameters): self.logger.info(parameters) keys = set(parameters.keys()) - extrinsic = {"phase", "psi", "ra", "dec", "theta_jn"} + extrinsic = {"phase", "psi", "ra", "dec", "theta_jn", "gpstime", "geocent_time"} conversions = {"mass_ratio", "total_mass", "luminosity_distance"} bad_keys = keys - set(self.waveform._args.keys()) - extrinsic - conversions if len(bad_keys) > 0: diff --git a/heron/models/__init__.py b/heron/models/__init__.py index f3ba0fd..c311280 100644 --- a/heron/models/__init__.py +++ b/heron/models/__init__.py @@ -26,16 +26,22 @@ def _convert_luminosity_distance(self, args): return args def _convert_mass_ratio_total_mass(self, args): - - args["m1"] = (args["total_mass"] / (1 + args["mass_ratio"])).to_value( - u.kilogram - ) - args["m2"] = (args["total_mass"] / (1 + 1 / args["mass_ratio"])).to_value( - u.kilogram - ) + args["m1"] = (args["total_mass"] / (1 + args["mass_ratio"])) + args["m2"] = (args["total_mass"] / (1 + (1 / args["mass_ratio"]))) + # Do these have units? + # If not then we can skip some relatively expensive operations and apply a heuristic. + if isinstance(args["m1"], u.Quantity): + args["m1"] = args["m1"].to_value(u.kilogram) + args["m2"] = args["m2"].to_value(u.kilogram) + if (not isinstance(args["m1"], u.Quantity)) and (args["m1"] < 1000): + # This appears to be in solar masses + args["m1"] *= MSUN_SI + if (not isinstance(args["m2"], u.Quantity)) and (args["m2"] < 1000): + # This appears to be in solar masses + args["m2"] *= MSUN_SI + args.pop("total_mass") args.pop("mass_ratio") - return args From 1efcd4c29ec7bcef1bd751a5cb42bb780e93fdce Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Mon, 26 Aug 2024 13:59:44 +0100 Subject: [PATCH 21/25] Add MSUN_SI --- heron/models/__init__.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/heron/models/__init__.py b/heron/models/__init__.py index c311280..171c0f6 100644 --- a/heron/models/__init__.py +++ b/heron/models/__init__.py @@ -1,5 +1,6 @@ import torch -from lal import antenna +import numpy as np +from lal import antenna, MSUN_SI from astropy import units as u From 07889ea54c4c0499a9d6b4bb7dcc68b2dbeb696d Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 3 Sep 2024 10:56:46 +0100 Subject: [PATCH 22/25] Added some bug fixes. --- heron/inference.py | 11 +++++++---- heron/injection.py | 2 +- heron/likelihood.py | 7 +++++-- heron/models/lalsimulation.py | 1 - heron/priors.py | 1 + heron/types.py | 3 +++ requirements.txt | 6 +++--- 7 files changed, 20 insertions(+), 11 deletions(-) diff --git a/heron/inference.py b/heron/inference.py index f68d95d..f680c16 100644 --- a/heron/inference.py +++ b/heron/inference.py @@ -6,7 +6,7 @@ import click -from gwpy.timeseries import TimeSeries +from .types import TimeSeries import astropy.units as u from nessai.flowsampler import FlowSampler @@ -84,6 +84,7 @@ def heron_inference(settings): if "data files" in settings.get("data", {}): # Load frame files from disk for ifo in settings["interferometers"]: + print(f"Loading {ifo} data") logger.info( f"Loading {ifo} data from " f"{settings['data']['data files'][ifo]}/{settings['data']['channels'][ifo]}" @@ -93,14 +94,16 @@ def heron_inference(settings): channel=settings["data"]["channels"][ifo], format="gwf", ) - elif "injection" in other_settings: - pass + #elif "injection" in other_settings: + # pass # Make Likelihood if len(settings["interferometers"]) > 1: likelihoods = [] + print("Creating likelihoods") waveform_model = KNOWN_WAVEFORMS[settings["waveform"]["model"]]() for ifo in settings["interferometers"]: + print(f"\t {ifo}") likelihoods.append( KNOWN_LIKELIHOODS[settings.get("likelihood").get("function")]( data[ifo], @@ -113,7 +116,7 @@ def heron_inference(settings): ), ) ) - likelihood = MultiDetector(*likelihoods) + likelihood = MultiDetector(*likelihoods) priors = heron.priors.PriorDict() priors.from_dictionary(settings["priors"]) diff --git a/heron/injection.py b/heron/injection.py index c6f2e6d..a0e357f 100644 --- a/heron/injection.py +++ b/heron/injection.py @@ -115,7 +115,7 @@ def injection_parameters_add_units(parameters): UNITS = {"luminosity_distance": u.megaparsec, "m1": u.solMass, "m2": u.solMass} for parameter, value in parameters.items(): - if not isinstance(value, u.Quantity): + if not isinstance(value, u.Quantity) and parameter in UNITS: parameters[parameter] = value * UNITS[parameter] return parameters diff --git a/heron/likelihood.py b/heron/likelihood.py index 3c8fd05..968d504 100644 --- a/heron/likelihood.py +++ b/heron/likelihood.py @@ -60,7 +60,6 @@ def __init__( self.timeseries = data self.data = np.array(data.data) self.times = data.times - self.C = self.psd.covariance_matrix(times=self.times) self.inverse_C = np.linalg.inv(self.C) @@ -96,7 +95,11 @@ def snr(self, waveform): return np.sqrt(np.abs(h_h)) def log_likelihood(self, waveform, norm=True): - a, b = self.timeseries.determine_overlap(self, waveform) + w = self.timeseries.determine_overlap(self, waveform) + if w is not None: + (a,b) = w + else: + return -np.inf residual = np.array(self.data.data[a[0]:a[1]]) - np.array(waveform.data[b[0]:b[1]]) weighted_residual = ( (residual) @ self.solve(self.C[a[0]:a[1],b[0]:b[1]], residual) * (self.dt * self.dt / 4) / 4 diff --git a/heron/models/lalsimulation.py b/heron/models/lalsimulation.py index 559b90d..df963ec 100644 --- a/heron/models/lalsimulation.py +++ b/heron/models/lalsimulation.py @@ -91,7 +91,6 @@ def _convert_units(self, args): args[name] = argument.to_value(units[mappings[name]]) elif name in mappings.keys() and argument: # This is commented out as it causes problems if e.g. lalnative values are passed - print(f"Performing a mapping on {name}, {argument}") args[name] = (argument * default_units[mappings[name]]).to_value( units[mappings[name]] ) diff --git a/heron/priors.py b/heron/priors.py index bca4049..d5b6265 100644 --- a/heron/priors.py +++ b/heron/priors.py @@ -13,6 +13,7 @@ "Uniform": bilby.prior.Uniform, "PowerLaw": bilby.prior.PowerLaw, "Sine": bilby.prior.Sine, + "Cosine": bilby.prior.Cosine, "UniformSourceFrame": bilby.gw.prior.UniformSourceFrame, "UniformInComponentsMassRatio": bilby.gw.prior.UniformInComponentsMassRatio, } diff --git a/heron/types.py b/heron/types.py index 42765d4..0a36cba 100644 --- a/heron/types.py +++ b/heron/types.py @@ -50,6 +50,9 @@ def is_in(time, timeseries): overlap = timeseries_a.times[0], timeseries_a.times[-1] else: overlap = None + #print("No overlap found") + #print(timeseries_a.times[0], timeseries_a.times[-1]) + #print(timeseries_b.times[0], timeseries_b.times[-1]) return None start_a = np.argmin(np.abs(timeseries_a.times - overlap[0])) diff --git a/requirements.txt b/requirements.txt index 4ba302f..1f1671b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,6 +16,6 @@ click asimov pesummary nessai -gpytorch==1.0.1 -torch==2.4.0 -torchvision==0.5.0 +gpytorch +torch +torchvision From f7df371b52425d5264350729cdc355e07dda87d6 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Tue, 1 Oct 2024 10:33:48 +0100 Subject: [PATCH 23/25] Move to internal timeseries --- heron/inference.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heron/inference.py b/heron/inference.py index f68d95d..6a7eadc 100644 --- a/heron/inference.py +++ b/heron/inference.py @@ -6,7 +6,7 @@ import click -from gwpy.timeseries import TimeSeries +from .types import TimeSeries import astropy.units as u from nessai.flowsampler import FlowSampler From 377ecbb54a9a95587014dac24ee2aac167ebcaf2 Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 4 Oct 2024 10:44:08 +0100 Subject: [PATCH 24/25] Minor fix for GPStime in injection --- heron/injection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heron/injection.py b/heron/injection.py index a0e357f..e0d0222 100644 --- a/heron/injection.py +++ b/heron/injection.py @@ -115,7 +115,7 @@ def injection_parameters_add_units(parameters): UNITS = {"luminosity_distance": u.megaparsec, "m1": u.solMass, "m2": u.solMass} for parameter, value in parameters.items(): - if not isinstance(value, u.Quantity) and parameter in UNITS: + if not isinstance(value, u.Quantity) and (parameter in UNITS.keys()): parameters[parameter] = value * UNITS[parameter] return parameters From 7112445ad6a6c45769e88c08c0181d5fabea8a0b Mon Sep 17 00:00:00 2001 From: Daniel Williams Date: Fri, 13 Dec 2024 14:35:04 +0000 Subject: [PATCH 25/25] Changes to allow bilby to read frame files. --- heron/injection.py | 15 +++++++++++---- heron/models/lalnoise.py | 1 + heron/models/lalsimulation.py | 12 ++++++------ 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/heron/injection.py b/heron/injection.py index e0d0222..4cd041a 100644 --- a/heron/injection.py +++ b/heron/injection.py @@ -22,6 +22,8 @@ def make_injection( waveform=IMRPhenomPv2, injection_parameters={}, + duration=32, + sample_rate=4096, times=None, detectors=None, framefile=None, @@ -33,7 +35,7 @@ def make_injection( waveform = waveform() if times is None: - times = np.linspace(-0.5, 0.1, int(0.6 * 4096)) + times = np.linspace(parameters['gpstime']-duration+2, parameters['gpstime']+2, int(duration * sample_rate)) waveform = waveform.time_domain( parameters, times=times, @@ -44,16 +46,19 @@ def make_injection( logger.info(f"Making injection for {detector}") psd_model = KNOWN_PSDS[psd_model]() detector = KNOWN_IFOS[detector]() + if times is None: + times = waveform['plus'].times.value data = psd_model.time_series(times) + print(data) channel = f"{detector.abbreviation}:Injection" injection = data + waveform.project(detector) injection.channel = channel injections[detector.abbreviation] = injection - likelihood = TimeDomainLikelihood(injection, psd=psd_model) - snr = likelihood.snr(waveform.project(detector)) + # likelihood = TimeDomainLikelihood(injection, psd=psd_model) + # snr = likelihood.snr(waveform.project(detector)) - logger.info(f"Optimal Filter SNR: {snr}") + #logger.info(f"Optimal Filter SNR: {snr}") if framefile: filename = f"{detector.abbreviation}_{framefile}.gwf" @@ -146,6 +151,8 @@ def injection(settings): } injections = make_injection( waveform=IMRPhenomPv2, + duration=settings["duration"], + sample_rate=settings["sample rate"], injection_parameters=parameters, detectors=detector_dict, framefile="injection", diff --git a/heron/models/lalnoise.py b/heron/models/lalnoise.py index ddd3783..f73aa81 100644 --- a/heron/models/lalnoise.py +++ b/heron/models/lalnoise.py @@ -78,6 +78,7 @@ def time_series(self, times): dt = times[1] - times[0] N = len(times) + print(N) T = times[-1] - times[0] df = 1 / T frequencies = torch.arange(len(times) // 2 + 1) * df diff --git a/heron/models/lalsimulation.py b/heron/models/lalsimulation.py index df963ec..4d644bb 100644 --- a/heron/models/lalsimulation.py +++ b/heron/models/lalsimulation.py @@ -121,9 +121,9 @@ def time_domain(self, parameters, times=None): """ Retrieve a time domain waveform for a given set of parameters. """ + epoch = parameters.get("gpstime", parameters.get("epoch", 0)) self._args.update(parameters) - epoch = parameters.get("gpstime", 0) - + print("epoch is ", epoch) if not (self._args == self._cache_key): self.logger.info(f"Generating new waveform at {self.args}") self._cache_key = self.args.copy() @@ -162,17 +162,17 @@ def time_domain(self, parameters, times=None): spl_hx = CubicSpline(times_wf, hx.data.data) hp_data = spl_hp(times) hx_data = spl_hx(times) - hp_ts = Waveform(data=hp_data, times=times) - hx_ts = Waveform(data=hx_data, times=times) + hp_ts = Waveform(data=hp_data, times=times + epoch) + hx_ts = Waveform(data=hx_data, times=times + epoch) parameters.pop("time") else: hp_data = hp.data.data hx_data = hx.data.data hp_ts = Waveform(data=hp_data, dt=hp.deltaT, t0=hp.epoch + epoch) hx_ts = Waveform(data=hx_data, dt=hx.deltaT, t0=hx.epoch + epoch) - + self._cache = WaveformDict(parameters=parameters, plus=hp_ts, cross=hx_ts) - + print("written epoch is ", hp_ts.times[0]) return self._cache