feat: K2-140 benchmark against PyORBIT, seed runs for the evidence scatter,...

| Data set | Code | Sampler | Laptop CPU | Laptop GPU | Server CPU | Server GPU | Evaluations |

|---|---|---|---|---|---|---|---|

| K2-140 | juliet | NS (dynesty, 500 live points) | 119 s (20) | – | 63 s (60) | – | 1.06 M |

| K2-140 | PyORBIT | NS (dynesty, 500 live points), PyAETNA's model and priors | 87 s (24) | – | 37 s (60) | – | 1.09 M |

| K2-140 | PyAETNA | NS (slice, 500 live points) | 8 min (24) | 56 s | 4 min (60) | 41 s | 5.71 M |

| Pi Men (circ) | allesfitter | NS (dynesty, 500 live points) | 62 min (24) | – | 21 min (60) | – | 3.16 M |

| Pi Men (circ) | PyAETNA | NS (slice, 500 live points) | 75 s (24) | 27 s | 67 s (60) | 22 s | 5.15 M |

- Against **pyaneti**, with the same ensemble MCMC, PyAETNA is 1.2 to 1.8 times faster.

- A **GPU** makes PyAETNA up to 8.5 times faster than the CPU of the same machine, and it is never

  slower (the smallest ratio, 1.02, is for K2-100 (NS, laptop)). Against the reference codes on their

  CPUs, the GPU runs are 1.1 to 137 times faster, except for TOI-5398 against PyORBIT.

  CPUs, the GPU runs are 1.1 to 137 times faster, except for K2-140 against PyORBIT, TOI-5398 against PyORBIT.

**Where PyAETNA is slower.**

- **Nested sampling against juliet and PyORBIT**: PyAETNA is 3.5 to 4.0 times slower than juliet and

  3.4 to 3.8 times slower than PyORBIT on the same CPUs, because its slice sampler needs

  3.4 to 6.0 times slower than PyORBIT on the same CPUs, because its slice sampler needs

  2.7 to 5.4 times more likelihood evaluations than dynesty's random walk. The next section shows what the

  extra evaluations buy.

