Integrate harmonic mode selection and waveform plugin support

pycbc/waveform/__init__.py

@@ -7,5 +7,6 @@
                                            get_fd_waveform_modes)
 from pycbc.waveform.plugin import (retrieve_waveform_plugins,
                                    add_custom_waveform,
+                                   add_custom_waveform_modes as add_custom_waveform_modes,
                                    add_length_estimator)
 retrieve_waveform_plugins()

pycbc/waveform/plugin.py

@@ -57,6 +57,52 @@ def add_custom_waveform(approximant, function, domain,
                          "'time' or 'frequency'".format(domain))
 
 
+def add_custom_waveform_modes(approximant, function, domain, force=False):
+    """Make a custom mode-by-mode waveform generator available to pycbc's
+    get_td_waveform_modes/get_fd_waveform_modes.
+
+    Unlike add_custom_waveform (which registers a generator returning the already-summed
+    (h_+, h_x) polarizations), this registers a generator returning a pair of dicts
+    (ulm, vlm), each mapping a mode label -> TimeSeries/FrequencySeries, following the
+    convention documented in pycbc.waveform.waveform_modes.get_td_waveform_modes /
+    get_fd_waveform_modes. The mode label is usually an (l, m) tuple, but plugins whose
+    modes carry an extra harmonic index (e.g. eccentric sub-harmonics; see the note in
+    parse_mode_array's docstring) may use longer tuples instead -- pycbc itself does not
+    interpret the label, it is only used as a dict key.
+
+    Parameters
+    ----------
+    approximant : str
+        The name of the waveform
+    function : function
+        The function to generate the modes. Must accept the same keyword arguments as
+        get_td_waveform_modes/get_fd_waveform_modes and return (ulm, vlm) as described above.
+    domain : str
+        Either 'frequency' or 'time' to indicate the domain of the waveform.
+    force : bool, False
+        Overwrite an existing registration for this approximant/domain instead of raising.
+    """
+    from pycbc.waveform.waveform_modes import _mode_waveform_fd, _mode_waveform_td
+
+    used = RuntimeError(
+        "Can't load plugin waveform modes generator {}, the name is"
+        " already in use.".format(approximant)
+    )
+
+    if domain == "time":
+        if not force and (approximant in _mode_waveform_td):
+            raise used
+        _mode_waveform_td[approximant] = function
+    elif domain == "frequency":
+        if not force and (approximant in _mode_waveform_fd):
+            raise used
+        _mode_waveform_fd[approximant] = function
+    else:
+        raise ValueError(
+            "Invalid domain ({}), should be 'time' or 'frequency'".format(domain)
+        )
+
+
 def add_length_estimator(approximant, function):
     """ Add length estimator for an approximant
 
@@ -123,6 +169,14 @@ def retrieve_waveform_plugins():
     for plugin in entry_points(group='pycbc.waveform.td'):
         add_custom_waveform(plugin.name, plugin.load(), 'time')
 
+    # Check for mode-by-mode fd waveforms (feed get_fd_waveform_modes)
+    for plugin in entry_points(group='pycbc.waveform.fd_modes'):
+        add_custom_waveform_modes(plugin.name, plugin.load(), 'frequency')
+
+    # Check for mode-by-mode td waveforms (feed get_td_waveform_modes)
+    for plugin in entry_points(group='pycbc.waveform.td_modes'):
+        add_custom_waveform_modes(plugin.name, plugin.load(), 'time')
+
     # Check for waveform length estimates
     for plugin in entry_points(group='pycbc.waveform.length'):
         add_length_estimator(plugin.name, plugin.load())

pycbc/waveform/waveform.py

@@ -129,6 +129,18 @@ def _check_lal_pars(p):
     if p['side_bands']:
         lalsimulation.SimInspiralWaveformParamsInsertSideband(lal_pars, p['side_bands'])
     if p['mode_array'] is not None:
+        # LAL's mode array only understands (l, m) pairs: it has no concept of the eccentric/sub-harmonic
+        # index n that some (non-LAL) plugin waveforms use to sub-divide a mode further, see parse_mode_array.
+        # Most LAL approximants have no eccentricity content at all (at most higher-order modes), so fail
+        # loudly here instead of silently dropping n or letting the "for l, m in ..." unpack raise a
+        # confusing ValueError.
+        bad = [entry for entry in p["mode_array"] if len(entry) != 2]
+        if bad:
+            raise ValueError(
+                "mode_array entries %s have more than (l, m); this LAL-based approximant only "
+                "supports selecting modes by (l, m), not by an additional harmonic index n. "
+                "Pass a plain (l, m) mode_array for this approximant." % (bad,)
+            )
         ma = lalsimulation.SimInspiralCreateModeArray()
         for l,m in p['mode_array']:
             lalsimulation.SimInspiralModeArrayActivateMode(ma, l, m)
@@ -400,6 +412,14 @@ def parse_mode_array(input_params):
     ints (e.g., ``[(2, 2), (3, 3), (4, 4)]``), a space-separated string giving
     the modes (e.g., ``22 33 44``), or an array of ints or floats (e.g.,
     ``[22., 33., 44.]``.
+
+    Some (non-LAL) mode-by-mode plugin waveforms further sub-divide a mode by an integer harmonic
+    index n (e.g. eccentric waveforms, where a given (l, m) multipole has contributions at several
+    harmonics of the orbital frequency). For those, ``mode_array`` entries may instead be 3-tuples
+    ``(l, m, n)``; these are passed through unchanged (only the string/scalar shorthand above is
+    restricted to plain (l, m)). Approximants that do not support this extra index (which is most of
+    them; at most LAL approximants support higher-order modes, not eccentric sub-harmonics) will raise
+    a clear error if given a 3-tuple, rather than silently ignoring n or failing an unpack.
     """
     if 'mode_array' in input_params and input_params['mode_array'] is not None:
         mode_array = input_params['mode_array']
@@ -1081,7 +1101,8 @@ def seobnrv4hm_length_in_time(**kwargs):
 def get_hm_length_in_time(lor_approx, maxm_default, **kwargs):
     kwargs = parse_mode_array(kwargs)
     if 'mode_array' in kwargs and kwargs['mode_array'] is not None:
-        maxm = max(m for _, m in kwargs['mode_array'])
+        # entries may be (l, m) or (l, m, n) (see parse_mode_array); m is always the second element
+        maxm = max(entry[1] for entry in kwargs["mode_array"])
     else:
         maxm = maxm_default
     try:

pycbc/waveform/waveform_modes.py

@@ -325,6 +325,16 @@ def get_imrphenomxh_modes(**params):
     mode_array = params.pop('mode_array', None)
     if mode_array is None:
         mode_array = default_modes(approx)
+    else:
+        # IMRPhenomXHM has no eccentric content, so it only knows how to select modes by (l, m);
+        # see the note on the extra harmonic index n in parse_mode_array's docstring.
+        bad = [entry for entry in mode_array if len(entry) != 2]
+        if bad:
+            raise ValueError(
+                "mode_array entries %s have more than (l, m); %s only supports selecting "
+                "modes by (l, m), not by an additional harmonic index n."
+                % (bad, approx)
+            )
     if 'f_final' not in params:
         # setting to 0 will default to ringdown frequency
         params['f_final'] = 0.