Merge branch 'marginalisation' into 'master'

Marginalisation

See merge request Monash/peyote!27

peyote/detector.py

 from __future__ import division, print_function, absolute_import
 import numpy as np
+import peyote
 import logging
 import os
 from scipy.interpolate import interp1d
@@ -210,7 +211,16 @@ class Interferometer(object):
         :param waveform_polarizations: dict, polarizations of the waveform
         :param parameters: dict, parameters describing position and time of arrival of the signal
         """
-        self.data += self.get_detector_response(waveform_polarizations, parameters)
+        signal_ifo = self.get_detector_response(waveform_polarizations, parameters)
+        self.data += signal_ifo
+        opt_snr = np.sqrt(peyote.utils.optimal_snr_squared(signal=signal_ifo, interferometer=self,
+                                                           time_duration=1 / (self.frequency_array[1]
+                                                                              - self.frequency_array[0])).real)
+        mf_snr = np.sqrt(peyote.utils.matched_filter_snr_squared(signal=signal_ifo, interferometer=self,
+                                                                 time_duration=1 / (self.frequency_array[1]
+                                                                                    - self.frequency_array[0])).real)
+        logging.info("Injection found with optimal SNR = {:.2f} and matched filter SNR = {:.2f} in {}".format(
+            opt_snr, mf_snr, self.name))
 
     def unit_vector_along_arm(self, arm):
         """

peyote/likelihood.py

+from __future__ import division, print_function
 import numpy as np
+try:
+    from scipy.special import logsumexp
+except ImportError:
+    from scipy.misc import logsumexp
+from scipy.special import i0, i0e
+from scipy.interpolate import interp1d
+import peyote
+import logging
 
 
 class Likelihood(object):
@@ -30,3 +39,120 @@ class Likelihood(object):
 
     def log_likelihood_ratio(self):
         return self.log_likelihood() - self.noise_log_likelihood()