- **NumPyro ensemble samplers** (pyro-ppl/numpyro#2189): releases 0.14.0–0.21 updated the second half of the ensemble

  against itself, so ESS/AIES posteriors were too narrow. PyAETNA refuses to run them on those versions.

## Peak search, ensembles and single precision

- **Peak search** (`sampling/peaks.py`): it runs in the unbounded coordinates of NumPyro, entirely in float64: the

  float32 physical models are smooth enough for sampling, but an optimiser's line searches stall on their rounding,

  and so do second derivatives. L-BFGS brings each start into the basin of a peak; on badly scaled energies (a

  normalisation known to 1e-5 next to an amplitude known to 1e-1) it then converges slowly, and 30 damped Newton

  (Levenberg-Marquardt) steps with the exact Hessian reach the top. Where automatic differentiation returns undefined

  second derivatives (the occultation model of WASP-18), the Hessian comes from central differences of the gradient.

  Two end points belong to the same peak when one lies within the 99.9 % region of the Gaussian approximation of the

  other, or within three marginal widths of it in every parameter, so that optima scattered along a curved ridge

  (b-rho* degeneracy) are one peak while aliases stay separate. A point whose Hessian is not positive definite is

  never a separate significant peak; along a direction of negative or vanishing curvature its Gaussian approximation

  takes the width of the widest well-measured direction.

- **Walker order of the ensemble slice sampler**: NumPyro's ESS shuffles the walkers at every step to split the

  ensemble at random, so each stored chain is a different walker at every step and its autocorrelation time is

  meaningless (1-3 steps whatever the model; 37 steps for a 20-dimensional Gaussian once the walkers are kept in

  place). PyAETNA restores their order after each step. The convergence test thins the chains by up to a quarter of

  the previous estimate of tau, which leaves tau unchanged within its noise.

- **Unbounded coordinates in single precision**: started from the prior, the ESS draws slices of enormous width and

  walkers reach unbounded coordinates of 1e3-1e5, where the map back to an interval prior rounds to the bound itself

  in float32; they then return slowly or not at all. PyAETNA confines the unbounded coordinates of parameters with

  interval priors to at most 15 (which excludes 3e-7 of each prior interval at each end) and keeps a walker in place

  when a move returns a non-finite position.

- **NUTS and narrow priors**: NUTS needs float64 for a parameter whose posterior spans only a few hundred float32

  values (the period of K2-140, with a prior 15 parts per million wide): its leapfrog steps round to zero and the

  chains do not move. The automatic precision therefore selects the mixed precision for NUTS when a prior is narrower

  than 1e-4 of its value.

## Planets without a transit

An RV-only planet (radius and phase-curve amplitudes fixed to zero) must not be passed to the light-curve model: a

   The paper cites the tag and its commit id. If the referee process changes code or results, repeat the same steps

   and move the paper to the new tag.

   Decided on 2026-10-11: steps 2–3 (the full clean rerun and the tag that the final paper cites) are made after the

   first referee report, so that they include the referee's changes; at submission the text rests on targeted reruns,

   each run recording its commit and uncommitted changes.

5. **Zenodo archive at acceptance.** Once the paper is accepted:

   1. deposit the tagged release (code plus this folder) on Zenodo;

   2. cite the DOI in the final version of the paper.

| Folder | What it shows | Reference |

|---|---|---|

| `k2140_juliet/` | Joint K2 + LCOGT + 3 RV instruments, 20 parameters: posteriors, ln Z and run time against juliet; nested-sampling settings | juliet tutorial data and priors (Espinoza et al. 2019), run here with dynesty on 16 threads |

| `k2140_pyorbit/` | K2-140 with PyAETNA's model and priors in PyORBIT (batman, dynesty's random walk): likelihood check, posteriors, ln Z over 29 unseeded runs against PyAETNA's 16 seeds, run time | data of `k2140_juliet/` converted by `make_inputs.py`, run here with PyORBIT 11.2.7 |

| `pimen_allesfitter/` | An RV-only planet (2093 d) and a transiting one, TESS + 3 RV instruments; circular vs eccentric model comparison | allesfitter paper runs (Günther & Daylan 2021), published outputs |

| `toi216_ttv/` | Free transit times of two planets near resonance, Matérn-3/2 GP, 37 parameters | allesfitter paper run |

| `wasp18_phase_curve/` | Phase curve with thermal emission, beaming and ellipsoidal variation | allesfitter paper run |

| `k2100_pyaneti/` | Multi-dimensional GP of HARPS-N RVs, log R'HK and BIS, against pyaneti | pyaneti's K2-100 example (Barragán et al. 2019), run here with pyaneti |

| `toi5398_rm_pyorbit/` | Rossiter-McLaughlin effect with PyAETNA's three RM observables, against PyORBIT's simulated CCF | Mantovan et al. (2024) HARPS-N data and priors, run here with PyORBIT |

| `eiv_validation/` | Errors-in-variables decorrelation on simulated RVs: slope recovery and coverage | simulated here |

| `ttv_scaling/` | Nested sampling against NUTS on a synthetic system with 12–100 free transit times (18–106 parameters) | simulated here |

| `ns_calibration/` | Bias and error calibration of the evidence on 10–20-D problems with known ln Z; ln Z scatter of the K2-140 benchmark over 8 more seeds (`k2140_seeds/`) | analytic; K2-140 as in `k2140_juliet/` |

| `ttv_scaling/` | The synthetic planet with 12–100 free transit times (18–106 parameters): data and configurations; its 106-parameter case is run in `sampler_comparison/` | simulated here |

| `ns_calibration/` | Bias and error calibration of the evidence on 10–20-D problems with known ln Z; ln Z scatter of the K2-140 benchmark over 16 seeds (`k2140_seeds/`) | analytic; K2-140 as in `k2140_juliet/` |

