Commit 9ad20bf9 authored by Gaetano Scandariato's avatar Gaetano Scandariato
Browse files

docs: three samplers and the peak search in the run-log guide, configurator and reference pages

parent 6aeb1c18
Loading
Loading
Loading
Loading
Loading
+28 −8
Changes for CHANGELOG.md: 28 added lines, 8 removed lines.
Original line number Diff line number Diff line
@@ -49,11 +49,15 @@ generate it or edit an existing one.
    are shared by the RVs and activity indicators.
  - an optional light-curve dilution factor (`star.dilution`), global or
    per instrument, for apertures contaminated by other stars.
- **Inference, in two branches:**
  - Branch 1: GPU-vectorised nested sampling (BlackJAX nested slice
    sampling), which returns the Bayesian evidence and the posterior;
  - Branch 2: an ensemble warmup (ESS or AIES) followed by NUTS or ensemble
    production, reporting AIC/BIC.
- **Inference, with three samplers chosen explicitly:**
  - GPU-vectorised nested sampling (BlackJAX nested slice sampling), which
    returns the Bayesian evidence and the posterior;
  - NUTS, and the ESS and AIES ensemble samplers, reporting AIC/BIC; they
    can run until a target number of independent draws is reached;
  - before any of them, a peak search finds the peaks of the posterior: NUTS
    and the ensembles start from the best one, and a posterior with several
    peaks stops the run with one configuration per peak, or is sampled from
    its best peak with a warning.

  Every automatic choice can be overridden from the YAML file.
