waveformrs for Python

Python bindings of the Rust crate waveformrs (the parent directory; see its CHANGELOG.md and report/), gravitational-wave models that reproduce LALSimulation 6.2.0 to rounding: TaylorF2, IMRPhenomD, IMRPhenomXAS, IMRPhenomXHM, IMRPhenomXP, IMRPhenomXPHM and IMRPhenomXPNR in the frequency domain, the binary-neutron-star models IMRPhenomD_NRTidalv2, IMRPhenomXAS_NRTidalv2/v3 and IMRPhenomXP_NRTidalv2/v3, and NRSur7dq4 in the time domain, named as in LALSimulation.

The bindings use PyO3 and rust-numpy. Every function releases the GIL while it computes and returns numpy arrays; the batch functions spread many binaries over threads and write straight into numpy-owned memory.

Installation

Wheels for Linux and macOS, for CPython 3.10 and later and free-threaded CPython 3.14 and later, are on PyPI:

pip install waveformrs

To build from source, from this directory (python/), in the environment where the package should go (a virtualenv or an activated conda environment; for an environment that is not activated, set CONDA_PREFIX to it):

pip install maturin           # >= 1.14.1
maturin develop --release     # build and install in editable mode

maturin build --release builds a wheel into target/wheels/ instead, and pip install . builds and installs one in a step. The package needs CPython 3.10 or later. The extension is built against Python's stable ABI (abi3), so one wheel serves CPython 3.10 and later on a platform.

Free-threaded CPython (3.14t and later) is supported: the extension declares that it does not need the GIL, so importing it leaves the GIL disabled, and waveforms can be computed from several threads at once. Free-threaded 3.15 and later share one wheel per platform (the free-threaded stable ABI, abi3t); 3.14t has a wheel of its own. (PyO3 0.30 does not support free-threaded 3.13.)

For this repository's development environment:

cd python
CONDA_PREFIX=/Users/colmtalbot/mambaforge/envs/waveformrs \
    /Users/colmtalbot/mambaforge/envs/waveformrs/bin/maturin develop --release

Usage

import numpy as np
import waveformrs as wf

wf.approximants()
# ['TaylorF2', 'IMRPhenomD', 'IMRPhenomXAS', 'IMRPhenomXHM', 'IMRPhenomXP',
#  'IMRPhenomXPHM', 'IMRPhenomXPNR', 'NRSur7dq4', 'IMRPhenomD_NRTidalv2',
#  'IMRPhenomXAS_NRTidalv2', 'IMRPhenomXAS_NRTidalv3', 'IMRPhenomXP_NRTidalv2',
#  'IMRPhenomXP_NRTidalv3']

binary = dict(
    mass_1=36.0,                 # solar masses
    mass_2=29.0,
    spin_1=(0.4, 0.2, 0.3),      # dimensionless (x, y, z) at f_ref, L along z
    spin_2=(-0.2, 0.3, -0.1),
    luminosity_distance=500.0,   # Mpc
    inclination=0.7,             # rad
    phi_ref=1.2,                 # orbital reference phase, rad
    f_ref=20.0,                  # Hz (0: LAL's default of the model)
)

# SimInspiralChooseFDWaveformSequence
frequencies = np.linspace(20.0, 1024.0, 8033)
hplus, hcross = wf.polarizations("IMRPhenomXPHM", frequencies, **binary, precession_version=223)

The binary is given by keyword arguments with the names of the Rust Binary: mass_1 and mass_2 are required; the spins default to zero, the distance to 1 Mpc, the angles and f_ref to zero. f_ref = 0 selects LAL's default reference frequency of each model (the first frequency for the IMRPhenom models, coalescence for TaylorF2).

Options

The settings of LAL's LALDict are WaveformOptions, given as keyword arguments of any waveform function or as an object (options=); keywords override the object.

options = wf.WaveformOptions(precession_version="MSA", mode_array=[(2, 2), (3, 3)])
hp, hc = wf.polarizations("IMRPhenomXPHM", frequencies, **binary, options=options)
hp, hc = wf.polarizations("IMRPhenomXPHM", frequencies, **binary, options=options, version_fallback=False)

