PyAETNA 4.0.0__py3-none-any.whl

PyAETNA/body_models.py

@@ -0,0 +1,2336 @@
+from __future__ import annotations
+
+from functools import partial
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+from astropy.constants import R_sun, c
+from jax import jit
+from jaxoplanet.light_curves import limb_dark_light_curve
+from jaxoplanet.light_curves.emission import light_curve as emission_lc
+from jaxoplanet.orbits.keplerian import Body, Central, System
+from jaxoplanet.orbits.ttv import TTVOrbit
+from jaxoplanet.starry.core.basis import A1 as _jxp_A1
+from jaxoplanet.starry.core.basis import A2_inv as _jxp_A2_inv
+from jaxoplanet.starry.core.basis import U as _jxp_U
+from jaxoplanet.starry.core.polynomials import Pijk as _jxp_Pijk
+from jaxoplanet.starry.core.rotation import dot_rz as _jxp_dot_rz
+from jaxoplanet.starry.core.solution import rT as _jxp_rT
+from jaxoplanet.starry.core.solution import solution_vector as _jxp_solution_vector
+from jaxoplanet.starry.doppler import rv_map_expansion as _jxp_rv_map_expansion
+from jaxoplanet.starry.orbit import SurfaceSystem
+from jaxoplanet.starry.surface import Surface
+
+# Import physical constants from centralized location
+from .utils.constants import G_CGS, RHO_SUN_CGS, SECONDS_PER_DAY
+
+# Speed of light in R_sun per day.
+# Use astropy constants for the speed of light and solar radius so that
+# the value is not hard-coded.  c.cgs.value is in cm/s and
+# R_sun.cgs.value is in cm.
+# Used for light-travel-time correction in phase curves.
+_C_RSUN_PER_DAY = (c.cgs.value / R_sun.cgs.value) * SECONDS_PER_DAY  # ≈ 37228 R_sun/day
+
+# Conversion factor: 1 R_sun / day → m/s
+# R_sun.cgs.value [cm] / SECONDS_PER_DAY [s] / 100 [cm→m]
+_RSUN_DAY_TO_MS = float(R_sun.cgs.value / SECONDS_PER_DAY / 100.0)  # ≈ 8052.08 m/s
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Pre-computed constants for the Rossiter-McLaughlin (RM) model
+# ──────────────────────────────────────────────────────────────────────────────
+# These constants are evaluated ONCE at module import time from scipy/jaxoplanet
+# basis matrices and stored as float32 JAX arrays.  They encode the polynomial
+# change-of-basis and occultation geometry for our fixed configuration:
+#   ydeg = 0  (no explicit stellar surface map — uniform disk)
+#   udeg = 2  (quadratic limb darkening, 2 free coefficients)
+#   vdeg = 1  (solid-body rotation velocity map, no differential rotation)
+#   deg  = max(ydeg, udeg) = 2   →  n_deg  = (deg+1)^2  = 9 coefficients
+#   total_deg = deg + vdeg  = 3   →  n_total = (total_deg+1)^2 = 16 coefficients
+#
+# Background:
+#   jaxoplanet.starry uses the Pijk polynomial basis from Luger et al. (2019).
+#   The change-of-basis matrix A2 converts from the occultation (sT) basis to the
+#   Pijk polynomial basis.  A1(deg) converts from the spherical-harmonic (Ylm) basis.
+#   These matrices depend only on integer degrees, so they are Python-level constants.
+#
+# WHY we precompute here (not inside the function):
+#   surface_radial_velocity calls scipy.sparse.linalg.inv(A2_inv(deg)) at TRACE
+#   TIME every time JAX compiles the model.  Pre-computing avoids repeated expensive
+#   scipy calls and produces a compile-time constant that XLA can fold into the
+#   computation, yielding faster and more predictable JIT compilation.
+import scipy.sparse as _scipy_sparse  # imported here to keep the namespace clean
+
+_RM_YDEG = 0  # stellar surface map degree (uniform)
+_RM_UDEG = 2  # limb-darkening degree (quadratic)
+_RM_VDEG = 1  # velocity map degree (solid-body rotation)
+_RM_DEG = max(_RM_YDEG, _RM_UDEG)  # = 2
+_RM_TOTAL_DEG = _RM_DEG + _RM_VDEG  # = 3
+_RM_N_DEG = (_RM_DEG + 1) ** 2  # = 9   (surface + LD polynomial size)
+_RM_N_TOTAL = (_RM_TOTAL_DEG + 1) ** 2  # = 16  (surface + LD + velocity size)
+
+# A2 change-of-basis matrix: converts occultation solution vector sT to polynomial
+# coefficients.  Shape: (_RM_N_TOTAL, _RM_N_TOTAL) = (16, 16).
+_RM_A2 = np.array(
+    _scipy_sparse.linalg.inv(_jxp_A2_inv(_RM_TOTAL_DEG)).todense(),
+    dtype=np.float32,
+)
+
+# A1 basis matrix for the velocity map (vdeg=1).
+# Converts velocity-map spherical harmonics to polynomial coefficients.
+# Shape: (_RM_VDEG_N, _RM_VDEG_N) = (4, 4).
+_RM_A1_VDEG = np.array(_jxp_A1(_RM_VDEG).todense(), dtype=np.float32)
+
+# A1(0) scalar: the degree-0 surface-map basis factor (= 1/π).
+# Used to scale the ydeg=0 surface amplitude into polynomial units.
+_RM_A1_YDEG_SCALAR = float(np.array(_jxp_A1(_RM_YDEG).todense()).flatten()[0])  # ≈ 1/π
+
+# Limb-darkening basis matrix U(2): maps [1, u1, u2] → 9 polynomial coefficients.
+# Shape: (3, 9).
+_RM_U = np.array(_jxp_U(_RM_UDEG), dtype=np.float32)
+
+# Rotation vectors rT(n): used for normalisation (total flux integrals).
+# rT(2) shape (9,) — for limb-darkening normalisation (udeg=2).
+# rT(3) shape (16,) — for out-of-transit design matrix (total_deg=3).
+_RM_rT_UDEG = np.array(_jxp_rT(_RM_UDEG), dtype=np.float32)  # (9,)
+_RM_rT_TOTAL = np.array(_jxp_rT(_RM_TOTAL_DEG), dtype=np.float32)  # (16,)
+
+
+@jit
+def reparametrize_LD(
+    q1: float | jax.Array, q2: float | jax.Array
+) -> tuple[jax.Array, jax.Array]:
+    """Convert (q1, q2) to (u1, u2) limb darkening coefficients. JIT-compiled."""
+    u1 = 2.0 * jnp.sqrt(q1) * q2
+    u2 = jnp.sqrt(q1) * (1 - 2 * q2)
+    return (u1, u2)
+
+
+@jit
+def reparametrize_LD_inverse(
+    u1: float | jax.Array, u2: float | jax.Array
+) -> tuple[float | jax.Array, float | jax.Array]:
+    """Convert (u1, u2) to (q1, q2) limb darkening parameters. JIT-compiled."""
+    q1 = jnp.square(u1 + u2)
+    q2 = 0.5 * u1 / (u1 + u2)
+    return (q1, q2)
+
+
+@jit
+def reparametrize_ecc_omega(
+    ec: float | jax.Array, es: float | jax.Array
+) -> tuple[float | jax.Array, float | jax.Array]:
+    """
+    Convert Cartesian eccentricity to (e, ω).
+
+    ec = e * cos(ω)
+    es = e * sin(ω)
+
+    ec and es are each sampled from a uniform prior on [-1, 1] (the widest
+    range consistent with e <= 1), so points outside the unit disk
+    (ec**2 + es**2 >= 1) do not correspond to any physical (e, ω) and must be
+    excluded upstream by an explicit validity gate -- see
+    ``ecc_physical_<planet>`` in
+    ``PyAETNA.sampling.engine.numpyro_model``, which rejects exactly those
+    points via ``numpyro.factor``.
+
+    Returns: (ecc, omega) where ecc ∈ [0, 1], omega ∈ [-π, π]
+
+    JIT-compiled for performance.
+    """
+    ecc = jnp.sqrt(ec**2 + es**2)
+    omega = jnp.arctan2(es, ec)
+    return (ecc, omega)
+
+
+@jit
+def compute_semimajor_axis_scaled(
+    rho_star: float | jax.Array, period: float | jax.Array
+) -> jax.Array:
+    """
+    Compute scaled semi-major axis (a/R_star) from stellar density and orbital period.
+
+    Uses Kepler's third law with dimensionless stellar parameters (M_star = R_star = 1.0):
+    a/R_star = [(G * ρ_star * P²) / (3π)]^(1/3)
+
+    Parameters:
+    -----------
+    rho_star : float
+        Stellar density [g/cm³]
+    period : float or array
+        Orbital period [days]
+
+    Returns:
+    --------
+    a_over_rstar : float or array
+        Scaled semi-major axis [dimensionless]
+    """
+    # Convert period to seconds
+    P_sec = period * SECONDS_PER_DAY
+
+    # Compute a/R_star using Kepler's third law
+    a_over_rstar = jnp.power((G_CGS * rho_star * P_sec**2) / (3.0 * jnp.pi), 1.0 / 3.0)
+
+    return a_over_rstar
+
+
+@jit
+def compute_inclination(
+    b_param: float | jax.Array,
+    a_over_rstar: float | jax.Array,
+    ecc: float | jax.Array = 0.0,
+    omega: float | jax.Array = 0.0,
+) -> jax.Array:
+    """
+    Compute orbital inclination from transit impact parameter.
+
+    Inverts the impact-parameter definition
+    (Winn 2010, Eq. 7, https://arxiv.org/pdf/1001.2010.pdf):
+
+        b ≡ (a/R*) · cos(i) · (1 − e²) / (1 + e · sin ω)
+
+    to obtain
+
+        i = arccos[ b / ((a/R*) · ρ_c) ]
+
+    where ρ_c = (1 − e²) / (1 + e · sin ω).
+
+    Parameters
+    ----------
+    b_param : float or array
+        Transit impact parameter [dimensionless].
+    a_over_rstar : float or array
+        Scaled semi-major axis a/R* [dimensionless].
+    ecc : float or array, optional
+        Orbital eccentricity [0, 1).
+    omega : float or array, optional
+        Argument of periastron [radians].
+
+    Returns
+    -------
+    incl : float or array
+        Orbital inclination [radians].
+    """
+
+    # NOTE: We use jnp.where to branch between circular (e=0) and
+    # eccentric orbits. This handles both scalar and array inputs.
+    # (Previously used jax.lax.cond which failed for array inputs).
+    incl_circ = jnp.arccos(jnp.clip(b_param / a_over_rstar, 0.0, 1.0 - 1e-10))
+    rho_c = (1.0 - ecc**2) / (1.0 + ecc * jnp.sin(omega))
+    incl_ecc = jnp.arccos(jnp.clip(b_param / (a_over_rstar * rho_c), 0.0, 1.0 - 1e-10))
+
+    return jnp.where(ecc == 0, incl_circ, incl_ecc)
+
+
+@jit
+def compute_transit_duration(
+    period: float | jax.Array,
+    rho_star: float | jax.Array,
+    b_param: float | jax.Array,
+    rp_by_rs: float | jax.Array,
+    ecc: float | jax.Array = 0.0,
+    omega: float | jax.Array = 0.0,
+    T23: bool = False,
+) -> jax.Array:
+    """
+    Compute transit duration (T14 or T23) from orbital parameters and stellar density.
+
+    Implements Eqs. 14–16 of Winn (2010, https://arxiv.org/pdf/1001.2010.pdf):
+
+        T = (P/π) · arcsin[ √(k_eff² − b²) / ((a/R*) · sin i) ]
+            · √(1 − e²) / (1 + e · sin ω)
+
+    where k_eff = (1 + k) for the total (first-to-fourth contact, T14)
+    duration, and k_eff = (1 − k) for the full (second-to-third contact,
+    T23) duration, with k = Rp/R*.
+
+    Parameters
+    ----------
+    period : float or array
+        Orbital period [days].
+    rho_star : float
+        Stellar density [g/cm³].
+    b_param : float or array
+        Transit impact parameter [dimensionless].
+    rp_by_rs : float or array
+        Planet-to-star radius ratio k = Rp/R* [dimensionless].
+    ecc : float or array, optional
+        Orbital eccentricity [0, 1).
+    omega : float or array, optional
+        Argument of periastron [radians].
+    T23 : bool, optional
+        If True, compute T23 (full duration, second-to-third contact)
+        using (1 − k) instead of (1 + k).  Default is False (T14).
+
+    Returns
+    -------
+    duration : float or array
+        Transit duration [days].  T14 by default, T23 if ``T23=True``.
+    """
+    a_over_rstar = compute_semimajor_axis_scaled(rho_star, period)
+    incl = compute_inclination(b_param, a_over_rstar, ecc, omega)
+    k = rp_by_rs
+
+    # T14 → (1+k),  T23 → (1−k)
+    k_eff = jnp.where(T23, 1.0 - k, 1.0 + k)
+
+    transit_occurs = b_param < k_eff
+    b_safe = jnp.where(transit_occurs, b_param, 0.0)
+
+    numerator = jnp.sqrt(jnp.maximum(k_eff**2 - b_safe**2, 0.0))
+    denominator = a_over_rstar * jnp.sin(incl)
+
+    duration = (period / jnp.pi) * jnp.arcsin(numerator / denominator)
+
+    # Eccentricity correction (Eq. 16 in Winn 2010)
+    ecc_factor = jnp.sqrt(1.0 - ecc**2) / (1.0 + ecc * jnp.sin(omega))
+    duration_ecc = duration * ecc_factor
+
+    duration = jnp.where(ecc == 0, duration, duration_ecc)
+
+    return jnp.where(transit_occurs, duration, 0.0)
+
+
+@jit
+def compute_eclipse_duration(
+    period: float | jax.Array,
+    rho_star: float | jax.Array,
+    b_param: float | jax.Array,
+    rp_by_rs: float | jax.Array,
+    ecc: float | jax.Array = 0.0,
+    omega: float | jax.Array = 0.0,
+    T23: bool = False,
+) -> jax.Array:
+    """
+    Compute secondary eclipse duration (total T14 or full T23).
+
+    Same geometry as :func:`compute_transit_duration` but the eccentricity
+    correction uses the eclipse-phase factor:
+
+        √(1 − e²) / (1 − e · sin ω)
+
+    (note the sign flip compared to the transit formula which has 1 + e · sin ω).
+
+    Parameters
+    ----------
+    period : float or array
+        Orbital period [days].
+    rho_star : float
+        Stellar density [g/cm³].
+    b_param : float or array
+        Impact parameter [dimensionless] — approximate; eclipse impact
+        parameter differs slightly for eccentric orbits, but this is
+        standard practice for moderate eccentricities.
+    rp_by_rs : float or array
+        Planet-to-star radius ratio [dimensionless].
+    ecc : float or array, optional
+        Orbital eccentricity [0, 1).
+    omega : float or array, optional
+        Argument of periastron [radians].
+    T23 : bool, optional
+        If True, compute T23 (full duration, second-to-third contact)
+        using (1 − k) instead of (1 + k).  Default is False (T14).
+
+    Returns
+    -------
+    duration : float or array
+        Secondary eclipse duration [days].
+    """
+    a_over_rstar = compute_semimajor_axis_scaled(rho_star, period)
+    incl = compute_inclination(b_param, a_over_rstar, ecc, omega)
+    k = rp_by_rs
+
+    # T14 → (1+k),  T23 → (1−k)
+    k_eff = jnp.where(T23, 1.0 - k, 1.0 + k)
+
+    eclipse_occurs = b_param < k_eff
+    b_safe = jnp.where(eclipse_occurs, b_param, 0.0)
+
+    numerator = jnp.sqrt(jnp.maximum(k_eff**2 - b_safe**2, 0.0))
+    denominator = a_over_rstar * jnp.sin(incl)
+
+    duration = (period / jnp.pi) * jnp.arcsin(numerator / denominator)
+
+    # Eccentricity correction at eclipse phase: sign flip on e·sin(ω)
+    ecc_factor = jnp.sqrt(1.0 - ecc**2) / (1.0 - ecc * jnp.sin(omega))
+    duration_ecc = duration * ecc_factor
+
+    duration = jnp.where(ecc == 0, duration, duration_ecc)
+
+    return jnp.where(eclipse_occurs, duration, 0.0)
+
+
+@jit
+def compute_eclipse_phase(
+    ecc: float | jax.Array, omega: float | jax.Array
+) -> jax.Array:
+    """
+    Compute orbital phase of secondary eclipse for eccentric orbits.
+
+    For circular orbits, the eclipse occurs at phase 0.5 (half period after transit).
+    For eccentric orbits, the eclipse phase shifts due to Kepler's equation.
+
+    The computation uses:
+    - Transit occurs at true anomaly f_transit = π/2 - ω
+    - Eclipse occurs at true anomaly f_eclipse = 3π/2 - ω
+    - Convert both to mean anomaly via Kepler's equation
+    - Phase difference = (M_eclipse - M_transit) / (2π)
+
+    Parameters:
+    -----------
+    ecc : float or array
+        Orbital eccentricity [0, 1)
+    omega : float or array
+        Argument of periastron [radians]
+
+    Returns:
+    --------
+    eclipse_phase : float or array
+        Orbital phase at eclipse [0, 1], where phase 0 = transit
+        For e=0: returns 0.5
+        For e>0: returns value shifted from 0.5
+    """
+
+    # NOTE: We use jnp.where to branch between circular (e=0) and
+    # eccentric orbits. This handles both scalar and array inputs.
+    # (Previously used jax.lax.cond which failed for array inputs).
+    phase_circ = jnp.array(0.5)
+
+    # True anomaly at transit and eclipse
+    f_transit = jnp.pi / 2.0 - omega
+    f_eclipse = 3.0 * jnp.pi / 2.0 - omega
+
+    # Convert true anomaly to eccentric anomaly using:
+    # tan(E/2) = sqrt((1-e)/(1+e)) * tan(f/2)
+    # Rearranged: E = 2 * arctan(sqrt((1-e)/(1+e)) * tan(f/2))
+    sqrt_ecc_factor = jnp.sqrt((1.0 - ecc) / (1.0 + ecc))
+
+    E_transit = 2.0 * jnp.arctan(sqrt_ecc_factor * jnp.tan(f_transit / 2.0))
+    E_eclipse = 2.0 * jnp.arctan(sqrt_ecc_factor * jnp.tan(f_eclipse / 2.0))
+
+    # Convert eccentric anomaly to mean anomaly: M = E - e*sin(E)
+    M_transit = E_transit - ecc * jnp.sin(E_transit)
+    M_eclipse = E_eclipse - ecc * jnp.sin(E_eclipse)
+
+    # Phase difference (handle wrapping: eclipse is ~π ahead of transit)
+    # For small e, this is close to π, so phase ≈ 0.5
+    phase_ecc = ((M_eclipse - M_transit) / (2.0 * jnp.pi)) % 1.0
+
+    return jnp.where(ecc == 0, phase_circ, phase_ecc)
+
+
+@jit
+def compute_phase_angle_fast(x: jax.Array, y: jax.Array, z: jax.Array) -> jax.Array:
+    """
+    Calculate phase angle for jaxoplanet coordinate system.
+
+    In jaxoplanet's coordinate system:
+    - +Z points towards the observer
+    - X is the north axis
+    - Y is the east axis
+
+    Phase angle: angle at the planet between star and observer directions.
+
+    Phase angle convention:
+    - α = 0: Eclipse (planet behind star, z > 0, fully illuminated)
+    - α = π/2: Quadrature (planet at side)
+    - α = π: Transit (planet in front of star, z < 0, dark side toward observer)
+
+    NOTE: This returns UNSIGNED phase angle [0, π]. For phase curves with
+    hotspot offsets, use compute_orbital_phase_signed() instead.
+
+    Parameters:
+    -----------
+    x, y, z : array, shape [n_planets, n_timestamps]
+        Planet coordinates relative to the star from system.relative_position()
+
+    Returns:
+    --------
+    phase_angle : array, shape [n_planets, n_timestamps]
+        Phase angle in radians [0, π]
+    """
+    # Vector from planet to star: -r_planet = (-x, -y, -z)
+    # Vector from planet to observer: observer is at +Z direction = (0, 0, 1)
+
+    # Dot product between (-x, -y, -z) and (0, 0, 1):
+    # dot = (-x)*(0) + (-y)*(0) + (-z)*(1) = -z
+
+    # Magnitude of planet position vector
+    r_norm = jnp.sqrt(x * x + y * y + z * z)
+
+    # cos(phase_angle) = -z / r_norm
+    # When z > 0 (planet away from observer): cos > 0, phase < π/2 (eclipse, α→0)
+    # When z < 0 (planet toward observer): cos < 0, phase > π/2 (transit, α→π)
+
+    cos_phase = jnp.clip(-z / r_norm, -1.0, 1.0)
+
+    return jnp.arccos(cos_phase)
+
+
+@jit
+def compute_orbital_phase_signed(x: jax.Array, y: jax.Array, z: jax.Array) -> jax.Array:
+    """
+    Calculate SIGNED orbital phase for phase curve modeling with hotspot offsets.
+
+    In jaxoplanet's coordinate system (same as exoplanet package):
+    - +Z points towards the observer
+    - z > 0: Planet is in FRONT of star (primary transit)
+    - z < 0: Planet is BEHIND star (secondary eclipse)
+    - X is in the orbital plane (north/south direction of orbital motion)
+    - Y is perpendicular to the orbital plane (impact parameter direction)
+
+    The orbit is primarily in the X-Z plane (Y is nearly constant = impact parameter).
+
+    IMPORTANT: For phase curve modeling:
+    - At primary transit (z > 0): We see the planet's NIGHTSIDE (facing away from star)
+    - At secondary eclipse (z < 0): We see the planet's DAYSIDE (facing the star)
+
+    Returns a SIGNED phase that distinguishes pre-eclipse from post-eclipse:
+    - φ = 0: Secondary eclipse (planet behind star, z < 0 min, x = 0) - DAYSIDE visible
+    - φ = +π/2: First quadrature (x > 0, z = 0)
+    - φ = ±π: Primary transit (planet in front of star, z > 0 max, x = 0) - NIGHTSIDE visible
+    - φ = -π/2: Second quadrature (x < 0, z = 0)
+
+    This is essential for modeling asymmetric phase curves due to hotspot offsets.
+
+    Parameters:
+    -----------
+    x, y, z : array, shape [n_planets, n_timestamps]
+        Planet coordinates relative to the star from system.relative_position()
+
+    Returns:
+    --------
+    orbital_phase : array, shape [n_planets, n_timestamps]
+        Signed orbital phase in radians [-π, π], 0 at secondary eclipse
+    """
+    # Use atan2 to get signed angle. The orbit is in the x-z plane.
+    # With jaxoplanet convention (z > 0 = transit, z < 0 = eclipse):
+    #
+    # We want phase = 0 at secondary eclipse (z minimum < 0).
+    # Use atan2(-x, -z) so that:
+    # - At secondary eclipse: z < 0 (min), x = 0 → atan2(0, positive) = 0
+    # - Approaching eclipse (x < 0): phase > 0 (positive)
+    # - Receding from eclipse (x > 0): phase < 0 (negative)
+    #
+    # This convention means:
+    # - phase_offset > 0: hotspot shifted AHEAD in orbital motion (eastward, evening)
+    #   → We see it BEFORE reaching secondary eclipse
+    #   → But it's brightest AFTER eclipse (when moving away, phase < 0)
+    # - phase_offset < 0: hotspot shifted BEHIND in orbital motion (westward, morning)
+    #   → We see it AFTER passing secondary eclipse
+    #   → But it's brightest BEFORE eclipse (when approaching, phase > 0)
+
+    return -jnp.arctan2(x, -z)
+
+
+@jit
+def getLambertPhaseCurve(alpha: jax.Array) -> jax.Array:
+    """
+    Exact phase function for a reflective Lambertian sphere.
+
+    Formula: (sin(α) + (π - α)cos(α)) / π
+
+    Valid for α ∈ [0, π]. JIT-compiled.
+
+    Parameters:
+    -----------
+    alpha : array
+        Phase angle in radians [0, π]
+
+    Returns:
+    --------
+    phase_modulation : array
+        Modulation factor [0, 1], where 1 = fully illuminated, 0 = fully dark
+    """
+    return (jnp.sin(alpha) + (jnp.pi - alpha) * jnp.cos(alpha)) / jnp.pi
+
+
+@jit
+def getIsotropicPhaseCurve(alpha: jax.Array) -> jax.Array:
+    """
+    Phase function for isotropic scattering.
+
+    Formula: (1 + cos(α)) / 2
+
+    Simpler than Lambert, peaks at α=0 (full illumination). JIT-compiled.
+
+    Parameters:
+    -----------
+    alpha : array
+        Phase angle in radians [0, π]
+
+    Returns:
+    --------
+    phase_modulation : array
+        Modulation factor [0, 1]
+    """
+    return 0.5 * (1 + jnp.cos(alpha))
+
+
+def _compute_ltt_positions(
+    system: Any, times: jax.Array, z_obs: jax.Array
+) -> tuple[jax.Array, jax.Array, jax.Array]:
+    """Re-evaluate planet positions at light-travel-time corrected epochs.
+
+    Photons observed at *t_obs* were emitted at  t_em = t_obs + z/c  where
+    z is the planet's line-of-sight displacement from the star (+z towards
+    observer, in R_sun) and c is in R_sun/day.
+
+    A single iteration is sufficient: the residual
+    |delta_t_2 / delta_t_1| ~ v_z * a / c**2  <  1e-4
+    for all known exoplanet systems.
+
+    Parameters
+    ----------
+    system : jaxoplanet.orbits.keplerian.System
+        Built orbital system.
+    times : array (n_times,)
+        Observed timestamps [days].
+    z_obs : array (n_planets, n_times)
+        Planet z-coordinates at *observed* times [R_sun].
+
+    Returns
+    -------
+    x, y, z : arrays (n_planets, n_times)
+        Positions evaluated at emission times [R_sun].
+    """
+    # Light-travel-time offset  [days]  (+z closer → emission later)
+    ltt = z_obs / _C_RSUN_PER_DAY
+
+    # Per-planet corrected positions  (loop unrolled at JIT trace time)
+    corrected = [
+        body.relative_position(times + ltt[i]) for i, body in enumerate(system.bodies)
+    ]
+    return (
+        jnp.stack([pos[0] for pos in corrected]),
+        jnp.stack([pos[1] for pos in corrected]),
+        jnp.stack([pos[2] for pos in corrected]),
+    )
+
+
+@jit
+def compute_radial_velocity(
+    times: jax.Array,
+    t0s: jax.Array,
+    periods: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    K_rv: jax.Array,
+) -> tuple[jax.Array, jax.Array]:
+    """
+    Compute Keplerian radial velocity model for multiple planets.
+
+    Builds a jaxoplanet orbital system and returns the stellar reflex RV
+    time series.  The returned values represent the *stellar* radial velocity
+    (reflex motion) — positive values mean the star moves **away** from the
+    observer.
+
+    When the observed (inclination-projected) semi-amplitude K is provided,
+    jaxoplanet does NOT need stellar mass/density, radius or inclination —
+    only the orbital ephemeris (period, τ) and eccentricity.
+
+    Parameters
+    ----------
+    times : array, shape (n_times,)
+        Observation times [days], centred on ``timeCenter``
+    t0s : array, shape (n_planets,)
+        Mid-transit reference epochs [days]
+    periods : array, shape (n_planets,)
+        Orbital periods [days]
+    ecosomega, esinomega : arrays, shape (n_planets,)
+        Cartesian eccentricity parameterisation
+    K_rv : array, shape (n_planets,)
+        Observed (projected) radial-velocity semi-amplitudes [m/s].
+        Since this is the *observed* (inclination-projected) K, neither
+        stellar mass nor orbital inclination are required.
+
+    Returns
+    -------
+    rv_total : array, shape (n_times,)
+        Total stellar radial velocity [m/s] (sum over all planets, **not**
+        including systemic velocity γ).
+    rv_per_planet : array, shape (n_planets, n_times)
+        Per-planet radial-velocity contribution [m/s].
+    """
+    eccentricities, omegas = reparametrize_ecc_omega(ecosomega, esinomega)
+    n_planets = len(t0s)
+
+    # Empty Central — mass/radius not needed when K is provided directly
+    star = Central()
+
+    system = System(star)
+    for i in range(n_planets):
+        system = system.add_body(
+            period=periods[i],
+            time_transit=t0s[i],
+            eccentricity=eccentricities[i],
+            omega_peri=omegas[i],
+            radial_velocity_semiamplitude=K_rv[i],
+        )
+
+    # jaxoplanet returns shape (n_planets, n_times)
+    rv_per_planet = system.radial_velocity(times)
+
+    # Sum over planets
+    rv_total = jnp.sum(rv_per_planet, axis=0)
+
+    return rv_total, rv_per_planet
+
+
+def _rm_disc_terms(
+    times: jax.Array,
+    rp_by_rs: jax.Array,
+    b_param: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    period: jax.Array,
+    rho_star: jax.Array,
+    v_sin_i: jax.Array,
+    lambda_spin: jax.Array,
+    inc_star: jax.Array,
+    q1: jax.Array,
+    q2: jax.Array,
+) -> tuple[jax.Array, jax.Array, jax.Array]:
+    """Disc integrals of the RM observables, with full-precision matrix products.
+
+    On GPUs, JAX's default precision for float32 matrix products is TF32,
+    which keeps only 10 bits of mantissa.  The occultation integrals below are
+    differences of matrix products, and with TF32 the occulted flux fraction
+    was wrong by up to 3e-4 (3 % of a 1 % transit) and the CCF anomaly by up
+    to 2 m/s near ingress and egress, while the CPU was exact.  The products
+    are few and small, so ``highest`` precision costs nothing measurable.
+    See :func:`_rm_disc_terms_tf32_unsafe` for the physics.
+    """
+    with jax.default_matmul_precision("highest"):
+        return _rm_disc_terms_tf32_unsafe(
+            times, rp_by_rs, b_param, ecosomega, esinomega, period, rho_star,
+            v_sin_i, lambda_spin, inc_star, q1, q2,
+        )  # fmt: skip
+
+
+def _rm_disc_terms_tf32_unsafe(
+    times: jax.Array,
+    rp_by_rs: jax.Array,
+    b_param: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    period: jax.Array,
+    rho_star: jax.Array,
+    v_sin_i: jax.Array,
+    lambda_spin: jax.Array,
+    inc_star: jax.Array,
+    q1: jax.Array,
+    q2: jax.Array,
+) -> tuple[jax.Array, jax.Array, jax.Array]:
+    """Disc integrals behind both RM observables (flux-weighted and CCF).
+
+    Returns ``(rv_fw, visible, r_sky)`` per time step: the flux-weighted RV
+    anomaly of the visible disc [m/s] (the classic RM model, below), the visible
+    fraction of the stellar flux (1 - f, with f the occulted fraction), and the
+    sky-plane distance of the planet centre from the disc centre [R_star].
+    All args must be JAX arrays.
+    """
+    """
+    Compute the Rossiter-McLaughlin (RM) radial-velocity anomaly for a single planet.
+
+    The RM effect is the apparent RV shift that arises during a planetary transit
+    because the planet's disk occults different Doppler-shifted strips of the rotating
+    stellar surface.  The anomaly is zero outside transit and follows the rotation
+    velocity profile of the occulted chord.
+
+    This implementation reproduces the physics of
+    ``jaxoplanet.starry.doppler.surface_radial_velocity``
+    (Luger et al. 2019, AJ 157 64) but avoids that function's use of BCOO sparse
+    matrix multiplications inside ``jax.vmap``.  Those sparse ops fail in the
+    backward pass (VJP) required by NumPyro's NUTS sampler: the batch rule for
+    ``bcoo_dot_general`` does not support a batched LHS, causing the error::
+
+        ValueError: Input primal JAX type float32[1,9].
+        Expected cotangent float32[1,9] but got float32[N,9].
+
+    Our fix pre-computes the change-of-basis matrices (A2, A1_vdeg) as dense
+    float32 constants at module import time, replaces the problematic
+    ``p_yu.tosparse() @`` with ``jnp.dot(p_yu_dense, dm[:n])`` (dense dot), and
+    uses ``Pijk.todense()`` (scatter-to-dense, vmap-compatible) instead of
+    ``Pijk.tosparse() @`` (matrix-multiply-with-BCOO, NOT vmap-compatible).
+    The result is functionally identical but compiles quickly (≈70 s for 117
+    data points) and produces correct gradients through ``jax.vmap``.
+
+    Unit system
+    -----------
+    jaxoplanet uses R_star = 1 R_sun.  All orbital positions are in solar radii.
+    The function returns values in R_sun/day, which are converted to m/s via
+    ``_RSUN_DAY_TO_MS`` (≈ 8052.08 m/s per R_sun/day).
+
+    Outside transit the function automatically returns 0 via the ``b_rot`` flag
+    (planet behind star → z≤0, or outside disk → b≥1+r), so no explicit mask is
+    needed.
+
+    Multi-planet RM
+    ---------------
+    Call once per planet and sum the results.  Simultaneous mutual transits are
+    not modelled (additive approximation).
+
+    Parameters
+    ----------
+    times : array_like, shape (n_times,)
+        Observation times [days], **folded so that mid-transit is at t = 0**.
+    rp_by_rs : float
+        Planet-to-star radius ratio (dimensionless).
+    b_param : float
+        Transit impact parameter b = (a/R*)·cos(i_orb) (dimensionless).
+    ecosomega : float
+        Cartesian eccentricity component e·cos(ω).
+    esinomega : float
+        Cartesian eccentricity component e·sin(ω).
+    period : float
+        Orbital period [days].
+    rho_star : float
+        Mean stellar density [g/cm³].
+    v_sin_i : float
+        Sky-projected stellar rotation velocity [m/s].
+    lambda_spin : float
+        Sky-projected spin-orbit angle λ [rad], in the convention of Ohta et al.
+        (2005): the angle from the sky projection of the stellar spin axis to
+        that of the orbit normal, as in their Eqs. (9)-(10), with the impact
+        parameter b >= 0.  λ=0 → aligned; |λ|=π → retrograde; for λ > 0 the
+        planet covers more of the receding (redshifted) hemisphere.
+    inc_star : float
+        Stellar inclination [rad].  π/2 → stellar pole in plane of sky.
+    q1 : float, optional
+        Quadratic limb-darkening parameter q1 (Kipping 2013).  Default: 0.5.
+    q2 : float, optional
+        Quadratic limb-darkening parameter q2 (Kipping 2013).  Default: 0.1.
+
+    Returns
+    -------
+    rv_rm : array, shape (n_times,)
+        Rossiter-McLaughlin radial-velocity anomaly [m/s].  Positive values mean
+        the star appears to recede (planet occuiting the blue-shifted hemisphere).
+
+    Notes
+    -----
+    The stellar rotation period is derived as::
+
+        v_eq  = v_sin_i / (sin(inc_star) * _RSUN_DAY_TO_MS)    [R_sun/day]
+        P_rot = 2π / v_eq                                        [days]
+
+    Polynomial degrees: ydeg=0, udeg=2, vdeg=1 (solid-body rotation, no alpha).
+    The module-level ``_RM_*`` constants encode the pre-computed basis matrices.
+
+    References
+    ----------
+    * Luger et al. (2019), AJ 157 64  — starry polynomial framework
+    * Ohta et al. (2005), ApJ 622 1118 — analytic RM formula
+    * Kipping (2013), MNRAS 435 2152  — q1/q2 LD parameterisation
+    """
+    # ------------------------------------------------------------------
+    # 1. Derive stellar rotation period from the observable v_sin_i
+    # ------------------------------------------------------------------
+    # jaxoplanet's Surface uses rotation period [days] as its argument.
+    # The relationship is:
+    #   v_eq [R_sun/day]  = v_sin_i [m/s] / (_RSUN_DAY_TO_MS × sin(inc_star))
+    #   P_rot [days]      = 2π / v_eq [R_sun/day]
+    # (since jaxoplanet sets R_star = 1 R_sun, so v_eq = 2π R_star / P_rot
+    #  is in R_sun/day and 1 R_sun/day = _RSUN_DAY_TO_MS m/s)
+    v_eq_rsun_per_day = v_sin_i / (_RSUN_DAY_TO_MS * jnp.sin(inc_star))
+    p_rot_star = 2.0 * jnp.pi / v_eq_rsun_per_day  # [days]
+
+    # ------------------------------------------------------------------
+    # 2. Compute polynomial representations of the stellar surface
+    #    (all depend on surface params only, NOT on planet position)
+    # ------------------------------------------------------------------
+    # These are computed ONCE per model call and closed over in the per-step
+    # vmap lambda below.  They do NOT use Pijk.tosparse() because that creates
+    # a BCOO sparse matrix whose VJP (backward pass) fails inside jax.vmap.
+    # Instead we use:
+    #   - Pijk.todense()  for the polynomial→dense-array conversion (works in vmap)
+    #   - jnp.dot()       for all dot products (standard dense, fully vmappable)
+    # The module-level constants (_RM_U, _RM_A1_VDEG, …) are pre-computed scipy
+    # matrices stored as float32 numpy arrays; converting them to JAX arrays here
+    # makes them compile-time constants.
+
+    # (a) Limb-darkening (LD) polynomial coefficients.
+    # u_vec = [1, u1, u2]; p_u_coeffs = u_vec @ U(2)  →  shape (9,)
+    u1, u2 = reparametrize_LD(q1, q2)
+    u_vec = jnp.array([1.0, u1, u2])
+    _U_jax = jnp.array(_RM_U)  # (3, 9) constant
+    p_u_coeffs = u_vec @ _U_jax  # (9,)  depends on u1, u2 (traced)
+
+    # (b) Surface map polynomial (ydeg=0 → uniform disk = constant 1/π scaled).
+    # A1(0) = [[1/π]]; p_y_coeff = [1/π].
+    p_y_coeff = jnp.array([_RM_A1_YDEG_SCALAR])  # [1/π], shape (1,)
+
+    # (c) Combined LD × surface polynomial (p_yu = p_y * p_u).
+    # Since p_y = [1/π] (degree-0 scalar), multiplication by p_u just scales:
+    #   p_yu_coeffs = (1/π) * p_u_coeffs  →  shape (9,)
+    p_yu_coeffs = p_y_coeff[0] * p_u_coeffs  # (9,)
+
+    # (d) Normalization scalar: norm = π / (p_u @ rT_2).
+    # This is the reciprocal of the total flux integral over the stellar disk.
+    _rT_udeg_jax = jnp.array(_RM_rT_UDEG)  # (9,) constant
+    norm = jnp.pi / jnp.dot(p_u_coeffs, _rT_udeg_jax)
+
+    # (e) RV map expansion coefficients (solid-body rotation, no differential rotation).
+    # y_rv = rv_map_expansion(inc, obl, veq, alpha=None)
+    # Returns 4 coefficients: [0, veq*√3*B/3, 0, veq*√3*A/3] * π
+    # where A = sin(i)*cos(λ), B = sin(i)*sin(λ).
+    # We use veq = v_eq_rsun_per_day (derived in step 1 above).
+    # lambda follows Ohta et al. (2005, Eqs. 9-10), the convention of most codes
+    # (e.g. PyAstronomy's RmcL, PyORBIT). starry's map obliquity is oriented the
+    # other way, so it is -lambda (tests/unit/core/test_rm_convention.py).
+    y_rv_jax = jnp.array(
+        _jxp_rv_map_expansion(
+            inc=inc_star, obl=-lambda_spin, veq=v_eq_rsun_per_day, alpha=None
+        ),
+        dtype=jnp.float32,
+    )  # shape (4,)
+
+    # (f) Pre-cast the basis matrices to JAX arrays for use inside the vmap.
+    _A2_jax = jnp.array(_RM_A2)  # (16, 16) constant
+    _A1_vdeg_jax = jnp.array(_RM_A1_VDEG)  # (4, 4) constant
+    _rT_total_jax = jnp.array(_RM_rT_TOTAL)  # (16,) constant
+
+    # (g) Pijk objects for the polynomial product inside the vmap.
+    # These carry the traced LD and surface map coefficients.
+    p_u_pk = _jxp_Pijk.from_dense(p_u_coeffs, degree=_RM_UDEG)  # degree-2 poly
+    p_y_pk = _jxp_Pijk.from_dense(p_y_coeff, degree=_RM_YDEG)  # degree-0 poly
+
+    # ------------------------------------------------------------------
+    # 3. Build Keplerian orbit to obtain planet (x, y, z) positions
+    # ------------------------------------------------------------------
+    # We use the same Central setup as in compute_lightcurve:
+    #   R_star = 1 R_sun  →  positions are in R_sun = R_star units.
+    #   M_star [M_sun] = rho_star [g/cm³] / rho_sun [g/cm³]
+    # (Kepler's 3rd law then gives the correct a/R_star from the period.)
+    eccentricity, omega = reparametrize_ecc_omega(ecosomega, esinomega)
+    star_mass = rho_star / RHO_SUN_CGS  # [M_sun], with R_star = 1 R_sun
+    star = Central(mass=star_mass, radius=1.0)
+
+    # Clip impact parameter to avoid degenerate geometry inside jaxoplanet.
+    b_safe = jnp.clip(b_param, 0.0, 0.99)
+
+    # Build a single-planet system.  We only need orbital positions here,
+    # not the Keplerian RV amplitude.  times are folded to t0 = 0.
+    system = System(star).add_body(
+        period=period,
+        time_transit=0.0,  # folded: mid-transit at t = 0
+        radius=rp_by_rs,
+        impact_param=b_safe,
+        eccentricity=eccentricity,
+        omega_peri=omega,
+    )
+
+    # relative_position returns (x, y, z) each shape (n_planets, n_times).
+    # Coordinate convention (jaxoplanet):
+    #   +x: transverse direction (parallel to stellar equator for aligned orbit)
+    #   +y: along the impact parameter direction
+    #   +z: toward the observer (> 0 during transit)
+    x_all, y_all, z_all = system.relative_position(times)
+    x_p = x_all[0]  # (n_times,)
+    y_p = y_all[0]
+    z_p = z_all[0]
+
+    # ------------------------------------------------------------------
+    # 4. Evaluate RM anomaly for each time step via jax.vmap
+    # ------------------------------------------------------------------
+    # WHY vmap WORKS HERE (despite jaxoplanet's sparse-op issues):
+    # The original jaxoplanet surface_radial_velocity computes:
+    #   intensity = p_yu.tosparse() @ design_matrix[:n_deg]
+    # This FAILS inside vmap during the backward pass (VJP) because the
+    # bcoo_dot_general VJP rule does not support batched LHS (BCOO matrices
+    # can't have a batch dimension in the backward rule of JAX ≤0.4.x).
+    # The error manifests as:
+    #   "Input primal float32[1,9]: expected cotangent float32[1,9] got float32[N,9]"
+    #
+    # FIX: We replace the two problematic operations:
+    #   OLD: intensity = (p_yu.tosparse() @ dm[:9]) * norm  # BCOO @ dense → FAILS in vmap
+    #   NEW: intensity = jnp.dot(p_yu_coeffs, dm[:9]) * norm  # dense @ dense → OK
+    # AND:
+    #   OLD: design_matrix = sT @ A2_bcoo   # BCOO matmul → may also FAIL in vmap
+    #   NEW: design_matrix = sT @ _A2_jax   # dense matmul → OK
+    #
+    # The Pijk polynomial product (p_y * p_rv * p_u).todense() also uses
+    # BCOO internally, but ONLY for BCOO.todense() (scatter-to-dense), which
+    # DOES support vmap (it's a scatter with static indices and dynamic values).
+
+    def _rm_per_step(
+        xi: jax.Array, yi: jax.Array, zi: jax.Array
+    ) -> tuple[jax.Array, jax.Array, jax.Array]:
+        """RM RV [R_sun/day], visible flux fraction and sky distance for one time step.
+
+        All ops are dense.
+        """
+        # Angular position of the planet on the stellar disk (azimuth in sky plane).
+        # This determines which part of the rotating stellar surface is occulted.
+        theta_z = jnp.arctan2(xi, yi)
+
+        # Sky-plane separation from the stellar disk centre.
+        b_dist = jnp.sqrt(xi**2 + yi**2)
+
+        # Out-of-transit flag: planet behind star (z≤0) or outside disk (b≥1+r).
+        # When True the RM signal is zero (we use safe dummy values to avoid NaN
+        # inside solution_vector / the jnp.where below).
+        b_rot = jnp.logical_or(b_dist >= 1.0 + rp_by_rs, zi <= 0.0)
+
+        # Clamp to safe non-degenerate values for the occultation solver.
+        b_safe_step = jnp.where(b_rot, jnp.array(1.0), b_dist)
+        r_safe_step = jnp.where(b_rot, jnp.array(1.0), rp_by_rs)
+
+        # Occultation solution vector sT (shape (16,)).
+        # solution_vector evaluates the integral over the occulted crescent.
+        sT = _jxp_solution_vector(_RM_TOTAL_DEG)(b_safe_step, r_safe_step)
+
+        # Design matrix: converts sT to polynomial-basis coefficients.
+        # Uses the pre-computed dense A2 matrix to avoid BCOO inside vmap.
+        # Out-of-transit: use rT_total (full-disk rotation vector).
+        design_matrix = jnp.where(b_rot, _rT_total_jax, sT @ _A2_jax)  # (16,)
+
+        # Rotate the RV map by the planet's azimuthal angle theta_z.
+        # dot_rz(vdeg=1, angle) returns a function that rotates a degree-1
+        # polynomial vector by the given angle.
+        rotated_rv_y = _jxp_dot_rz(_RM_VDEG, -theta_z)(y_rv_jax)  # (4,)
+
+        # Apply the A1 basis transform to get the velocity polynomial in the
+        # Pijk polynomial basis (degree 1, 4 coefficients).
+        p_rv_coeffs = _A1_vdeg_jax @ rotated_rv_y  # (4,)
+        p_rv_pk = _jxp_Pijk.from_dense(p_rv_coeffs, degree=_RM_VDEG)  # degree-1 Pijk
+
+        # Combined polynomial product (p_y × p_rv × p_u) in the Pijk basis.
+        # This is the velocity-weighted LD-attenuated surface brightness.
+        # .todense() calls .tosparse() then BCOO.todense() — the BCOO.todense()
+        # (scatter-to-dense) DOES work inside vmap, unlike .tosparse() @.
+        rv_coeffs = (p_y_pk * p_rv_pk * p_u_pk).todense()  # (16,)
+
+        # RV numerator: dot product of the combined polynomial with the design matrix.
+        rv_val = jnp.dot(rv_coeffs, design_matrix)  # scalar
+
+        # Intensity denominator: dot product of the LD+surface polynomial with the
+        # first n_deg elements of the design matrix.  Uses the pre-computed dense
+        # p_yu_coeffs (= p_y_coeff * p_u_coeffs) to avoid tosparse() @.
+        intensity = jnp.dot(p_yu_coeffs, design_matrix[:_RM_N_DEG]) * norm  # scalar
+
+        # RM RV in R_sun/day (surface amplitude = 1.0 for a uniform disk).
+        # Matches the return value of surface_radial_velocity(surface, x,y,z,r).
+        # ``intensity`` is the visible fraction of the stellar flux (1 out of
+        # transit); the CCF observable also needs it and the planet's distance.
+        return rv_val * norm / intensity, intensity, b_dist
+
+    # Vectorise over all time steps in parallel (vmap is safe because
+    # all matrix operations inside _rm_per_step are dense).
+    rv_rm_rsun_per_day, visible, r_sky = jax.vmap(_rm_per_step)(x_p, y_p, z_p)
+
+    # ------------------------------------------------------------------
+    # 5. Convert R_sun/day → m/s and return
+    # ------------------------------------------------------------------
+    # 1 R_sun/day = _RSUN_DAY_TO_MS m/s ≈ 8052.08 m/s
+    return rv_rm_rsun_per_day * _RSUN_DAY_TO_MS, visible, r_sky
+
+
+@jax.jit
+def _compute_rossiter_mclaughlin_jit(
+    times: jax.Array,
+    rp_by_rs: jax.Array,
+    b_param: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    period: jax.Array,
+    rho_star: jax.Array,
+    v_sin_i: jax.Array,
+    lambda_spin: jax.Array,
+    inc_star: jax.Array,
+    q1: jax.Array,
+    q2: jax.Array,
+) -> jax.Array:
+    """JIT-compiled flux-weighted RM anomaly [m/s] (see :func:`_rm_disc_terms`)."""
+    return _rm_disc_terms(
+        times, rp_by_rs, b_param, ecosomega, esinomega, period, rho_star,
+        v_sin_i, lambda_spin, inc_star, q1, q2,
+    )[0]  # fmt: skip
+
+
+def compute_rossiter_mclaughlin(
+    times: Any,
+    rp_by_rs: Any,
+    b_param: Any,
+    ecosomega: Any,
+    esinomega: Any,
+    period: Any,
+    rho_star: Any,
+    v_sin_i: Any,
+    lambda_spin: Any,
+    inc_star: Any,
+    q1: Any = 0.5,
+    q2: Any = 0.1,
+) -> jax.Array:
+    """Public wrapper: converts scalar inputs to JAX float32 arrays, then calls the JIT core."""
+
+    def _f(x: Any) -> jax.Array:
+        return jnp.asarray(x, dtype=jnp.float32)
+
+    return _compute_rossiter_mclaughlin_jit(
+        jnp.asarray(times, dtype=jnp.float32),
+        _f(rp_by_rs),
+        _f(b_param),
+        _f(ecosomega),
+        _f(esinomega),
+        _f(period),
+        _f(rho_star),
+        _f(v_sin_i),
+        _f(lambda_spin),
+        _f(inc_star),
+        _f(q1),
+        _f(q2),
+    )
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Rossiter-McLaughlin effect as measured by cross-correlation (CCF observable)
+# ──────────────────────────────────────────────────────────────────────────────
+# The flux-weighted anomaly above is the first moment of the distorted line.
+# RVs measured by cross-correlating the spectrum with a template (or by fitting
+# the CCF with a Gaussian) respond differently when v sin i is comparable to the
+# line width (Hirano et al. 2010).  Hirano et al. (2011, ApJ 742, 69, Eq. 16)
+# give the anomaly measured by maximising the cross-correlation with a template
+# (the iodine-cell and template-matching techniques; Boue et al. 2013), for a Voigt
+# intrinsic line (Gaussian width beta, Lorentzian width gamma), the exact
+# rotational kernel with quadratic limb darkening, and radial-tangential
+# macroturbulence (zeta):
+#
+#   dv = -(f / 2 pi) N / D
+#   N  = int E(s) M(s) T(s) sin(2 pi s v_p) s ds
+#   D  = int E(s) M(s) [M(s) - f T(s) cos(2 pi s v_p)] s^2 ds
+#   E  = exp(-2 pi^2 beta^2 s^2 - 4 pi gamma s)
+#
+# where s is the Fourier frequency [s/km], f the occulted flux fraction, v_p
+# the intensity-weighted rotational velocity of the occulted region, M(s) the
+# Fourier transform of the disc-integrated rotation+macroturbulence kernel
+# (their Eq. 17) and T(s) that of the macroturbulence at the planet (Eq. 18).
+# Hirano's widths follow exp(-v^2/beta^2): beta is sqrt(2) times a standard
+# deviation.  f and v_p come from _rm_disc_terms, so the geometry, the limb
+# darkening and the lambda convention are those of the flux-weighted model.
+
+_HIRANO_N_FREQ = 1025  # Simpson's rule: odd number of points
+_HIRANO_N_MU = 128  # Gauss-Legendre nodes for the disc integral of M(s)
+_HIRANO_S_MAX_BETA = 1.2  # s_max = 1.2 / beta: exp(-2 pi^2 beta^2 s^2) < 1e-12
+_mu_nodes, _mu_weights = np.polynomial.legendre.leggauss(_HIRANO_N_MU)
+_HIRANO_MU = 0.5 * (_mu_nodes + 1.0)  # on [0, 1]
+_HIRANO_MU_W = 0.5 * _mu_weights
+_simpson = np.ones(_HIRANO_N_FREQ)
+_simpson[1:-1:2], _simpson[2:-1:2] = 4.0, 2.0
+_HIRANO_SIMPSON = _simpson / 3.0
+
+
+def _bessel_j0(x: jax.Array) -> jax.Array:
+    """Bessel function J0, Abramowitz & Stegun 9.4.1 and 9.4.3 (|error| < 1e-7)."""
+    ax = jnp.abs(x)
+    y = (ax / 3.0) ** 2
+    small = 1.0 + y * (
+        -2.2499997
+        + y * (1.2656208 + y * (-0.3163866 + y * (0.0444479 + y * (-0.0039444 + y * 0.0002100))))
+    )  # fmt: skip
+    big = jnp.maximum(ax, 3.0)
+    z = 3.0 / big
+    f0 = 0.79788456 + z * (
+        -0.00000077
+        + z * (-0.00552740 + z * (-0.00009512 + z * (0.00137237 + z * (-0.00072805 + z * 0.00014476))))
+    )  # fmt: skip
+    theta0 = big - 0.78539816 + z * (
+        -0.04166397
+        + z * (-0.00003954 + z * (0.00262573 + z * (-0.00054125 + z * (-0.00029333 + z * 0.00013558))))
+    )  # fmt: skip
+    return jnp.where(ax <= 3.0, small, f0 * jnp.cos(theta0) / jnp.sqrt(big))
+
+
+def _hirano2011_anomaly(
+    f: jax.Array,
+    v_p: jax.Array,
+    cos_theta: jax.Array,
+    v_sin_i: Any,
+    beta: Any,
+    gamma: Any,
+    zeta: Any,
+    u1: Any,
+    u2: Any,
+) -> jax.Array:
+    """Hirano et al. (2011) Eq. (16): CCF RM anomaly [km/s].
+
+    ``f``, ``v_p`` [km/s] and ``cos_theta`` (of the occulted region) have shape
+    (n_times,); the other arguments are scalars, velocities in km/s.
+    """
+    s_max = _HIRANO_S_MAX_BETA / jnp.maximum(beta, 0.05)
+    s = jnp.linspace(0.0, 1.0, _HIRANO_N_FREQ) * s_max  # (N,)
+    ds_w = jnp.asarray(_HIRANO_SIMPSON) * s_max / (_HIRANO_N_FREQ - 1)
+    mu = jnp.asarray(_HIRANO_MU)
+    mu_w = jnp.asarray(_HIRANO_MU_W)
+    t = jnp.sqrt(jnp.clip(1.0 - mu**2, 0.0, 1.0))
+    limb = (1.0 - u1 * (1.0 - mu) - u2 * (1.0 - mu) ** 2) / (1.0 - u1 / 3.0 - u2 / 6.0)
+    pz2 = (jnp.pi * zeta) ** 2
+    s2 = s[:, None] ** 2
+    macro = jnp.exp(-pz2 * s2 * mu**2) + jnp.exp(-pz2 * s2 * t**2)
+    # M(s), Eq. (17), written in mu = cos(theta): t dt = -mu dmu
+    m_s = jnp.sum(
+        mu_w * limb * macro * _bessel_j0(2.0 * jnp.pi * s[:, None] * v_sin_i * t) * mu,
+        axis=1,
+    )  # (N,)
+    envelope = jnp.exp(-2.0 * jnp.pi**2 * beta**2 * s**2 - 4.0 * jnp.pi * gamma * s)
+    sin2 = jnp.clip(1.0 - cos_theta**2, 0.0, 1.0)
+    theta_s = 0.5 * (
+        jnp.exp(-pz2 * cos_theta[:, None] ** 2 * s**2)
+        + jnp.exp(-pz2 * sin2[:, None] * s**2)
+    )  # (n_times, N), Eq. (18) at the planet position
+    phase = 2.0 * jnp.pi * s * v_p[:, None]
+    num = jnp.sum(ds_w * envelope * m_s * theta_s * jnp.sin(phase) * s, axis=1)
+    den = jnp.sum(
+        ds_w * envelope * m_s * (m_s - f[:, None] * theta_s * jnp.cos(phase)) * s**2,
+        axis=1,
+    )
+    return -f / (2.0 * jnp.pi) * num / den
+
+
+@jax.jit
+def _compute_rossiter_mclaughlin_hirano2011_jit(
+    times: jax.Array,
+    rp_by_rs: jax.Array,
+    b_param: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    period: jax.Array,
+    rho_star: jax.Array,
+    v_sin_i: jax.Array,
+    lambda_spin: jax.Array,
+    inc_star: jax.Array,
+    q1: jax.Array,
+    q2: jax.Array,
+    line_beta: jax.Array,
+    line_gamma: jax.Array,
+    line_zeta: jax.Array,
+) -> jax.Array:
+    """JIT-compiled CCF RM anomaly [m/s]; velocity inputs in m/s."""
+    rv_fw, visible, r_sky = _rm_disc_terms(
+        times, rp_by_rs, b_param, ecosomega, esinomega, period, rho_star,
+        v_sin_i, lambda_spin, inc_star, q1, q2,
+    )  # fmt: skip
+    f = jnp.clip(1.0 - visible, 0.0, 1.0)
+    in_transit = f > 1e-9
+    f_safe = jnp.where(in_transit, f, 1.0)
+    v_sin_i_kms = v_sin_i / 1000.0
+    # flux-weighted anomaly = -(occulted flux x v_p) / visible flux, since the
+    # whole rotating disc integrates to zero velocity
+    v_p = -(rv_fw / 1000.0) * (1.0 - f) / f_safe
+    v_p = jnp.where(in_transit, jnp.clip(v_p, -v_sin_i_kms, v_sin_i_kms), 0.0)
+    cos_theta = jnp.sqrt(jnp.clip(1.0 - r_sky**2, 0.0, 1.0))
+    u1, u2 = reparametrize_LD(q1, q2)
+    dv = _hirano2011_anomaly(
+        jnp.where(in_transit, f, 0.0), v_p, cos_theta, v_sin_i_kms,
+        line_beta / 1000.0, line_gamma / 1000.0, line_zeta / 1000.0, u1, u2,
+    )  # fmt: skip
+    return jnp.where(in_transit, dv * 1000.0, 0.0)
+
+
+def compute_rossiter_mclaughlin_hirano2011(
+    times: Any,
+    rp_by_rs: Any,
+    b_param: Any,
+    ecosomega: Any,
+    esinomega: Any,
+    period: Any,
+    rho_star: Any,
+    v_sin_i: Any,
+    lambda_spin: Any,
+    inc_star: Any,
+    q1: Any = 0.5,
+    q2: Any = 0.1,
+    line_beta: Any = 2000.0,
+    line_gamma: Any = 0.0,
+    line_zeta: Any = 0.0,
+) -> jax.Array:
+    """RM anomaly [m/s] as measured by cross-correlation with a template (Hirano et al. 2011).
+
+    The RV is the maximum of the cross-correlation of the spectrum with the
+    out-of-transit template (their Eq. 16): the iodine-cell and template-matching
+    techniques, ``rm_observable: template``.
+
+    Arguments as :func:`compute_rossiter_mclaughlin`, plus the intrinsic line
+    widths in **m/s** (like ``v_sin_i``), in Hirano's convention
+    exp(-v^2/beta^2): ``line_beta`` (Gaussian: thermal, microturbulent and
+    instrumental broadening, sqrt(2) x their standard deviation added in
+    quadrature), ``line_gamma`` (Lorentzian: pressure and natural broadening) and
+    ``line_zeta`` (radial-tangential macroturbulence).
+    """
+
+    def _f(x: Any) -> jax.Array:
+        return jnp.asarray(x, dtype=jnp.float32)
+
+    return _compute_rossiter_mclaughlin_hirano2011_jit(
+        jnp.asarray(times, dtype=jnp.float32),
+        _f(rp_by_rs), _f(b_param), _f(ecosomega), _f(esinomega), _f(period),
+        _f(rho_star), _f(v_sin_i), _f(lambda_spin), _f(inc_star), _f(q1), _f(q2),
+        _f(line_beta), _f(line_gamma), _f(line_zeta),
+    )  # fmt: skip
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# CCF pipelines (HARPS, HARPS-N, ESPRESSO) measure the RV as the centre of a
+# Gaussian fitted to the CCF, which is not the maximum of the cross-correlation
+# modelled above.  Boué et al. (2013, A&A 550, A53, Eq. 14) give the shift of
+# that centre for a Gaussian line of standard deviation beta_0 (thermal,
+# microturbulent and instrumental broadening):
+#
+#   dv = -(1 / a0) [2 sigma0^2 / (sigma0^2 + beta_p^2)]^(3/2)
+#        f v_p exp(-v_p^2 / (2 (sigma0^2 + beta_p^2)))
+#
+# sigma0 and a0 are the width and area of the best Gaussian fit to the
+# out-of-transit line (the rotationally broadened, limb-darkened profile); they
+# depend on v sin i, beta_0 and the limb darkening, so they are refitted at
+# every call.  beta_p^2 = beta_0^2 + (v sin i k / 2)^2 is the width of the line
+# emitted behind the planet, broadened by the rotation across the planet disc
+# (their Sect. 4.3, for a fully occulting planet).  The linear response of the
+# centre is exact to first order in f: at the best-fit width the residual term
+# of the Hessian vanishes (a Gaussian satisfies the heat equation), so the
+# Gauss-Newton expression of Boué et al. needs no correction.
+
+_BOUE_N_U = 256  # Gauss-Legendre nodes on the rotation kernel
+_BOUE_NEWTON = 8  # Newton steps for the fitted width (4 suffice)
+_u_nodes, _u_weights = np.polynomial.legendre.leggauss(_BOUE_N_U)
+_BOUE_X = _u_nodes  # x = v / v sin i, on [-1, 1]
+_BOUE_W = _u_weights
+
+
+def _boue2013_fit(v_sin_i: Any, beta0: Any, u1: Any, u2: Any) -> tuple[Any, Any]:
+    """Width sigma0 and area a0 of the least-squares Gaussian fit to the stellar line.
+
+    The line is a Gaussian of standard deviation ``beta0`` convolved with the
+    rotation kernel of a quadratically limb-darkened disc; velocities in km/s.
+    """
+    x = jnp.asarray(_BOUE_X)
+    root = jnp.sqrt(jnp.clip(1.0 - x**2, 0.0, 1.0))
+    # rotation kernel: the intensity integrated along each strip of constant v
+    kernel = (
+        (1.0 - u1 - u2) * 2.0 * root
+        + (u1 + 2.0 * u2) * 0.5 * jnp.pi * root**2
+        - u2 * 4.0 / 3.0 * root**3
+    ) * jnp.asarray(_BOUE_W)
+    kernel = kernel / jnp.sum(kernel)
+    v2 = (x * v_sin_i) ** 2
+
+    def overlap(log_sigma: Any) -> Any:
+        """sigma * int line(v) exp(-v^2 / 2 sigma^2) dv / sqrt(2 pi) sigma."""
+        t2 = jnp.exp(2.0 * log_sigma) + beta0**2
+        return (
+            jnp.exp(log_sigma)
+            * jnp.sum(kernel * jnp.exp(-0.5 * v2 / t2))
+            / jnp.sqrt(t2)
+        )
+
+    def objective(log_sigma: Any) -> Any:
+        # least squares with the amplitude profiled out: maximise overlap^2 / sigma
+        return 2.0 * jnp.log(overlap(log_sigma)) - log_sigma
+
+    grad = jax.grad(objective)
+    hess = jax.grad(grad)
+
+    def newton(log_sigma: Any) -> Any:
+        step = grad(log_sigma) / (jnp.abs(hess(log_sigma)) + 1e-6)  # ascent
+        return log_sigma + jnp.clip(step, -0.5, 0.5)
+
+    # start from the second moment of the line; iterate without gradients, then
+    # one differentiable step, which carries the implicit-function derivative
+    log_sigma = 0.5 * jnp.log(beta0**2 + jnp.sum(kernel * v2))
+    log_sigma = jax.lax.fori_loop(0, _BOUE_NEWTON, lambda _, s: newton(s), log_sigma)
+    log_sigma = newton(jax.lax.stop_gradient(log_sigma))
+    sigma0 = jnp.exp(log_sigma)
+    # area of the fitted Gaussian relative to the unit-area line (Boué Eq. 13)
+    a0 = jnp.sqrt(2.0) * overlap(log_sigma)
+    return sigma0, a0
+
+
+def _boue2013_anomaly(
+    f: jax.Array, v_p: jax.Array, v_sin_i: Any, beta0: Any, k: Any, u1: Any, u2: Any
+) -> jax.Array:
+    """Boué et al. (2013) Eq. (14): Gaussian-fit RM anomaly [km/s].
+
+    ``f`` and ``v_p`` [km/s] have shape (n_times,); ``beta0`` is the standard
+    deviation of the non-rotating line [km/s], ``k`` the radius ratio.
+    """
+    sigma0, a0 = _boue2013_fit(v_sin_i, beta0, u1, u2)
+    s2 = sigma0**2 + beta0**2 + (0.5 * v_sin_i * k) ** 2  # sigma0^2 + beta_p^2
+    return (
+        -(1.0 / a0)
+        * (2.0 * sigma0**2 / s2) ** 1.5
+        * f
+        * v_p
+        * jnp.exp(-0.5 * v_p**2 / s2)
+    )
+
+
+@jax.jit
+def _compute_rossiter_mclaughlin_boue2013_jit(
+    times: jax.Array,
+    rp_by_rs: jax.Array,
+    b_param: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    period: jax.Array,
+    rho_star: jax.Array,
+    v_sin_i: jax.Array,
+    lambda_spin: jax.Array,
+    inc_star: jax.Array,
+    q1: jax.Array,
+    q2: jax.Array,
+    line_beta: jax.Array,
+) -> jax.Array:
+    """JIT-compiled Gaussian-fit RM anomaly [m/s]; velocity inputs in m/s."""
+    rv_fw, visible, _ = _rm_disc_terms(
+        times, rp_by_rs, b_param, ecosomega, esinomega, period, rho_star,
+        v_sin_i, lambda_spin, inc_star, q1, q2,
+    )  # fmt: skip
+    f = jnp.clip(1.0 - visible, 0.0, 1.0)
+    in_transit = f > 1e-9
+    f_safe = jnp.where(in_transit, f, 1.0)
+    v_sin_i_kms = v_sin_i / 1000.0
+    v_p = -(rv_fw / 1000.0) * (1.0 - f) / f_safe
+    v_p = jnp.where(in_transit, jnp.clip(v_p, -v_sin_i_kms, v_sin_i_kms), 0.0)
+    u1, u2 = reparametrize_LD(q1, q2)
+    # the width is given in Hirano's convention, exp(-v^2/beta^2)
+    beta0 = line_beta / 1000.0 / jnp.sqrt(2.0)
+    dv = _boue2013_anomaly(
+        jnp.where(in_transit, f, 0.0), v_p, v_sin_i_kms, beta0, rp_by_rs, u1, u2
+    )
+    return jnp.where(in_transit, dv * 1000.0, 0.0)
+
+
+def compute_rossiter_mclaughlin_boue2013(
+    times: Any,
+    rp_by_rs: Any,
+    b_param: Any,
+    ecosomega: Any,
+    esinomega: Any,
+    period: Any,
+    rho_star: Any,
+    v_sin_i: Any,
+    lambda_spin: Any,
+    inc_star: Any,
+    q1: Any = 0.5,
+    q2: Any = 0.1,
+    line_beta: Any = 2000.0,
+) -> jax.Array:
+    """RM anomaly [m/s] as the centre of a Gaussian fit to the CCF (Boué et al. 2013).
+
+    The RV of HARPS-like pipelines, ``rm_observable: ccf``.
+
+    Arguments as :func:`compute_rossiter_mclaughlin_hirano2011`: ``line_beta`` [m/s] is
+    the width of the non-rotating line in Hirano's convention exp(-v^2/beta^2),
+    i.e. sqrt(2) times its standard deviation.
+    """
+
+    def _f(x: Any) -> jax.Array:
+        return jnp.asarray(x, dtype=jnp.float32)
+
+    return _compute_rossiter_mclaughlin_boue2013_jit(
+        jnp.asarray(times, dtype=jnp.float32),
+        _f(rp_by_rs), _f(b_param), _f(ecosomega), _f(esinomega), _f(period),
+        _f(rho_star), _f(v_sin_i), _f(lambda_spin), _f(inc_star), _f(q1), _f(q2),
+        _f(line_beta),
+    )  # fmt: skip
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Transit Timing Variation (TTV) lightcurve model
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def compute_transit_ttv(
+    times: jax.Array,
+    transit_times: jax.Array,
+    duration: jax.Array,
+    impact_param: jax.Array,
+    radius_ratio: jax.Array,
+    u_limb_dark: jax.Array,
+) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
+    """Compute the transit light curve for a single planet using TTV parameterization.
+
+    Unlike :func:`compute_lightcurve`, which is parameterized by a global
+    period + reference epoch (suitable for systems with a strictly linear
+    ephemeris), this function accepts an *explicit array of individual
+    mid-transit times*.  Each transit epoch can deviate from a linear
+    ephemeris by an arbitrary amount (the TTV signal), and
+    :class:`jaxoplanet.orbits.ttv.TTVOrbit` computes the correct transit
+    geometry for every epoch independently.
+
+    **What is modeled here:**
+      * The primary transit depth and shape at the correct TTV-corrected time
+        for every observed transit.  The transit profile is described by a
+        quadratic limb-darkening law (triangular parameterization via
+        ``u_limb_dark``).
+      * No secondary eclipse, no phase curve, no beaming, no ellipsoidal
+        variations.  Those components should be computed separately using the
+        *linear* ephemeris (``linear_t0``, ``linear_period``) returned by this
+        function, and then summed with the transit offset.
+
+    **When to use this instead of compute_lightcurve:**
+      Use this function when the planet's transit times are free parameters
+      (one per observed transit epoch) rather than a single (t0, period) pair.
+      This is the correct model whenever TTVs are large enough to be
+      measurable (amplitude ≳ 2 minutes for typical photometric precision).
+
+    **Precision notes:**
+      The function works in whatever precision JAX is currently configured
+      for (float32 or float64).  In PyAETNA's default float32 mode the transit
+      times should be in the *centered* frame (i.e., subtracted by
+      ``timeCenter``) to keep their magnitude well within the float32 range.
+      A centered time of ≲ 15 d has float32 precision of ≲ 1.8e-6 d ≈ 0.15 s,
+      which is sufficient for typical photometric TTV measurements.
+
+    **JIT compatibility:**
+      The number of transit epochs ``n_transits = len(transit_times)`` is
+      fixed at compile time (it is a Python-level shape, not a traced value).
+      The *values* of ``transit_times`` are dynamic JAX tracers.  This
+      design is fully compatible with ``jax.jit`` and ``numpyro.MCMC``.
+
+    Parameters
+    ----------
+    times : jax.Array, shape (N,)
+        Observation times in the centered frame [days].
+        Must be the same reference frame as ``transit_times``
+        (typically ``TIME - timeCenter``).
+    transit_times : jax.Array, shape (n_transits,)
+        Mid-transit time for each observed transit epoch [days, centered frame].
+        Element ``i`` is the time of the *i*-th transit that falls within the
+        observation window, ordered chronologically.  Typically constructed as:
+            ``expected_t + ttv_offset``
+        where ``expected_t = tau_ref_centered + epoch_index * period_ref`` is
+        the linear-ephemeris prediction and ``ttv_offset`` is the free
+        TTV parameter for that epoch.
+    duration : jax.Array, shape (1,)
+        Transit duration T14 [days] — from first to last contact.
+        Fixed (not TTV-perturbed) across all epochs for this implementation.
+    impact_param : jax.Array, shape (1,)
+        Impact parameter b = a cos(i) / R_star (dimensionless, 0 = central).
+    radius_ratio : jax.Array, shape (1,)
+        Planet-to-star radius ratio Rp/Rs (dimensionless).
+    u_limb_dark : jax.Array, shape (2,)
+        Limb-darkening coefficients [u1, u2] in the triangular (q1, q2)
+        reparameterization *already converted* to the standard Kipping (2013)
+        coefficients u1, u2 by :func:`reparametrize_LD`.
+
+    Returns
+    -------
+    transit_offset : jax.Array, shape (N,)
+        Additive flux offset from stellar baseline due to the transit.
+        The value is ≤ 0 (flux deficit during transit), and exactly 0
+        out of transit.  Units: same as ``y`` in the data file (normalized
+        flux, so the transit depth is ≈ (Rp/Rs)²).
+    linear_t0 : jax.Array, shape (1,)
+        Best-fit linear reference epoch from the TTVOrbit internal least-squares
+        fit through all ``transit_times`` [days, centered frame].  This is the
+        constant term of the linear trend T_i = linear_t0 + i * linear_period.
+        Use this to fold the RV data or to compute the phase curve with the
+        linear ephemeris.
+    linear_period : jax.Array, shape (1,)
+        Best-fit linear period from the TTVOrbit internal least-squares fit
+        [days].  Use this together with ``linear_t0`` for any computation
+        that requires a global period (phase curves, RV, corner-plot labels).
+    ttvs : jax.Array, shape (n_transits,)
+        Observed minus Computed (O-C) residuals for every transit epoch
+        [days].  These are the TTV signals:
+            ttvs[i] = transit_times[i] - (linear_t0 + i * linear_period)
+        where ``i`` is the epoch index inferred from the linear fit.
+        Positive = transit arrived late; negative = transit arrived early.
+
+    Notes
+    -----
+    TTVOrbit attributes used:
+        orbit.linear_t0[0]     → linear_t0 (1-element array)
+        orbit.linear_period[0] → linear_period (1-element array)
+        orbit.ttvs[0]          → ttvs (n_transits-element array)
+
+    Example
+    -------
+    >>> # Assume tau_ref=2459000.0, period_ref=3.5, timeCenter=2459015.0
+    >>> tau_ref_c = 2459000.0 - 2459015.0  # = -15.0  (centered)
+    >>> epoch_indices = jnp.array([0, 1, 2, 3, 4])
+    >>> expected_times = tau_ref_c + epoch_indices * 3.5    # linear ephem
+    >>> ttv_offsets    = jnp.array([0.01, -0.02, 0.015, -0.01, 0.005])  # TTV
+    >>> transit_times  = expected_times + ttv_offsets
+    >>> offset, lt0, lP, ttvs = compute_transit_ttv(
+    ...     times, transit_times,
+    ...     duration=jnp.array([0.12]),
+    ...     impact_param=jnp.array([0.3]),
+    ...     radius_ratio=jnp.array([0.1]),
+    ...     u_limb_dark=jnp.array([0.4, 0.2]),
+    ... )
+    """
+    # TTVOrbit accepts transit_times as a *tuple* of 1-D arrays, one per planet.
+    # For a single-planet call (PyAETNA evaluates one planet at a time), we
+    # wrap the array in a length-1 tuple.
+    orbit = TTVOrbit(
+        transit_times=(transit_times,),  # 1-tuple: 1 planet, n_transits epochs
+        duration=duration,  # shape (1,) — one planet
+        impact_param=impact_param,  # shape (1,) — one planet
+        radius_ratio=radius_ratio,  # shape (1,) — one planet
+    )
+
+    # limb_dark_light_curve returns a function f(times) → array of shape (N, n_planets)
+    # For n_planets=1 the shape is (N, 1); we squeeze the planet axis.
+    lc_raw = limb_dark_light_curve(orbit, u_limb_dark)(times)  # (N, 1)
+    transit_offset = lc_raw[:, 0]  # (N,) — additive offset, ≤ 0 at transit
+
+    # Extract derived linear ephemeris from the internal least-squares fit.
+    # TTVOrbit fits T_i = linear_t0 + i * linear_period + ttvs[i] to the
+    # supplied transit_times, so:
+    #   orbit.linear_t0[0]    : reference epoch (first transit in the fit)
+    #   orbit.linear_period[0]: mean orbital period
+    #   orbit.ttvs[0]         : O-C residuals array, shape (n_transits,)
+    linear_t0 = orbit.linear_t0[0]  # scalar (float32 or float64)
+    linear_period = orbit.linear_period[0]  # scalar
+    ttvs = orbit.ttvs[0]  # shape (n_transits,)
+
+    return transit_offset, linear_t0, linear_period, ttvs
+
+
+@partial(jit, static_argnames=("transit_only", "circular"))
+def compute_lightcurve(
+    times: jax.Array,
+    q_limb_dark: jax.Array,
+    rho_star: float | jax.Array,
+    t0s: jax.Array,
+    periods: jax.Array,
+    b_params: jax.Array,
+    rp_by_rs: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    amp_night: jax.Array,
+    amp_day_minus_amp_night: jax.Array,
+    amp_beam: jax.Array,
+    amp_ellip: jax.Array,
+    phase_offsets: jax.Array,
+    tidal_lag: jax.Array,
+    K_amp: jax.Array,
+    transit_only: bool = False,
+    circular: bool = False,
+) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
+    """
+    Compute multi-planet lightcurve with phase curves, beaming, and ellipsoidal variations.
+
+    ``transit_only=True`` (static) returns the transits alone, with zero
+    phase-curve, beaming, ellipsoidal and RV components: the caller sets it
+    when every phase-curve amplitude is fixed to zero.  The secondary
+    eclipse, light-travel-time positions, velocities and RVs are then never
+    evaluated (about 1.5x faster on GPU, 1.7x on CPU at 1e5 points).
+
+    ``circular=True`` (static) builds circular orbits, ignoring *ecosomega*
+    and *esinomega*: the caller sets it when both are fixed to zero for every
+    planet.  jaxoplanet then computes positions without solving Kepler's
+    equation, which with an eccentricity array of zeros takes about a third of
+    a transit evaluation (K2-140, CPU); the result is the same.
+
+    Uses dimensionless stellar units (M_star = R_star = 1.0) with stellar density
+    to set the orbital scale via Kepler's third law.
+
+    The key observable is a/R_star = (G*M_star/(4π²))^(1/3) * P^(2/3) / R_star
+    which can be rewritten as: a/R_star ∝ (ρ_star * P²)^(1/3)
+
+    PERFORMANCE NOTES:
+    - JIT-compiled for speed
+    - Fully vectorized over time and planets
+    - System initialization is efficient
+    - Phase offset interpretation: shifts RADIATING hemisphere relative to ILLUMINATED hemisphere
+
+    LIGHT TRAVEL TIME (LTT):
+    Photons observed at t_obs were emitted at t_em = t_obs + z(t_obs)/c, where
+    z is the planet's line-of-sight displacement in R_sun (+z towards observer).
+    For typical hot Jupiters (a ~ 10 R_sun), |Δt| ~ 25 s.
+    The LTT correction is applied to the PHASE CURVE
+    orbital phase only; transit/eclipse geometry, beaming, and ellipsoidal
+    variations use observed-time positions (physically correct).
+    One Newton iteration suffices (residual < 1e-4 for all known systems).
+    Eccentric orbits are handled exactly via re-solving Kepler's equation.
+
+    Phase Offset Physics:
+    - orbital_phase_angle = 0 at secondary eclipse (planet fully illuminated from observer's view)
+    - orbital_phase_angle = π at primary transit (planet's night side toward observer)
+    - phase_offset > 0: hotspot shifted EAST (evening terminator, trailing limb)
+    - phase_offset < 0: hotspot shifted WEST (morning terminator, leading limb)
+    - shifted_angle = orbital_phase_angle + phase_offset defines where peak emission occurs
+
+    The modulation uses the shifted angle, so:
+    - If phase_offset = 0: peak brightness at secondary eclipse (α=0)
+    - If phase_offset = π/4: peak brightness happens at α=π/4 (after secondary eclipse)
+
+    Parameters:
+    -----------
+    times : array
+        Times at which to evaluate the model [days]
+    q_limb_dark : array, shape (2,)
+        Limb darkening parameters [q1, q2] for reparametrization
+    rho_star : float
+        Stellar density [g/cm³]. Used to set stellar mass (M = 4/3 π R³ ρ) so that
+        jaxoplanet derives the orbital scale from period and stellar mass.
+        Typical values: 0.5-3.0 g/cm³ for main-sequence stars. Solar value: 1.4 g/cm³
+    t0s : array, shape (n_planets,)
+        Mid-transit times [days]
+    periods : array, shape (n_planets,)
+        Orbital periods [days]
+    b_params : array, shape (n_planets,)
+        Impact parameters (dimensionless, b = a*cos(i)/R_star)
+    rp_by_rs : array, shape (n_planets,)
+        Planet-to-star radius ratios (dimensionless)
+    ecosomega, esinomega : arrays, shape (n_planets,)
+        Eccentricity parameterization (dimensionless)
+    amp_night, amp_day_minus_amp_night, amp_beam, amp_ellip : arrays, shape (n_planets,)
+        Phase curve amplitudes [ppm]
+    phase_offsets : array, shape (n_planets,)
+        Phase offset for each planet's emission peak [radians]
+        Positive = hotspot shifted toward trailing hemisphere (evening)
+        Negative = hotspot shifted toward leading hemisphere (morning)
+    tidal_lag : array, shape (n_planets,)
+        Tidal lag angle for ellipsoidal variations [radians]
+        The stellar tidal bulge lags behind the planet due to:
+        - Finite response time of stellar material (viscous dissipation)
+        - Stellar rotation relative to orbital period
+        Positive = bulge lags behind planet (typical for slow stellar rotation)
+        Negative = bulge leads planet (rare, for super-synchronous rotation)
+        Typical values: 0-10 degrees for hot Jupiter systems
+    K_amp : array, shape (n_planets,)
+        RV semi-amplitudes [m/s] (for beaming calculation)
+
+    Returns:
+    --------
+    planet_offset : array
+        Normalized planetary light curve offset from baseline (total_flux - 1.0)
+        This is the additive contribution: flux_obs = 1.0 + planet_offset + detrending
+    lc_phase : array, shape (n_planets, n_times)
+        Phase curve component per planet [ppm]
+    lc_beam : array, shape (n_planets, n_times)
+        Doppler beaming component per planet [ppm]
+    lc_ellip : array, shape (n_planets, n_times)
+        Ellipsoidal variation component per planet [ppm]
+    lc_transit : array
+        Transit light curve component (relative flux)
+    rvs : array, shape (n_planets, n_times)
+        Radial velocities per planet [m/s]
+    """
+    # ==========================================
+    # 1. Create System with Stellar Density
+    # ==========================================
+
+    # Convert eccentricity parameterization
+    eccentricities, omegas = reparametrize_ecc_omega(ecosomega, esinomega)
+    n_planets = len(t0s)
+
+    # Create star with given density.
+    # rho_star is in g/cm³; jaxoplanet uses solar units (M_sun, R_sun, day).
+    # For R_star = 1 R_sun: M_star [M_sun] = ρ_star [g/cm³] / ρ_sun [g/cm³]
+    star_radius = 1.0
+    star_mass = rho_star / RHO_SUN_CGS  # Stellar mass in M_sun for R_star = 1 R_sun
+
+    star = Central(mass=star_mass, radius=star_radius)
+
+    # ------------------------------------------------------------------
+    # Physics guard for b >= 1 (no transit / eclipse)
+    # ------------------------------------------------------------------
+    # When b >= 1 + rp_by_rs the planet does not cross the stellar disc,
+    # so no transit or eclipse occurs.  However jaxoplanet's internal
+    # geometry (limb_dark_light_curve, emission_lc) can produce NaN for
+    # such configurations because it takes sqrt(1 - b²) etc.
+    #
+    # Strategy (fully JIT-compatible via jnp.where):
+    #   1. Build the orbital system with b_safe = min(b, 0.99) so that
+    #      jaxoplanet geometry never sees b >= 1.  This keeps positions,
+    #      velocities AND transit/eclipse calculations finite.
+    #   2. After computing transit & eclipse light-curves, multiply them
+    #      by a per-planet mask that zeros them out when the real b
+    #      implies no transit (b >= 1 + k).
+    #   3. Phase curve, beaming, and ellipsoidal variations depend only
+    #      on orbital phase (not on b) and are always valid.
+    # ------------------------------------------------------------------
+    transit_occurs = b_params < (1.0 + rp_by_rs)  # (n_planets,)
+    b_safe = jnp.where(transit_occurs, b_params, 0.99)
+
+    system = System(star)
+
+    # Add planets using period directly (jaxoplanet computes a/R_star internally)
+    for i in range(n_planets):
+        system = system.add_body(
+            period=periods[i],  # Use period directly
+            time_transit=t0s[i],
+            radius=rp_by_rs[i],
+            impact_param=b_safe[i],
+            eccentricity=None if circular else eccentricities[i],
+            omega_peri=None if circular else omegas[i],
+            radial_velocity_semiamplitude=K_amp[i],
+        )
+
+    # ==========================================
+    # 2. Calculate Geometry & Kinematics (Vectorized)
+    # ==========================================
+
+    # Get relative position and velocity
+    # Shape: (3, n_planets, n_times) - vectorized by jaxoplanet
+    pos = system.relative_position(times)
+    vel = system.relative_velocity(times)
+
+    x, y, z = pos[0], pos[1], pos[2]
+    vx, vy, vz = vel[0], vel[1], vel[2]
+
+    # --- Light Travel Time Correction (phase curve only) ---
+    # Photons we observe at t_obs were emitted at t_em = t_obs + z/c.
+    # Re-evaluate planet positions at emission times so the Lambert phase
+    # function sees the correct orbital geometry.  Transit, eclipse-mask,
+    # beaming, and ellipsoidal use *observed* positions (no LTT).
+    x_ltt, y_ltt, z_ltt = _compute_ltt_positions(system, times, z)
+
+    # ==========================================
+    # 3. Compute Flux Components (Fully Vectorized)
+    # ==========================================
+
+    # --- Component 1: Primary Transits ---
+    u_limb_dark = reparametrize_LD(q_limb_dark[0], q_limb_dark[1])
+    # Raw transit computed with b_safe — always finite
+    lc_transit_raw = limb_dark_light_curve(system, u_limb_dark)(
+        times
+    )  # (n_times, n_planets)
+    # Zero out planets that have no transit (b >= 1 + k)
+    lc_transit_raw = lc_transit_raw * transit_occurs[None, :]  # broadcast mask
+    lc_transit = jnp.sum(lc_transit_raw, axis=-1)
+
+    if transit_only:
+        # Positions, velocities and LTT positions computed above do not reach
+        # these outputs, so XLA removes them from the compiled graph.
+        zeros = jnp.zeros((t0s.shape[0], times.shape[0]), dtype=lc_transit.dtype)
+        return lc_transit, zeros, zeros, zeros, lc_transit, zeros
+
+    # --- Component 2: Phase Curve (Vectorized over planets and time) ---
+    #
+    # HOTSPOT OFFSET PHYSICS:
+    # -----------------------
+    # The substellar point is where the star is directly overhead on the planet.
+    # For a tidally locked planet, this point is fixed on the planet surface.
+    #
+    # The hotspot (peak thermal emission) may be OFFSET from the substellar point
+    # due to atmospheric circulation (winds carry heat eastward/westward).
+    #
+    # phase_offset defines WHERE on the planet the hotspot is located:
+    # - phase_offset = 0: Hotspot at substellar point (no offset)
+    # - phase_offset > 0: Hotspot shifted EAST of substellar (evening terminator)
+    #                     Peak brightness occurs AFTER secondary eclipse
+    # - phase_offset < 0: Hotspot shifted WEST of substellar (morning terminator)
+    #                     Peak brightness occurs BEFORE secondary eclipse
+    #
+    # The key insight: we need a SIGNED orbital phase to distinguish pre/post eclipse.
+    # The orbital_phase_signed gives us the planet's position in its orbit:
+    # - φ_orb = 0 at secondary eclipse (substellar point faces us)
+    # - φ_orb > 0 after secondary eclipse (planet moved east)
+    # - φ_orb < 0 before secondary eclipse (planet approaching from west)
+    #
+    # To see maximum brightness from the hotspot, we need:
+    #   φ_orb = phase_offset
+    # i.e., the planet must have moved past secondary eclipse by angle phase_offset
+    # for the hotspot to directly face us.
+    #
+    # The "viewing angle" to the hotspot is: |φ_orb - phase_offset|
+    # This is the angle between our line of sight and the hotspot direction.
+
+    # OBSERVED orbital phase — used for ellipsoidal variations (tidal
+    # bulge follows the planet's instantaneous position, no LTT needed)
+    # Shape: (n_planets, n_times)
+    orbital_phase_signed = compute_orbital_phase_signed(x, y, z)
+
+    # LTT-CORRECTED orbital phase — used for thermal phase curve.
+    # Evaluates where the planet *was* when it emitted the photons we
+    # receive at t_obs.  This shifts the peak by ~a/c ≈ 25 s for a
+    # typical hot Jupiter, detectable with JWST-class precision.
+    orbital_phase_ltt = compute_orbital_phase_signed(x_ltt, y_ltt, z_ltt)
+
+    # Compute viewing angle to hotspot with PROPER ANGLE WRAPPING
+    # This is the angle between our line of sight and the hotspot on the planet
+    # phase_offsets shape: (n_planets,) -> broadcast to (n_planets, 1)
+    #
+    # The raw difference can be in [-2π, 2π], which causes discontinuities
+    # when orbital_phase_signed crosses ±π at transit.
+    # We use atan2(sin, cos) to properly wrap to [-π, π] before taking abs.
+    angle_diff = orbital_phase_ltt - phase_offsets[:, None]
+
+    # Wrap angle difference to [-π, π] using atan2 (handles all edge cases)
+    # This ensures smooth transitions across the transit (where φ jumps from +π to -π)
+    angle_diff_wrapped = jnp.arctan2(jnp.sin(angle_diff), jnp.cos(angle_diff))
+
+    # Convert to [0, π] for the Lambert phase function
+    # The Lambert function expects:
+    #   α = 0: hotspot directly facing us (maximum brightness)
+    #   α = π: hotspot on far side (minimum brightness)
+    hotspot_viewing_angle = jnp.abs(angle_diff_wrapped)
+
+    # Compute phase curve modulation using Lambert law
+    # Shape: (n_planets, n_times)
+    modulation = getLambertPhaseCurve(hotspot_viewing_angle)
+
+    # Vectorized computation of phase curve flux
+    # amp_day calculated once, broadcast properly
+    amp_day = amp_day_minus_amp_night + amp_night
+    lc_phase = amp_night[:, None] + (amp_day[:, None] - amp_night[:, None]) * modulation
+
+    # --- Component 2a: Smooth Secondary Eclipse Masking ---
+    # During secondary eclipse, the STAR occults the PLANET'S thermal emission.
+    # We use jaxoplanet's emission model which computes smooth ingress/egress
+    # profiles based on geometric occultation.
+    #
+    # The emission API calculates limb geometry, producing continuous and
+    # differentiable eclipse profiles suitable for gradient-based inference.
+
+    # Build emission system with planet surface properties
+    star_surface = Surface(u=u_limb_dark)
+
+    # Create planet bodies and surfaces for emission calculation
+    # Uses b_safe to keep jaxoplanet geometry finite; masking applied below
+    planet_bodies_with_surfaces = []
+    for i in range(n_planets):
+        # Define planet body with orbital parameters
+        planet_body = Body(
+            period=periods[i],
+            time_transit=t0s[i],
+            radius=rp_by_rs[i],
+            impact_param=b_safe[i],
+            eccentricity=None if circular else eccentricities[i],
+            omega_peri=None if circular else omegas[i],
+        )
+        # A planet surface of unit emission: its light curve is the visible
+        # fraction of the planetary disc (1 out of occultation, 0 when hidden),
+        # the eclipse mask itself.  The day-side amplitude multiplies lc_phase
+        # above, so no division by it is needed: dividing by amp_day * 1e-6 gave
+        # 0/0 = NaN in float32 when the sampler drove amp_day to ~1e-35, whose
+        # product with 1e-6 is a denormal that the hardware flushes to zero.
+        planet_surface = Surface(
+            amplitude=1.0,
+            u=(),  # No limb darkening for planet emission
+            normalize=False,  # keep the unit amplitude
+        )
+        planet_bodies_with_surfaces.append((planet_body, planet_surface))
+
+    # Build emission system for smooth eclipse calculation
+    emission_system = SurfaceSystem(
+        central=star, central_surface=star_surface, bodies=planet_bodies_with_surfaces
+    )
+
+    # Compute emission light curves with smooth secondary eclipse profiles
+    # Output shape: (n_times, n_planets+1) where:
+    #   - Column 0: stellar baseline (not used)
+    #   - Column i+1: planet i's emission * smooth_eclipse_mask
+    emission_flux = emission_lc(emission_system)(times)
+
+    # Eclipse masks, shape (n_planets, n_times): 0 = fully eclipsed, 1 = fully visible
+    eclipse_masks = emission_flux[:, 1:].T
+
+    # Zero out eclipse masks for planets with no transit/eclipse (b >= 1+k)
+    # When transit_occurs is False, the eclipse mask should be 1.0 (no occultation)
+    eclipse_masks = jnp.where(transit_occurs[:, None], eclipse_masks, 1.0)
+
+    # Apply smooth secondary eclipse masks to phase curves (element-wise multiplication)
+    lc_phase = lc_phase * eclipse_masks
+
+    # --- Component 3: Doppler Beaming ---
+    # The star brightens when its reflex motion is TOWARD the observer
+    # (negative stellar RV), so beaming = -amp_beam * RV_star / K.  With
+    # jaxoplanet's conventions the planet's relative vz equals -RV_star (both
+    # scaled by K): vz / K = -RV_star / K, i.e. +1 at phase 0.25 after transit,
+    # where the star approaches.
+    # K_amp may be zero for LC-only cases; avoid division-by-zero by using a
+    # safe denominator and masking the ratio to zero when K_amp == 0.
+    K_amp_safe = jnp.where(K_amp[:, None] != 0.0, K_amp[:, None], 1.0)
+    ratio = vz / K_amp_safe
+    ratio = jnp.where(K_amp[:, None] != 0.0, ratio, 0.0)
+    lc_beam = amp_beam[:, None] * ratio
+
+    # --- Component 4: Ellipsoidal Variations (Fully Vectorized) ---
+    # The planet's gravity deforms the star into an ellipsoid, with brightness maxima
+    # at quadrature (twice per orbit, when viewing the elongated axis).
+    #
+    # TIDAL LAG PHYSICS:
+    # The tidal bulge follows the planet's orbital motion, but lags behind due to:
+    # - Finite response time of stellar material (viscous dissipation)
+    # - Stellar rotation (if rotating faster/slower than orbital period)
+    #
+    # The lag is typically small (few degrees) and follows the planet's motion
+    # (not a uniform rotation) because the tidal force is always directed
+    # along the star-planet line.
+    #
+    # IMPORTANT: Must use SIGNED orbital phase (orbital_phase_signed) to properly
+    # account for the directional lag. Using unsigned phase would make the lag
+    # symmetric (same before/after eclipse), which is physically incorrect.
+    #
+    # Formula: EV ~ cos(2 * (φ - δ_lag))
+    # where φ = SIGNED orbital phase, δ_lag = tidal lag angle
+    # We SUBTRACT the lag because the bulge lags BEHIND the planet's position
+    # Factor of 2: two maxima per orbit (both quadratures)
+    lc_ellip = -amp_ellip[:, None] * jnp.cos(
+        2 * (orbital_phase_signed - tidal_lag[:, None])
+    )
+
+    # ==========================================
+    # 4. Combine and Return (Vectorized Summation)
+    # ==========================================
+
+    # Sum over all planets (axis=0)
+    total_phase_effects = jnp.sum(lc_phase + lc_beam + lc_ellip, axis=0)
+
+    # Return NORMALIZED offset (flux - 1.0, i.e. normalized stellar baseline removed)
+    # for consistent additive model structure
+    # This allows uniform treatment in the MC sampler: observed_flux = 1.0 + planet_offset + poly + fourier + GP
+
+    planet_offset = lc_transit + total_phase_effects * 1e-6
+
+    # Radial velocities (vectorized)
+    rvs = system.radial_velocity(times)
+
+    return planet_offset, lc_phase, lc_beam, lc_ellip, lc_transit, rvs
+
+
+# ==============================================================================
+# Exposure-time integration (per-point supersampling)
+# ==============================================================================
+
+# Default number of uniformly-spaced sub-samples per exposure for the Kipping
+# (2010) resampling scheme.  Must be odd; 7 is the standard
+# (used by batman, exoplanet, jaxoplanet's built-in scalar integrator).
+SUPERSAMPLE_NUM = 7
+
+
+def exposure_average(
+    fn: Any, times: jax.Array, exptime_seconds: Any, **kwargs: Any
+) -> jax.Array:
+    """Average ``fn(times=..., **kwargs)`` over each exposure (Kipping 2010).
+
+    Each time is replaced by ``SUPERSAMPLE_NUM`` equally spaced sub-samples over
+    ``[t - exptime/2, t + exptime/2]`` and the results are averaged, as for the
+    light curves (:func:`compute_lightcurve_supersampled`). Used for the RM
+    anomaly of RV instruments with ``EXPTIME > 0``: through ingress and egress
+    the anomaly changes within a 10-minute exposure. Points with zero exposure
+    time get the instantaneous value.
+    """
+    ns = SUPERSAMPLE_NUM
+    dt = jnp.linspace(-0.5, 0.5, ns, dtype=times.dtype)
+    exp_days = jnp.asarray(exptime_seconds, dtype=times.dtype) / SECONDS_PER_DAY
+    t_sub = times[:, None] + dt[None, :] * exp_days[:, None]
+    out = fn(times=t_sub.reshape(-1), **kwargs)
+    return jnp.mean(out.reshape(times.shape[0], ns), axis=1)
+
+
+@partial(jit, static_argnames=("transit_only", "circular"))
+def compute_lightcurve_supersampled(
+    times: jax.Array,
+    exptime_seconds: jax.Array,
+    q_limb_dark: jax.Array,
+    rho_star: float | jax.Array,
+    t0s: jax.Array,
+    periods: jax.Array,
+    b_params: jax.Array,
+    rp_by_rs: jax.Array,
+    ecosomega: jax.Array,
+    esinomega: jax.Array,
+    amp_night: jax.Array,
+    amp_day_minus_amp_night: jax.Array,
+    amp_beam: jax.Array,
+    amp_ellip: jax.Array,
+    phase_offsets: jax.Array,
+    tidal_lag: jax.Array,
+    K_amp: jax.Array,
+    transit_only: bool = False,
+    circular: bool = False,
+) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, jax.Array]:
+    """
+    Compute multi-planet lightcurve with **per-point** exposure-time integration.
+
+    Wraps :func:`compute_lightcurve` with the Kipping (2010) resampling scheme:
+    each data point is evaluated at ``SUPERSAMPLE_NUM`` equally-spaced sub-sample
+    times spanning ``[t − exptime/2, t + exptime/2]``, then averaged (rectangle
+    rule).  This convolves the model with a top-hat kernel of width ``exptime``.
+
+    **Pre-condition (enforced by the caller):** all ``N`` rows passed here must
+    have ``exptime_seconds > 0``.  Filtering out zero-exptime rows before
+    calling this function is the responsibility of :func:`_compute_planet_model`
+    in ``PyAETNA/sampling/engine/runner.py``, which pre-computes static per-instrument per-point split
+    indices and routes zero-exptime rows to :func:`compute_lightcurve` instead.
+    Calling this function with ``exptime == 0`` rows wastes
+    ``SUPERSAMPLE_NUM − 1`` redundant evaluations per such point — that
+    inefficiency has been eliminated at the call site.
+
+    Because the integration is vectorised over all ``N`` data points
+    simultaneously (single call to ``compute_lightcurve`` with
+    ``N × SUPERSAMPLE_NUM`` times), the function is fully JIT-compatible and
+    GPU-friendly.
+
+    .. note:: **Why per-point rather than per-visit?**
+
+       In practice, all points within a visit usually share the same
+       exposure time, so one might expect a scalar-per-visit approach to be
+       faster.  However, for JAX the difference is negligible: the
+       ``compute_lightcurve`` call at ``N × 7`` points dominates the budget,
+       and the per-point broadcast that builds the sub-sample grid is
+       essentially free.  A per-visit loop would *reduce* vectorisation
+       (more, smaller calls), and expanding a scalar back to an ``(N,)``
+       array would be identical to the current design.  Keeping a single
+       ``(N,)`` shape also avoids shape-dependent JIT recompilations.
+
+    Parameters
+    ----------
+    times : array, shape (N,)
+        Mid-exposure times [days], same reference frame as orbital elements.
+        All rows must have a corresponding positive ``exptime_seconds`` value.
+    exptime_seconds : array, shape (N,)
+        Per-point exposure duration [seconds] — must be **> 0** for every row
+        (caller must filter out zero-exptime points before calling this
+        function).
+    q_limb_dark … K_amp :
+        Same parameters as :func:`compute_lightcurve`.
+
+    Returns
+    -------
+    Identical 6-tuple as :func:`compute_lightcurve`, with each output averaged
+    over the exposure window:
+
+    * ``planet_offset`` — shape ``(N,)``
+    * ``lc_phase``      — shape ``(n_planets, N)``
+    * ``lc_beam``       — shape ``(n_planets, N)``
+    * ``lc_ellip``      — shape ``(n_planets, N)``
+    * ``lc_transit``    — shape ``(N,)``
+    * ``rvs``           — shape ``(n_planets, N)``
+    """
+    n_times = times.shape[0]
+    ns = SUPERSAMPLE_NUM
+    exptime_days = exptime_seconds / SECONDS_PER_DAY
+
+    # Sub-sample offsets: equally spaced in [-0.5, 0.5] × exptime_days.
+    # The dtype must follow ``times``: a default-dtype linspace is float64
+    # whenever x64 is enabled, and would silently promote the whole transit
+    # model (Kepler solver included) out of the float32 firewall, making it
+    # ~20x slower on consumer GPUs.
+    dt = jnp.linspace(-0.5, 0.5, ns, dtype=times.dtype)  # (ns,)
+    t_sub = times[:, None] + dt[None, :] * exptime_days[:, None]  # (N, ns)
+    t_flat = t_sub.reshape(-1)  # (N * ns,)
+
+    # Evaluate model at all sub-sample times in a single vectorised call
+    (
+        planet_offset_flat,
+        lc_phase_flat,
+        lc_beam_flat,
+        lc_ellip_flat,
+        lc_transit_flat,
+        rvs_flat,
+    ) = compute_lightcurve(
+        t_flat,
+        q_limb_dark,
+        rho_star,
+        t0s,
+        periods,
+        b_params,
+        rp_by_rs,
+        ecosomega,
+        esinomega,
+        amp_night,
+        amp_day_minus_amp_night,
+        amp_beam,
+        amp_ellip,
+        phase_offsets,
+        tidal_lag,
+        K_amp,
+        transit_only=transit_only,
+        circular=circular,
+    )
+
+    # Average over sub-samples — 1-D arrays (n_times * ns,) → (n_times,)
+    planet_offset = jnp.mean(planet_offset_flat.reshape(n_times, ns), axis=1)
+    lc_transit = jnp.mean(lc_transit_flat.reshape(n_times, ns), axis=1)
+
+    # 2-D per-planet arrays: (n_planets, n_times * ns) → (n_planets, n_times)
+    n_planets = lc_phase_flat.shape[0]
+    lc_phase = jnp.mean(lc_phase_flat.reshape(n_planets, n_times, ns), axis=2)
+    lc_beam = jnp.mean(lc_beam_flat.reshape(n_planets, n_times, ns), axis=2)
+    lc_ellip = jnp.mean(lc_ellip_flat.reshape(n_planets, n_times, ns), axis=2)
+    rvs = jnp.mean(rvs_flat.reshape(n_planets, n_times, ns), axis=2)
+
+    return planet_offset, lc_phase, lc_beam, lc_ellip, lc_transit, rvs
+
+
+if __name__ == "__main__":
+    q_limb_dark = jnp.array([0.1, 0.3])  # Limb darkening
+    t0s = jnp.array([0.0, 0.5])
+    periods = jnp.array([1.5, 4.0])  # Orbital periods [days]
+    b_params = jnp.array([0.1, 0.0])
+    rp_by_rs = jnp.array([0.01, 0.1])
+    ecosomega = jnp.array([0.0, 0.0])  # between -1 and 1
+    esinomega = jnp.array([0.0, 0.0])  # between -1 and 1
+    amp_night = jnp.array([0, 0])  # in ppm
+    amp_day_minus_amp_night = jnp.array([0, 0])  # in ppm
+    amp_beam = jnp.array([0, 1])  # in ppm
+    amp_ellip = jnp.array([0, 5])  # in ppm
+    phase_offsets = jnp.array([0.0, jnp.pi / 4])  # in radians
+    tidal_lag = jnp.array([0.0, jnp.pi / 4])  # in radians
+    K_amp = jnp.array([1e-10, 100])  # in m/s
+
+    # ==========================================
+    # Run and Visualize - Demonstrate Flexibility
+    # ==========================================
+    # Time grid: 4 days (covers 1 orbit of Planet c)
+    t_grid = jnp.linspace(0, 4.0, 1000)
+
+    # Method 1: Use stellar density (can be sampled in MCMC)
+    rho_star = 1.4  # g/cm³ (solar-like)
+    planet_offset1, phases1, beams1, ellips1, transits1, rvs1 = compute_lightcurve(
+        t_grid,
+        q_limb_dark,
+        rho_star,
+        t0s,
+        periods,
+        b_params,
+        rp_by_rs,
+        ecosomega,
+        esinomega,
+        amp_night,
+        amp_day_minus_amp_night,
+        amp_beam,
+        amp_ellip,
+        phase_offsets,
+        tidal_lag,
+        K_amp,
+    )
+    flux1 = 1.0 * 0 + planet_offset1
+
+    # Setup Plot
+    fig, axes = plt.subplots(6, 1, figsize=(10, 12), sharex=True)
+
+    # 1. Total Lightcurve (with flexible stellar density)
+    axes[0].plot(t_grid, flux1, "k", lw=1.5)
+    axes[0].set_title(
+        f"Total Lightcurve (ρ_star = {rho_star:.2f} g/cm³, Dimensionless & Flexible)",
+        fontsize=12,
+        weight="bold",
+    )
+    axes[0].set_ylabel("Relative Flux")
+
+    # 2. Phase Curves (Thermal/Reflected) - use first model's components
+    # Sum over planets to show total signal for that component
+    axes[1].plot(t_grid, jnp.sum(phases1, axis=0), "C1")
+    axes[1].set_ylabel("Phase Curve\n(ppm)")
+    axes[1].text(
+        0.02,
+        0.85,
+        "Note: Cuts to 0 during Secondary Eclipse",
+        transform=axes[1].transAxes,
+        fontsize=9,
+    )
+
+    # 3. Doppler Beaming
+    axes[2].plot(t_grid, jnp.sum(beams1, axis=0), "C2")
+    axes[2].set_ylabel("Doppler Beaming\n(ppm)")
+
+    # 4. Ellipsoidal Variations
+    axes[3].plot(t_grid, jnp.sum(ellips1, axis=0), "C3")
+    axes[3].set_ylabel("Ellipsoidal\n(ppm)")
+
+    # 5. Transits
+    axes[4].plot(t_grid, transits1 * 1e6, "C0")
+    axes[4].set_ylabel("Transits\n(ppm)")
+    axes[4].set_xlabel("Time [days]")
+
+    # 6. RVs
+    axes[5].plot(t_grid, jnp.sum(rvs1, axis=0), "C4")
+    axes[5].set_ylabel("Radial Velocity\n(m/s)")
+    axes[5].set_xlabel("Time [days]")
+
+    for ax in axes:
+        ax.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.show()
+
+    # ==========================================
+    # Additional Test: Phase Offset Comparison
+    # ==========================================
+    # print("\n=== Phase Offset Test ===")
+
+    # Single planet test with different offsets
+    t_test = jnp.linspace(0, 4.0, 500)  # One orbit of 4-day planet
+
+    fig2, axes2 = plt.subplots(2, 1, figsize=(12, 8), sharex=True)
+
+    # Test 1: Circular orbit, different phase offsets
+    offsets_to_test = [0.0, jnp.pi / 6, jnp.pi / 4, jnp.pi / 3]
+    colors = ["C0", "C1", "C2", "C3"]
+
+    for offset, color in zip(offsets_to_test, colors):
+        _, phases_test, _, _, _, _ = compute_lightcurve(
+            t_test,
+            q_limb_dark,
+            rho_star,
+            jnp.array([0.0]),
+            jnp.array([4.0]),
+            jnp.array([0.1]),
+            jnp.array([0.1]),
+            jnp.array([0.0]),
+            jnp.array([0.0]),  # Circular
+            jnp.array([100.0]),
+            jnp.array([800.0]),
+            jnp.array([0.0]),
+            jnp.array([0.0]),
+            jnp.array([offset]),
+            jnp.array([0.0]),
+            jnp.array([100.0]),
+        )
+        label = f"offset = {offset * 180 / jnp.pi:.0f}°"
+        axes2[0].plot(t_test, phases_test[0], color=color, label=label, lw=1.5)
+
+    axes2[0].axvline(2.0, color="k", ls="--", alpha=0.5, label="Secondary Eclipse")
+    axes2[0].set_ylabel("Phase Curve (ppm)")
+    axes2[0].set_title(
+        "Circular Orbit: Phase Offset Comparison", fontsize=12, weight="bold"
+    )
+    axes2[0].legend(loc="upper right")
+    axes2[0].grid(True, alpha=0.3)
+
+    # Test 2: Eccentric orbit (e=0.3) with same offsets
+    ecc = 0.3
+    omega = jnp.pi / 4  # Argument of periastron
+    ec_test = jnp.sqrt(2 * ecc) * jnp.cos(omega)
+    es_test = jnp.sqrt(2 * ecc) * jnp.sin(omega)
+
+    for offset, color in zip(offsets_to_test, colors):
+        _, phases_test, _, _, _, _ = compute_lightcurve(
+            t_test,
+            q_limb_dark,
+            rho_star,
+            jnp.array([0.0]),
+            jnp.array([4.0]),
+            jnp.array([0.1]),
+            jnp.array([0.1]),
+            jnp.array([ec_test]),
+            jnp.array([es_test]),  # Eccentric!
+            jnp.array([100.0]),
+            jnp.array([800.0]),
+            jnp.array([0.0]),
+            jnp.array([0.0]),
+            jnp.array([offset]),
+            jnp.array([0.0]),
+            jnp.array([100.0]),
+        )
+        label = f"offset = {offset * 180 / jnp.pi:.0f}°"
+        axes2[1].plot(t_test, phases_test[0], color=color, label=label, lw=1.5)
+
+    axes2[1].set_ylabel("Phase Curve (ppm)")
+    axes2[1].set_xlabel("Time [days]")
+    axes2[1].set_title(
+        f"Eccentric Orbit (e={ecc:.1f}, ω={omega * 180 / jnp.pi:.0f}°): Phase Offset Comparison",
+        fontsize=12,
+        weight="bold",
+    )
+    axes2[1].legend(loc="upper right")
+    axes2[1].grid(True, alpha=0.3)
+
+    plt.tight_layout()
+    plt.show()
+
+    print("Phase offset = 0°: Peak at secondary eclipse")
+    print("Phase offset > 0°: Peak shifted AFTER secondary eclipse (eastward hotspot)")
+    print("Phase offset < 0°: Peak shifted BEFORE secondary eclipse (westward hotspot)")