jaxtfwaves

JAXTFwaves is the sibling project of tfwaves that is using JAX as backend.

tfwaves

TFWaves is a python package to compute fast time-frequency gravitational-wave response for LISA data analysis. It takes as imput phase and amplitude functions of the gravitational strain in the pseudo-inertial solar system barycenter frame, and outputs time-delay interferometry responses in the time-frequency domain.

Installation

The core package is available on PyPI:

pip install jaxtfwaves

To install a specific plugin, use its named extra:

pip install "jaxtfwaves[sobhbwaveform]"   # SOBHB waveforms via ripplegw
# pip install "jaxtfwaves[emriwaveform]"  # EMRI waveforms (coming soon)

Or install all available plugins at once:

pip install "jaxtfwaves[plugins]"

Typed API (TLT-First)

jaxtfwaves provides a typed output layer powered by typed-lisa-toolkit (TLT). This is the direction of the public API: the interface is being normalized around TLT-native objects rather than long-term raw sparse arrays.

Raw sparse outputs remain available for compatibility during the transition, but new code should prefer the typed methods and adapters below.

High-level typed methods

Use these methods directly on waveform objects:

• compute_hphc_tlt(...)
• compute_tdi_tlt(...)

from tfwaves.stft import TimeToTFWaveform

# Build your waveform object as usual
wvf = TimeToTFWaveform(...)

stft_plus, stft_cross = wvf.compute_hphc_tlt(di=5)
tdi_data = wvf.compute_tdi_tlt(di=5)  # defaults to XYZ or AET from wvf.channel

Adapter-level typed helpers

If you already work with raw sparse outputs, convert them explicitly with:

• normalize_sparse_stft_output
• sparse_stft_from_tracks
• hphc_stfts_from_tracks
• stftdata_from_tdi

from tfwaves import hphc_stfts_from_tracks, stftdata_from_tdi

h_plus, h_cross, indices_f, indices_t = wvf.compute_hphc(di=5)
stft_plus, stft_cross = hphc_stfts_from_tracks(
	h_plus,
	h_cross,
	indices_f,
	indices_t,
	wvf.f_bins,
	wvf.t_bins,
)

tdi_values, indices_f, indices_t = wvf.compute_tdi_tf_response(di=5)
tdi_tlt = stftdata_from_tdi(
	tdi_values,
	indices_f,
	indices_t,
	wvf.f_bins,
	wvf.t_bins,
)

Channel naming behavior

• If waveform channel contains XYZ, compute_tdi_tlt() defaults to X, Y, Z.
• If waveform channel contains AET, compute_tdi_tlt() defaults to A, E, T.
• You can override names with channel_names=(...).

tdi_custom = wvf.compute_tdi_tlt(channel_names=("C1", "C2", "C3"))

Transition policy

The typed API is the canonical interface for new development. Transitional raw-output pathways are retained for one release cycle and will be progressively de-emphasized as TLT-normalized workflows become the default.

Contributing

Development environment

We use uv for dependency management. First clone the repository on your machine and go in the directory:

git clone git@gitlab.in2p3.fr:servignat/jaxtfwaves.git
cd jaxtfwaves

Create the lockfile and synchronize dependencies (including development tools):

uv sync --all-extras --group dev

We recommend you install pre-commit hooks to detect errors before you even commit.

uv run pre-commit install

Then you can run any script using

uv run python your_script.py

Lint and format checks are handled with Ruff:

uv run ruff check src tests
uv run ruff format --check src tests

Alternative local installation

As an alternative to uv-managed environments, you can install the package in development mode when you are in the jaxtfwaves directory:

pip install -e .            # core only
pip install -e ".[plugins]" # core + plugins