feat: optional NUTS mass-matrix adaptation (adapt_mass_matrix_if_nuts) and its...

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

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

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

  instead of freezing the diagonal one taken from the ensemble warmup. It makes NUTS converge on posteriors with

  about a hundred parameters (13× more effective samples per second at 106 parameters) but is slower on other models,

  so it stays off by default; see `docs/config/sampling.md`.

### Fixed before the release (2026-10-03)

- **Phase curves in float32:** the eclipse mask divided the planet's emission by its day-side amplitude; when a

several separate solutions, use **Branch 1**. Otherwise use **Branch 2**, whose

`auto` mode is the fastest route to a well-sampled posterior with a single peak.

```

                                   [ Start ]

                                       │

            ┌──────────────────────────┴──────────────────────────────┐

            ▼                                                         ▼

  BRANCH 1: production_sampler: ns                  BRANCH 2: production_sampler: auto

            │                                                         │

            ▼                                              Phase 1: ensemble warmup

   Nested sampling from the prior                              (ESS or AIES)

   (no warmup, no MCMC)                                               │

            │                                          stuck-walker re-seeding

            │                                                         │

            │                                      Phase 2: energy-spread test

            │                                          ┌──────────────┴──────────────┐

            │                                          ▼                             ▼

            │                                  unimodal (≤ threshold)      multimodal (> threshold)

            │                                          │                             │

            │                                  Phase 3a: NUTS              Phase 3b: ensemble

            │                                    production                   production

            │                                          │                             │

            │     Direct variants (skip the warmup):   │                             │

            │     production_sampler: nuts → NUTS only │                             │

            │     production_sampler: ess | aies → ensemble only                     │

            ▼                                          ▼                             ▼

   posterior + ln Z                          posterior (+ AIC/BIC in the report)

            └───────────────────────────┬─────────────────────────────────────────────┘

                                        ▼

                               Reports and physics

```

![The two sampling branches](img/branches.png)

The direct variants skip the warmup: `production_sampler: nuts` runs NUTS only, and `ess` or `aies` the ensemble

only. The figure is drawn by `make_branches_figure.py` in this folder; rerun it when the pipeline changes.

### Pros and cons of each route

   the region of the maximum likelihood, from which NUTS starts. Walkers

   that are still stuck after the last round (the log then warns) are

   harmless when NUTS follows, because only the lowest-energy walkers are used.

   The re-seeding is needed also when NUTS follows: the spread of the walkers

   decides between NUTS and ensemble production (step 3) and sets the NUTS mass

   matrix, and stuck walkers would distort both. A re-seeded walker, however,

   continues the history of the walker it was copied from, so NUTS never starts

   two chains from walkers of the same lineage (step 4).

3. **Phase 2, energy-spread test.** Compares the spread of walker potential

   energies with `chain_energy_threshold × n_free_params`.

4. **Phase 3a, unimodal → NUTS.** If the spread is ≤ threshold, NUTS chains

   start from the lowest-energy walkers, one each, with a frozen diagonal mass matrix estimated

   start from the lowest-energy walkers, one each and at most one per lineage, with a frozen diagonal mass matrix estimated

   from the ensemble. It runs `internal_warmup_if_nuts` adaptation steps, then

   `num_samples_if_nuts` samples per chain.

