chore: WASP-18 runs with binned errors, NUTS rerun queues after the lineage...

# (results/mcmc_efficiency_*.json, results/warmup_test.json, results/summary.json) are; rerun to regenerate them.

k2100_pyaneti/runs/aies_stretch*/PyAETNAouts/*_samples.npz

wasp18_phase_curve/runs/aies_de*/PyAETNAouts/*_samples.npz

*/runs/warmup_test/*/PyAETNAouts/*_samples.npz

*/runs/warmup_test/*/*.dat

*/runs/warmup_test*/*/PyAETNAouts/*_samples.npz

*/runs/warmup_test*/*/*.dat

# Posterior chains and checkpoints of the TTV-scaling NUTS runs (up to 33 MB each); summaries are kept.

ttv_scaling/runs/nuts*/PyAETNAouts/*_samples.npz

ttv_scaling/runs/*/PyAETNAouts/*.pckl

#!/bin/bash

# 2026-10-05 evening, laptop (the machine of the paper's tables): (1) WASP-18 with the binned errors of Sect. 7.4 for

# the timings and Table `efficiency` (run folders from wasp18_phase_curve/make_err_runs.py); (2) every NUTS run of the

# laptop tables again after [code 49] (NUTS starts one per lineage): K2-140 and the TTV benchmark, frozen and adapted

# mass. The previous run folders are archived first. One job at a time; each run refuses to start with capped clocks,

# outside the performance profile or off mains, and records them in power_profile.txt.

set -u

B=$(realpath "$(dirname "$0")/..")

LOG=/dati/software/pyLib/PyAETNA_bench_work/logs/lineage_err_laptop.log

ARCH=/dati/software/pyLib/PyAETNA_bench_work/archive_2026-10-05_pre_lineage

mkdir -p "$(dirname "$LOG")" "$ARCH"

export PATH=/dati/pyenvs/pyaetnaEnv/bin:$PATH

exec 9> /dati/software/pyLib/PyAETNA_bench_work/heavy.lock; flock 9

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

check() {

  local f p ac

  f=$(cat /sys/devices/system/cpu/cpu*/cpufreq/scaling_max_freq | sort -n | tail -1); p=$(powerprofilesctl get)

  ac=$(cat /sys/class/power_supply/A*/online 2>/dev/null | head -1)

  if [ "$f" -lt 5000000 ] || [ "$p" != performance ] || [ "$ac" != 1 ]; then

    say "REFUSED: max_freq=$f profile=$p ac=$ac"; exit 1

  fi

  echo "profile=$p max_freq_khz=$f ac=$ac $(date '+%F %T')" > "$B/$1/power_profile.txt"

}

run() {  # folder platforms

  if [ -d "$B/$1/PyAETNAouts" ] && [ ! -d "$ARCH/$1" ]; then mkdir -p "$ARCH/$(dirname "$1")"; cp -a "$B/$1" "$ARCH/$1"; fi

  check "$1"

  say "start $1 ($2)"

  (cd "$B" && taskset -c 0-23 nice -n 19 ./run.sh "$1" "$2" 24G > /dev/null)

  check "$1"

  say "end $1 $(cat "$B/$1/done.txt") $(grep 'Elapsed (wall' "$B/$1/time.out")"

}

say "start queue, PyAETNA $(cd "$B/.." && git rev-parse --short HEAD) dirty=$(cd "$B/.." && git status --porcelain src | wc -l)"

w=wasp18_phase_curve/runs

run $w/ns_err_cpu24 cpu

run $w/ns_err_gpu_rtx cuda,cpu

run $w/aies_de_init_err_cpu24 cpu

run $w/aies_de_init_err_rtx cuda,cpu

run $w/nuts_auto_err_rtx cuda,cpu

run $w/nuts_auto_err_cpu24 cpu

run $w/nuts_adiag_err_rtx cuda,cpu

run $w/nuts_adiag_err_cpu24 cpu

run k2140_juliet/runs/nuts_rtx cuda,cpu

