feat: referee-round benchmarks (reference evidences, injection and recovery,...

# paper_benchmarks: keep inputs, configs, portable outputs; drop regenerable or fragile files

paper_benchmarks/scaling/data/

paper_benchmarks/**/runs/**/*.pckl

paper_benchmarks/**/*.pckl

paper_benchmarks/**/juliet_out_*/

paper_benchmarks/**/runs/**/*.dat

paper_benchmarks/**/runs/**/run.out

# PyORBIT outputs (pickles; the portable export is in <benchmark>/reference/)

# Node.js dependencies of the configurator tests

node_modules/

prompt.md

# Local instructions for the coding assistant (paths of the author's machines): local only

CLAUDE.md

Whenever you send me messages, report the current date and time in local time, compact format, boldface, and in parentheses. For example: **(2024-06-05 14:30)**. This will help me keep track of the timeline of our discussion.

# PyAETNA: instructions for the code and the benchmarks (the software paper: `paper/CLAUDE.md`)

## Working with the author

- The author is an exoplanet astronomer, not a software developer. Make the standard technical choices yourself

  and explain the outcome in plain language; do not ask them to write code or to choose between technical options.

  Still ask before committing, unless they said to go ahead for that piece of work.

- Never print or propagate access tokens (e.g. the token file in `/dati/software/pyLib/`).

## Git commits

- **Python environment**: use `/dati/pyenvs/pyaetnaEnv`.

- **Commit often**: small, focused commits are easier to review and understand. Ask for confirmation before

  committing.

- **Pre-commit tests**: run all pre-commit hooks before committing; recipes are in `.pre-commit-config.yaml`

  (black, ruff, mypy).

- **Repository test suite**: run it before committing, with `pytest -v --maxfail=1 --disable-warnings`. Fix all

  failures, including those unrelated to your changes, before committing.

- **Format**: Conventional Commits (feat:, fix:, docs:, refactor:, test:, chore:).

- **Subject line**: at most 100 characters, no period.

- **Body**: only if truly necessary; omit it if the subject says enough. No lengthy explanations unless the change

  warrants them.

- **Examples**: `fix: handle NaN in RV uncertainties`; `feat: add Bayesian model comparison for exoplanet

  candidates`; `refactor: simplify posterior sampling loop`.

- **Not committed**: `prompt.md`, the author's own edits to `.gitignore` and `CLAUDE.md` unless asked, and large

  outputs. Checkpoints (`.pckl`) and the posterior chains of ensemble-MCMC runs stay out of git

  (`paper_benchmarks/.gitignore`); their summaries are kept.

## Heavy runs

- **One heavy job at a time** on each machine: a GPU or CPU PyAETNA run, the full pytest suite and benchmark

  scripts never overlap. Only light commands (grep, git, lint, single small tests) may run alongside. On the laptop,

  check `nvidia-smi` and `free -g` first. On the server, every queue takes the lock `heavy.lock`.

- **Laptop timings** only in the `performance` power profile, on mains, with uncapped CPU clocks. A reboot brings

  back the quiet profile, and a suspend/resume caps every core at 2 GHz until the next reboot. Check all three

  before each timed run; each run records them in `power_profile.txt`. Ask the author before long laptop runs.

- **Server** (ExoREC-Ganymede, `ssh ganymede-cuda13`, workspace `/DATA/ExoREC/pyaetna_bench`): if unreachable,

  activate the OACT VPN (`nmcli connection up OACT`). GitHub clones fail there, so copy sources from the laptop.

- **Nothing that must survive goes in `/tmp`** or the session scratchpad (a tmpfs wiped at reboot). Environments,

  sources, run folders and queue logs go in `/dati/software/pyLib/PyAETNA_bench_work/`. Copy results into the

  repository as soon as a run ends.

- When timing a JAX computation, wait for it to finish (`jax.block_until_ready`): JAX returns before the work is

  done.

- Whenever a long job is not required to run on the laptop, run it on the server

## Paper benchmarks

`paper_benchmarks/` (public, in this repository) reproduces every result of the software paper. A benchmark folder

`paper_benchmarks/<case>/` holds `data/`, `configs/`, `runs/` and `results/summary.json`. Before changing anything

there, read the **Reproducibility policy** in `paper_benchmarks/README.md`:

- runs record their provenance (commit, dirty flag, package versions), and each `results/*.json` records the

  hardware, software versions, wall-clock time and settings;

- outputs are kept as portable `.npz`/`.json`, not pickles;

- the paper cites a release tag made at submission, from clean reruns (`./rerun_all.sh`);

- the tagged release is archived on Zenodo, with a DOI, at acceptance.

Numbers in the paper come only from `*/results/*.json`, through the generated tables or by reading the JSON. If a

result changes, rerun the generators; never retype numbers.

**Fair comparisons with other codes.**

- Same machine; each code at its own best number of CPUs (from a scan); one BLAS thread per worker. Timings from

  other machines are not quoted.

- Samplers and models as close as possible: if the reference code runs an ensemble MCMC, run PyAETNA with the same

  MCMC (moves, walkers, steps, start) in direct mode, without its warmup. If the reference code offers a closed-form

  model equivalent to PyAETNA's, compare with that one.

- Remaining large differences go in the paper constructively: strengths and weaknesses of each approach, and

  guidance on which code to choose. Every case where PyAETNA is slower is tracked, with its cause and possible

  improvements, in `localDocs/performance/PYAETNA_SLOWER_CASES.md`.

## The software paper

The paper is in `paper/`, its own repository (synced with Overleaf; ignored here). Everything about writing and

revising it (scope, language and style, layout, conventions, building, the live documents and the revision log) is in

`paper/CLAUDE.md`. The paper takes its numbers only from `paper_benchmarks/*/results/*.json`, through the generated

tables and figures: if a result changes, rerun the generators; never retype numbers.

| `n_live_if_ns` | `500` | Live points. The error on $\ln Z$ scales as $\sqrt{H/n_\mathrm{live}}$ (with $H$ the information, below); the cost grows proportionally. |

| `dlogz_if_ns` | `0.1` | Stop when the remaining live points could raise $\ln Z$ by less than this (same meaning as in *dynesty* and *juliet*). Must be in $(0, 10]$. |

| `num_samples_if_ns` | `4000` | Equally weighted posterior draws kept for the reports. |

| `inner_steps_per_dim_if_ns` | `4` | Slice-sampling steps per new point, per free parameter. Too few bias $\ln Z$ upwards: on a strongly correlated 20-parameter test, with the default adaptive bracket the bias was +0.7 nats with 2 steps (8 seeds), +0.21 ± 0.07 with 3 and +0.09 ± 0.06 with 4 (24 seeds each). The default is 4, which favours an unbiased evidence over speed: it costs 1.34 times the evaluations of 3 (7.9 against 5.9 million on that test). Use 3 for faster runs whose evidence differences are well above 0.2 nats; with the stepping-out bracket +1.1 with 1 step and +0.3 with 2 or 3. On K2-140 (20 parameters) every setting tried gave the same $\ln Z$ and posteriors; fewer steps are a reasonable choice when you need posteriors rather than precise evidences: `inner_steps_per_dim_if_ns: 1` halves the cost of K2-140 (2.6 against 5.2 million evaluations) with the same posterior (medians within 0.09σ, widths within 5 %), but do not compare its evidences between models. |

| `inner_steps_per_dim_if_ns` | `4` | Slice-sampling steps per new point, per free parameter. Too few bias $\ln Z$ upwards: on a strongly correlated 20-parameter test with the default adaptive bracket the bias is +0.55 ± 0.07 nats with 2 steps, +0.21 ± 0.07 with 3 and +0.09 ± 0.06 with 4 (24 seeds each). The default is 4, which favours an unbiased evidence over speed: it costs 1.34 times the evaluations of 3. On strongly curved posteriors (e.g. the impact parameter–stellar density degeneracy of a grazing transit) the bias is larger and falls with more steps: +3.9, +2.1 and +1.2 nats with 4, 8 and 12 steps on a curved 20-parameter test, and on K2-140 ln Z drops by 0.18 nats from 4 to 8 steps. **Check for curvature before a model comparison of a few nats:** rerun with `8` and over a few seeds; if ln Z drops by more than its reported error, or the scatter over seeds exceeds it, use the 8-step value and repeated runs. Fewer steps are a reasonable choice when you need posteriors rather than precise evidences: with the stepping-out bracket, 1 step instead of 2 halves the cost of K2-140 (2.6 against 5.2 million evaluations) with the same posterior (medians within 0.17σ), but do not compare its evidences between models. |

| `slice_bracket_if_ns` | `"adaptive"` | How each slice step brackets the likelihood contour along its direction. `"adaptive"`: a wide bracket, shrunk until a point inside the contour is found, whose width adapts during the run so that 30 % of the first proposals succeed: ≈2.7 likelihood evaluations per step. `"stepping_out"` (the only option before October 2026): a bracket of two standard deviations of the live points, widened until both ends are outside the contour, then shrunk: ≈5 evaluations per step, but each step moves further, so 2 steps per parameter suffice. With 3 adaptive steps the evidence is as accurate as with 2 stepping-out steps, for 1.3 times fewer evaluations. |

| `slice_loop_if_ns` | `"fused"` | How the slice steps of each chain are looped (adaptive bracket only). `"fused"`: one loop per chain over all its steps, so that a chain that has finished a step goes on with the next instead of waiting for the slowest chain of the batch; same algorithm, about twice as fast on a GPU. `"stepwise"`: BlackJAX's loop, one batched step at a time (the only option before October 2026; reproduces earlier runs exactly). |

| `space_if_ns` | `"constrained"` | Coordinates of the slice moves: `"constrained"` (the parameters themselves, the prior support a hard wall) or `"unconstrained"` (NumPyro's transformed space, as used by MCMC). The transform bends thin, correlated posterior ridges, which biased $\ln Z$ by ~1 nat on the 20-parameter test; keep the default unless you have a reason. |

k2140_pyorbit/runs/*/k2140_pyorbit/

k2140_pyorbit/runs/*/k2140_*.dat

k2140_pyorbit/runs/*/run_pyorbit.py

# Full nested-sampling runs of the synthetic light curves: their data files are copies of data/ (large)

scaling/runs/*/lc_*.dat

"""Injection and recovery on K2-140-like data: is the posterior of the full joint model calibrated?

Referee report of 2026-10-11 (M6.1). N synthetic datasets with the times, cadences, exposure times and uncertainties of

the K2-140 benchmark (k2140_juliet/runs/k2c_e8_rtx). The true parameters of each dataset are a random draw of the

posterior of that run (all 20 free parameters, jitters included), and the observations are drawn from the likelihood

of PyAETNA's model at those parameters (the distribution of every observed site of the NumPyro model, so the noise

model is exactly the fitted one). Each dataset is then fitted with the benchmark configuration and its default nested

sampling (summarise.py checks how often the 68 % and 95 % intervals contain the truth).

usage (PyAETNA's environment): python make_data.py [N, default 50]   (writes runs/inj_<k>/ and runs/inj_<k>/truth.json)

"""

from __future__ import annotations

import json

import os

import re

import sys

import tempfile

from pathlib import Path

import jax

import numpy as np

from numpyro import handlers

HERE = Path(__file__).resolve().parent

SRC = HERE.parent / "k2140_juliet" / "runs" / "k2c_e8_rtx"

INSTRUMENT_CODE = {

    "K2": 1,

    "LCOGT": 2,

    "CORALIE": 3,

    "FIES": 4,

    "HARPS": 5,

}  # column INSTRUMENT of k2140c.dat

def main(n: int) -> None:

    from PyAETNA.sampling.engine.sampling_utils import initialize_sampling

    tmp = Path(tempfile.mkdtemp())

    cwd = os.getcwd()

    os.chdir(SRC)

    ctx = initialize_sampling(

        "k2140.yaml", "k2140c.dat", str(tmp / "x.log"), outDir=str(tmp)

    )

    os.chdir(cwd)

    post = np.load(SRC / "PyAETNAouts" / "k2140_samples.npz")

    names = [k for k in post.files if k != "potential_energy"]

    draws = {k: np.asarray(post[k]).ravel() for k in names}

    lines = (SRC / "k2140c.dat").read_text().splitlines()

    header = [l for l in lines if l.startswith("#") or l.startswith("BJD")]

    rows = [

        l.split() for l in lines if l and not (l.startswith("#") or l.startswith("BJD"))

    ]

    instr = np.array([int(r[5]) for r in rows])

    cfg = (SRC / "k2140.yaml").read_text()

    rng = np.random.default_rng(20261011)

    picks = rng.choice(len(draws[names[0]]), size=n, replace=False)

    for k, i in enumerate(picks):

        theta = {p: float(draws[p][i]) for p in names}

        tr = handlers.trace(

            handlers.seed(

                handlers.substitute(ctx["model_fn"], theta),

                jax.random.PRNGKey(1000 + k),

            )

        ).get_trace(ctx["plData"], ctx["plModel"])

        y = np.array([float(r[1]) for r in rows])

        for site, code in INSTRUMENT_CODE.items():

            s = tr[f"obs_{site}"]

            new = np.asarray(

                s["fn"].sample(jax.random.PRNGKey(2000 + 10 * k + code)), dtype=float

            )

            assert new.shape == (int(np.sum(instr == code)),), site

            y[instr == code] = new

        run = HERE / "runs" / f"inj_{k:03d}"

        run.mkdir(parents=True, exist_ok=True)

        out = header + [

            " ".join([r[0], f"{y[j]:.10e}", *r[2:]]) for j, r in enumerate(rows)

        ]

        (run / "k2140c.dat").write_text("\n".join(out) + "\n")

        (run / "k2140.yaml").write_text(

            re.sub(r"random_seed: *\d+", "random_seed: 1", cfg)

        )

        (run / "truth.json").write_text(

            json.dumps({"posterior_draw": int(i), **theta}, indent=1)

        )

    print(f"{n} datasets in {HERE / 'runs'}")

if __name__ == "__main__":

    main(int(sys.argv[1]) if len(sys.argv) > 1 else 50)