fix: extending a float64 NUTS run (plain tuples in the checkpoint upcast);...

  interval priors (`ensemble_kernels.BoundedESS`; the excluded prior mass is 3e-7 of the interval at each end).

- **NUTS with narrow priors:** `use_float64: auto` now selects the mixed precision for Branch 2 when a free

  parameter's prior is narrower than 1e-4 of its value (`precision.NUTS_RELATIVE_WIDTH_LIMIT`).

- **Extending a NUTS run in float64** (`num_samples_if_nuts` raised with `reset: false`) failed with "cond branches

  must have equal output types": the step-size adaptation state is stored in plain tuples, which the upcast of the

  checkpoint to float64 did not enter.

- The log warning about walkers still stuck after the re-seeding rounds now says that it is harmless when NUTS follows.

- These were found by the paper's NUTS benchmarks: until then the default route had run only on models with GPs,

  which are always in float64.

   inflate ensemble posteriors and trigger a spurious "multimodal" verdict. The

   log line `Re-seeding N stuck walker(s)` reports it. Switch off with

   `reseed_stuck_walkers: false`; change the cut with `stuck_walker_threshold`.

   The warmup does not have to converge: its job is to let a few walkers find

   the region of the maximum likelihood, from which NUTS starts. Walkers

   that are still stuck after the last round (the log then warns) are

   harmless when NUTS follows, because only the lowest-energy walkers are used.

3. **Phase 2, energy-spread test.** Compares the spread of walker potential

   energies with `chain_energy_threshold × n_free_params`.

4. **Phase 3a, unimodal → NUTS.** If the spread is ≤ threshold, NUTS chains

- **Resuming and extending MCMC runs**: with `reset: false`, an interrupted run

  continues from its checkpoint, and a completed one is extended when

  `num_samples_if_nuts` (or `num_samples_if_ensemble`) is raised,

  e.g. when R-hat shows the chains have not converged. The run continues with

  e.g. when R-hat shows the chains have not converged (the usual way to

  converge a hard posterior: raise it in steps and check R-hat and the effective

  sample size in the log or in the results file; the checkpoint holds the adapted

  step size and mass matrix, so nothing is lost). The run continues with

  the model and priors **stored in the checkpoint**: if the data, planets, priors

  or noise sections of the YAML have changed, the log warns and names them; set

  `reset: true` to sample the new configuration.

#!/bin/bash

# 2026-10-03: extend the TOI-216 warm-up -> NUTS run (runs/branch2_h100) until its chains have converged: the

# checkpoint is continued (reset: false) with more draws per chain, in steps of 3000, until the largest R-hat is

# below 1.01 and the smallest effective sample size above 400 (at most 18000 draws per chain). H100, heavy.lock.

# usage: round5_toi216_server.sh <work dir>

set -u

W=$(realpath "$1"); B=$W/PyAETNA/paper_benchmarks; LOG=$W/logs/round5.log

exec 9> "$W/heavy.lock"; flock 9

export PATH=$W/envs/pyaetna/bin:$PATH GNU_TIME=$W/tools/gnutime/bin/time

say() { echo "$(date '+%F %T') $*" >> "$LOG"; }

R=toi216_ttv/runs/branch2_h100

cd "$B/$R" || exit 1

[ -f toi216.yaml.initial ] || cp toi216.yaml toi216.yaml.initial

sed -i 's/^  reset: true/  reset: false/' toi216.yaml

n=1500

while [ $n -lt 18000 ]; do

  n=$((n + 3000))

  sed -i "s/^  num_samples_if_nuts: .*/  num_samples_if_nuts: $n/" toi216.yaml

  say "start extension to $n draws per chain"

  JAX_PLATFORMS=cuda,cpu nice -n 19 "$GNU_TIME" -v pyaetna toi216.yaml > run_$n.out 2> time_$n.out

  say "end extension to $n exit=$? $(grep 'Elapsed (wall' time_$n.out)"

  conv=$(python - <<'PY'

import json, glob

d = json.load(open(glob.glob("PyAETNAouts/*_results.json")[0]))["parameters"]

free = {k: v for k, v in d.items() if not k.startswith(("transit_times", "linear_", "ttvs_"))}

print(int(max(v["rhat"] for v in free.values()) < 1.01 and min(v["ess"] for v in free.values()) > 400),

      max(v["rhat"] for v in free.values()), min(v["ess"] for v in free.values()))

PY

)

  say "after $n draws: converged/maxRhat/minESS = $conv"

  [ "${conv%% *}" = 1 ] && break

done

say "done"

def upcast_value_to_float64(val: Any) -> Any:

    """Recursively upcast JAX arrays and nested structures to float64/int64 (x64 mode).

    Supported container types: ``dict``, ``list``, named-tuples (``_replace``

    Supported container types: ``dict``, ``list``, plain tuples, named-tuples (``_replace``

    protocol), and any object with a ``dtype`` + ``shape`` attribute (JAX/NumPy

    arrays).  Unrecognised types are returned unchanged.

    elif isinstance(val, list):

        # Recurse into lists so that lists of JAX arrays are upcast correctly.

        return [upcast_value_to_float64(v) for v in val]

    elif isinstance(val, tuple) and not hasattr(val, "_fields"):

        # Plain tuples, such as the step-size adaptation state of NUTS (``ss_state``:

        # log step size, its average, the gradient average, the step counter, mu).

        return tuple(upcast_value_to_float64(v) for v in val)

    elif hasattr(val, "dtype") and hasattr(val, "shape"):

        if val.dtype == jnp.float32:

            return jnp.array(val, dtype=jnp.float64)

    _log(

        plModel,

        f"  Walkers still stuck after {max_rounds} re-seeding rounds: the posterior "

        f"may have a genuine secondary mode, or the warmup is too short.",

        f"  Walkers still stuck after {max_rounds} re-seeding rounds. This does not matter "

        f"when NUTS follows (it starts from the lowest-energy walkers, which only need to have "

        f"found the region of the maximum likelihood); for an ensemble production run, "

        f"it means a genuine secondary mode or a warmup that is too short.",

        level="WARNING",

    )

    return total