run k2140_juliet/runs/nuts_adiag_rtx cuda,cpu

for n in 12 25 50 100; do

  run ttv_scaling/runs/nuts_${n}_rtx cuda,cpu

  run ttv_scaling/runs/nuts_${n}_adiag_rtx cuda,cpu

done

# juliet's K2-140 fit with dynesty's slice proposals on the laptop (20 threads, its best count), five runs, so that

# Sect. 9.1 can quote the laptop instead of the server (cpu_scaling/juliet_rslice_laptop)

J=$B/cpu_scaling/juliet_rslice_laptop; mkdir -p "$J"; cp "$B"/k2140_juliet/data/K2-140_*.dat "$B/k2140_juliet/run_juliet.py" "$J/"

for i in 1 2 3 4 5; do

  [ -f "$J/juliet_20_rslice_$i.json" ] && continue

  check cpu_scaling/juliet_rslice_laptop; say "start juliet rslice $i"

  (cd "$J" && env OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 taskset -c 0-19 nice -n 19 /usr/bin/time -v \

     /dati/software/pyLib/PyAETNA_bench_work/envs/juliet/bin/python run_juliet.py 20 rslice $i > juliet_rslice_$i.out 2> juliet_rslice_$i.time)

  say "end juliet rslice $i $(grep 'Elapsed (wall' "$J/juliet_rslice_$i.time")"

done

say "done"

#!/bin/bash

# 2026-10-05 evening, server (H100 and 60 CPU threads): the server side of lineage_err_laptop.sh, for the rescaling

# of Sect. 8.5, the server rows of efficiency.json, the TOI-216 NUTS run of Sect. 8.3 and the warmup test of App. B.2:

# (1) WASP-18 with the binned errors (wasp18_phase_curve/make_err_runs.py); (2) every server NUTS run again after

# [code 49] (NUTS starts one per lineage). The TOI-216 frozen-mass run is extended as in round5_toi216_server.sh (3000

# draws per chain at a time, until R-hat < 1.01 and ESS > 400). heavy.lock; source copied from the laptop.

# usage: lineage_err_server.sh <work dir>

set -u

W=$(realpath "$1"); B=$W/PyAETNA/paper_benchmarks; LOG=$W/logs/lineage_err_server.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"; }

say "PyAETNA src sha256 $(cd "$B/../src" && find . -name '*.py' -print0 | sort -z | xargs -0 cat | sha256sum | cut -c1-16)"

run() {  # folder platforms

  say "start $1 ($2)"

  if [ "$2" = cpu ]; then (cd "$B" && taskset -c 0-59 nice -n 19 ./run.sh "$1" cpu 64G > /dev/null)

  else (cd "$B" && nice -n 19 ./run.sh "$1" cuda,cpu 80G > /dev/null); fi

  say "end $1 $(cat "$B/$1/done.txt") $(grep 'Elapsed (wall' "$B/$1/time.out")"

}

w=wasp18_phase_curve/runs

run $w/ns_err_gpu_h100 cuda

run $w/ns_err_cpu_epyc60 cpu

run $w/aies_de_init_err_h100 cuda

run $w/aies_de_init_err_epyc60 cpu

run $w/nuts_auto_err_h100 cuda

run $w/nuts_auto_err_epyc60 cpu

run $w/nuts_adiag_err_h100 cuda

for v in warmup_s1 direct_s1 warmup_s2 direct_s2 warmup_s3 direct_s3; do run $w/warmup_test_err/$v cpu; done

run k2140_juliet/runs/nuts_h100 cuda

run k2140_juliet/runs/nuts_adiag_h100 cuda

for n in 12 25 50 100; do

  run ttv_scaling/runs/nuts_${n}_h100 cuda

  run ttv_scaling/runs/nuts_${n}_adiag_h100 cuda

done

run ttv_scaling/runs/nuts_100_long_h100 cuda

run toi216_ttv/runs/branch2_adiag_h100 cuda