+
+
+class MarginalizedLikelihood(Likelihood):
+    def __init__(self, interferometers, waveform_generator, distance_marginalization=False, phase_marginalization=False,
+                 time_marginalization=False, prior=None):
+        Likelihood.__init__(self, interferometers, waveform_generator)
+        self.distance_marginalization = distance_marginalization
+        self.phase_marginalization = phase_marginalization
+        self.time_marginalization = time_marginalization
+        self.prior = prior
+
+        if self.distance_marginalization:
+            self.distance_array = np.array([])
+            self.delta_distance = 0
+            self.distance_prior_array = np.array([])
+            self.setup_distance_marginalization()
+            prior['luminosity_distance'] = 1
+
+        if self.phase_marginalization:
+            self.bessel_function_interped = None
+            self.setup_phase_marginalization()
+            prior['psi'] = 0
+
+        if self.time_marginalization:
+            logging.warning('Time marginalisation not yet implemented.')
+            self.time_marginalization = False
+
+    def setup_distance_marginalization(self):
+        if 'luminosity_distance' not in self.prior.keys():
+            logging.info('No prior provided for distance, using default prior.')
+            self.prior['luminosity_distance'] = peyote.prior.create_default_prior('luminosity_distance')
+        self.distance_array = np.linspace(self.prior['luminosity_distance'].minimum,
+                                          self.prior['luminosity_distance'].maximum, 1e4)
+        self.delta_distance = self.distance_array[1] - self.distance_array[0]
+        self.distance_prior_array = np.array([self.prior['luminosity_distance'].prob(distance)
+                                              for distance in self.distance_array])
+
+    def setup_phase_marginalization(self):
+        if 'psi' not in self.prior.keys() or not isinstance(prior['psi'], peyote.prior.Prior):
+            logging.info('No prior provided for polarization, using default prior.')
+            self.prior['psi'] = peyote.prior.create_default_prior('psi')
+        self.bessel_function_interped = interp1d(np.linspace(0, 1e6, 1e5),
+                                                 np.log([i0e(snr) for snr in np.linspace(0, 1e6, 1e5)])
+                                                 + np.linspace(0, 1e6, 1e5),
+                                                 bounds_error=False, fill_value=-np.inf)
+
+    def noise_log_likelihood(self):
+        log_l = 0
+        for interferometer in self.interferometers:
+            log_l -= peyote.utils.noise_weighted_inner_product(interferometer.data, interferometer.data,
+                                                               interferometer.power_spectral_density_array,
+                                                               self.waveform_generator.time_duration) / 2
+        return log_l.real
+
+    def log_likelihood_ratio(self):
+        waveform_polarizations = self.waveform_generator.frequency_domain_strain()
+
+        if waveform_polarizations is None:
+            return np.nan_to_num(-np.inf)
+
+        if self.time_marginalization:
+            signal_times_data = 0 * 1j
+        else:
+            matched_filter_snr_squared = 0
+        optimal_snr_squared = 0
+
+        for interferometer in self.interferometers:
+            signal_ifo = interferometer.get_detector_response(waveform_polarizations,
+                                                              self.waveform_generator.parameters)
+            if self.time_marginalization:
+                signal_times_data += np.conj(signal_ifo) * interferometer.data\
+                                     / interferometer.power_spectral_density_array
+            else:
+                matched_filter_snr_squared += peyote.utils.matched_filter_snr_squared(
+                    signal_ifo, interferometer, self.waveform_generator.time_duration)
+            optimal_snr_squared += peyote.utils.optimal_snr_squared(signal_ifo, interferometer,
+                                                                    self.waveform_generator.time_duration)
+
+        if self.time_marginalization:
+            matched_filter_snr_squared = np.abs(np.fft.ifft(signal_times_data))\
+                                         * 4 / self.waveform_generator.time_duration
+
+        if self.phase_marginalization and not self.distance_marginalization and not self.time_marginalization:
+            matched_filter_snr_squared = self.bessel_function_interped(abs(matched_filter_snr_squared))
+
+            log_l = matched_filter_snr_squared - optimal_snr_squared / 2
+
+        elif self.distance_marginalization and not self.phase_marginalization and not self.time_marginalization:
+            log_l = logsumexp(matched_filter_snr_squared * self.waveform_generator.parameters['luminosity_distance']
+                              / self.distance_array
+                              - optimal_snr_squared * self.waveform_generator.parameters['luminosity_distance']**2
+                              / self.distance_array**2 / 2, b=self.distance_prior_array * self.delta_distance)
+
+        elif self.distance_marginalization and self.phase_marginalization and not self.time_marginalization:
+            log_l = logsumexp(self.bessel_function_interped(abs(matched_filter_snr_squared)
+                                                            * self.waveform_generator.parameters['luminosity_distance']
+                                                            / self.distance_array)
+                              - optimal_snr_squared * self.waveform_generator.parameters['luminosity_distance']**2
+                              / self.distance_array**2 / 2, b=self.distance_prior_array * self.delta_distance)
+
+        elif self.distance_marginalization and self.phase_marginalization and self.time_marginalization:
+            log_l = logsumexp(np.array([sum(self.bessel_function_interped(
+                matched_filter_snr_squared * self.waveform_generator.parameters['luminosity_distance']
+                / distance)) for distance in self.distance_array])
+                              - optimal_snr_squared * self.waveform_generator.parameters['luminosity_distance']**2
+                              / self.distance_array**2 / 2, b=self.distance_prior_array * self.delta_distance)\
+                    - np.log(self.waveform_generator.sampling_frequency)
+            print(log_l)
+        else:
+            log_l = matched_filter_snr_squared - optimal_snr_squared / 2
+
+        return log_l.real
+
+    def log_likelihood(self):
+        return self.log_likelihood() + self.noise_log_likelihood()
+
+

peyote/source.py

@@ -15,7 +15,7 @@ def lal_binary_black_hole(
         geocent_time, psi):
     """ A Binary Black Hole waveform model using lalsimulation """
     if mass_2 > mass_1:
-        return {'plus': 0, 'cross': 0}
+        return None
 
     luminosity_distance = luminosity_distance * 1e6 * utils.parsec
     mass_1 = mass_1 * utils.solar_mass

