Add support for explicit mode-by-mode waveform generators

pycbc/inference/models/gated_gaussian_noise.py

@@ -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
         """
@@ -805,7 +821,7 @@ def get_gated_waveforms(self):
                 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
@@ -964,4 +980,4 @@ def multi_loglikelihood(self, models):
                                  for x in wfs]))
 
         self._current_wfs = combine
-        return self._loglikelihood()
+        return self._loglikelihood()

pycbc/waveform/generator.py

@@ -343,7 +343,13 @@ class FDomainMassSpinRingdownGenerator(BaseGenerator):
 
     """
     def __init__(self, variable_args=(), **frozen_params):
-        super(FDomainMassSpinRingdownGenerator, self).__init__(ringdown.get_fd_from_final_mass_spin,
+        if frozen_params['approximant'] == 'FdQNMfromFinalMassSpin':
+            approximant = ringdown.get_fd_from_final_mass_spin
+        elif frozen_params['approximant'] == 'FdModesfromFinalMassSpin':
+            approximant = ringdown.get_fd_modes_from_final_mass_spin
+        else:
+            raise ValueError(f"Invalid approximant name: {frozen_params['approximant']}")
+        super(FDomainMassSpinRingdownGenerator, self).__init__(approximant,
             variable_args=variable_args, **frozen_params)
 
 
@@ -371,7 +377,13 @@ class FDomainFreqTauRingdownGenerator(BaseGenerator):
 
     """
     def __init__(self, variable_args=(), **frozen_params):
-        super(FDomainFreqTauRingdownGenerator, self).__init__(ringdown.get_fd_from_freqtau,
+        if frozen_params['approximant'] == 'FdQNMfromFreqTau':
+            approximant = ringdown.get_fd_from_freqtau
+        elif frozen_params['approximant'] == 'FdModesfromFreqTau':
+            approximant = ringdown.get_fd_modes_from_freqtau
+        else:
+            raise ValueError(f"Invalid approximant name: {frozen_params['approximant']}")
+        super(FDomainFreqTauRingdownGenerator, self).__init__(approximant,
             variable_args=variable_args, **frozen_params)
 
 
@@ -399,7 +411,13 @@ class TDomainMassSpinRingdownGenerator(BaseGenerator):
 
     """
     def __init__(self, variable_args=(), **frozen_params):
-        super(TDomainMassSpinRingdownGenerator, self).__init__(ringdown.get_td_from_final_mass_spin,
+        if frozen_params['approximant'] == 'TdQNMfromFinalMassSpin':
+            approximant = ringdown.get_td_from_final_mass_spin
+        elif frozen_params['approximant'] == 'TdModesfromFinalMassSpin':
+            approximant = ringdown.get_td_modes_from_final_mass_spin
+        else:
+            raise ValueError(f"Invalid approximant name: {frozen_params['approximant']}")
+        super(TDomainMassSpinRingdownGenerator, self).__init__(approximant,
             variable_args=variable_args, **frozen_params)
 
 
@@ -427,7 +445,13 @@ class TDomainFreqTauRingdownGenerator(BaseGenerator):
 
     """
     def __init__(self, variable_args=(), **frozen_params):
-        super(TDomainFreqTauRingdownGenerator, self).__init__(ringdown.get_td_from_freqtau,
+        if frozen_params['approximant'] == 'TdQNMfromFreqTau':
+            approximant = ringdown.get_td_from_freqtau
+        elif frozen_params['approximant'] == 'TdModesfromFreqTau':
+            approximant = ringdown.get_td_modes_from_freqtau
+        else:
+            raise ValueError(f"Invalid approximant name: {frozen_params['approximant']}")
+        super(TDomainFreqTauRingdownGenerator, self).__init__(approximant,
             variable_args=variable_args, **frozen_params)
 
 
@@ -1220,32 +1244,38 @@ def get_td_generator(approximant, modes=False):
         if modes:
             return TDomainCBCModesGenerator
         return TDomainCBCGenerator
+
+    if approximant in waveform_modes._mode_waveform_td:
+        return TDomainCBCModesGenerator
 
     if approximant in ringdown.ringdown_td_approximants:
-        if approximant == 'TdQNMfromFinalMassSpin':
+        if approximant in ['TdQNMfromFinalMassSpin', 'TdModesfromFinalMassSpin']:
             return TDomainMassSpinRingdownGenerator
         return TDomainFreqTauRingdownGenerator
 
     if approximant in supernovae.supernovae_td_approximants:
         return TDomainSupernovaeGenerator
 
-    raise ValueError(f"No time-domain generator found for "
-                       "approximant: {approximant}")
+    raise ValueError(f"No time-domain generator found for" 
+                     "approximant: {approximant}")
 
 def get_fd_generator(approximant, modes=False):
     """Returns the frequency-domain generator for the given approximant."""
     if approximant in waveform.fd_approximants():
         if modes:
             return FDomainCBCModesGenerator
         return FDomainCBCGenerator
+
+    if approximant in waveform_modes._mode_waveform_fd:
+        return FDomainCBCModesGenerator
 
     if approximant in ringdown.ringdown_fd_approximants:
-        if approximant == 'FdQNMfromFinalMassSpin':
+        if approximant in ['FdQNMfromFinalMassSpin', 'FdModesfromFinalMassSpin']:
             return FDomainMassSpinRingdownGenerator
         return FDomainFreqTauRingdownGenerator
 
-    raise ValueError(f"No frequency-domain generator found for "
-                       "approximant: {approximant}")
+    raise ValueError(f"No frequency-domain generator found for"
+                     "approximant: {approximant}")
 
 def select_waveform_generator(approximant, domain=None):
     """Returns the single-IFO generator for the approximant.

pycbc/waveform/plugin.py

@@ -3,7 +3,7 @@
 
 
 def add_custom_waveform(approximant, function, domain,
-                        sequence=False, has_det_response=False,
+                        sequence=False, has_det_response=False, modes=False,
                         force=False,):
     """ Make custom waveform available to pycbc
 
@@ -23,14 +23,18 @@ def add_custom_waveform(approximant, function, domain,
     """
     from pycbc.waveform.waveform import (cpu_fd, cpu_td, fd_sequence,
                                          fd_det, fd_det_sequence)
+    from pycbc.waveform.waveform_modes import _mode_waveform_td
 
     used = RuntimeError("Can't load plugin waveform {}, the name is"
                         " already in use.".format(approximant))
 
     if domain == 'time':
         if not force and (approximant in cpu_td):
             raise used
-        cpu_td[approximant] = function
+        if modes:
+            _mode_waveform_td[approximant] = function
+        else:
+            cpu_td[approximant] = function
     elif domain == 'frequency':
         if sequence:
             if not has_det_response:
@@ -121,6 +125,11 @@ def retrieve_waveform_plugins():
     for plugin in entry_points(group='pycbc.waveform.td'):
         add_custom_waveform(plugin.name, plugin.load(), 'time')
 
+    # Check for td modal waveforms
+    for plugin in entry_points(group='pycbc.waveform.td_modes'):
+        add_custom_waveform(plugin.name, plugin.load(), 'time',
+                            modes=True)
+
     # Check for waveform length estimates
     for plugin in entry_points(group='pycbc.waveform.length'):
         add_length_estimator(plugin.name, plugin.load())