feat!: three explicit samplers after a peak search; remove the auto warmup pipeline

Requirements: Python ≥ 3.10, JAX, NumPyro ≥ 0.22, BlackJAX ≥ 1.6.2,

jaxoplanet, tinygp; see `docs/dependencies.md`.

### Changed before the release: three explicit samplers and a peak search (2026-10-06)

- **No more `auto` pipeline.** `production_sampler` is mandatory: `ns`, `nuts`, `ess` or `aies`. The ensemble

  warmup, its energy-spread test, the restarting of stuck walkers and the hand-over to NUTS are gone, with their keys

  (`warmup_sampler`, `num_warmup`, `chain_energy_threshold`, `reseed_stuck_walkers`, `stuck_walker_threshold`,

  `ensemble_init_values`, `ensemble_init_radius`, `num_chains`). `internal_warmup_if_nuts` is now

  `num_adapt_if_nuts`. Configurations of 3.x must be updated; PyAETNA stops with the accepted keys.

- **Peak search before every sampler** (`sampling/peaks.py`, which also holds the MAP search of the reports): BFGS

  from `peak_search_starts` prior draws, grouped into distinct peaks with their Gaussian (Laplace) approximation,

  saved in `<config>_peaks.json`. NUTS chains and ensemble walkers start from draws of the best peak; NUTS starts with

  the peak's curvature as mass matrix and adapts it without NumPyro's regularisation, which spoilt parameters known

  far better than 1e-3 (the K2-140 period). `on_multiple_peaks` (`stop` or `continue`) is mandatory: with `stop`,

  several significant peaks give one configuration per peak, `<config>_peak<k>.yaml`, with `prior_bounds` isolating

  it, and the run stops (exit status 3). `peak_search: false` skips the search (start from a saved peak with

  `init_from_peaks`, or from the prior).

- **`sampling.prior_bounds`**: truncate priors to a box, keeping their shape.

- **Runs until converged** (`target_ess`, `max_rhat`) for NUTS and the ensembles.

- **Fixed**: the potential energy and the divergences of NUTS were never recorded during sampling (the batched loop

  looked for the kernel under the wrong attribute), so the MAP search of the reports started from chain endpoints.

### Added before the release (2026-10-05)

- **`sampling.adapt_mass_matrix_if_nuts`** (default `false`): lets NUTS adapt its mass matrix in the `auto` branch

### Features

- **JAX/NumPyro Backend**: JIT-compiled, GPU-accelerated MCMC sampling

- **Flexible MCMC Sampling**:

  - NUTS (gradient-based, default)

  - AIES and ESS ensemble samplers (gradient-free)

  - Automatic multimodal detection and restart strategies

- **Nested sampling with Bayesian evidence**: `production_sampler: ns` runs

  GPU-vectorised nested sampling instead of MCMC, returning the posterior and

  ln Z (Bayes factors for model comparison) in one run

- **Three samplers, chosen explicitly** (`production_sampler`):

  - `ns`: GPU-vectorised nested sampling, returning the posterior and ln Z (Bayes factors for model comparison)

  - `nuts`: the No-U-Turn Sampler (gradient-based)

  - `ess`, `aies`: ensemble slice sampling and the affine-invariant ensemble sampler (gradient-free)

- **Peak search before sampling**: optimisations from many prior draws find the peaks of the posterior; NUTS and

  the ensembles start from the best one with its curvature as scales, and several peaks (e.g. period aliases) are

  reported with one configuration per peak to run separately

- **Runs until converged**: NUTS and the ensembles can stop as soon as the slowest parameter has the requested number

  of independent draws (`target_ess`)

- **Named Priors & Global Tying**: tie parameters across different model components (e.g., sharing a stellar rotation period between multiple instrument and corresponding detrending schemes)

- **Comprehensive Light Curve Modeling**:

  - Primary transits, phase curves and secondary eclipses (via `jaxoplanet`)

