pespace.inference._interface_bilby_jax

Bilby interface for the jax-based Whittle likelihood.

This module is kept separate from pespace.inference.interface_bilby so that importing the taichi interface does not pull in jax (and, through the waveform interfaces, ripplegw). LikelihoodBilbyInterfaceJax is re-exported lazily from pespace.inference.interface_bilby for backwards compatibility.

Classes

LikelihoodBilbyInterfaceJax(waveform, ...[, ...])

Bilby-compatible likelihood class using the jax backend.
class pespace.inference._interface_bilby_jax.LikelihoodBilbyInterfaceJax(waveform, detector, channels, parameter_conversion=None)[source]

Bases: WhittleLikelihoodJax, Likelihood

Bilby-compatible likelihood class using the jax backend.

This class combines the jax-based Whittle likelihood computation in WhittleLikelihoodJax with the bilby bilby.core.likelihood.Likelihood interface, so that the external samplers supported in bilby can be used for parameter estimation with the jax backend.

The parameters are passed to log_likelihood() as a dictionary (the new bilby parameters-as-argument convention), which is converted to the namedtuple expected by log_likelihood(); keys that are not required by the waveform/response computation are ignored.

Parameters:

  • waveform (WaveformRippleBase) – Waveform interface instance for waveform generation. See WhittleLikelihoodJax.

  • detector (InterferometerAntenna or tuple of InterferometerAntenna) – Detector instance(s) configured with backend='jax'.

  • channels (tuple of str) – TDI channels used in the likelihood computation, a subset of (‘A’, ‘E’, ‘T’).

  • parameter_conversion (callable, optional) – Optional callable that takes a dict of parameters and returns a dict. If given, it is applied to the parameters dict in log_likelihood() before constructing the namedtuple expected by the jax backend.

log_likelihood(parameters)[source]

Calculate the log-likelihood value for the given parameters.

Parameters:

parameters (dict) – Dictionary of the parameter names and values. Only the keys required by the waveform and the detector response computation are used, other keys are ignored. If parameter_conversion is set, it is applied to parameters first.

Returns:

The log-likelihood value, summed over all detectors and channels. The evaluation of the underlying jax computation is jit-compiled in __init__.

Return type:

float