jaxwave API Reference

This section provides the auto-generated reference documentation for JAX-based waveform generation, JIT compilation, shape-bucketing, and inference modules in jaxwave.

For physical coverage and execution restrictions, see model support. For numerical qualification, see accuracy and validation.

---

Core Waveform Functions

jaxwave.get_fd_waveform(approximant, mass1, mass2, chi1=_UNSET, chi2=_UNSET, spin1x=0.0, spin1y=0.0, spin1z=None, spin2x=0.0, spin2y=0.0, spin2z=None, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, lambda1=None, lambda2=None, f_ref=None, phasing_convention=None, reference_frame=None, f_lower=None, f_final=None, delta_f=None, sample_frequencies=None, modes=None, dtype=jnp.float64, frequency_support=None, **kwargs)

Generate frequency-domain gravitational-wave polarizations h_plus(f) and h_cross(f) in JAX.

Source code in src/jaxwave/waveforms.py

def get_fd_waveform(
    approximant: str,
    mass1: Union[float, jnp.ndarray],
    mass2: Union[float, jnp.ndarray],
    chi1: Union[float, jnp.ndarray, object] = _UNSET,
    chi2: Union[float, jnp.ndarray, object] = _UNSET,
    spin1x: Union[float, jnp.ndarray] = 0.0,
    spin1y: Union[float, jnp.ndarray] = 0.0,
    spin1z: Optional[Union[float, jnp.ndarray]] = None,
    spin2x: Union[float, jnp.ndarray] = 0.0,
    spin2y: Union[float, jnp.ndarray] = 0.0,
    spin2z: Optional[Union[float, jnp.ndarray]] = None,
    distance: Union[float, jnp.ndarray] = 100.0,
    inclination: Union[float, jnp.ndarray] = 0.0,
    phic: Union[float, jnp.ndarray] = 0.0,
    tc: Union[float, jnp.ndarray] = 0.0,
    lambda1: Optional[Union[float, jnp.ndarray]] = None,
    lambda2: Optional[Union[float, jnp.ndarray]] = None,
    f_ref: Optional[Union[float, jnp.ndarray]] = None,
    phasing_convention: Optional[str] = None,
    reference_frame: Optional[str] = None,
    f_lower: Optional[float] = None,
    f_final: Optional[float] = None,
    delta_f: Optional[float] = None,
    sample_frequencies: Optional[jnp.ndarray] = None,
    modes: Optional[Sequence[Tuple[int, int]]] = None,
    dtype: jnp.dtype = jnp.float64,
    frequency_support: Optional[str] = None,
    **kwargs,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Generate frequency-domain gravitational-wave polarizations h_plus(f) and h_cross(f) in JAX."""
    out = _get_fd_waveform_raw(
        approximant=approximant,
        mass1=mass1,
        mass2=mass2,
        chi1=chi1,
        chi2=chi2,
        spin1x=spin1x,
        spin1y=spin1y,
        spin1z=spin1z,
        spin2x=spin2x,
        spin2y=spin2y,
        spin2z=spin2z,
        distance=distance,
        inclination=inclination,
        phic=phic,
        tc=tc,
        lambda1=lambda1,
        lambda2=lambda2,
        f_ref=f_ref,
        phasing_convention=phasing_convention,
        reference_frame=reference_frame,
        f_lower=f_lower,
        f_final=f_final,
        delta_f=delta_f,
        sample_frequencies=sample_frequencies,
        frequency_support=frequency_support,
        modes=modes,
        dtype=dtype,
        **kwargs,
    )
    if isinstance(out, tuple) and len(out) >= 2 and isinstance(out[0], jnp.ndarray) and isinstance(out[1], jnp.ndarray):
        hp, hc = out[0], out[1]
        f = sample_frequencies
        if f is None and delta_f is not None and f_lower is not None and f_final is not None:
            f = jnp.arange(f_lower, f_final + delta_f / 2.0, delta_f, dtype=dtype)
        if f is not None:
            f_eval = as_array(f, dtype=hp.real.dtype)
            zero_mask = f_eval <= 0.0
            if f_lower is not None:
                zero_mask = zero_mask | (f_eval < f_lower)
            if f_final is not None:
                zero_mask = zero_mask | (f_eval >= f_final)
            hp = jnp.where(zero_mask, jnp.zeros_like(hp), hp)
            hc = jnp.where(zero_mask, jnp.zeros_like(hc), hc)
        if len(out) == 2:
            return hp, hc
        return (hp, hc) + out[2:]
    return out

jaxwave.WaveformGenerator

Class-based interface for gravitational waveform generation.

Configured with static waveform parameters and grid specification, allowing repeated waveform evaluation across varying binary parameters.

Source code in src/jaxwave/waveforms.py

class WaveformGenerator:
    """Class-based interface for gravitational waveform generation.

    Configured with static waveform parameters and grid specification,
    allowing repeated waveform evaluation across varying binary parameters.
    """

    def __init__(
        self,
        approximant: str = "IMRPhenomD",
        f_lower: Optional[float] = None,
        f_final: Optional[float] = None,
        delta_f: Optional[float] = None,
        sample_frequencies: Optional[Any] = None,
        modes: Optional[Sequence[Tuple[int, int]]] = None,
        use_jit: bool = False,
        **default_kwargs,
    ):
        self.approximant = approximant
        self.f_lower = f_lower
        self.f_final = f_final
        self.delta_f = delta_f
        self.sample_frequencies = sample_frequencies
        self.modes = modes
        self.use_jit = use_jit
        self.default_kwargs = default_kwargs
        self._compiled_fd_fn = None
        self._compiled_param_order = None

    def get_compiled_fd_waveform(
        self,
        parameter_order: Tuple[str, ...] = ("mass1", "mass2", "spin1z", "spin2z", "distance", "inclination", "phic", "tc"),
    ):
        """Return a cached @jax.jit compiled function evaluating (hp, hc) for an ordered parameter vector.

        The callable closes over grid specifications and approximant metadata.
        Repeated calls with compatible parameter shapes and dtypes can reuse
        compilation; new signatures may require recompilation.
        """
        if self._compiled_fd_fn is not None and self._compiled_param_order == parameter_order:
            return self._compiled_fd_fn

        approx = self.approximant
        f_low = self.f_lower
        f_fin = self.f_final
        df = self.delta_f
        s_freqs = self.sample_frequencies
        m_modes = self.modes
        d_kwargs = dict(self.default_kwargs)

        def _raw_eval(*param_values):
            kw = dict(d_kwargs)
            for k, val in zip(parameter_order, param_values):
                kw[k] = val
            return get_fd_waveform(
                approximant=approx,
                f_lower=f_low,
                f_final=f_fin,
                delta_f=df,
                sample_frequencies=s_freqs,
                modes=m_modes,
                **kw,
            )

        jitted_eval = jax.jit(_raw_eval)
        self._compiled_fd_fn = jitted_eval
        self._compiled_param_order = parameter_order
        return jitted_eval

    def generate_fd_waveform(
        self,
        parameters: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """Generate frequency-domain polarizations for given parameters."""
        params = dict(self.default_kwargs)
        if parameters is not None:
            params.update(parameters)
        params.update(kwargs)

        name_map = {
            "mass_1": "mass1",
            "mass_2": "mass2",
            "chi_1": "chi1",
            "chi_2": "chi2",
            "spin_1z": "spin1z",
            "spin_2z": "spin2z",
            "spin_1x": "spin1x",
            "spin_1y": "spin1y",
            "spin_2x": "spin2x",
            "spin_2y": "spin2y",
            "luminosity_distance": "distance",
            "theta_jn": "inclination",
            "phase": "phic",
            "geocent_time": "tc",
            "lambda_1": "lambda1",
            "lambda_2": "lambda2",
        }
        normalized_params = {}
        for k, v in params.items():
            norm_k = name_map.get(k, k)
            normalized_params[norm_k] = v

        if self.use_jit:
            # Check if parameters match standard order for cached JIT execution
            std_order = ("mass1", "mass2", "chi1", "chi2", "spin1z", "spin2z", "distance", "inclination", "phic", "tc", "lambda1", "lambda2", "quad_monopole1", "quad_monopole2")
            if all(k in normalized_params for k in ("mass1", "mass2")):
                order = tuple(k for k in std_order if k in normalized_params)
                compiled_fn = self.get_compiled_fd_waveform(parameter_order=order)
                args = tuple(jnp.asarray(normalized_params[k], dtype=jnp.float64) for k in order)
                return compiled_fn(*args)

        call_kwargs = {
            "approximant": normalized_params.pop("approximant", self.approximant),
            "f_lower": normalized_params.pop("f_lower", self.f_lower),
            "f_final": normalized_params.pop("f_final", self.f_final),
            "delta_f": normalized_params.pop("delta_f", self.delta_f),
            "sample_frequencies": normalized_params.pop("sample_frequencies", self.sample_frequencies),
            "modes": normalized_params.pop("modes", self.modes),
        }
        call_kwargs.update(normalized_params)

        return get_fd_waveform(**call_kwargs)


    def generate_td_waveform(
        self,
        parameters: Optional[Dict[str, Any]] = None,
        **kwargs,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """Generate time-domain polarizations for given parameters."""
        params = dict(self.default_kwargs)
        if parameters is not None:
            params.update(parameters)
        params.update(kwargs)

        name_map = {
            "mass_1": "mass1",
            "mass_2": "mass2",
            "spin_1z": "spin1z",
            "spin_2z": "spin2z",
            "spin_1x": "spin1x",
            "spin_1y": "spin1y",
            "spin_2x": "spin2x",
            "spin_2y": "spin2y",
            "luminosity_distance": "distance",
            "theta_jn": "inclination",
            "phase": "phic",
            "geocent_time": "tc",
            "lambda_1": "lambda1",
            "lambda_2": "lambda2",
        }
        normalized_params = {}
        for k, v in params.items():
            norm_k = name_map.get(k, k)
            normalized_params[norm_k] = v

        call_kwargs = {
            "approximant": normalized_params.pop("approximant", self.approximant),
            "f_lower": normalized_params.pop("f_lower", self.f_lower if self.f_lower is not None else 20.0),
            "f_final": normalized_params.pop("f_final", self.f_final),
        }
        call_kwargs.update(normalized_params)

        return get_td_waveform(**call_kwargs)

    def __call__(self, *args, **kwargs):
        if args and isinstance(args[0], dict):
            return self.generate_fd_waveform(parameters=args[0], **kwargs)
        return self.generate_fd_waveform(**kwargs)

generate_fd_waveform(parameters=None, **kwargs)

Generate frequency-domain polarizations for given parameters.

Source code in src/jaxwave/waveforms.py

def generate_fd_waveform(
    self,
    parameters: Optional[Dict[str, Any]] = None,
    **kwargs,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Generate frequency-domain polarizations for given parameters."""
    params = dict(self.default_kwargs)
    if parameters is not None:
        params.update(parameters)
    params.update(kwargs)

    name_map = {
        "mass_1": "mass1",
        "mass_2": "mass2",
        "chi_1": "chi1",
        "chi_2": "chi2",
        "spin_1z": "spin1z",
        "spin_2z": "spin2z",
        "spin_1x": "spin1x",
        "spin_1y": "spin1y",
        "spin_2x": "spin2x",
        "spin_2y": "spin2y",
        "luminosity_distance": "distance",
        "theta_jn": "inclination",
        "phase": "phic",
        "geocent_time": "tc",
        "lambda_1": "lambda1",
        "lambda_2": "lambda2",
    }
    normalized_params = {}
    for k, v in params.items():
        norm_k = name_map.get(k, k)
        normalized_params[norm_k] = v

    if self.use_jit:
        # Check if parameters match standard order for cached JIT execution
        std_order = ("mass1", "mass2", "chi1", "chi2", "spin1z", "spin2z", "distance", "inclination", "phic", "tc", "lambda1", "lambda2", "quad_monopole1", "quad_monopole2")
        if all(k in normalized_params for k in ("mass1", "mass2")):
            order = tuple(k for k in std_order if k in normalized_params)
            compiled_fn = self.get_compiled_fd_waveform(parameter_order=order)
            args = tuple(jnp.asarray(normalized_params[k], dtype=jnp.float64) for k in order)
            return compiled_fn(*args)

    call_kwargs = {
        "approximant": normalized_params.pop("approximant", self.approximant),
        "f_lower": normalized_params.pop("f_lower", self.f_lower),
        "f_final": normalized_params.pop("f_final", self.f_final),
        "delta_f": normalized_params.pop("delta_f", self.delta_f),
        "sample_frequencies": normalized_params.pop("sample_frequencies", self.sample_frequencies),
        "modes": normalized_params.pop("modes", self.modes),
    }
    call_kwargs.update(normalized_params)

    return get_fd_waveform(**call_kwargs)

generate_td_waveform(parameters=None, **kwargs)

Generate time-domain polarizations for given parameters.

Source code in src/jaxwave/waveforms.py

def generate_td_waveform(
    self,
    parameters: Optional[Dict[str, Any]] = None,
    **kwargs,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Generate time-domain polarizations for given parameters."""
    params = dict(self.default_kwargs)
    if parameters is not None:
        params.update(parameters)
    params.update(kwargs)

    name_map = {
        "mass_1": "mass1",
        "mass_2": "mass2",
        "spin_1z": "spin1z",
        "spin_2z": "spin2z",
        "spin_1x": "spin1x",
        "spin_1y": "spin1y",
        "spin_2x": "spin2x",
        "spin_2y": "spin2y",
        "luminosity_distance": "distance",
        "theta_jn": "inclination",
        "phase": "phic",
        "geocent_time": "tc",
        "lambda_1": "lambda1",
        "lambda_2": "lambda2",
    }
    normalized_params = {}
    for k, v in params.items():
        norm_k = name_map.get(k, k)
        normalized_params[norm_k] = v

    call_kwargs = {
        "approximant": normalized_params.pop("approximant", self.approximant),
        "f_lower": normalized_params.pop("f_lower", self.f_lower if self.f_lower is not None else 20.0),
        "f_final": normalized_params.pop("f_final", self.f_final),
    }
    call_kwargs.update(normalized_params)

    return get_td_waveform(**call_kwargs)

get_compiled_fd_waveform(parameter_order=('mass1', 'mass2', 'spin1z', 'spin2z', 'distance', 'inclination', 'phic', 'tc'))

Return a cached @jax.jit compiled function evaluating (hp, hc) for an ordered parameter vector.

The callable closes over grid specifications and approximant metadata. Repeated calls with compatible parameter shapes and dtypes can reuse compilation; new signatures may require recompilation.

Source code in src/jaxwave/waveforms.py

def get_compiled_fd_waveform(
    self,
    parameter_order: Tuple[str, ...] = ("mass1", "mass2", "spin1z", "spin2z", "distance", "inclination", "phic", "tc"),
):
    """Return a cached @jax.jit compiled function evaluating (hp, hc) for an ordered parameter vector.

    The callable closes over grid specifications and approximant metadata.
    Repeated calls with compatible parameter shapes and dtypes can reuse
    compilation; new signatures may require recompilation.
    """
    if self._compiled_fd_fn is not None and self._compiled_param_order == parameter_order:
        return self._compiled_fd_fn

    approx = self.approximant
    f_low = self.f_lower
    f_fin = self.f_final
    df = self.delta_f
    s_freqs = self.sample_frequencies
    m_modes = self.modes
    d_kwargs = dict(self.default_kwargs)

    def _raw_eval(*param_values):
        kw = dict(d_kwargs)
        for k, val in zip(parameter_order, param_values):
            kw[k] = val
        return get_fd_waveform(
            approximant=approx,
            f_lower=f_low,
            f_final=f_fin,
            delta_f=df,
            sample_frequencies=s_freqs,
            modes=m_modes,
            **kw,
        )

    jitted_eval = jax.jit(_raw_eval)
    self._compiled_fd_fn = jitted_eval
    self._compiled_param_order = parameter_order
    return jitted_eval

jaxwave.get_td_waveform(approximant='IMRPhenomD', mass1=30.0, mass2=25.0, chi1=_UNSET, chi2=_UNSET, spin1x=0.0, spin1y=0.0, spin1z=None, spin2x=0.0, spin2y=0.0, spin2z=None, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, lambda1=None, lambda2=None, delta_t=1.0 / 2048.0, f_lower=20.0, f_final=None, duration=None, t_pad=0.2, center_coalescence=True, return_times=False, **kwargs)

Generate conditioned time-domain gravitational-wave polarizations h_plus(t) and h_cross(t) in JAX.

Literature: - Allen et al. (2012), Phys. Rev. D 85, 122006 (arXiv:gr-qc/0509116) - Buonanno et al. (2009), Phys. Rev. D 80, 084043 (arXiv:0907.0700)

The IMRPhenomT family is evaluated natively in time; f_ref fixes its orbital phase. Other models use the inverse Fourier path, which conserves signal energy: int |h(t)|^2 dt = 2 * int |h_tilde(f)|^2 df

Source code in src/jaxwave/waveforms.py

def get_td_waveform(
    approximant: str = "IMRPhenomD",
    mass1: Union[float, jnp.ndarray] = 30.0,
    mass2: Union[float, jnp.ndarray] = 25.0,
    chi1: Union[float, jnp.ndarray, object] = _UNSET,
    chi2: Union[float, jnp.ndarray, object] = _UNSET,
    spin1x: Union[float, jnp.ndarray] = 0.0,
    spin1y: Union[float, jnp.ndarray] = 0.0,
    spin1z: Optional[Union[float, jnp.ndarray]] = None,
    spin2x: Union[float, jnp.ndarray] = 0.0,
    spin2y: Union[float, jnp.ndarray] = 0.0,
    spin2z: Optional[Union[float, jnp.ndarray]] = None,
    distance: Union[float, jnp.ndarray] = 100.0,
    inclination: Union[float, jnp.ndarray] = 0.0,
    phic: Union[float, jnp.ndarray] = 0.0,
    tc: Union[float, jnp.ndarray] = 0.0,
    lambda1: Optional[Union[float, jnp.ndarray]] = None,
    lambda2: Optional[Union[float, jnp.ndarray]] = None,
    delta_t: float = 1.0 / 2048.0,
    f_lower: float = 20.0,
    f_final: Optional[float] = None,
    duration: Optional[float] = None,
    t_pad: float = 0.2,
    center_coalescence: bool = True,
    return_times: bool = False,
    **kwargs,
) -> Union[
    Tuple[jnp.ndarray, jnp.ndarray],
    Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
]:
    """Generate conditioned time-domain gravitational-wave polarizations h_plus(t) and h_cross(t) in JAX.

    Literature:
    - Allen et al. (2012), Phys. Rev. D 85, 122006 (arXiv:gr-qc/0509116)
    - Buonanno et al. (2009), Phys. Rev. D 80, 084043 (arXiv:0907.0700)

    The IMRPhenomT family is evaluated natively in time; f_ref fixes its orbital phase.
    Other models use the inverse Fourier path, which conserves signal energy:
        int |h(t)|^2 dt = 2 * int |h_tilde(f)|^2 df
    """
    key = approximant.lower().replace("_", "").replace("-", "")
    if key in ("imrphenomt", "imrphenomthm", "imrphenomtp", "imrphenomtphm"):
        from ._phenomt_td import generate_td

        return generate_td(
            approximant=key, mass1=mass1, mass2=mass2,
            chi1=None if chi1 is _UNSET else chi1,
            chi2=None if chi2 is _UNSET else chi2,
            spin1z=spin1z, spin2z=spin2z, spin1x=spin1x, spin1y=spin1y,
            spin2x=spin2x, spin2y=spin2y, distance=distance,
            inclination=inclination, phic=phic, tc=tc, lambda1=lambda1, lambda2=lambda2,
            delta_t=delta_t, f_lower=f_lower, f_final=f_final, duration=duration,
            t_pad=t_pad, center_coalescence=center_coalescence,
            return_times=return_times, **kwargs,
        )

    m1_t = as_array(mass1, dtype=jnp.float64)
    m2_t = as_array(mass2, dtype=jnp.float64)

    if duration is None:
        m_chirp = component_masses_to_chirp_mass(m1_t, m2_t)
        m_chirp_min = float(jnp.min(m_chirp))
        m_chirp_sec = m_chirp_min * MSUN_SEC
        tau0 = (5.0 / 256.0) * (m_chirp_sec ** (-5.0 / 3.0)) * ((PI * f_lower) ** (-8.0 / 3.0))
        est_duration = max(2.0, 1.2 * tau0 + 2.0 * t_pad + 1.0)
        n_samples_raw = int(math.ceil(est_duration / delta_t))
        n_samples = 1 << (n_samples_raw - 1).bit_length()
        dur = n_samples * delta_t
    else:
        dur = float(duration)
        n_samples = int(math.ceil(dur / delta_t))

    df = 1.0 / dur
    f_nyquist = 1.0 / (2.0 * delta_t)
    f_max = f_final if f_final is not None else f_nyquist

    freqs = jnp.arange(0, n_samples // 2 + 1, dtype=jnp.float64) * df
    valid_mask = (freqs >= f_lower) & (freqs <= f_max)
    active_freqs = freqs[valid_mask]

    hp_fd_active, hc_fd_active = get_fd_waveform(
        approximant=approximant,
        mass1=mass1,
        mass2=mass2,
        chi1=chi1,
        chi2=chi2,
        spin1x=spin1x,
        spin1y=spin1y,
        spin1z=spin1z,
        spin2x=spin2x,
        spin2y=spin2y,
        spin2z=spin2z,
        distance=distance,
        inclination=inclination,
        phic=phic,
        tc=0.0,
        lambda1=lambda1,
        lambda2=lambda2,
        sample_frequencies=active_freqs,
        **kwargs,
    )

    is_batched = m1_t.ndim >= 1
    if is_batched:
        batch_shape = m1_t.shape
        full_shape = batch_shape + (len(freqs),)
        hp_fd = jnp.zeros(full_shape, dtype=jnp.complex128).at[..., valid_mask].set(hp_fd_active)
        hc_fd = jnp.zeros(full_shape, dtype=jnp.complex128).at[..., valid_mask].set(hc_fd_active)
    else:
        hp_fd = jnp.zeros(len(freqs), dtype=jnp.complex128).at[valid_mask].set(hp_fd_active)
        hc_fd = jnp.zeros(len(freqs), dtype=jnp.complex128).at[valid_mask].set(hc_fd_active)

    df_taper = max(2.0, 0.1 * f_lower)
    f_start = f_lower
    f_end = f_lower + df_taper
    z = jnp.clip((freqs - f_start) / (f_end - f_start), 1e-7, 1.0 - 1e-7)
    sigma = (1.0 / z) + (1.0 / (z - 1.0))
    w_smooth = 1.0 / (jnp.exp(jnp.clip(sigma, -50.0, 50.0)) + 1.0)
    w_taper = jnp.where(freqs <= f_start, 0.0, w_smooth)
    w_taper = jnp.where(freqs >= f_end, 1.0, w_taper)

    if is_batched:
        w_taper = jnp.expand_dims(w_taper, 0)

    hp_fd = hp_fd * w_taper
    hc_fd = hc_fd * w_taper

    scale_factor = 1.0 / delta_t
    hp_td = scale_factor * jnp.fft.irfft(hp_fd, n=n_samples, axis=-1)
    hc_td = scale_factor * jnp.fft.irfft(hc_fd, n=n_samples, axis=-1)

    if center_coalescence:
        tc_val = float(tc) if isinstance(tc, (int, float)) else float(jnp.asarray(tc).reshape(-1)[0])
        target_t = (dur - t_pad) + tc_val
        shift_samples = int(round(target_t / delta_t)) % n_samples
        hp_td = jnp.roll(hp_td, shift=shift_samples, axis=-1)
        hc_td = jnp.roll(hc_td, shift=shift_samples, axis=-1)

    times = jnp.arange(n_samples, dtype=jnp.float64) * delta_t
    if center_coalescence:
        times = times - (dur - t_pad)

    if return_times:
        return hp_td, hc_td, times
    return hp_td, hc_td

Native Time-Domain Models

See the IMRPhenomT guide for examples and conventions.

jaxwave.IMRPhenomT

Aligned-spin dominant (2,2) IMRPhenomT waveform in JAX.

Source code in src/jaxwave/approximants/imrphenomt.py

class IMRPhenomT:
    """Aligned-spin dominant ``(2,2)`` IMRPhenomT waveform in JAX."""

    @staticmethod
    def _state(m1, m2, chi1, chi2):
        M = m1 + m2
        eta = m1 * m2 / M**2
        delta = (m1 - m2) / M
        # The fit variable is sTotR, normalized by m1^2+m2^2=1-2 eta.
        S = (chi1 * (m1 / M) ** 2 + chi2 * (m2 / M) ** 2) / (1 - 2 * eta)
        dchi = chi1 - chi2
        # The helper provides the published 2017 remnant and Kerr fits.
        af = _fit_any(("IMRPhenomX_FinalSpin2017", "final_spin"), eta, chi1, chi2, delta)
        mf = _fit_any(("IMRPhenomX_FinalMass2017", "final_mass"), eta, chi1, chi2, delta)
        af, mf = _t(af, eta), _t(mf, eta)
        return M, eta, delta, S, dchi, af, mf

    @staticmethod
    def _coeff(m1, m2, chi1, chi2, final_spin=None):
        """Reconstruct the carrier, optionally using a precessing remnant spin."""
        M, eta, delta, S, dchi, af, mf = IMRPhenomT._state(m1, m2, chi1, chi2)
        t0 = _t(_fit_any(("IMRPhenomT_Inspiral_TaylorT3_t0",), eta, S, dchi, delta), eta)
        freq_cp = [_t(_fit_any(("IMRPhenomT_Inspiral_Freq_CP",), eta, S, dchi, delta, i), eta) for i in range(1, 6)]
        peak_amp = _t(_fit_any(("IMRPhenomT_PeakAmp",), eta, S, dchi, delta, 22), eta)
        amp_mid = _t(_fit_any(("IMRPhenomT_Intermediate_Amp_CP1",), eta, S, dchi, delta, 22), eta)
        tshift = jnp.zeros_like(eta)
        c3a = _t(_fit_any(("IMRPhenomT_Ringdown_Amp_C3",), eta, S, dchi, 22), eta)
        omega_peak = _t(_fit_any(("IMRPhenomT_PeakFrequency",), eta, S, dchi, delta, 22), eta)
        c2f = _t(_fit_any(("IMRPhenomT_RD_Freq_C2",), eta, S, dchi, delta, 22), eta)
        c3f = _t(_fit_any(("IMRPhenomT_RD_Freq_C3",), eta, S, dchi, delta, 22), eta)
        mid_freq = _t(_fit_any(("IMRPhenomT_Intermediate_Freq_CP1",), eta, S, dchi, delta, 22), eta)
        # The QNM fits are in units of final mass; the ansatz time is in
        # units of total mass, so only division by mf is required here.
        fr = _t(_fit_any(("IMRPhenomT_fringfit",), af, 22), eta) / mf
        fd = _t(_fit_any(("IMRPhenomT_fdampfit",), af, 22), eta) / mf
        fd2 = _t(_fit_any(("IMRPhenomT_fdampn2fit",), af, 22), eta) / mf
        # The NNLO precessing remnant changes only ringdown (TPHM Sec. II C).
        # Preserve the calibrated aligned-spin inspiral and merger joins.
        rd_fr, rd_fd, rd_fd2 = fr, fd, fd2
        if final_spin is not None:
            rd_af = jnp.asarray(final_spin, dtype=eta.dtype) + jnp.zeros_like(eta)
            rd_fr = _t(_fit_any(("IMRPhenomT_fringfit",), rd_af, 22), eta) / mf
            rd_fd = _t(_fit_any(("IMRPhenomT_fdampfit",), rd_af, 22), eta) / mf
            rd_fd2 = _t(_fit_any(("IMRPhenomT_fdampn2fit",), rd_af, 22), eta) / mf

        th = jnp.asarray((0.33, 0.45, 0.55, 0.65, 0.75, 0.82), dtype=eta.dtype)
        qmat = jnp.stack([th ** (8 + j) for j in range(6)], -1)
        tearly = -5.0 / (eta * th[0] ** 8)
        theta0 = jnp.power(jnp.maximum(eta * (t0 - tearly) / 5.0, 1.0e-30), -1 / 8)
        vals = jnp.stack([_poly_pn_frequency(theta0, eta, chi1, chi2, delta)] + freq_cp, -1)
        pvals = _poly_pn_frequency(th, eta[..., None], chi1[..., None], chi2[..., None], delta[..., None])
        rhs = 4.0 / th**3 * (vals - pvals)
        wc = jnp.linalg.solve(jnp.broadcast_to(qmat, rhs.shape[:-1] + qmat.shape), rhs[..., None])[..., 0]

        # Merger frequency coefficients, Eqs. (28-31) of the reference.
        tcut = -26.982976386771437 / eta
        tcp = -5.0 / (eta * jnp.asarray(0.95, dtype=eta.dtype) ** 8)
        wcut = IMRPhenomT._omega_insp(tcut[..., None], eta, delta, chi1, chi2, wc)[..., 0]
        dwcut = _domega_insp(tcut[..., None], eta, chi1, chi2, delta, wc)[..., 0]
        wr = 2 * PI * fr
        alpha = 2 * PI * fd
        c1f = (1 + c3f) * (wr - omega_peak) / (c2f * c3f)
        dwpeak = -c1f * c2f**2 * c3f / (1 + c3f) ** 2 / wr
        acut = jnp.arcsinh(alpha * tcut)
        amer = jnp.arcsinh(alpha * tcp)
        dencut = jnp.sqrt(1 + (alpha * tcut) ** 2)
        b0 = 1 - (wcut / wr) - (1 - omega_peak / wr) - dwpeak / alpha * acut
        b1 = 1 - mid_freq / wr - (1 - omega_peak / wr) - dwpeak / alpha * amer
        b2 = -dwcut / wr - dwpeak / dencut
        amat = jnp.stack(
            (
                jnp.stack((acut**2, acut**3, acut**4), -1),
                jnp.stack((amer**2, amer**3, amer**4), -1),
                jnp.stack((2 * alpha * acut / dencut, 3 * alpha * acut**2 / dencut, 4 * alpha * acut**3 / dencut), -1),
            ),
            -2,
        )
        mc = jnp.linalg.solve(amat, jnp.stack((b0, b1, b2), -1)[..., None])[..., 0]

        # Amplitude reconstruction is kept in the MIT-derived private helper;
        # it consumes analytic piecewise omega and domega callbacks.
        from ._phenomt_amplitude import amplitude_coefficients

        def omega_cb(tt):
            ti = jnp.minimum(tt, tcut[..., None] - 1.0e-10)
            wi = IMRPhenomT._omega_insp(ti, eta, delta, chi1, chi2, wc)
            tm = jnp.clip(tt, tcut[..., None], 0.0)
            x = jnp.arcsinh(alpha[..., None] * tm)
            wm = wr[..., None] * (
                1
                - (
                    1
                    - omega_peak[..., None] / wr[..., None]
                    + x
                    * (
                        dwpeak[..., None] / alpha[..., None]
                        + x * (mc[..., 0, None] + x * (mc[..., 1, None] + x * mc[..., 2, None]))
                    )
                )
            )
            tr = jnp.maximum(tt, jnp.zeros_like(tt))
            er = jnp.exp(-c2f[..., None] * tr)
            wd = wr[..., None] - c1f[..., None] * c2f[..., None] * c3f[..., None] * er / (1 + c3f[..., None] * er)
            return jnp.where(tt < tcut[..., None], wi, jnp.where(tt < 0, wm, wd))

        def domega_cb(tt):
            ti = jnp.minimum(tt, tcut[..., None] - 1.0e-10)
            di = _domega_insp(ti, eta, chi1, chi2, delta, wc)
            tm = jnp.clip(tt, tcut[..., None], 0.0)
            x = jnp.arcsinh(alpha[..., None] * tm)
            dx = alpha[..., None] / jnp.sqrt(1 + (alpha[..., None] * tm) ** 2)
            dqx = (
                dwpeak[..., None] / alpha[..., None]
                + 2 * mc[..., 0, None] * x
                + 3 * mc[..., 1, None] * x**2
                + 4 * mc[..., 2, None] * x**3
            )
            dm = -wr[..., None] * dqx * dx
            tr = jnp.maximum(tt, jnp.zeros_like(tt))
            er = jnp.exp(-c2f[..., None] * tr)
            dd = c1f[..., None] * c2f[..., None] ** 2 * c3f[..., None] * er / (1 + c3f[..., None] * er) ** 2
            return jnp.where(tt < tcut[..., None], di, jnp.where(tt < 0, dm, dd))

        amp_state = amplitude_coefficients(eta, delta, chi1, chi2, fd, fd2, omega_cb, domega_cb, rd_fd, rd_fd2)

        # Analytic phase offsets make all three regions C0 continuous.
        p5 = jnp.asarray(5.0, dtype=eta.dtype) ** (jnp.arange(8, dtype=eta.dtype) / 8)
        theta_i = jnp.power(jnp.maximum(eta * (t0 - tearly), 1.0e-30), -1 / 8)
        phoff_i = (
            _pn_phase(theta_i, eta, delta, chi1, chi2)
            - _phase_insp_raw(tearly[..., None], eta, chi1, chi2, delta, wc, p5)[..., 0]
        )
        phoff_m = (
            phoff_i
            + _phase_insp_raw(tcut[..., None], eta, chi1, chi2, delta, wc, p5)[..., 0]
            - IMRPhenomT._phase_inter_raw(tcut, alpha, wr, omega_peak, dwpeak, mc)
        )
        phoff_r = IMRPhenomT._phase_inter_raw(jnp.zeros_like(eta), alpha, wr, omega_peak, dwpeak, mc) + phoff_m
        return (
            M,
            eta,
            delta,
            S,
            dchi,
            chi1,
            chi2,
            wc,
            amp_state,
            peak_amp,
            amp_mid,
            tshift,
            c3a,
            omega_peak,
            c2f,
            c3f,
            mid_freq,
            fr,
            fd,
            fd2,
            alpha,
            wr,
            c1f,
            dwpeak,
            mc,
            tcut,
            phoff_i,
            phoff_m,
            phoff_r,
            p5,
            2 * PI * rd_fr,
            (1 + c3f) * (2 * PI * rd_fr - omega_peak) / (c2f * c3f),
        )

    @staticmethod
    def _omega_insp(t, eta, delta, s1, s2, wc):
        if t.ndim <= eta.ndim:
            t = t.reshape((1,) * eta.ndim + (-1,))
        th = jnp.power(jnp.maximum(-eta[..., None] * t / 5, 1e-30), -1 / 8)
        return _poly_pn_frequency(th, eta[..., None], s1[..., None], s2[..., None], delta[..., None]) + th**3 / 4 * sum(
            wc[..., j, None] * th ** (8 + j) for j in range(6)
        )

    @staticmethod
    def _phase_inter_raw(t, alpha, wr, op, dwp, mc):
        x = jnp.arcsinh(alpha * t)
        x2, x3, x4 = x * x, x * x * x, x * x * x * x
        root = jnp.sqrt(1 + (alpha * t) ** 2)
        return wr * t * (
            1
            - (
                2 * mc[..., 0]
                + 24 * mc[..., 2]
                + (6 * mc[..., 1] + dwp / alpha) * x
                + (1 - op / wr)
                + (mc[..., 0] + 12 * mc[..., 2]) * x2
                + mc[..., 1] * x3
                + mc[..., 2] * x4
            )
        ) - wr / alpha * root * (
            -dwp / alpha
            - 6 * mc[..., 1]
            - x * (2 * mc[..., 0] + 24 * mc[..., 2])
            - 3 * mc[..., 1] * x2
            - 4 * mc[..., 2] * x3
        )

    @staticmethod
    def _evaluate(u, state):
        (
            M,
            eta,
            delta,
            S,
            dchi,
            s1,
            s2,
            wc,
            amp_state,
            peak,
            amid,
            tshift,
            c3a,
            op,
            c2f,
            c3f,
            midf,
            fr,
            fd,
            fd2,
            alpha,
            wr,
            c1f,
            dwp,
            mc,
            tcut,
            phoi,
            phom,
            phor,
            p5,
            rd_wr,
            rd_c1f,
        ) = state
        ui = jnp.minimum(u, tcut[..., None])
        win = IMRPhenomT._omega_insp(ui, eta, delta, s1, s2, wc)
        um = jnp.clip(u, tcut[..., None], 0.0)
        ur = jnp.maximum(u, jnp.zeros_like(u))
        ex = jnp.exp(-c2f[..., None] * ur)
        wring = rd_wr[..., None] - rd_c1f[..., None] * c2f[..., None] * c3f[..., None] * ex / (1 + c3f[..., None] * ex)
        x = jnp.arcsinh(alpha[..., None] * um)
        wm = wr[..., None] * (
            1
            - (
                1
                - op[..., None] / wr[..., None]
                + x
                * (
                    dwp[..., None] / alpha[..., None]
                    + x * (mc[..., 0, None] + x * (mc[..., 1, None] + x * mc[..., 2, None]))
                )
            )
        )
        omega = jnp.where(u < tcut[..., None], win, jnp.where(u < 0, wm, wring))
        pi = _phase_insp_raw(ui, eta, s1, s2, delta, wc, p5) + phoi[..., None]
        pm = (
            IMRPhenomT._phase_inter_raw(
                um, alpha[..., None], wr[..., None], op[..., None], dwp[..., None], mc[..., None, :]
            )
            + phom[..., None]
        )
        er = jnp.exp(-c2f[..., None] * ur)
        pr = (
            rd_c1f[..., None] * jnp.log((1 + c3f[..., None] * er) / (1 + c3f[..., None]))
            + rd_wr[..., None] * ur
            + phor[..., None]
        )
        phase = jnp.where(u < tcut[..., None], pi, jnp.where(u < 0, pm, pr))
        return omega, phase, amp_state

    @staticmethod
    def _domega_full(u, state):
        """Analytic derivative of the piecewise IMR angular frequency."""
        (
            M,
            eta,
            delta,
            S,
            dchi,
            s1,
            s2,
            wc,
            amp_state,
            peak,
            amid,
            tshift,
            c3a,
            op,
            c2f,
            c3f,
            midf,
            fr,
            fd,
            fd2,
            alpha,
            wr,
            c1f,
            dwp,
            mc,
            tcut,
            phoi,
            phom,
            phor,
            p5,
            rd_wr,
            rd_c1f,
        ) = state
        ui = jnp.minimum(u, tcut[..., None])
        um = jnp.clip(u, tcut[..., None], 0.0)
        ur = jnp.maximum(u, jnp.zeros_like(u))
        di = _domega_insp(ui, eta, s1, s2, delta, wc)
        x = jnp.arcsinh(alpha[..., None] * um)
        dx = alpha[..., None] / jnp.sqrt(1 + (alpha[..., None] * um) ** 2)
        dqx = (
            dwp[..., None] / alpha[..., None]
            + 2 * mc[..., 0, None] * x
            + 3 * mc[..., 1, None] * x**2
            + 4 * mc[..., 2, None] * x**3
        )
        dm = -wr[..., None] * dqx * dx
        er = jnp.exp(-c2f[..., None] * ur)
        dd = rd_c1f[..., None] * c2f[..., None] ** 2 * c3f[..., None] * er / (1 + c3f[..., None] * er) ** 2
        return jnp.where(u < tcut[..., None], di, jnp.where(u < 0, dm, dd))

    def time_domain(
        self,
        times: TensorLike,
        mass1: TensorLike,
        mass2: TensorLike,
        chi1: TensorLike = 0.0,
        chi2: TensorLike = 0.0,
        distance: TensorLike = 100.0,
        inclination: TensorLike = 0.0,
        phic: TensorLike = 0.0,
        tc: TensorLike = 0.0,
        f_ref: TensorLike = 20.0,
        *,
        _final_spin_override=None,
    ) -> tuple[jnp.ndarray, jnp.ndarray]:
        """Return h_plus, h_cross on times in seconds relative to the peak.

        Masses are solar masses, distance is Mpc, and f_ref is Hz. Parameter
        batches precede the shared time axis. This method applies no start mask.
        """
        t = jnp.asarray(times, dtype=jnp.float64)
        m1, m2, c1, c2, d, inc, pc, tcenter, fr = [
            jnp.asarray(v, dtype=t.dtype) for v in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref)
        ]
        m1, m2, c1, c2, d, inc, pc, tcenter, fr = jnp.broadcast_arrays(m1, m2, c1, c2, d, inc, pc, tcenter, fr)
        state = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)
        M = state[0]
        u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
        omega, phase, amp_state = self._evaluate(u, state)
        from ._phenomt_amplitude import amplitude

        amp = amplitude(u, amp_state, lambda tt: self._evaluate(tt, state)[0])
        # f_ref=0 follows the public convention and means the default 20 Hz
        # lower/reference frequency.  Subtract the analytic phase at that
        # frequency, as the author implementation does.
        fphase = jnp.where(fr > 0, fr, jnp.full_like(fr, 20.0))
        tref = self.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=_final_spin_override)
        uref = tref / (M * MSUN_SEC)
        phref = self._evaluate(uref[..., None], state)[1][..., 0]
        ph = phase - phref[..., None] + 2.0 * pc[..., None]
        scale = (M[..., None] * MSUN_SEC) / (d[..., None] * MPC_SI / C_SI)
        y22 = math.sqrt(5.0 / (4.0 * math.pi))
        ci = jnp.cos(inc[..., None])
        pref = -scale * y22 * amp
        return pref * (1 + ci**2) * 0.5 * jnp.cos(ph), pref * ci * jnp.sin(ph)

    forward = time_domain
    __call__ = time_domain

    def frequency(self, times, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
        """Instantaneous gravitational-wave frequency in Hz at peak-relative times."""
        t = jnp.asarray(times, dtype=jnp.float64)
        m1 = jnp.asarray(mass1, dtype=t.dtype)
        m2 = jnp.asarray(mass2, dtype=t.dtype)
        c1 = jnp.asarray(chi1, dtype=t.dtype)
        c2 = jnp.asarray(chi2, dtype=t.dtype)
        m1, m2, c1, c2 = jnp.broadcast_arrays(m1, m2, c1, c2)
        M = m1 + m2
        u = t.reshape((1,) * M.ndim + (-1,)) / (M[..., None] * MSUN_SEC)
        return self._evaluate(u, self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override))[0] / (2 * PI * M[..., None] * MSUN_SEC)

    def time_at_frequency(self, frequency, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
        """Return peak-relative time at a requested instantaneous frequency.

        The bracket is grown from the Newtonian inspiral estimate so that it
        remains useful for low-frequency BNS requests.  Bisection itself is
        stopped from the differentiation graph; an implicit Newton correction
        then restores the derivative through the frequency equation while
        keeping JIT and batched execution fully functional.
        """
        f = jnp.asarray(frequency, dtype=jnp.float64)
        m1 = jnp.asarray(mass1, dtype=f.dtype)
        m2 = jnp.asarray(mass2, dtype=f.dtype)
        c1 = jnp.asarray(chi1, dtype=f.dtype)
        c2 = jnp.asarray(chi2, dtype=f.dtype)
        m1, m2, c1, c2, f = jnp.broadcast_arrays(m1, m2, c1, c2, f)
        M = m1 + m2
        target = 2 * PI * f * M * MSUN_SEC
        st = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)

        # Eager calls give a useful exception.  Under jit these conditions are
        # represented by ``valid`` below, yielding NaN rather than selecting a
        # clipped endpoint when the requested root does not exist.
        concrete = not any(isinstance(v, jax.core.Tracer) for v in (f, m1, m2, c1, c2, _final_spin_override))
        valid = jnp.isfinite(f) & (f > 0) & (target < st[30])
        if concrete:
            if not bool(jnp.all(jnp.isfinite(f) & (f > 0))):
                raise ValueError("frequency must be finite and positive")
            if bool(jnp.any(target >= st[30])):
                raise ValueError("frequency is above the supported IMRPhenomT range")

        # Newtonian leading-order time, buffered as in the author solver.
        lo0 = -5.0 / (256.0 * st[1] * jnp.power(target / 2.0, 8.0 / 3.0)) * 1.2
        hi0 = jnp.zeros_like(target)

        def grow_hi(_, hi):
            om_hi = self._evaluate(hi[..., None], st)[0][..., 0]
            need = om_hi < target
            return jnp.where(need, jnp.where(hi == 0, jnp.ones_like(hi), 2 * hi), hi)

        def grow_lo(_, lo):
            om_lo = self._evaluate(lo[..., None], st)[0][..., 0]
            return jnp.where(om_lo > target, lo * 2.0, lo)

        hi = jax.lax.fori_loop(0, 16, grow_hi, hi0)
        lo = jax.lax.fori_loop(0, 12, grow_lo, lo0)

        def bisect(_, bounds):
            lo_i, hi_i = bounds
            mid = (lo_i + hi_i) / 2
            om = self._evaluate(mid[..., None], st)[0][..., 0]
            below = om < target
            return jnp.where(below, mid, lo_i), jnp.where(below, hi_i, mid)

        lo, hi = jax.lax.fori_loop(0, 80, bisect, (lo, hi))
        root = jax.lax.stop_gradient((lo + hi) / 2)
        om = self._evaluate(root[..., None], st)[0][..., 0]
        dom = self._domega_full(root[..., None], st)[..., 0]
        root = root + (target - om) / jnp.maximum(dom, jnp.finfo(root.dtype).tiny)
        root = jnp.where(valid, root, jnp.nan)
        return root * M * MSUN_SEC

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, *, _final_spin_override=None)

Return h_plus, h_cross on times in seconds relative to the peak.

Masses are solar masses, distance is Mpc, and f_ref is Hz. Parameter batches precede the shared time axis. This method applies no start mask.

Source code in src/jaxwave/approximants/imrphenomt.py

def time_domain(
    self,
    times: TensorLike,
    mass1: TensorLike,
    mass2: TensorLike,
    chi1: TensorLike = 0.0,
    chi2: TensorLike = 0.0,
    distance: TensorLike = 100.0,
    inclination: TensorLike = 0.0,
    phic: TensorLike = 0.0,
    tc: TensorLike = 0.0,
    f_ref: TensorLike = 20.0,
    *,
    _final_spin_override=None,
) -> tuple[jnp.ndarray, jnp.ndarray]:
    """Return h_plus, h_cross on times in seconds relative to the peak.

    Masses are solar masses, distance is Mpc, and f_ref is Hz. Parameter
    batches precede the shared time axis. This method applies no start mask.
    """
    t = jnp.asarray(times, dtype=jnp.float64)
    m1, m2, c1, c2, d, inc, pc, tcenter, fr = [
        jnp.asarray(v, dtype=t.dtype) for v in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref)
    ]
    m1, m2, c1, c2, d, inc, pc, tcenter, fr = jnp.broadcast_arrays(m1, m2, c1, c2, d, inc, pc, tcenter, fr)
    state = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)
    M = state[0]
    u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
    omega, phase, amp_state = self._evaluate(u, state)
    from ._phenomt_amplitude import amplitude

    amp = amplitude(u, amp_state, lambda tt: self._evaluate(tt, state)[0])
    # f_ref=0 follows the public convention and means the default 20 Hz
    # lower/reference frequency.  Subtract the analytic phase at that
    # frequency, as the author implementation does.
    fphase = jnp.where(fr > 0, fr, jnp.full_like(fr, 20.0))
    tref = self.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=_final_spin_override)
    uref = tref / (M * MSUN_SEC)
    phref = self._evaluate(uref[..., None], state)[1][..., 0]
    ph = phase - phref[..., None] + 2.0 * pc[..., None]
    scale = (M[..., None] * MSUN_SEC) / (d[..., None] * MPC_SI / C_SI)
    y22 = math.sqrt(5.0 / (4.0 * math.pi))
    ci = jnp.cos(inc[..., None])
    pref = -scale * y22 * amp
    return pref * (1 + ci**2) * 0.5 * jnp.cos(ph), pref * ci * jnp.sin(ph)

frequency(times, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None)

Instantaneous gravitational-wave frequency in Hz at peak-relative times.

Source code in src/jaxwave/approximants/imrphenomt.py

def frequency(self, times, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
    """Instantaneous gravitational-wave frequency in Hz at peak-relative times."""
    t = jnp.asarray(times, dtype=jnp.float64)
    m1 = jnp.asarray(mass1, dtype=t.dtype)
    m2 = jnp.asarray(mass2, dtype=t.dtype)
    c1 = jnp.asarray(chi1, dtype=t.dtype)
    c2 = jnp.asarray(chi2, dtype=t.dtype)
    m1, m2, c1, c2 = jnp.broadcast_arrays(m1, m2, c1, c2)
    M = m1 + m2
    u = t.reshape((1,) * M.ndim + (-1,)) / (M[..., None] * MSUN_SEC)
    return self._evaluate(u, self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override))[0] / (2 * PI * M[..., None] * MSUN_SEC)

time_at_frequency(frequency, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None)

Return peak-relative time at a requested instantaneous frequency.

The bracket is grown from the Newtonian inspiral estimate so that it remains useful for low-frequency BNS requests. Bisection itself is stopped from the differentiation graph; an implicit Newton correction then restores the derivative through the frequency equation while keeping JIT and batched execution fully functional.

Source code in src/jaxwave/approximants/imrphenomt.py

def time_at_frequency(self, frequency, mass1, mass2, chi1=0.0, chi2=0.0, *, _final_spin_override=None):
    """Return peak-relative time at a requested instantaneous frequency.

    The bracket is grown from the Newtonian inspiral estimate so that it
    remains useful for low-frequency BNS requests.  Bisection itself is
    stopped from the differentiation graph; an implicit Newton correction
    then restores the derivative through the frequency equation while
    keeping JIT and batched execution fully functional.
    """
    f = jnp.asarray(frequency, dtype=jnp.float64)
    m1 = jnp.asarray(mass1, dtype=f.dtype)
    m2 = jnp.asarray(mass2, dtype=f.dtype)
    c1 = jnp.asarray(chi1, dtype=f.dtype)
    c2 = jnp.asarray(chi2, dtype=f.dtype)
    m1, m2, c1, c2, f = jnp.broadcast_arrays(m1, m2, c1, c2, f)
    M = m1 + m2
    target = 2 * PI * f * M * MSUN_SEC
    st = self._coeff(m1, m2, c1, c2, final_spin=_final_spin_override)

    # Eager calls give a useful exception.  Under jit these conditions are
    # represented by ``valid`` below, yielding NaN rather than selecting a
    # clipped endpoint when the requested root does not exist.
    concrete = not any(isinstance(v, jax.core.Tracer) for v in (f, m1, m2, c1, c2, _final_spin_override))
    valid = jnp.isfinite(f) & (f > 0) & (target < st[30])
    if concrete:
        if not bool(jnp.all(jnp.isfinite(f) & (f > 0))):
            raise ValueError("frequency must be finite and positive")
        if bool(jnp.any(target >= st[30])):
            raise ValueError("frequency is above the supported IMRPhenomT range")

    # Newtonian leading-order time, buffered as in the author solver.
    lo0 = -5.0 / (256.0 * st[1] * jnp.power(target / 2.0, 8.0 / 3.0)) * 1.2
    hi0 = jnp.zeros_like(target)

    def grow_hi(_, hi):
        om_hi = self._evaluate(hi[..., None], st)[0][..., 0]
        need = om_hi < target
        return jnp.where(need, jnp.where(hi == 0, jnp.ones_like(hi), 2 * hi), hi)

    def grow_lo(_, lo):
        om_lo = self._evaluate(lo[..., None], st)[0][..., 0]
        return jnp.where(om_lo > target, lo * 2.0, lo)

    hi = jax.lax.fori_loop(0, 16, grow_hi, hi0)
    lo = jax.lax.fori_loop(0, 12, grow_lo, lo0)

    def bisect(_, bounds):
        lo_i, hi_i = bounds
        mid = (lo_i + hi_i) / 2
        om = self._evaluate(mid[..., None], st)[0][..., 0]
        below = om < target
        return jnp.where(below, mid, lo_i), jnp.where(below, hi_i, mid)

    lo, hi = jax.lax.fori_loop(0, 80, bisect, (lo, hi))
    root = jax.lax.stop_gradient((lo + hi) / 2)
    om = self._evaluate(root[..., None], st)[0][..., 0]
    dom = self._domega_full(root[..., None], st)[..., 0]
    root = root + (target - om) / jnp.maximum(dom, jnp.finfo(root.dtype).tiny)
    root = jnp.where(valid, root, jnp.nan)
    return root * M * MSUN_SEC

jaxwave.IMRPhenomTHM

Aligned-spin dominant and subdominant time-domain modes.

Source code in src/jaxwave/approximants/imrphenomthm.py

class IMRPhenomTHM:
    """Aligned-spin dominant and subdominant time-domain modes."""
    supported_modes = _POSITIVE

    def __init__(self, mode_array: Iterable[tuple[int, int]] | None = None):
        self.mode_array = _validate_modes(mode_array) if mode_array is not None else _POSITIVE + tuple((l, -m) for l, m in _POSITIVE)
        self._tmodel = IMRPhenomT()

    def _positive(self, u, state22, mode, eta, delta, S, dchi, chi1, chi2, final_spin=None):
        l, m = mode
        om22, ph22, a22 = self._tmodel._evaluate(u, state22)
        if mode == (2, 2):
            from ._phenomt_amplitude import amplitude
            return amplitude(u, a22, lambda tt: self._tmodel._evaluate(tt, state22)[0]), ph22, om22
        # The published HM inspiral phase is m/2 times the cached 22 phase.
        phase = m / 2 * ph22 - _OFF[10 * l + m]
        om = m / 2 * om22
        # THM's PN amplitude uses the orbital (22) frequency.
        x = jnp.clip(om22 / 2, min=1e-30) ** (2 / 3)
        coeff = _mode_pn(10 * l + m, eta, delta, chi1, chi2)
        cp = jnp.stack([_fit("_IMRPhenomT_Inspiral_Amp_CP", eta, S, dchi, delta, 10 * l + m, i) for i in (1, 2, 3)], -1)
        # Collocation times are fixed in total-mass units; evaluate from 22.
        uc = eta[..., None] * 0 + jnp.asarray((-2000., -250., -150.), dtype=eta.dtype)
        # The PN amplitude uses x=(omega_22/2)^(2/3), even for m != 2.
        oc = self._tmodel._evaluate(uc, state22)[0]
        xx = jnp.clip(oc / 2, min=1e-30) ** (2 / 3)
        b0 = _pn_amp(xx, coeff, jnp.zeros(eta.shape + (3,), dtype=eta.dtype)).real
        A = jnp.stack((xx**4, xx**4.5, xx**5), -1)
        den = coeff[2]
        den = jnp.where(jnp.abs(den) < 1e-30, jnp.where(den < 0, -1e-30, 1e-30), den)
        pseudo = jnp.linalg.solve(A, ((cp - b0) / (den[..., None] * xx))[..., None])[..., 0]
        # The augmented spin supplied by precession changes the QNM rates in
        # the ringdown only.  Keep the aligned remnant for the merger
        # collocation system so its inspiral/merger portion remains identical
        # to the aligned model; amplitude and phase are then joined at their
        # respective fitted amplitude and frequency cuts.
        af_merger = _fit("IMRPhenomX_FinalSpin2017", eta, chi1, chi2, delta)
        af_rd = af_merger if final_spin is None else final_spin
        mf = _fit("IMRPhenomX_FinalMass2017", eta, chi1, chi2, delta)
        fd = _fit("IMRPhenomT_fdampfit", af_merger, 10 * l + m) / mf
        alpha = 2 * PI * fd
        fd_rd = _fit("IMRPhenomT_fdampfit", af_rd, 10 * l + m) / mf
        alpha_rd = 2 * PI * fd_rd
        peak = _fit("_IMRPhenomT_PeakAmp", eta, S, dchi, delta, 10 * l + m)
        mid = _fit("_IMRPhenomT_Intermediate_Amp_CP1", eta, S, dchi, delta, 10 * l + m)
        c3 = _fit("_IMRPhenomT_Ringdown_Amp_C3", eta, S, dchi, 10 * l + m)
        tshift = _fit("_IMRPhenomT_tshift", eta, S, dchi, 10 * l + m)
        c2 = (2 * PI * _fit("IMRPhenomT_fdampn2fit", af_merger, 10 * l + m) / mf - alpha) / 2
        c2 = jnp.where(c2 > jnp.abs(-.5 * alpha / jnp.tanh(c3)), -.5 * alpha / jnp.tanh(c3), c2)
        c2_rd = (2 * PI * _fit("IMRPhenomT_fdampn2fit", af_rd, 10 * l + m) / mf - alpha_rd) / 2
        c2_rd = jnp.where(c2_rd > jnp.abs(-.5 * alpha_rd / jnp.tanh(c3)), -.5 * alpha_rd / jnp.tanh(c3), c2_rd)
        c1, c4 = peak * alpha_rd * jnp.cosh(c3)**2 / c2_rd, peak - peak * alpha_rd * jnp.cosh(c3)**2 / c2_rd * jnp.tanh(c3)
        mt = eta[..., None] * 0 + jnp.asarray((-150., -25., 0.), dtype=eta.dtype)
        amp_cut = _pn_amp(xx[..., -1:], coeff, pseudo)[..., 0]
        ic = jnp.sign(jnp.real(amp_cut)) * jnp.abs(amp_cut)
        A4 = jnp.concatenate((_basis(mt - tshift[..., None], alpha), _basis((mt - tshift[..., None])[..., :1], alpha, True)), -2)
        # The published matching system fixes the peak row's quadratic basis
        # entry to zero (the ansatz itself still evaluates (t-tshift)^2).
        A4 = A4.at[..., 2, :].set(jnp.concatenate((jnp.ones_like(A4[..., 2:3, :3]), jnp.zeros_like(A4[..., 2:3, 3:])), -1)[..., 0, :])
        omega_cut = self._tmodel._evaluate(jnp.asarray(-150., dtype=u.dtype), state22)[0][..., 0]
        domega_cut = self._tmodel._domega_full(jnp.asarray(-150., dtype=u.dtype), state22)[..., 0]
        x_cut = jnp.clip(omega_cut / 2, min=1e-30) ** (2 / 3)
        # The source ansatz is signed by Re(A_PN) before matching its slope.
        deriv = jnp.sign(jnp.real(amp_cut)) * _pn_amp_abs_derivative(x_cut, coeff, pseudo, omega_cut, domega_cut)
        rhs = jnp.stack((ic, mid, peak, deriv), -1)
        mc = jnp.linalg.solve(A4, rhs[..., None])[..., 0]
        # The amplitude ansatz switches to ringdown at tshift (the fitted
        # ringdown cut), while the phase/frequency switch is at zero.
        tshift_n = tshift[..., None]
        tm = jnp.minimum(jnp.maximum(u, -150.), tshift_n)
        tr = jnp.maximum(u, tshift_n)
        am = jnp.sum(_basis(tm - tshift[..., None], alpha) * mc[..., None, :], -1)
        tr = tr - tshift[..., None]
        ar = jnp.exp(-alpha_rd[..., None] * tr) * (c1[..., None] * jnp.tanh(c2_rd[..., None] * tr + c3[..., None]) + c4[..., None])
        amp_complex = _pn_amp(jnp.clip(x, min=1e-30), coeff, pseudo)
        amp = jnp.where(u < -150, amp_complex, jnp.where(u < tshift[..., None], am, ar))

        # Independent HM merger/ringdown frequency ansatz and its analytic
        # primitive (THM Eqs. 15, 25, 28-31), joined to m/2 times the 22 phase.
        op = _fit("_IMRPhenomT_PeakFrequency", eta, S, dchi, delta, 10 * l + m)
        c2f = _fit("_IMRPhenomT_RD_Freq_C2", eta, S, dchi, delta, 10 * l + m)
        c3f = _fit("_IMRPhenomT_RD_Freq_C3", eta, S, dchi, delta, 10 * l + m)
        wr = 2 * PI * _fit("IMRPhenomT_fringfit", af_merger, 10 * l + m) / mf
        aa = 2 * PI * _fit("IMRPhenomT_fdampfit", af_merger, 10 * l + m) / mf
        c1f = (1 + c3f) * (wr - op) / (c2f * c3f)
        wr_rd = 2 * PI * _fit("IMRPhenomT_fringfit", af_rd, 10 * l + m) / mf
        c1f_rd = (1 + c3f) * (wr_rd - op) / (c2f * c3f)
        cut = jnp.asarray(-150., dtype=u.dtype)
        tcp = jnp.asarray(-25., dtype=u.dtype)
        w22cut = self._tmodel._evaluate(cut, state22)[0][..., 0]
        dw22cut = self._tmodel._domega_full(cut, state22)[..., 0]
        x_cut = jnp.clip(w22cut / 2, min=1e-30) ** (2 / 3)
        omega_cut_amp = -_pn_amp_phase_derivative(x_cut, coeff, pseudo, w22cut, dw22cut)
        wcut = m / 2 * w22cut + omega_cut_amp
        dwcut = m / 2 * dw22cut
        dwp = -c1f * c2f**2 * c3f / (1 + c3f)**2 / wr
        ac, am = jnp.asinh(aa * cut), jnp.asinh(aa * tcp)
        wtcp = _fit("_IMRPhenomT_Intermediate_Freq_CP1", eta, S, dchi, delta, 10 * l + m)
        b = jnp.stack((1 - wcut / wr - (1 - op / wr) - dwp / aa * ac,
                         1 - wtcp / wr - (1 - op / wr) - dwp / aa * am,
                         -dwcut / wr - dwp / jnp.sqrt(1 + (aa * cut)**2)), -1)
        matf = jnp.stack((jnp.stack((ac**2, ac**3, ac**4), -1), jnp.stack((am**2, am**3, am**4), -1), jnp.stack((2*aa*ac/jnp.sqrt(1+(aa*cut)**2), 3*aa*ac**2/jnp.sqrt(1+(aa*cut)**2), 4*aa*ac**3/jnp.sqrt(1+(aa*cut)**2)), -1)), -2)
        cf = jnp.linalg.solve(matf, b[..., None])[..., 0]
        tm = jnp.clip(u, -150., 0.)
        tr = jnp.clip(u, min=0.)
        xxm = jnp.asinh(aa[..., None] * tm)
        wm = wr[..., None] * (1 - (1 - op[..., None]/wr[..., None] + xxm * (dwp[..., None]/aa[..., None] + xxm * (cf[...,0,None] + xxm*(cf[...,1,None] + xxm*cf[...,2,None])))))
        wrg = wr_rd[..., None] - c1f_rd[..., None]*c2f[..., None]*c3f[..., None]*jnp.exp(-c2f[..., None]*tr)/(1+c3f[..., None]*jnp.exp(-c2f[..., None]*tr))
        om = jnp.where(u < -150, m/2*om22, jnp.where(u < 0, wm, wrg))
        # The join phase is evaluated at the scalar cut, rather than borrowed
        # from whichever time samples the caller happened to request.
        phcut22 = self._tmodel._evaluate(cut, state22)[1][..., 0]
        # atan2 has an undefined derivative at (0, 0).  Odd modes vanish at
        # equal mass/equal spin, so evaluate the inactive branch at a safe
        # unit real amplitude before masking its phase.
        amp_nonzero = jnp.abs(amp_cut) > 1e-30
        amp_re = jnp.where(amp_nonzero, jnp.real(amp_cut), jnp.ones_like(jnp.real(amp_cut)))
        amp_im = jnp.where(amp_nonzero, jnp.imag(amp_cut), jnp.zeros_like(jnp.imag(amp_cut)))
        phicut = jnp.arctan2(amp_im, amp_re) + jnp.where(
            amp_re < 0, jnp.asarray(PI, dtype=u.dtype), jnp.zeros_like(amp_re))
        phicut = jnp.where(amp_nonzero, phicut, jnp.zeros_like(phicut))
        rawcut = m / 2 * phcut22 - phicut - _OFF[10 * l + m]
        from .imrphenomt import IMRPhenomT as _T
        inter = _T._phase_inter_raw(tm, aa[..., None], wr[..., None], op[..., None], dwp[..., None], cf[..., None, :])
        inter0 = _T._phase_inter_raw(jnp.zeros_like(eta), aa, wr, op, dwp, cf)
        offm = rawcut - _T._phase_inter_raw(cut, aa, wr, op, dwp, cf)
        rphase = c1f_rd[..., None] * jnp.log((1 + c3f[..., None]*jnp.exp(-c2f[..., None]*tr))/(1+c3f[..., None])) + wr_rd[..., None]*tr + inter0[..., None] + offm[..., None]
        phase = jnp.where(u < -150, phase, jnp.where(u < 0, inter + offm[..., None], rphase))
        odd_zero = (jnp.abs(delta) < 1e-14) & (jnp.abs(chi1 - chi2) < 1e-14) & (m % 2 == 1)
        amp = jnp.where(odd_zero[..., None], jnp.zeros_like(amp), amp)
        phase = jnp.where(odd_zero[..., None], jnp.zeros_like(phase), phase)
        om = jnp.where(odd_zero[..., None], jnp.zeros_like(om), om)
        return amp, phase, om

    def modes(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
        t = jnp.asarray(times, dtype=getattr(times, "dtype", jnp.float64) if jnp.issubdtype(jnp.asarray(times).dtype, jnp.inexact) else jnp.float64)
        vals = [jnp.asarray(v, dtype=t.dtype) for v in (mass1, mass2, chi1, chi2, distance, phic, tc)]
        fref = jnp.asarray(f_ref, dtype=t.dtype)
        m1, m2, c1, c2, dist, pc, tcenter, fref = jnp.broadcast_arrays(*vals, fref)
        if final_spin is not None:
            final_spin = jnp.broadcast_arrays(jnp.asarray(final_spin, dtype=t.dtype), m1)[0]
        if modes is not None:
            if mode_array is not None:
                raise ValueError("pass only one of modes and mode_array")
            mode_array = modes
        swap = m1 < m2
        m1, m2 = jnp.where(swap, m2, m1), jnp.where(swap, m1, m2)
        c1, c2 = jnp.where(swap, c2, c1), jnp.where(swap, c1, c2)
        state = self._tmodel._coeff(m1, m2, c1, c2, final_spin)
        M, eta, delta, S, dchi = state[0], state[1], state[2], state[3], state[4]
        u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
        sel = _validate_modes(mode_array) if mode_array is not None else self.mode_array
        scale = (M[..., None] * MSUN_SEC) / (dist[..., None] * MPC_SI / C_SI)
        # Match the 22 reference phase convention used by IMRPhenomT.  A
        # zero reference frequency follows the public 20-Hz convention.
        fphase = jnp.where(fref > 0, fref, jnp.full_like(fref, 20.0))
        tref = self._tmodel.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=final_spin)
        uref = tref / (M * MSUN_SEC)
        phref = self._tmodel._evaluate(uref[..., None], state)[1][..., 0]
        out = {}
        positive = [(l, m) for l, m in sel if m > 0]
        for l, m in positive:
            if m < 0:
                continue
            amp, phase, _ = self._positive(u, state, (l, m), eta, delta, S, dchi, c1, c2, final_spin)
            out[(l, m)] = scale * amp * jnp.exp(-1j * (phase - (m / 2) * phref[..., None] + m * pc[..., None]))
        for l, m in sel:
            if m < 0:
                if (l, -m) not in out:
                    amp, phase, _ = self._positive(u, state, (l, -m), eta, delta, S, dchi, c1, c2, final_spin)
                    out[(l, -m)] = scale * amp * jnp.exp(-1j * (phase - ((-m) / 2) * phref[..., None] + (-m) * pc[..., None]))
                out[(l, m)] = (-1) ** l * out[(l, -m)].conj()
        return out

    def time_domain(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., inclination=0., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
        # Polarizations require the equatorial-symmetry partners.  A public
        # positive-mode selection therefore denotes the corresponding pair.
        if modes is not None:
            if mode_array is not None:
                raise ValueError("pass only one of modes and mode_array")
            mode_array = modes
        selected = None if mode_array is None else tuple(mode_array)
        if selected is not None:
            selected = tuple(dict.fromkeys(selected + tuple((l, -m) for l, m in selected if m > 0)))
        modes_out = self.modes(times, mass1, mass2, chi1, chi2, distance, phic, tc, f_ref=f_ref, mode_array=selected, final_spin=final_spin)
        inc = jnp.asarray(inclination, dtype=jnp.float64)
        # s=-2 harmonics for the five supported l,m; closed forms avoid a scipy dependency.
        c = jnp.cos(inc[..., None])
        out = jnp.zeros_like(next(iter(modes_out.values())).real), jnp.zeros_like(next(iter(modes_out.values())).real)
        # Factorial Wigner-d form of the s=-2 spherical harmonics, evaluated
        # at the author's azimuth pi/2.  Summing signed modes preserves the
        # exact 22 projection while retaining every selected harmonic.
        for (l, m), h in modes_out.items():
            d = jnp.zeros_like(c)
            for k in range(0, 20):
                a, b, q, r = l + m - k, k, 2 - m + k, l - 2 - k
                if min(a, b, q, r) < 0:
                    continue
                den = math.factorial(a)*math.factorial(b)*math.factorial(q)*math.factorial(r)
                num = math.factorial(l + 2) * math.factorial(l - 2) * math.factorial(l + m) * math.factorial(l - m)
                coef = math.sqrt(num) / den
                d = d + (-1.)**(k + m - 2) * coef * jnp.cos(inc[...,None]/2)**(2*l + m - 2*k - 2) * jnp.sin(inc[...,None]/2)**(2 - m + 2*k)
            y = math.sqrt((2*l+1)/(4*math.pi)) * d * ((-1) ** m) * ((1j) ** m)
            out = (out[0] + (h*y).real, out[1] - (h*y).imag)
        return out

    forward = time_domain
    __call__ = time_domain

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, mode_array=None, final_spin=None, modes=None)

Source code in src/jaxwave/approximants/imrphenomthm.py

def time_domain(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., inclination=0., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
    # Polarizations require the equatorial-symmetry partners.  A public
    # positive-mode selection therefore denotes the corresponding pair.
    if modes is not None:
        if mode_array is not None:
            raise ValueError("pass only one of modes and mode_array")
        mode_array = modes
    selected = None if mode_array is None else tuple(mode_array)
    if selected is not None:
        selected = tuple(dict.fromkeys(selected + tuple((l, -m) for l, m in selected if m > 0)))
    modes_out = self.modes(times, mass1, mass2, chi1, chi2, distance, phic, tc, f_ref=f_ref, mode_array=selected, final_spin=final_spin)
    inc = jnp.asarray(inclination, dtype=jnp.float64)
    # s=-2 harmonics for the five supported l,m; closed forms avoid a scipy dependency.
    c = jnp.cos(inc[..., None])
    out = jnp.zeros_like(next(iter(modes_out.values())).real), jnp.zeros_like(next(iter(modes_out.values())).real)
    # Factorial Wigner-d form of the s=-2 spherical harmonics, evaluated
    # at the author's azimuth pi/2.  Summing signed modes preserves the
    # exact 22 projection while retaining every selected harmonic.
    for (l, m), h in modes_out.items():
        d = jnp.zeros_like(c)
        for k in range(0, 20):
            a, b, q, r = l + m - k, k, 2 - m + k, l - 2 - k
            if min(a, b, q, r) < 0:
                continue
            den = math.factorial(a)*math.factorial(b)*math.factorial(q)*math.factorial(r)
            num = math.factorial(l + 2) * math.factorial(l - 2) * math.factorial(l + m) * math.factorial(l - m)
            coef = math.sqrt(num) / den
            d = d + (-1.)**(k + m - 2) * coef * jnp.cos(inc[...,None]/2)**(2*l + m - 2*k - 2) * jnp.sin(inc[...,None]/2)**(2 - m + 2*k)
        y = math.sqrt((2*l+1)/(4*math.pi)) * d * ((-1) ** m) * ((1j) ** m)
        out = (out[0] + (h*y).real, out[1] - (h*y).imag)
    return out

modes(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, phic=0.0, tc=0.0, f_ref=20.0, mode_array=None, final_spin=None, modes=None)

Source code in src/jaxwave/approximants/imrphenomthm.py

def modes(self, times, mass1, mass2, chi1=0., chi2=0., distance=100., phic=0., tc=0., f_ref=20., mode_array=None, final_spin=None, modes=None):
    t = jnp.asarray(times, dtype=getattr(times, "dtype", jnp.float64) if jnp.issubdtype(jnp.asarray(times).dtype, jnp.inexact) else jnp.float64)
    vals = [jnp.asarray(v, dtype=t.dtype) for v in (mass1, mass2, chi1, chi2, distance, phic, tc)]
    fref = jnp.asarray(f_ref, dtype=t.dtype)
    m1, m2, c1, c2, dist, pc, tcenter, fref = jnp.broadcast_arrays(*vals, fref)
    if final_spin is not None:
        final_spin = jnp.broadcast_arrays(jnp.asarray(final_spin, dtype=t.dtype), m1)[0]
    if modes is not None:
        if mode_array is not None:
            raise ValueError("pass only one of modes and mode_array")
        mode_array = modes
    swap = m1 < m2
    m1, m2 = jnp.where(swap, m2, m1), jnp.where(swap, m1, m2)
    c1, c2 = jnp.where(swap, c2, c1), jnp.where(swap, c1, c2)
    state = self._tmodel._coeff(m1, m2, c1, c2, final_spin)
    M, eta, delta, S, dchi = state[0], state[1], state[2], state[3], state[4]
    u = (t.reshape((1,) * M.ndim + (-1,)) - tcenter[..., None]) / (M[..., None] * MSUN_SEC)
    sel = _validate_modes(mode_array) if mode_array is not None else self.mode_array
    scale = (M[..., None] * MSUN_SEC) / (dist[..., None] * MPC_SI / C_SI)
    # Match the 22 reference phase convention used by IMRPhenomT.  A
    # zero reference frequency follows the public 20-Hz convention.
    fphase = jnp.where(fref > 0, fref, jnp.full_like(fref, 20.0))
    tref = self._tmodel.time_at_frequency(fphase, m1, m2, c1, c2, _final_spin_override=final_spin)
    uref = tref / (M * MSUN_SEC)
    phref = self._tmodel._evaluate(uref[..., None], state)[1][..., 0]
    out = {}
    positive = [(l, m) for l, m in sel if m > 0]
    for l, m in positive:
        if m < 0:
            continue
        amp, phase, _ = self._positive(u, state, (l, m), eta, delta, S, dchi, c1, c2, final_spin)
        out[(l, m)] = scale * amp * jnp.exp(-1j * (phase - (m / 2) * phref[..., None] + m * pc[..., None]))
    for l, m in sel:
        if m < 0:
            if (l, -m) not in out:
                amp, phase, _ = self._positive(u, state, (l, -m), eta, delta, S, dchi, c1, c2, final_spin)
                out[(l, -m)] = scale * amp * jnp.exp(-1j * (phase - ((-m) / 2) * phref[..., None] + (-m) * pc[..., None]))
            out[(l, m)] = (-1) ** l * out[(l, -m)].conj()
    return out

jaxwave.IMRPhenomTP

Bases: IMRPhenomT

Precessing dominant-harmonic PhenomT with selectable Euler-angle prescriptions.

Source code in src/jaxwave/approximants/imrphenomtp.py

class IMRPhenomTP(IMRPhenomT):
    """Precessing dominant-harmonic PhenomT with selectable Euler-angle prescriptions."""

    def time_domain(
        self,
        times,
        mass1,
        mass2,
        chi1=0.0,
        chi2=0.0,
        distance=100.0,
        inclination=0.0,
        phic=0.0,
        tc=0.0,
        f_ref=20.0,
        spin1x=0.0,
        spin1y=0.0,
        spin2x=0.0,
        spin2y=0.0,
        precession_method="nnlo",
        modes=None,
        final_spin=None,
        precession_version=None,
        _precession_state=None,
    ):
        method, version = resolve_precession(precession_method, precession_version)
        if modes is not None:
            requested = tuple((int(mode[0]), int(mode[1])) for mode in modes)
            if not requested or any((l, abs(m)) != (2, 2) for l, m in requested):
                raise ValueError("IMRPhenomTP supports only mode (2, 2)")
        t = jnp.asarray(times, dtype=jnp.float64)
        vals = [
            jnp.asarray(x, dtype=t.dtype)
            for x in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref, spin1x, spin1y, spin2x, spin2y)
        ]
        if final_spin is None:
            m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2 = jnp.broadcast_arrays(*vals)
            fs = None
        else:
            vals.append(jnp.asarray(final_spin, dtype=t.dtype))
            m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2, fs = jnp.broadcast_arrays(*vals)
        m1, m2, x1, y1, z1, x2, y2, z2 = _canonicalize_components(m1, m2, x1, y1, z1, x2, y2, z2)
        fs, prepared = prepare_precession(
            method, m1, m2, x1, y1, z1, x2, y2, z2, fr,
            final_spin=fs, prepared=_precession_state,
        )
        hp0, hc0 = IMRPhenomT.time_domain(
            self,
            t,
            m1,
            m2,
            z1,
            z2,
            dist,
            0.0,
            0.0,
            tcenter,
            fr,
            _final_spin_override=fs,
        )
        h22 = -(hp0 - 1j * hc0) / math.sqrt(5.0 / (4.0 * math.pi))
        shifted = t.reshape((1,) * m1.ndim + (-1,)) - tcenter[..., None]
        # Eq. 14 uses the aligned carrier for the angle series; the augmented
        # remnant enters only the post-peak QNM continuation.
        state_angle = self._coeff(m1, m2, z1, z2)
        freq = self._evaluate(shifted / ((m1 + m2)[..., None] * MSUN_SEC), state_angle)[0]
        freq = freq / (2.0 * math.pi * (m1 + m2)[..., None] * MSUN_SEC)
        fpeak = state_angle[13] / (2.0 * math.pi * (m1 + m2) * MSUN_SEC)
        a, b, g, thjn, zeta = precession_geometry(
            shifted,
            freq,
            m1,
            m2,
            x1,
            y1,
            z1,
            x2,
            y2,
            z2,
            inc,
            pc,
            fr,
            frequency_peak=fpeak,
            final_spin=fs,
            precession_method=method,
            precession_version=version,
            prepared=prepared,
            frequency_at_time=lambda query: carrier_frequency(self, state_angle, m1, m2, query),
        )
        hp_rot, hc_rot = project_modes({(2, 2): h22}, a, b, g, thjn[..., None], zeta[..., None])
        hp_al, hc_al = IMRPhenomT.time_domain(
            self,
            t,
            m1,
            m2,
            z1,
            z2,
            dist,
            inc,
            pc,
            tcenter,
            fr,
            _final_spin_override=fs,
        )
        aligned = (x1**2 + y1**2 + x2**2 + y2**2) == 0
        return jnp.where(aligned[..., None], hp_al, hp_rot), jnp.where(aligned[..., None], hc_al, hc_rot)

    __call__ = time_domain
    forward = time_domain

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, spin1x=0.0, spin1y=0.0, spin2x=0.0, spin2y=0.0, precession_method='nnlo', modes=None, final_spin=None, precession_version=None, _precession_state=None)

Source code in src/jaxwave/approximants/imrphenomtp.py

def time_domain(
    self,
    times,
    mass1,
    mass2,
    chi1=0.0,
    chi2=0.0,
    distance=100.0,
    inclination=0.0,
    phic=0.0,
    tc=0.0,
    f_ref=20.0,
    spin1x=0.0,
    spin1y=0.0,
    spin2x=0.0,
    spin2y=0.0,
    precession_method="nnlo",
    modes=None,
    final_spin=None,
    precession_version=None,
    _precession_state=None,
):
    method, version = resolve_precession(precession_method, precession_version)
    if modes is not None:
        requested = tuple((int(mode[0]), int(mode[1])) for mode in modes)
        if not requested or any((l, abs(m)) != (2, 2) for l, m in requested):
            raise ValueError("IMRPhenomTP supports only mode (2, 2)")
    t = jnp.asarray(times, dtype=jnp.float64)
    vals = [
        jnp.asarray(x, dtype=t.dtype)
        for x in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref, spin1x, spin1y, spin2x, spin2y)
    ]
    if final_spin is None:
        m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2 = jnp.broadcast_arrays(*vals)
        fs = None
    else:
        vals.append(jnp.asarray(final_spin, dtype=t.dtype))
        m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2, fs = jnp.broadcast_arrays(*vals)
    m1, m2, x1, y1, z1, x2, y2, z2 = _canonicalize_components(m1, m2, x1, y1, z1, x2, y2, z2)
    fs, prepared = prepare_precession(
        method, m1, m2, x1, y1, z1, x2, y2, z2, fr,
        final_spin=fs, prepared=_precession_state,
    )
    hp0, hc0 = IMRPhenomT.time_domain(
        self,
        t,
        m1,
        m2,
        z1,
        z2,
        dist,
        0.0,
        0.0,
        tcenter,
        fr,
        _final_spin_override=fs,
    )
    h22 = -(hp0 - 1j * hc0) / math.sqrt(5.0 / (4.0 * math.pi))
    shifted = t.reshape((1,) * m1.ndim + (-1,)) - tcenter[..., None]
    # Eq. 14 uses the aligned carrier for the angle series; the augmented
    # remnant enters only the post-peak QNM continuation.
    state_angle = self._coeff(m1, m2, z1, z2)
    freq = self._evaluate(shifted / ((m1 + m2)[..., None] * MSUN_SEC), state_angle)[0]
    freq = freq / (2.0 * math.pi * (m1 + m2)[..., None] * MSUN_SEC)
    fpeak = state_angle[13] / (2.0 * math.pi * (m1 + m2) * MSUN_SEC)
    a, b, g, thjn, zeta = precession_geometry(
        shifted,
        freq,
        m1,
        m2,
        x1,
        y1,
        z1,
        x2,
        y2,
        z2,
        inc,
        pc,
        fr,
        frequency_peak=fpeak,
        final_spin=fs,
        precession_method=method,
        precession_version=version,
        prepared=prepared,
        frequency_at_time=lambda query: carrier_frequency(self, state_angle, m1, m2, query),
    )
    hp_rot, hc_rot = project_modes({(2, 2): h22}, a, b, g, thjn[..., None], zeta[..., None])
    hp_al, hc_al = IMRPhenomT.time_domain(
        self,
        t,
        m1,
        m2,
        z1,
        z2,
        dist,
        inc,
        pc,
        tcenter,
        fr,
        _final_spin_override=fs,
    )
    aligned = (x1**2 + y1**2 + x2**2 + y2**2) == 0
    return jnp.where(aligned[..., None], hp_al, hp_rot), jnp.where(aligned[..., None], hc_al, hc_rot)

jaxwave.IMRPhenomTPHM

Bases: IMRPhenomT

Precessing multipolar PhenomT with selectable Euler-angle prescriptions.

Source code in src/jaxwave/approximants/imrphenomtphm.py

class IMRPhenomTPHM(IMRPhenomT):
    """Precessing multipolar PhenomT with selectable Euler-angle prescriptions."""

    def time_domain(
        self,
        times,
        mass1,
        mass2,
        chi1=0.0,
        chi2=0.0,
        distance=100.0,
        inclination=0.0,
        phic=0.0,
        tc=0.0,
        f_ref=20.0,
        spin1x=0.0,
        spin1y=0.0,
        spin2x=0.0,
        spin2y=0.0,
        precession_method="nnlo",
        modes=None,
        final_spin=None,
        precession_version=None,
        _precession_state=None,
    ):
        method, version = resolve_precession(precession_method, precession_version)
        requested = tuple(_mode_key(mode) for mode in modes) if modes is not None else DEFAULT_MODES
        positive = tuple(dict.fromkeys((l, abs(m)) for l, m in requested))
        if not positive:
            raise ValueError("modes must contain at least one mode")
        if any(key not in _SUPPORTED_MODES for key in positive):
            raise ValueError(f"unsupported THM mode; supported positive modes are {sorted(_SUPPORTED_MODES)}")
        t = jnp.asarray(times, dtype=jnp.float64)
        vals = [
            jnp.asarray(x, dtype=t.dtype)
            for x in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref, spin1x, spin1y, spin2x, spin2y)
        ]
        if final_spin is None:
            m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2 = jnp.broadcast_arrays(*vals)
            fs = None
        else:
            vals.append(jnp.asarray(final_spin, dtype=t.dtype))
            m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2, fs = jnp.broadcast_arrays(*vals)
        m1, m2, x1, y1, z1, x2, y2, z2 = _canonicalize_components(m1, m2, x1, y1, z1, x2, y2, z2)
        fs, prepared = prepare_precession(
            method, m1, m2, x1, y1, z1, x2, y2, z2, fr,
            final_spin=fs, prepared=_precession_state,
        )
        from .imrphenomthm import IMRPhenomTHM

        thm = IMRPhenomTHM()
        selected = tuple(positive) + tuple((l, -m) for l, m in positive)
        raw = thm.modes(
            t,
            m1,
            m2,
            z1,
            z2,
            dist,
            phic=0.0,
            tc=tcenter,
            mode_array=selected,
            final_spin=fs,
            f_ref=fr,
        )
        cp = {}
        for key, value in raw.items():
            lm = _mode_key(key)
            if lm in selected:
                cp[lm] = _as_complex(value, t.dtype)
        if not cp:
            raise ValueError("THM mode provider returned no requested positive modes")
        shifted = t.reshape((1,) * m1.ndim + (-1,)) - tcenter[..., None]
        # Eq. 14 uses the aligned carrier for the angle series; the augmented
        # remnant enters only the post-peak QNM continuation.
        state_angle = self._coeff(m1, m2, z1, z2)
        freq = self._evaluate(shifted / ((m1 + m2)[..., None] * MSUN_SEC), state_angle)[0]
        freq = freq / (2.0 * math.pi * (m1 + m2)[..., None] * MSUN_SEC)
        fpeak = state_angle[13] / (2.0 * math.pi * (m1 + m2) * MSUN_SEC)
        a, b, g, thjn, zeta = precession_geometry(
            shifted,
            freq,
            m1,
            m2,
            x1,
            y1,
            z1,
            x2,
            y2,
            z2,
            inc,
            pc,
            fr,
            frequency_peak=fpeak,
            final_spin=fs,
            precession_method=method,
            precession_version=version,
            prepared=prepared,
            frequency_at_time=lambda query: carrier_frequency(self, state_angle, m1, m2, query),
        )
        hp_rot, hc_rot = project_modes(cp, a, b, g, thjn[..., None], zeta[..., None])
        hp_al, hc_al = thm.time_domain(
            t,
            m1,
            m2,
            z1,
            z2,
            dist,
            inc,
            phic=pc,
            tc=tcenter,
            f_ref=fr,
            mode_array=selected,
            final_spin=fs,
        )
        aligned = (x1**2 + y1**2 + x2**2 + y2**2) == 0
        return jnp.where(aligned[..., None], hp_al, hp_rot), jnp.where(aligned[..., None], hc_al, hc_rot)

    __call__ = time_domain
    forward = time_domain

time_domain(times, mass1, mass2, chi1=0.0, chi2=0.0, distance=100.0, inclination=0.0, phic=0.0, tc=0.0, f_ref=20.0, spin1x=0.0, spin1y=0.0, spin2x=0.0, spin2y=0.0, precession_method='nnlo', modes=None, final_spin=None, precession_version=None, _precession_state=None)

Source code in src/jaxwave/approximants/imrphenomtphm.py

def time_domain(
    self,
    times,
    mass1,
    mass2,
    chi1=0.0,
    chi2=0.0,
    distance=100.0,
    inclination=0.0,
    phic=0.0,
    tc=0.0,
    f_ref=20.0,
    spin1x=0.0,
    spin1y=0.0,
    spin2x=0.0,
    spin2y=0.0,
    precession_method="nnlo",
    modes=None,
    final_spin=None,
    precession_version=None,
    _precession_state=None,
):
    method, version = resolve_precession(precession_method, precession_version)
    requested = tuple(_mode_key(mode) for mode in modes) if modes is not None else DEFAULT_MODES
    positive = tuple(dict.fromkeys((l, abs(m)) for l, m in requested))
    if not positive:
        raise ValueError("modes must contain at least one mode")
    if any(key not in _SUPPORTED_MODES for key in positive):
        raise ValueError(f"unsupported THM mode; supported positive modes are {sorted(_SUPPORTED_MODES)}")
    t = jnp.asarray(times, dtype=jnp.float64)
    vals = [
        jnp.asarray(x, dtype=t.dtype)
        for x in (mass1, mass2, chi1, chi2, distance, inclination, phic, tc, f_ref, spin1x, spin1y, spin2x, spin2y)
    ]
    if final_spin is None:
        m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2 = jnp.broadcast_arrays(*vals)
        fs = None
    else:
        vals.append(jnp.asarray(final_spin, dtype=t.dtype))
        m1, m2, z1, z2, dist, inc, pc, tcenter, fr, x1, y1, x2, y2, fs = jnp.broadcast_arrays(*vals)
    m1, m2, x1, y1, z1, x2, y2, z2 = _canonicalize_components(m1, m2, x1, y1, z1, x2, y2, z2)
    fs, prepared = prepare_precession(
        method, m1, m2, x1, y1, z1, x2, y2, z2, fr,
        final_spin=fs, prepared=_precession_state,
    )
    from .imrphenomthm import IMRPhenomTHM

    thm = IMRPhenomTHM()
    selected = tuple(positive) + tuple((l, -m) for l, m in positive)
    raw = thm.modes(
        t,
        m1,
        m2,
        z1,
        z2,
        dist,
        phic=0.0,
        tc=tcenter,
        mode_array=selected,
        final_spin=fs,
        f_ref=fr,
    )
    cp = {}
    for key, value in raw.items():
        lm = _mode_key(key)
        if lm in selected:
            cp[lm] = _as_complex(value, t.dtype)
    if not cp:
        raise ValueError("THM mode provider returned no requested positive modes")
    shifted = t.reshape((1,) * m1.ndim + (-1,)) - tcenter[..., None]
    # Eq. 14 uses the aligned carrier for the angle series; the augmented
    # remnant enters only the post-peak QNM continuation.
    state_angle = self._coeff(m1, m2, z1, z2)
    freq = self._evaluate(shifted / ((m1 + m2)[..., None] * MSUN_SEC), state_angle)[0]
    freq = freq / (2.0 * math.pi * (m1 + m2)[..., None] * MSUN_SEC)
    fpeak = state_angle[13] / (2.0 * math.pi * (m1 + m2) * MSUN_SEC)
    a, b, g, thjn, zeta = precession_geometry(
        shifted,
        freq,
        m1,
        m2,
        x1,
        y1,
        z1,
        x2,
        y2,
        z2,
        inc,
        pc,
        fr,
        frequency_peak=fpeak,
        final_spin=fs,
        precession_method=method,
        precession_version=version,
        prepared=prepared,
        frequency_at_time=lambda query: carrier_frequency(self, state_angle, m1, m2, query),
    )
    hp_rot, hc_rot = project_modes(cp, a, b, g, thjn[..., None], zeta[..., None])
    hp_al, hc_al = thm.time_domain(
        t,
        m1,
        m2,
        z1,
        z2,
        dist,
        inc,
        phic=pc,
        tc=tcenter,
        f_ref=fr,
        mode_array=selected,
        final_spin=fs,
    )
    aligned = (x1**2 + y1**2 + x2**2 + y2**2) == 0
    return jnp.where(aligned[..., None], hp_al, hp_rot), jnp.where(aligned[..., None], hc_al, hc_rot)

---

Shape Bucketing & MCMC Evaluators

jaxwave.inference.bucketing.BucketedWaveformEvaluator

Manages JIT-compiled waveform evaluators across fixed shape buckets.

Routes evaluations to the nearest bucket size and caches each evaluator. Each new bucket compiles on first use. Changed dtypes or other traced shapes can trigger additional compilations.

Source code in src/jaxwave/inference/bucketing.py

class BucketedWaveformEvaluator:
    """Manages JIT-compiled waveform evaluators across fixed shape buckets.

    Routes evaluations to the nearest bucket size and caches each evaluator.
    Each new bucket compiles on first use. Changed dtypes or other traced shapes
    can trigger additional compilations.
    """

    def __init__(
        self,
        approximant: str = "IMRPhenomD",
        f_lower: Optional[float] = None,
        f_final: Optional[float] = None,
        delta_f: Optional[float] = None,
        sample_frequencies: Optional[jnp.ndarray] = None,
        buckets: Sequence[int] = DEFAULT_BATCH_BUCKETS,
        parameter_order: Tuple[str, ...] = ("mass1", "mass2", "spin1z", "spin2z", "distance", "inclination", "phic", "tc"),
        **default_kwargs,
    ):
        self.approximant = approximant
        self.f_lower = f_lower
        self.f_final = f_final
        self.delta_f = delta_f
        self.sample_frequencies = sample_frequencies
        self.buckets = sorted(buckets)
        self.parameter_order = parameter_order
        self.default_kwargs = default_kwargs
        self._compiled_evaluators: Dict[int, Callable] = {}

    def _get_or_compile_evaluator(self, bucket_size: int) -> Callable:
        """Compile or retrieve the JIT evaluator for a specific bucket size."""
        if bucket_size in self._compiled_evaluators:
            return self._compiled_evaluators[bucket_size]

        approx = self.approximant
        f_low = self.f_lower
        f_fin = self.f_final
        df = self.delta_f
        s_freqs = self.sample_frequencies
        param_keys = self.parameter_order
        d_kwargs = dict(self.default_kwargs)

        def _raw_eval(*param_arrays):
            kw = dict(d_kwargs)
            for k, arr in zip(param_keys, param_arrays):
                kw[k] = arr
            return get_fd_waveform(
                approximant=approx,
                f_lower=f_low,
                f_final=f_fin,
                delta_f=df,
                sample_frequencies=s_freqs,
                **kw,
            )

        jitted_eval = jax.jit(_raw_eval)
        self._compiled_evaluators[bucket_size] = jitted_eval
        return jitted_eval

    def __call__(
        self,
        parameters: Dict[str, Any],
        **kwargs,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """Evaluate waveforms for an arbitrary batch size with automatic bucketing."""
        norm_params = {}
        name_map = {
            "mass_1": "mass1",
            "mass_2": "mass2",
            "spin_1z": "spin1z",
            "spin_2z": "spin2z",
            "luminosity_distance": "distance",
            "theta_jn": "inclination",
            "phase": "phic",
            "geocent_time": "tc",
        }
        for k, v in parameters.items():
            norm_params[name_map.get(k, k)] = v
        norm_params.update(kwargs)

        # Determine original batch size
        first_val = next(iter(norm_params.values()))
        first_arr = jnp.asarray(first_val)
        orig_batch = first_arr.shape[0] if first_arr.ndim > 0 else 1

        bucket_size = get_bucketed_batch_size(orig_batch, self.buckets)
        evaluator = self._get_or_compile_evaluator(bucket_size)

        # Pad parameters to bucket
        padded_params, _ = pad_parameters_to_bucket(norm_params, bucket_size)

        # Build positional arguments according to parameter_order
        args = []
        for k in self.parameter_order:
            if k in padded_params:
                args.append(jnp.asarray(padded_params[k], dtype=jnp.float64))
            else:
                default_v = self.default_kwargs.get(k, 0.0)
                args.append(jnp.full((bucket_size,), float(default_v), dtype=jnp.float64))

        hp_padded, hc_padded = evaluator(*args)
        return unpad_waveform_from_bucket(hp_padded, hc_padded, orig_batch)

__call__(parameters, **kwargs)

Evaluate waveforms for an arbitrary batch size with automatic bucketing.

Source code in src/jaxwave/inference/bucketing.py

def __call__(
    self,
    parameters: Dict[str, Any],
    **kwargs,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Evaluate waveforms for an arbitrary batch size with automatic bucketing."""
    norm_params = {}
    name_map = {
        "mass_1": "mass1",
        "mass_2": "mass2",
        "spin_1z": "spin1z",
        "spin_2z": "spin2z",
        "luminosity_distance": "distance",
        "theta_jn": "inclination",
        "phase": "phic",
        "geocent_time": "tc",
    }
    for k, v in parameters.items():
        norm_params[name_map.get(k, k)] = v
    norm_params.update(kwargs)

    # Determine original batch size
    first_val = next(iter(norm_params.values()))
    first_arr = jnp.asarray(first_val)
    orig_batch = first_arr.shape[0] if first_arr.ndim > 0 else 1

    bucket_size = get_bucketed_batch_size(orig_batch, self.buckets)
    evaluator = self._get_or_compile_evaluator(bucket_size)

    # Pad parameters to bucket
    padded_params, _ = pad_parameters_to_bucket(norm_params, bucket_size)

    # Build positional arguments according to parameter_order
    args = []
    for k in self.parameter_order:
        if k in padded_params:
            args.append(jnp.asarray(padded_params[k], dtype=jnp.float64))
        else:
            default_v = self.default_kwargs.get(k, 0.0)
            args.append(jnp.full((bucket_size,), float(default_v), dtype=jnp.float64))

    hp_padded, hc_padded = evaluator(*args)
    return unpad_waveform_from_bucket(hp_padded, hc_padded, orig_batch)

jaxwave.inference.bucketing.get_bucketed_batch_size(batch_size, buckets=DEFAULT_BATCH_BUCKETS)

Find the smallest bucket capacity greater than or equal to batch_size.

If batch_size exceeds all predefined buckets, rounds up to the next power of 2.

Source code in src/jaxwave/inference/bucketing.py

def get_bucketed_batch_size(batch_size: int, buckets: Sequence[int] = DEFAULT_BATCH_BUCKETS) -> int:
    """Find the smallest bucket capacity greater than or equal to batch_size.

    If batch_size exceeds all predefined buckets, rounds up to the next power of 2.
    """
    if batch_size <= 0:
        raise ValueError(f"batch_size must be positive, got {batch_size}")
    for b in buckets:
        if b >= batch_size:
            return b
    # Round up to next power of 2
    return 1 << (batch_size - 1).bit_length()

jaxwave.inference.bucketing.pad_parameters_to_bucket(parameters, bucket_size)

Pad parameter arrays along axis 0 up to bucket_size.

Pads using the last value of each parameter array to reuse an existing parameter row.

Returns:

Name	Type	Description
padded_parameters	Dict[str, ndarray]	Dict with arrays of length bucket_size.
valid_mask	ndarray	Boolean array of shape (bucket_size,) indicating valid (non-padded) entries.

Source code in src/jaxwave/inference/bucketing.py

def pad_parameters_to_bucket(
    parameters: Dict[str, jnp.ndarray],
    bucket_size: int,
) -> Tuple[Dict[str, jnp.ndarray], jnp.ndarray]:
    """Pad parameter arrays along axis 0 up to bucket_size.

    Pads using the last value of each parameter array to reuse an existing parameter row.

    Returns:
        padded_parameters: Dict with arrays of length bucket_size.
        valid_mask: Boolean array of shape (bucket_size,) indicating valid (non-padded) entries.
    """
    padded = {}
    valid_mask = None

    for k, v in parameters.items():
        arr = jnp.asarray(v)
        if arr.ndim == 0:
            arr = arr[None]
        n_orig = arr.shape[0]
        if n_orig > bucket_size:
            raise ValueError(f"Parameter '{k}' length ({n_orig}) exceeds bucket_size ({bucket_size})")

        pad_len = bucket_size - n_orig
        if pad_len > 0:
            # Pad with the last value along axis 0
            last_val = arr[-1:]
            pad_shape = (pad_len,) + arr.shape[1:]
            fill_arr = jnp.broadcast_to(last_val, pad_shape)
            padded[k] = jnp.concatenate([arr, fill_arr], axis=0)
        else:
            padded[k] = arr

        if valid_mask is None:
            valid_mask = jnp.arange(bucket_size) < n_orig

    if valid_mask is None:
        valid_mask = jnp.ones((bucket_size,), dtype=bool)

    return padded, valid_mask

jaxwave.inference.bucketing.unpad_waveform_from_bucket(hp, hc, original_batch_size)

Slice evaluated batched waveforms back to the original proposal batch size.

Source code in src/jaxwave/inference/bucketing.py

def unpad_waveform_from_bucket(
    hp: jnp.ndarray,
    hc: jnp.ndarray,
    original_batch_size: int,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Slice evaluated batched waveforms back to the original proposal batch size."""
    if hp.ndim >= 2:
        return hp[:original_batch_size], hc[:original_batch_size]
    return hp, hc

---

Bayesian Inference & Fisher Matrix

jaxwave.inference.fisher.fisher_information_matrix(parameters_dict, waveform_fn, psd_dict, detector_network=None, f_low=20.0, f_high=1024.0, df=None, frequencies=None, parameter_order=None, method='autodiff', step_sizes=None)

Compute the differentiable Fisher Information Matrix Gamma_ij using JAX autodiff or parameter differences.

Literature: Cutler & Flanagan (1994), Vallisneri (2008).

Formula

Gamma_ij = sum_{k in network} 4 Re int_{f_low}^{f_high} (1 / Sn^(k)(f)) * (dh_k(f)/dtheta_i) * (dh_k(f)/dtheta_j)* df

Parameters:

Name	Type	Description	Default
parameters_dict	Dict[str, Union[float, ndarray]]	Dict mapping parameter names to parameter values/arrays	required
waveform_fn	Callable[[Dict[str, ndarray]], Union[ndarray, Dict[str, ndarray], Tuple[ndarray, ndarray]]]	Function mapping parameter dictionary to frequency-domain strain(s)	required
psd_dict	Union[ndarray, Dict[str, ndarray]]	Real PSD array or dict of PSDs for each detector {det: psd_det}	required
detector_network	Optional[Any]	Optional DetectorNetwork instance to project (h_plus, h_cross)	None
f_low	float	Low-frequency cutoff in Hz	20.0
f_high	float	High-frequency cutoff in Hz	1024.0
df	Optional[float]	Frequency resolution in Hz	None
frequencies	Optional[ndarray]	Optional frequency vector in Hz	None
parameter_order	Optional[Sequence[str]]	Optional explicit order of parameter keys	None
method	str	"autodiff" (forward-mode JVP), "hessian" (autodiff Whittle Hessian), or "finite_diff"	'autodiff'
step_sizes	Optional[Dict[str, float]]	Optional dict of step sizes for finite differences	None

Returns:

Type	Description
(fisher_matrix, param_names)	(P, P) symmetric Fisher matrix and list of parameter names.
List[str]	Degenerate parameters can make the matrix singular.

Source code in src/jaxwave/inference/fisher.py

def fisher_information_matrix(
    parameters_dict: Dict[str, Union[float, jnp.ndarray]],
    waveform_fn: Callable[[Dict[str, jnp.ndarray]], Union[jnp.ndarray, Dict[str, jnp.ndarray], Tuple[jnp.ndarray, jnp.ndarray]]],
    psd_dict: Union[jnp.ndarray, Dict[str, jnp.ndarray]],
    detector_network: Optional[Any] = None,
    f_low: float = 20.0,
    f_high: float = 1024.0,
    df: Optional[float] = None,
    frequencies: Optional[jnp.ndarray] = None,
    parameter_order: Optional[Sequence[str]] = None,
    method: str = "autodiff",
    step_sizes: Optional[Dict[str, float]] = None,
) -> Tuple[jnp.ndarray, List[str]]:
    """Compute the differentiable Fisher Information Matrix Gamma_ij using JAX autodiff or parameter differences.

    Literature: Cutler & Flanagan (1994), Vallisneri (2008).

    Formula:
        Gamma_ij = sum_{k in network} 4 Re int_{f_low}^{f_high} (1 / Sn^(k)(f)) * (dh_k(f)/dtheta_i) * (dh_k(f)/dtheta_j)* df

    Args:
        parameters_dict: Dict mapping parameter names to parameter values/arrays
        waveform_fn: Function mapping parameter dictionary to frequency-domain strain(s)
        psd_dict: Real PSD array or dict of PSDs for each detector {det: psd_det}
        detector_network: Optional DetectorNetwork instance to project (h_plus, h_cross)
        f_low: Low-frequency cutoff in Hz
        f_high: High-frequency cutoff in Hz
        df: Frequency resolution in Hz
        frequencies: Optional frequency vector in Hz
        parameter_order: Optional explicit order of parameter keys
        method: "autodiff" (forward-mode JVP), "hessian" (autodiff Whittle Hessian), or "finite_diff"
        step_sizes: Optional dict of step sizes for finite differences

    Returns:
        (fisher_matrix, param_names): (P, P) symmetric Fisher matrix and list of parameter names.
        Degenerate parameters can make the matrix singular.
    """
    if parameter_order is None:
        param_names = list(parameters_dict.keys())
    else:
        param_names = list(parameter_order)

    base_params = {
        k: as_array(parameters_dict[k], dtype=jnp.float64)
        for k in param_names
    }
    num_params = len(param_names)

    def eval_strains(p_dict: Dict[str, jnp.ndarray]) -> Dict[str, jnp.ndarray]:
        wf_out = waveform_fn(p_dict)
        if isinstance(wf_out, dict):
            return wf_out
        elif isinstance(wf_out, tuple) and len(wf_out) == 2:
            hp, hc = wf_out
            if detector_network is not None:
                ra = p_dict.get("ra", jnp.array(0.0, dtype=jnp.float64))
                dec = p_dict.get("dec", jnp.array(0.0, dtype=jnp.float64))
                psi = p_dict.get("psi", jnp.array(0.0, dtype=jnp.float64))
                tg = p_dict.get("t_gps", p_dict.get("tc", jnp.array(0.0, dtype=jnp.float64)))
                freq_vec = frequencies if frequencies is not None else jnp.array(0.0, dtype=jnp.float64)
                return detector_network.project_fd_strain(freq_vec, hp, hc, ra, dec, psi, tg)
            else:
                return {"det": hp}
        elif isinstance(wf_out, (jnp.ndarray, np.ndarray)):
            return {"det": wf_out}
        else:
            raise TypeError(f"Unsupported waveform_fn output type: {type(wf_out)}")

    if method == "hessian":
        # Compute Fisher matrix via the autodiff Hessian of the Whittle log-likelihood at fiducial theta
        def neg_log_lik_vec(theta_vec: jnp.ndarray) -> jnp.ndarray:
            p_dict = {name: theta_vec[idx] for idx, name in enumerate(param_names)}
            strains = eval_strains(p_dict)
            base_strains = eval_strains(base_params)
            total = jnp.array(0.0, dtype=jnp.float64)
            for det_k, h_k in strains.items():
                if isinstance(psd_dict, dict):
                    if det_k not in psd_dict:
                        if len(psd_dict) == 1:
                            psd_k = next(iter(psd_dict.values()))
                        else:
                            continue
                    else:
                        psd_k = psd_dict[det_k]
                else:
                    psd_k = psd_dict

                d_k = base_strains[det_k]
                term_dh = noise_weighted_inner_product(
                    d_k, h_k, psd_k, df=df, frequencies=frequencies, f_low=f_low, f_high=f_high
                )
                term_hh = noise_weighted_inner_product(
                    h_k, h_k, psd_k, df=df, frequencies=frequencies, f_low=f_low, f_high=f_high
                )
                total = total + (term_dh - 0.5 * term_hh)
            return -total

        theta_0 = jnp.array([base_params[k] for k in param_names], dtype=jnp.float64)
        fisher = jax.hessian(neg_log_lik_vec)(theta_0)
        fisher = 0.5 * (fisher + fisher.T)
        return fisher, param_names

    det_jacobians: Dict[str, List[jnp.ndarray]] = {}

    if method == "autodiff":
        for i, p_name in enumerate(param_names):
            tangent = {
                k: (jnp.array(1.0, dtype=jnp.float64) if k == p_name else jnp.array(0.0, dtype=jnp.float64))
                for k in param_names
            }
            _, dh_dict = jax.jvp(eval_strains, (base_params,), (tangent,))
            for det_k, dh_i in dh_dict.items():
                if det_k not in det_jacobians:
                    det_jacobians[det_k] = []
                det_jacobians[det_k].append(dh_i)
    elif method == "finite_diff":
        for i, p_name in enumerate(param_names):
            val = float(base_params[p_name])
            if step_sizes is not None and p_name in step_sizes:
                delta = float(step_sizes[p_name])
            else:
                delta = 1.0e-6 * max(abs(val), 1.0)

            p_plus = dict(base_params)
            p_plus[p_name] = base_params[p_name] + delta
            p_minus = dict(base_params)
            p_minus[p_name] = base_params[p_name] - delta

            s_plus = eval_strains(p_plus)
            s_minus = eval_strains(p_minus)

            for det_k in s_plus:
                dh_i = (s_plus[det_k] - s_minus[det_k]) / (2.0 * delta)
                if det_k not in det_jacobians:
                    det_jacobians[det_k] = []
                det_jacobians[det_k].append(dh_i)
    else:
        raise ValueError(f"Unknown method '{method}', supported: 'autodiff', 'hessian', 'finite_diff'")

    fisher = jnp.zeros((num_params, num_params), dtype=jnp.float64)

    for det_k, jac_list in det_jacobians.items():
        if isinstance(psd_dict, dict):
            if det_k not in psd_dict:
                if len(psd_dict) == 1:
                    psd_k = next(iter(psd_dict.values()))
                else:
                    continue
            else:
                psd_k = psd_dict[det_k]
        else:
            psd_k = psd_dict

        psd_k = as_array(psd_k, dtype=jnp.float64)
        psd_safe = jnp.maximum(psd_k, 1.0e-50)
        weight = 4.0 / psd_safe

        if frequencies is not None:
            freqs = as_array(frequencies, dtype=jnp.float64)
            mask = jnp.ones(freqs.shape, dtype=bool)
            if f_low is not None:
                mask = mask & (freqs >= f_low)
            if f_high is not None:
                mask = mask & (freqs <= f_high)
            weight = jnp.where(mask, weight, 0.0)
            integration_mode = "trapezoid"
        else:
            df_val = df if df is not None else 1.0
            weight = weight * df_val
            integration_mode = "sum"

        J = jnp.stack(jac_list, axis=-1)

        if integration_mode == "trapezoid":
            for i in range(num_params):
                for j in range(i, num_params):
                    integrand = weight * jnp.real(jac_list[i] * jnp.conj(jac_list[j]))
                    integral = _trapezoid(integrand, freqs, axis=-1)
                    fisher = fisher.at[i, j].add(integral)
                    if i != j:
                        fisher = fisher.at[j, i].add(integral)
        else:
            W = weight[:, None]
            W_J = W * J
            Gamma_det = jnp.real(jnp.matmul(jnp.conj(J).T, W_J))
            Gamma_det = 0.5 * (Gamma_det + Gamma_det.T)
            fisher = fisher + Gamma_det

    return fisher, param_names

jaxwave.inference.relative_binning.RelativeBinningLikelihood

Approximate Whittle log-likelihood using relative binning (Zackay et al. 2018).

Precomputes summary coefficients from dense data and a fiducial waveform, then evaluates candidate waveforms at M bin frequencies instead of N_f dense frequencies. The approximation uses local linear waveform ratios.

Accuracy and evaluation cost depend on binning, the reference waveform, signal properties, and parameter displacement. Validate likelihood and gradient errors against dense evaluation over the intended parameter region. Autodiff differentiates this approximation, not the dense likelihood.

Source code in src/jaxwave/inference/relative_binning.py

@jax.tree_util.register_pytree_node_class
class RelativeBinningLikelihood:
    """Approximate Whittle log-likelihood using relative binning (Zackay et al. 2018).

    Precomputes summary coefficients from dense data and a fiducial waveform,
    then evaluates candidate waveforms at M bin frequencies instead of N_f dense
    frequencies. The approximation uses local linear waveform ratios.

    Accuracy and evaluation cost depend on binning, the reference waveform,
    signal properties, and parameter displacement. Validate likelihood and
    gradient errors against dense evaluation over the intended parameter region.
    Autodiff differentiates this approximation, not the dense likelihood.
    """

    def __init__(
        self,
        data: Dict[str, jnp.ndarray],
        psds: Dict[str, jnp.ndarray],
        h0: Dict[str, jnp.ndarray],
        frequencies: jnp.ndarray,
        bin_edges: Optional[jnp.ndarray] = None,
        n_bins: int = 50,
        f_min: Optional[float] = None,
        f_max: Optional[float] = None,
        df: Optional[float] = None,
    ):
        self.detectors = tuple(sorted(list(data.keys())))
        f_dense = as_array(frequencies, dtype=jnp.float64)
        self.frequencies = f_dense

        min_f = float(f_min) if f_min is not None else float(f_dense[0])
        max_f = float(f_max) if f_max is not None else float(f_dense[-1])

        if bin_edges is None:
            edges = compute_relative_binning_bins(
                f_min=min_f,
                f_max=max_f,
                n_bins=n_bins,
                frequencies=f_dense,
                dtype=jnp.float64,
            )
        else:
            edges = as_array(bin_edges, dtype=jnp.float64)

        self.bin_edges = edges
        self.n_bins = len(edges) - 1

        # Use the same dense-grid frequencies for fiducial and candidate evaluations
        raw_centers = 0.5 * (edges[:-1] + edges[1:])
        center_indices = jnp.argmin(jnp.abs(f_dense[:, None] - raw_centers[None, :]), axis=0)
        centers = f_dense[center_indices]
        self.bin_centers = centers
        self.center_indices = center_indices

        # Precompute summary coefficients for each detector
        self.a0: Dict[str, jnp.ndarray] = {}
        self.a1: Dict[str, jnp.ndarray] = {}
        self.b0: Dict[str, jnp.ndarray] = {}
        self.b1: Dict[str, jnp.ndarray] = {}
        self.h0_centers: Dict[str, jnp.ndarray] = {}
        self.dd: Dict[str, jnp.ndarray] = {}

        df_val = float(df) if df is not None else float(f_dense[1] - f_dense[0])

        for det in self.detectors:
            d_tensor = as_array(data[det], dtype=jnp.complex128 if jnp.iscomplexobj(data[det]) else jnp.float64)
            psd_tensor = as_array(psds[det], dtype=jnp.float64)
            h0_tensor = as_array(h0[det], dtype=jnp.complex128 if jnp.iscomplexobj(h0[det]) else jnp.float64)

            a0, _, b0, _, _, d_d_val = compute_relative_binning_summary(
                data=d_tensor,
                psd=psd_tensor,
                h0=h0_tensor,
                frequencies=f_dense,
                bin_edges=edges,
                df=df_val,
            )

            # Re-evaluate a1 and b1 with exact snapped centers
            safe_psd = jnp.maximum(psd_tensor, 1.0e-50)
            integrand_a = 4.0 * ((d_tensor * jnp.conj(h0_tensor)) / safe_psd) * df_val
            integrand_b = 4.0 * (jnp.real(h0_tensor * jnp.conj(h0_tensor)) / safe_psd) * df_val

            f_row = f_dense[None, :]
            mask = (f_row >= edges[:-1, None]) & (f_row < edges[1:, None])
            mask = mask.at[-1].set((f_dense >= edges[-2]) & (f_dense <= edges[-1]))
            delta_f = f_row - centers[:, None]

            a1 = jnp.sum(jnp.where(mask, integrand_a[None, :] * delta_f, 0.0), axis=-1)
            b1 = jnp.sum(jnp.where(mask, integrand_b[None, :] * delta_f, 0.0), axis=-1)

            self.a0[det] = a0
            self.a1[det] = a1
            self.b0[det] = b0
            self.b1[det] = b1
            self.h0_centers[det] = h0_tensor[center_indices]
            self.dd[det] = d_d_val

    def tree_flatten(self):
        children = (
            self.frequencies,
            self.bin_edges,
            self.bin_centers,
            self.center_indices,
            self.a0,
            self.a1,
            self.b0,
            self.b1,
            self.h0_centers,
            self.dd,
        )
        aux_data = (self.detectors, self.n_bins)
        return children, aux_data

    @classmethod
    def tree_unflatten(cls, aux_data, children):
        obj = cls.__new__(cls)
        obj.detectors, obj.n_bins = aux_data
        (
            obj.frequencies,
            obj.bin_edges,
            obj.bin_centers,
            obj.center_indices,
            obj.a0,
            obj.a1,
            obj.b0,
            obj.b1,
            obj.h0_centers,
            obj.dd,
        ) = children
        return obj

    @staticmethod
    def compute_r0_r1(
        h_centers: jnp.ndarray,
        h0_centers: jnp.ndarray,
        bin_centers: jnp.ndarray,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """Compute the ratio r_0 and slope r_1 at bin centers in JAX.

        Args:
            h_centers: Test waveform evaluated at bin centers, shape (..., M).
            h0_centers: Reference waveform evaluated at bin centers, shape (M,) or (..., M).
            bin_centers: Bin center frequencies, shape (M,).

        Returns:
            Tuple of (r0, r1) with same batch dimensions as h_centers.
        """
        safe_h0 = jnp.where(
            jnp.abs(h0_centers) < 1e-40,
            jnp.ones_like(h0_centers) * 1e-40,
            h0_centers,
        )
        r0 = h_centers / safe_h0

        m = bin_centers.shape[-1]
        if m >= 3:
            df_interior = bin_centers[2:] - bin_centers[:-2]
            r1_interior = (r0[..., 2:] - r0[..., :-2]) / df_interior

            df_0 = bin_centers[1] - bin_centers[0]
            r1_0 = (r0[..., 1:2] - r0[..., 0:1]) / df_0

            df_end = bin_centers[-1] - bin_centers[-2]
            r1_end = (r0[..., -1:] - r0[..., -2:-1]) / df_end

            r1 = jnp.concatenate([r1_0, r1_interior, r1_end], axis=-1)
        elif m == 2:
            df_step = bin_centers[1] - bin_centers[0]
            diff = (r0[..., 1:2] - r0[..., 0:1]) / df_step
            r1 = jnp.concatenate([diff, diff], axis=-1)
        else:
            r1 = jnp.zeros_like(r0)

        return r0, r1

    def inner_product_dh(
        self,
        r0: jnp.ndarray,
        r1: jnp.ndarray,
        detector: str,
    ) -> jnp.ndarray:
        """Compute noise-weighted inner product <d | h(theta)> from summary coefficients.

        <d | h> = Re sum_{k=0}^{M-1} [ r_0^*(f_bar_k) * A_{0, k} + r_1^*(f_bar_k) * A_{1, k} ]
        """
        a0 = self.a0[detector]
        a1 = self.a1[detector]

        r0_conj = jnp.conj(r0)
        r1_conj = jnp.conj(r1)

        term0 = r0_conj * a0
        term1 = r1_conj * a1

        sum_terms = jnp.sum(term0 + term1, axis=-1)
        return jnp.real(sum_terms)

    def inner_product_hh(
        self,
        r0: jnp.ndarray,
        r1: jnp.ndarray,
        detector: str,
    ) -> jnp.ndarray:
        """Compute waveform self inner product <h(theta) | h(theta)> from summary coefficients.

        <h | h> = sum_{k=0}^{M-1} [ |r_0(f_bar_k)|^2 * B_{0, k} + 2 * Re(r_0 * r_1^*) * B_{1, k} ]
        """
        b0 = self.b0[detector]
        b1 = self.b1[detector]

        r0_sq = jnp.real(r0 * jnp.conj(r0))
        cross_term = 2.0 * jnp.real(r0 * jnp.conj(r1))

        sum_terms = jnp.sum(r0_sq * b0 + cross_term * b1, axis=-1)
        return sum_terms

    def log_likelihood_from_r(
        self,
        r0_dict: Dict[str, jnp.ndarray],
        r1_dict: Dict[str, jnp.ndarray],
    ) -> jnp.ndarray:
        """Compute relative binning log-likelihood from pre-computed ratios (r0, r1).

        ln L = sum_{det} [ <d | h> - 0.5 * <h | h> - 0.5 * <d | d> ]
        """
        total_ln_l = jnp.array(0.0, dtype=jnp.float64)

        for det in self.detectors:
            r0 = r0_dict[det]
            r1 = r1_dict[det]
            dh = self.inner_product_dh(r0, r1, det)
            hh = self.inner_product_hh(r0, r1, det)
            dd = self.dd[det]

            det_ln_l = dh - 0.5 * hh - 0.5 * dd
            total_ln_l = total_ln_l + det_ln_l

        return total_ln_l

    def log_likelihood(
        self,
        templates_at_bin_centers: Union[Dict[str, jnp.ndarray], jnp.ndarray],
        h0_at_bin_centers: Optional[Dict[str, jnp.ndarray]] = None,
        slopes: Optional[Dict[str, jnp.ndarray]] = None,
    ) -> jnp.ndarray:
        """Compute relative binning log-likelihood from waveforms evaluated at bin centers.

        Args:
            templates_at_bin_centers: Dictionary mapping detector name to waveform
                array h(theta; f_bar) evaluated ONLY at bin centers, or single array
                if only one detector is configured.
            h0_at_bin_centers: Optional fiducial waveform h0(f_bar).
            slopes: Optional explicit slopes r_1(f_bar). If None, calculated via
                centered finite differences.

        Returns:
            Scalar or batched log-likelihood array.
        """
        r0_dict: Dict[str, jnp.ndarray] = {}
        r1_dict: Dict[str, jnp.ndarray] = {}

        for det in self.detectors:
            if isinstance(templates_at_bin_centers, dict):
                h_det = templates_at_bin_centers[det]
            else:
                h_det = templates_at_bin_centers

            if h0_at_bin_centers is not None:
                h0_det = h0_at_bin_centers[det]
            else:
                h0_det = self.h0_centers[det]

            if slopes is not None and det in slopes:
                safe_h0 = jnp.where(
                    jnp.abs(h0_det) < 1e-40,
                    jnp.ones_like(h0_det) * 1e-40,
                    h0_det,
                )
                r0 = h_det / safe_h0
                r1 = slopes[det]
            else:
                r0, r1 = self.compute_r0_r1(h_det, h0_det, self.bin_centers)

            r0_dict[det] = r0
            r1_dict[det] = r1

        return self.log_likelihood_from_r(r0_dict, r1_dict)

    def forward(
        self,
        templates_at_bin_centers: Union[Dict[str, jnp.ndarray], jnp.ndarray],
        **kwargs,
    ) -> jnp.ndarray:
        return self.log_likelihood(templates_at_bin_centers, **kwargs)

    def __call__(self, *args, **kwargs) -> jnp.ndarray:
        return self.forward(*args, **kwargs)

compute_r0_r1(h_centers, h0_centers, bin_centers) staticmethod

Compute the ratio r_0 and slope r_1 at bin centers in JAX.

Parameters:

Name	Type	Description	Default
h_centers	ndarray	Test waveform evaluated at bin centers, shape (..., M).	required
h0_centers	ndarray	Reference waveform evaluated at bin centers, shape (M,) or (..., M).	required
bin_centers	ndarray	Bin center frequencies, shape (M,).	required

Returns:

Type	Description
Tuple[ndarray, ndarray]	Tuple of (r0, r1) with same batch dimensions as h_centers.

Source code in src/jaxwave/inference/relative_binning.py

@staticmethod
def compute_r0_r1(
    h_centers: jnp.ndarray,
    h0_centers: jnp.ndarray,
    bin_centers: jnp.ndarray,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """Compute the ratio r_0 and slope r_1 at bin centers in JAX.

    Args:
        h_centers: Test waveform evaluated at bin centers, shape (..., M).
        h0_centers: Reference waveform evaluated at bin centers, shape (M,) or (..., M).
        bin_centers: Bin center frequencies, shape (M,).

    Returns:
        Tuple of (r0, r1) with same batch dimensions as h_centers.
    """
    safe_h0 = jnp.where(
        jnp.abs(h0_centers) < 1e-40,
        jnp.ones_like(h0_centers) * 1e-40,
        h0_centers,
    )
    r0 = h_centers / safe_h0

    m = bin_centers.shape[-1]
    if m >= 3:
        df_interior = bin_centers[2:] - bin_centers[:-2]
        r1_interior = (r0[..., 2:] - r0[..., :-2]) / df_interior

        df_0 = bin_centers[1] - bin_centers[0]
        r1_0 = (r0[..., 1:2] - r0[..., 0:1]) / df_0

        df_end = bin_centers[-1] - bin_centers[-2]
        r1_end = (r0[..., -1:] - r0[..., -2:-1]) / df_end

        r1 = jnp.concatenate([r1_0, r1_interior, r1_end], axis=-1)
    elif m == 2:
        df_step = bin_centers[1] - bin_centers[0]
        diff = (r0[..., 1:2] - r0[..., 0:1]) / df_step
        r1 = jnp.concatenate([diff, diff], axis=-1)
    else:
        r1 = jnp.zeros_like(r0)

    return r0, r1

inner_product_dh(r0, r1, detector)

Compute noise-weighted inner product from summary coefficients.

= Re sum_{k=0}^{M-1} [ r_0^(f_bar_k) * A_{0, k} + r_1^(f_bar_k) * A_{1, k} ]

Source code in src/jaxwave/inference/relative_binning.py

def inner_product_dh(
    self,
    r0: jnp.ndarray,
    r1: jnp.ndarray,
    detector: str,
) -> jnp.ndarray:
    """Compute noise-weighted inner product <d | h(theta)> from summary coefficients.

    <d | h> = Re sum_{k=0}^{M-1} [ r_0^*(f_bar_k) * A_{0, k} + r_1^*(f_bar_k) * A_{1, k} ]
    """
    a0 = self.a0[detector]
    a1 = self.a1[detector]

    r0_conj = jnp.conj(r0)
    r1_conj = jnp.conj(r1)

    term0 = r0_conj * a0
    term1 = r1_conj * a1

    sum_terms = jnp.sum(term0 + term1, axis=-1)
    return jnp.real(sum_terms)

inner_product_hh(r0, r1, detector)

Compute waveform self inner product from summary coefficients.

= sum_{k=0}^{M-1} [ |r_0(f_bar_k)|^2 * B_{0, k} + 2 * Re(r_0 * r_1^*) * B_{1, k} ]

Source code in src/jaxwave/inference/relative_binning.py

def inner_product_hh(
    self,
    r0: jnp.ndarray,
    r1: jnp.ndarray,
    detector: str,
) -> jnp.ndarray:
    """Compute waveform self inner product <h(theta) | h(theta)> from summary coefficients.

    <h | h> = sum_{k=0}^{M-1} [ |r_0(f_bar_k)|^2 * B_{0, k} + 2 * Re(r_0 * r_1^*) * B_{1, k} ]
    """
    b0 = self.b0[detector]
    b1 = self.b1[detector]

    r0_sq = jnp.real(r0 * jnp.conj(r0))
    cross_term = 2.0 * jnp.real(r0 * jnp.conj(r1))

    sum_terms = jnp.sum(r0_sq * b0 + cross_term * b1, axis=-1)
    return sum_terms

log_likelihood(templates_at_bin_centers, h0_at_bin_centers=None, slopes=None)

Compute relative binning log-likelihood from waveforms evaluated at bin centers.

Parameters:

Name	Type	Description	Default
templates_at_bin_centers	Union[Dict[str, ndarray], ndarray]	Dictionary mapping detector name to waveform array h(theta; f_bar) evaluated ONLY at bin centers, or single array if only one detector is configured.	required
h0_at_bin_centers	Optional[Dict[str, ndarray]]	Optional fiducial waveform h0(f_bar).	None
slopes	Optional[Dict[str, ndarray]]	Optional explicit slopes r_1(f_bar). If None, calculated via centered finite differences.	None

Returns:

Type	Description
ndarray	Scalar or batched log-likelihood array.

Source code in src/jaxwave/inference/relative_binning.py

def log_likelihood(
    self,
    templates_at_bin_centers: Union[Dict[str, jnp.ndarray], jnp.ndarray],
    h0_at_bin_centers: Optional[Dict[str, jnp.ndarray]] = None,
    slopes: Optional[Dict[str, jnp.ndarray]] = None,
) -> jnp.ndarray:
    """Compute relative binning log-likelihood from waveforms evaluated at bin centers.

    Args:
        templates_at_bin_centers: Dictionary mapping detector name to waveform
            array h(theta; f_bar) evaluated ONLY at bin centers, or single array
            if only one detector is configured.
        h0_at_bin_centers: Optional fiducial waveform h0(f_bar).
        slopes: Optional explicit slopes r_1(f_bar). If None, calculated via
            centered finite differences.

    Returns:
        Scalar or batched log-likelihood array.
    """
    r0_dict: Dict[str, jnp.ndarray] = {}
    r1_dict: Dict[str, jnp.ndarray] = {}

    for det in self.detectors:
        if isinstance(templates_at_bin_centers, dict):
            h_det = templates_at_bin_centers[det]
        else:
            h_det = templates_at_bin_centers

        if h0_at_bin_centers is not None:
            h0_det = h0_at_bin_centers[det]
        else:
            h0_det = self.h0_centers[det]

        if slopes is not None and det in slopes:
            safe_h0 = jnp.where(
                jnp.abs(h0_det) < 1e-40,
                jnp.ones_like(h0_det) * 1e-40,
                h0_det,
            )
            r0 = h_det / safe_h0
            r1 = slopes[det]
        else:
            r0, r1 = self.compute_r0_r1(h_det, h0_det, self.bin_centers)

        r0_dict[det] = r0
        r1_dict[det] = r1

    return self.log_likelihood_from_r(r0_dict, r1_dict)

log_likelihood_from_r(r0_dict, r1_dict)

Compute relative binning log-likelihood from pre-computed ratios (r0, r1).

ln L = sum_{det} [ - 0.5 * - 0.5 * ]

Source code in src/jaxwave/inference/relative_binning.py

def log_likelihood_from_r(
    self,
    r0_dict: Dict[str, jnp.ndarray],
    r1_dict: Dict[str, jnp.ndarray],
) -> jnp.ndarray:
    """Compute relative binning log-likelihood from pre-computed ratios (r0, r1).

    ln L = sum_{det} [ <d | h> - 0.5 * <h | h> - 0.5 * <d | d> ]
    """
    total_ln_l = jnp.array(0.0, dtype=jnp.float64)

    for det in self.detectors:
        r0 = r0_dict[det]
        r1 = r1_dict[det]
        dh = self.inner_product_dh(r0, r1, det)
        hh = self.inner_product_hh(r0, r1, det)
        dd = self.dd[det]

        det_ln_l = dh - 0.5 * hh - 0.5 * dd
        total_ln_l = total_ln_l + det_ln_l

    return total_ln_l