Skip to content

nrcatalogtools.surrogate

NRSur7dq4 loading, evaluation, and prior check utilities.

gwsurrogate is an optional dependency. Importing this module succeeds even when gwsurrogate is not installed; the ImportError is raised only on the first call to load_nrsur7dq4() or generate_surrogate_modes().

Conceptual background: See Package Internals § 11 for notes on the f_ref clipping strategy and spin-epoch conventions.


Functions

surrogate

Surrogate waveform utilities for NR vs NRSur7dq4 comparisons.

Handles loading the surrogate, calling it with NR-compatible parameters, and wrapping the output as physical-unit pycbc TimeSeries objects.

NRSur7dq4 notes
  • Precessing model — takes full 3-vector spins, returns all modes up to ell=4.
  • The ellMax parameter caps the output (max is 4); mode_list is not supported for precessing models.
  • f_low controls waveform truncation only (start frequency); for NRSur7dq4 the recommended value is 0 (return the entire waveform).
  • f_ref sets the reference epoch at which the spins are defined; it is in cycles/M (= M * f_GW where M is in seconds). For aligned-spin or non-spinning systems f_ref = f_lower_NR * M_seconds is sufficient. For precessing SXS systems Phase 2 extracts the instantaneous spins from Horizons.h5 at the epoch corresponding to the chosen f_ref.
  • dt argument is in dimensionless M units.
  • Output h[(ell,em)] is a complex numpy array representing the spin-weight -2 spherical harmonic mode h_lm (same convention as SXS/RIT/MAYA NR data: h_+ - i h_× decomposed as Σ h_lm Y_{-2,lm}).
  • Time array t_sur is in dimensionless M units; t_sur[-1] ≈ +100 M is near merger (peak amplitude).
Spin epoch conventions

Two cases arise depending on whether the NR waveform starts before or after the surrogate's minimum training frequency (~0.0165 M·Ω_orbital):

  1. NR shorter than surrogate (f_lower_NR > f_min_sur): Pass f_low=0 (full waveform) and f_ref=f_lower_NR_dimless. The surrogate backward-evolves the spins from the NR epoch to its start.

  2. NR longer than surrogate (f_lower_NR < f_min_sur): The surrogate cannot extrapolate before its minimum frequency, so f_ref is clipped to f_min_sur. For aligned-spin or non-spinning systems the spin components are constant (no precession), so the metadata spins remain valid regardless of epoch. For precessing SXS systems Phase 2 automatically re-extracts the spin vectors from Horizons.h5 at the clipped epoch so that the spin-epoch is always physically consistent.

load_nrsur7dq4

load_nrsur7dq4()

Load and cache the NRSur7dq4 surrogate model.

Returns:

Type Description
gwsurrogate surrogate object
Source code in nrcatalogtools/surrogate.py
def load_nrsur7dq4():
    """Load and cache the NRSur7dq4 surrogate model.

    Returns
    -------
    gwsurrogate surrogate object
    """
    global _nrsur7dq4
    if _nrsur7dq4 is None:
        import gwsurrogate as gws

        _nrsur7dq4 = gws.LoadSurrogate("NRSur7dq4")
    return _nrsur7dq4

generate_surrogate_modes

generate_surrogate_modes(params: dict, total_mass: float, distance: float = 1.0, delta_t_seconds: float = 1.0 / 4096, sim_name: str | None = None, catalog=None, nr_wfm=None) -> tuple[dict, float]

Call NRSur7dq4 and return physical-unit modes as a pycbc TimeSeries dict.

Parameters:

Name Type Description Default
params dict

PyCBC-compatible binary parameter dict as returned by CatalogBase.get_parameters(). Must contain: mass1, mass2, spin1x/y/z, spin2x/y/z, f_lower.

required
total_mass float

Total binary mass in solar masses (sets the physical time/frequency scale for the surrogate call).

required
distance float

Luminosity distance in Mpc for amplitude scaling (default 1).

1.0
delta_t_seconds float

Desired sample spacing in physical seconds (default 1/4096).

1.0 / 4096
sim_name str

Simulation name, used for epoch-aligned spin extraction on precessing SXS runs.

None
catalog CatalogBase

Catalog instance; enables Phase 2 epoch alignment when the catalog is SXS.

None
nr_wfm WaveformModes

Unused; kept for API compatibility.

None

Returns:

Type Description
tuple[dict, float]

({(ell, em): pycbc.types.TimeSeries}, f_lower_effective) — complex mode time series in physical units with epoch at peak (2,2) amplitude, plus the effective starting GW frequency in Hz.

Raises:

Type Description
ValueError

If f_lower is not positive or the surrogate call fails.

