Add gated gaussian noise model that marginalizes over phase

pycbc/inference/models/__init__.py

@@ -34,7 +34,8 @@
 from .marginalized_gaussian_noise import MarginalizedTime
 from .brute_marg import BruteParallelGaussianMarginalize
 from .brute_marg import BruteLISASkyModesMarginalize
-from .gated_gaussian_noise import (GatedGaussianNoise, GatedGaussianMargPol)
+from .gated_gaussian_noise import (GatedGaussianNoise, GatedGaussianMargPol,
+                                   GatedGaussianMargPhase)
 from .single_template import SingleTemplate
 from .relbin import Relative, RelativeTime, RelativeTimeDom
 from .hierarchical import (HierarchicalModel, MultiSignalModel,
@@ -202,6 +203,7 @@ def read_from_config(cp, **kwargs):
     BruteLISASkyModesMarginalize,
     GatedGaussianNoise,
     GatedGaussianMargPol,
+    GatedGaussianMargPhase,
     SingleTemplate,
     Relative,
     RelativeTime,

pycbc/inference/models/gated_gaussian_noise.py

@@ -53,7 +53,7 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None,
         self._overwhitened_data = {}
         # cache the current gated data
         self._gated_data = {}
-        # cache terms related to normalization and gating 
+        # cache terms related to normalization and gating
         self._invasds = {}
         self._Rss = {}
         self._lognorm = {}
@@ -77,7 +77,7 @@ def __init__(self, variable_params, data, low_frequency_cutoff, psds=None,
     def from_config(cls, cp, data_section='data', data=None, psds=None,
                     **kwargs):
         """Adds addiotional keyword arguments based on config file.
-        
+
         Additional keyword arguments are:
 
            * ``highpass_waveforms`` : waveforms will be highpassed.
@@ -142,10 +142,10 @@ def psds(self, psds):
             if self.normalize:
                 self._set_covfit(det)
         self._overwhitened_data = self.whiten(self.data, 2, inplace=False)
- 
+
     def _set_covfit(self, det):
         """Sets the fit function for estimating the covariance determinant.
-        
+
         This must be called after the PSDs have been set, otherwise a
         ValueError will be raised.
         """
@@ -163,7 +163,7 @@ def _set_covfit(self, det):
     def logdet_fit(self, cov, p):
         """Construct a linear regression from a sample of truncated covariance
         matrices.
-        
+
         Returns the sample points used for linear fit generation as well as the
         linear fit parameters.
         """
@@ -197,7 +197,7 @@ def logdet_fit(self, cov, p):
         x = numpy.vstack([sample_sizes, numpy.ones(len(sample_sizes))]).T
         m, b = numpy.linalg.lstsq(x, sample_dets, rcond=None)[0]
         return (sample_sizes, sample_dets), (m, b)
-            
+
     @BaseGaussianNoise.normalize.setter
     def normalize(self, normalize):
         """Clears the current stats if the normalization state is changed.
@@ -226,11 +226,11 @@ def gate_indices(self, det):
         gt = gate_start + window
         lindex, rindex = ts.get_gate_indices(gt, window)
         return lindex, rindex
-    
+
     def det_lognorm(self, det, start_index=None, end_index=None):
         """Calculate the normalization term from the truncated covariance
         matrix.
-        
+
         Determinant is estimated using a linear fit to logdet vs truncated
         matrix size.
         """
@@ -414,14 +414,14 @@ def _get_gate_times(self, gatestart, gateend, ra, dec):
         Parameters
         ----------
         gatestart : float
-            Geocentric start time of the gate.
+            Start time of the gate.
         gateend : float
-            Geocentric end time of the gate.
+            End time of the gate.
         ra : float
             Right ascension of the signal.
         dec : float
             Declination of the signal.
-            
+
         Returns
         -------
         dict :
@@ -432,10 +432,9 @@ def _get_gate_times(self, gatestart, gateend, ra, dec):
             thisdet = Detector(det)
             # account for the time delay between the waveforms of the
             # different detectors
-            gatestartdelay = gatestart + thisdet.time_delay_from_earth_center(
-                ra, dec, gatestart)
-            gateenddelay = gateend + thisdet.time_delay_from_earth_center(
-                ra, dec, gateend)
+            refdet = self.current_params.get('tc_ref_frame', 'geocentric')
+            gatestartdelay = thisdet.arrival_time(gatestart, ra, dec, refdet)
+            gateenddelay = thisdet.arrival_time(gateend, ra, dec, refdet)
             dgatedelay = gateenddelay - gatestartdelay
             gatetimes[det] = (gatestartdelay, dgatedelay)
         return gatetimes
@@ -672,14 +671,31 @@ def _loglikelihood(self):
             rr = 4 * invpsd.delta_f * rtilde[slc].inner(gated_rtilde[slc]).real
             logl += norm - 0.5*rr
         return float(logl)
-
+
+    @property
+    def _extra_stats(self):
+        """Adds ``loglr``, plus ``cplx_loglr`` and ``optimal_snrsq`` in each
+        detector."""
+        return ['loglr', 'maxl_phase'] + ['{}_optimal_snrsq'.format(det) for det in self._data]
+
+    def _nowaveform_loglr(self):
+        """Convenience function to set loglr values if no waveform generated.
+        """
+        setattr(self._current_stats, 'loglikelihood', -numpy.inf)
+        # maxl phase doesn't exist, so set it to nan
+        setattr(self._current_stats, 'maxl_phase', numpy.nan)
+        for det in self._data:
+            # snr can't be < 0 by definition, so return 0
+            setattr(self._current_stats, '{}_optimal_snrsq'.format(det), 0.)
+        return -numpy.inf
+
     @property
     def multi_signal_support(self):
         """ The list of classes that this model supports in a multi-signal
         likelihood
         """
         return [type(self)]
-    
+
     def multi_loglikelihood(self, models):
         """ Calculate a multi-model (signal) likelihood
         """
@@ -797,15 +813,15 @@ def get_gated_waveforms(self):
             # the waveforms are a dictionary of (hp, hc)
             pols = []
             for h in wfs[det]:
-                ht = h.to_timeseries() 
+                ht = h.to_timeseries()
                 ht = ht.gate(gatestartdelay + dgatedelay/2,
                             window=dgatedelay/2, copy=False,
                             invpsd=invpsd, method='paint')
                 h = ht.to_frequencyseries()
                 pols.append(h)
             out[det] = tuple(pols)
         return out
-    
+
     def get_gate_times_hmeco(self):
         """Gets the time to apply a gate based on the current sky position.
         Returns
@@ -931,14 +947,14 @@ def _loglikelihood(self):
         # compute the marginalized log likelihood
         marglogl = special.logsumexp(loglr) + lognl - numpy.log(len(self.pol))
         return float(marglogl)
-    
+
     @property
     def multi_signal_support(self):
         """ The list of classes that this model supports in a multi-signal
         likelihood
         """
         return [type(self)]
-    
+
     @catch_waveform_error
     def multi_loglikelihood(self, models):
         """ Calculate a multi-model (signal) likelihood
@@ -965,3 +981,211 @@ def multi_loglikelihood(self, models):
 
         self._current_wfs = combine
         return self._loglikelihood()
+
+
+class GatedGaussianMargPhase(BaseGatedGaussian):
+    r"""Gated Gaussian noise model that analytically marginalizes over the
+    phase of a signal.
+
+    The phase to be marginalized over is specified by the user using the 
+    `ref_phase` argument. If a model consists of multiple modes each with their
+    own phase, only the reference phase is marginalized over. All phases must
+    be specified with the `phase_names` argument. This can be passed as a list
+    or a string delimited by spaces (e.g. 'phase1 phase2 phase3').
+
+    Marginalization is done using explicit numerical integration over 500
+    thousand integration points by default. This method assumes that the
+    waveform h can be written in terms of an overall phase phi as
+
+        h = h_c * cos(phi) + h_s * sin(phi),
+
+    where h_c and h_s are the waveform with phi set to zero and pi/2
+    respectively. The number of integration points can be controlled via the
+    `phase_samples` argument.
+    """
+    name = 'gated_gaussian_margphase'
+
+    def __init__(self, variable_params, data, low_frequency_cutoff, psds=None,
+                 high_frequency_cutoff=None, normalize=False,
+                 static_params=None, phase_samples=500000, phase_names=None,
+                 ref_phase=None, **kwargs):
+        # set up the boiler-plate attributes
+        super().__init__(
+            variable_params, data, low_frequency_cutoff, psds=psds,
+            high_frequency_cutoff=high_frequency_cutoff, normalize=normalize,
+            static_params=static_params, **kwargs)
+        self.det_names = list(self.data.keys())
+        self.dets = {}
+        # phase marginalization parameters
+        self.phase_samples = int(phase_samples)
+        self.phases = numpy.linspace(0, 2*numpy.pi, self.phase_samples)
+        if ref_phase is None:
+            raise KeyError('ref_phase is set to None. Please specify the '
+                           'name of the phase parameter to marginalize '
+                           'over')
+        self.ref_phase = ref_phase
+        if phase_names is None:
+            logging.warning('No phase_names provided. Assuming single mode '
+                            f'specified by ref_phase {ref_phase}')
+            self.phase_names = [ref_phase]
+        elif type(phase_names) == list:
+            self.phase_names = phase_names
+        elif type(phase_names) == str:
+            self.phase_names = phase_names.split(' ')
+        else:
+            raise TypeError('Unrecognized format for phase_names arg. Accepts '
+                            'string, list, or None')
+        # create the waveform generator
+        self.waveform_generator = create_waveform_generator(
+            self.variable_params, self.data,
+            waveform_transforms=self.waveform_transforms,
+            recalibration=self.recalibration,
+            generator_class=generator.FDomainDetFrameTwoPhaseGenerator,
+            **self.static_params)
+
+    def get_waveforms(self):
+        r"""Generate the waveforms.
+        """
+        if self._current_wfs is None:
+            params = self.current_params
+            # generate the cosine and sine terms
+            wfs = self.waveform_generator.generate(phases=self.phase_names, 
+                                                   ref_phase=self.ref_phase,
+                                                   **params)
+            for det, (hc, hs) in wfs.items():
+                # make the same length as the data
+                hc.resize(len(self.data[det]))
+                hs.resize(len(self.data[det]))
+                # apply high pass
+                if self.highpass_waveforms:
+                    hc = highpass(
+                         hc.to_timeseries(),
+                         frequency=self.highpass_waveforms).to_frequencyseries()
+                    hs = highpass(
+                         hs.to_timeseries(),
+                         frequency=self.highpass_waveforms).to_frequencyseries()
+                wfs[det] = (hc, hs)
+            self._current_wfs = wfs
+        return self._current_wfs
+
+    def get_gated_waveforms(self):
+        r"""Generate the gated waveforms.
+        """
+        wfs = self.get_waveforms()
+        out = {}
+        # apply the gate
+        for det, (hc, hs) in wfs.items():
+            hct = hc.to_timeseries()
+            hst = hs.to_timeseries()
+            invpsd = self._invpsds[det]
+            gate_times = self.get_gate_times()
+            gatestartdelay, dgatedelay = gate_times[det]
+            hct = hct.gate(gatestartdelay + dgatedelay/2,
+                           window=dgatedelay/2, copy=False,
+                           invpsd=invpsd, method='paint')
+            hst = hst.gate(gatestartdelay + dgatedelay/2,
+                           window=dgatedelay/2, copy=False,
+                           invpsd=invpsd, method='paint')
+            hc = hct.to_frequencyseries()
+            hs = hst.to_frequencyseries()
+            out[det] = (hc, hs)
+        return out
+
+    @property
+    def _extra_stats(self):
+        """Adds the maxL phase and corresponding likelihood."""
+        return ['maxl_phase', 'maxl_logl']
+
+    @catch_waveform_error
+    def _loglikelihood(self):
+        r"""Computes the log likelihood.
+        """
+        # get waveforms
+        wfs = self.get_waveforms()
+        gated_wfs = self.get_gated_waveforms()
+        # get data
+        data = self.get_data()
+        gated_data = self.get_gated_data()
+        # cycle over all detectors
+        norm = 0.
+        hchc = 0.
+        hchs = 0.
+        hshc = 0.
+        hshs = 0.
+        dhc = 0.
+        dhs = 0.
+        hcd = 0.
+        hsd = 0.
+        dd = 0.
+        for det in self.det_names:
+            if det not in self.dets:
+                self.dets[det] = Detector(det)
+            # we always filter the entire segment starting from kmin, since the
+            # gated series may have high frequency components
+            slc = slice(self._kmin[det], self._kmax[det])
+            invpsd = self._invpsds[det]
+            d = data[det].copy()
+            hc, hs = wfs[det]
+            gated_d = gated_data[det].copy()
+            gated_hc, gated_hs = gated_wfs[det]
+            # overwhiten gated waveforms and data
+            gated_hc *= 2 * invpsd.delta_f * invpsd
+            gated_hs *= 2 * invpsd.delta_f * invpsd
+            gated_d *= 2 * invpsd.delta_f * invpsd
+            # evaluate the inner products
+            hchc += hc[slc].inner(gated_hc[slc]).real
+            hchs += hc[slc].inner(gated_hs[slc]).real
+            hshc += hs[slc].inner(gated_hc[slc]).real
+            hshs += hs[slc].inner(gated_hs[slc]).real
+            dhc += d[slc].inner(gated_hc[slc]).real
+            dhs += d[slc].inner(gated_hs[slc]).real
+            hcd += hc[slc].inner(gated_d[slc]).real
+            hsd += hs[slc].inner(gated_d[slc]).real
+            dd += d[slc].inner(gated_d[slc]).real
+            # get the normalization in this detector
+            if self.normalize:
+                start_index, end_index = self.gate_indices(det)
+            else:
+                start_index = end_index = None
+            norm += self.det_lognorm(det, start_index, end_index)
+        # numerical marginalization over phases
+        cphi = numpy.cos(self.phases)
+        sphi = numpy.sin(self.phases)
+        hh = cphi*cphi*hchc + sphi*sphi*hshs + cphi*sphi*(hchs+hshc)
+        dh = cphi*dhc + sphi*dhs
+        hd = cphi*hcd + sphi*hsd
+        loglr = -(hh-dh-hd)
+        lognl = -dd
+        # get the maxL phase
+        maxlidx = loglr.argmax()
+        setattr(self._current_stats, 'maxl_phase', self.phases[maxlidx])
+        setattr(self._current_stats, 'maxl_logl', loglr[maxlidx] + lognl + norm)
+        # get the marginalized log likelihood ratio
+        marglogl = special.logsumexp(loglr) + lognl + norm - numpy.log(self.phase_samples)
+        return marglogl
+
+    @property
+    def multi_signal_support(self):
+        """ The list of classes that this model supports in a multi-signal
+        likelihood
+        """
+        return [type(self)]
+
+    @catch_waveform_error
+    def multi_loglikelihood(self, models):
+        """ Calculate a multi-model (signal) likelihood
+        """
+        # Generate the waveforms for each submodel
+        wfs = []
+        for m in models + [self]:
+            wf = m.get_waveforms()
+            wfs.append(wf)
+        # combine into a single waveform
+        combine = {}
+        for det in self.data:
+            # get max waveform length
+            mlen = max([len(x[det]) for x in wfs])
+            [x[det].resize(mlen) for x in wfs]
+            combine[det] = sum([x[det] for x in wfs])
+        self._current_wfs = combine
+        return self._loglikelihood()