fix: peak search in double precision with L-BFGS and Newton steps; saddles are not peaks

## The peak search (all samplers)

Before sampling, PyAETNA optimises the posterior from `peak_search_starts` prior draws at once (BFGS in double

Before sampling, PyAETNA optimises the posterior from `peak_search_starts` prior draws at once (L-BFGS, then damped Newton steps with the exact Hessian, in double

precision, in the unbounded coordinates the samplers use, on the same posterior energy), and groups the optima into

distinct peaks. Each peak is described by its energy, its maximum log-likelihood, its position and its curvature,

i.e. a Gaussian (Laplace) approximation, which also gives an approximate evidence and so a rough share of the

| **Peak search** | | | | |

| `sampling.peak_search` | no | `true` | bool | Search the peaks of the posterior before sampling |

| `sampling.peak_search_starts` | no | `32` | int ($\ge 1$) | Prior draws from which the optimisations start |

| `sampling.peak_search_maxiter` | no | `200` | int ($\ge 1$) | BFGS iterations per start |

| `sampling.peak_search_maxiter` | no | `500` | int ($\ge 1$) | L-BFGS iterations per start, at most (then 50 damped Newton steps) |

| `sampling.peak_weight_threshold` | no | `0.001` | float ($\ge 0$) | Approximate evidence, relative to the best peak, above which a peak is significant |

| `sampling.init_from_peaks` | no | — | string | With `peak_search: false`: the `_peaks.json` of an earlier run, to start NUTS or the ensemble from |

| `sampling.init_peak` | no | `1` | int ($\ge 1$) | Rank of the saved peak to start from |

    "tinygp>=0.3.0",

    "beartype>=0.18.0",

    "blackjax>=1.6.2",

    "optax>=0.2.4",  # L-BFGS of the peak search before sampling

]

[project.optional-dependencies]

    init_from_peaks: str | None = None

    init_peak: int = 1

    peak_search_starts: int = 32

    peak_search_maxiter: int = 200

    peak_search_maxiter: int = 500

    peak_weight_threshold: float = 1e-3

    num_chains_if_nuts: int = 4

    num_adapt_if_nuts: int = 1000

        "init_from_peaks": init_from,

        "init_peak": _number(cfg, "init_peak", int, 1, 1),

        "peak_search_starts": _number(cfg, "peak_search_starts", int, 32, 1),

        "peak_search_maxiter": _number(cfg, "peak_search_maxiter", int, 200, 1),

        "peak_search_maxiter": _number(cfg, "peak_search_maxiter", int, 500, 1),

        "peak_weight_threshold": _number(cfg, "peak_weight_threshold", float, 1e-3, 0),

        "target_ess": _number(cfg, "target_ess", int, None, 1),

        "max_rhat": _number(cfg, "max_rhat", float, 1.01, 1),

