Installing nrhjsurrogate
How to set up and evaluate the AA surrogate (nrhjsurrogate.NRHJSurAA)
on a fresh machine. Three tiers, in increasing effort:
- Python-only (numpy backend) — works anywhere, no C++.
- CPU-OpenMP kokkos backend — the compiled evaluator, ~2.4x faster
single-call, ~4-5x at batch. Built automatically by
pip installof the source distribution, Kokkos included (section 2a). - CUDA kokkos backend — for NVIDIA GPUs; includes the honest caveats for GeForce cards (GTX 1650 Ti / Turing sm_75 specifically).
Which of the three you get is chosen by NRHJ_BACKEND (section 1b); by
default you get tier 2 where the compiler allows it and tier 1 where it
does not.
The CPU tiers are the well-trodden ones. The CUDA tier builds and links, but no waveform has yet been evaluated on the device and its float32 route is untested, so treat it as the least exercised path here.
The acceptance test for any tier is one command (section 6):
nrhjsur-smoke
To check the install AND the downloaded model files against shipped reference values (load, waveform, derivatives, error estimate, hybridisation, dynamics, remnant, file integrity), run the verification suite; it ships inside the package:
nrhjsur-verify # or: python -m nrhjsurrogate.verify
0a. Installing from PyPI — and where the models come from
There is no public git repository for this release. The package goes to
PyPI; the model artifacts are distributed separately as an archive deposit.
git lfs pull is therefore a developer step, not a user one.
pip install nrhjsurrogate
spectools(gitlab.com/vaishakp/spectools, MIT) is on PyPI and is a declared dependency, sopip install nrhjsurrogatepulls it and its own dependencies automatically.The pin is
spectools>=2026.8.8. Do not lower it: earlier releases were missing a module this package imports at import time, soimport nrhjsurrogatefails against them.
Then get the models. Either point at a directory you filled yourself:
export NRHJ_MODEL_DIR=/path/to/models
or let the package fetch and verify them by name:
import nrhjsurrogate
nrhjsurrogate.fetch_model("NRHJSur3dq8_AA_v3") # sha256-checked before use
The two released models are NRHJSur3dq8_AA_v3 (action-angle) and
NRHJSur3dq8_AdA_v3 (adiabatic-angle, 14-element), together with the
shared merger-ringdown arm NRHJSur3dq8_MR_v3 that both of them name.
The version digit is the model generation, not a patch level; the previous
generation (_v2) is deposited alongside and still loads by name, so
published v2 numbers stay reproducible. Both models are 3dq8 —
aligned-spin (q, chi1z, chi2z) with q ≤ 8, matching the NRHybSur3dq8
convention — and both are built on the N2 NR extrapolation, which the
names do not encode: each HDF5 records it as the root attribute
extrapolation = "N2".
Loaders take a name or a path:
from nrhjsurrogate import NRHJSurAA
model = NRHJSurAA.load_h5("NRHJSur3dq8_AA_v3")
out = model(2.0, 0.3, -0.1)
Resolution order: explicit path → $NRHJ_MODEL_DIR → $LAL_DATA_PATH →
per-user cache ($NRHJ_MODEL_CACHE_DIR, else platformdirs if installed,
else $XDG_CACHE_HOME/nrhjsurrogate, else ~/.cache/nrhjsurrogate) → a
source checkout's model_release/ → a download from the deposit into that
cache, checked against the registry's SHA-256 before the file is moved into
place. Failing all of it, a FileNotFoundError naming the exact file, its
SHA-256, every directory searched, and how to obtain it.
NRHJ_AUTO_FETCH=0 turns the download off. Nothing degrades silently —
the same rule as the merger-ringdown arm.
Verify the install:
nrhjsur-smoke
0. What the package needs — and what it does NOT
Run-path dependencies (traced, then verified by running in a clean env)
| dependency | why | verified version |
|---|---|---|
| Python ≥ 3.12 | pyproject.toml |
3.12.13 |
| numpy | everything | 2.3.5 |
| scipy | splines, quadrature, windows | 1.16.3 |
| h5py | reading the SWSpH merger-ringdown HDF5 | 3.16.0 |
spectools (spectools.chebyshev only) |
Chebyshev basis/diagnostics used by surrogate_prediction.py and spectral.py |
2026.8.8 (PyPI) |
spectools is pure numpy on this path (basis.py, chebyshev.py,
diagnostics.py). It is on PyPI and declared as a dependency with a
floor of 2026.8.8 — the first published version whose wheel ships
chebyshev/diagnostics.py (older wheels were missing it and broke
import nrhjsurrogate). A plain pip install resolves it; no
checkout is needed. See section 0a.
Explicitly NOT needed to evaluate waveforms
Verified by running the full smoke test under env -i with a PYTHONPATH
containing only this repo and spectools:
- lal / LALSimulation /
LAL_DATA_PATH— nothing innrhjsurrogateimports lal; the physical-units layer (NRHJSurAA.physical) carries its own constants (fisher.py:MTSUN,MRSUN,MPC). (If you ever do want lal for comparisons, it pip-installs fine — it is simply not a dependency of generation.) - sxs, pandas — imported only by TRAIN-time modules, which are not part
of the released evaluation package (
NRHJSurAA.build()unpickles the 1.3 GB training cache; running never touches them). - waveformtools, tdanalysis — mismatch scoring and BMS analysis in
scripts/, not waveform generation. No run-path module imports either. - bilby, pyseobnr, gwsurrogate — PE demos and benchmark scripts only.
- juliacall / PostNewtonian.jl — PN transcription oracle tests only (they skip cleanly when absent).
- matplotlib, scikit-learn — figures / legacy pickles only.
Artifact manifest
Released (distributed separately from the package — see section 0a):
| model | file | size (bytes) | sha256 | extrapolation |
|---|---|---|---|---|
NRHJSur3dq8_AA_v3 |
NRHJSur3dq8_AA_v3.h5 |
31,952,616 | b6d2d300035f5c9f7dac005980e8f01536ee59c08f8efaf57643f46811ebf483 |
N2 |
NRHJSur3dq8_AdA_v3 |
NRHJSur3dq8_AdA_v3.h5 |
15,734,472 | c0ca99222f8972e3d13f8273782e37fe84292ef9741e06893d2a7f18ff549a1a |
N2 |
NRHJSur3dq8_MR_v3 |
NRHJSur3dq8_MR_v3.h5 |
19,606,296 | dd1bf00860d01e16cf740bd227f5f7fb163284cb42309febd7194bc8d235059a |
N2 |
The previous generation is deposited too, so that published v2 numbers stay reproducible. It is not the default; name it explicitly to get it.
| model | file | size (bytes) | sha256 | extrapolation |
|---|---|---|---|---|
NRHJSur3dq8_AA_v2 |
NRHJSur3dq8_AA_v2.h5 |
30,743,544 | a05f6e0288e9896d2becff83b7eac0121c39285a4905c3d6f62741be97bd06b4 |
N2 |
NRHJSur3dq8_AdA_v2 |
NRHJSur3dq8_AdA_v2.h5 |
15,830,464 | 20dad0f5de344d347b605540fb0320570c6672411bffa23d093adee858bc61ce |
N2 |
NRHJSur3dq8_MR_v2 |
NRHJSur3dq8_MR_v2.h5 |
19,690,832 | da0ea1b031a5ff4ab1805566d2c3d2047948dee778929eb446fa8ce59f67546b |
N2 |
Both inspiral models are de-embedded: they reference the shared
merger-ringdown arm of their own generation — NRHJSur3dq8_MR_v3.h5 for the
v3 pair — by name rather than each carrying a copy, so all three files of a
generation are required. Do not mix generations: the arm is checked by
content digest at load, and a crossed pair is refused rather than served.
The same table lives in code as nrhjsurrogate.models.MODELS, which is what
fetch_model and nrhjsur-verify check a file against.
NRHJSur3dq8_AA_v3.h5 is the trained AA surrogate (dynamics + radiative
GPRs, the phase grid, blend calibration, and the predicted (J, E) action
content). NRHJSurAA.load_h5 resolves the shared merger-ringdown arm by
name at load and verifies it by content digest — a missing or wrong arm
raises loudly; there is no silent inspiral-only fallback.
NRHJSur3dq8_AdA_v3.h5 composes the same way — a 14-element inspiral
plus the shared arm — not a merger-ringdown-only file.
1. Python-only quickstart (numpy backend)
Works on any OS/CPU; no compiler.
# 1. code — Python 3.12 or newer
pip install nrhjsurrogate
# 4. artifacts — already in the checkout, nothing to copy (section 5)
# (git lfs pull above already fetched the released *_v3.h5 files,
# including the shared merger-ringdown arm the models need)
# 5. acceptance test
nrhjsur-smoke
First use:
from nrhjsurrogate import NRHJSurAA
model = NRHJSurAA.load_h5("NRHJSur3dq8_AA_v3") # a NAME; see section 0a
out = model(2.0, 0.3, -0.1) # geometric units, native grid
h22 = out[(2, 2)] # t = 0 at the (2,2) peak
r = model.physical(40.0, 25.0, 0.3, -0.1, # LAL-style physical units
delta_t=1.0 / 4096, f_min=20.0,
distance=440.0, inclination=0.7)
hp, hc = r["hp"], r["hc"]
The legacy NRHJSurAA.load(pkl, mr_path=...) path still exists — for a
locally-trained model (NRHJSurAA.build / .save), not the shipped
artifact — but its behavior on an unresolvable mr_path="auto" changed:
it now raises a clear FileNotFoundError naming what to pass,
instead of silently degrading to an inspiral-only model.
mr_path, updated: NRHJSurAA.load_h5 needs no mr_path in the normal
case, but not because the arm is inside the model file — it is not. The
shipped h5 records the arm's name, and the loader resolves
NRHJSur3dq8_MR_v3.h5 down the same search path as the model itself and
checks its content digest. Put the arm anywhere on that path (or fetch it,
which happens automatically) and you pass nothing; leave it out and the load
raises FileNotFoundError naming the file and every directory searched,
rather than returning a model that stops at the merger. merger_ringdown_path=
overrides the resolved arm when you deliberately want a different one.
For the legacy NRHJSurAA.load(pkl, ...) path,
mr_path="auto" now raises that FileNotFoundError (naming what to
pass) rather than silently falling back to inspiral-only;
mr_path=None is unchanged and still means an explicit, intentional
inspiral-only load.
With 8 threads: model load 0.08 s; full IMR generation ≈ 114 ms per
waveform single-call; physical() ≈ 120 ms. The
numpy backend is a ~0.1 s-per-call class — fine for interactive work and
small studies. (The ~16 ms/waveform figure quoted elsewhere is the
compiled batched path, not this one; see sections 2–3.) For
sampling-scale throughput use the compiled backend.
1b. Choosing the backend: NRHJ_BACKEND
One variable picks what gets compiled. It is orthogonal to
NRHJ_BUILD_KOKKOS, which only decides whether to compile at all, so
NRHJ_BUILD_KOKKOS=require NRHJ_BACKEND=cuda means "a CUDA build or no
install" — what a GPU CI job wants, and what neither variable can say alone.
NRHJ_BACKEND |
what you get | notes |
|---|---|---|
auto (default) |
Serial + OpenMP if the compiler has it, else Serial | never fails the install |
numpy |
no compile at all, pure-Python wheel | same as NRHJ_BUILD_KOKKOS=0 |
serial |
Serial only, single-threaded | for reproducing a single-thread number |
openmp |
Serial + OpenMP, fails the configure if OpenMP is missing | turns a silent degradation into a loud error |
cuda |
Serial + OpenMP + CUDA | needs NRHJ_CUDA_ARCH |
NRHJ_BACKEND=openmp pip install --no-binary nrhjsurrogate nrhjsurrogate
NRHJ_BACKEND=cuda NRHJ_CUDA_ARCH=sm_75 \
pip install --no-binary nrhjsurrogate nrhjsurrogate
NRHJ_CUDA_ARCH takes either spelling — sm_75 or the Kokkos name
TURING75. It is required and never guessed, because nvcc_wrapper's
own default is sm_70 and CUDA 13 rejects that outright, which would surface
as a confusing mid-build failure instead of a clean one.
CMAKE_ARGS is still honoured verbatim and is applied after these, so
anything you spell out by hand overrides the shorthand.
What each mode leaves in the build record, which you can read back with
python -m nrhjsurrogate.compiled_backend:
| mode | status |
serial / openmp / cuda |
|---|---|---|
numpy |
skipped |
— |
serial |
built |
ON / OFF / OFF |
openmp |
built |
ON / ON / OFF |
cuda |
built |
ON / ON / ON, plus the architecture |
A CUDA build carries the host spaces too: it is CPU and GPU in one module, not GPU instead of CPU.
2. CPU-OpenMP kokkos backend (compiled)
The compiled evaluator (cpp/ → nrhjkokkos) accelerates the
element-evaluation and assembly stages; the Python package auto-detects it
(the import is deferred, so nothing breaks if it is absent).
2a. The default route: pip install of the source distribution
Since 2026-09-02 (release checklist item B1) a source install builds the compiled backend itself, including Kokkos. Nothing has to be installed beforehand except a C++17 compiler:
pip install --no-binary nrhjsurrogate nrhjsurrogate # from PyPI, forcing the sdist
pip install nrhjsurrogate-1.0.0.tar.gz # or the tarball itself
What happens, in order:
- pip installs the build tools (
pybind11,cmake) into its isolated build environment; these are wheels, no compile. setup.pyruns CMake oncpp/. CMake looks for an installed Kokkos first (Kokkos_ROOT,Kokkos_DIR,CMAKE_PREFIX_PATH, or the conda environment it runs in). If none is found it downloads Kokkos 5.2.0 (tag5.2.0, SHA-256 checked, 1.8 MB) and builds it into the module as a static library, with the Serial backend always and the OpenMP backend when the compiler has OpenMP. 5.2.0 is the version every measurement in this repository was taken on; the pin and its reason are incpp/cmake/kokkos_dependency.cmake.- The module lands at the root of
site-packagesas a top-levelnrhjkokkos, which the package imports before it looks anywhere else.
On an 8-core machine the CMake step takes about 30 s from download to
link, and the whole pip install about two minutes, most of it
dependency downloads. With an installed Kokkos found through Kokkos_ROOT
nothing is downloaded and only the module's seven files compile.
What it needs, and what Kokkos 5.2.0 itself insists on: CMake ≥ 3.22 (pip
provides it), GCC ≥ 10.4, Clang ≥ 14, Apple Clang, or MSVC 19.30 (Visual
Studio 2022), and network access the first time unless Kokkos_ROOT
points at an install. A compiler below the floor is refused by Kokkos's
own configure step, which brings us to:
A failed build does not fail the install. If the compile cannot happen
(no compiler, one that is too old, no network), the install still
succeeds with the complete NumPy backend. It says so once: a banner in the
pip install -v transcript, and one RuntimeWarning the first time
import nrhjsurrogate runs in a process. The reason and the tail of the
CMake log are recorded inside the installed package; read them with
python -m nrhjsurrogate.compiled_backend
After fixing the cause, rebuild with
pip install --no-cache-dir --force-reinstall --no-binary nrhjsurrogate nrhjsurrogate
--no-cache-dir is not decoration: pip caches the wheel it built, degraded
or not, and a plain reinstall would reuse it without compiling.
Controls, all environment variables so that they pass through pip's build isolation unchanged:
| variable | effect |
|---|---|
NRHJ_BUILD_KOKKOS=0 |
do not attempt the compiled backend at all; the install is pure Python and the warning never appears. Use this when you want the NumPy path and do not want to wait for a compile. |
NRHJ_BUILD_KOKKOS=1 |
the compile must succeed or the install fails (for CI and the release acceptance test). Default when unset: try, and degrade to NumPy on failure. |
Kokkos_ROOT=<prefix> |
use that installed Kokkos; nothing is downloaded. Kokkos_DIR and CMAKE_PREFIX_PATH work the same way. |
CXX=<compiler> |
which C++ compiler CMake uses (the standard CMake variable). |
CMAKE_ARGS="-D... -D..." |
any option of cpp/CMakeLists.txt, for example -DNRHJ_MARCH_NATIVE=OFF for a build that does not target the build host's ISA, or -DFETCHCONTENT_SOURCE_DIR_KOKKOS=/path/to/kokkos-5.2.0 for an offline build from an unpacked Kokkos source tree. |
CMAKE_BUILD_PARALLEL_LEVEL=N |
compile jobs; default min(cpu count, 8). |
NRHJ_KOKKOS_DIR is a different variable: it tells the loader where an
already built nrhjkokkos module is, and is not involved in building one.
Building on a cluster: read this before you install on a login node
By default the compiled backend is built with -march=native, which targets
the instruction set of the machine running the compile. On a workstation
that is what you want. On a cluster it is a trap:
- you
pip installon a login node, - the build succeeds and
nrhjsur-smokepasses there, - your job lands on a compute node with an older or simply different CPU,
- and the module dies with an illegal instruction (SIGILL) — with nothing in the error naming a build-time flag as the cause.
The install now prints a warning at configure time whenever this flag is
applied, and records the flag together with the CPU it was chosen for, so
python -m nrhjsurrogate.compiled_backend will show you both after the fact.
If the nodes you will run on are not the node you are building on, choose one:
# portable baseline: runs anywhere, gives up AVX2/FMA tuning
CMAKE_ARGS="-DNRHJ_MARCH_NATIVE=OFF" \
pip install --no-binary nrhjsurrogate nrhjsurrogate
# or name the floor your compute nodes actually meet (usually better)
CMAKE_ARGS="-DNRHJ_MARCH_NATIVE=OFF -DCMAKE_CXX_FLAGS=-march=x86-64-v3" \
pip install --no-binary nrhjsurrogate nrhjsurrogate
x86-64-v3 is AVX2 + FMA, which every server CPU since roughly 2015 has.
Building on a compute node, in an interactive job, avoids the question
entirely and is the simplest answer where the queue allows it.
Note that NRHJ_MARCH_NATIVE=OFF is also the configuration the bitwise
reference gates are defined on (cpp/CMakeLists.txt), so a portable build is
if anything the more conservative choice numerically, not less.
Platforms. Verified here: Linux x86-64 with GCC. macOS and Windows are
reasoned, not run (no such machine was available): the CMake logic
chooses -march=native on x86-64, -mcpu=native on Apple Silicon and no
ISA flag on MSVC (cl.exe has no "native" option; -DNRHJ_MSVC_ARCH_FLAG=/arch:AVX2
through CMAKE_ARGS opts in), confirms each flag with the compiler before
using it, and builds Kokkos Serial-only on MSVC and on an Apple Clang
without libomp. cpp/cmake/compiler_flags.cmake states each branch and
its status.
Check what you got:
python -m nrhjsurrogate.compiled_backend # built / skipped / failed, with the reason
nrhjsur-smoke # the kokkos arm must say PASS, not SKIP
2b. Developer route: build cpp/ by hand in a checkout
The recipe every recorded measurement used, with the conda-forge toolchain
(all of it comes from setup_env.yml):
# an environment providing kokkos 5.2, pybind11, cmake and a C++17 compiler
ENV=$CONDA_PREFIX
cmake -S cpp -B cpp/build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_CXX_COMPILER=$ENV/bin/x86_64-conda-linux-gnu-g++ \
-DKokkos_ROOT=$ENV \
-Dpybind11_DIR=$($ENV/bin/python -m pybind11 --cmakedir) \
-DPYTHON_EXECUTABLE=$(which python)
cmake --build cpp/build -j 8
Notes:
-DPYTHON_EXECUTABLEis ALL-CAPS. That (legacy-style) spelling is what the pybind11 config actually consumes — it is the variable every verified build cache records, including this one. The mixed-case-DPython_EXECUTABLEform seen in older notes is not what got picked up; use the all-caps form so the module is pinned to the right interpreter (passing both is harmless).- Use the conda g++ together with the conda kokkos package (ABI
match). Any Kokkos ≥ 4 with a host backend works via
-DKokkos_ROOT=; leave-DKokkos_ROOTout and CMake fetches the pinned 5.2.0 exactly as the pip route does. - No install step: the module stays in
cpp/build/, andnrhjsurrogate.execution.compiled_module_loader.load_module()adds that directory tosys.pathitself. (Internal path, named here only to say where thesys.pathentry comes from; it is not part of the public surface and you should not need to call it.) - Build time ≈ 1–2 min on 8 cores.
Use it:
model = NRHJSurAA.load_h5(h5_path, backend="kokkos") # or model.backend = "kokkos"
Verify: nrhjsur-smoke also exercises the kokkos backend and gates
numpy-vs-kokkos parity (worst-mode relative L2 below 1e-6; it typically
comes out around 3e-9).
For what the compiled backend buys you, see docs/USAGE.md section 9, which
carries the current measured numbers for both backends taken on one machine at
OMP_NUM_THREADS=8. They are one machine's numbers, not a specification;
measure on yours.
3. CUDA build — GTX 1650 Ti (Turing, sm_75) target
This route has been built and run on a Turing-class card (sm_75) with CUDA 13.0; the compile takes a couple of minutes. It is the least exercised path in this document — see the note in the introduction.
The short way. The backend selector does all of this for you; you should not need the manual recipe below unless you are debugging it:
export PATH=/path/to/cuda/bin:$PATH # nvcc must be findable
NRHJ_BUILD_KOKKOS=require NRHJ_BACKEND=cuda NRHJ_CUDA_ARCH=sm_75 \
pip install --no-binary nrhjsurrogate nrhjsurrogate
setup.py fetches the pinned Kokkos itself to obtain nvcc_wrapper, makes
it the project's C++ compiler, and hands the same unpacked tree back to CMake
through FETCHCONTENT_SOURCE_DIR_KOKKOS so nothing is downloaded twice.
Override the wrapper with NRHJ_NVCC_WRAPPER=/path/to/nvcc_wrapper, or let
an installed Kokkos supply it via Kokkos_ROOT.
Two things worth knowing if you build this by hand.
- A CUDA build does not require an installed Kokkos — the fetched one
supplies
nvcc_wrapper. nvcc_wrapperhas to be the compiler for the whole project, not just for Kokkos. EnablingKokkos_ENABLE_CUDAwhile leaving your hostg++asCMAKE_CXX_COMPILERcompiles everything and then fails at link withsymbol 'fatbinData' is already defined, because the toolchain's link-time-optimisation flags collide with the CUDA fatbins. Usingnvcc_wrapperthroughout avoids it.
3.0 Read this first: is the GPU worth it on a 1650 Ti?
- FP64 caveat (honest). GeForce Turing runs float64 at 1/32 of
its float32 rate. The deployed evaluation path — including the public
backend="kokkos"route — is float64 end-to-end, so double-precision device evaluation on a 1650 Ti will be slow, plausibly no better than the CPU-OpenMP build on the same box. GPU numbers recorded on datacenter parts do not transfer to this card; those have full-rate FP64. - Recommended default on this card: CPU-OpenMP (section 2) for the
public API; use CUDA only if you opt into float32. float32 is a
documented accuracy class, not a silent downgrade: max strain
deviation vs the float64 reference ≈ 1.9e-4 of peak (measured;
~1e-8 mismatch class). Gates:
tests/test_kokkos.py:: test_evaluator_single_precision, and the tightertests/test_aa_batch.py::test_float32_envelope(5e-4 of peak), which is the float32 row of theACCURACY.mdcontract. It is opted into on the low-level evaluator —kokkos_eval.make_evaluator(..., single_precision=True)/aa_kokkos.AAKokkosEvaluator(model, single_precision=True). The publicNRHJSurAAAPI does not yet expose asingle_precisionkeyword, so float32 means driving the evaluator directly. - 4 GB VRAM. Ample for realistic batches: the in-span packing keeps
the device assembly footprint small — measured 34 MB at B = 256 on
segment grids, down 16x from the 537 MB an unpacked assembly needed.
Watch
nvidia-smithe first time anyway, since the footprint depends on how long the waveforms you ask for are.
3.1 Build Kokkos with CUDA (once)
# needs: CUDA toolkit (nvcc in PATH) and a gcc compatible with that CUDA
git clone https://github.com/kokkos/kokkos && cd kokkos
cmake -S . -B build \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_INSTALL_PREFIX=$HOME/Installations/kokkos/cuda \
-DKokkos_ENABLE_CUDA=ON \
-DKokkos_ENABLE_OPENMP=ON \
-DKokkos_ARCH_TURING75=ON \
-DCMAKE_CXX_STANDARD=17
cmake --build build -j && cmake --install build
-DKokkos_ARCH_TURING75=ON is the sm_75 architecture flag — the right
one for the GTX 1650 Ti. Keeping OpenMP on in the same install lets you
fall back to the CPU backend without a second Kokkos.
3.2 Build nrhjkokkos against it
From the recipe in cpp/CMakeLists.txt (adjust paths):
export PATH=/usr/local/cuda/bin:$PATH # nvcc must be found
export NVCC_WRAPPER_DEFAULT_COMPILER=$(which g++)
cmake -S cpp -B cpp/build-cuda \
-DCMAKE_BUILD_TYPE=Release \
-DCMAKE_CXX_COMPILER=$HOME/Installations/kokkos/cuda/bin/nvcc_wrapper \
-DCMAKE_CXX_FLAGS=-arch=sm_75 \
-DKokkos_ROOT=$HOME/Installations/kokkos/cuda \
-Dpybind11_DIR=$(python -m pybind11 --cmakedir) \
-DPYTHON_EXECUTABLE=$(which python)
cmake --build cpp/build-cuda -j
ln -sf $PWD/cpp/build-cuda/nrhjkokkos.*.so cpp/build/ # where load_module looks
Known sharp edges (already handled in-tree):
-DCMAKE_CXX_FLAGS=-arch=sm_75is required with CUDA ≥ 13: its nvcc rejects nvcc_wrapper's sm_70 default.- glibc 2.42+ vs CUDA ≤ 12.x
noexceptclash onsinpi/cospi/rsqrt: worked around bycpp/glibc_cuda_shim/math.h, included automatically by the CMake target. - Same
-DPYTHON_EXECUTABLEall-caps caveat as section 2. - No fast-math, ever: the accuracy contract forbids reassociation and FP contraction on the anchor path.
Then run nrhjsur-smoke — the parity gate is
backend-agnostic and must pass on CUDA exactly as on CPU in float64.
4. Artifact download / copy manifest
Installed users (no repository): you need not do anything. Loading a released model by name downloads what is missing and verifies it. To do it ahead of time, by name, which verifies them the same way —
import nrhjsurrogate
nrhjsurrogate.fetch_model("NRHJSur3dq8_AA_v3")
nrhjsurrogate.fetch_model("NRHJSur3dq8_AdA_v3")
nrhjsurrogate.fetch_model("NRHJSur3dq8_MR_v3")
— or download them yourself and check them:
export NRHJ_MODEL_DIR=/path/to/models
sha256sum -c <<'EOF'
b6d2d300035f5c9f7dac005980e8f01536ee59c08f8efaf57643f46811ebf483 NRHJSur3dq8_AA_v3.h5
c0ca99222f8972e3d13f8273782e37fe84292ef9741e06893d2a7f18ff549a1a NRHJSur3dq8_AdA_v3.h5
dd1bf00860d01e16cf740bd227f5f7fb163284cb42309febd7194bc8d235059a NRHJSur3dq8_MR_v3.h5
EOF
nrhjsur-verify does the same check against the package's own registry,
for whichever released files it finds, and names any that are missing.
Developers with a checkout: everything is fetched straight into the checkout — no separate copy step onto a cache directory:
git lfs install && git lfs pull # model_release/*.h5
The checkout copies carry construction provenance that is stripped from the deposited ones, so their sha256 does not match the table above and is not meant to. The digests above are for files obtained from the deposit.
Not needed on a run-only machine: the training-pickle cache the released
models were built from, loo_all260_hybrid.npz, the SXS catalogue cache,
LAL data.
5. Environment variables that matter
Build-time (read by pip install, not at run time)
| variable | values | effect |
|---|---|---|
NRHJ_BUILD_KOKKOS |
auto (default), 0/skip, 1/require |
whether to compile at all; require fails the install rather than degrading |
NRHJ_BACKEND |
auto (default), numpy, serial, openmp, cuda |
which backend to compile (section 1b) |
NRHJ_CUDA_ARCH |
sm_75, TURING75, … |
required by NRHJ_BACKEND=cuda; never guessed |
NRHJ_NVCC_WRAPPER |
path | use this nvcc_wrapper instead of fetching Kokkos for one |
CMAKE_ARGS |
any CMake flags | honoured verbatim, applied last so it overrides the selector |
CMAKE_BUILD_PARALLEL_LEVEL |
integer | build parallelism; default min(cpu, 8) (the cap is memory) |
Kokkos_ROOT / Kokkos_DIR |
prefix | use an installed Kokkos instead of fetching |
CXX |
path | the host compiler |
What the install tells you afterwards
pip discards build output on success and offers no post-install hook, so the outcome is written into the installed package and surfaced three ways:
python -m nrhjsurrogate.compiled_backend # the full build record, any time
- Once per install, on the first
import nrhjsurrogateafter installing, a one-line notice on stderr naming the module, the backend spaces (device space first when there is one, with its architecture) and the compiler. Keyed on the build record's timestamp, so a reinstall speaks again and a reimport does not. Silence it withNRHJ_QUIET_INSTALL_NOTICE=1. - A
RuntimeWarningwhen the build FAILED, or when it succeeded but got no OpenMP on a multicore machine — a Serial-only build is a success to pip and silently single-threaded to you. describe(), which every "compiled backend not found" message quotes.
The record carries the compiler version (not just its path), the CMake
build type and flags, Kokkos_ENABLE_{SERIAL,OPENMP,CUDA,HIP}, the device
Kokkos_ARCH_*, the backend you requested, and the verbatim CMAKE_ARGS.
A binary without its flags cannot be reasoned about, and a performance or
bit-exactness claim has to name what produced it.
Thread binding — set for you, automatically
Importing nrhjsurrogate sets OMP_PROC_BIND=close and
OMP_PLACES=cores, unless you have already chosen. You do not need to
set them by hand.
This happens in the first lines of nrhjsurrogate/__init__.py, ahead of
every other import, because OMP_PROC_BIND and OMP_PLACES are read
once, when the OpenMP runtime initialises — setting them later is a
silent no-op.
What was measured, and what was not. On an AMD EPYC 7742, paired arms on a
genuinely exclusive whole-node allocation (20 rounds, byte-identical null
control, 4-thread negative control), binding gives 1.08–1.11x on the WHOLE
CALL at OMP_NUM_THREADS=32 — 1.083 at B=1, 1.090 at B=256, 1.088 at B=1024,
against a null-control floor of ~1.3%. Kokkos' own printed advice (spread)
measured worse than unbound, which is why the default here is close rather
than what the warning suggests.
It is bitwise-inert by construction: an environment variable cannot change floating-point arithmetic, so nothing about accuracy is at stake.
Overriding:
- Set
OMP_PROC_BINDand/orOMP_PLACESyourself — any explicit setting of either wins, and the package then sets neither (the two are one policy; filling in half of a decision you did not make would be worse than doing nothing). - Or switch the whole thing off:
export NRHJ_NO_THREAD_DEFAULTS=1.
Check what the runtime actually did — this reads back
omp_get_proc_bind(), it does not echo the environment:
import nrhjsurrogate
print(nrhjsurrogate.thread_binding())
# {'env': {...}, 'applied_by_us': {...}, 'proc_bind': 'close', ...}
OMP_NUM_THREADS is deliberately NOT defaulted. A binding policy
generalises across machines; a thread count does not (OMP=32 was the
wall on that box — OMP=64 never faster and 13% slower at B=1, a property
of that hardware). Set it yourself per the table below.
Note: scripts/bench_memo.slurm, the preamble much of this project copies,
exports OMP_PROC_BIND=false. That is respected — and it means the
in-repo benchmark preamble opts out of the faster setting. Drop that line
to pick it up.
RESOLVED — the two measurements never disagreed; one of them was contaminated and the other was about a different regime.
- The single-process claim of 1.25–1.28x is withdrawn. It was taken on a node whose load climbed from 0.33 to 24 during the run. Retaken on a whole-node exclusive allocation it is 1.08–1.11x whole-call, with a ~1.3% null floor.
- The fork-pool warning stands and is now handled in code, not by advice: the library declines to bind when
OMP_NUM_THREADS<=1, which is exactly the fork-pool mode where a bound master's one-CPU mask is inherited byforkand can collapse an N-worker pool onto one core.So the shipped default no longer trades one regime against the other. A further measured result: the explicit
close/corespair buys essentially nothing over the library default (1.085 vs 1.083 at B=1) — so setting these variables by hand is not worth doing, and the value is in the library default alone.If you fork worker pools you should still find the default harmless, but
NRHJ_NO_THREAD_DEFAULTS=1(or an explicitOMP_PROC_BIND=false) remains the escape hatch.
| variable | run value | why |
|---|---|---|
OMP_PROC_BIND, OMP_PLACES |
close, cores — set for you on import |
measured 1.08–1.11x whole-call at OMP=32 on a dedicated node (B=1 1.083, B=256 1.090, B=1024 1.088; null-control floor ~1.3%). An earlier 1.25–1.28x was taken on a shared node as its load climbed from 0.33 to 24 and is withdrawn — binding works by exclusive occupancy of the cores it binds, so under contention you measure contention. Any explicit value of either yours wins |
NRHJ_NO_THREAD_DEFAULTS |
unset | set to 1 to stop the package touching affinity at all |
NRHJ_MODEL_DIR |
directory (or $PATH-style list) holding the released *.h5 |
the one users need. Searched by nrhjsurrogate.models.resolve_model right after an explicit path, so it wins over the per-user cache and over a checkout's model_release/ |
NRHJ_MODEL_CACHE_DIR |
download destination | overrides the per-user cache (platformdirs if installed, else $XDG_CACHE_HOME/nrhjsurrogate, else ~/.cache/nrhjsurrogate) that fetch_model writes into |
NRHJ_MODEL_BASE_URL |
deposit base URL, a local directory, or a file:// URL |
where the download comes from. Defaults to the Zenodo deposit, so it needs setting only to point at a mirror, a local directory or a file:// URL |
NRHJ_AUTO_FETCH |
unset, meaning on | set to 0, off, no or false and a missing artifact raises with instructions instead of downloading. resolve_model(..., fetch=False) does the same for one call |
NRHJ_CACHE_DIR |
default ~/nrhj_cache |
build-time cache root; the four NRHJSurAA.build/load defaults hang off it. Replaced a set of hardcoded absolute cluster paths |
NRHJ_MODEL_PKL |
path to the training pkl | legacy-path only: default artifact for NRHJSurAA.load() and for nrhjsur-smoke --pkl. Not read by NRHJSurAA.load_h5 at all |
NRHJ_V4_H5 |
path to nrhjsur_production_v4_swsph.h5 |
legacy-path only: MR-arm default for nrhjsur-smoke --pkl and for NRHJSurAA.load(..., mr_path="auto")'s fallback candidate. Irrelevant to the default load_h5 flow, which needs no MR path at all |
OMP_NUM_THREADS |
desktops: the physical core count. Server CPUs: 32–64 (measured plateau on EPYC 7742; all 128 threads THRASH ~10x) | kokkos OpenMP + numpy BLAS threading |
OMP_PROC_BIND |
leave it unset | The library sets close (+ OMP_PLACES=cores) for you when the process is multi-threaded, in kokkos_eval.load_module() before the compiled module is imported — libgomp reads these once at its own initialisation, so setting them later is a silent no-op. Setting the variable yourself to anything, including false, is treated as a deliberate choice and is never overridden. The library declines to bind when OMP_NUM_THREADS<=1: that is production PE's fork-pool mode, where a bound master thread has a one-CPU mask that fork inherits, so binding can collapse an N-worker pool onto one core — the hazard this row previously recommended false to avoid. NRHJ_PROC_BIND=off disables the default; =force applies it even at one thread. |
LAL_DATA_PATH |
not needed | nothing on the run path imports lal |
NRHJ_PROD_PKL, NRHJ_LOO_NPZ |
train-time only | NRHJSurAA.build() inputs |
Renamed variables. Every variable this package reads uses the NRHJ_
prefix. The older names still work and emit a DeprecationWarning once
per process: a variable set in a shell profile or job script cannot be
updated by an upgrade, and a hard rename would not raise — it would silently
fall back to a default and run the wrong model.
The warning names the replacement, so you do not need a table here: set the
variable it tells you to and remove the old one. One pairing is worth calling
out because the names are close and the directories are not the same — the
old model-cache variable maps to NRHJ_MODEL_CACHE_DIR, not to
NRHJ_CACHE_DIR, which is the build-cache root.
The per-user artifact cache is ~/.cache/nrhjsurrogate. If you have
artifacts under a cache path from an earlier version, move them across or
point NRHJ_MODEL_DIR at it; otherwise they are simply downloaded again.
6. Acceptance testing, deeper layers
nrhjsur-smoke— the install acceptance test (a couple of seconds on the NumPy backend). Checks: artifact load and merger-ringdown attach, geometric andphysical()generation, the below-nativef_lowPN-extension arm, a reference slice keyed to the artifact's sha256 at a tolerance of 5e-5 of peak (the cross-machine libm bound), and, when the compiled module imports, numpy-vs-kokkos parity. Loads the released action-angle model by NAME (NRHJSur3dq8_AA_v3) viaNRHJSurAA.load_h5by default, letting the resolver find it and the merger-ringdown arm it names (nomr_path, no pkl);--pklswitches to the legacy pkl +mr_pathflow (--model/--mr) for testing a locally-trained model. Other flags:--no-kokkos,--no-hybrid,--write-golden(rekey after deploying a new artifact). Against the shipped h5 the committed golden slice (generated from the older pkl artifact) reportsSKIP, notFAIL— expected until it is regenerated.nrhjsur-verify(equivalentlypython -m nrhjsurrogate.verify) is the verification suite that SHIPS in the wheel, so it runs on a pip install with no checkout. About 8 s on the numpy backend. It reports one PASS/FAIL/SKIP line per check and exits 1 on any failure: package imports; the shipped reference values; which released model files are present and where they were searched for; each present file's SHA-256 and size against the package's registry (a renamed file of another generation is named as such); the action-angle model loads with its merger-ringdown arm; a waveform has the seven modes and the recorded span and matches the recorded values at fixed sample times; parameter derivatives cover the same span and match; the model's own error estimate returns all seven modes over the whole span and matches; a start frequency below the artifact's reach extends the waveform and matches; dynamics; remnant mass and spin; physical-unit polarisations; the adiabatic-angle model and its waveform; the previous-generation files when present; and, when the compiled backend is built, its agreement with numpy. Without model files every model check is a SKIP naming the file and the variable to set, and the exit status is 0. Flags:--model-dir DIR,--only TEXT,--strict(a SKIP counts as a failure, for a machine where everything is expected),--list. The same checks run under pytest withpytest --pyargs nrhjsurrogate.verify(pip install nrhjsurrogate[test]). Tolerances: 1e-6 of the per-mode peak inside the native span, 1e-4 of the in-window peak for the post-Newtonian extension, 1e-8 absolute on the remnant; the reasoning is innrhjsurrogate/verify/checks.py. Reference values are regenerated only when a released artifact changes:python scripts/generate_verify_reference_values.pyin a checkout.pytest tests -q— the unit layer; cluster-only tests (training cache, julia oracles) skip cleanly on a fresh machine.NOT IN THE DISTRIBUTION. The test suite ships in neither the wheel nor the sdist: the released package is evaluation code only. Everything in this bullet and the three caveats under it applies to a source checkout. From an installed package, the acceptance test is the
nrhjsur-smokecommand in section 1.Three caveats, all MEASURED on a fresh worktree, all pre-existing:
- 23 tests FAIL rather than skip when the compiled
nrhjkokkosis absent —tests/test_exact_times.py,test_generic_grids.py, the[kokkos]parameters oftest_reach_authority.py,test_values_spline.py::test_kokkos_backend_ab, andtest_workspace_views.py::test_fork_pool_is_bitwise. They raiseImportError: could not import the compiled nrhjkokkosinstead of skipping, so a numpy-only install sees a wall of red that means nothing more than "section 2 not done yet". Build the module and all of them pass. Nothing is wrong with the model. - Always pass
-p no:randomly.pytest-randomlyis installed in this environment and both shuffles test order and reseedsrandom/numpy.randomper test, so two runs are not comparable without it. - Run the two
*_julia.pyfiles in a separatepytestinvocation. Imported late in a long run,juliacallaborts the interpreter (exit 127, no traceback) and every test after it is lost.
- 23 tests FAIL rather than skip when the compiled
tests/test_anchor_canonical.py+paper_data/anchor_golden.npz+ACCURACY.md— the four-layer canonical-anchor accuracy contract, whose golden layer runs artifact-free.