5. **Phase 3b, multimodal → ensemble.** If the spread is > threshold (always,

  walker costs one likelihood evaluation per step (NumPyro's own AIES recomputes the current one, two

  evaluations); the chains are the same.

- **`aies_moves`**: default `{DEMove: 0.5, StretchMove: 0.5}`. Supported:

  `DEMove`, `StretchMove`.

  `DEMove`, `StretchMove`. By default the AIES therefore alternates the stretch move of Goodman & Weare (2010), the

  algorithm of `emcee`, with differential-evolution moves (ter Braak 2006; Nelson et al. 2014), which help the

  ensemble along narrow, correlated posteriors; `{StretchMove: 1.0}` reproduces a plain `emcee` run.

### Walker initialisation and safeguards

  chains that did not share their starting point make R-hat meaningful: R-hat

  compares the spread between chains with the spread within them, and chains

  started from one point agree with each other for a while even when they have

  not converged. Until October 2026 all chains started from the single best

  walker. R-hat above 1.01 is flagged in the diagnostics table (`*`, and `**`

  not converged. For the same reason the starts are taken at most one per

  lineage: a walker re-seeded from another during the warmup and the walker it

  was copied from count as one, and the next lowest energy is taken instead

  (only if there are fewer lineages than chains do two chains share one). Until

  October 2026 all chains started from the single best walker, and until

  2026-10-05 a re-seeded walker and its donor could both be chosen. R-hat above 1.01 is flagged in the diagnostics table (`*`, and `**`

  above 1.1).

- **Frozen inverse mass matrix**: a diagonal inverse mass matrix from the

  variance of the walker positions (`adapt_mass_matrix=False`). Only the step

  size adapts during `internal_warmup_if_nuts`. This prevents step-size

  collapse when parameter scales differ by many orders of magnitude (e.g.

  `delta_period` ~ $10^{-8}$ vs `amp_day` ~ $10^2$).

- **Adapted mass matrix** (`adapt_mass_matrix_if_nuts: true`, default `false`): NUTS starts from the walkers' diagonal

  inverse mass matrix but adapts it during `internal_warmup_if_nuts` (with `dense_mass: true`, a dense one). Use it

  for posteriors with many parameters, where the walkers have not converged by the end of the ensemble warmup and

  the frozen mass matrix makes the step sizes collapse: on a synthetic system with 106 free transit times it made NUTS

  converge ($\hat R$ 1.067 → 1.009, 13× more effective samples per second), and it gained 1.4–5× on WASP-18, TOI-216

  and a system of 56 parameters. It is not the default because it can be much slower: 50–100× on K2-140 (20

  parameters, a narrow period prior in mixed precision) and about 4× on a system of 31 parameters; a dense matrix was

  slower still (`paper_benchmarks/ttv_scaling`). Use 1000 adaptation steps or more.

- **Minimum adaptation**: at least 200 steps in `auto` mode and 500 in direct

  `nuts` mode; lower values are raised with a warning.

| `sampling.target_accept_prob` | when `auto` or `nuts` | — | float ($[0.5, 0.99]$) | NUTS target acceptance probability |

| `sampling.max_tree_depth` | when `auto` or `nuts` | — | int ($\ge 5$) | NUTS maximum tree depth |

| `sampling.dense_mass` | no | `false` | bool | Dense NUTS mass matrix |

| `sampling.adapt_mass_matrix_if_nuts` | no | `false` | bool | Adapt the NUTS mass matrix during `internal_warmup_if_nuts`, starting from the walkers' variances (`auto`); see "Adapted mass matrix" above |

| `sampling.num_samples_if_ensemble` | when `auto`, `ess`, `aies` | — | int ($> 0$) | Samples per walker in ensemble production |

---

| | `ccf` | `template` |

|---|---|---|

| Line model | Gaussian (macroturbulence folded into $\beta$) | Voigt ($\beta$, $\gamma$) and radial-tangential macroturbulence ($\zeta$) |

| How | the width and amplitude of the Gaussian fitted to the out-of-transit line are refitted at every evaluation (so $v \sin i_*$ stays free), then the linear response of its centre to the occultation | two Fourier integrals per point on a fixed grid |

| How | the width and amplitude of the Gaussian fitted to the out-of-transit line are refitted at every evaluation by a few Newton steps (so $v \sin i_*$ stays free and the model stays differentiable), then the linear response of its centre to the occultation | two Fourier integrals per point on a fixed grid |

| Accuracy | ≤ 1 % of the amplitude against a brute-force Gaussian fit to simulated line profiles | ≤ 4 % of the amplitude against a brute-force cross-correlation |

| Cost | small against the disc integration: the TOI-5398 fit took 14 s on a laptop GPU | 18 s |

| Approximations | linear in $f$; the Gaussian fit is assumed without a free continuum level (~1 % effect) | small occulted region: a few % for large planets or near the limb |

GPU and exhausted its memory at batches of 512. The physical models are therefore traced with double precision

switched off, and regression tests check that no float64 array appears inside them. Keep both when adding a model.

The reason for the firewall is the hardware: consumer GPUs run float64 arithmetic 32 to 64 times more slowly than

float32 (64 times on the RTX 5070 Ti Laptop GPU of the benchmarks), whereas times need double precision (a transit

epoch of ~2 460 000 BJD is known to ~0.1 d in float32). Folded times, whose magnitude never exceeds half a period,

are exact enough in float32 (0.005 s at half a period of 1.5 d, 0.3 s at 100 d), so only the ephemeris arithmetic runs

in float64. Models with GPs run their likelihood in float64 too, so a consumer GPU gains little on them (the K2-100

nested sampling takes as long on the laptop GPU as on its CPU); a data-centre GPU, with fast float64, does.

## RM disc integrals: dense, batched, highest precision

The occulted-disc terms of the Rossiter-McLaughlin effect follow the spherical-harmonic formalism of `starry`

(Luger et al. 2019) as implemented in `jaxoplanet`. Its change-of-basis matrices are sparse, which cannot be

differentiated in batches (vmap) on a GPU; `body_models.py` therefore precomputes them once as dense float32 matrices

for the fixed polynomial degree and applies them as matrix products. On NVIDIA GPUs, JAX's default float32 matrix

products are TF32 (10-bit mantissa), which put the anomaly off by up to 2 m/s near ingress and egress; the disc

integrals run under `jax.default_matmul_precision("highest")`, and `test_rm_matmul_precision.py` checks the precision

of every dot product in the traced program.

## AIES keeps the walkers' log densities

NumPyro's `AIES` recomputes the log density of every walker's current position at each step, doubling the

evaluations (pyro-ppl/numpyro#2301). PyAETNA's `CachedAIES` (`sampling/engine/ensemble_kernels.py`) carries the

walkers' log densities between steps, as `emcee` does, with unchanged chains. The cache is an internal state: it is

not stored with the samples (potential energies are recomputed after sampling, `sampling/energy.py`).

## CPU runs: one chain per host device, one BLAS thread

On a CPU, nested sampling gives each slice chain its own JAX host device (`sampling/nested_sampling.py`,

atmosphere; compare total harmonic amplitudes, not individual terms, with

codes that use a sinusoidal atmospheric model.

**Occultation.** The planetary (thermal + reflected) term is multiplied by the visible fraction of the planetary disc,

computed by jaxoplanet, which describes the ingress and egress of the occultation; beaming and ellipsoidal terms are

not occulted.

**Light-travel time and tidal lag.** The phase $\phi$ of the planetary and ellipsoidal terms is computed at the

position of the planet corrected for the light-travel time across the orbit: photons observed at $t$ were emitted at

$t + z/c$, with $z$ the distance of the planet from the sky plane, positive towards the observer (one iteration, whose

residual is below $10^{-4}$ for known systems). The ellipsoidal term can be delayed by a tidal lag angle

`planet_<X>.tidallag`, $-A_{\rm ellip}\cos 2(\phi - \delta_{\rm ell})$. With `amp_beam: from_K` the beaming is tied to

the orbit, $A_{\rm beam} = 4\alpha_{\rm beam}K/c$, with the factor $\alpha_{\rm beam}$ of the star in the passband

(`star.beaming_alpha`; Loeb & Gaudi 2003, Bloemen et al. 2011).

---

## 2. Lambert sphere phase function