| `scaling/` | Likelihood throughput vs number of points (1e3–3e6) and model class, CPU and GPU, against batman/celerite/george | measured here |

| `gpu_precision/` | Throughput in float32/float64, the float32 firewall, the XLA fusion-autotuner issue (profiling scripts; values transcribed into `results/summary.json`) | internal |

| Benchmark | Files here | Downloaded from | Original analysis |

|---|---|---|---|

| `k2140_juliet/` | `data/K2-140_lc.dat`, `data/K2-140_rvs.dat`, `data/K2-140_priors_fullcircular.dat` | https://github.com/nespinoza/juliet, folders `tutorial1/` and `tutorial1/priors/` (raw files under `https://raw.githubusercontent.com/nespinoza/juliet/master/tutorial1/`) | Espinoza, Kossakowski & Brahm (2019), MNRAS 490, 2262; K2-140 b: Giles et al. (2018), Korth et al. (2019) |

| `k2140_pyorbit/` | `data/k2140_*.dat` (written by `make_inputs.py` from `k2140_juliet/data/`; RVs converted to m/s) | as `k2140_juliet/` | as `k2140_juliet/`; PyORBIT: Malavolta et al. (2016, 2018) |

| | `reference/juliet_16.json`, `juliet_post.npz`, `juliet16_summary.log` | **run here** with `run_juliet.py` (juliet 2.2.10, dynesty 3.1.0, 16 threads; `environment/juliet_env.txt`) | — |

| `pimen_allesfitter/` | `data/TESS.csv`, `AAT.csv`, `HARPS1.csv`, `HARPS2.csv`, `data/README.txt`; `reference/allesfit_*/` (settings, params), `reference/af*_ns_table.csv`, `af*_*.log` | https://github.com/MNGuenther/allesfitter, folder `paper/Pi_Mensae/` (and its `allesfit_*/results/`) | Günther & Daylan (2021), ApJS 254, 13; Huang et al. (2018), Gandolfi et al. (2018) |

| `toi216_ttv/` | `data/TESS.csv`; `reference/` (settings, params, `ns_table.csv`, logs, `run.py`) | allesfitter repository, folder `paper/TOI-216/allesfit/` (and `results/`) | Günther & Daylan (2021); Dawson et al. (2019, 2021), Kipping et al. (2019) |

"""Does PyAETNA's ensemble warmup help? Summary of the direct-against-warmup test of one benchmark.

Each benchmark has runs/warmup_test/{direct,warmup}_s{1,2,3}: the same ensemble MCMC (move, walkers, total steps),

run directly from a broad start, or after PyAETNA's ensemble warmup with re-seeding of stuck walkers. For every run

this script records the wall time, the walkers in the main mode (criterion of mcmc_efficiency.main_mode), the

effective sample size of the slowest parameter over the main-mode walkers, the independent draws per minute and the

walkers re-seeded by the warmup; then the largest shift and the range of width ratios of each run against the first

warmup run, with all the walkers and with the main-mode walkers only (widths: half the 16-84 % interval).

With an optional converged reference run (samples NPZ of nested sampling, or of a run started near the solution),

each run is also compared with it, to tell whether either mode converged at all.

usage: python warmup_test.py <benchmark dir> <config name> <burn steps of the direct runs> [<reference samples.npz>]

writes <benchmark dir>/results/warmup_test.json (with WARMUP_TEST_DIR=<folder>: runs/<folder>, results/<folder>.json)

"""

from __future__ import annotations

import json

import os

import re

import sys

from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))

from mcmc_efficiency import main_mode, summarise

def wall_s(time_out: Path) -> float:

    """Wall-clock seconds from GNU time -v output."""

    text = time_out.read_text()

    value = re.search(r"Elapsed \(wall clock\) time .*: ([\d:.]+)", text).group(1)

    seconds = 0.0

    for part in value.split(":"):

        seconds = 60 * seconds + float(part)

    return seconds

