likelihood.py 7.12 KB
Newer Older
1
from __future__ import division, print_function
RorySmith's avatar
RorySmith committed
2
import numpy as np
3

4 5 6 7
try:
    from scipy.special import logsumexp
except ImportError:
    from scipy.misc import logsumexp
8
from scipy.special import i0e
9
from scipy.interpolate import interp1d
moritz's avatar
moritz committed
10
import tupak
11
import logging
RorySmith's avatar
RorySmith committed
12

Gregory Ashton's avatar
Gregory Ashton committed
13

14
class GravitationalWaveTransient(object):
Gregory Ashton's avatar
Gregory Ashton committed
15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44
    """ A gravitational-wave transient likelihood object

    This is the usual likelihood object to use for transient gravitational
    wave parameter estimation. It computes the log-likelihood in the frequency
    domain assuming a colored Gaussian noise model described by a power
    spectral density


    Parameters
    ----------
    interferometers: list
        A list of `tupak.detector.Interferometer` instances - contains the
        detector data and power spectral densities
    waveform_generator: `tupak.waveform_generator.WaveformGenerator`
        An object which computes the frequency-domain strain of the signal,
        given some set of parameters
    distance_marginalization: bool
        If true, analytic distance marginalization
    phase_marginalization: bool
        If true, analytic phase marginalization
    prior: dict
        If given, used in the distance and phase marginalization.

    Returns
    -------
    Likelihood: `tupak.likelihood.Likelihood`
        A likehood object, able to compute the likelihood of the data given
        some model parameters

    """
45 46
    def __init__(self, interferometers, waveform_generator, distance_marginalization=False, phase_marginalization=False,
                 prior=None):
47
        # GravitationalWaveTransient.__init__(self, interferometers, waveform_generator)
Gregory Ashton's avatar
Gregory Ashton committed
48
        self.interferometers = interferometers
49
        self.waveform_generator = waveform_generator
50
        self.parameters = self.waveform_generator.parameters
51
        self.non_standard_sampling_parameter_keys = self.waveform_generator.non_standard_sampling_parameter_keys
52 53 54
        self.distance_marginalization = distance_marginalization
        self.phase_marginalization = phase_marginalization
        self.prior = prior
Colm Talbot's avatar
Colm Talbot committed
55

56
        if self.distance_marginalization:
Colm Talbot's avatar
Colm Talbot committed
57 58 59 60
            self.distance_array = np.array([])
            self.delta_distance = 0
            self.distance_prior_array = np.array([])
            self.setup_distance_marginalization()
61
            prior['luminosity_distance'] = 1
Colm Talbot's avatar
Colm Talbot committed
62

63
        if self.phase_marginalization:
Colm Talbot's avatar
Colm Talbot committed
64 65
            self.bessel_function_interped = None
            self.setup_phase_marginalization()
66 67 68 69 70
            prior['psi'] = 0

    def noise_log_likelihood(self):
        log_l = 0
        for interferometer in self.interferometers:
moritz's avatar
moritz committed
71 72 73
            log_l -= tupak.utils.noise_weighted_inner_product(interferometer.data, interferometer.data,
                                                              interferometer.power_spectral_density_array,
                                                              self.waveform_generator.time_duration) / 2
74 75 76 77 78
        return log_l.real

    def log_likelihood_ratio(self):
        waveform_polarizations = self.waveform_generator.frequency_domain_strain()

79
        if waveform_polarizations is None:
80
            return np.nan_to_num(-np.inf)
81

82
        matched_filter_snr_squared = 0
83 84 85 86 87
        optimal_snr_squared = 0

        for interferometer in self.interferometers:
            signal_ifo = interferometer.get_detector_response(waveform_polarizations,
                                                              self.waveform_generator.parameters)
moritz's avatar
moritz committed
88
            matched_filter_snr_squared += tupak.utils.matched_filter_snr_squared(
Colm Talbot's avatar
Colm Talbot committed
89
                signal_ifo, interferometer, self.waveform_generator.time_duration)
90

moritz's avatar
moritz committed
91
            optimal_snr_squared += tupak.utils.optimal_snr_squared(
Colm Talbot's avatar
Colm Talbot committed
92
                signal_ifo, interferometer, self.waveform_generator.time_duration)
Colm Talbot's avatar
Colm Talbot committed
93

Colm Talbot's avatar
Colm Talbot committed
94
        if self.phase_marginalization:
95
            matched_filter_snr_squared = self.bessel_function_interped(abs(matched_filter_snr_squared))
96

Colm Talbot's avatar
Colm Talbot committed
97 98 99 100 101 102
        else:
            matched_filter_snr_squared = matched_filter_snr_squared.real

        if self.distance_marginalization:

            optimal_snr_squared_array = optimal_snr_squared \
103 104
                                        * self.waveform_generator.parameters['luminosity_distance'] ** 2 \
                                        / self.distance_array ** 2
Colm Talbot's avatar
Colm Talbot committed
105

106 107
            matched_filter_snr_squared_array = matched_filter_snr_squared * \
                self.waveform_generator.parameters['luminosity_distance'] / self.distance_array
108

Colm Talbot's avatar
Colm Talbot committed
109 110
            log_l = logsumexp(matched_filter_snr_squared_array - optimal_snr_squared_array / 2,
                              b=self.distance_prior_array * self.delta_distance)
111 112 113 114 115 116 117 118
        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()

119 120 121 122 123
    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'] = tupak.prior.create_default_prior('luminosity_distance')
        self.distance_array = np.linspace(self.prior['luminosity_distance'].minimum,
124
                                          self.prior['luminosity_distance'].maximum, int(1e4))
125 126 127 128 129 130 131 132
        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(self.prior['psi'], tupak.prior.Prior):
            logging.info('No prior provided for polarization, using default prior.')
            self.prior['psi'] = tupak.prior.create_default_prior('psi')
133 134 135
        self.bessel_function_interped = interp1d(np.linspace(0, 1e6, int(1e5)),
                                                 np.log([i0e(snr) for snr in np.linspace(0, 1e6, int(1e5))])
                                                 + np.linspace(0, 1e6, int(1e5)),
136 137 138
                                                 bounds_error=False, fill_value=-np.inf)


139 140 141 142 143 144 145 146 147
def get_binary_black_hole_likelihood(interferometers):
    """ A rapper to quickly set up a likelihood for BBH parameter estimation

    Parameters
    ----------
    interferometers: list
        A list of `tupak.detector.Interferometer` instances, typically the
        output of either `tupak.detector.get_interferometer_with_open_data`
        or `tupak.detector.get_interferometer_with_fake_noise_and_injection`
Gregory Ashton's avatar
Gregory Ashton committed
148

149
    Returns
Gregory Ashton's avatar
Gregory Ashton committed
150
    -------
151
    likelihood: tupak.likelihood.GravitationalWaveTransient
152 153
        The likelihood to pass to `run_sampler`
    """
Colm Talbot's avatar
Colm Talbot committed
154 155 156 157
    waveform_generator = tupak.waveform_generator.WaveformGenerator(
        time_duration=interferometers[0].duration, sampling_frequency=interferometers[0].sampling_frequency,
        frequency_domain_source_model=tupak.source.lal_binary_black_hole,
        parameters={'waveform_approximant': 'IMRPhenomPv2', 'reference_frequency': 50})
158
    likelihood = tupak.likelihood.GravitationalWaveTransient(interferometers, waveform_generator)
159
    return likelihood