chore: benchmarks of the paper review: TTV scaling (NUTS against nested...

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

*/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

| `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/` |

| `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 |

"""PyAETNA's own model sampled by nautilus (Lange 2023) instead of the nested slice sampler of BlackJAX.

A prototype for the paper's discussion of nautilus: the PyAETNA pipeline runs as usual (configuration, data, model,

device and precision) up to Branch 1, where ``nested_sampling_production`` is replaced by a function that

- traces the NumPyro model and builds nautilus' unit-cube prior transform from the inverse CDFs of the priors;

- hands nautilus the log-likelihood of the same model, batched with ``jax.vmap`` and jitted, with the factors that

  restrict the prior (e.g. b < 1 + k) added to it; on a CPU with several host devices (PyAETNA gives each slice chain

  its own), the batch is spread over them;

- runs nautilus with PyORBIT's settings (500 live points, 4 networks, exploration discarded) and stops.

The reference is the PyAETNA run in the same folder (its results.json and time.out). One row per run is appended to

machines/results/nautilus_pyaetna.json.

usage (PyAETNA's environment with nautilus-sampler installed; JAX_PLATFORMS=cpu for a CPU run):

    python nautilus_pyaetna.py <benchmark>/runs/<PyAETNA run folder> <seed>

"""

from __future__ import annotations

import glob

import hashlib

import json

import os

import re

import shutil

import subprocess

import sys

import time

from pathlib import Path

import numpy as np

import PyAETNA._device_config  # noqa: F401  (as the pyaetna command: host devices and BLAS threads, before NumPy)

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

B = HERE.parent

OUT = B / "machines" / "results" / "nautilus_pyaetna.json"

WORK = Path(

    os.environ.get(

        "NAUTILUS_WORK", "/dati/software/pyLib/PyAETNA_bench_work/scratch_nautilus"

    )

)

T_START = time.time()

def wall_seconds(time_out: Path) -> float | None:

    if not time_out.exists():

        return None

    m = re.search(

        r"Elapsed \(wall clock\) time \(h:mm:ss or m:ss\): ([\d:.]+)",

        time_out.read_text(),

    )

    if not m:

        return None

    t = [float(x) for x in m.group(1).split(":")]

    return sum(v * 60**i for i, v in enumerate(reversed(t)))

def provenance() -> dict:

    git = ["git", "-C", str(B.parent)]

    head = subprocess.run(

        git + ["rev-parse", "--short", "HEAD"], capture_output=True, text=True

    ).stdout.strip()

    diff = subprocess.run(

        git + ["diff", "HEAD", "--", "src"], capture_output=True, text=True

    ).stdout

    prof = subprocess.run(

        ["powerprofilesctl", "get"], capture_output=True, text=True

    ).stdout.strip()

    return {

        "pyaetna_commit": head or "unknown",

        "src_diff_sha256": (

            hashlib.sha256(diff.encode()).hexdigest()[:16] if diff else None

        ),

        "power_profile": prof or None,

        "cpus_allowed": len(os.sched_getaffinity(0)),

        "date": time.strftime("%Y-%m-%d %H:%M"),

    }

def main(run_dir: Path, seed: int) -> None:

    cfg = next(run_dir.glob("*.yaml"))

    work = WORK / f"{run_dir.parent.parent.name}_{run_dir.name}_s{seed}"

    shutil.rmtree(work, ignore_errors=True)

    work.mkdir(parents=True)

    import yaml

    c = yaml.safe_load(cfg.read_text())

    data = c["data"]["file_path"]

    shutil.copy(

        (

            run_dir / data

            if (run_dir / data).exists()

            else run_dir.parent.parent / "data" / data

        ),

        work,

    )

    for extra in run_dir.glob("*.dat"):

        shutil.copy(extra, work)

    c["sampling"]["reset"] = True

    c.setdefault("reports", {})["run"] = False

    (work / cfg.name).write_text(yaml.safe_dump(c, sort_keys=False))

    res = glob.glob(str(run_dir / "PyAETNAouts" / "*_results.json"))

    ref = json.loads(Path(res[0]).read_text())["evidence"] if res else {}

    row: dict = {

        "benchmark": run_dir.parent.parent.name,

        "pyaetna_run": f"{run_dir.parent.parent.name}/runs/{run_dir.name}",

        "seed": seed,

        "pyaetna_log_z": ref.get("log_evidence"),

        "pyaetna_log_z_err": ref.get("log_evidence_err"),

        "pyaetna_sampling_s": ref.get("elapsed_s"),

        "pyaetna_evals": ref.get("n_likelihood_evals_approx"),

        "pyaetna_wall_s": wall_seconds(run_dir / "time.out"),

    }

    def nautilus_production(settings, config_file, model_fn, plData, plModel, *args, **kwargs):  # type: ignore

        import jax

        import jax.numpy as jnp

        import nautilus

        from jax.sharding import Mesh, NamedSharding, PartitionSpec

        from numpyro import handlers

        from PyAETNA.sampling.nested_sampling import split_prior_likelihood

        split = split_prior_likelihood(

            model_fn, (plData, plModel), jax.random.PRNGKey(seed), "constrained"

        )

        tr = handlers.trace(

            handlers.seed(model_fn, jax.random.PRNGKey(seed))

        ).get_trace(plData, plModel)

        sites = [

            (n, s["fn"])

            for n, s in tr.items()

            if s["type"] == "sample" and not s["is_observed"]

        ]

        shapes = [tuple(np.shape(tr[n]["value"])) for n, _ in sites]

        sizes = [int(np.prod(s)) for s in shapes]

        n_dim = sum(sizes)

        dtype = jnp.asarray(tr[sites[0][0]]["value"]).dtype

        def to_dict(u):  # type: ignore[no-untyped-def]

            out, i = {}, 0

            for (n, d), shp, k in zip(sites, shapes, sizes):

                out[n] = d.icdf(u[i : i + k].reshape(shp).astype(dtype))

                i += k

            return out

        def loglike(u):  # type: ignore[no-untyped-def]

            x = to_dict(u)

            restrict = split.logprior_fn(x) - sum(

                jnp.sum(d.log_prob(x[n])) for n, d in sites

            )

            return split.loglikelihood_fn(x) + restrict

        devices = jax.local_devices()

        n_dev = len(devices)

        if devices[0].platform == "cpu" and n_dev > 1:

            sharding = NamedSharding(

                Mesh(np.array(devices), ("b",)), PartitionSpec("b")

            )

            batched = jax.jit(jax.vmap(loglike), in_shardings=sharding)

        else:

            n_dev = 1

            batched = jax.jit(jax.vmap(loglike))

        def call(u: np.ndarray) -> np.ndarray:

            n = len(u)

            pad = (-n) % n_dev

            if pad:

                u = np.concatenate([u, np.repeat(u[-1:], pad, 0)])

            return np.asarray(batched(jnp.asarray(u, dtype)), float)[:n]

        t_setup = time.time()

        sampler = nautilus.Sampler(

            lambda u: u,

            call,

            n_dim=n_dim,

            n_live=500,

            n_networks=4,

            vectorized=True,

            seed=seed,

        )

        sampler.run(discard_exploration=True, verbose=False)

        points, log_w, log_l = (

            sampler.posterior()

        )  # weighted posterior, kept outside the repository

        values = jax.vmap(to_dict)(

            jnp.asarray(points, dtype)

        )  # unit cube -> parameters

        np.savez_compressed(

            work / "nautilus_posterior.npz",

            log_w=log_w,

            log_l=log_l,

            **{n: np.asarray(v) for n, v in values.items()},

        )

        row.update(

            {

                "device": devices[0].platform,

                "host_devices": n_dev,

                "precision": str(dtype),

                "n_dim": n_dim,

                "nautilus": nautilus.__version__,

                "jax": jax.__version__,

                "log_z": float(sampler.log_z),

                "evals": int(sampler.n_like),

                "sampling_s": time.time() - t_setup,

                "wall_s": time.time() - T_START,

                "threads_env": {

                    k: os.environ.get(k)

                    for k in ("OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS")

                },

                "provenance": provenance(),

            }

        )

        rows = json.loads(OUT.read_text()) if OUT.exists() else []

        OUT.write_text(json.dumps(rows + [row], indent=1))

        print(

            json.dumps({k: v for k, v in row.items() if k != "provenance"}), flush=True

        )

        raise SystemExit(0)

    import PyAETNA.sampling.nested_sampling as nsm

    nsm.nested_sampling_production = nautilus_production

    from PyAETNA.run import run

    os.chdir(work)

    run(str(work / cfg.name))

if __name__ == "__main__":

    main(Path(sys.argv[1]).resolve(), int(sys.argv[2]))

{"cpus": 1, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 29.252159556999686, "s_per_iteration": 27.093501113449747, "evals_per_s": 360.5514089558058}

{"cpus": 1, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 24.248103778998484, "s_per_iteration": 11.476807774249755, "evals_per_s": 851.1600256926477}

{"cpus": 2, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 15.78125661399099, "s_per_iteration": 14.356703935949918, "evals_per_s": 680.4208015698456}

{"cpus": 2, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 7.647485834997497, "s_per_iteration": 5.579512430849718, "evals_per_s": 1750.7981425022679}

{"cpus": 4, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 9.398072500000126, "s_per_iteration": 7.925213492149487, "evals_per_s": 1232.597709787443}

{"cpus": 4, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 4.781266692007193, "s_per_iteration": 2.871130487749906, "evals_per_s": 3402.3531990897477}

{"cpus": 8, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 9.387506993996794, "s_per_iteration": 7.631543645649799, "evals_per_s": 1280.0293693620406}

{"cpus": 8, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 3.6995045059884433, "s_per_iteration": 1.8384104499506067, "evals_per_s": 5313.612093677153}

{"cpus": 12, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 8.646050101000583, "s_per_iteration": 6.971191648949753, "evals_per_s": 1401.281228794175}

{"cpus": 12, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 3.5758241469884524, "s_per_iteration": 1.4065952017001109, "evals_per_s": 6944.855199415564}

{"cpus": 16, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 8.919945809000637, "s_per_iteration": 7.036325247599597, "evals_per_s": 1388.3098998774258}

{"cpus": 16, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 3.2676085549901472, "s_per_iteration": 1.3785763373998634, "evals_per_s": 7086.005856175207}

{"cpus": 24, "host_devices": 1, "chain_devices": 1, "iterations": 20, "compile_and_first_s": 9.927420876992983, "s_per_iteration": 8.454728329500359, "evals_per_s": 1155.4008147033255}

{"cpus": 24, "host_devices": 64, "chain_devices": 50, "iterations": 20, "compile_and_first_s": 3.2880595349997748, "s_per_iteration": 1.281011489299999, "evals_per_s": 7625.6927292182145}

      "scatter_over_reported": 0.6606854039355358

    }

  },

  "evaluations_per_new_point": {

    "dynesty_defaults": {

      "rwalk_walks": "ndim + 20",

      "rslice_slices": "3 + ndim"

    },

    "juliet_k2140_rwalk": {

      "ndim": 19,

      "range": [

        37.230731416131185,

        38.11827630424512

      ],

      "n_runs": 8

    },

    "juliet_k2140_rslice": {

      "ndim": 19,

      "slices": 22,

      "range": [

        118.01958459105937,

        119.09961919816065

      ],

      "per_slice_range": [

        5.36452657232088,

        5.413619054461848

      ],

      "n_runs": 5

    },

    "pyorbit_toi5398_rwalk": {

      "ndim": 12,

      "per_point": 29.82164875814691

    },

    "pyaetna_k2140": {

      "ndim": 20,

      "slice_steps_per_point": 80,

      "per_point": 213.9425468164794,

      "per_slice_step": 2.6742818352059925

    },

    "pyaetna_toi5398": {

      "ndim": 12,

      "slice_steps_per_point": 48,

      "per_point": 126.32841379310345,

      "per_slice_step": 2.6318419540229887

    }

  },

  "k2100_blas_threads": [

    {

      "cpus": 8,