"""Posterior peaks: where the posterior is high, before and after sampling.

Before sampling, every branch runs :func:`find_posterior_peaks`: BFGS from many prior draws at once, in double

precision and in NumPyro's unconstrained space, on the potential energy that the samplers use; the optima are grouped

Before sampling, every branch runs :func:`find_posterior_peaks`: L-BFGS from many prior draws at once, refined by

damped Newton steps with the exact Hessian, in double precision (physical models included) and in NumPyro's

unconstrained space, on the potential energy that the samplers use; the optima are grouped

into distinct peaks, each described by its energy, position and curvature (a Gaussian, or Laplace, approximation).

NUTS and the ensemble samplers start from the best peak, with the parameter scales of its curvature

(:func:`draw_from_peak`). When several peaks carry a significant share of the approximate evidence,

    @property

    def significant(self) -> list[Peak]:

        """Peaks whose approximate evidence is at least ``threshold`` times that of the best peak."""

        return [

            p for p in self.peaks if p.weight >= self.threshold * self.peaks[0].weight

        """The best peak, and the others that are true maxima (positive-definite Hessian; a saddle or an unconverged

        point is not a peak) with an approximate evidence at least ``threshold`` times that of the best.

        """

        best = self.peaks[0]

        return [best] + [

            p

            for p in self.peaks[1:]

            if p.hessian_ok and p.weight >= self.threshold * best.weight

        ]

    def to_dict(self) -> dict[str, Any]:

                    "log_likelihood": p.log_likelihood,

                    "log_evidence_laplace": p.log_evidence,

                    "weight": p.weight,

                    "significant": p.weight >= self.threshold * self.peaks[0].weight,

                    "significant": any(p is q for q in self.significant),

                    "n_starts": p.n_starts,

                    "hessian_ok": p.hessian_ok,

                    "params": p.params,

        }

NEWTON_STEPS = 30  # damped Newton steps after L-BFGS, for every start

SAME_PEAK_WIDTHS = (

    3.0  # optima closer than this many marginal widths in every parameter are one peak

)

def _same_peak_threshold(n_dim: int) -> float:

    """Squared Mahalanobis distance below which two optima are the same peak: the 99.9 % region of a Gaussian in

    ``n_dim`` dimensions (Wilson-Hilferty approximation of the chi-square quantile)."""

def _laplace_covariance(hessian: np.ndarray) -> tuple[np.ndarray, float, bool]:

    """Covariance and log-determinant of the Hessian, with eigenvalues clipped when it is not positive definite."""

    h = 0.5 * (hessian + hessian.T)

    finite = np.isfinite(hessian)

    h = np.where(finite, hessian, 0.0)

    h = 0.5 * (h + h.T)

    # a parameter whose curvature could not be computed gets unit variance (in unconstrained units)

    diag = np.diag(h).copy()

    diag[~np.diag(finite) | (diag == 0)] = 1.0

    np.fill_diagonal(h, diag)

    eigval, eigvec = np.linalg.eigh(h)

    ok = bool(np.all(eigval > 0) and np.all(np.isfinite(eigval)))

    ok = bool(finite.all() and np.all(eigval > 0) and np.all(np.isfinite(eigval)))

    floor = (

        1e-12 * float(np.nanmax(np.abs(eigval))) if np.any(np.isfinite(eigval)) else 1.0

    )

    plData: Any,

    plModel: Any,

    n_starts: int = 32,

    maxiter: int = 200,

    maxiter: int = 500,

    seed: int = 0,

    threshold: float = 1e-3,

    log: Callable[[str], Any] = print,

    Args:

        model_fn: The NumPyro model.

        plData, plModel: Its arguments.

        n_starts: Prior draws from which BFGS starts, all at once (vectorised).

        maxiter: BFGS iterations per start.

        n_starts: Prior draws from which L-BFGS starts, all at once (vectorised).

        maxiter: L-BFGS iterations per start, at most.

        seed: Random seed of the prior draws.

        threshold: Relative approximate evidence above which a secondary peak counts as significant.

        log: Logging function taking one string.

    Raises:

        RuntimeError: If no optimisation ends at a finite energy.

    """

    import jax.scipy.optimize as jsopt

    from jax.flatten_util import ravel_pytree

    from PyAETNA.sampling.nested_sampling import (

        draw_live_points,

        split_prior_likelihood,

    )

    from PyAETNA.sampling.precision import full_double_precision

    start_time = time.time()

    prev_x64 = jax.config.jax_enable_x64

    jax.config.update("jax_enable_x64", True)

    firewall = full_double_precision()

    firewall.__enter__()

    try:

        key = jax.random.PRNGKey(seed)

        key_trace, key_draw = jax.random.split(key)

        sec_like = _timed(energy_batch)

        sec_grad = _timed(grad_batch)

        log(

            f"Peak search: {n_starts} BFGS runs from prior draws, {n_dim} free parameters, double precision; "

            f"Peak search: {n_starts} L-BFGS runs from prior draws, {n_dim} free parameters, double precision; "

            f"one likelihood {sec_like * 1e3:.3g} ms, one gradient {sec_grad * 1e3:.3g} ms (batched)"

        )

        @jax.jit

        def _optimise(x_start: Any) -> Any:

            # L-BFGS with a zoom line search (optax), which copes with the badly scaled energies of light curves

            # (a normalisation known to 1e-5 next to an amplitude known to 1e-1); JAX's BFGS stalls there

            import optax

            import optax.tree_utils as otu

            opt = optax.lbfgs()

            value_and_grad = optax.value_and_grad_from_state(energy)

            def one(x: Any) -> Any:

                return jsopt.minimize(

                    energy, x, method="BFGS", options={"maxiter": maxiter}

                ).x

                def step(carry: Any) -> Any:

                    x, state = carry

                    value, grad = value_and_grad(x, state=state)

                    updates, state = opt.update(

                        grad, state, x, value=value, grad=grad, value_fn=energy

                    )

                    return optax.apply_updates(x, updates), state

                def go_on(carry: Any) -> Any:

                    _, state = carry

                    k = otu.tree_get(state, "count")

                    g = otu.tree_get(state, "grad")

                    return (k == 0) | ((k < maxiter) & (otu.tree_l2_norm(g) > 1e-6))

                x_end, _ = jax.lax.while_loop(go_on, step, (x, opt.init(x)))

                return x_end

            return jax.vmap(one)(x_start)

        x_opt = _optimise(x0)

        grad = jax.grad(energy)

        hess_ad = jax.hessian(energy)

        def hess_fd(x: Any, step: Any) -> Any:

            """Hessian by central differences of the gradient with ``step``."""

            def column(i: Any) -> Any:

                e = jnp.zeros_like(x).at[i].set(step[i])

                return (grad(x + e) - grad(x - e)) / (2 * step[i])

            h_fd = jax.vmap(column)(jnp.arange(x.size))

            return 0.5 * (h_fd + h_fd.T)

        use_fd = False

        def hess(x: Any, step: Any) -> Any:

            return hess_fd(x, step) if use_fd else hess_ad(x)

        def _fd_step(h: Any) -> Any:

            """A tenth of the width implied by the curvature, within [1e-8, 1e-1] (unconstrained units)."""

            curv = jnp.abs(jnp.diag(h))

            return jnp.clip(0.1 / jnp.sqrt(jnp.where(curv > 0, curv, 1.0)), 1e-8, 1e-1)

        @jax.jit

        def _newton(x_start: Any) -> Any:

            # damped Newton (Levenberg-Marquardt) with the exact Hessian: L-BFGS reaches the basin, but on these

            # ill-conditioned energies it converges slowly; a few Newton steps reach the optimum

            def one(x: Any) -> Any:

                def step(_: int, carry: Any) -> Any:

                    x, f, lam, fd = carry

                    g = grad(x)

                    h = hess(x, fd)

                    d = jnp.maximum(jnp.abs(jnp.diag(h)), 1e-12)

                    dx = jnp.linalg.solve(h + lam * jnp.diag(d), g)

                    xn = x - dx

                    fn = energy(xn)

                    ok = jnp.isfinite(fn) & (fn < f) & jnp.all(jnp.isfinite(xn))

                    return (

                        jnp.where(ok, xn, x),

                        jnp.where(ok, fn, f),

                        jnp.clip(jnp.where(ok, lam / 3, lam * 10), 1e-12, 1e12),

                        _fd_step(h),

                    )

                init = (x, energy(x), jnp.asarray(1e-3), jnp.full(x.shape, 1e-4))

                x_end, _, _, fd = jax.lax.fori_loop(0, NEWTON_STEPS, step, init)

                return x_end, fd

            return jax.vmap(one)(x_start)

        x_lbfgs = _optimise(x0)

        # the second derivatives of some transit and occultation models are NaN by automatic differentiation:

        # then the Hessian comes from central differences of the gradient (decided once, at the best start)

        e_lbfgs = np.asarray(energy_batch(x_lbfgs))

        best_start = x_lbfgs[

            int(np.argmin(np.where(np.isfinite(e_lbfgs), e_lbfgs, np.inf)))

        ]

        use_fd = not bool(np.all(np.isfinite(np.asarray(jax.jit(hess_ad)(best_start)))))

        if use_fd:

            log(

                "  the Hessian is not finite by automatic differentiation: central differences of the gradient"

            )

        x_opt, fd_steps = _newton(x_lbfgs)

        x_opt = np.asarray(x_opt)

        fd_steps = np.asarray(fd_steps)

        e_opt = np.asarray(energy_batch(jnp.asarray(x_opt)))

        finite = np.isfinite(e_opt) & np.all(np.isfinite(x_opt), axis=1)

        n_failed = int((~finite).sum())

                "The peak search found no finite posterior density: check the priors and the model"

            )

        hessian = jax.jit(jax.hessian(energy))

        hessian = jax.jit(hess)

        threshold_d2 = _same_peak_threshold(n_dim)

        found: list[

            tuple[np.ndarray, float, np.ndarray, np.ndarray, float, bool, int]

                continue

            for j, (z, e, h, cov, logdet, ok, n) in enumerate(found):

                d = x_opt[i] - z

                if float(d @ h @ d) < threshold_d2:

                # the same peak: inside its Gaussian region, or within SAME_PEAK_WIDTHS marginal widths in every

                # parameter (optima scattered along a curved ridge, such as the impact parameter-stellar density

                # degeneracy, are not separate peaks; aliases differ by many widths in some parameter)

                sd = np.sqrt(np.maximum(np.diag(cov), 0.0))

                if float(d @ h @ d) < threshold_d2 or bool(

                    np.all(np.abs(d) < SAME_PEAK_WIDTHS * sd)

                ):

                    found[j] = (z, e, h, cov, logdet, ok, n + 1)

                    break

            else:

                h = np.asarray(hessian(jnp.asarray(x_opt[i])))

                h = np.asarray(hessian(jnp.asarray(x_opt[i]), jnp.asarray(fd_steps[i])))

                cov, logdet, ok = _laplace_covariance(h)

                found.append((x_opt[i], float(e_opt[i]), h, cov, logdet, ok, 1))

        for p in peaks:

            p.weight = float(np.exp(p.log_evidence - log_total))

    finally:

        firewall.__exit__(None, None, None)

        jax.config.update("jax_enable_x64", prev_x64)

    result = PeakSearch(

from __future__ import annotations

import contextlib

import contextvars

import functools

import os

from collections.abc import Callable

        raise RuntimeError(msg) from e

# Switched on by full_double_precision(): the firewall lets everything through in float64 (the peak search)

_FIREWALL_OFF: contextvars.ContextVar[bool] = contextvars.ContextVar(

    "pyaetna_firewall_off", default=False

)

@contextlib.contextmanager

def full_double_precision() -> Any:

    """Trace the physical models in float64 too, i.e. switch the float32 firewall off, within this block.

    For the peak search before sampling (:mod:`PyAETNA.sampling.peaks`): its optimiser and its Hessian need a smooth

    energy, which the float32 rounding of the physical models (about 1e-7 relative) does not give: BFGS line

    searches stall and second derivatives of the transit model come out as NaN. Models traced inside the block keep

    the float64 path; the samplers, traced elsewhere, are not affected.

    """

    token = _FIREWALL_OFF.set(True)

    try:

        yield

    finally:

        _FIREWALL_OFF.reset(token)

def _fold_times_f32(

    times: jax.Array, t0: float | jax.Array, period: float | jax.Array

) -> jax.Array:

    """

    phase = (times - t0) / period

    phase_folded = phase - jnp.round(phase)  # [-0.5, +0.5]

    if _FIREWALL_OFF.get():

        return phase_folded * period

    return jnp.float32(phase_folded * period)  # magnitude < period/2

    @functools.wraps(fn)

    def wrapper(*args: Any, **kwargs: Any) -> Any:

        if _FIREWALL_OFF.get():

            return fn(*args, **kwargs)

        with jax.enable_x64(False):

            args = jax.tree.map(_as_float32, args)

            kwargs = jax.tree.map(_as_float32, kwargs)