peyote/utils.py

@@ -344,3 +344,38 @@ def inner_product(aa, bb, frequency, PSD):
     product = 4. * np.real(integral)
 
     return product
+
+
+def noise_weighted_inner_product(aa, bb, power_spectral_density, time_duration):
+    """
+    Calculate the noise weighted inner product between two arrays.
+
+    Parameters
+    ----------
+    aa: array
+        Array to be complex conjugated
+    bb: array
+        Array not to be complex conjugated
+    power_spectral_density: array
+        Power spectral density
+    time_duration: float
+        time_duration of the data
+
+    Return
+    ------
+    Noise-weighted inner product.
+    """
+
+    # caluclate the inner product
+    integrand = np.conj(aa) * bb / power_spectral_density
+    product = 4 / time_duration * np.sum(integrand)
+    return product
+
+
+def matched_filter_snr_squared(signal, interferometer, time_duration):
+    return noise_weighted_inner_product(signal, interferometer.data, interferometer.power_spectral_density_array,
+                                        time_duration)
+
+
+def optimal_snr_squared(signal, interferometer, time_duration):
+    return noise_weighted_inner_product(signal, signal, interferometer.power_spectral_density_array, time_duration)

tutorials/BasicTutorial.py

@@ -8,14 +8,16 @@ comoving volume prior on luminosity distance, the cosmology is WMAP7.
 from __future__ import division, print_function
 import peyote
 
-peyote.utils.setup_logger()
+peyote.utils.setup_logger(log_level="info")
 
 time_duration = 4.
 sampling_frequency = 2048.
 outdir = 'outdir'
 