# TOI-216, frozen mass: the initial 1500 draws per chain, then extensions of the checkpoint

R=toi216_ttv/runs/branch2_h100

(cd "$B/$R" && cp toi216.yaml.initial toi216.yaml && rm -f run_*.out time_*.out)

run $R cuda

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

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

n=1500

while [ $n -lt 18000 ]; do

  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

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

done

say "done"

def wasp18() -> tuple[dict, dict]:

    comp = load("wasp18_phase_curve")["pyaetna"]["ns_cpu"]["posterior_vs_allesfitter"]

    # the run with allesfitter's binned errors; its white noise is a jitter added to them, so it has no

    # counterpart in allesfitter and is not shown

    comp = load("wasp18_phase_curve")["pyaetna"]["ns_err_cpu24"][

        "posterior_vs_allesfitter"

    ]

    labels = {

        "b_rr": r"$R_p/R_\star$",

        "b_rsuma": r"$(R_\star+R_p)/a$",

        "nightside_flux_ppt": r"$F_{\rm night}$",

        "b_phase_curve_beaming_TESS": r"$A_{\rm beam}$",

        "b_phase_curve_ellipsoidal_TESS": r"$\cos 2\phi$ amplitude",

        "log_err_flux_TESS": r"$\ln\sigma_{\rm TESS}$",

    }

    return comp, labels

    print(f"wrote {out.relative_to(PAPER)}")

def toi216() -> None:

    """Timing uncertainty of every transit: PyAETNA (both branches) and allesfitter (published run and reruns)."""

    d = load("toi216_ttv")

    ns = d["pyaetna"]["ns_h100"]["posterior_vs_allesfitter"]

    nuts = d["pyaetna"]["branch2_h100"]["posterior_vs_allesfitter"]

    reruns = [r["posterior"] for r in d["allesfitter_reruns"]["runs"].values()]

    fig, axes = plt.subplots(2, 1, figsize=(COLUMN_WIDTH, 2.9), sharex=True)

    for ax, planet in zip(axes, ("b", "c")):

        names = sorted((n for n in d["ttv_epochs"] if n.startswith(f"{planet}_")),

                       key=lambda n: d["ttv_epochs"][n])  # fmt: skip

        period = {"b": 17.0714, "c": 34.5555}[planet]

        t = np.array([d["ttv_epochs"][n] for n in names], float) * period

        for k, post in enumerate(reruns):

            ax.plot(t, [post[n]["sigma"] * 1440 for n in names], color=ORANGE, lw=0.6, alpha=0.5,

                    label="allesfitter, reruns" if k == 0 else None)  # fmt: skip

        series = (

            (

                [ns[n]["reference_sigma"] for n in names],

                ORANGE,

                "s",

                "allesfitter, published",

            ),

            (

                [ns[n]["pyaetna_sigma"] for n in names],

                BLUE,

                "o",

                "PyAETNA, nested sampling",

            ),

            ([nuts[n]["pyaetna_sigma"] for n in names], AQUA, "^", "PyAETNA, NUTS"),

        )

        for sig, color, marker, label in series:

            ax.plot(

                t,

                np.array(sig) * 1440,

                marker,

                color=color,

                ms=3.5,

                mec="white",

                mew=0.6,

                label=label,

            )

        ax.grid(axis="y", color=GRID, lw=0.5)

        ax.set_ylabel(f"$\\sigma_{{T}}$, TOI-216{planet} [min]")

    axes[0].legend(

        frameon=False,

        loc="upper left",

        fontsize=5.5,

        ncol=2,

        columnspacing=0.8,

        handletextpad=0.3,

    )

    axes[0].set_ylim(top=axes[0].get_ylim()[1] * 1.25)

    axes[1].set_xlabel("time from reference transit [d]")

    fig.tight_layout(h_pad=0.4)

    out = FIG / "benchmarks" / "toi216_timing.pdf"

    out.parent.mkdir(parents=True, exist_ok=True)

    fig.savefig(out)

    plt.close(fig)

    print(f"wrote {out.relative_to(PAPER)}")

