Add support for SEOBNRv5HM and SEOBNRv5PHM as PMR piece (#10)

* Add support for SEOBNRv5HM and SEOBNRv5PHM for PMR piece
  - added support for using SEOBNRv5HM and SEOBNRv5PHM as
  the PMR piece in ESIGMA and ESIGMASur via pyseobnr
  - pyseobnr kept as an optional dependency; code
  uses lazy imports to avoid import errors if pyseobnr
  waveform models aren't used and pyseobnr is not installed
  - added exception handling for MR generation failure
  - Update README for additional merger-ringdown choices
  - Update ESIGMA usage notebook

---------

Co-authored-by: Prayush Kumar <prayush@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: 3 people

README.md

@@ -20,7 +20,7 @@ Then to use:
  * `ESIGMAHM`: requires custom `LALSuite` installation and `NRSur7dq4` surrogate data download ([installation instructions below](https://github.com/gwnrtools/esigmapy/tree/master?tab=readme-ov-file#blue_square-esigmahm)).
  * `ESIGMASur`: requires downloading `ESIGMASur` surrogate data ([installation instructions below](https://github.com/gwnrtools/esigmapy/tree/master?tab=readme-ov-file#green_square-esigmasur))
 
-Usage instructions at:
+Usage instructions:
  * `ESIGMAHM` [tutorial notebook](https://github.com/gwnrtools/esigmapy/blob/master/notebooks/ESIGMA_tutorial.ipynb)
  * `ESIGMASur` [tutorial notebook](https://github.com/gwnrtools/esigmapy/blob/master/notebooks/ESIGMASur_tutorial.ipynb)
 ***
@@ -31,7 +31,7 @@ Usage instructions at:
 * **Inspiral piece (`InspiralESIGMAHM`):** It comes from a combination of post-Newtonian theory, self-force, and black hole perturbation theory.
 * **Plunge-merger-ringdown piece**: Assuming moderate starting eccentricities that decay by the late inspiral, we use the quasi-circular (QC) NR surrogate `NRSur7dq4` for the plunge-merger-ringdown piece for `ESIGMAHM`.
 
-  (**Note:** We also allow other QC `LALSuite` waveform models to be used as the plunge-merger-ringdown piece; see the optional argument `merger_ringdown_approximant` in the generation functions `get_imr_esigma_waveform` and `get_imr_esigma_modes` [here](https://github.com/gwnrtools/esigmapy/blob/master/esigmapy/generator.py). However, the default and the most tested choice is `NRSur7dq4`.)
+  (**Note:** We also allow other QC waveform models to be used as the plunge-merger-ringdown piece; see the optional argument `merger_ringdown_approximant` in the generation functions `get_imr_esigma_waveform` and `get_imr_esigma_modes` [here](https://github.com/gwnrtools/esigmapy/blob/master/esigmapy/generator.py). The models currently supported are: `NRSur7dq4`, `SEOBNRv4PHM` (via `LALSuite`), `SEOBNRv5HM`, `SEOBNRv5PHM` (via `pyseobnr`). However, the default and the most tested choice is `NRSur7dq4`.)
 
 The full IMR waveform `ESIGMAHM` is produced by smoothly attaching the inspiral piece `InspiralESIGMAHM` to the plunge-merger-ringdown piece `NRSur7dq4`. This attachment is facilitated by `ESIGMAPy`.
 
@@ -73,8 +73,8 @@ Using `ESIGMAHM` therefore requires installing 1. `InspiralESIGMAHM`, 2. `NRSur7
 * To avoid performing the above step in every new terminal session, either add the above command to your `.bashrc` file, or follow the instructions [here](http://gitlab.icts.res.in/akash.maurya/Installation-instructions/wikis/conda-tricks), replacing `PYTHONPATH` with `LAL_DATA_PATH`, to set this environment variable automatically on activating your `conda` environment.
 
 ### Installing `ESIGMAPy`
-* Activate your `conda` environment and install `ESIGMAPy` by running: `pip install esigmapy`.
-
+* Activate your `conda` environment and install `ESIGMAPy` by running: `pip install git+https://github.com/gwnrtools/esigmapy.git`.
+* (Optional) If you want to use `SEOBNRv5HM` or `SEOBNRv5PHM` as the merger-ringdown piece in `ESIGMAHM`, install `ESIGMAPy` with `pyseobnr` option by running: `pip install "esigmapy[pyseobnr] @ git+https://github.com/gwnrtools/esigmapy.git"`
 ***
 ### Trying out `ESIGMAHM`
 The instructions to generate `ESIGMAHM` waveforms and the various functionalities that it offers are detailed in [this tutorial notebook](https://github.com/gwnrtools/esigmapy/blob/master/notebooks/ESIGMA_tutorial.ipynb). 
@@ -99,7 +99,7 @@ The instructions to generate `ESIGMAHM` waveforms and the various functionalitie
     
 * Download the surrogate data files of `ESIGMASur`, which can be found [on the repo](https://github.com/gwnrtools/esigmapy/tree/master/esigmapy/surrogate/data). Next, set the shell environment variable `ESIGMASUR_DATA_PATH` to the directory where you keep these surrogate data files by running: `export ESIGMASUR_DATA_PATH="/path/to/ESIGMASur"`.
 
-* Using the inspiral-only surrogate `InspiralESIGMASur` would not require any further dependencies. However, its hybridized IMR version `IMRESIGMASur` will require downloading the NR surrogate data file ([installation instructions above](https://github.com/gwnrtools/esigmapy/tree/master#installing-nrsur7dq4)). 
+* Using the inspiral-only surrogate `InspiralESIGMASur` would not require any further dependencies. However, its hybridized IMR version `IMRESIGMASur` will require downloading the NR surrogate data file ([installation instructions above](https://github.com/gwnrtools/esigmapy/tree/master#installing-nrsur7dq4)). Or, if you want to use `SEOBNRv5HM` or `SEOBNRv5PHM` as the merger-ringdown piece in `IMRESIGMASur`, add `pyseobnr` option during the installation: `pip install "esigmapy[surrogate,pyseobnr] @ git+https://github.com/gwnrtools/esigmapy.git"`.
 
 ### Trying out `ESIGMASur`
 The usage instructions and the various functionalities of `ESIGMASur` are detailed in [this tutorial notebook](https://github.com/gwnrtools/esigmapy/blob/master/notebooks/ESIGMASur_tutorial.ipynb).

esigmapy/__init__.py

@@ -1,6 +1,6 @@
 from __future__ import absolute_import
 
-from . import blend, generator, legacy, utils
+from . import blend, generator, mr_generator, legacy, utils
 from .generator import *
 
 
@@ -19,23 +19,32 @@ def get_version_information():
 
 # pycbc wrapper
 import inspect
+
 _params = inspect.signature(get_imr_esigma_waveform).parameters.keys()
+
+
 def pycbc_esigma(**params):
     from pycbc.waveform.waveform import parse_mode_array
 
     # Pass all parameters the model suppports
     pt = {p: params[p] for p in _params if p in params}
-    pt['f_ref'] = pt['f_lower'] if pt['f_ref'] == 0 else pt['f_ref']
+    pt["f_ref"] = pt["f_lower"] if pt["f_ref"] == 0 else pt["f_ref"]
 
     gen_wav = get_imr_esigma_waveform
-    if 'skip_merger' in params:
+    if "skip_merger" in params:
         gen_wav = get_inspiral_esigma_waveform
 
-    modes = params.pop('mode_array', [(2,2), (2,-2)])
-    modes = parse_mode_array({'mode_array':modes})['mode_array']
+    modes = params.pop("mode_array", [(2, 2), (2, -2)])
+    modes = parse_mode_array({"mode_array": modes})["mode_array"]
+
+    return gen_wav(
+        pt.pop("mass1"),
+        pt.pop("mass2"),
+        pt.pop("f_lower"),
+        pt.pop("delta_t"),
+        **pt,
+        modes_to_use=modes
+    )
 
-    return gen_wav(pt.pop('mass1'), pt.pop('mass2'),
-                   pt.pop('f_lower'), pt.pop('delta_t'),
-                   **pt, modes_to_use=modes)
 
 __version__ = get_version_information()

esigmapy/generator.py

@@ -11,6 +11,7 @@
 import lalsimulation as ls
 import pycbc.types as pt
 from .utils import f_ISCO_spin
+from .mr_generator import check_available_mr_approximants, get_mr_modes
 
 ECCENTRICITY_LEVEL_ISCO_WARNING = 0.02
 ECCENTRICITY_LEVEL_ISCO_ERROR = 0.1
@@ -45,7 +46,7 @@ def find_x_for_y(x, y, y0):
     t.data.data *= lal.MTSUN_SI
 
     if verbose:
-        print(f"Orbital evolution took: {time.perf_counter() - itime} seconds")
+        print(f"Inspiral orbital evolution took: {time.perf_counter() - itime} seconds")
 
     omega = pt.TimeSeries(phidot.data.data, delta_t=t.data.data[1] - t.data.data[0])
     if extremum == "periastron":
@@ -131,7 +132,7 @@ def eccentricity_at_reference_frequency(
     t.data.data *= lal.MTSUN_SI
 
     if verbose:
-        print(f"Orbital evolution took: {time.perf_counter() - itime} seconds")
+        print(f"Inspiral orbital evolution took: {time.perf_counter() - itime} seconds")
 
     x_reference = (np.pi * (mass1 + mass2) * lal.MTSUN_SI * f_reference) ** (2.0 / 3.0)
 
@@ -189,7 +190,8 @@ def get_inspiral_esigma_modes(
                                    parameters.
                                    We require f_ref <= f_lower. f_ref = f_lower by default.
         delta_t                 -- Waveform's time grid-spacing (in s)
-        spin1z, spin2z          -- z-components of component dimensionless spins (lies in [0,1))
+        spin1z, spin2z          -- z-components of component dimensionless spins
+                                   (lies in (-1, 1))
         eccentricity            -- Initial eccentricity
         mean_anomaly            -- Mean anomaly of the periastron (in rad)
         distance                -- Luminosity distance to the binary (in Mpc)
@@ -278,7 +280,7 @@ def get_inspiral_esigma_modes(
     ) * lal.MTSUN_SI  # Time from geometrized units to seconds
 
     if verbose:
-        print(f"Orbital evolution took: {time.perf_counter() - itime} seconds")
+        print(f"Inspiral orbital evolution took: {time.perf_counter() - itime} seconds")
 
     # Include conjugate modes in the mode list
     if include_conjugate_modes:
@@ -319,7 +321,7 @@ def get_inspiral_esigma_modes(
         modes = {k: np.asarray(modes[k].data.data) for k in modes}
 
     if verbose:
-        print(f"Modes generation took: {time.perf_counter() - itime} seconds")
+        print(f"Inspiral modes generation took: {time.perf_counter() - itime} seconds")
 
     if return_orbital_params:
         orbital_var_dict = {}
@@ -372,12 +374,13 @@ def get_inspiral_esigma_waveform(
                                    parameters.
                                    We require f_ref <= f_lower. f_ref = f_lower by default.
         delta_t                 -- Waveform's time grid-spacing (in s)
-        spin1z, spin2z          -- z-components of component dimensionless spins (lies in [0,1))
+        spin1z, spin2z          -- z-components of component dimensionless spins
+                                   (lies in (-1, 1))
         eccentricity            -- Initial eccentricity
         mean_anomaly            -- Mean anomaly of the periastron (in rad)
         inclination             -- Inclination (in rad), defined as the angle between the orbital
                                    angular momentum L and the line-of-sight
-        coa_phase               -- Coalesence phase of the binary (in rad)
+        coa_phase               -- Coalescence phase of the binary (in rad)
         distance                -- Luminosity distance to the binary (in Mpc)
         modes_to_use            -- GW modes to use. List of tuples (l, |m|)
         return_orbital_params   -- If True, returns the orbital evolution of all the orbital elements (in
@@ -675,7 +678,7 @@ def get_imr_esigma_modes(
 ):
     """
     Returns IMR GW modes constructed using ESIGMA for inspiral and
-    NRSur7dq4/SEOBNRv4PHM for merger-ringdown
+    NRSur7dq4/SEOBNRv4PHM/SEOBNRv5HM/SEOBNRv5PHM for merger-ringdown
 
     Parameters:
     -----------
@@ -687,10 +690,10 @@ def get_imr_esigma_modes(
                                      f_ref = f_lower by default.
         delta_t                   -- Waveform's time grid-spacing (in s)
         spin1z, spin2z            -- z-components of component dimensionless
-                                     spins (lies in [0,1))
+                                     spins (lies in (-1,1))
         eccentricity              -- Initial eccentricity
         mean_anomaly              -- Mean anomaly of the periastron (radians)
-        coa_phase                 -- Coalesence phase of the binary (in rad)
+        coa_phase                 -- Coalescence phase of the binary (in rad)
         distance                  -- Luminosity distance to the binary (in Mpc)
         modes_to_use              -- GW modes to use. List of tuples (l, |m|)
         mode_to_align_by          -- GW mode to use to align inspiral and merger
@@ -731,8 +734,10 @@ def get_imr_esigma_modes(
                                      the center of the hybridization window.
                                      Otherwise, it's kept at the end of the
                                      window (default).
-        merger_ringdown_approximant    -- Choose merger-ringdown model. Tested
-                                     choices: [NRSur7dq4, SEOBNRv4PHM]
+        merger_ringdown_approximant    -- Choose merger-ringdown model.
+                                        Available choices:
+                                        NRSur7dq4, SEOBNRv4PHM  (requires `lalsimulation`)
+                                        SEOBNRv5HM, SEOBNRv5PHM (requires `pyseobnr`)
         return_hybridization_info -- If True, returns hybridization related data
         return_orbital_params     -- If True, returns the orbital evolution of
                                      all the orbital elements (in
@@ -757,11 +762,7 @@ def get_imr_esigma_modes(
         retval            -- Hybridization related data. Returned only if the
                              flag `return_hybridization_info` is set
     """
-    if not hasattr(ls, merger_ringdown_approximant):
-        raise IOError(
-            """We cannot generate individual modes for {merger_ringdown_approximant}.
-                      Try one of: [NRSur7dq4, SEOBNRv4PHM]"""
-        )
+    check_available_mr_approximants(merger_ringdown_approximant)
     if (mean_anomaly is None) and (coa_phase is None):
         raise IOError(
             f"""Please specify one of the phase angles, either of
@@ -965,47 +966,42 @@ def get_imr_esigma_modes(
     if not keep_f_mr_transition_at_center:
         f_mr_transition -= f_window_mr_transition / 2.0
 
-    # Generate NR surrogate waveform that will be our merger-ringdown, starting
+    # Generate plunge-merger-ringdown waveform, starting
     # from a frequency = 90% of
     max_retries = 20
     f_lower_mr = (f_mr_transition - f_window_mr_transition / 2) * (
         1.8 / mode_to_align_by_em
     )
+    modes_to_use = list(modes_inspiral_numpy.keys())
     for _ in range(max_retries):
         try:
             if verbose:
                 print(f"Generating MR waveform from {f_lower_mr}Hz...")
-            hlm_mr = ls.SimInspiralChooseTDModes(
-                coa_phase,  # phiRef
-                delta_t,  # deltaT
-                mass1 * lal.MSUN_SI,
-                mass2 * lal.MSUN_SI,
-                0,  # spin1x
-                0,  # spin1y
-                spin1z,
-                0,  # spin2x
-                0,  # spin2y
-                spin2z,
-                f_lower_mr,  # f_min
-                f_lower_mr,  # f_ref
-                distance * lal.PC_SI * 1.0e6,
-                None,  # LALpars
-                4,  # lmax
-                getattr(ls, merger_ringdown_approximant),
+            modes_mr_numpy = get_mr_modes(
+                mass1=mass1,
+                mass2=mass2,
+                f_lower=f_lower_mr,
+                f_ref=f_lower_mr,
+                delta_t=delta_t,
+                spin1z=spin1z,
+                spin2z=spin2z,
+                coa_phase=coa_phase,
+                distance=distance,
+                modes_to_use=modes_to_use,
+                approximant=merger_ringdown_approximant,
+                verbose=verbose,
             )
             break
         except:
             f_lower_mr *= 0.8
             continue
-    # Extracting only the modes we need
-    modes_to_use = list(modes_inspiral_numpy.keys())
-    modes_mr_numpy = {}
-    while hlm_mr is not None:
-        key = (hlm_mr.l, hlm_mr.m)
-        if key in modes_to_use:
-            modes_mr_numpy[key] = hlm_mr.mode.data.data
-        hlm_mr = hlm_mr.next
-
+    # else clause in a for-else block executes only if the
+    # for-loop is not terminated by a break statement
+    else:
+        raise RuntimeError(
+            f"""Failed to generate merger-ringdown waveform after {max_retries} retries.
+Last f_lower tried: {f_lower_mr/0.8:.4f}Hz."""
+        )
     try:
         retval = esigmapy.blend.blend_modes(
             modes_inspiral_numpy,
@@ -1103,13 +1099,13 @@ def get_imr_esigma_waveform(
                                      `f_ref = f_lower` by default.
         delta_t                   -- Waveform's time grid-spacing (in s)
         spin1z, spin2z            -- z-components of component dimensionless
-                                     spins (lies in [0,1))
+                                     spins (lies in (-1,1))
         eccentricity              -- Initial eccentricity
         mean_anomaly              -- Mean anomaly of the periastron (in rad)
         inclination               -- Inclination (in rad), defined as the angle
                                      between the orbital angular momentum L and
                                      the line-of-sight
-        coa_phase                 -- Coalesence phase of the binary (in rad)
+        coa_phase                 -- Coalescence phase of the binary (in rad)
         distance                  -- Luminosity distance to the binary (in Mpc)
         modes_to_use              -- GW modes to use. List of tuples (l, |m|)
         mode_to_align_by          -- GW mode to use to align inspiral and merger
@@ -1139,8 +1135,10 @@ def get_imr_esigma_waveform(
                                      the center of the hybridization window.
                                      Otherwise, it's kept at the end of the
                                      window (default).
-        merger_ringdown_approximant    -- Choose merger-ringdown model. Tested
-                                     choices: [NRSur7dq4, SEOBNRv4PHM]
+        merger_ringdown_approximant    -- Choose merger-ringdown model.
+                                    Available choices:
+                                    NRSur7dq4, SEOBNRv4PHM  (requires `lalsimulation`)
+                                    SEOBNRv5HM, SEOBNRv5PHM (requires `pyseobnr`)
         return_hybridization_info -- If True, returns hybridization related data
         return_orbital_params     -- If True, returns the orbital evolution of
                                      all the orbital elements (in