Quickstart

Install the package, point it at the model files, get a waveform.

1. Install

pip install nrhjsurrogate

The package ships code only. Its runtime dependencies are NumPy, SciPy, h5py and spectools, and pip installs them for you. Everything on this page works with those.

That already includes the compiled backend -- the one section 6 uses -- wherever a wheel is published:

3.14 gained a wheel in 1.4. Before that requires-python admitted 3.14 with no wheel to match, so every 3.14 user compiled Kokkos at install time. The free-threaded builds (3.14t) are deliberately not published and still compile from source.

Windows wheels carry the scalar evaluator. Every platform gets the same numbers, and Linux and macOS additionally get 1.4's sample-axis vector kernels; Windows does not, because cl.exe rejects the OpenMP simd directive those kernels use unless it is put into an experimental mode that this project's Kokkos build has not been tested against. A Windows wheel is therefore as accurate as any other and as fast as 1.3 was.

On anything else pip compiles the source distribution for you. That needs a C++17 compiler, and on macOS an OpenMP runtime that Apple's clang does not ship: conda install -c conda-forge llvm-openmp, or brew install libomp (keg-only, so CMake may also need CMAKE_ARGS="-DOpenMP_ROOT=$(brew --prefix libomp)"). Without one the build succeeds single-threaded, which is a success to pip and slow to you.

Ask for a source build yourself only for -march=native, for CUDA, or to turn a missing OpenMP into an error:

NRHJ_BUILD_KOKKOS=require NRHJ_BACKEND=openmp \
  pip install --no-binary nrhjsurrogate nrhjsurrogate
NRHJ_BACKEND=cuda NRHJ_CUDA_ARCH=sm_75 \
  pip install --no-binary nrhjsurrogate nrhjsurrogate        # NVIDIA GPU

A compile can fail, or succeed without OpenMP and be single-threaded, and neither of those fails the install. So it tells you: once on the first import after installing, in the header of nrhjsur-bench --quick, and any time you ask —

python -m nrhjsurrogate.compiled_backend

INSTALL.md section 1b has the full table.

2. Get the model files

The model artifacts are HDF5 files distributed separately from the code, because they are tens of megabytes each. Download them and put them in one directory.

For the released action-angle model you need two files:

The second one is not optional. The model resolves it by name every time it loads, and a directory without it gives a FileNotFoundError rather than a model that quietly stops at the merger.

3. Point the package at them

export NRHJ_MODEL_DIR=/path/to/the/directory

If you already keep LALSuite waveform data in $LAL_DATA_PATH, you can drop the files there instead and set nothing. That directory is searched too. You can also hand a full path straight to the loader.

You can also skip this step. If you have not downloaded anything, the first load fetches what it needs from the deposit into your per-user cache and verifies its SHA-256 before using it. Set NRHJ_AUTO_FETCH=0 if you would rather it never reached the network.

4. A waveform

from nrhjsurrogate import NRHJSurAA

model = NRHJSurAA.load_h5("NRHJSur3dq8_AA_v3")

# Geometric units: time in units of the total mass M, modes are r h/M.
times, modes = model.get_td_waveform_modes_geometric(4.0, 0.3, -0.2)

print(times.shape)                      # (15775,)
print(sorted(modes))
# [(2, 1), (2, 2), (3, 2), (3, 3), (4, 3), (4, 4), (5, 5)]

The three arguments are mass_ratio >= 1 and the two spin components along the orbital angular momentum. The model covers roughly 1 <= mass_ratio <= 8 and spin components from about -0.95 to +1.0. Check any point with model.in_hull(mass_ratio, spin1z, spin2z).

There are four waveform methods: get_td_waveform_modes_geometric and get_td_waveform_geometric in geometric units, get_td_waveform_modes and get_td_waveform in physical ones. The plain model(4.0, 0.3, -0.2) call is retained and returns the same content as one dictionary.

5. A waveform in physical units

import numpy as np

times, modes = model.get_td_waveform_modes(
    mass1=50.0, mass2=12.5, spin1z=0.3, spin2z=-0.2,
    delta_t=1.0 / 4096.0, f_min=30.0, distance=400.0)

h_plus, h_cross = model.get_td_waveform(
    mass1=50.0, mass2=12.5, spin1z=0.3, spin2z=-0.2,
    delta_t=1.0 / 4096.0, f_min=30.0, distance=400.0,
    inclination=np.pi / 3.0)

times                  # seconds, zero at the amplitude peak
modes[(2, 2)]          # strain in the (2, 2) mode

Masses in solar masses, delta_t in seconds, f_min in Hz of the (2, 2) gravitational-wave frequency, distance in Mpc, angles in radians.

Two things catch people out here:

• get_td_waveform returns two arrays and no time axis, which is the pycbc spelling this name comes from. Take the axis from get_td_waveform_modes at the same arguments.
• Without distance the modes are the geometric r h_lm / M, not strain. get_td_waveform requires a distance, because a polarization without one is not a strain.

6. Make it fast

model = NRHJSurAA.load_h5("NRHJSur3dq8_AA_v3", backend="kokkos")

Since 1.1 that is what an ordinary load() gives you anyway; the line above forces it, and raises rather than falling back if the compiled backend cannot serve the artifact.

On one test machine, with OMP_NUM_THREADS=8, this took 2.2 ms per waveform against 87.5 ms for the NumPy backend, and the two agreed to 2.2e-08 of the peak on the worst mode. The compiled module is looked for beside the package; set $NRHJ_KOKKOS_DIR to point somewhere else.

This needs a compiled backend that actually has OpenMP. An install that fell back to NumPy, or built Serial-only, will run this line without complaint and give you none of that speedup — python -m nrhjsurrogate.compiled_backend is how you tell which you have.

Next

docs/USAGE.md covers starting frequencies below what the artifact carries, the derived waveform routes, parameter derivatives, the Fisher matrix, and the model's own uncertainty estimate. Runnable versions of everything are in examples/usage/.