wf.WaveformOptions.lalsimulation_defaults()   # LAL 6.2.0 with an empty LALDict: MSA angles
options.replace(mode_array=None)              # a changed copy
options.as_dict()                             # {'precession_version': 223, ...}

Note that the default precession version is the crate's (SpinTaylor, 320), not LAL 6.2.0's (MSA, 223), and that LAL's multibanding is not implemented: the models reproduce LAL with PhenomXHMThresholdMband = PhenomXPHMThresholdMband = 0.

Frequency grids

For many waveforms on the same frequencies, a FrequencyGrid precomputes the powers of the frequency that the models need. The grid paths agree with polarizations to about 1e-13 relative rather than bit for bit.

grid = wf.FrequencyGrid.uniform(20.0, 1024.0, 0.125)    # or wf.FrequencyGrid(frequencies)
aligned = dict(binary, spin_1=(0.0, 0.0, 0.3), spin_2=(0.0, 0.0, -0.1))   # XAS takes aligned spins only
hp, hc = wf.polarizations_on_grid("IMRPhenomXAS", grid, **aligned)
grid.frequencies, len(grid)

Modes

modes returns the modes h_lm without the spherical harmonics at +f and -f, in the conventions of SimInspiralChooseFDModes, and choose_fd_modes reproduces SimInspiralChooseFDModes itself (IMRPhenomXHM, IMRPhenomXPHM, IMRPhenomXPNR) on its uniform grid. Both return a FrequencyModes, laid out as LAL's SphHarmFrequencySeries: each mode on the two-sided frequencies -f[::-1], f (with f = 0 once).

two_sided, hlm = wf.modes("IMRPhenomXHM", np.linspace(20.0, 512.0, 1000), mass_1=30.0, mass_2=20.0)
hlm[2, -2]                   # h_{2,-2} on the frequencies two_sided

result = wf.choose_fd_modes("IMRPhenomXPHM", 0.25, 20.0, 1024.0, **binary)
result.frequencies, result.modes, result.theta, result.phi, result.polarization_angle
hp, hc = result.polarizations()        # summed with the spherical harmonics (edit result.modes to leave modes out)

Time domain and conversions

NRSur7dq4 needs its data file NRSur7dq4_v1.0.h5 (zenodo.org/records/14999310): set NRSUR7DQ4_DATA to the file or LAL_DATA_PATH to its directory before the first call (wf.nrsur7dq4_data_file() shows what is found).

# SimInspiralChooseTDWaveform (also for the frequency-domain models but TaylorF2,
# through XLALSimInspiralTDFromFD as in LAL)
h = wf.td_polarizations("NRSur7dq4", 1 / 4096, 20.0, **binary)
hplus, hcross = h
h.epoch, h.delta_t, h.times         # first sample (s, relative to the peak), sampling, sample times

modes = wf.td_modes("NRSur7dq4", 1 / 4096, 20.0, **binary)   # SimInspiralChooseTDModes
times, hlm = modes

# SimInspiralTD and SimInspiralFD: conditioned waveforms of any approximant
td = wf.simulate_td("IMRPhenomXPHM", 1 / 4096, 20.0, **binary)
fd = wf.simulate_fd("NRSur7dq4", 0.0, 20.0, 2048.0, **binary)    # delta_f = 0: automatic
fd.hplus, fd.frequencies, fd.delta_f, fd.epoch

td_polarizations keeps an NRSur7dq4 workspace per thread: repeated calls allocate nothing, and calls that change only the inclination, phase, distance or sampling skip the surrogate evaluation.

Batches

batch_polarizations evaluates many binaries on one grid and returns (n_binaries, n_frequencies) arrays. The binaries are a dict of arrays (or scalars, broadcast), a numpy structured array, a pandas DataFrame or an (n, 12) table (wf.parameter_table); the spins are (n, 3) arrays or their components spin_1x, spin_1y, ... The results do not depend on threads (0 for all cores) and each row equals the single polarizations_on_grid call (exact=True: the single polarizations call) bit for bit.