def load(run: Path, config: str, burn: int) -> dict[str, np.ndarray]:

    d = np.load(run / "PyAETNAouts" / f"{config}_samples.npz")

    n_walkers = max(d[k].shape[0] for k in d.files if d[k].ndim == 2)

    return {

        k: d[k][:, burn:]

        for k in d.files

        if d[k].ndim == 2 and d[k].shape[0] == n_walkers and np.ptp(d[k]) > 0

    }

def stats(chains: dict[str, np.ndarray], mask: np.ndarray) -> dict[str, tuple]:

    out = {}

    for k, x in chains.items():

        lo, med, hi = np.percentile(x[mask], [16, 50, 84])

        out[k] = (med, (hi - lo) / 2)

    return out

def main() -> None:

    bench, config, burn_direct = Path(sys.argv[1]), sys.argv[2], int(sys.argv[3])

    folder = os.environ.get(

        "WARMUP_TEST_DIR", "warmup_test"

    )  # e.g. warmup_test_err (WASP-18, binned errors)

    base = bench / "runs" / folder

    runs = sorted(p for p in base.iterdir() if (p / "done.txt").exists())

    chains, out = {}, {

        "benchmark": bench.name,

        "burn_steps_direct": burn_direct,

        "runs": {},

    }

    for run in runs:

        burn = burn_direct if run.name.startswith("direct") else 0

        c = load(run, config, burn)

        ok = main_mode(c)

        eff = summarise({k: x[ok] for k, x in c.items()}, 1, 0)

        log = (run / "PyAETNAouts" / f"{config}_run.log").read_text()

        reseeded = re.findall(r"Re-seeding (\d+) stuck walker", log)

        wall = wall_s(run / "time.out")

        chains[run.name] = (c, ok)

        out["runs"][run.name] = {

            "wall_s": wall,

            "walkers": int(ok.size),

            "walkers_in_main_mode": int(ok.sum()),

            "kept_steps": eff["kept_steps"],

            "max_tau_steps": eff["max_tau_steps"],

            "slowest_parameter": eff["slowest_parameter"],

            "min_ess_main_mode": eff["min_ess"],

            "ess_per_minute": eff["min_ess"] / (wall / 60),

            "walkers_reseeded_per_round": [int(n) for n in reseeded],

        }

    ref_c, ref_ok = chains[min(k for k in chains if k.startswith("warmup"))]

    refs = {"first_warmup": stats(ref_c, ref_ok)}

    if len(sys.argv) > 4:

        d = np.load(sys.argv[4])

        flat = {k: np.ravel(d[k])[None, :] for k in ref_c if k in d.files}

        refs["reference"] = stats(flat, np.ones(1, dtype=bool))

        out["reference_samples"] = sys.argv[4]

    for name, (c, ok) in chains.items():

        for label, mask in (("all_walkers", np.ones_like(ok)), ("main_mode", ok)):

            s = stats(c, mask)

            for rname, ref in refs.items():

                worst = max(ref, key=lambda k: abs(s[k][0] - ref[k][0]) / ref[k][1])

                out["runs"][name][f"vs_{rname}_{label}"] = {

                    "max_abs_shift_sigma": abs(s[worst][0] - ref[worst][0])

                    / ref[worst][1],

                    "worst_parameter": worst,

                    "width_ratio_min": min(s[k][1] / ref[k][1] for k in ref),

                    "width_ratio_max": max(s[k][1] / ref[k][1] for k in ref),

                }

    dest = bench / "results" / f"{folder}.json"

    dest.write_text(json.dumps(out, indent=1))

    for name, r in out["runs"].items():

        print(

            f"{name}: {r['wall_s']:.0f} s, main mode {r['walkers_in_main_mode']}/{r['walkers']}, "

            f"tau {r['max_tau_steps']:.0f}, ESS {r['min_ess_main_mode']:.0f} ({r['ess_per_minute']:.0f}/min), "

            f"re-seeded {r['walkers_reseeded_per_round']}, all-walker widths up to "

            f"{r['vs_first_warmup_all_walkers']['width_ratio_max']:.2f}"

        )

if __name__ == "__main__":

    main()