fix: ESS walkers returned shuffled, pilot drawing a block, quadratic run...

#!/bin/bash

# Regenerate the outputs of the example configurations (examples/*.yaml) on the server, under heavy.lock, in a copy

# of examples/ inside the refactor/samplers clone (sync_refactor_server.sh). Copy examples_run/PyAETNAouts back to

# examples/PyAETNAouts on the laptop, then run docs/outputs/make_images.py there.

# usage: examples_server.sh <work dir> [clone folder]

set -u

W=$(realpath "$1"); C=${2:-PyAETNA_refactor}

LOG=$W/logs/examples_server.log

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

export PATH=$W/envs/pyaetna/bin:$PATH PYTHONPATH=$W/$C/src

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

E=$W/$C/examples_run

rm -rf "$E"; mkdir -p "$E/PyAETNAouts"

cp "$W/$C"/examples/*.yaml "$W/$C"/examples/*.dat "$E/"

cp "$W/$C"/examples/PyAETNAouts/*.dat "$W/$C"/examples/PyAETNAouts/README.md "$E/PyAETNAouts/"

cd "$E" || exit 1

say "start examples, clone $C"

for cfg in *.yaml; do

  JAX_PLATFORMS=cuda,cpu nice -n 19 pyaetna "$cfg" > "${cfg%.yaml}.out" 2>&1

  say "end $cfg exit=$?"

done

say "done"

# Copy the working tree of the laptop's repository (the refactor/samplers branch, committed or not) to the server

# clone used for tests and benchmarks (GitHub and GitLab clones fail on the server). Run outputs, checkpoints and

# node_modules of the configurator test are synced too, except the large posterior chains. What the server produces

# itself (run folders, results) is protected from deletion: an earlier version of this script, with --delete, wiped

# the outputs of a finished queue (2026-10-06).

# itself (run folders, results, examples_run/) is protected from deletion: an earlier version of this script, with

# --delete, wiped the outputs of a finished queue (2026-10-06).

# usage: sync_refactor_server.sh [server clone, default /DATA/ExoREC/pyaetna_bench/PyAETNA_refactor]

set -u

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

D=${1:-/DATA/ExoREC/pyaetna_bench/PyAETNA_refactor}

rsync -a --delete --filter='P /paper_benchmarks/*/runs/**' --filter='P /paper_benchmarks/*/results/**' \

  --filter='P /paper_benchmarks/*/runs/*/**' --exclude .git --exclude '__pycache__' --exclude '*.egg-info' --exclude 'paper/' \

  --filter='P /paper_benchmarks/*/runs/*/**' --filter='P /examples_run' --filter='P /examples_run/**' --exclude .git --exclude '__pycache__' --exclude '*.egg-info' --exclude 'paper/' \

  --exclude 'juliet_out_*' --exclude '*.pckl' --exclude 'prompt.md' \

  "$R/" "ganymede-cuda13:$D/" 2>&1 | grep -v 'bind\|channel\|forward'

echo "synced $(git -C "$R" rev-parse --short HEAD) + working tree → $D"

  archivePrefix = {arXiv}

}""")

                # ── NUTS and the ensemble samplers ────────────────────────

                if _ps.lower() == "nuts":

                    _log(_sep)

                    _log("No-U-Turn Sampler:")

                    _log(_sep)

                    _log_bib("""

@ARTICLE{Hoffman2014,

  author  = {{Hoffman}, Matthew D. and {Gelman}, Andrew},

  title   = {{The No-U-Turn Sampler: Adaptively Setting Path Lengths

              in Hamiltonian Monte Carlo}},

  journal = {Journal of Machine Learning Research},

  year    = 2014,

  volume  = {15},

  pages   = {1593--1623}

}""")

                elif _ps.lower() == "ess":

                    _log(_sep)

                    _log("Ensemble slice sampling:")

                    _log(_sep)

                    _log_bib("""

@ARTICLE{Karamanis2020,

  author        = {{Karamanis}, Minas and {Beutler}, Florian},

  title         = {{Ensemble Slice Sampling}},

  year          = {2020},

  eprint        = {2002.06212},

  archivePrefix = {arXiv}

}""")

                    _log_bib("""

@ARTICLE{Karamanis2021,

  author  = {{Karamanis}, Minas and {Beutler}, Florian and {Peacock}, John A.},

  title   = {{Preparing for precision cosmology with fast, robust and efficient

              Bayesian inference using the ensemble slice sampling method}},

  journal = {\\mnras},

  year    = {2021},

  volume  = {516},

  pages   = {1644--1653},

  doi     = {10.1093/mnras/stac2272},

  eprint  = {2111.02702},

  archivePrefix = {arXiv}

}""")

                elif _ps.lower() == "aies":

                    _log(_sep)

                    _log("Affine-invariant ensemble sampler:")

                    _log(_sep)

                    _log_bib("""

@ARTICLE{Goodman2010,

  author  = {{Goodman}, Jonathan and {Weare}, Jonathan},

  title   = {{Ensemble samplers with affine invariance}},

  journal = {Communications in Applied Mathematics and Computational Science},

  year    = {2010},

  volume  = {5},

  number  = {1},

  pages   = {65--80},

  doi     = {10.2140/camcos.2010.5.65}

}""")

                # ── Physics ───────────────────────────────────────────────

                _phy = config_check.get("physics", {})

                if isinstance(_phy, dict) and _phy:

def _free(samples: dict[str, Any], names: Iterable[str]) -> dict[str, np.ndarray]:

    return {

        k: np.asarray(samples[k], dtype=float)

        k: np.asarray(samples[k])

        for k in names

        if k in samples and np.asarray(samples[k]).ndim == 2

    }

    """Smallest effective sample size and largest split R-hat over the free parameters ``names``."""

    from numpyro.diagnostics import effective_sample_size, split_gelman_rubin

    chains = _free(samples, names)

    chains = {k: np.asarray(v, dtype=float) for k, v in _free(samples, names).items()}

    ess = min(float(effective_sample_size(v)) for v in chains.values())

    rhat = max(float(split_gelman_rubin(v)) for v in chains.values())

    return ess, rhat

def ensemble_status(samples: dict[str, Any], names: Iterable[str]) -> dict[str, float]:

    """Autocorrelation time, burn-in, effective sample size and R-hat over the walkers of an ensemble run."""

def ensemble_status(

    samples: dict[str, Any], names: Iterable[str], tau_hint: float | None = None

) -> dict[str, float]:

    """Autocorrelation time, burn-in, effective sample size and R-hat over the walkers of an ensemble run.

    ``tau_hint``, an earlier estimate of ``tau`` in the same run, makes the estimate cheaper: the chains are thinned

    by up to ``tau_hint / 4`` (keeping at least 1000 steps), which leaves ``tau`` unchanged within its noise; on a

    long run (148 walkers, 50 000 steps, 37 parameters) the full estimate took 11 s, more than the batch it tested.

    """

    from numpyro.diagnostics import split_gelman_rubin

    chains = _free(samples, names)

    n_walkers, n_steps = next(iter(chains.values())).shape

    tau = max(autocorr_time(v) for v in chains.values())

    thin = 1

    if tau_hint is not None and np.isfinite(tau_hint):

        thin = max(1, min(int(tau_hint / 4), n_steps // 1000))

    tau = max(autocorr_time(v[:, ::thin]) * thin for v in chains.values())

    burn = min(int(np.ceil(BURN_TAU * tau)), n_steps - 1) if np.isfinite(tau) else 0

    kept = n_steps - burn

    ess = n_walkers * kept / tau if np.isfinite(tau) and tau > 0 else 0.0

    rhat = (

        max(float(split_gelman_rubin(v[:, burn:])) for v in chains.values())

        if kept >= 4

        max(

            float(split_gelman_rubin(np.asarray(v[:, burn::thin], dtype=float)))

            for v in chains.values()

        )

        if kept >= 4 * thin

        else np.inf

    )

    return {

) -> Callable[[dict[str, Any]], bool]:

    """The ``stop_when`` test of :func:`PyAETNA.sampling.engine.sampler_utils.run_batched`."""

    names = list(names)

    last_tau: list[float] = []

    def test(samples: dict[str, Any]) -> bool:

        if sampler == "nuts":

                f"  convergence: smallest ESS {ess:.0f} (target {target_ess}), largest R-hat {rhat:.4f} (max {max_rhat})"

            )

            return ess >= target_ess and rhat <= max_rhat

        st = ensemble_status(samples, names)

        st = ensemble_status(samples, names, last_tau[-1] if last_tau else None)

        last_tau.append(st["tau"])

        log(

            f"  convergence: tau {st['tau']:.1f} steps, run {st['steps']:.0f} steps ({st['steps'] / st['tau']:.1f} tau, "

            f"need {TAU_FACTOR}), ESS {st['ess']:.0f} (target {target_ess}), R-hat {st['rhat']:.4f} (max {max_rhat})"

            (z_flat, log_density, inner_state._replace(log_density=log_density)),

        )

        if self._randomize_split:

            # keep the cache aligned with the walkers' new order

            # each walker back in its place (see BoundedESS.sample), with its cached log density

            inverse = jnp.argsort(perm)

            z_flat, log_density = z_flat[inverse], log_density[inverse]

            inner_state = inner_state._replace(log_density=log_density)

        return EnsembleSamplerState(unravel_fn(z_flat), inner_state, rng_key)

        self._batch_log_density = bounded_log_density

        return init_params

    def sample(self, state, model_args, model_kwargs):  # type: ignore[no-untyped-def]

        """``numpyro.infer.ensemble.EnsembleSampler.sample``, with every walker returned in its place.

        NumPyro's ESS (``randomize_split=True``) shuffles the walkers to split the ensemble at random, and returns

        them shuffled: a "chain" of the output was then a different walker at every step, so its autocorrelation

        time, effective sample size and R-hat said nothing (the autocorrelation time came out as 1-3 steps whatever

        the model). The split stays random; the walkers are put back in their order after the update.

        """

        z, inner_state, rng_key = state

        rng_key, _ = jax.random.split(rng_key)

        z_flat, unravel_fn = batch_ravel_pytree(z)

        perm = None

        if self._randomize_split:

            perm = jax.random.permutation(rng_key, z_flat.shape[0])

            z_flat = z_flat[perm]

        half = self._num_chains // 2

        def body_fn(k, carry):  # type: ignore[no-untyped-def]

            z_flat, inner = carry

            active, inactive = jax.lax.cond(

                k == 0,

                lambda x: (x[:half], x[half:]),

                lambda x: (x[half:], x[:half]),

                z_flat,

            )

            z_updates, inner = self.update_active_chains(active, inactive, inner)

            z_flat = jax.lax.cond(

                k == 0,

                lambda x: x.at[:half].set(z_updates),

                lambda x: x.at[half:].set(z_updates),

                z_flat,

            )

            return z_flat, inner

        z_flat, inner_state = jax.lax.fori_loop(0, 2, body_fn, (z_flat, inner_state))

        if perm is not None:

            z_flat = z_flat[jnp.argsort(perm)]

        return EnsembleSamplerState(unravel_fn(z_flat), inner_state, rng_key)