Commit 930280f6 authored by Gregory Ashton's avatar Gregory Ashton
Browse files

Merge branch 'implementing_ROQ' into 'master'

ROQ implementation into bilby

See merge request !247
parents a2c84e56 0a28644c
Pipeline #45865 passed with stages
in 11 minutes and 26 seconds
......@@ -29,6 +29,7 @@ Changes currently on master, but not under a tag.
- Added method to result to get injection recovery credible levels
- Added function to generate a pp-plot from many results to core/result.py
- Fixed a bug which caused `Interferometer.detector_tensor` not to update when `latitude`, `longitude`, `xarm_azimuth`, `yarm_azimuth`, `xarm_tilt`, `yarm_tilt` were updated.
- Added implementation of the ROQ likelihood. The basis needs to be specified by the user.
- Extracted time and frequency series behaviour from `WaveformGenerator` and `InterferometerStrainData` and moved it to `series.gw.CoupledTimeAndFrequencySeries`
### Changes
......
......@@ -358,8 +358,8 @@ class InterferometerStrainData(object):
-------
array_like: An array of boolean values
"""
return ((self.frequency_array > self.minimum_frequency) &
(self.frequency_array < self.maximum_frequency))
return ((self.frequency_array >= self.minimum_frequency) &
(self.frequency_array <= self.maximum_frequency))
@property
def alpha(self):
......
......@@ -14,8 +14,9 @@ from ..core.prior import Prior, Uniform
from .detector import InterferometerList
from .prior import BBHPriorDict
from .source import lal_binary_black_hole
from .utils import noise_weighted_inner_product
from .utils import noise_weighted_inner_product, build_roq_weights, blockwise_dot_product
from .waveform_generator import WaveformGenerator
from math import ceil
class GravitationalWaveTransient(likelihood.Likelihood):
......@@ -385,12 +386,194 @@ class BasicGravitationalWaveTransient(likelihood.Likelihood):
return log_l.real
class ROQGravitationalWaveTransient(GravitationalWaveTransient):
"""A reduced order quadrature likelihood object
This uses the method described in Smith et al., (2016) Phys. Rev. D 94,
044031. A public repository of the ROQ data is available from
https://git.ligo.org/lscsoft/ROQ_data.
Parameters
----------
interferometers: list, bilby.gw.detector.InterferometerList
A list of `bilby.detector.Interferometer` instances - contains the
detector data and power spectral densities
waveform_generator: `bilby.waveform_generator.WaveformGenerator`
An object which computes the frequency-domain strain of the signal,
given some set of parameters
linear_matrix: str, array
Either a string point to the file from which to load the linear_matrix
array, or the array itself.
quadratic_matrix: str, array
Either a string point to the file from which to load the quadratic_matrix
array, or the array itself.
priors: dict, bilby.prior.PriorDict
A dictionary of priors containing at least the geocent_time prior
"""
def __init__(self, interferometers, waveform_generator,
linear_matrix, quadratic_matrix, priors):
GravitationalWaveTransient.__init__(
self, interferometers=interferometers,
waveform_generator=waveform_generator, priors=priors)
if isinstance(linear_matrix, str):
logger.info("Loading linear matrix from {}".format(linear_matrix))
linear_matrix = np.load(linear_matrix).T
if isinstance(quadratic_matrix, str):
logger.info("Loading quadratic_matrix from {}".format(quadratic_matrix))
quadratic_matrix = np.load(quadratic_matrix).T
self.linear_matrix = linear_matrix
self.quadratic_matrix = quadratic_matrix
self.time_samples = None
self.weights = dict()
self._set_weights()
self.frequency_nodes_linear =\
waveform_generator.waveform_arguments['frequency_nodes_linear']
def log_likelihood_ratio(self):
optimal_snr_squared = 0.
matched_filter_snr_squared = 0.
indices, in_bounds = self._closest_time_indices(
self.parameters['geocent_time'] - self.interferometers.start_time)
if not in_bounds:
return np.nan_to_num(-np.inf)
waveform = self.waveform_generator.frequency_domain_strain(
self.parameters)
if waveform is None:
return np.nan_to_num(-np.inf)
for ifo in self.interferometers:
f_plus = ifo.antenna_response(
self.parameters['ra'], self.parameters['dec'],
self.parameters['geocent_time'], self.parameters['psi'], 'plus')
f_cross = ifo.antenna_response(
self.parameters['ra'], self.parameters['dec'],
self.parameters['geocent_time'], self.parameters['psi'], 'cross')
dt = ifo.time_delay_from_geocenter(
self.parameters['ra'], self.parameters['dec'],
ifo.strain_data.start_time)
ifo_time = self.parameters['geocent_time'] + dt - \
ifo.strain_data.start_time
h_plus_linear = f_plus * waveform['linear']['plus']
h_cross_linear = f_cross * waveform['linear']['cross']
h_plus_quadratic = f_plus * waveform['quadratic']['plus']
h_cross_quadratic = f_cross * waveform['quadratic']['cross']
indices, in_bounds = self._closest_time_indices(ifo_time)
if not in_bounds:
return np.nan_to_num(-np.inf)
matched_filter_snr_squared_array = np.einsum(
'i,ji->j', np.conjugate(h_plus_linear + h_cross_linear),
self.weights[ifo.name + '_linear'][indices])
matched_filter_snr_squared += interp1d(
self.time_samples[indices],
matched_filter_snr_squared_array, kind='quadratic')(ifo_time)
optimal_snr_squared += \
np.vdot(np.abs(h_plus_quadratic + h_cross_quadratic)**2,
self.weights[ifo.name + '_quadratic'])
log_l = matched_filter_snr_squared - optimal_snr_squared / 2
return log_l.real
def _closest_time_indices(self, time):
"""
Get the closest an two neighbouring times
Parameters
----------
time: float
Time to check
Returns
-------
indices: list
Indices nearest to time.
in_bounds: bool
Whether the indices are for valid times.
"""
closest = np.argmin(np.absolute(self.time_samples - time))
indices = [closest + ii for ii in [-1, 0, 1]]
in_bounds = (indices[0] >= 0) & (indices[2] < self.time_samples.size)
return indices, in_bounds
def _set_weights(self):
"""
Setup the time-dependent ROQ weights.
This follows FIXME: Smith et al.
The times are chosen to allow all the merger times allows in the time
prior.
"""
for ifo in self.interferometers:
# only get frequency components up to maximum_frequency
self.linear_matrix = \
self.linear_matrix[:, :sum(ifo.frequency_mask)]
self.quadratic_matrix = \
self.quadratic_matrix[:, :sum(ifo.frequency_mask)]
# array of relative time shifts to be applied to the data
# 0.045s comes from time for GW to traverse the Earth
self.time_samples = np.linspace(
self.priors['geocent_time'].minimum - 0.045,
self.priors['geocent_time'].maximum + 0.045,
int(ceil((self.priors['geocent_time'].maximum -
self.priors['geocent_time'].minimum + 0.09) *
ifo.strain_data.sampling_frequency)))
self.time_samples -= ifo.strain_data.start_time
time_space = self.time_samples[1] - self.time_samples[0]
# array to be filled with data, shifted by discrete time_samples
tc_shifted_data = np.zeros([
len(self.time_samples),
len(ifo.frequency_array[ifo.frequency_mask])], dtype=complex)
# shift data to beginning of the prior
# increment by the time step
shifted_data =\
ifo.frequency_domain_strain[ifo.frequency_mask] * \
np.exp(2j * np.pi * ifo.frequency_array[ifo.frequency_mask] *
self.time_samples[0])
single_time_shift = np.exp(
2j * np.pi * ifo.frequency_array[ifo.frequency_mask] *
time_space)
for j in range(len(self.time_samples)):
tc_shifted_data[j] = shifted_data
shifted_data *= single_time_shift
# to not kill all computers this minimises the memory usage of the
# required inner products
max_block_gigabytes = 4
max_elements = int((max_block_gigabytes * 2 ** 30) / 8)
self.weights[ifo.name + '_linear'] = blockwise_dot_product(
tc_shifted_data /
ifo.power_spectral_density_array[ifo.frequency_mask],
self.linear_matrix, max_elements) * 4 / ifo.strain_data.duration
del tc_shifted_data
self.weights[ifo.name + '_quadratic'] = build_roq_weights(
1 / ifo.power_spectral_density_array[ifo.frequency_mask],
self.quadratic_matrix.real, 1 / ifo.strain_data.duration)
def get_binary_black_hole_likelihood(interferometers):
""" A rapper to quickly set up a likelihood for BBH parameter estimation
Parameters
----------
interferometers: list
interferometers: {bilby.gw.detector.InterferometerList, list}
A list of `bilby.detector.Interferometer` instances, typically the
output of either `bilby.detector.get_interferometer_with_open_data`
or `bilby.detector.get_interferometer_with_fake_noise_and_injection`
......
......@@ -12,6 +12,7 @@ from .utils import (lalsim_SimInspiralTransformPrecessingNewInitialConditions,
try:
import lal
import lalsimulation as lalsim
except ImportError:
logger.warning("You do not have lalsuite installed currently. You will"
" not be able to use some of the prebuilt functions.")
......@@ -333,3 +334,112 @@ def lal_binary_neutron_star(
h_cross = h_cross[:len(frequency_array)]
return {'plus': h_plus, 'cross': h_cross}
def roq(frequency_array, mass_1, mass_2, luminosity_distance, a_1, tilt_1,
phi_12, a_2, tilt_2, phi_jl, iota, phase, **waveform_arguments):
"""
See https://git.ligo.org/lscsoft/lalsuite/blob/master/lalsimulation/src/LALSimInspiral.c#L1460
Parameters
----------
frequency_array: np.array
This input is ignored for the roq source model
mass_1: float
The mass of the heavier object in solar masses
mass_2: float
The mass of the lighter object in solar masses
luminosity_distance: float
The luminosity distance in megaparsec
a_1: float
Dimensionless primary spin magnitude
tilt_1: float
Primary tilt angle
phi_12: float
a_2: float
Dimensionless secondary spin magnitude
tilt_2: float
Secondary tilt angle
phi_jl: float
iota: float
Orbital inclination
phase: float
The phase at coalescence
Waveform arguments
------------------
Non-sampled extra data used in the source model calculation
frequency_nodes_linear: np.array
frequency_nodes_quadratic: np.array
reference_frequency: float
version: str
Note: for the frequency_nodes_linear and frequency_nodes_quadratic arguments,
if using data from https://git.ligo.org/lscsoft/ROQ_data, this should be
loaded as `np.load(filename).T`.
Returns
-------
waveform_polarizations: dict
Dict containing plus and cross modes evaluated at the linear and
quadratic frequency nodes.
"""
if mass_2 > mass_1:
return None
frequency_nodes_linear = waveform_arguments['frequency_nodes_linear']
frequency_nodes_quadratic = waveform_arguments['frequency_nodes_quadratic']
reference_frequency = getattr(waveform_arguments,
'reference_frequency', 20.0)
versions = dict(IMRPhenomPv2=lalsim.IMRPhenomPv2_V)
version = versions[getattr(waveform_arguments, 'version', 'IMRPhenomPv2')]
luminosity_distance = luminosity_distance * 1e6 * utils.parsec
mass_1 = mass_1 * utils.solar_mass
mass_2 = mass_2 * utils.solar_mass
if tilt_1 == 0 and tilt_2 == 0:
spin_1x = 0
spin_1y = 0
spin_1z = a_1
spin_2x = 0
spin_2y = 0
spin_2z = a_2
else:
iota, spin_1x, spin_1y, spin_1z, spin_2x, spin_2y, spin_2z = \
lalsim.SimInspiralTransformPrecessingNewInitialConditions(
iota, phi_jl, tilt_1, tilt_2, phi_12, a_1, a_2, mass_1, mass_2,
reference_frequency, phase)
chi_1_l, chi_2_l, chi_p, theta_jn, alpha, phase_aligned, zeta =\
lalsim.SimIMRPhenomPCalculateModelParametersFromSourceFrame(
mass_1, mass_2, reference_frequency, phase, iota, spin_1x,
spin_1y, spin_1z, spin_2x, spin_2y, spin_2z, version)
waveform_polarizations = dict()
h_linear_plus, h_linear_cross = lalsim.SimIMRPhenomPFrequencySequence(
frequency_nodes_linear, chi_1_l, chi_2_l, chi_p, theta_jn,
mass_1, mass_2, luminosity_distance,
alpha, phase_aligned, reference_frequency, version, None)
h_quadratic_plus, h_quadratic_cross = lalsim.SimIMRPhenomPFrequencySequence(
frequency_nodes_quadratic, chi_1_l, chi_2_l, chi_p, theta_jn,
mass_1, mass_2, luminosity_distance,
alpha, phase_aligned, reference_frequency, version, None)
waveform_polarizations['linear'] = dict(
plus=(np.cos(2 * zeta) * h_linear_plus.data.data +
np.sin(2 * zeta) * h_linear_cross.data.data),
cross=(np.cos(2 * zeta) * h_linear_cross.data.data -
np.sin(2 * zeta) * h_linear_plus.data.data))
waveform_polarizations['quadratic'] = dict(
plus=(np.cos(2 * zeta) * h_quadratic_plus.data.data +
np.sin(2 * zeta) * h_quadratic_cross.data.data),
cross=(np.cos(2 * zeta) * h_quadratic_cross.data.data -
np.sin(2 * zeta) * h_quadratic_plus.data.data))
return waveform_polarizations
......@@ -632,6 +632,88 @@ def plot_skymap(result, center='120d -40d', nside=512):
fig.savefig('{}/{}_skymap.png'.format(result.outdir, result.label))
def build_roq_weights(data, basis, deltaF):
"""
for a data array and reduced basis compute roq weights
basis: (reduced basis element)*invV (the inverse Vandermonde matrix)
data: data set
PSD: detector noise power spectral density (must be same shape as data)
deltaF: integration element df
"""
weights = np.dot(data, np.conjugate(basis)) * deltaF * 4.
return weights
def blockwise_dot_product(matrix_a, matrix_b, max_elements=int(2 ** 27),
out=None):
"""
Memory efficient
Computes the dot product of two matrices in a block-wise fashion.
Only blocks of `matrix_a` with a maximum size of `max_elements` will be
processed simultaneously.
Parameters
----------
matrix_a, matrix_b: array-like
Matrices to be dot producted, matrix_b is complex conjugated.
max_elements: int
Maximum number of elements to consider simultaneously, should be memory
limited.
out: array-like
Output array
Return
------
out: array-like
Dot producted array
"""
def block_slices(dim_size, block_size):
"""Generator that yields slice objects for indexing into
sequential blocks of an array along a particular axis
Useful for blockwise dot
"""
count = 0
while True:
yield slice(count, count + block_size, 1)
count += block_size
if count > dim_size:
return
matrix_b = np.conjugate(matrix_b)
m, n = matrix_a.shape
n1, o = matrix_b.shape
if n1 != n:
raise ValueError(
'Matrices are not aligned, matrix a has shape ' +
'{}, matrix b has shape {}.'.format(matrix_a.shape, matrix_b.shape))
if matrix_a.flags.f_contiguous:
# prioritize processing as many columns of matrix_a as possible
max_cols = max(1, max_elements // m)
max_rows = max_elements // max_cols
else:
# prioritize processing as many rows of matrix_a as possible
max_rows = max(1, max_elements // n)
max_cols = max_elements // max_rows
if out is None:
out = np.empty((m, o), dtype=np.result_type(matrix_a, matrix_b))
elif out.shape != (m, o):
raise ValueError('Output array has incorrect dimensions.')
for mm in block_slices(m, max_rows):
out[mm, :] = 0
for nn in block_slices(n, max_cols):
a_block = matrix_a[mm, nn].copy() # copy to force a read
out[mm, :] += np.dot(a_block, matrix_b[nn, :])
del a_block
return out
def convert_args_list_to_float(*args_list):
""" Converts inputs to floats, returns a list in the same order as the input"""
try:
......
#!/usr/bin/env python
"""
Example of how to use the Reduced Order Quadrature method (see Smith et al.,
(2016) Phys. Rev. D 94, 044031) for a Binary Black hole simulated signal in
Gaussian noise.
This requires files specifying the appropriate basis weights.
These aren't shipped with Bilby, but are available on LDG clusters and
from the public repository https://git.ligo.org/lscsoft/ROQ_data.
"""
from __future__ import division, print_function
import numpy as np
import bilby
outdir = 'outdir'
label = 'roq'
# Load in the pieces for the linear part of the ROQ. Note you will need to
# adjust the filenames here to the correct paths on your machine
basis_matrix_linear = np.load("B_linear.npy").T
freq_nodes_linear = np.load("fnodes_linear.npy")
# Load in the pieces for the quadratic part of the ROQ
basic_matrix_quadratic = np.load("B_quadratic.npy").T
freq_nodes_quadratic = np.load("fnodes_quadratic.npy")
np.random.seed(170808)
duration = 4
sampling_frequency = 2048
injection_parameters = dict(
chirp_mass=36., mass_ratio=0.9, a_1=0.4, a_2=0.3, tilt_1=0.0, tilt_2=0.0,
phi_12=1.7, phi_jl=0.3, luminosity_distance=1000., iota=0.4, psi=0.659,
phase=1.3, geocent_time=1126259642.413, ra=1.375, dec=-1.2108)
waveform_arguments = dict(waveform_approximant='IMRPhenomPv2',
reference_frequency=20., minimum_frequency=20.)
waveform_generator = bilby.gw.WaveformGenerator(
duration=duration, sampling_frequency=sampling_frequency,
frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
waveform_arguments=waveform_arguments,
parameter_conversion=bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters)
ifos = bilby.gw.detector.InterferometerList(['H1', 'L1', 'V1'])
ifos.set_strain_data_from_power_spectral_densities(
sampling_frequency=sampling_frequency, duration=duration,
start_time=injection_parameters['geocent_time'] - 3)
ifos.inject_signal(waveform_generator=waveform_generator,
parameters=injection_parameters)
# make ROQ waveform generator
search_waveform_generator = bilby.gw.waveform_generator.WaveformGenerator(
duration=duration, sampling_frequency=sampling_frequency,
frequency_domain_source_model=bilby.gw.source.roq,
waveform_arguments=dict(frequency_nodes_linear=freq_nodes_linear,
frequency_nodes_quadratic=freq_nodes_quadratic,
reference_frequency=20., minimum_frequency=20.,
approximant='IMRPhenomPv2'),
parameter_conversion=bilby.gw.conversion.convert_to_lal_binary_black_hole_parameters)
priors = bilby.gw.prior.BBHPriorDict()
for key in ['a_1', 'a_2', 'tilt_1', 'tilt_2', 'iota', 'phase', 'psi', 'ra',
'dec', 'phi_12', 'phi_jl', 'luminosity_distance']:
priors[key] = injection_parameters[key]
priors.pop('mass_1')
priors.pop('mass_2')
priors['chirp_mass'] = bilby.core.prior.Uniform(
15, 40, latex_label='$\\mathcal{M}$')
priors['mass_ratio'] = bilby.core.prior.Uniform(0.5, 1, latex_label='$q$')
priors['geocent_time'] = bilby.core.prior.Uniform(
injection_parameters['geocent_time'] - 0.1,
injection_parameters['geocent_time'] + 0.1, latex_label='$t_c$', unit='s')
likelihood = bilby.gw.likelihood.ROQGravitationalWaveTransient(
interferometers=ifos, waveform_generator=search_waveform_generator,
linear_matrix=basis_matrix_linear, quadratic_matrix=basic_matrix_quadratic,
prior=priors)
result = bilby.run_sampler(
likelihood=likelihood, priors=priors, sampler='dynesty', npoints=500,
injection_parameters=injection_parameters, outdir=outdir, label=label)
# Make a corner plot.
result.plot_corner()
......@@ -36,14 +36,13 @@ class TestBasicGWTransient(unittest.TestCase):
def test_noise_log_likelihood(self):
"""Test noise log likelihood matches precomputed value"""
self.likelihood.noise_log_likelihood()
self.assertAlmostEqual(self.likelihood.noise_log_likelihood(),
-4036.1064342687155, 3)
self.assertAlmostEqual(-4037.0994372143414, self.likelihood.noise_log_likelihood(), 3)
def test_log_likelihood(self):
"""Test log likelihood matches precomputed value"""
self.likelihood.log_likelihood()
self.assertAlmostEqual(self.likelihood.log_likelihood(),
-4054.2229111227016, 3)
-4055.2526551677647, 3)
def test_log_likelihood_ratio(self):
"""Test log likelihood ratio returns the correct value"""
......@@ -106,14 +105,13 @@ class TestGWTransient(unittest.TestCase):
def test_noise_log_likelihood(self):
"""Test noise log likelihood matches precomputed value"""
self.likelihood.noise_log_likelihood()
self.assertAlmostEqual(self.likelihood.noise_log_likelihood(),
-4036.1064342687155, 3)
self.assertAlmostEqual(-4037.0994372143414, self.likelihood.noise_log_likelihood(), 3)
def test_log_likelihood(self):
"""Test log likelihood matches precomputed value"""
self.likelihood.log_likelihood()
self.assertAlmostEqual(self.likelihood.log_likelihood(),
-4054.2229111227016, 3)
-4055.2526551677647, 3)
def test_log_likelihood_ratio(self):
"""Test log likelihood ratio returns the correct value"""
......@@ -457,6 +455,79 @@ class TestTimePhaseMarginalization(unittest.TestCase):
delta=0.5)
class TestROQLikelihood(unittest.TestCase):
def setUp(self):
self.duration = 4
self.sampling_frequency = 2048
roq_dir = '/roq_basis'
linear_matrix_file = "{}/B_linear.npy".format(roq_dir)
quadratic_matrix_file = "{}/B_quadratic.npy".format(roq_dir)
fnodes_linear_file = "{}/fnodes_linear.npy".format(roq_dir)
fnodes_linear = np.load(fnodes_linear_file).T
fnodes_quadratic_file = "{}/fnodes_quadratic.npy".format(roq_dir)
fnodes_quadratic = np.load(fnodes_quadratic_file).T
self.test_parameters = dict(
mass_1=36.0, mass_2=36.0, a_1=0.0, a_2=0.0, tilt_1=0.0,
tilt_2=0.0, phi_12=1.7, phi_jl=0.3, luminosity_distance=5000.,
iota=0.4, psi=0.659, phase=1.3, geocent_time=1.2, ra=1.3, dec=-1.2)
ifos = bilby.gw.detector.InterferometerList(['H1'])
ifos.set_strain_data_from_power_spectral_densities(
sampling_frequency=self.sampling_frequency, duration=self.duration)
self.priors = bilby.gw.prior.BBHPriorDict()
self.priors['geocent_time'] = bilby.core.prior.Uniform(1.1, 1.3)
non_roq_wfg = bilby.gw.WaveformGenerator(
duration=self.duration, sampling_frequency=self.sampling_frequency,
frequency_domain_source_model=bilby.gw.source.lal_binary_black_hole,
waveform_arguments=dict(
reference_frequency=20.0, minimum_frequency=20.0,
approximant='IMRPhenomPv2'))
ifos.inject_signal(
parameters=self.test_parameters, waveform_generator=non_roq_wfg)
roq_wfg = bilby.gw.waveform_generator.WaveformGenerator(
duration=self.duration, sampling_frequency=self.sampling_frequency,