"""The No-U-Turn Sampler engine (exact for every prior)."""
from __future__ import annotations
import copy
import os
import pickle
import warnings
from collections.abc import Iterator
from typing import Any
import numpy as np
from .._utils import RNGLike, progress, spawn, warn
from ..model import PopulationModel
from ..results import PopulationResult
from ..warnings import ConvergenceWarning
from ._curvature import covariance_from_hessian, fd_hessian, fd_hessian_diag
from ._nuts_core import Metric, adaptation_windows, find_reasonable_eps, nuts_step
from .base import Engine
__all__ = ["NUTS"]
#: ``target_accept`` when it is ``None``, except for logistic priors.
_TARGET_ACCEPT = 0.8
#: ``target_accept`` for logistic priors when it is ``None``.
_TARGET_ACCEPT_LOGISTIC = 0.9
#: ``metric="auto"`` is dense for logistic priors up to this many dimensions, and up to
#: this many the warm-up curvature is averaged over every chain's point, not taken at one.
_DENSE_MAX_DIM = 200
#: Up to this many dimensions, warm-up starts the metric from the curvature at the mode.
_HESSIAN_MAX_DIM = 500
# Partial centring re-creates a funnel when the field itself pins the scale
# hyperparameters down poorly (few modes per hyperparameter). In validation runs 0.28
# separated K <= 6, N <= 40 (0.33-0.62; divergences when centred) from K >= 10
# (0.11-0.25; 3-20x more ESS/s when centred).
#: Logistic priors stay non-centred when the median spread of the modes' log-scales given
#: the field exceeds this.
_CENTRING_MAX_SD = 0.28
#: Warm-up iteration of the partial centring of logistic priors, or the first slow
#: window's start if that is earlier.
_GEOMETRY_ITER = 25
#: ``t0``, ``gamma`` and ``kappa`` of the dual-averaging step-size adaptation.
_DUAL_AVERAGING = (10.0, 0.05, 0.75)
def _centred(model: PopulationModel, lam: np.ndarray) -> PopulationModel:
"""The model with its logistic prior partially centred by weights ``lam``."""
return model._with_bound_prior(model.prior.with_centering(lam))
def _recentre(prior, x: np.ndarray, lam: np.ndarray) -> np.ndarray:
"""``prior.recentre`` on the prior's block of ``x``; an inferred selection's
logit S after it is unchanged (DECISIONS #80)."""
d = prior.dim
return np.concatenate([prior.recentre(x[:d], lam), x[d:]])
def _conditional_log_scale_sd(model, prior, pts: list[np.ndarray]) -> np.ndarray:
"""Spread of each mode's log-scale given the field, from the hyperparameter curvature.
In the centred form the field stays put while the hyperparameters move, so the
curvature of the hyperparameter block alone says how well the field pins them.
(The marginal spread from the full Hessian is unreliable: weakly identified
length scales can make it arbitrarily large.) Modes with fixed scales get 0.
"""
n_act = prior.n_active
if prior.n_free == 0:
return np.zeros(n_act)
zeros = np.zeros(n_act)
fn = _centred(model, zeros).log_density_unconstrained_and_grad
hyper = range(n_act, prior.dim) # the hyperparameters, before any logit S
H = np.mean(
[fd_hessian(fn, _recentre(prior, x, zeros), cols=hyper)[hyper] for x in pts], axis=0
)
cov = covariance_from_hessian(H)
if cov is None:
return np.full(n_act, np.inf)
_, dlogs = prior.log_scales(prior.theta_from_x(np.mean(pts, axis=0)[: prior.dim]))
G = dlogs[prior.free] # (n_free, n_active): d log s_j / d log theta
return np.sqrt(np.maximum(np.einsum("aj,ab,bj->j", G, cov, G), 0.0))
class _Chain:
"""State of one NUTS chain, including its step-size and metric adaptation."""
def __init__(self, x0, fn, rng, n_warmup, target_accept, dense):
self.fn = fn
self.rng = rng
self.x = np.asarray(x0, dtype=float)
self.lp, self.g = fn(self.x)
d = self.x.size
self.metric = Metric(np.eye(d) if dense else np.ones(d))
self.informed = False # metric initialised from curvature rather than the identity
self.dense = dense
self.eps = find_reasonable_eps(self.x, self.lp, self.g, self.metric, fn, rng)
self.target = target_accept
self._restart_dual_averaging()
self.windows = adaptation_windows(n_warmup)
self.win_ends = {e for _, e in self.windows}
self.buf: list[np.ndarray] = []
self.n_warmup = n_warmup
self.it = 0
def reset_geometry(self, fn, x, inv_m) -> None:
"""Move to a new parametrisation and/or metric mid warm-up."""
self.fn = fn
self.x = np.asarray(x, dtype=float)
self.lp, self.g = fn(self.x)
if inv_m is not None:
self.metric = Metric(inv_m)
self.informed = True
self.buf = []
self.eps = find_reasonable_eps(self.x, self.lp, self.g, self.metric, self.fn, self.rng)
self._restart_dual_averaging()
def _restart_dual_averaging(self):
self.mu = np.log(10.0 * self.eps)
self.Hbar = 0.0
self.log_eps_bar = 0.0
self.m = 0
def step(self, max_depth, adapt):
x, lp, g, acc, depth, div, nl, energy = nuts_step(
self.x, self.lp, self.g, self.eps, self.metric, self.fn, self.rng, max_depth
)
self.x, self.lp, self.g = x, lp, g
if adapt:
self._adapt(acc)
self.it += 1
return acc, depth, div, nl, energy
def _adapt(self, acc):
t0, gamma, kappa = _DUAL_AVERAGING
self.m += 1
m = self.m
self.Hbar = (1.0 - 1.0 / (m + t0)) * self.Hbar + (self.target - acc) / (m + t0)
log_eps = self.mu - np.sqrt(m) / gamma * self.Hbar
eta = m ** (-kappa)
self.log_eps_bar = eta * log_eps + (1.0 - eta) * self.log_eps_bar
self.eps = float(np.exp(log_eps))
it = self.it
if any(s <= it < e for s, e in self.windows):
self.buf.append(self.x.copy())
if (it + 1) in self.win_ends:
arr = np.asarray(self.buf)
n = len(arr)
d = self.x.size
prev = self.metric.inv_m
if self.dense:
cov = np.atleast_2d(np.cov(arr.T)) if n > 1 else prev
if self.informed: # shrink toward the curvature-based metric
k = d + 5.0
inv_m = (n / (n + k)) * cov + (k / (n + k)) * prev
else: # Stan's regularisation
inv_m = (n / (n + 5.0)) * cov + 1e-3 * (5.0 / (n + 5.0)) * np.eye(d)
else:
var = arr.var(axis=0, ddof=1) if n > 1 else prev
if self.informed:
inv_m = (n / (n + 5.0)) * var + (5.0 / (n + 5.0)) * prev
else:
inv_m = (n / (n + 5.0)) * var + 1e-3 * (5.0 / (n + 5.0))
self.metric = Metric(inv_m)
self.buf = []
self.eps = find_reasonable_eps(self.x, self.lp, self.g, self.metric, self.fn, self.rng)
self._restart_dual_averaging()
if it == self.n_warmup - 1 and self.m > 0:
self.eps = float(np.exp(self.log_eps_bar))
_STAT_KEYS = ("accept_stat", "tree_depth", "diverging", "n_leapfrog", "energy")
def _advance_chain(chain: _Chain, n_iter: int, adapt: bool, max_depth: int, record: bool):
"""Advance one chain; returns ``(draws, log_probs, stats, n_divergent)``."""
d = chain.x.size
xs = np.empty((n_iter, d)) if record else None
lps = np.empty(n_iter) if record else None
st = np.empty((len(_STAT_KEYS), n_iter)) if record else None
n_div = 0
for i in range(n_iter):
acc, depth, div, nl, energy = chain.step(max_depth, adapt)
n_div += int(div)
if record:
xs[i] = chain.x
lps[i] = chain.lp
st[:, i] = (acc, depth, div, nl, energy)
return xs, lps, st, n_div
def _cpu_count() -> int:
# the cores this process may use (CPU affinity, e.g. under a batch scheduler)
return os.process_cpu_count() or 1
def _worker_threads(n_workers: int) -> int:
"""BLAS threads per worker: the cores shared between workers, never above the
limit already in force in this process (e.g. ``OMP_NUM_THREADS=1``; audit N-3)."""
from threadpoolctl import threadpool_info
share = max(1, _cpu_count() // n_workers)
info = threadpool_info()
# the BLAS pools' limit (another runtime, e.g. a single-threaded OpenMP, does not cap
# it); where no BLAS pool is visible (e.g. Accelerate), every reported limit counts
blas = [lib["num_threads"] for lib in info if lib.get("user_api") == "blas"]
limits = [max(blas)] if blas else [lib["num_threads"] for lib in info]
return min([share, *limits])
_WORKER_MODEL: PopulationModel | None = None # the base model, sent once per worker
def _worker_model(model: PopulationModel) -> PopulationModel:
"""The model without the data kept for value draws (DECISIONS #84), which workers
never use: kept samples need not be copied to each of them."""
lk = model.likelihoods
if "within" not in lk.meta:
return model
out = copy.copy(model)
out.likelihoods = lk._copy(meta={k: v for k, v in lk.meta.items() if k != "within"})
return out
def _init_worker(model: PopulationModel, n_threads: int) -> None:
global _WORKER_MODEL
from threadpoolctl import threadpool_limits
_WORKER_MODEL = model
threadpool_limits(limits=n_threads)
def _chain_worker(payload):
"""Run in a worker process: advance one chain and send it back without the model.
An error from the model comes back as a value, so the parent can tell it from a
failure of the pool itself (audit N-7).
"""
chain, lam, n_iter, adapt, max_depth, record = payload
try:
model = _WORKER_MODEL if lam is None else _centred(_WORKER_MODEL, lam)
chain.fn = model.log_density_unconstrained_and_grad
out = _advance_chain(chain, n_iter, adapt, max_depth, record)
except Exception as err:
return err
chain.fn = None
return chain, out
[docs]
class NUTS(Engine):
"""No-U-Turn Sampler on the exact posterior, for every prior.
Samples the unconstrained coordinates of the model (isometric log-ratios
for simplex priors, whitened fields and log-hyperparameters for smooth
priors, log-rates for rates) with multinomial NUTS, dual-averaging step
size adaptation and Stan-style windowed metric adaptation. Chains run in
lock-step, so the engine supports progress reporting, a generator
interface (:meth:`sample`) and resuming (call :meth:`run` again to draw
more samples without re-running warm-up).
Parameters
----------
model : PopulationModel
n_chains : int
Number of independent chains (default 4; needed for R-hat).
n_warmup : int
Warm-up (adaptation) iterations per chain.
target_accept : float, optional
Target mean acceptance statistic for step-size adaptation. Defaults to
0.8, or 0.9 for logistic (smoothness) priors, whose hyperparameters
make the geometry harder.
max_depth : int
Maximum tree depth (at least 1; at most ``2**max_depth`` gradient
evaluations per draw).
metric : {"auto", "diag", "dense"}
Adapted Euclidean metric. ``"auto"`` uses a dense metric for logistic
(smoothness) priors with at most 200 unconstrained dimensions, whose
hyperparameters correlate with the field, and a diagonal one otherwise.
reparametrize : bool
Partially centre the basis coefficients of logistic priors during warm-up.
For logistic priors, switch each basis coefficient between the
non-centred and centred forms according to how strongly the data
constrain it (partial non-centring with weights ``1 / -H_jj`` from the
curvature at the chains' positions early in warm-up). If the field
pins the hyperparameters down poorly (few basis coefficients per
hyperparameter, e.g. a handful of bins), the coefficients stay
non-centred, which avoids funnels. This changes the sampling
coordinates, not the posterior.
init : array_like (dim,) or (n_chains, dim), optional
Starting points in the model's unconstrained coordinates; a single
point gets ``jitter`` added for each chain. Every start must have a
finite log density. By default chains start at the posterior mode (in
unconstrained coordinates) plus ``jitter``.
jitter : float
Standard deviation of the Gaussian jitter around the starting point.
n_jobs : int
Number of worker processes for the chains (``-1``: one per CPU this
process may use). The default ``1`` runs the chains in this process.
The workers share the cores' BLAS threads instead of oversubscribing
them, never exceeding a thread limit already set. Parallel and serial
runs use the same random streams, so their draws agree up to rounding
differences in multithreaded linear algebra. Parallel runs show no
progress bar.
rng : int, Generator or None
Random number generator (SPEC 7); each chain gets an independent child.
progress : bool
Show a progress bar.
Notes
-----
For logistic (smoothness) priors, after ``min(25, first window start)``
warm-up iterations (none when ``n_warmup < 20``, which also skips metric
adaptation) the engine reparametrises the prior (``reparametrize``) and,
for at most 500 dimensions, initialises the metric from the
finite-difference curvature of the log posterior at the chains'
positions. Each later adaptation window shrinks its estimate toward the
previous window's metric, starting from that curvature-based one.
Simplex and rate priors use Stan's windowed adaptation unchanged.
Results (``x``, ``log_density_unconstrained``, ``model``) are in the
prior's non-centred coordinates whatever the sampling coordinates;
``info["centring"]`` holds the partial-centring weights used, and
``info["inv_metric"]`` is in those sampling coordinates.
"""
name = "nuts"
joint = True
markov = True
[docs]
def __init__(
self,
model: PopulationModel,
*,
n_chains: int = 4,
n_warmup: int = 1000,
target_accept: float | None = None,
max_depth: int = 10,
metric: str = "auto",
reparametrize: bool = True,
init: Any = None,
jitter: float = 0.1,
n_jobs: int = 1,
rng: RNGLike = None,
progress: bool = False,
) -> None:
logistic = model.prior.has_centering
# results are in the non-centred form, so only engines ever hold a centred prior
if logistic and np.any(model.prior.lam != 1.0):
raise ValueError(
"model: the logistic prior must be in its non-centred form (prior.lam all 1), "
"as the model is built; NUTS chooses the centring itself."
)
super().__init__(model, rng=rng, progress=progress)
if n_chains < 1:
raise ValueError("n_chains must be at least 1.")
if max_depth < 1:
raise ValueError("max_depth must be at least 1.")
if n_warmup < 0 or jitter < 0:
raise ValueError("n_warmup and jitter must be non-negative.")
if n_jobs == -1:
n_jobs = _cpu_count()
if n_jobs < 1:
raise ValueError("n_jobs must be a positive integer or -1.")
self.n_jobs = int(n_jobs)
if metric not in ("auto", "diag", "dense"):
raise ValueError("metric must be 'auto', 'diag' or 'dense'.")
if target_accept is None:
target_accept = _TARGET_ACCEPT_LOGISTIC if logistic else _TARGET_ACCEPT
if not 0 < target_accept < 1:
raise ValueError("target_accept must be in (0, 1).")
self.n_chains = int(n_chains)
self.n_warmup = int(n_warmup)
self.target_accept = float(target_accept)
self.max_depth = int(max_depth)
if metric == "auto":
metric = "dense" if logistic and model.dim <= _DENSE_MAX_DIM else "diag"
self.metric = metric
self.reparametrize = bool(reparametrize)
self.init = init
self.jitter = float(jitter)
self._chains: list[_Chain] | None = None
self._base_model = model
self._lam = None # partial-centring weights of the sampling coordinates, once chosen
self._pool = None # worker processes, kept for one sample() (audit N-13)
self._warm = False
self._x: list[np.ndarray] = []
self._lp: list[np.ndarray] = []
self._stats: dict[str, list[np.ndarray]] = {k: [] for k in _STAT_KEYS}
self.warmup_divergences = 0
# ---------------------------------------------------------- set-up
def _initial_points(self) -> np.ndarray:
d = self.model.dim
child = spawn(self.rng, 1)[0]
if self.init is not None:
x0 = np.asarray(self.init, dtype=float)
if x0.shape == (d,):
x0 = np.tile(x0, (self.n_chains, 1))
x0 = x0 + self.jitter * child.standard_normal(x0.shape)
if x0.shape != (self.n_chains, d):
raise ValueError(f"init must have shape ({d},) or ({self.n_chains}, {d}).")
if not all(np.isfinite(self.model.log_density_unconstrained(x)) for x in x0):
raise ValueError("init has a starting point with zero posterior density.")
return x0
from .optimize import MAP
with warnings.catch_warnings():
warnings.simplefilter("ignore", ConvergenceWarning)
mode = MAP(self.model, space="unconstrained").run()
x_mode = mode.x
pts = []
for _ in range(self.n_chains):
for scale in (self.jitter, self.jitter / 10, 0.0):
x = x_mode + scale * child.standard_normal(d)
if np.isfinite(self.model.log_density_unconstrained(x)):
break
pts.append(x)
return np.asarray(pts)
def _setup(self) -> None:
if self._chains is not None:
return
x0 = self._initial_points()
rngs = spawn(self.rng, self.n_chains)
fn = self.model.log_density_unconstrained_and_grad
self._chains = [
_Chain(x0[c], fn, rngs[c], self.n_warmup, self.target_accept, self.metric == "dense")
for c in range(self.n_chains)
]
def _adapt_geometry(self) -> None:
"""Reparametrise a logistic prior and initialise the metric from curvature.
Only for priors with a ``with_centering`` method (the logistic family):
for simplex priors the curvature at single points is a poor guide to
the posterior scales (sparse bins), so they keep Stan's adaptation.
"""
assert self._chains is not None
base = self._base_model # non-centred (see __init__)
prior0 = base.prior
if not prior0.has_centering:
return
chains = self._chains
d = base.dim
pts = [ch.x for ch in chains]
# curvature at every chain's position for small problems, one chain otherwise
n_pts = len(chains) if d <= _DENSE_MAX_DIM else 1
model = base
if self.reparametrize and prior0.n_active > 0:
fn0 = base.log_density_unconstrained_and_grad
n_act = prior0.n_active
# only the field's diagonal curvature is needed here (audit N-10)
diag = np.mean([fd_hessian_diag(fn0, x, range(n_act)) for x in pts[:n_pts]], axis=0)
lam = np.clip(1.0 / np.maximum(-diag, 1.0), 0.0, 1.0)
# one decision for all modes that share the hyperparameters: mixing centred
# and non-centred modes under one length scale mixed worst in validation
sd_c = _conditional_log_scale_sd(base, prior0, pts)
scaled = sd_c > 0.0
if np.any(scaled) and np.median(sd_c[scaled]) > _CENTRING_MAX_SD:
lam[scaled] = 1.0
model = _centred(base, lam)
pts = [_recentre(prior0, x, lam) for x in pts]
self._lam = lam
self.model = model
fn = model.log_density_unconstrained_and_grad
inv_m = None
if d <= _HESSIAN_MAX_DIM:
H = np.mean([fd_hessian(fn, x) for x in pts[:n_pts]], axis=0)
cov = covariance_from_hessian(H)
if cov is not None:
inv_m = cov if self.metric == "dense" else np.diag(cov).copy()
for ch, x in zip(chains, pts, strict=True):
ch.reset_geometry(fn, x, inv_m)
# ---------------------------------------------------------- running
def _parallel(self, n_iter: int, adapt: bool, record: bool):
"""Advance every chain ``n_iter`` iterations in worker processes.
Returns per-chain outputs, or ``None`` if no worker pool could be used
(the caller then runs serially; the chains are untouched). An error raised
by the model in a worker comes back as a value and is raised here.
"""
from concurrent.futures import ProcessPoolExecutor
from concurrent.futures.process import BrokenProcessPool
assert self._chains is not None
chains = self._chains
fn = self.model.log_density_unconstrained_and_grad
for ch in chains:
ch.fn = None # the model goes to each worker once, not with every chain
payloads = [(ch, self._lam, n_iter, adapt, self.max_depth, record) for ch in chains]
try:
if self._pool is None:
model = _worker_model(self._base_model)
try: # e.g. a lambda in a prior cannot be sent to workers
pickle.dumps((model, payloads[0]))
except (pickle.PicklingError, AttributeError, TypeError) as err:
return self._serial_from_now(err)
n_workers = min(self.n_jobs, len(chains))
self._pool = ProcessPoolExecutor(
max_workers=n_workers,
initializer=_init_worker,
initargs=(model, _worker_threads(n_workers)),
)
# errors from the model come back as values (see _chain_worker), so only a
# failure of the pool itself falls back to serial (audit N-7)
results = list(self._pool.map(_chain_worker, payloads))
except (OSError, BrokenProcessPool, NotImplementedError) as err:
return self._serial_from_now(err)
finally:
for ch in chains:
ch.fn = fn
for r in results:
if isinstance(r, BaseException):
raise r
self._chains = []
outs = []
for ch, out in results:
ch.fn = fn
self._chains.append(ch)
outs.append(out)
return outs
def _serial_from_now(self, err: BaseException) -> None:
warn(
f"Could not run chains in parallel ({type(err).__name__}: {err}); "
"running them serially.",
RuntimeWarning,
)
self.n_jobs = 1 # do not try again for this engine
self._close_pool()
def _close_pool(self) -> None:
if self._pool is not None:
self._pool.shutdown()
self._pool = None
def _advance(self, n_iter: int, adapt: bool, desc: str) -> None:
"""Warm-up iterations for every chain (parallel or interleaved serial)."""
if n_iter <= 0:
return
assert self._chains is not None
if self.n_jobs > 1 and len(self._chains) > 1:
outs = self._parallel(n_iter, adapt, record=False)
if outs is not None:
self.warmup_divergences += sum(o[3] for o in outs)
return
for _ in progress(range(n_iter), n_iter, desc, self.progress):
for ch in self._chains:
_, _, div, _, _ = ch.step(self.max_depth, adapt=adapt)
self.warmup_divergences += int(div)
[docs]
def warmup(self) -> None:
"""Run warm-up (adaptation) if it has not been run yet."""
self._setup()
if self._warm:
return
assert self._chains is not None
windows = self._chains[0].windows
logistic = self._base_model.prior.has_centering
# adapt the geometry once the chains have left the mode for the typical set
switch = min(_GEOMETRY_ITER, windows[0][0]) if (windows and logistic) else 0
try:
self._advance(switch, adapt=True, desc="NUTS warm-up")
if switch > 0:
self._adapt_geometry()
self._advance(self.n_warmup - switch, adapt=True, desc="NUTS warm-up")
finally:
self._close_pool() # warm-up phases share one pool; sampling starts its own
self._warm = True
def _record(self, xs, lps, st, n_done: int) -> None:
if n_done:
self._x.append(xs[:, :n_done])
self._lp.append(lps[:, :n_done])
for j, k in enumerate(_STAT_KEYS):
self._stats[k].append(st[:, j, :n_done])
def _report(self, X: np.ndarray, LP: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
"""Draws and log densities from the sampling coordinates to the model's (audit N-1)."""
if self._lam is None:
return X, LP
d = self.model.prior.dim # an inferred selection's logit S follows unchanged
x, lp = self.model.prior.to_noncentred(X[..., :d].reshape(-1, d), LP.reshape(-1))
x = np.concatenate([x.reshape(*X.shape[:-1], d), X[..., d:]], axis=-1)
return x, lp.reshape(LP.shape)
[docs]
def sample(self, n_samples: int) -> Iterator[np.ndarray]:
"""Generator over post-warm-up iterations, yielding ``(n_chains, dim)`` positions.
Positions are in the model's unconstrained coordinates (the prior's
non-centred form). Draws are also recorded and included in the next
:meth:`run` / :attr:`result`. With ``n_jobs > 1`` the chains advance in
worker processes and the positions are yielded once all have finished.
Parameters
----------
n_samples : int
Iterations per chain.
"""
try:
self.warmup()
assert self._chains is not None
C, d = self.n_chains, self.model.dim
n_samples = int(n_samples)
if self.n_jobs > 1 and C > 1 and n_samples > 0:
outs = self._parallel(n_samples, adapt=False, record=True)
if outs is not None:
xs, lps = self._report(
np.stack([o[0] for o in outs]), np.stack([o[1] for o in outs])
)
self._record(xs, lps, np.stack([o[2] for o in outs]), n_samples)
yield from (xs[:, i] for i in range(n_samples))
return
xs = np.empty((C, n_samples, d))
lps = np.empty((C, n_samples))
st = np.empty((C, len(_STAT_KEYS), n_samples))
n_done = 0
try:
for i in progress(range(n_samples), n_samples, "NUTS sampling", self.progress):
for c, ch in enumerate(self._chains):
acc, depth, div, nl, energy = ch.step(self.max_depth, adapt=False)
xs[c, i] = ch.x
lps[c, i] = ch.lp
st[c, :, i] = (acc, depth, div, nl, energy)
xs[:, i], lps[:, i] = self._report(xs[:, i], lps[:, i])
n_done = i + 1
yield xs[:, i]
finally:
self._record(xs, lps, st, n_done)
finally:
self._close_pool()
[docs]
def run(self, n_samples: int = 1000) -> PopulationResult:
"""Warm up (first call only), draw ``n_samples`` per chain, and return all draws so far.
Parameters
----------
n_samples : int, optional
Draws per chain after warm-up.
"""
for _ in self.sample(int(n_samples)):
pass
return self.result
@property
def result(self) -> PopulationResult:
"""A :class:`~binference.PopulationResult` with every draw recorded so far."""
if not self._x:
raise RuntimeError("No draws yet; call run().")
X = np.concatenate(self._x, axis=1)
LP = np.concatenate(self._lp, axis=1)
C, n, d = X.shape
base = self._base_model
info: dict[str, Any] = {} if self._lam is None else {"centring": self._lam}
vals = base.transform(X.reshape(-1, d)).reshape(C, n, -1)
stats = {k: np.concatenate(v, axis=1) for k, v in self._stats.items()}
info.update(stats)
info["n_divergent"] = int(stats["diverging"].sum())
info["warmup_divergences"] = self.warmup_divergences
info["max_depth_fraction"] = float(np.mean(stats["tree_depth"] >= self.max_depth))
info["step_size"] = np.array([ch.eps for ch in self._chains])
info["inv_metric"] = [ch.metric.inv_m for ch in self._chains]
E = stats["energy"]
with np.errstate(invalid="ignore", divide="ignore"):
ebfmi = np.sum(np.diff(E, axis=1) ** 2, axis=1) / np.sum(
(E - E.mean(axis=1, keepdims=True)) ** 2, axis=1
)
info["ebfmi"] = ebfmi
info.update(base._joint_info(X))
hyper = getattr(base.prior, "hyper_names", [])
if hyper:
u = X[..., base.prior.n_active : base.prior.dim]
info["hyperparameters"] = {name: np.exp(u[..., j]) for j, name in enumerate(hyper)}
return PopulationResult(
base,
vals,
method="nuts",
x=X,
log_density_unconstrained=LP,
info=info,
)