Source code in nrcatalogtools/surrogate.py
def generate_surrogate_modes(
    params: dict,
    total_mass: float,
    distance: float = 1.0,
    delta_t_seconds: float = 1.0 / 4096,
    sim_name: str | None = None,
    catalog=None,
    nr_wfm=None,
) -> tuple[dict, float]:
    """Call NRSur7dq4 and return physical-unit modes as a pycbc TimeSeries dict.

    Parameters
    ----------
    params : dict
        PyCBC-compatible binary parameter dict as returned by
        ``CatalogBase.get_parameters()``.  Must contain:
        ``mass1``, ``mass2``, ``spin1x/y/z``, ``spin2x/y/z``, ``f_lower``.
    total_mass : float
        Total binary mass in solar masses (sets the physical time/frequency
        scale for the surrogate call).
    distance : float, optional
        Luminosity distance in Mpc for amplitude scaling (default 1).
    delta_t_seconds : float, optional
        Desired sample spacing in physical seconds (default 1/4096).
    sim_name : str, optional
        Simulation name, used for epoch-aligned spin extraction on precessing SXS runs.
    catalog : CatalogBase, optional
        Catalog instance; enables Phase 2 epoch alignment when the catalog is SXS.
    nr_wfm : WaveformModes, optional
        Unused; kept for API compatibility.

    Returns
    -------
    tuple[dict, float]
        ``({(ell, em): pycbc.types.TimeSeries}, f_lower_effective)`` —
        complex mode time series in physical units with epoch at peak (2,2)
        amplitude, plus the effective starting GW frequency in Hz.

    Raises
    ------
    ValueError
        If ``f_lower`` is not positive or the surrogate call fails.
    """
    q = params["mass1"] / params["mass2"]  # ≥ 1 by PyCBC convention
    chiA = [params["spin1x"], params["spin1y"], params["spin1z"]]
    chiB = [params["spin2x"], params["spin2y"], params["spin2z"]]
    f_lower_hz = params["f_lower"]

    if f_lower_hz <= 0:
        raise ValueError(
            f"f_lower = {f_lower_hz} Hz is not positive. "
            "Cannot determine surrogate start frequency."
        )

    m_secs = utils.time_to_physical(total_mass)  # M * MTSUN_SI  [seconds]

    # f_ref in cycles/M (= M * f_GW_hz, where M is in seconds).
    f_ref_dimless = f_lower_hz * m_secs  # cycles/M
    dt_dimless = delta_t_seconds / m_secs  # dimensionless time step

    sur = load_nrsur7dq4()

    # Query the surrogate's training window start from its internal metadata so
    # that epoch-aligned spin extraction targets the correct NR time regardless
    # of which surrogate model is loaded.  Fall back to a conservative value if
    # the attribute is absent (e.g. a future surrogate with a different layout).
    try:
        _sur_t_start_M = abs(sur._sur_dimless.t_0)
    except AttributeError:
        _sur_t_start_M = 4500.0

    chi1_perp = np.sqrt(chiA[0] ** 2 + chiA[1] ** 2)
    chi2_perp = np.sqrt(chiB[0] ** 2 + chiB[1] ** 2)
    is_precessing = (chi1_perp > 1e-4) or (chi2_perp > 1e-4)

    # --- Phase 2: Proper epoch alignment for precessing SXS binaries ---
    # For aligned-spin/non-spinning systems the spins are constant, so
    # the metadata values are valid at any epoch.  For precessing systems
    # we must extract the instantaneous spin vectors from the NR dynamics
    # at a common epoch and rotate them into the surrogate's coprecessing
    # reference frame before calling NRSur7dq4.
    _phase2_sim_obj = None
    if (
        is_precessing
        and catalog is not None
        and getattr(catalog, "CATALOG_TYPE", None) == "SXS"
        and sim_name
    ):
        try:
            import sxs as _sxs

            _phase2_sim_obj = _sxs.load(sim_name, auto_supersede=True, download=False)
            chiA, chiB, f_ref_dimless, _phase2_R = _epoch_align_spins(
                _phase2_sim_obj, target_t_before_merger=_sur_t_start_M
            )
        except Exception as _exc:
            import traceback

            traceback.print_exc()
            print(
                f"      [Phase 2 Warning] Epoch alignment failed ({_exc}); "
                "falling back to metadata spins."
            )
    # -----------------------------------------------------------------------

    try:
        t_sur, h_sur, _ = sur(
            q, chiA, chiB, ellMax=4, dt=dt_dimless, f_low=0, f_ref=f_ref_dimless
        )
    except Exception as exc:
        err_str = str(exc)
        if (
            "too small" in err_str.lower()
            or "omega_ref" in err_str
            or "omega_0" in err_str
        ):
            import re

            m_match = re.search(r"([0-9.]+)\s*=\s*omega_0", err_str)
            if m_match:
                omega_min = float(m_match.group(1)) * 1.01
            else:
                omega_min = 0.0170
            f_ref_clipped = omega_min / np.pi

            if is_precessing and _phase2_sim_obj is not None:
                # The epoch implied by f_ref_clipped differs from the one we
                # used above — re-extract spin vectors at the correct NR epoch.
                try:
                    chiA, chiB, f_ref_clipped, _phase2_R = _epoch_align_spins(
                        _phase2_sim_obj, f_ref_target=f_ref_clipped
                    )
                    print(
                        f"      [Phase 2] f_ref clipped to {f_ref_clipped:.5f} cycles/M; "
                        "re-extracted epoch-aligned spins."
                    )
                except Exception as _exc2:
                    print(
                        f"      [Phase 2 Warning] Re-alignment at clipped f_ref failed "
                        f"({_exc2}); keeping previously aligned spins."
                    )
            elif is_precessing:
                print(
                    f"      WARNING: f_ref={f_ref_dimless:.5f} cycles/M is below "
                    f"surrogate minimum; clipping to {f_ref_clipped:.5f}.  "
                    "For precessing systems the NR spins should be extracted from "
                    "NR dynamics at f_min_sur — using metadata spins instead "
                    "(may introduce a spin-epoch error)."
                )
            t_sur, h_sur, _ = sur(
                q, chiA, chiB, ellMax=4, dt=dt_dimless, f_low=0, f_ref=f_ref_clipped
            )
        else:
            raise

    amp_scale = utils.amp_to_physical(total_mass, distance)
    t_physical = t_sur * m_secs

    peak_idx = int(np.argmax(np.abs(h_sur[(2, 2)])))
    peak_time_phys = t_physical[peak_idx]
    epoch = t_physical[0] - peak_time_phys

    phase22 = np.unwrap(np.angle(h_sur[(2, 2)]))
    omega_gw_start = abs(phase22[1] - phase22[0]) / dt_dimless
    f_lower_effective = omega_gw_start / (2.0 * np.pi) / m_secs

    result = {}
    for (ell, em) in SURROGATE_MODES:
        if (ell, em) not in h_sur:
            continue
        h_phys = h_sur[(ell, em)] * amp_scale
        result[(ell, em)] = TimeSeries(
            h_phys.astype(np.complex128),
            delta_t=delta_t_seconds,
            epoch=epoch,
        )

    return result, f_lower_effective

check_surrogate_prior

check_surrogate_prior(params: dict, q_max: float = 4.0, chi_max: float = 0.8) -> bool

Return True if the binary parameters fall within the NRSur7dq4 prior volume.

NRSur7dq4 is valid for: - q = m1/m2 ∈ [1, 4] - |χ₁|, |χ₂| ≤ 0.8 - Waveform length ≥ ~4350 M (not checked here — handled by the surrogate call).

Parameters:

Name Type Description Default
params dict

PyCBC-compatible parameter dict.

required
q_max float

Maximum mass ratio (default 4).

4.0
chi_max float

Maximum spin magnitude (default 0.8).

0.8

Returns:

Type Description
bool
Source code in nrcatalogtools/surrogate.py
def check_surrogate_prior(
    params: dict, q_max: float = 4.0, chi_max: float = 0.8
) -> bool:
    """Return True if the binary parameters fall within the NRSur7dq4 prior volume.

    NRSur7dq4 is valid for:
    - ``q = m1/m2 ∈ [1, 4]``
    - ``|χ₁|, |χ₂| ≤ 0.8``
    - Waveform length ≥ ~4350 M (not checked here — handled by the surrogate call).

    Parameters
    ----------
    params : dict
        PyCBC-compatible parameter dict.
    q_max : float, optional
        Maximum mass ratio (default 4).
    chi_max : float, optional
        Maximum spin magnitude (default 0.8).

    Returns
    -------
    bool
    """
    q = params["mass1"] / params["mass2"]
    if not (1.0 <= q <= q_max):
        return False
    chi1 = np.sqrt(
        params["spin1x"] ** 2 + params["spin1y"] ** 2 + params["spin1z"] ** 2
    )
    chi2 = np.sqrt(
        params["spin2x"] ** 2 + params["spin2y"] ** 2 + params["spin2z"] ** 2
    )
    return chi1 <= chi_max and chi2 <= chi_max

Constants

Name Description
SURROGATE_MODES Mode list output by NRSur7dq4: (2,1), (2,2), (3,2), (3,3), (4,3), (4,4)
NR_MODES Superset of SURROGATE_MODES — all modes present in the NR catalogs, including (5,5)
_SURROGATE_F_MIN_CYCLES_PER_M Minimum \(M f_{GW}\) (cycles/M) for NRSur7dq4 training (0.0165)