- **Precision:** the ephemeris is computed in float64 and the physical
@@ -79,18 +83,34 @@ jaxoplanet, tinygp; see `docs/dependencies.md`.
  (`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,
- **Peak search before every sampler** (`sampling/peaks.py`, which also holds the MAP search of the reports): L-BFGS
  and damped Newton steps, in double precision, 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).
  `init_from_peaks`, near given values with `init_values`, 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.
- **Fixed**: ensemble slice sampling (`ess`) returned its walkers in a new random order at every step (NumPyro's
  random split of the ensemble), so a "chain" of the output mixed walkers: its autocorrelation time came out as 1–3
  steps whatever the model, and the effective sample sizes and R-hat of every ESS run were meaningless (about 30 times
  too optimistic on a 20-parameter Gaussian, whose true autocorrelation time is 37 steps). The split stays random;
  the walkers are put back in their places. The posterior draws themselves were valid.
- **Fixed**: an ensemble started from given positions drew a whole block of `checkpoint_interval` steps in its
  one-step compilation pilot, uncounted, so runs were longer than requested and the log counted draws wrongly.
- **Faster long runs**: the draws grow in place instead of being copied at every block, the checkpoint is written
  after a block at most every 5 minutes instead of at every block, and the convergence test of the ensembles thins
  the chains by up to a quarter of the autocorrelation time. On a 148-walker, 50 000-step AIES run these overheads
  took as long as the sampling. The end-of-run diagnostics skip vector sites (transit times), whose ESS computation
  ran that run out of memory (80 GB).
- **Fixed**: along directions where the curvature of a peak is negative or nearly zero, its Gaussian approximation
  was up to 10⁶ times wider than along the narrowest one; NUTS chains of the 106-parameter TTV benchmark started
  far apart and one of them never moved. The width is now set by the absolute curvature, at most 1000 times the
  narrowest. The log warns when a run reaches its maximum length without converging.

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

+1 −0
Changes for README.md: 1 added line, 0 removed lines.
Original line number Diff line number Diff line
@@ -172,6 +172,7 @@ Comprehensive documentation is available:

  - **Log (`<config>_run.log`)**: the record of the run: provenance and device, data and priors, the sampler's progress, the posterior summary with R-hat and effective sample sizes (or the Bayesian evidence), goodness of fit, derived quantities, warnings and the BibTeX entries to cite. **[Reading the run log](docs/outputs/run_log.md)** goes through it block by block and lists what to check first.
  - **Results (`<config>_results.json`)**: the run settings (sampler, device, precision, chains, draws, seed), the provenance (code revision, package versions), R-hat and ESS for every parameter, the fixed parameters, the MAP, the goodness of fit (χ², maximum likelihood, AIC, AICc, BIC), the nested-sampling evidence, the instruments and the configuration.
  - **Peaks (`<config>_peaks.json`)**: the peaks of the posterior found by the peak search before sampling, best first: energy, maximum log-likelihood, approximate (Laplace) evidence and weight, whether the peak is significant, position and widths of the parameters, and its Gaussian approximation (`z`, `cov`), from which a later run can start (`sampling.init_from_peaks`). With several significant peaks and `sampling.on_multiple_peaks: stop`, the run stops here and writes one configuration per peak, `<config>_peak<k>.yaml`, next to the original ([details](docs/config/sampling.md#the-peak-search-all-samplers)).
  - **Posterior samples (`<config>_samples.npz`)**: one array per parameter, shaped `(chain, draw)`, and the `potential_energy` of every draw when the sampler records it (NUTS, nested sampling).
  - **Model components (`<config>_components.npz`)**: the data (`bjd`, `y`, `dy`, `instrument`, `visit`) and, on every data point, the MAP model and its components (`map/total`, `map/planet`, `map/gp_mean`, …), the posterior median (`median/…`) and the 16th–84th percentile band (`q16_84/…`).
  - **Report (`<config>_report.pdf`)** and, with a `physics:` section, **physics report (`<config>_physics.pdf`)**: phase-folded fits, residuals, corner plots and derived quantities. **[Reading the PDF reports](docs/outputs/report.md)** shows an example of every page and how to read it.
+22 −3
Changes for docs/CITATIONS.md: 22 added lines, 3 removed lines.
Original line number Diff line number Diff line
@@ -50,9 +50,28 @@ PyAETNA uses NumPyro for Bayesian MCMC sampling (NUTS, ESS, AIES).
}
```

### No-U-Turn Sampler (NUTS)

PyAETNA uses NUTS, the gradient-based sampler of NumPyro, with `production_sampler: nuts`.

- Hoffman, M. D. & Gelman, A. (2014). *The No-U-Turn Sampler: Adaptively Setting Path Lengths in Hamiltonian
  Monte Carlo.* Journal of Machine Learning Research, **15**, 1593–1623.

```bibtex
@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}
}
```

### Ensemble Slice Sampling (ESS)

PyAETNA uses ESS for gradient-free ensemble warmup and multimodal production sampling.
PyAETNA uses ESS, a gradient-free ensemble sampler, with `production_sampler: ess`.

- Karamanis, M. & Beutler, F. (2020). *Ensemble Slice Sampling.*
  arXiv:[2002.06212](https://arxiv.org/abs/2002.06212)
@@ -87,7 +106,7 @@ PyAETNA uses ESS for gradient-free ensemble warmup and multimodal production sam

### Affine Invariant Ensemble Sampler (AIES)

PyAETNA uses AIES (Goodman–Weare algorithm) as an alternative ensemble sampler.
PyAETNA uses AIES (Goodman–Weare algorithm), the other ensemble sampler, with `production_sampler: aies`.

- Goodman, J. & Weare, J. (2010). *Ensemble samplers with affine invariance.*
  Comm. App. Math. Comp. Sci., **5**(1), 65–80.
@@ -109,7 +128,7 @@ PyAETNA uses AIES (Goodman–Weare algorithm) as an alternative ensemble sampler

### Nested Sampling (`production_sampler: ns`)

Cite these when you use PyAETNA's nested sampling (Branch 1), e.g. to report a Bayesian evidence.
Cite these when you use PyAETNA's nested sampling, e.g. to report a Bayesian evidence.

- Skilling, J. (2006). *Nested sampling for general Bayesian computation.*
  Bayesian Analysis, **1**, 833–859.
+27 −2
Changes for docs/config-generator.html: 27 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -149,6 +149,16 @@
    <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="noSearch mcmc">
    <label><span class="key-name">init_values</span> <span class="badge b-opt">Optional</span></label>
    <input type="text" id="m-iv" data-key="init_values" data-emit="set" placeholder="{rp_by_rs_b: 0.097, impact_b: 0.35, ...}">
    <div class="help">Instead of a saved peak: start near these values, one per free parameter (by the names of the samples file), as codes such as allesfitter start their walkers. Not together with <code>init_from_peaks</code>.</div>
  </div>
  <div class="fg" data-show-if="noSearch mcmc">
    <label><span class="key-name">init_radius</span> <span class="badge b-opt">Optional</span></label>
    <input type="number" id="m-ir" data-key="init_radius" data-emit="set" placeholder="1e-4" min="0" step="any">
    <div class="help">Half-width of the uniform ball around <code>init_values</code>, in the sampler's unbounded coordinates. Default: <code>1e-4</code>.</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">
@@ -173,16 +183,31 @@
      <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="search">
      <label><span class="key-name">peak_search_maxiter</span> <span class="badge b-opt">Optional</span></label>
      <input type="number" id="m-psm" data-key="peak_search_maxiter" data-emit="set" placeholder="500" min="1" step="1">
      <div class="help">L-BFGS iterations per start, at most; damped Newton steps with the exact curvature follow. Default: 500.</div>
    </div>
    <div class="fg">
      <label><span class="key-name">prior_bounds</span> <span class="badge b-opt">Optional</span></label>
      <input type="text" id="m-pb" data-key="prior_bounds" data-emit="set" placeholder="{period_b: [3.0, 4.5]}">
      <div class="help">Truncate the priors of these free parameters to <code>[low, high]</code>, keeping their shape: one peak of a posterior with several. The configurations written by <code>on_multiple_peaks: stop</code> set it for each peak.</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 class="fg" data-show-if="mcmc">
      <label><span class="key-name">max_rhat</span> <span class="badge b-opt">Optional</span></label>
      <input type="number" id="m-mrh" data-key="max_rhat" data-emit="set" placeholder="1.01" min="1" step="0.001">
      <div class="help">Largest R-hat accepted by <code>target_ess</code>. Default: 1.01.</div>
    </div>
  </div>

  <!-- ── Nested sampling: ns only ── -->
  <div data-show-if="ns">
    <h3 class="subsec">Nested Sampling (Branch 1)</h3>
    <h3 class="subsec">Nested Sampling</h3>
    <div class="fg">
      <label><span class="key-name">n_live_if_ns</span> &mdash; live points</label>
      <input type="number" id="m-nlive" data-key="n_live_if_ns" value="500" min="50">
@@ -340,7 +365,7 @@
  <div class="fg">
    <label><span class="key-name">checkpoint_interval</span></label>
    <input type="number" id="m-ci" data-key="checkpoint_interval" value="500" min="1">
    <div class="help">Save checkpoint every N samples. Allows resuming interrupted runs. Required for every sampler.</div>
    <div class="help">Draws per block for NUTS and the ensembles: progress (and, with <code>target_ess</code>, the convergence test) is reported after each block, and the checkpoint, which allows resuming interrupted runs, is written after a block at most every 5 minutes. Required for every sampler; nested sampling checkpoints every minute on its own.</div>
  </div>
  <div class="fg chk">
    <input type="checkbox" id="m-reset" data-key="reset" checked>
+1 −1
Changes for docs/config-generator.js: 1 added line, 1 removed line.
Original line number Diff line number Diff line
@@ -1704,7 +1704,7 @@ var DEFAULTS = {
  'sampling.reset': true, 'sampling.random_seed': null,
  'sampling.dense_mass': false, 'sampling.adapt_mass_matrix_if_nuts': true, 'sampling.peak_search': true,
  'sampling.peak_search_starts': 32, 'sampling.peak_search_maxiter': 200, 'sampling.peak_weight_threshold': 0.001,
  'sampling.init_peak': 1, 'sampling.max_rhat': 1.01, 'sampling.target_accept_prob': 0.8, 'sampling.max_tree_depth': 10,
  'sampling.init_peak': 1, 'sampling.init_radius': 1e-4, 'sampling.max_rhat': 1.01, 'sampling.target_accept_prob': 0.8, 'sampling.max_tree_depth': 10,
  'sampling.ess_moves': {DifferentialMove: 0.7, KDEMove: 0.3}, 'sampling.aies_moves': {DEMove: 0.5, StretchMove: 0.5},
  'sampling.n_live_if_ns': 500, 'sampling.dlogz_if_ns': 0.1, 'sampling.num_samples_if_ns': 4000,
  'sampling.inner_steps_per_dim_if_ns': 4, 'sampling.delete_fraction_if_ns': 0.1,
Loading