def scaling() -> None:

    d = load("scaling")

    rows = d["rows"]

            "o",

            "-",

        ),

        (

            lambda r: r["code"] == "pyaetna" and r["device"] == "gpu (H100)",

            "PyAETNA, H100 GPU",

            VIOLET,

            "D",

            "-",

        ),

        (

            lambda r: r["code"] == "pyaetna" and r["device"] == "cpu",

            "PyAETNA, CPU",

            "--",

        ),

    ]

    fig, axs = plt.subplots(2, 3, figsize=(2 * COLUMN_WIDTH, 3.5), sharey="row", sharex="col",

                            gridspec_kw={"height_ratios": [2, 1]})  # fmt: skip

    axes, rax = axs

    fig, axes = plt.subplots(1, 3, figsize=(2 * COLUMN_WIDTH, 2.4), sharey=True)

    run_hours = {"1 h": 1.0, "1 day": 24.0}

    for ax, ra, (model, title) in zip(axes, rax, models):

    for ax, (model, title) in zip(axes, models):

        ref = sorted(

            (r for r in rows if r["model"] == model and r["code"] == "juliet stack"),

            key=lambda r: r["n"],

        ):  # the reference: one thread (lower edge) to 16 threads with perfect scaling (the line)

            ax.fill_between([r["n"] for r in ref], [r["evals_per_s_1thread"] for r in ref],

                            [r["evals_per_s"] for r in ref], color=AQUA, alpha=0.15, lw=0)  # fmt: skip

            refrate = {r["n"]: r["evals_per_s"] for r in ref}

            for pick, _label, color, marker, ls in series[:3]:

                sel = sorted((r for r in rows if r["model"] == model and pick(r) and r["n"] in refrate),

                             key=lambda r: r["n"])  # fmt: skip

                if sel:

                    ra.plot([r["n"] for r in sel], [r["evals_per_s"] / refrate[r["n"]] for r in sel], ls=ls, lw=1.0,

                            marker=marker, color=color, ms=3, mfc=color, mec="white", mew=0.8)  # fmt: skip

        ra.axhline(1, color=AQUA, lw=0.8, ls="--")

        ra.set_xscale("log")

        ra.set_yscale("log")

        ra.grid(color=GRID, lw=0.5)

        ra.set_xlabel("photometric points $N$")

        for pick, label, color, marker, ls in series:

            sel = sorted(

                (r for r in rows if r["model"] == model and pick(r)),

        ax.set_yscale("log")

        ax.set_title(title, fontsize=8, color=INK)

        ax.grid(color=GRID, lw=0.5)

        ax.set_xlabel("photometric points $N$")

    axes[0].set_ylabel("likelihood evaluations / s")

    rax[0].set_ylabel("ratio to reference")

    axes[0].legend(frameon=False, loc="lower left", fontsize=6)

    fig.tight_layout(w_pad=0.5)

    out = FIG / "performance" / "scaling.pdf"

    print(f"wrote {out.relative_to(PAPER)}")

def branches() -> None:

    """Flowchart of the two sampling branches (Sect. 4, Fig. branches)."""

    from matplotlib.patches import FancyBboxPatch

    fig, ax = plt.subplots(figsize=(COLUMN_WIDTH, 3.6))

    ax.set_xlim(0, 10)

    ax.set_ylim(0, 12.4)

    ax.axis("off")

    fs = 5.9

    def box(

        x: float,

        y: float,

        text: str,

        w: float = 4.5,

        h: float = 0.95,

        color: str = GRID,

        bold: bool = False,

        key: str = "",

    ) -> None:

        ax.add_patch(

            FancyBboxPatch(

                (x - w / 2, y - h / 2),

                w,

                h,

                boxstyle="round,pad=0.02,rounding_size=0.18",

                fc=color,

                ec="none",

            )

        )

        if key:

            ax.text(

                x,

                y + 0.2,

                text,

                ha="center",

                va="center",

                fontsize=fs,

                color=INK,

                fontweight="bold",

            )

            ax.text(

                x,

                y - 0.22,

                key,

                ha="center",

                va="center",

                fontsize=fs - 0.4,

                color=INK,

                family="monospace",

            )

        else:

            ax.text(

                x,

                y,

                text,

                ha="center",

                va="center",

                fontsize=fs,

                color=INK,

                fontweight="bold" if bold else "normal",

                linespacing=1.15,

            )

    def arrow(x0: float, y0: float, x1: float, y1: float) -> None:

        ax.annotate(

            "",

            xy=(x1, y1),

            xytext=(x0, y0),

            arrowprops={

                "arrowstyle": "-|>",

                "color": INK2,

                "lw": 0.9,

                "shrinkA": 0,

                "shrinkB": 0,

                "mutation_scale": 7,

            },

        )

    blue, aqua = "#d6e6f9", "#d1efe3"

    box(5.0, 11.8, "YAML configuration + data table", w=6.4, bold=True)

    box(2.4, 10.3, "Branch 1", color=blue, key="production_sampler: ns")

    box(7.6, 10.3, "Branch 2", color=aqua, key="production_sampler: auto")

    arrow(4.0, 11.3, 2.6, 10.8)

    arrow(6.0, 11.3, 7.4, 10.8)

    box(

        2.4,

        8.2,

        "nested slice sampling\nfrom the prior,\nvectorised over\nlive points",

        h=2.1,

        color=blue,

    )

    arrow(2.4, 9.8, 2.4, 9.25)

    box(7.6, 8.75, "Phase 1: ensemble warm-up\n(ESS or AIES)", color=aqua)

    box(7.6, 7.45, "re-seed stuck walkers;\nPhase 2: energy-spread test", color=aqua)

    arrow(7.6, 9.8, 7.6, 9.25)

    arrow(7.6, 8.25, 7.6, 7.95)

    box(6.45, 5.65, "Phase 3a\nNUTS", w=2.2, color=aqua)

    box(8.75, 5.65, "Phase 3b\nensemble", w=2.2, color=aqua)

    arrow(7.3, 6.95, 6.65, 6.15)

    arrow(7.9, 6.95, 8.55, 6.15)

    ax.text(

        6.9, 6.5, "unimodal", fontsize=fs - 0.6, color=INK2, ha="right", va="center"

    )

    ax.text(

        8.2, 6.5, "multimodal", fontsize=fs - 0.6, color=INK2, ha="left", va="center"

    )

    box(2.4, 5.65, "posterior + $\\ln\\mathcal{Z}$", color=blue)

    arrow(2.4, 7.15, 2.4, 6.15)

    box(7.6, 4.05, "posterior; AIC, BIC", color=aqua)

    arrow(6.45, 5.15, 7.2, 4.55)

    arrow(8.75, 5.15, 8.0, 4.55)

    box(5.0, 2.4, "report: fit, residuals, noise budget,\ndiagnostics, MAP", w=7.4)

    arrow(2.4, 5.15, 4.0, 2.9)

    arrow(7.6, 3.55, 6.0, 2.9)

    box(

        5.0,

        0.75,

        "optional physics stage: derived quantities,\ntemperatures, albedos",

        w=7.4,

    )

    arrow(5.0, 1.9, 5.0, 1.25)

    out = FIG / "inference" / "branches.pdf"

    out.parent.mkdir(parents=True, exist_ok=True)

    fig.savefig(out, bbox_inches="tight", pad_inches=0.02)

    plt.close(fig)

    print(f"wrote {out.relative_to(PAPER)}")

def eiv_validation() -> None:

    """Slope recovered in each realisation, EIV vs ordinary detrending (Sect. validation)."""

    d = load("eiv_validation")

def main() -> None:

    branches()

    ns_calibration()

    posterior_comparisons()

    if (BENCH / "toi216_ttv" / "results" / "summary.json").exists():

        toi216()

    if (BENCH / "eiv_validation" / "results" / "summary.json").exists():

        eiv_validation()

    if (BENCH / "scaling" / "results" / "summary.json").exists():

"""Create the binned-error twin of every WASP-18 run folder used by the paper.

Sect. 7.4 compares PyAETNA with allesfitter on the binned errors of allesfitter's data file

(``configs/wasp18_err.yaml``, data ``wasp18_err.dat``, written by ``make_data.py``). The timing, NUTS and warmup-test

runs were made on ``wasp18.dat`` (equal errors). For each of them this script writes ``wasp18_err.yaml`` into a new run folder

(``ns_err_<machine>``, ``<name>_err_<machine>``; the warmup test: ``runs/warmup_test_err/<name>/``): the same configuration with the data file and the jitter prior

of ``configs/wasp18_err.yaml``; for the runs started near the solution (``ensemble_init_values``), the starting values

are the posterior means of the nested-sampling run with the binned errors (``runs/ns_err_cpu24``).

usage: python make_err_runs.py

"""

from __future__ import annotations

import re

from pathlib import Path

import numpy as np

import yaml

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

RUNS = HERE / "runs"

SOURCES = [

    "aies_de_init_cpu24",

    "aies_de_init_rtx",

    "aies_de_init_epyc60",

    "aies_de_init_h100",

    "nuts_auto_cpu24",

    "nuts_auto_rtx",

    "nuts_auto_epyc60",

    "nuts_auto_h100",

    "nuts_adiag_cpu24",

    "nuts_adiag_rtx",

    "nuts_adiag_h100",

    "ns_gpu_h100",

    "ns_cpu_epyc60",

] + [f"warmup_test/{m}_s{s}" for m in ("direct", "warmup") for s in (1, 2, 3)]

def err_config(text: str, init: dict[str, float]) -> str:

    ref = yaml.safe_load((HERE / "configs" / "wasp18_err.yaml").read_text())

    lo, hi = ref["log_jitter"][2:4]

    text = text.replace("file_path: wasp18.dat", "file_path: wasp18_err.dat")

    text, n = re.subn(

        r"log_jitter:\n- global\n- uniform\n- [-\d.e]+\n- [-\d.e]+",

        f"log_jitter:\n- global\n- uniform\n- {lo}\n- {hi}",

        text,

    )

    assert n == 1, "jitter prior not found"

    if "ensemble_init_values:" in text:

        block = "".join(f"    {k}: {v}\n" for k, v in init.items())

        text = re.sub(

            r"  ensemble_init_values:\n(    .*\n)+",

            "  ensemble_init_values:\n" + block,

            text,

        )

        text = text.replace("# posterior means of PyAETNA's nested-sampling run on the same machine (runs/ns_cpu_epyc60).",

                            "# posterior means of PyAETNA's nested-sampling run with the binned errors (runs/ns_err_cpu24).")  # fmt: skip

    return text

def main() -> None:

    post = np.load(RUNS / "ns_err_cpu24" / "PyAETNAouts" / "wasp18_err_samples.npz")

    init = {k: float(np.mean(post[k])) for k in post.files if k != "potential_energy"}

    for name in SOURCES:

        src = RUNS / name / "wasp18.yaml"

        if not src.exists():

            print(f"skip {name}: no configuration")

            continue

        if name.startswith("warmup_test/"):

            dst = RUNS / name.replace("warmup_test/", "warmup_test_err/")

        elif name.startswith(

            "ns_"

        ):  # ns_gpu_h100 -> ns_err_gpu_h100, as the existing ns_err_* runs

            dst = RUNS / name.replace("ns_", "ns_err_", 1)

        else:  # aies_de_init_rtx -> aies_de_init_err_rtx: the machine suffix stays last

            base, machine = name.rsplit("_", 1)

            dst = RUNS / f"{base}_err_{machine}"

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

        (dst / "wasp18_err.yaml").write_text(err_config(src.read_text(), init))

        print(f"wrote {dst.relative_to(HERE)}/wasp18_err.yaml")