-injection_parameters = dict(mass_1=36., mass_2=29., a_1=0, a_2=0, tilt_1=0, tilt_2=0, phi_12=0, phi_jl=0,
-                            luminosity_distance=2500., iota=0.4, psi=2.659, phase=1.3, geocent_time=1126259642.413,
+np.random.seed(170809)
+
+injection_parameters = dict(mass_1=36., mass_2=29., a_1=0.4, a_2=0.3, tilt_1=0.5, tilt_2=1.0, phi_12=1.7, phi_jl=0.3,
+                            luminosity_distance=4000., iota=0.4, psi=2.659, phase=1.3, geocent_time=1126259642.413,
                             waveform_approximant='IMRPhenomPv2', reference_frequency=50., ra=1.375, dec=-1.2108)
 
 # Create the waveform_generator using a LAL BinaryBlackHole source function
@@ -30,9 +32,6 @@ IFOs = [peyote.detector.get_inteferometer_with_fake_noise_and_injection(
     name, injection_polarizations=hf_signal, injection_parameters=injection_parameters, time_duration=time_duration,
     sampling_frequency=sampling_frequency, outdir=outdir) for name in ['H1', 'L1', 'V1']]
 
-# Set up likelihood
-likelihood = peyote.likelihood.Likelihood(IFOs, waveform_generator)
-
 # Set up prior
 priors = dict()
 # These parameters will not be sampled

tutorials/Injection.py

@@ -4,43 +4,59 @@ Tutorial to show signal injection using new features of detector.py
 from __future__ import division, print_function
 import numpy as np
 import peyote
-
-peyote.utils.setup_logger()
-
-outdir = 'outdir'
-label = 'injection'
-
-# Create the waveform generator
-waveform_generator = peyote.waveform_generator.WaveformGenerator(
-    peyote.source.lal_binary_black_hole, sampling_frequency=4096, time_duration=4,
-    parameters={'reference_frequency': 50.0, 'waveform_approximant': 'IMRPhenomPv2'})
-
-# Define the prior
-# Merger time is some time in 2018, shame LIGO will never see it...
-time_of_event = np.random.uniform(1198800018, 1230336018)
-prior = dict()
-prior['geocent_time'] = peyote.prior.Uniform(time_of_event-0.01, time_of_event+0.01, name='geocent_time')
-peyote.prior.fill_priors(prior, waveform_generator)
-
-# Create signal injection
-injection_parameters = {name: prior[name].sample() for name in prior}
-print(injection_parameters)
-for parameter in injection_parameters:
-    waveform_generator.parameters[parameter] = injection_parameters[parameter]
-injection_polarizations = waveform_generator.frequency_domain_strain()
-
-# Create interferometers and inject signal
-IFOs = [peyote.detector.get_inteferometer_with_fake_noise_and_injection(
-    name, injection_polarizations=injection_polarizations, injection_parameters=injection_parameters,
-    sampling_frequency=4096, time_duration=4, outdir=outdir) for name in ['H1', 'L1', 'V1', 'GEO600']]
-
-# Define a likelihood
-likelihood = peyote.likelihood.Likelihood(IFOs, waveform_generator)
-
-# Run the sampler
-result = peyote.sampler.run_sampler(
-    likelihood, prior, label='injection', sampler='nestle',
-    npoints=200, resume=False, outdir=outdir, use_ratio=True, injection_parameters=injection_parameters)
-truth = [injection_parameters[parameter] for parameter in result.search_parameter_keys]
-result.plot_corner(truth=truth)
-print(result)
+import logging
+
+
+def main():
+
+    peyote.utils.setup_logger()
+
+    outdir = 'outdir'
+    label = 'injection'
+
+    # Create the waveform generator
+    waveform_generator = peyote.waveform_generator.WaveformGenerator(
+        peyote.source.lal_binary_black_hole, sampling_frequency=2048, time_duration=4,
+        parameters={'reference_frequency': 50.0, 'waveform_approximant': 'IMRPhenomPv2'})
+
+    # Define the prior
+    # Merger time is some time in 2018, shame LIGO will never see it...
+    time_of_event = np.random.uniform(1198800018, 1230336018)
+    prior = dict()
+    prior['luminosity_distance'] = peyote.prior.PowerLaw(alpha=2, minimum=100, maximum=5000, name='luminosity_distance')
+    prior['geocent_time'] = peyote.prior.Uniform(time_of_event-0.01, time_of_event+0.01, name='geocent_time')
+    prior['mass_1'] = peyote.prior.Gaussian(mu=40, sigma=5, name='mass_1')
+    prior['mass_2'] = peyote.prior.Gaussian(mu=40, sigma=5, name='mass_2')
+    peyote.prior.fill_priors(prior, waveform_generator)
+
+    # Create signal injection
+    injection_parameters = {name: prior[name].sample() for name in prior}
+    if injection_parameters['mass_1'] < injection_parameters['mass_2']:
+        injection_parameters['mass_1'], injection_parameters['mass_2'] =\
+            injection_parameters['mass_2'], injection_parameters['mass_1']
+    logging.info("Injection parameters:\n{}".format("\n".join(["{}: {}".format(key, injection_parameters[key])
+                                                               for key in injection_parameters])))
+    for parameter in injection_parameters:
+        waveform_generator.parameters[parameter] = injection_parameters[parameter]
+    injection_polarizations = waveform_generator.frequency_domain_strain()
+
+    # Create interferometers and inject signal
+    interferometers = [peyote.detector.get_inteferometer_with_fake_noise_and_injection(
+            name, injection_polarizations=injection_polarizations, injection_parameters=injection_parameters,
+            sampling_frequency=2048, time_duration=4, outdir=outdir) for name in ['H1', 'L1', 'V1']]
+
+    # Define a likelihood
+    likelihood = peyote.likelihood.MarginalizedLikelihood(interferometers, waveform_generator, prior=prior,
+                                                          distance_marginalization=True, phase_marginalization=True)
+
+    # Run the sampler
+    result = peyote.sampler.run_sampler(
+        likelihood, prior, label=label, sampler='dynesty', npoints=500, resume=False, outdir=outdir, use_ratio=True,
+        injection_parameters=injection_parameters)
+    truth = [injection_parameters[parameter] for parameter in result.search_parameter_keys]
+    result.plot_corner(truth=truth)
+    print(result)
+
+
+if __name__ == "__main__":
+    main()