fix: full-precision matrix products in the RM disc integrals (TF32 on GPUs biased the anomaly)

  `ln L (maximum)` and stored with χ² in the new `goodness_of_fit` section of `<config>_results.json`. AICc uses

  the standard correction 2k(k+1)/(N−k−1) (it divided by N−k). The `map` section's score is now called

  `map_logposterior` (was `map_loglike`).

- **Rossiter–McLaughlin model on GPUs (2026-10-04):** JAX runs float32 matrix products on NVIDIA GPUs in TF32

  (10-bit mantissa) by default. In the disc integrals of the RM observables the occulted flux fraction was then

  wrong by up to 3e-4 and the CCF anomaly by up to 2 m/s near ingress and egress (CPUs were exact); on TOI-5398 the

  GPU runs gave ln Z 0.45 nats above the CPU runs. These products now use full precision. Transit, phase-curve and

  GP models are not affected (checked: 0.2 ppm between GPU and CPU on a transit).

- The log's summary tables print the effective samples per second (the column read 0.00); with nested sampling

  the table is titled POSTERIOR SUMMARY and has no ESS and R-hat columns, which say nothing about independent

  draws.

    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).

"""The RM disc integrals must not use TF32 matrix products on GPUs.

With JAX's default precision, float32 matrix products on NVIDIA GPUs run in

TF32 (10-bit mantissa): the occulted flux fraction was then wrong by up to

3e-4 and the CCF anomaly by up to 2 m/s near ingress, while the CPU was exact.

The check is on the traced program, so it runs on any device.

"""

import jax

import jax.numpy as jnp

from PyAETNA import body_models

def _dot_precisions(jaxpr: jax.extend.core.ClosedJaxpr) -> list:

    found = []

    def walk(jx: jax.extend.core.Jaxpr) -> None:

        for eqn in jx.eqns:

            if eqn.primitive.name == "dot_general":

                found.append(eqn.params["precision"])

            for v in eqn.params.values():

                if isinstance(v, jax.extend.core.ClosedJaxpr):

                    walk(v.jaxpr)

                elif isinstance(v, jax.extend.core.Jaxpr):

                    walk(v)

    walk(jaxpr.jaxpr)

    return found

def test_rm_disc_products_use_highest_precision() -> None:

    f32 = jnp.float32

    args = (

        jnp.linspace(-0.05, 0.05, 5, dtype=f32),

        *(

            f32(v)

            for v in (0.1, 0.3, 0.0, 0.0, 3.0, 1.4, 5000.0, -0.5, 1.5708, 0.45, 0.25)

        ),

    )

    jaxpr = jax.make_jaxpr(body_models._rm_disc_terms)(*args)

    precisions = _dot_precisions(jaxpr)

    assert precisions, "no matrix product found: the test no longer checks anything"

    highest = jax.lax.Precision.HIGHEST

    for p in precisions:

        assert p is not None and all(q == highest for q in p), p