Commit 50326134 authored by Daniel Williams's avatar Daniel Williams 🤖
Browse files

Added frequency mixins.

parent d3603f17
Pipeline #254429 failed with stages
in 16 seconds
"""
This module contains objects which provide the specifically-GW parts of waveform surrogate models.
"""
from elk.waveform import Waveform, Timeseries
from elk.waveform import Waveform, Timeseries, FrequencySeries
import astropy.constants as c
import numpy as np
from math import log
import torch
from ..utils import diag_cuda
class HofTSurrogate(object):
def _to_frequency(self, timeseries, *args):
"""Convert an elk timeseries output from heron into an elk frequency series output."""
return timeseries.to_fseries(*args)
def time_domain_waveform(self, p, times=np.linspace(-2, 2, 1000)):
"""
Return the timedomain waveform.
"""
return self.mean(times, p)
def bilby(self, time, mass_1, mass_2, luminosity_distance):
"""
Return a waveform from the GPR in a format expected by the Bilby ecosystem
......@@ -65,3 +74,30 @@ class BBHNonSpinSurrogate(object):
}
parameters = ("mass ratio",)
c_ind = {j:i for i,j in columns.items()}
class FrequencyMixin:
def frequency_domain_waveform(self, p, window, times=np.linspace(-2, 2, 1000)):
"""
Return the frequency domain waveform.
"""
data = {}
for polarisation in self.polarisations:
mean, _, cov = self._predict(times, p, polarisation=polarisation)
strain_f = torch.view_as_complex((window*mean.double()).rfft(1))
cov_f = torch.view_as_complex(cov.rfft(2))
dt = times[-1] - times[-2]
uncert_f = diag_cuda(cov_f)
#Complex(torch.stack([torch.diag(cov_f[:, :, 0]), torch.diag(cov_f[:, :, 1])]).T)
if not isinstance(times, type(None)):
srate = 1/np.diff(times).mean()
nf = int(np.floor(len(times)/2))+1
frequencies = np.linspace(0, srate, nf)
data[polarisation] = FrequencySeries(data=strain_f,
frequencies=frequencies,
covariance=cov_f,
variance=uncert_f)
return data
......@@ -9,6 +9,8 @@ import pkg_resources
import numpy as np
import os
import torch
import gpytorch
from gpytorch.kernels import RBFKernel
......@@ -24,7 +26,7 @@ from elk.catalogue import PPCatalogue
from . import Model
from ..utils import diag_cuda
from .gw import BBHSurrogate, HofTSurrogate, BBHNonSpinSurrogate
from .gw import BBHSurrogate, HofTSurrogate, BBHNonSpinSurrogate, FrequencyMixin
from heron.models import Model
import matplotlib.pyplot as plt
......@@ -63,7 +65,7 @@ def train(model, iterations=1000):
# Our loss object. We're using the VariationalELBO
mll = gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model.model_plus)
epochs_iter = tqdm.tqdm_notebook(range(training_iterations), desc="Epoch")
epochs_iter = tqdm.tqdm(range(training_iterations), desc="Epoch")
if hasattr(model, "state_vector"):
state_vector = model.state_vector
......@@ -365,7 +367,7 @@ class HeronCUDA(CUDAModel, BBHSurrogate, HofTSurrogate):
class HeronCUDAIMR(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
class HeronCUDAIMR(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate, FrequencyMixin):
"""
A GPR BBH waveform model which is capable of using CUDA resources.
"""
......@@ -378,6 +380,7 @@ class HeronCUDAIMR(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
Construct a CUDA-based waveform model with pyTorch
"""
self.device = device
self.polarisations = ["plus", "cross"]
# super(HeronCUDA, self).__init__()
#assert torch.cuda.is_available() # This is a bit of a kludge
(self.model_plus, self.model_cross), self.likelihood = self.build()
......@@ -568,38 +571,7 @@ class HeronCUDAIMR(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
for sample in y_preds.sample_n(samples)])
return return_samples
def frequency_domain_waveform(self, p, window, times=np.linspace(-2, 2, 1000)):
"""
Return the frequency domain waveform.
"""
data = {}
for polarisation in ["plus", "cross"]:
mean, _, cov = self._predict(times, p, polarisation=polarisation)
strain_f = torch.view_as_complex((window*mean.double()).rfft(1))
cov_f = torch.view_as_complex(cov.rfft(2))
dt = times[-1] - times[-2]
uncert_f = diag_cuda(cov_f)
#Complex(torch.stack([torch.diag(cov_f[:, :, 0]), torch.diag(cov_f[:, :, 1])]).T)
if not isinstance(times, type(None)):
srate = 1/np.diff(times).mean()
nf = int(np.floor(len(times)/2))+1
frequencies = np.linspace(0, srate, nf)
data[polarisation] = FrequencySeries(data=strain_f,
frequencies=frequencies,
covariance=cov_f,
variance=uncert_f)
return data
def time_domain_waveform(self, p, times=np.linspace(-2, 2, 1000)):
"""
Return the timedomain waveform.
"""
return self.mean(times, p)
def plot_td(self, parameters, times=np.linspace(-2, 2, 1000), f=None):
"""
......@@ -622,7 +594,7 @@ class HeronCUDAIMR(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
class HeronCUDAMix(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
class HeronCUDAMix(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate, FrequencyMixin):
"""
A GPR BBH waveform model which is capable of using CUDA resources.
"""
......@@ -643,14 +615,17 @@ class HeronCUDAMix(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
self.specification = specification
self.data = np.genfromtxt(training_data)
super().__init__()
self.polarisations = ["plus"]
self.x_dimensions = 2
self.time_factor = 100
self.strain_input_factor = 1e21
self.total_mass = specification['total mass']
self.state_vector = f"{self.specification['name']}.pth"
self.model_plus, self.likelihood = self.build()
self.eval()
def build(self):
"""
"""
......@@ -697,17 +672,21 @@ class HeronCUDAMix(CUDAModel, BBHNonSpinSurrogate, HofTSurrogate):
return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
data = self.data
idx = data[:, 1] == 0
tx = data[:, [2, 3]]*self.time_factor
tx = data[idx][:, [2, 3]]*self.time_factor
training_x = self.training_x = torch.tensor(tx, device=self.device).float()
training_y = self.training_y = torch.tensor(data[:, -1]*self.strain_input_factor, device=self.device).float()
training_y = self.training_y = torch.tensor(data[idx, -1]*self.strain_input_factor, device=self.device).float()
likelihood = gpytorch.likelihoods.GaussianLikelihood()#noise_constraint=LessThan(1))
model = ExactGPModel(training_x, training_y, likelihood)
if not (self.device.type == "cpu") and torch.cuda.is_available():
model = model.cuda()
likelihood = likelihood.cuda()
model.load_state_dict(torch.load(self.state_vector))
if os.path.exists(self.state_vector):
model.load_state_dict(torch.load(self.state_vector))
return model, likelihood
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment