add rational_fit and nr_precessing tests

Author: md-arif-shaikh

gw_eccentricity/load_data.py

@@ -743,6 +743,78 @@ def load_lvcnr_waveform(**kwargs):
     return return_dict
 
 
+def download_sxs_waveform(sxs_id, lev, data_dir, extrap_order=2):
+    """Download SXS catalog waveform files to a local directory.
+
+    Downloads the files required by ``load_waveform(origin="SXSCatalog")``:
+
+    - ``Strain_N{extrap_order}.h5``
+    - ``Strain_N{extrap_order}.json``
+    - ``metadata.json``
+    - ``Horizons.h5``
+
+    Parameters
+    ----------
+    sxs_id : str
+        SXS simulation identifier, e.g. ``"SXS:BBH:3726"``.
+    lev : int
+        Resolution level, e.g. ``3``.
+    data_dir : str
+        Root directory.  Files are placed inside
+        ``<data_dir>/<sxs_id_safe>/Lev<lev>/``, where ``sxs_id_safe``
+        is ``sxs_id`` with ``":"`` replaced by ``"_"``
+        (e.g. ``"SXS_BBH_3726"``).
+    extrap_order : int, optional
+        Waveform extrapolation order.  Default is 2, which selects the
+        ``Strain_N2.h5`` / ``Strain_N2.json`` file pair.
+
+    Returns
+    -------
+    str
+        Full path to the directory that now contains the downloaded
+        files.  Pass this as ``data_dir`` to
+        ``load_waveform(origin="SXSCatalog", data_dir=...)``.
+
+    Notes
+    -----
+    Files that already exist on disk are skipped, that is not re-downloaded. 
+    The SXS simulations catalogue is loaded via the ``sxs`` Python package to
+    retrieve direct download URLs for each file.
+    """
+    from sxs.utilities import download_file
+
+    # Build target directory
+    sxs_id_safe = sxs_id.replace(":", "_")
+    target_dir = os.path.join(
+        os.path.expanduser(data_dir), sxs_id_safe, f"Lev{lev}")
+    os.makedirs(target_dir, exist_ok=True)
+
+    # Load the simulations catalogue to resolve file URLs
+    simulations = sxs.Simulations.load(download=True)
+    if sxs_id not in simulations:
+        raise ValueError(
+            f"'{sxs_id}' not found in the SXS simulations catalogue.  "
+            "Check the SXS ID and that the catalogue is up to date.")
+
+    sim_files = simulations[sxs_id].get("files", {})
+
+    filenames = [
+        f"Strain_N{extrap_order}.h5",
+        f"Strain_N{extrap_order}.json",
+        "metadata.json",
+        "Horizons.h5",
+    ]
+
+    for filename in filenames:
+        dest = os.path.join(target_dir, filename)
+        if os.path.exists(dest):
+            continue  # already present, skip
+        url = sim_files[f"Lev{lev}:{filename}"]["link"]
+        download_file(url, dest, progress=False)
+
+    return target_dir
+
+
 def load_sxs_catalogformat(**kwargs):
     """Load modes from sxs waveform files in sxs catalog format.
 

gw_eccentricity/rational_fit.py

@@ -104,7 +104,8 @@ def _eval_arnoldi_basis(x, H, degree, v0_norm):
     # Build higher degree basis functions using the stored Hessenberg
     # matrix H
     for k in range(degree):
-        v = V[:, k] * x - V[:, :k+1] @ H[:k+1, k]
+        with np.errstate(divide='ignore', over='ignore', invalid='ignore'):
+            v = V[:, k] * x - V[:, :k+1] @ H[:k+1, k]
         if H[k + 1, k] < 1e-14:
             break
         V[:, k + 1] = v / H[k + 1, k]
@@ -405,8 +406,9 @@ def predict(self, x_new):
 
         # Compute the numerator and denominator at x_new using the
         # fitted coefficients
-        p = P @ self.a
-        q = Q @ self.b
+        with np.errstate(divide='ignore', over='ignore', invalid='ignore'):
+            p = P @ self.a
+            q = Q @ self.b
 
         # Return the rational function values at x_new
         r = p / q

pytest.ini

@@ -1 +1,5 @@
-[pytest]
+[pytest]
+filterwarnings =
+    ignore::DeprecationWarning:pyseobnr
+    ignore::DeprecationWarning:scipy
+    ignore::DeprecationWarning:qnm

test/test_nr_precessing.py

@@ -0,0 +1,201 @@
+"""Tests for precessing=True with frame="inertial" and frame="coprecessing"."""
+import os
+import json
+import numpy as np
+import pytest
+from gw_eccentricity import measure_eccentricity
+from gw_eccentricity.load_data import load_waveform, get_coprecessing_data_dict
+
+# ---------------------------------------------------------------------------
+# SXS test cases
+#   SXS:BBH:2859  — nearly quasi-circular, precessing  [active]
+#   SXS:BBH:xxxx  — high-eccentricity, precessing      [placeholder; all
+#                   eccentric tests skipped until a suitable simulation is
+#                   available in the public SXS catalog]
+# ---------------------------------------------------------------------------
+
+DATA_DIR = "/tmp/gwecc_test_data"
+SXS_CASES = {
+    "eccentric":     ("SXS:BBH:xxxx", 0),
+    "quasicircular": ("SXS:BBH:2859", 4),
+}
+ALL_ELL2_MODES = [(2, -2), (2, -1), (2, 0), (2, 1), (2, 2)]
+
+
+def _data_dir(sxs_id, lev):
+    sxs_id_safe = sxs_id.replace(":", "_")
+    return os.path.join(DATA_DIR, sxs_id_safe, f"Lev{lev}")
+
+
+def _load(sxs_id, lev):
+    """Load all ell=2 modes from a downloaded SXS waveform."""
+    return load_waveform(
+        origin="SXSCatalog",
+        data_dir=_data_dir(sxs_id, lev),
+        mode_array=ALL_ELL2_MODES,
+        num_orbits_to_remove_as_junk=4)
+
+
+def _ref_ecc(sxs_id, lev):
+    """Read reference eccentricity from the downloaded metadata.json."""
+    path = os.path.join(_data_dir(sxs_id, lev), "metadata.json")
+    with open(path) as f:
+        return float(json.load(f)["reference_eccentricity"])
+
+
+# ---------------------------------------------------------------------------
+# Fixtures — check data is present, skip if not
+# ---------------------------------------------------------------------------
+
+def _make_fixture(key):
+    """Return a module-scoped fixture that loads one SXS case by key."""
+    @pytest.fixture(scope="module")
+    def _fixture():
+        sxs_id, lev = SXS_CASES[key]
+        data_dir = _data_dir(sxs_id, lev)
+        if not os.path.isfile(os.path.join(data_dir, "Strain_N2.h5")):
+            pytest.skip(
+                f"SXS data not found at {data_dir} — run download_sxs_waveform first")
+        return _load(sxs_id, lev), _ref_ecc(sxs_id, lev)
+    return _fixture
+
+
+eccentric_dataDict = _make_fixture("eccentric")
+quasicircular_dataDict = _make_fixture("quasicircular")
+
+
+@pytest.fixture(scope="module", params=list(SXS_CASES.keys()))
+def sxs_dataDict(request):
+    """Module-scoped dataDict for each SXS case; skips if data not downloaded."""
+    key = request.param
+    sxs_id, lev = SXS_CASES[key]
+    data_dir = _data_dir(sxs_id, lev)
+    if not os.path.isfile(os.path.join(data_dir, "Strain_N2.h5")):
+        pytest.skip(f"SXS data not found at {data_dir} — run download_sxs_waveform first")
+    return key, _load(sxs_id, lev), _ref_ecc(sxs_id, lev)
+
+
+# ---------------------------------------------------------------------------
+# Test: frame="inertial" triggers internal coprecessing transform
+# ---------------------------------------------------------------------------
+
+def test_inertial_frame_gives_finite_eccentricity(sxs_dataDict):
+    """measure_eccentricity with precessing=True, frame='inertial' returns finite result."""
+    key, dataDict, _ = sxs_dataDict
+    t = dataDict["t"]
+    t_mid = t[len(t) // 2]
+
+    result = measure_eccentricity(
+        tref_in=t_mid,
+        method="AmplitudeFits",
+        dataDict=dataDict,
+        precessing=True,
+        frame="inertial")
+    assert np.isfinite(result["eccentricity"]).all(), (
+        f"Non-finite eccentricity for {key} with frame='inertial'")
+
+
+# ---------------------------------------------------------------------------
+# Test: frame="coprecessing" (pre-rotated modes) matches frame="inertial"
+# ---------------------------------------------------------------------------
+
+def test_coprecessing_frame_matches_inertial(sxs_dataDict):
+    """frame='coprecessing' and frame='inertial' agree to within numerical precision."""
+    key, dataDict, _ = sxs_dataDict
+    t = dataDict["t"]
+    t_mid = t[len(t) // 2]
+
+    result_inertial = measure_eccentricity(
+        tref_in=t_mid,
+        method="AmplitudeFits",
+        dataDict=dataDict,
+        precessing=True,
+        frame="inertial")
+
+    coprec_dataDict = get_coprecessing_data_dict(dataDict)
+    result_coprec = measure_eccentricity(
+        tref_in=t_mid,
+        method="AmplitudeFits",
+        dataDict=coprec_dataDict,
+        precessing=True,
+        frame="coprecessing")
+
+    np.testing.assert_allclose(
+        result_coprec["eccentricity"],
+        result_inertial["eccentricity"],
+        rtol=1e-10,
+        err_msg=f"coprecessing vs inertial mismatch for {key}")
+
+
+# ---------------------------------------------------------------------------
+# Test: spin_filter is exercised for quasicircular precessing data.
+#       Eccentricity evolution should be monotonically decreasing and the
+#       initial value should be within a factor of 2 of the metadata reference.
+#
+#       NOTE: spin_filter is not triggered for high-eccentricity waveforms
+#       because the eccentricity signal dominates spin-induced oscillations.
+# ---------------------------------------------------------------------------
+
+def test_spin_filter_gives_monotonic_eccentricity(quasicircular_dataDict):
+    """For quasicircular precessing NR data (SXS:BBH:2859), verify that:
+
+    1. The eccentricity evolution is monotonically decreasing (spin_filter path
+       exercised — spin-induced oscillations are removed before measurement).
+    2. The initial eccentricity is within a factor of 2 of the SXS metadata
+       reference value.
+    """
+    dataDict, ecc_ref = quasicircular_dataDict
+    t = dataDict["t"]
+
+    result = measure_eccentricity(
+        tref_in=t,
+        method="AmplitudeFits",
+        dataDict=dataDict,
+        precessing=True,
+        frame="inertial")
+
+    ecc = result["eccentricity"]
+    assert np.isfinite(ecc).all(), "Non-finite eccentricity for quasicircular case"
+
+    # use the built-in monotonicity check which populates decc_dt_for_checks
+    obj = result["gwecc_object"]
+    obj.check_monotonicity_and_convexity()
+    assert np.all(obj.decc_dt_for_checks <= 0), (
+        "Eccentricity is not monotonically decreasing for quasicircular case")
+
+    # initial eccentricity vs SXS metadata reference (within factor of 2)
+    ecc_initial = float(ecc[0])
+    assert 0.5 * ecc_ref <= ecc_initial <= 2.0 * ecc_ref, (
+        f"Initial eccentricity {ecc_initial:.6f} not within factor of 2 "
+        f"of metadata reference {ecc_ref:.6f}")
+
+
+@pytest.mark.skip(
+    reason="No suitable high-eccentricity precessing SXS simulation is currently "
+           "available in the public catalog. TODO: re-enable and tighten to rtol=0.05 "
+           "once such a simulation is added.")
+def test_eccentric_precessing_eccentricity_value(eccentric_dataDict):
+    """For high-eccentricity precessing NR data, verify that the initial
+    eccentricity is within 5% of the SXS metadata reference value.
+
+    spin_filter is not triggered for high-eccentricity waveforms; only the
+    eccentricity value is checked here.
+    """
+    dataDict, ecc_ref = eccentric_dataDict
+    t = dataDict["t"]
+
+    result = measure_eccentricity(
+        tref_in=t,
+        method="AmplitudeFits",
+        dataDict=dataDict,
+        precessing=True,
+        frame="inertial")
+
+    ecc = result["eccentricity"]
+    assert np.isfinite(ecc).all(), "Non-finite eccentricity for eccentric case"
+
+    ecc_initial = float(ecc[0])
+    np.testing.assert_allclose(
+        ecc_initial, ecc_ref, rtol=0.05,
+        err_msg=f"Initial eccentricity {ecc_initial:.4f} differs from "
+                f"metadata reference {ecc_ref:.4f} by more than 5%")