pespace.inference.likelihood_jax

Likelihood computation with jax.

Functions

noise_weighted_inner_product(observed_data, ...)

The noise-weighted inner product under the Gaussian stationary noise assumption (Whittle likelihood kernel -2 * df * sum(|data-model|^2 / psd)).

Classes

DetectorData(data, psd, df, channel_indices)

Static per-detector data for the likelihood.

WhittleLikelihoodJax(waveform, detector, ...)

The Whittle likelihood computation with jax.
class pespace.inference.likelihood_jax.DetectorData(data, psd, df, channel_indices)[source]

Bases: NamedTuple

Static per-detector data for the likelihood.

Parameters:

  • data (Array)

  • psd (Array)

  • df (float)

  • channel_indices (tuple | None)

data: Array

Alias for field number 0

psd: Array

Alias for field number 1

df: float

Alias for field number 2

channel_indices: tuple | None

Alias for field number 3

_asdict()

Return a new dict which maps field names to their values.

_field_defaults = {}
_fields = ('data', 'psd', 'df', 'channel_indices')
classmethod _make(iterable)

Make a new DetectorData object from a sequence or iterable

_replace(**kwds)

Return a new DetectorData object replacing specified fields with new values

pespace.inference.likelihood_jax.noise_weighted_inner_product(observed_data, response_data, psd, df)[source]

The noise-weighted inner product under the Gaussian stationary noise assumption (Whittle likelihood kernel -2 * df * sum(|data-model|^2 / psd)).

Parameters:

  • observed_data (jnp.ndarray) – Observed frequency-domain TDI channel data, modeled response, and noise power spectral density, each of shape (N_freq,) or (N_channels, N_freq).

  • response_data (jnp.ndarray) – Observed frequency-domain TDI channel data, modeled response, and noise power spectral density, each of shape (N_freq,) or (N_channels, N_freq).

  • psd (jnp.ndarray) – Observed frequency-domain TDI channel data, modeled response, and noise power spectral density, each of shape (N_freq,) or (N_channels, N_freq).

  • df (float) – Frequency bin spacing in Hz.

Returns:

Value for the Whittle log-likelihood, summed over all axes.

Return type:

float

class pespace.inference.likelihood_jax.WhittleLikelihoodJax(waveform, detector, channels)[source]

Bases: object

The Whittle likelihood computation with jax.

__init__(waveform, detector, channels)[source]
Parameters:

  • waveform (WaveformRippleBase) – Waveform interface instance for waveform generation, the function of waveform computation should be jit-able and grad-able. Its frequencies should match the frequency_samples of the TDI data of all detectors.

  • detector (InterferometerAntenna or tuple of InterferometerAntenna) – Detector instance(s) for computing detector responses, should be configured with backend='jax', with frequency-domain TDI data and noise power spectral densities properly set.

  • channels (tuple of str) – TDI channels used in likelihood computation. Must be a subset of (‘A’, ‘E’, ‘T’). Common choices are (‘A’, ‘E’, ‘T’) or (‘A’, ‘E’).

Raises:

  • TypeError – If waveform is not a WaveformRippleBase instance.

  • ValueError –

    • If any detector does not use backend='jax'; - If any channel in channels is not one of orthogonal ‘A’, ‘E’, or ‘T’; - If any requested channel is not provided by the TDI combination model; - If the frequency-domain data or noise power spectral density of a detector is not set or does not contain all requested channels.

Return type:

None

Notes

  • The waveform computation function returns a dict: {"plus": hp, "cross": hc, "tf": tf} for a single-mode waveform, or a mode-keyed dict for higher modes.

  • The coalescence time shift is applied in the response computation. Thus, the coalescence time should be set to zero in the waveform computation.

_static_data: tuple[DetectorData, ...]
_prepare_static_data()[source]

Prepare the static data of detectors for likelihood computation, facilitating capturing these data via closure in jit-compile of log_likelihood().

Observed frequency-domain TDI data and PSD data are stacked with the shape (N_channels, N_freq) and ordered as the requested channels. Along with the delta_frequency, and the indices for selecting the requested channels from the TDI combination output (None when no slicing is needed), are stored in a namedtuple for each detector.

The namedtuples for all detectors are collected in a tuple and set as the _static_data attribute.

Raises:

ValueError –

  • If any requested channel is not a label provided in the TDI combination model; - If the frequency-domain data or noise power spectral density is not set or does not contain all requested channels.

log_likelihood(params)[source]

log-likelihood evaluation with jax.

Parameters:

params (tuple) – Namedtuple of parameters, including ecliptic_longitude, ecliptic_latitude, polarization, coalescence_time for response computation, and parameters for waveform which depends on the waveform model.

Returns:

The value of log-likelihood, summed over all detectors and channels.

Return type:

float

Notes

The waveform and the per-stage outputs are wrapped in jax.lax.optimization_barrier (which natively accepts pytrees) to prevent XLA from fusing the whole pipeline into a single kernel, which on the CPU backend produces much slower code (~6x slower without the barriers).