n = 1000
rng = np.random.default_rng(1)
parameters = dict(
    mass_1=rng.uniform(20, 50, n),
    mass_2=rng.uniform(10, 20, n),
    spin_1=rng.uniform(-0.5, 0.5, (n, 3)),
    spin_2=rng.uniform(-0.5, 0.5, (n, 3)),
    luminosity_distance=400.0,
    inclination=rng.uniform(0, np.pi, n),
    phi_ref=0.0,
    f_ref=20.0,
)
result = wf.batch_polarizations("IMRPhenomXPHM", grid, parameters, threads=8)
hplus, hcross = result                    # complex128, (1000, 8033)
result.ok                                 # False for the binaries that failed (their rows are zero)
result.errors                             # {index: exception}

wf.batch_polarizations("IMRPhenomXPHM", grid, parameters, strict=True)   # raise the first error

# NRSur7dq4 (or any approximant) in the time domain: one workspace per thread
heavy = dict(parameters, mass_1=rng.uniform(50, 70, n), mass_2=rng.uniform(25, 35, n))
td = wf.batch_td_polarizations("NRSur7dq4", 1 / 4096, 20.0, heavy, threads=8)
td[0].hplus, td[0].epoch, td.ok, td.errors      # td[i] is None where td.ok[i] is False

Python threads also run in parallel, since the GIL is released:

from concurrent.futures import ThreadPoolExecutor

with ThreadPoolExecutor(4) as pool:
    waveforms = list(pool.map(lambda m: wf.polarizations("IMRPhenomXHM", frequencies, mass_1=m, mass_2=20.0),
                              np.linspace(20, 60, 100)))

Errors

The crate's errors are raised as subclasses of waveformrs.WaveformError (in waveformrs.exceptions) that also derive from the matching built-in exception:

Missing or unknown keyword arguments raise TypeError.

Reference

The result classes unpack like tuples of their main arrays (hplus, hcross = ..., frequencies, modes = ..., times, modes = ...) and keep the sampling metadata as attributes. Docstrings and the type stub _waveformrs.pyi give the details.

Tests

The tests compare with lalsimulation's Python API (python-lalsimulation 6.2.0), at the crate's tolerances: a relative error (max |difference| / max |LAL|) of 1e-10 for TaylorF2 and IMRPhenomD and 1e-9 for the IMRPhenomX models, and a mismatch of 1e-12, on macOS arm64; elsewhere the crate's cross-build tolerances. They also check that batches do not depend on the number of threads and equal the single calls, the exceptions, and that the GIL is released.

cd python
OMP_NUM_THREADS=1 NRSUR7DQ4_DATA=/path/to/NRSur7dq4_v1.0.h5 python -m pytest tests

The NRSur7dq4 tests are skipped without the data (found through NRSUR7DQ4_DATA, LAL_DATA_PATH or ../lal_data/); the comparisons with LAL are skipped without lalsimulation. LAL runs with OMP_NUM_THREADS=1 (the tests set it), since the OpenMP loop of IMRPhenomXP races in the MSA and PNR angles.

Benchmark

benchmarks/benchmark.py compares the Python-level call overhead and the throughput with lalsimulation's Python API, single-threaded and batched:

OMP_NUM_THREADS=1 NRSUR7DQ4_DATA=... LAL_DATA_PATH=... python benchmarks/benchmark.py [--quick] [--json out.json]

Indicative results (Apple M2 Ultra, 24 cores, with other jobs running; lalsimulation 6.2.0 single-threaded with multibanding off; median of 7 runs):

The speed-ups are those of the Rust models (see the crate's report); the Python layer adds about 2 µs per call. A batch on one thread is faster than single calls because it uses the grid paths.

Layout

• Cargo.toml: the extension crate waveformrs-python (library _waveformrs), depending on the crate with waveformrs = { path = ".." } (the root is not a workspace), PyO3 0.30 with abi3-py310 and abi3t-py315, and numpy 0.30, built with lto = "fat" as the crate recommends.
• src/: the bindings (lib.rs functions and FrequencyGrid, params.rs keyword arguments and WaveformOptions, batch.rs, results.rs, errors.rs, models.rs).
• python/waveformrs/: the Python package (__init__.py, the exceptions, the batch helpers, the stub _waveformrs.pyi and py.typed); maturin puts the extension module waveformrs/_waveformrs*.so next to them.
• tests/, benchmarks/.

A new option of the crate's WaveformOptions is bound by an arm of apply_option and an entry of option_items in src/params.rs (plus a getter and a line of the stub).