- **GPU-accelerated end-to-end.** The entire forward model, likelihood, and

  gradient computation runs inside JIT-compiled JAX, delivering order-of-

  magnitude speedups on GPU/TPU hardware without code changes.

- **Adaptive three-phase sampling.** The pipeline manages warmup and

  production automatically, with ensemble warmup (ESS or AIES) for

  exploration and NUTS for efficient local convergence. Warmup and

  production behavior is still highly configurable in the YAML file (sampler

  choice, number of chains, adaptation schedule, etc.).

- **Transparent sampling.** You choose the sampler; a peak search before it finds where the posterior is, so NUTS

  and the ensembles start there without a warmup, and a posterior with several peaks is never sampled silently as

  if it had one.

- **Checkpointed, restartable runs.** The sampler state is saved

  periodically, so that long runs can survive crashes, and subsequent

  launches resume exactly where they left off — including the ensemble

  warmup phase.

  launches resume exactly where they left off.

- **Portable results.** Samples, diagnostics, the MAP, the evidence and the

  model components are written as plain `.npz` and `.json` files that open

  with NumPy alone and stay readable as PyAETNA changes.

<div class="step-section" data-step="4">

  <div class="section-intro">

    <h2>Sampling <span class="badge b-req">Required</span> <span class="badge b-g">Global</span></h2>

    <p>Configure the <code>sampling:</code> section. Two branches:

    <strong>Branch 1</strong>, <code>production_sampler: ns</code>: nested sampling, no warmup, the only route that gives the Bayesian evidence (ln&nbsp;Z) for model comparison.

    <strong>Branch 2</strong>, <code>auto</code> (default): <strong>Phase 1</strong> ensemble warmup, <strong>Phase 2</strong> convergence test, then

    <strong>Phase 3a</strong> NUTS production (if unimodal) or <strong>Phase 3b</strong> ensemble production (if multimodal); <code>nuts</code>, <code>ess</code> and <code>aies</code> skip the warmup. Branch 2 reports AIC/BIC only.

    The form shows only the keys that apply to the chosen mode, and only those are written.</p>

    <p>Configure the <code>sampling:</code> section. Three branches, chosen explicitly:

    <code>ns</code>, nested sampling, the only route that gives the Bayesian evidence (ln&nbsp;Z) for model comparison;

    <code>nuts</code>, the No-U-Turn Sampler; <code>ess</code> or <code>aies</code>, an ensemble sampler. Every branch

    starts with a <strong>peak search</strong>: optimisations from many prior draws find the peaks of the posterior;

    NUTS and the ensemble start from the best one, with the parameter scales of its curvature. The MCMC branches

    report AIC/BIC only. The form shows only the keys that apply to the chosen sampler, and only those are written.</p>

  </div>

  <h3 class="subsec">Production Sampler Mode</h3>

  <h3 class="subsec">Sampler</h3>

  <div class="fg">

    <label><span class="key-name">production_sampler</span></label>

    <select id="m-ps" data-key="production_sampler">

      <option value="auto" selected>auto &mdash; full pipeline (warmup &rarr; convergence test &rarr; auto NUTS/ensemble)</option>

      <option value="nuts">nuts &mdash; skip warmup, go directly to NUTS (init_to_median, full adaptation)</option>

      <option value="ess">ess &mdash; skip warmup, go directly to ESS ensemble production</option>

      <option value="aies">aies &mdash; skip warmup, go directly to AIES ensemble production</option>

      <option value="ns">ns &mdash; nested sampling instead of MCMC (also gives ln&nbsp;Z)</option>

      <option value="ns" selected>ns &mdash; nested sampling (posterior and ln&nbsp;Z)</option>

      <option value="nuts">nuts &mdash; No-U-Turn Sampler, from the best peak</option>

      <option value="ess">ess &mdash; ensemble slice sampling, from the best peak</option>

      <option value="aies">aies &mdash; affine-invariant ensemble sampler, from the best peak</option>

    </select>

    <div class="help"><strong><code>auto</code></strong> (default): Phase&nbsp;1 ensemble warmup &rarr; Phase&nbsp;2 convergence test &rarr; auto Phase&nbsp;3a NUTS or 3b ensemble.

    <strong><code>nuts</code></strong>: no warmup; NUTS starts at prior medians (<code>init_to_median</code>) with full step-size &amp; mass-matrix adaptation.

    <strong><code>ess</code></strong>/<strong><code>aies</code></strong>: no warmup; goes directly to ensemble production.

    <strong><code>ns</code></strong>: nested sampling, no warmup and no MCMC; the only route that gives the Bayesian evidence (ln&nbsp;Z) for model comparison, and it weights separate solutions correctly. The other routes report AIC/BIC only.</div>

    <div class="help"><strong><code>ns</code></strong>: starts from the prior; gives the posterior and the evidence. It keeps only one of several separate peaks: run one peak at a time (see <code>on_multiple_peaks</code>).

    <strong><code>nuts</code></strong>: gradient-based; chains start from draws around the best peak, with its curvature as mass matrix, then adapt.

    <strong><code>ess</code></strong>/<strong><code>aies</code></strong>: gradient-free ensembles; walkers start from draws around the best peak; the burn-in is discarded in the reports (<code>reports.discard_extra_warmup</code>).</div>

  </div>

  <div class="fg chk">

    <input type="checkbox" id="m-psearch" data-key="peak_search" data-emit="nondefault" data-default="true" checked>

    <label for="m-psearch"><span class="key-name">peak_search</span> &mdash; search the peaks of the posterior before sampling</label>

    <div class="help" style="margin-left:1.6rem">Uncheck when the posterior is known to have a single peak or its peaks were found by an earlier run: NUTS and the ensemble then start from a saved peak (<code>init_from_peaks</code>), or else from the prior, which needs a longer adaptation or burn-in. Default: <code>true</code>.</div>

  </div>

  <div class="fg" data-show-if="noSearch mcmc">

    <label><span class="key-name">init_from_peaks</span> <span class="badge b-opt">Optional</span></label>

    <input type="text" id="m-ifp" data-key="init_from_peaks" data-emit="set" placeholder="PyAETNAouts/run_peaks.json">

    <div class="help">The <code>&lt;config&gt;_peaks.json</code> of an earlier run: start from one of its peaks (<code>init_peak</code>, 1 = the best).</div>

  </div>

  <div class="fg" data-show-if="noSearch mcmc">

    <label><span class="key-name">init_peak</span> <span class="badge b-opt">Optional</span></label>

    <input type="number" id="m-ipk" data-key="init_peak" data-emit="set" placeholder="1" min="1" step="1">

    <div class="help">Rank of the saved peak to start from. Default: 1.</div>

  </div>

  <div class="fg" data-show-if="search">

    <label><span class="key-name">on_multiple_peaks</span></label>

    <select id="m-omp" data-key="on_multiple_peaks">

      <option value="stop" selected>stop &mdash; write one configuration per peak and stop</option>

      <option value="continue">continue &mdash; sample from the best peak, with a warning</option>

    </select>

    <div class="help">Required: what to do when the peak search finds several significant peaks (e.g. period aliases). With <code>stop</code>, PyAETNA writes <code>&lt;config&gt;_peak&lt;k&gt;.yaml</code>, one per peak, whose <code>prior_bounds</code> isolate it; run each (with <code>ns</code>, their evidences add up to the total).</div>

  </div>

  <div class="opt-toggle" onclick="toggleOpt(this)">

    <span class="arrow">&#9654;</span>

    <h4>Peak Search and Convergence</h4>

    <span class="skip-hint">Optional &mdash; leave empty for the defaults</span>

  </div>

  <div class="opt-content">

    <div class="fg" data-show-if="search">

      <label><span class="key-name">peak_search_starts</span> <span class="badge b-opt">Optional</span></label>

      <input type="number" id="m-pss" data-key="peak_search_starts" data-emit="set" placeholder="32" min="1" step="1">

      <div class="help">Prior draws from which the optimisations start, all at once. Default: 32.</div>

    </div>

    <div class="fg" data-show-if="search">

      <label><span class="key-name">peak_weight_threshold</span> <span class="badge b-opt">Optional</span></label>

      <input type="number" id="m-pwt" data-key="peak_weight_threshold" data-emit="set" placeholder="0.001" min="0" step="0.001">

      <div class="help">A secondary peak counts as significant when its approximate evidence is at least this fraction of the best peak's. Default: 0.001.</div>

    </div>

    <div class="fg" data-show-if="mcmc">

      <label><span class="key-name">target_ess</span> <span class="badge b-opt">Optional</span></label>

      <input type="number" id="m-tess" data-key="target_ess" data-emit="set" placeholder="(run the given number of draws)" min="1" step="1">

      <div class="help">Run until the slowest parameter has this many independent draws and the chains agree (R-hat &le; <code>max_rhat</code>, default 1.01), at most <code>num_samples_if_*</code> draws. For an ensemble, the run must also exceed 50 autocorrelation times.</div>

    </div>

  </div>

  <!-- ── Nested sampling: ns only ── -->

    </div>

  </div>

  <!-- ── Phase 1: auto only ── -->

  <div data-show-if="auto">

    <h3 class="subsec">Phase 1: Ensemble Warmup <span class="badge b-opt" style="font-size:.7rem;vertical-align:middle">auto only</span></h3>

    <div class="help" style="margin-bottom:.5rem">Gradient-free ensemble exploration of the posterior before NUTS.

    The best walker position seeds NUTS initialization; the walker variance sets the mass matrix.</div>

    <div class="fg">

      <label><span class="key-name">warmup_sampler</span></label>

      <select id="m-ws" data-key="warmup_sampler">

        <option value="ess" selected>ess &mdash; Ensemble Slice Sampling (recommended)</option>

        <option value="aies">aies &mdash; Affine Invariant Ensemble Sampler</option>

      </select>

      <div class="help"><strong>ESS</strong> (Karamanis &amp; Beutler 2020; Karamanis, Beutler &amp; Peacock 2021) is recommended for most cases. <strong>AIES</strong> (Goodman &amp; Weare 2010) is an alternative. The same sampler runs Phase&nbsp;3b if the posterior is multimodal.</div>

    </div>

    <div class="fg">

      <label><span class="key-name">num_warmup</span></label>

      <input type="number" id="m-nw" data-key="num_warmup" value="1000" min="1">

      <div class="help">Number of ensemble warmup iterations.</div>

    </div>

  </div>

  <!-- ── Walkers and moves: auto, ess, aies ── -->

  <!-- ── Ensemble: ess, aies ── -->

  <div data-show-if="ens">

    <h3 class="subsec">Ensemble</h3>

    <div class="fg">

      <label><span class="key-name">walkers_per_free_parameter</span></label>

      <input type="number" id="m-wpfp" data-key="walkers_per_free_parameter" value="3" min="2" step="1">

      <div class="help">Coefficient multiplied by n_free_params to determine the number of ensemble walkers. Must be &ge; 2. The actual walker count is the first even number &ge; coefficient &times; n_free_params (at least 3&nbsp;&times;&nbsp;n_free_params when ESS uses <code>KDEMove</code>). Typical values: 3&ndash;5.</div>

    </div>

    <div class="fg" data-show-if="aiesUsed">

      <label><span class="key-name">ensemble_init_radius</span> <span class="badge b-opt">Optional</span></label>

      <input type="number" id="m-eir" data-key="ensemble_init_radius" data-emit="set" placeholder="2.0" step="0.1" min="0.01">

      <div class="help">AIES walkers start uniformly within &plusmn;radius in unconstrained space. Default: <code>2.0</code>. A smaller value (e.g. <code>0.5</code>) keeps them away from the edges of bounded priors, where walkers can stall. ESS walkers start at the prior medians instead, so this key does not apply to ESS.</div>

    <div class="fg">

      <label><span class="key-name">num_samples_if_ensemble</span></label>

      <input type="number" id="m-nse" data-key="num_samples_if_ensemble" value="5000" min="1">

      <div class="help">Steps per walker (the maximum, with <code>target_ess</code>).</div>

    </div>

    <div class="opt-toggle" onclick="toggleOpt(this)">

      <span class="arrow">&#9654;</span>

    </div>

  </div>

  <!-- ── Phase 2: auto only ── -->

  <div data-show-if="auto">

    <h3 class="subsec">Phase 2: Convergence Test <span class="badge b-opt" style="font-size:.7rem;vertical-align:middle">auto only</span></h3>

    <div class="fg">

      <label><span class="key-name">chain_energy_threshold</span></label>

      <input type="number" id="m-cet" data-key="chain_energy_threshold" value="5.0" min="0" step="0.5">

      <div class="help">Energy-range threshold factor for unimodal/multimodal detection after warmup. The effective threshold is this value &times; n_free_params. Recommended: <code>5.0</code> (scale-invariant). Set to <code>0.0</code> to always use ensemble production.</div>

    </div>

    <div class="fg chk">

      <input type="checkbox" id="m-rsw" data-key="reseed_stuck_walkers" data-emit="nondefault" data-default="true" checked>

      <label for="m-rsw"><span class="key-name">reseed_stuck_walkers</span> &mdash; restart walkers stuck far below the ensemble after warmup</label>

      <div class="help" style="margin-left:1.6rem">Walkers trapped at a prior boundary early in the warmup carry negligible probability but distort the posterior. Default: <code>true</code>.</div>

    </div>

    <div class="fg" data-show-if="reseed">

      <label><span class="key-name">stuck_walker_threshold</span> <span class="badge b-opt">Optional</span></label>

      <input type="text" id="m-swt" data-key="stuck_walker_threshold" data-emit="set" placeholder="auto">

      <div class="help">Energy above the ensemble median (nats) beyond which a walker counts as stuck. Default <code>auto</code> = max(20, 6&radic;(n/2)).</div>

    </div>

  </div>

  <!-- ── NUTS: auto, nuts ── -->

  <!-- ── NUTS ── -->

  <div data-show-if="nutsUsed">

    <h3 class="subsec" id="mcmc-nuts-heading">Phase 3a: NUTS Production (if unimodal)</h3>

    <div class="help" style="margin-bottom:.5rem" id="mcmc-nuts-help">When posterior is unimodal, NUTS initializes at the best ESS walker position with a frozen diagonal mass matrix from ESS walker variance. Only the step size adapts during NUTS warmup.</div>

    <h3 class="subsec">NUTS</h3>

    <div class="fg">

      <label><span class="key-name">num_chains_if_nuts</span></label>

      <input type="number" id="m-ncn" data-key="num_chains_if_nuts" value="4" min="1">

      <div class="help">Number of parallel NUTS chains. In <code>auto</code> mode it must be smaller than the number of walkers.</div>

      <div class="help">Number of NUTS chains; they start from different draws around the best peak, so R-hat compares independent starts. At least 2 with <code>target_ess</code>.</div>

    </div>

    <div class="fg">

      <label><span class="key-name">internal_warmup_if_nuts</span></label>

      <input type="number" id="m-iwn" data-key="internal_warmup_if_nuts" value="500" min="1">

      <div class="help">NUTS step-size adaptation iterations. At least 200 in <code>auto</code> mode and 500 in <code>nuts</code> mode (lower values are raised to these).</div>

      <label><span class="key-name">num_adapt_if_nuts</span></label>

      <input type="number" id="m-nan" data-key="num_adapt_if_nuts" value="1000" min="100">

      <div class="help">Adaptation steps (step size and, unless switched off, mass matrix), then discarded. At least 100.</div>

    </div>

    <div class="fg">

      <label><span class="key-name">num_samples_if_nuts</span></label>

      <input type="number" id="m-nsn" data-key="num_samples_if_nuts" value="5000" min="1">

      <div class="help">Production samples per NUTS chain. Recommended: 2000&ndash;10000 for production runs.</div>

      <input type="number" id="m-nsn" data-key="num_samples_if_nuts" value="2000" min="1">

      <div class="help">Draws per chain (the maximum, with <code>target_ess</code>).</div>

    </div>

  </div>

  <!-- ── Ensemble production: auto, ess, aies ── -->

  <div data-show-if="ens">

    <h3 class="subsec" id="mcmc-ens-heading">Phase 3b: Ensemble Production (if multimodal)</h3>

    <div class="fg">

      <label><span class="key-name">num_samples_if_ensemble</span></label>

      <input type="number" id="m-nse" data-key="num_samples_if_ensemble" value="5000" min="1">

      <div class="help">Production samples per walker.</div>

    </div>

  </div>

  <!-- ── NUTS tuning: auto, nuts ── -->

  <!-- ── NUTS tuning ── -->

  <div data-show-if="nutsUsed">

    <h3 class="subsec">NUTS Tuning</h3>

    <div class="fg">

      <input type="number" id="m-mtd" data-key="max_tree_depth" value="6" min="5">

      <div class="help">Maximum NUTS trajectory length (2^depth leapfrog steps). At least 5. Higher allows longer trajectories but uses more memory. Recommended: 8&ndash;12 for hard problems.</div>

    </div>

    <div class="fg chk">

      <input type="checkbox" id="m-amm" data-key="adapt_mass_matrix_if_nuts" data-emit="nondefault" data-default="true" checked>

      <label for="m-amm"><span class="key-name">adapt_mass_matrix_if_nuts</span> &mdash; adapt the mass matrix during the adaptation steps</label>

      <div class="help" style="margin-left:1.6rem">The mass matrix starts from the curvature of the best peak; unchecked, it is kept and only the step size adapts. Default: <code>true</code>.</div>

    </div>

    <div class="fg chk">

      <input type="checkbox" id="m-dm" data-key="dense_mass" data-emit="nondefault" data-default="false">

      <label for="m-dm"><span class="key-name">dense_mass</span> &mdash; use full mass matrix for NUTS</label>

      <label for="m-dm"><span class="key-name">dense_mass</span> &mdash; use a dense (full) mass matrix for NUTS</label>

      <div class="help" style="margin-left:1.6rem"><strong>Diagonal (default, unchecked):</strong> NUTS adapts one independent step-size scale per parameter. Memory-lean; works well when parameters are approximately uncorrelated.<br><strong>Dense/full (checked):</strong> NUTS learns the full parameter covariance during warmup. Strongly improves mixing when parameters are strongly correlated (e.g. <code>rp_by_rs</code>&ndash;<code>impact</code>). Cost: O(N&sup2;) memory. Optional, default <code>false</code>.</div>

    </div>

  </div>

      <li>Pad missing basis vectors for instruments that lack them with <code>NaN</code></li>

    </ul>

    <p>3. Run PyAETNA: <code>pyaetna your_config.yaml</code></p>

    <p style="margin-top:.5rem; font-size:.84rem;">Tip: For the first run, use shorter chains (<code>num_warmup: 200</code>, <code>num_samples_if_nuts: 500</code>) to verify the model compiles and runs correctly before a full production run.</p>

    <p style="margin-top:.5rem; font-size:.84rem;">Tip: For the first run, use shorter chains (<code>num_adapt_if_nuts: 200</code>, <code>num_samples_if_nuts: 500</code>) to verify the model compiles and runs correctly before a full production run.</p>

  </div>

</div>