From 1459734a4704dbf0b142a7d23ac102e1ce65c1d7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 16:28:11 +0000 Subject: [PATCH 1/2] feat: re-mainline af.NSS as a first-class search on mainline blackjax >=1.6 Restores the NSS nested slice sampler removed in #1356/#1357 (the removal reason -- unshippable git-fork pins -- is gone: nested sampling merged into mainline blackjax in release 1.6). Implementation restored verbatim from the autofit_workspace_developer/searches/nss/ stash with the mainline port: - from nss.ns import log_weights, finalise -> blackjax.ns.utils; the import guard now checks blackjax >= 1.6 instead of the fork packages. - SliceInfo eval counter: fork num_steps -> mainline num_expansions. - Chunked GPU-memory path recomposed from mainline public helpers (build_slice_kernel / slice_constrained_step / build_adaptive_kernel / live_covariance) since mainline build_kernel has no update_strategy kwarg; _chunked_update.py unchanged (already byte-equivalent to mainline's update_with_mcmc_take_last bar the vmap->lax.map swap). - Tests restored verbatim (one ImportError match-string update); NSS re-exported from autofit/__init__.py; optional extra blackjax floor bumped 1.2.0 -> 1.6.2 (no cap existed). Validation (CPU): 17/17 nss tests pass from the checkout; end-to-end 2D analytic toy through search.fit gives logZ -4.558 +/- 0.078 vs analytic -4.605 with truth recovered, and the chunked path is bit-identical to unchunked at fixed seed (same logZ/evals/ESS). Closes #1491. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_018dV9h5h76hZ6trGZKWJFHe --- autofit/__init__.py | 1 + .../non_linear/search/nest/nss/__init__.py | 1 + .../search/nest/nss/_chunked_nss.py | 138 ++++ .../search/nest/nss/_chunked_update.py | 117 ++++ autofit/non_linear/search/nest/nss/samples.py | 25 + autofit/non_linear/search/nest/nss/search.py | 607 ++++++++++++++++++ pyproject.toml | 2 +- .../non_linear/search/nest/nss/__init__.py | 0 .../search/nest/nss/test_checkpoint.py | 243 +++++++ .../non_linear/search/nest/nss/test_search.py | 252 ++++++++ 10 files changed, 1385 insertions(+), 1 deletion(-) create mode 100644 autofit/non_linear/search/nest/nss/__init__.py create mode 100644 autofit/non_linear/search/nest/nss/_chunked_nss.py create mode 100644 autofit/non_linear/search/nest/nss/_chunked_update.py create mode 100644 autofit/non_linear/search/nest/nss/samples.py create mode 100644 autofit/non_linear/search/nest/nss/search.py create mode 100644 test_autofit/non_linear/search/nest/nss/__init__.py create mode 100644 test_autofit/non_linear/search/nest/nss/test_checkpoint.py create mode 100644 test_autofit/non_linear/search/nest/nss/test_search.py diff --git a/autofit/__init__.py b/autofit/__init__.py index 4c6b2ad67..dd64afc99 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -95,6 +95,7 @@ from .non_linear.search.mcmc.emcee.search import Emcee from .non_linear.search.mcmc.zeus.search import Zeus from .non_linear.search.nest.nautilus.search import Nautilus +from .non_linear.search.nest.nss.search import NSS from .non_linear.search.nest.dynesty.search.dynamic import DynestyDynamic from .non_linear.search.nest.dynesty.search.static import DynestyStatic from .non_linear.search.mle.drawer.search import Drawer diff --git a/autofit/non_linear/search/nest/nss/__init__.py b/autofit/non_linear/search/nest/nss/__init__.py new file mode 100644 index 000000000..76f18ac49 --- /dev/null +++ b/autofit/non_linear/search/nest/nss/__init__.py @@ -0,0 +1 @@ +from .search import NSS diff --git a/autofit/non_linear/search/nest/nss/_chunked_nss.py b/autofit/non_linear/search/nest/nss/_chunked_nss.py new file mode 100644 index 000000000..b7f360392 --- /dev/null +++ b/autofit/non_linear/search/nest/nss/_chunked_nss.py @@ -0,0 +1,138 @@ +"""Local replica of ``blackjax.ns.nss.as_top_level_api`` with chunked vmap sites. + +PyAutoFit#1303 added ``make_chunked_update_strategy`` for the per-iteration +MCMC step path inside ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``. +That covers the inner vmap over ``num_delete`` particles, but **not** the +separate hardcoded ``jax.vmap(init_state_fn)`` inside +``blackjax.ns.nss.as_top_level_api``'s ``init_fn`` — and that's where +inversion-heavy lensing cells (PyAutoLens pixelization / Delaunay at HST +scale) OOM on A100 80 GB before the sampling loop even starts. + +This module replaces ``blackjax.nss(...)`` with a local builder so we +control both seams: + +- The chunked update strategy replaces the hardcoded + ``update_with_mcmc_take_last`` inside the kernel composition (mainline + blackjax >= 1.6 exposes no ``update_strategy=`` kwarg — the pre-1.6 fork + did — so the kernel is composed here from mainline's public helpers: + ``build_slice_kernel`` / ``slice_constrained_step`` / + ``build_adaptive_kernel``). +- The chunked init swaps ``jax.vmap(init_state_fn)`` for + ``jax.lax.map(init_state_fn, positions, batch_size=chunk_size)`` so peak + GPU memory during ``algo.init`` becomes ``chunk_size × per_particle_state`` + instead of ``n_live × per_particle_state``. + +When ``chunk_size`` is None the builder still uses ``jax.vmap`` and is +bit-identical to upstream. The builder otherwise produces a +``blackjax.SamplingAlgorithm`` with the same shape ``blackjax.nss(...)`` +returns, so ``af.NSS._fit`` is a one-line switch. + +See PyAutoFit#1304 for the diagnosis and A100 evidence (jobs 322605 / +322606 OOM at the same byte counts as before #1303 landed, because the +crash is in ``algo.init`` not ``algo.step``). +""" + +from __future__ import annotations + +from functools import partial +from typing import Callable, Optional + + +def build_chunked_nss_algorithm( + *, + logprior_fn: Callable, + loglikelihood_fn: Callable, + num_inner_steps: int, + num_delete: int, + chunk_size: Optional[int], +): + """Return a ``blackjax.SamplingAlgorithm`` with chunked init + step paths. + + Replicates the body of ``blackjax.ns.nss.as_top_level_api`` (the + handley-lab fork) so we can plug in chunked variants of the two + vmap sites — the inner MCMC step (via the existing + ``make_chunked_update_strategy``) and the n_live-wide init. + + Parameters + ---------- + logprior_fn + Log-prior callable, ``positions -> scalar log-prior``. + loglikelihood_fn + Log-likelihood callable, ``positions -> scalar log-L``. + num_inner_steps + Number of HRSS steps per particle replacement (matches + ``af.NSS.num_mcmc_steps``). + num_delete + Number of particles replaced per outer iteration (matches + ``af.NSS.num_delete``). + chunk_size + Optional GPU-memory knob. When None, both vmap sites use plain + ``jax.vmap`` and the result is bit-identical to upstream + ``blackjax.nss(...)``. When set, peak memory in each site becomes + ``chunk_size × per_particle_state``. + """ + # Local imports keep this module cheap to import when ``af.NSS`` is + # never used (blackjax + jax are optional deps). + import jax + from blackjax import SamplingAlgorithm + from blackjax.ns.adaptive import init as ns_init + from blackjax.ns.base import init_state_strategy + from blackjax.ns.from_mcmc import build_adaptive_kernel, default_delete_fn + from blackjax.ns.nss import ( + build_slice_kernel, + covariance_proposal, + live_covariance, + slice_constrained_step, + stepping_out, + ) + + from ._chunked_update import ( + make_chunked_update_strategy, + ) + + init_state_fn = partial( + init_state_strategy, + logprior_fn=logprior_fn, + loglikelihood_fn=loglikelihood_fn, + ) + + # Replicates mainline ``blackjax.ns.nss.build_kernel`` + + # ``ns.from_mcmc.build_kernel`` with the hardcoded + # ``update_with_mcmc_take_last`` swapped for the chunked variant. + # (The pre-1.6 fork exposed an ``update_strategy=`` kwarg for this; + # mainline does not, hence the explicit composition.) + slice_kernel = build_slice_kernel( + interval=stepping_out, + max_expansions=10, + max_shrinkage=100, + ) + constrained_step_fn = slice_constrained_step( + init_state_fn, slice_kernel, covariance_proposal + ) + inner_kernel = make_chunked_update_strategy(chunk_size)( + constrained_step_fn, num_inner_steps, num_delete + ) + kernel = build_adaptive_kernel( + partial(default_delete_fn, num_delete=num_delete), + inner_kernel, + update_inner_kernel_params_fn=live_covariance, + ) + + def init_fn(position, rng_key=None): + if chunk_size is None: + init_batcher = jax.vmap(init_state_fn) + else: + init_batcher = lambda p: jax.lax.map( + init_state_fn, p, batch_size=chunk_size + ) + return ns_init( + position, + init_state_fn=init_batcher, + update_inner_kernel_params_fn=live_covariance, + rng_key=rng_key, + ) + + def step_fn(rng_key, state): + return kernel(rng_key, state) + + return SamplingAlgorithm(init_fn, step_fn) diff --git a/autofit/non_linear/search/nest/nss/_chunked_update.py b/autofit/non_linear/search/nest/nss/_chunked_update.py new file mode 100644 index 000000000..4f9b57f81 --- /dev/null +++ b/autofit/non_linear/search/nest/nss/_chunked_update.py @@ -0,0 +1,117 @@ +"""Chunked replacement for ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``. + +Upstream blackjax fans out ``num_delete`` particles through ``jax.vmap`` +with no chunking: + + sample_keys = jax.random.split(sample_key, num_delete) + return jax.vmap(mcmc_kernel)(sample_keys, start_state) + +On inversion-heavy likelihoods (e.g. PyAutoLens pixelization / Delaunay +source models) the per-particle MCMC state plus scatter temp buffers +exceeds A100 80 GB even at ``num_delete=16``. See PyAutoFit#1301 for the +full diagnosis and per-cell evidence from ``autolens_profiling``. + +``chunked_update_with_mcmc_take_last`` accepts a ``chunk_size`` kwarg and +swaps the vmap for ``jax.lax.map(..., batch_size=chunk_size)`` when +``chunk_size < num_delete`` — same vmap parallelism within a chunk, sequential +chunks across. Peak memory becomes ``chunk_size × per_particle_state`` +instead of ``num_delete × per_particle_state``. + +When ``chunk_size`` is None or ``>= num_delete`` the function is +bit-identical to upstream. + +Mainline blackjax (>= 1.6) hardcodes ``update_with_mcmc_take_last`` +inside ``blackjax.ns.from_mcmc.build_kernel`` — there is no +``update_strategy=`` kwarg (the pre-1.6 fork had one). The chunked +variant is therefore wired in by ``_chunked_nss.build_chunked_nss_algorithm``, +which recomposes the kernel from mainline's public helpers. +""" + +from __future__ import annotations + +from functools import partial +from typing import Callable, Optional + + +def make_chunked_update_strategy(chunk_size: Optional[int]) -> Callable: + """Return a chunked replacement for ``update_with_mcmc_take_last``. + + Signature matches ``blackjax.ns.from_mcmc.update_with_mcmc_take_last`` + so ``_chunked_nss.build_chunked_nss_algorithm`` can drop it into the + kernel composition unmodified. + + Parameters + ---------- + chunk_size + Number of particles to vmap-batch per chunk. When None or + ``>= num_delete`` the chunked path is skipped and the function + falls through to a plain ``jax.vmap`` (matching upstream + behaviour bit-for-bit). + """ + + def chunked_update_with_mcmc_take_last( + constrained_mcmc_step_fn, + num_mcmc_steps, + num_delete, + ): + """Drop-in for ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``. + + Identical to upstream except the inner + ``jax.vmap(mcmc_kernel)(sample_keys, start_state)`` is replaced + with ``jax.lax.map(..., batch_size=chunk_size)`` when + ``chunk_size`` is set and smaller than ``num_delete``. + """ + import jax + import jax.numpy as jnp + + def update_function(rng_key, state, loglikelihood_0, **step_parameters): + choice_key, sample_key = jax.random.split(rng_key) + particles = state.particles + + # Select start particles from survivors (verbatim from upstream). + weights = (particles.loglikelihood > loglikelihood_0).astype(jnp.float32) + weights = jnp.where(weights.sum() > 0.0, weights, jnp.ones_like(weights)) + start_idx = jax.random.choice( + choice_key, + len(weights), + shape=(num_delete,), + p=weights / weights.sum(), + replace=True, + ) + start_state = jax.tree.map(lambda x: x[start_idx], particles) + + shared_mcmc_step_fn = partial( + constrained_mcmc_step_fn, + loglikelihood_0=loglikelihood_0, + **step_parameters, + ) + + def mcmc_kernel(rng_key, state): + keys = jax.random.split(rng_key, num_mcmc_steps) + + def body_fn(state, rng_key): + new_state, info = shared_mcmc_step_fn(rng_key, state) + return new_state, info + + final_state, infos = jax.lax.scan(body_fn, state, keys) + return final_state, infos + + sample_keys = jax.random.split(sample_key, num_delete) + + # Fall through to bit-identical upstream behaviour when the + # user hasn't asked for chunking, or when the requested chunk + # already covers every particle. + if chunk_size is None or chunk_size >= num_delete: + return jax.vmap(mcmc_kernel)(sample_keys, start_state) + + # Chunked path: jax.lax.map(batch_size=k) vmaps within each + # chunk-of-k particles and loops across chunks. + return jax.lax.map( + lambda xs: mcmc_kernel(xs[0], xs[1]), + (sample_keys, start_state), + batch_size=chunk_size, + ) + + return update_function + + return chunked_update_with_mcmc_take_last diff --git a/autofit/non_linear/search/nest/nss/samples.py b/autofit/non_linear/search/nest/nss/samples.py new file mode 100644 index 000000000..b63b78376 --- /dev/null +++ b/autofit/non_linear/search/nest/nss/samples.py @@ -0,0 +1,25 @@ +from autofit.non_linear.samples.nest import SamplesNest + + +class NSSamples(SamplesNest): + """Posterior samples from an ``af.NSS`` (Nested Slice Sampling) fit. + + Thin subclass of ``SamplesNest`` that exists primarily for type identity — + aggregator code and downstream consumers can distinguish an NSS run from + a Nautilus / Dynesty run by ``isinstance(samples, NSSamples)``. The + actual posterior + log-evidence wiring is inherited from ``SamplesNest`` + (and ultimately ``SamplesPDF``); ``af.NSS`` builds the ``sample_list`` in + ``NSS.samples_via_internal_from``. + """ + + @property + def log_evidence_error(self) -> float: + """The stochastic batch-error of the NS evidence estimate. + + NSS returns an array of log-evidence estimates ``logZs`` across the + live ensemble (the Monte Carlo simulation of ``log_dX``). This + property exposes the standard deviation as the natural per-sample + uncertainty. The mean is reported as ``log_evidence`` in + ``samples_info``. + """ + return float(self.samples_info.get("log_evidence_error", float("nan"))) diff --git a/autofit/non_linear/search/nest/nss/search.py b/autofit/non_linear/search/nest/nss/search.py new file mode 100644 index 000000000..1bdb889d5 --- /dev/null +++ b/autofit/non_linear/search/nest/nss/search.py @@ -0,0 +1,607 @@ +import logging +import os +import pickle +from pathlib import Path +from typing import Optional + +import numpy as np + +from autofit.database.sqlalchemy_ import sa +from autofit.mapper.prior_model.abstract import AbstractPriorModel +from autofit.non_linear.fitness import Fitness +from autofit.non_linear.paths.null import NullPaths +from autofit.non_linear.search.nest import abstract_nest +from .samples import NSSamples +from autofit.non_linear.samples.sample import Sample +from autofit.non_linear.test_mode import is_test_mode + + +try: + import blackjax as _blackjax + from blackjax.ns.utils import ( + log_weights as _nss_log_weights, + finalise as _nss_finalise, + ) + + _HAS_NSS = tuple( + int(x) for x in _blackjax.__version__.split(".")[:2] + ) >= (1, 6) +except ImportError: + _blackjax = None + _nss_log_weights = None + _nss_finalise = None + _HAS_NSS = False + + +logger = logging.getLogger(__name__) + + +_CHECKPOINT_FILENAME = "nss_checkpoint.pkl" + + +def _save_checkpoint(path, state, dead, run_key, iteration): + """Atomically pickle the resumable state of an in-flight NSS run. + + The blackjax ``state`` and each entry in ``dead`` are pytrees of JAX arrays. + JAX arrays do pickle directly, but to keep the on-disk format independent + of the JAX install (so a checkpoint written on one cluster can be loaded + on another) we round-trip through NumPy before writing. ``_load_checkpoint`` + reverses the conversion. + + A tmp-and-rename pattern guards against partial writes — a SLURM timeout + halfway through a pickle dump leaves the previous-good checkpoint intact. + """ + import jax + + to_numpy = lambda x: jax.tree_util.tree_map(np.asarray, x) + blob = { + "state": to_numpy(state), + "dead": [to_numpy(d) for d in dead], + "run_key": np.asarray(run_key), + "iteration": int(iteration), + } + tmp_path = Path(str(path) + ".tmp") + with open(tmp_path, "wb") as f: + pickle.dump(blob, f) + os.replace(tmp_path, path) + + +def _load_checkpoint(path): + """Reverse of ``_save_checkpoint`` — restore a saved blob to JAX pytrees.""" + import jax + import jax.numpy as jnp + + with open(path, "rb") as f: + blob = pickle.load(f) + to_jax = lambda x: jax.tree_util.tree_map(jnp.asarray, x) + return ( + to_jax(blob["state"]), + [to_jax(d) for d in blob["dead"]], + jnp.asarray(blob["run_key"]), + int(blob["iteration"]), + ) + + +class _NSSInternal: + """Container holding the post-run state of the blackjax NSS loop. + + The ``NonLinearSearch`` pipeline stores this object as ``search_internal`` + (via ``paths.save_search_internal``) and passes it back into + ``samples_via_internal_from`` to extract the posterior. Keeping it as a + plain dataclass-style holder rather than the bare ``nss`` return tuple + means we can unpickle it later without depending on JAX (the + ``final_state.particles.position`` array is stored as a NumPy view). + """ + + def __init__( + self, + positions: np.ndarray, + loglikelihoods: np.ndarray, + log_weights: np.ndarray, + logZs: np.ndarray, + wall_time: float, + sampling_time: float, + evals: int, + ess: int, + n_live: int, + num_mcmc_steps: int, + num_delete: int, + termination: float, + seed: int, + ): + self.positions = positions + self.loglikelihoods = loglikelihoods + self.log_weights = log_weights + self.logZs = logZs + self.wall_time = wall_time + self.sampling_time = sampling_time + self.evals = evals + self.ess = ess + self.n_live = n_live + self.num_mcmc_steps = num_mcmc_steps + self.num_delete = num_delete + self.termination = termination + self.seed = seed + + +class NSS(abstract_nest.AbstractNest): + __identifier_fields__ = ( + "n_live", + "num_mcmc_steps", + "num_delete", + "termination", + "seed", + ) + + def __init__( + self, + name: Optional[str] = None, + path_prefix: Optional[str] = None, + unique_tag: Optional[str] = None, + n_live: int = 200, + num_mcmc_steps: int = 5, + num_delete: int = 50, + chunk_size: Optional[int] = None, + termination: float = -3.0, + checkpoint_interval: int = 100, + iterations_per_quick_update: Optional[int] = None, + iterations_per_full_update: Optional[int] = None, + number_of_cores: int = 1, + seed: int = 42, + silence: bool = False, + session: Optional[sa.orm.Session] = None, + **kwargs, + ): + """ + A Nested Slice Sampling non-linear search (JAX-native, ``blackjax.nss``). + + ``af.NSS`` wraps mainline blackjax's nested slice sampler + (``blackjax.nss``, merged upstream in blackjax 1.6) into a first-class + ``NonLinearSearch`` so users can drop ``search = af.NSS(...)`` into any + production autolens / autogalaxy / autofit script alongside + ``af.Nautilus(...)``. The likelihood and prior closures both run inside + ``jax.jit``, leveraging the Phase 0 JAX-native priors (PyAutoFit#1262) + to make ``model.vector_from_unit_vector`` and + ``model.log_prior_list_from_vector`` traceable. + + Phases 1-3 of the ``nss_first_class_sampler`` roadmap are live: + - Phase 1: the wrapper itself (this class). + - Phase 2: checkpoint/resume via ``checkpoint_interval`` — a + ``nss_checkpoint.pkl`` is written to ``paths.search_internal_path`` + every N outer iterations and reloaded automatically on resume. + - Phase 3: on-the-fly visualization via ``iterations_per_quick_update`` + — every N outer iterations the current best live particle is fed to + ``analysis.visualize`` so partial results appear in the image_path + directory during long fits. + + Nested sampling merged into mainline blackjax in release 1.6, so the + only requirement is ``pip install 'blackjax>=1.6.2'`` (covered by the + ``autofit[optional]`` extra) — the pinned-git-fork install that forced + this search's removal in PyAutoFit#1356 is gone. + + Parameters + ---------- + name + The name of the search, controlling the last folder results are output. + path_prefix + The path of folders prefixing the name folder where results are output. + unique_tag + A unique tag for this model-fit, given a unique entry in the + sqlite database and used as the folder after the path prefix and + before the search name. + n_live + Number of live particles maintained by NSS. Production-tested + default of 200 corresponds to FINDINGS_v3's ``c3_big_delete`` + config on the HST MGE problem (recovered ``einstein_radius=1.5996`` + in 7 min wall time). + num_mcmc_steps + Number of slice-MCMC inner steps per dead-point batch. + num_delete + Number of particles killed per outer iteration. Larger + ``num_delete`` reduces JIT overhead per iteration at the cost of + slightly worse posterior coverage. + chunk_size + Optional GPU-memory knob. When set and ``< num_delete``, the + inner MCMC step vmap (which fans out ``num_delete`` particles + in parallel inside ``blackjax.ns.from_mcmc.update_with_mcmc_take_last``) + is replaced with ``jax.lax.map(..., batch_size=chunk_size)``. + Peak GPU memory becomes ``chunk_size × per_particle_state`` + instead of ``num_delete × per_particle_state`` — required to + run NSS on inversion-heavy likelihoods (PyAutoLens pixelization + / Delaunay) at A100 80 GB scale. Default ``None`` preserves the + upstream un-chunked behaviour (and is the right choice on CPU + or whenever ``num_delete`` already fits the device). + termination + Convergence criterion. The fit stops when + ``logZ_live - logZ < termination``. Default ``-3.0`` corresponds + to delta-logZ < 1e-3. + checkpoint_interval + Outer iterations between checkpoint writes. Default ``100`` writes + a ``nss_checkpoint.pkl`` (atomic via tmp-and-rename) every ~5000- + 10000 likelihood evaluations at typical ``num_delete=50``. Set to + a large value to effectively disable checkpointing on short runs. + iterations_per_quick_update + Outer iterations between on-the-fly visualizations. When non-None + the current best live particle is fed to ``analysis.visualize`` + every N iterations so partial results appear in the image_path + directory during long fits. ``analysis.visualize`` is wrapped in + try/except so a viz failure logs a warning but does not kill the + sampler. + iterations_per_full_update + Accepted for API parity. NSS does not have a full-update concept + separate from the outer-iteration cadence. + number_of_cores + Accepted for API parity only. NSS runs on whatever device JAX is + configured for (CPU, GPU, TPU) — multiprocessing parallelism is + not used. If ``number_of_cores > 1`` a warning is logged. + seed + JAX random seed used to initialise the unit-cube draws for the + ``n_live`` initial particles and the inner slice-sampling RNG. + silence + If True, suppresses NSS's progress bar output. + session + An SQLAlchemy session instance so the results of the model-fit + are written to an SQLite database. + + Notes + ----- + Sampling space: NSS samples in **physical** parameter space (initial + particles are drawn via ``model.vector_from_unit_vector`` per unit-cube + sample, then the slice walks proceed in physical coordinates). The + Nautilus / PocoMC convention is to sample in unit-cube space and + apply the prior transform inside the likelihood — that is **not** what + ``af.NSS`` does. The slice walks have the natural metric of the + problem and there is no per-step remapping cost. + """ + + if not _HAS_NSS: + raise ImportError( + "af.NSS requires blackjax >= 1.6 (nested sampling merged " + "into mainline blackjax in release 1.6). Install via:\n" + " pip install 'blackjax>=1.6.2'" + ) + + super().__init__( + name=name, + path_prefix=path_prefix, + unique_tag=unique_tag, + iterations_per_full_update=iterations_per_full_update, + iterations_per_quick_update=iterations_per_quick_update, + number_of_cores=number_of_cores, + silence=silence, + session=session, + **kwargs, + ) + + self.n_live = n_live + self.num_mcmc_steps = num_mcmc_steps + self.num_delete = num_delete + self.chunk_size = chunk_size + self.termination = termination + self.checkpoint_interval = checkpoint_interval + self.seed = seed + + if number_of_cores is not None and number_of_cores > 1: + logger.warning( + "af.NSS received number_of_cores=%d. NSS is JAX-native and " + "runs on whatever device JAX is configured for; the value is " + "ignored. Set number_of_cores=1 to silence this warning.", + number_of_cores, + ) + + if is_test_mode(): + self.apply_test_mode() + + self.logger.debug("Creating NSS Search") + + def apply_test_mode(self): + logger.warning( + "TEST MODE 1 (reduced iterations): NSS will run with a loose " + "termination criterion for faster completion." + ) + self.termination = -1.0 + + def _fit(self, model: AbstractPriorModel, analysis): + """ + Fit a model using NSS, with checkpoint/resume + on-the-fly visualization. + + Builds JAX-traceable ``log_likelihood`` and ``prior_logprob`` closures + threaded through Phase 0's ``xp=jnp`` plumbing, draws ``n_live`` initial + particles by mapping unit-cube samples through the prior transform, + and runs the NSS outer loop inline (mirroring upstream blackjax's + nested-sampling loop pattern). Between outer iterations the + loop can (a) pickle resumable state to ``nss_checkpoint.pkl`` and + (b) call ``analysis.visualize`` on the current best live particle. + + On entry, if a checkpoint exists at the expected path the loop resumes + from the saved ``(state, dead, run_key, iteration)``. On successful + exit the checkpoint is deleted — mirrors Nautilus's + ``output_search_internal`` post-success cleanup so the next fresh fit + doesn't accidentally resume from a stale checkpoint. + + Returns + ------- + (search_internal, fitness) + ``search_internal`` is a ``_NSSInternal`` holder (NumPy arrays + only). ``fitness`` is a ``Fitness`` instance that ``af.NSS`` does + not use for sampling (inline JAX closures handle that) but is + required by ``AbstractNest.perform_update`` for post-fit work + like latent-sample generation, which calls ``fitness.batch_size``. + """ + + import jax + import jax.numpy as jnp + import time + + self.logger.info("Starting NSS non-linear search.") + + ndim = model.prior_count + + def log_likelihood(params): + instance = model.instance_from_vector(vector=params, xp=jnp) + raw = analysis.log_likelihood_function(instance=instance) + return jnp.where(jnp.isfinite(raw), raw, -1e30) + + def prior_logprob(params): + log_priors = model.log_prior_list_from_vector(vector=params, xp=jnp) + return sum(log_priors) + + rng_key = jax.random.PRNGKey(self.seed) + rng_key, init_key, run_key = jax.random.split(rng_key, 3) + + unit_cube = jax.random.uniform(init_key, shape=(self.n_live, ndim)) + initial_samples = jnp.stack( + [ + model.vector_from_unit_vector(unit_cube[i], xp=jnp) + for i in range(self.n_live) + ] + ) + + # When ``chunk_size`` is set and below the wider of ``n_live`` / + # ``num_delete``, build the algorithm via the PyAutoFit-local + # ``build_chunked_nss_algorithm``. That replicates mainline + # ``blackjax.ns.nss.as_top_level_api`` with both vmap sites chunked: + # the inner MCMC step (PyAutoFit#1303's + # ``make_chunked_update_strategy``) **and** the n_live-wide + # ``algo.init`` (PyAutoFit#1304 — the OOM-causing site for + # inversion-heavy lensing cells). ``chunk_size=None`` or + # ``chunk_size >= max(n_live, num_delete)`` keeps using upstream + # ``blackjax.nss(...)`` bit-for-bit. + if self.chunk_size is not None and self.chunk_size < max( + self.n_live, self.num_delete + ): + from ._chunked_nss import ( + build_chunked_nss_algorithm, + ) + + algo = build_chunked_nss_algorithm( + logprior_fn=prior_logprob, + loglikelihood_fn=log_likelihood, + num_inner_steps=self.num_mcmc_steps, + num_delete=self.num_delete, + chunk_size=self.chunk_size, + ) + else: + algo = _blackjax.nss( + logprior_fn=prior_logprob, + loglikelihood_fn=log_likelihood, + num_delete=self.num_delete, + num_inner_steps=self.num_mcmc_steps, + ) + + @jax.jit + def one_step(carry, _): + state, k = carry + k, subk = jax.random.split(k, 2) + state, dead_point = algo.step(subk, state) + return (state, k), dead_point + + checkpoint_path = self._nss_checkpoint_path + if checkpoint_path is not None and checkpoint_path.exists(): + state, dead, run_key, iteration = _load_checkpoint(checkpoint_path) + self.logger.info( + "Resuming NSS from checkpoint at iteration %d (state file %s).", + iteration, + checkpoint_path, + ) + else: + state = algo.init(initial_samples) + dead = [] + iteration = 0 + self.logger.info( + "NSS configuration: n_live=%d, num_mcmc_steps=%d, num_delete=%d, " + "chunk_size=%s, termination=%s, ndim=%d, checkpoint_interval=%d. " + "JIT compile on first iteration may take 25-30 s.", + self.n_live, + self.num_mcmc_steps, + self.num_delete, + self.chunk_size, + self.termination, + ndim, + self.checkpoint_interval, + ) + + t_start = time.time() + while not state.integrator.logZ_live - state.integrator.logZ < self.termination: + (state, run_key), dead_info = one_step((state, run_key), None) + dead.append(dead_info) + iteration += 1 + + if ( + checkpoint_path is not None + and iteration % self.checkpoint_interval == 0 + ): + _save_checkpoint(checkpoint_path, state, dead, run_key, iteration) + + if ( + self.iterations_per_quick_update is not None + and iteration % self.iterations_per_quick_update == 0 + ): + self._fire_quick_update(state=state, model=model, analysis=analysis) + + wall_time = time.time() - t_start + + final_state = _nss_finalise(state, dead) + + rng_key, weight_key = jax.random.split(rng_key, 2) + log_w_mc = _nss_log_weights(weight_key, final_state, shape=100) + log_w_per_particle = log_w_mc.mean(axis=-1) + logZs = jax.scipy.special.logsumexp( + jnp.nan_to_num(log_w_mc, nan=jnp.nan_to_num(log_w_mc).min()), + axis=0, + ) + + def _safe_ess(log_w_mean): + log_w_mean = log_w_mean - log_w_mean.max() + weights = jnp.exp(log_w_mean) + return float(weights.sum() ** 2 / (weights**2).sum()) + + ess = int(_safe_ess(log_w_mc.mean(axis=-1))) + # blackjax >= 1.6 SliceInfo: stepping-out expansions + shrinkage + # contractions are the two likelihood-evaluation counters (the fork's + # ``num_steps`` field is ``num_expansions`` in mainline). + evals = int( + final_state.update_info.num_expansions.sum() + + final_state.update_info.num_shrink.sum() + ) + + search_internal = _NSSInternal( + positions=np.asarray(final_state.particles.position), + loglikelihoods=np.asarray(final_state.particles.loglikelihood), + log_weights=np.asarray(log_w_per_particle), + logZs=np.asarray(logZs), + wall_time=float(wall_time), + sampling_time=float(wall_time), + evals=evals, + ess=ess, + n_live=int(self.n_live), + num_mcmc_steps=int(self.num_mcmc_steps), + num_delete=int(self.num_delete), + termination=float(self.termination), + seed=int(self.seed), + ) + + if checkpoint_path is not None and checkpoint_path.exists(): + try: + checkpoint_path.unlink() + except OSError as exc: + self.logger.warning( + "Failed to delete completed-run checkpoint %s: %s. The " + "next fresh af.NSS fit at this path will attempt to resume " + "from it — delete manually if that is not desired.", + checkpoint_path, + exc, + ) + + fitness = Fitness( + model=model, + analysis=analysis, + paths=self.paths, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, + batch_size=1, + ) + + return search_internal, fitness + + @property + def _nss_checkpoint_path(self) -> Optional[Path]: + """Resolve the checkpoint location, or None when paths is NullPaths.""" + if isinstance(self.paths, NullPaths): + return None + try: + return Path(self.paths.search_internal_path) / _CHECKPOINT_FILENAME + except TypeError: + return None + + def _fire_quick_update(self, state, model, analysis): + """Push the current best live particle through ``analysis.visualize``. + + The Nautilus / Dynesty quick-update path goes through + ``Fitness.manage_quick_update``; ``af.NSS`` bypasses ``Fitness._call`` + for sampling so we invoke ``analysis.visualize`` directly between + outer-loop iterations. Wrapped in try/except — a visualization failure + logs a warning but does not kill a long sampler run. + """ + try: + best_idx = int(np.asarray(state.particles.loglikelihood).argmax()) + best_params = np.asarray(state.particles.position[best_idx]).tolist() + instance = model.instance_from_vector(vector=best_params) + analysis.visualize( + paths=self.paths, + instance=instance, + during_analysis=True, + ) + except Exception as exc: + self.logger.warning( + "af.NSS quick-update visualization failed: %s. Continuing the " + "fit — quick-update is best-effort, the final visualization " + "fires at the end of the run regardless.", + exc, + ) + + @property + def checkpoint_file(self): + """Path to the on-disk checkpoint written between outer-loop iterations. + + Returns the same value as ``_nss_checkpoint_path`` — exposed as a + public property for symmetry with ``af.Nautilus.checkpoint_file``. + """ + return self._nss_checkpoint_path + + def samples_info_from(self, search_internal: Optional[_NSSInternal] = None): + if search_internal is None: + search_internal = self.paths.load_search_internal() + return { + "log_evidence": float(np.asarray(search_internal.logZs).mean()), + "log_evidence_error": float(np.asarray(search_internal.logZs).std()), + "total_samples": int(search_internal.evals), + "total_accepted_samples": int(len(search_internal.positions)), + "time": float(search_internal.wall_time), + "sampling_time": float(search_internal.sampling_time), + "number_live_points": int(search_internal.n_live), + "num_mcmc_steps": int(search_internal.num_mcmc_steps), + "num_delete": int(search_internal.num_delete), + "termination": float(search_internal.termination), + "ess": int(search_internal.ess), + } + + def samples_via_internal_from( + self, + model: AbstractPriorModel, + search_internal: Optional[_NSSInternal] = None, + ): + """Convert the stored ``_NSSInternal`` holder into an ``NSSamples``.""" + + if search_internal is None: + search_internal = self.paths.load_search_internal() + + parameter_lists = np.asarray(search_internal.positions).tolist() + log_likelihood_list = np.asarray(search_internal.loglikelihoods).tolist() + + log_w = np.asarray(search_internal.log_weights) + log_w_norm = log_w - log_w.max() + weights = np.exp(log_w_norm) + weight_total = weights.sum() + if weight_total > 0: + weights = weights / weight_total + weight_list = weights.tolist() + + log_prior_list = [ + sum(model.log_prior_list_from_vector(vector=vector)) + for vector in parameter_lists + ] + + sample_list = Sample.from_lists( + model=model, + parameter_lists=parameter_lists, + log_likelihood_list=log_likelihood_list, + log_prior_list=log_prior_list, + weight_list=weight_list, + ) + + return NSSamples( + model=model, + sample_list=sample_list, + samples_info=self.samples_info_from(search_internal=search_internal), + ) diff --git a/pyproject.toml b/pyproject.toml index cfdb58ff2..7539f046a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ mcp = ["mcp"] optional = [ "autofit[jax]", "astropy>=5.0", - "blackjax>=1.2.0", + "blackjax>=1.6.2", "getdist==1.4", "nautilus-sampler==1.0.5", "zeus-mcmc==2.5.4", diff --git a/test_autofit/non_linear/search/nest/nss/__init__.py b/test_autofit/non_linear/search/nest/nss/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py new file mode 100644 index 000000000..5f32d72ea --- /dev/null +++ b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py @@ -0,0 +1,243 @@ +""" +Unit tests for the ``af.NSS`` checkpoint/resume helpers +(``_save_checkpoint`` / ``_load_checkpoint``). + +No real blackjax NSS loop calls — the heavy end-to-end resume +verification lives in +``autolens_workspace_developer/searches_minimal/nss_checkpoint_resume.py``. +""" + +from unittest.mock import patch + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.search.nest.nss import search as nss_search_module +from autofit.non_linear.search.nest.nss.search import ( + _load_checkpoint, + _save_checkpoint, +) + + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") +requires_nss = pytest.mark.skipif( + not nss_search_module._HAS_NSS, + reason="requires blackjax >= 1.6", +) + + +jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") + + +def _synthetic_state(): + """Plain-dict pytree mimicking the blackjax NSS state shape. + + We only need a pytree of JAX arrays — the round-trip serialiser doesn't + care about the type as long as ``jax.tree_util.tree_map`` can walk it. + Plain dicts are registered pytrees; SimpleNamespace is not. + """ + return { + "particles": { + "position": jnp.asarray([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), + "loglikelihood": jnp.asarray([-1.5, -0.5, -0.1]), + }, + "integrator": { + "logZ": jnp.float64(-12.3), + "logZ_live": jnp.float64(-11.8), + }, + } + + +def _synthetic_dead(n_iter=3): + return [ + { + "particles": { + "position": jnp.asarray([[float(i), float(i + 1)]]), + "loglikelihood": jnp.asarray([-float(i)]), + }, + } + for i in range(n_iter) + ] + + +def test__save_load_checkpoint_round_trip(tmp_path): + state = _synthetic_state() + dead = _synthetic_dead() + run_key = jax.random.PRNGKey(7) + + path = tmp_path / "nss_checkpoint.pkl" + _save_checkpoint(path, state, dead, run_key, iteration=42) + + assert path.exists() + loaded_state, loaded_dead, loaded_run_key, loaded_iter = _load_checkpoint(path) + + assert loaded_iter == 42 + assert np.array_equal(np.asarray(loaded_run_key), np.asarray(run_key)) + + np.testing.assert_array_equal( + np.asarray(loaded_state["particles"]["position"]), + np.asarray(state["particles"]["position"]), + ) + np.testing.assert_array_equal( + np.asarray(loaded_state["particles"]["loglikelihood"]), + np.asarray(state["particles"]["loglikelihood"]), + ) + np.testing.assert_array_equal( + np.asarray(loaded_state["integrator"]["logZ"]), + np.asarray(state["integrator"]["logZ"]), + ) + + assert len(loaded_dead) == len(dead) + for orig, loaded in zip(dead, loaded_dead): + np.testing.assert_array_equal( + np.asarray(loaded["particles"]["position"]), + np.asarray(orig["particles"]["position"]), + ) + + +def test__save_checkpoint_is_atomic(tmp_path): + """A partially-written checkpoint must not clobber the previous good one. + + The helper writes to ``.tmp`` then ``os.replace`` to the final path. + If the rename fails (simulated here by patching ``os.replace`` to raise), + the final path must be untouched — the previous-good blob still loads. + """ + state = _synthetic_state() + dead = _synthetic_dead() + run_key = jax.random.PRNGKey(3) + path = tmp_path / "nss_checkpoint.pkl" + + _save_checkpoint(path, state, dead, run_key, iteration=1) + good_blob = path.read_bytes() + + state["integrator"]["logZ"] = jnp.float64(-100.0) + with patch( + "autofit.non_linear.search.nest.nss.search.os.replace", + side_effect=OSError("simulated rename failure"), + ): + with pytest.raises(OSError): + _save_checkpoint(path, state, dead, run_key, iteration=2) + + assert path.read_bytes() == good_blob + + +@requires_nss +def test__nss_checkpoint_path_is_none_for_null_paths(): + """Without a real output dir (NullPaths), checkpoint resolution returns None. + + This means ``_fit``'s resume detection silently skips for NullPaths fits + (e.g. unit tests, in-memory aggregator round-trips) rather than blowing + up on the missing ``search_internal_path`` attribute. + """ + search = af.NSS() + assert search._nss_checkpoint_path is None + assert search.checkpoint_file is None + + +@requires_nss +def test__init_accepts_checkpoint_interval(): + search = af.NSS(checkpoint_interval=25) + assert search.checkpoint_interval == 25 + + default = af.NSS() + assert default.checkpoint_interval == 100 + + +@requires_nss +def test__init_iterations_per_quick_update_no_longer_warns(caplog): + """Phase 1's no-op log when the kwarg is set is gone in Phase 3 (the kwarg + is now actually wired). The warning text must not appear. + """ + with caplog.at_level("INFO"): + af.NSS(iterations_per_quick_update=10) + assert not any("not yet wired" in record.message for record in caplog.records) + + +@requires_nss +def test__load_checkpoint_called_when_file_exists(tmp_path): + """Verify ``_load_checkpoint`` is invoked from ``_fit`` when a checkpoint + file exists at the resolved path. + + We replace ``_blackjax.nss`` with a mock that fails loudly if ``algo.init`` + fires — i.e. if the fresh-init branch ran. The resume branch must call + ``_load_checkpoint`` first; we patch that helper to return a sentinel + that satisfies the immediate-termination logZ check, so the outer loop + exits before doing any real work. + """ + from types import SimpleNamespace + + fake_checkpoint = tmp_path / "nss_checkpoint.pkl" + fake_checkpoint.write_bytes(b"placeholder - actual contents replaced by mock") + + sentinel_state = SimpleNamespace( + integrator=SimpleNamespace(logZ=0.0, logZ_live=-100.0), + particles=SimpleNamespace( + position=jnp.zeros((4, 2)), + loglikelihood=jnp.zeros(4), + ), + ) + + # Minimum stub for model + analysis. _fit calls model.prior_count and + # both vector_from_unit_vector + log_prior_list_from_vector inside the + # closure construction (which isn't traced unless one_step fires). + mock_model = SimpleNamespace( + prior_count=2, + instance_from_vector=lambda **kw: SimpleNamespace(), + vector_from_unit_vector=lambda v, xp=None: jnp.asarray([0.0, 0.0]), + log_prior_list_from_vector=lambda **kw: [0.0, 0.0], + ) + mock_analysis = SimpleNamespace( + log_likelihood_function=lambda instance: 0.0, + ) + + search = af.NSS(n_live=4, num_mcmc_steps=1, num_delete=1, termination=-3.0) + # Force the checkpoint property to return our sentinel path even though + # paths is the default NullPaths. + with patch.object( + type(search), + "_nss_checkpoint_path", + new=fake_checkpoint, + ), patch.object( + nss_search_module, + "_load_checkpoint", + return_value=(sentinel_state, [], jax.random.PRNGKey(0), 17), + ) as mock_load, patch.object( + nss_search_module, + "_blackjax", + ) as mock_bjax, patch.object( + nss_search_module, + "_nss_finalise", + return_value=SimpleNamespace( + particles=SimpleNamespace( + position=np.zeros((1, 2)), + loglikelihood=np.zeros(1), + ), + update_info=SimpleNamespace( + num_steps=np.zeros(1, dtype=int), + num_shrink=np.zeros(1, dtype=int), + ), + ), + ), patch.object( + nss_search_module, + "_nss_log_weights", + return_value=jnp.zeros((1, 100)), + ): + mock_bjax.nss.return_value.init.side_effect = AssertionError( + "algo.init called from resume path — expected _load_checkpoint instead." + ) + mock_bjax.nss.return_value.step.side_effect = AssertionError( + "algo.step called even though logZ termination should fire immediately." + ) + + # The mocked downstream pipeline (Fitness construction, _NSSInternal + # repackaging) is intentionally not realistic — we only care that the + # resume branch was entered and ``_load_checkpoint`` was called. Catch + # any downstream stub-related failure; the assertion below is the gate. + try: + search._fit(model=mock_model, analysis=mock_analysis) + except (AttributeError, AssertionError, TypeError): + pass + + mock_load.assert_called_once_with(fake_checkpoint) diff --git a/test_autofit/non_linear/search/nest/nss/test_search.py b/test_autofit/non_linear/search/nest/nss/test_search.py new file mode 100644 index 000000000..e08bf40d4 --- /dev/null +++ b/test_autofit/non_linear/search/nest/nss/test_search.py @@ -0,0 +1,252 @@ +""" +Unit tests for ``af.NSS`` — initialisation, config, optional-import guard, +and samples round-trip via a synthetic ``_NSSInternal`` holder. + +No JAX in unit tests (library policy — cross-xp checks live in +``autofit_workspace_test``). The numerical-parity smoke against a real +``run_nested_sampling`` call lives in +``autolens_workspace_developer/searches_minimal/nss_first_class.py``. +""" + +import numpy as np +import pytest + +import autofit as af +from autofit.non_linear.search.nest.nss import search as nss_search_module +from autofit.non_linear.search.nest.nss.samples import NSSamples + + +pytestmark = pytest.mark.filterwarnings("ignore::FutureWarning") +requires_nss = pytest.mark.skipif( + not nss_search_module._HAS_NSS, + reason="requires blackjax >= 1.6", +) + + +@requires_nss +def test__explicit_params(): + search = af.NSS( + n_live=500, + num_mcmc_steps=10, + num_delete=20, + chunk_size=4, + termination=-2.0, + seed=7, + ) + + assert search.n_live == 500 + assert search.num_mcmc_steps == 10 + assert search.num_delete == 20 + assert search.chunk_size == 4 + assert search.termination == -2.0 + assert search.seed == 7 + + default = af.NSS() + assert default.n_live == 200 + assert default.num_mcmc_steps == 5 + assert default.num_delete == 50 + assert default.chunk_size is None + assert default.termination == -3.0 + assert default.seed == 42 + + +def test__chunked_update_strategy_factory(): + """``make_chunked_update_strategy`` returns a callable with the same + signature as blackjax's ``update_with_mcmc_take_last`` regardless of + whether ``chunk_size`` is set. This lets ``build_chunked_nss_algorithm`` + drop it into the kernel composition without further branching. + """ + from autofit.non_linear.search.nest.nss._chunked_update import ( + make_chunked_update_strategy, + ) + + strategy_none = make_chunked_update_strategy(None) + strategy_chunked = make_chunked_update_strategy(4) + # Both are callables with the upstream three-arg signature + # (constrained_mcmc_step_fn, num_mcmc_steps, num_delete). + import inspect + + for strategy in (strategy_none, strategy_chunked): + params = list(inspect.signature(strategy).parameters) + assert params == [ + "constrained_mcmc_step_fn", + "num_mcmc_steps", + "num_delete", + ] + + +@requires_nss +def test__chunked_nss_algorithm_factory(): + """``build_chunked_nss_algorithm`` returns a ``blackjax.SamplingAlgorithm``- + shape NamedTuple with ``init`` and ``step`` attributes. This is what + ``af.NSS._fit`` relies on as a drop-in for ``_blackjax.nss(...)``. + + No JAX execution — library policy keeps JAX-traced tests in + ``autofit_workspace_test``. + """ + blackjax = pytest.importorskip("blackjax") + from autofit.non_linear.search.nest.nss._chunked_nss import ( + build_chunked_nss_algorithm, + ) + + algo = build_chunked_nss_algorithm( + logprior_fn=lambda p: 0.0, + loglikelihood_fn=lambda p: 0.0, + num_inner_steps=3, + num_delete=4, + chunk_size=2, + ) + + assert isinstance(algo, blackjax.SamplingAlgorithm) + assert callable(algo.init) + assert callable(algo.step) + + # chunk_size=None → still returns the same shape (unchunked path). + algo_unchunked = build_chunked_nss_algorithm( + logprior_fn=lambda p: 0.0, + loglikelihood_fn=lambda p: 0.0, + num_inner_steps=3, + num_delete=4, + chunk_size=None, + ) + assert isinstance(algo_unchunked, blackjax.SamplingAlgorithm) + + +@requires_nss +def test__identifier_fields(): + search = af.NSS() + for field in ("n_live", "num_mcmc_steps", "num_delete", "termination", "seed"): + assert field in search.__identifier_fields__ + + +@requires_nss +def test__test_mode_loosens_termination(): + search = af.NSS(termination=-3.0) + search.apply_test_mode() + assert search.termination == -1.0 + + +def test__init_raises_when_nss_unavailable(monkeypatch): + monkeypatch.setattr(nss_search_module, "_HAS_NSS", False) + with pytest.raises(ImportError, match="af.NSS requires blackjax >= 1.6"): + af.NSS() + + +@requires_nss +def test__init_warns_when_number_of_cores_gt_one(caplog): + with caplog.at_level("WARNING"): + af.NSS(number_of_cores=4) + assert any( + "number_of_cores=4" in record.message and "ignored" in record.message + for record in caplog.records + ) + + +def _make_synthetic_internal(n_live=20, ndim=2, seed=0): + rng = np.random.default_rng(seed) + positions = rng.normal(loc=0.0, scale=1.0, size=(n_live, ndim)) + loglikelihoods = rng.normal(loc=0.0, scale=0.5, size=(n_live,)) + log_weights = rng.normal(loc=-5.0, scale=0.1, size=(n_live,)) + logZs = rng.normal(loc=-10.0, scale=0.1, size=(100,)) + + return nss_search_module._NSSInternal( + positions=positions, + loglikelihoods=loglikelihoods, + log_weights=log_weights, + logZs=logZs, + wall_time=12.5, + sampling_time=10.0, + evals=5000, + ess=150, + n_live=n_live, + num_mcmc_steps=5, + num_delete=10, + termination=-3.0, + seed=42, + ) + + +@requires_nss +def test__samples_info_from_synthetic_internal(): + search = af.NSS() + internal = _make_synthetic_internal() + + info = search.samples_info_from(search_internal=internal) + + assert info["log_evidence"] == pytest.approx( + float(internal.logZs.mean()), abs=1e-12 + ) + assert info["log_evidence_error"] == pytest.approx( + float(internal.logZs.std()), abs=1e-12 + ) + assert info["total_samples"] == 5000 + assert info["total_accepted_samples"] == 20 + assert info["number_live_points"] == 20 + assert info["num_mcmc_steps"] == 5 + assert info["num_delete"] == 10 + assert info["termination"] == -3.0 + assert info["ess"] == 150 + assert info["time"] == pytest.approx(12.5, abs=1e-12) + assert info["sampling_time"] == pytest.approx(10.0, abs=1e-12) + + +@requires_nss +def test__samples_via_internal_from_returns_nssamples(): + model = af.Model(af.m.MockClassx2) + model.one = af.UniformPrior(lower_limit=-3.0, upper_limit=3.0) + model.two = af.UniformPrior(lower_limit=-3.0, upper_limit=3.0) + + search = af.NSS() + internal = _make_synthetic_internal(n_live=20, ndim=model.prior_count) + + samples = search.samples_via_internal_from(model=model, search_internal=internal) + + assert isinstance(samples, NSSamples) + assert len(samples.sample_list) == 20 + assert samples.model is model + + weights = np.array([s.weight for s in samples.sample_list]) + assert weights.sum() == pytest.approx(1.0, abs=1e-10) + assert (weights >= 0).all() + + assert samples.samples_info["log_evidence"] == pytest.approx( + float(internal.logZs.mean()), abs=1e-12 + ) + + +@requires_nss +def test__samples_weight_normalisation_handles_zero_total(): + """When every log_weight is -inf (e.g. catastrophic underflow) we should + not divide by zero. The implementation clamps via max-subtraction so this + case should still produce zero-weights without crashing.""" + model = af.Model(af.m.MockClassx2) + model.one = af.UniformPrior(lower_limit=-3.0, upper_limit=3.0) + model.two = af.UniformPrior(lower_limit=-3.0, upper_limit=3.0) + + search = af.NSS() + internal = _make_synthetic_internal(n_live=10, ndim=model.prior_count) + # Sentinel: all log-weights identical → exp(log_w - max) = 1 for all, + # weights normalise uniformly to 1/n. + internal.log_weights = np.full_like(internal.log_weights, -3.14) + + samples = search.samples_via_internal_from(model=model, search_internal=internal) + + weights = np.array([s.weight for s in samples.sample_list]) + assert weights.sum() == pytest.approx(1.0, abs=1e-12) + assert np.allclose(weights, 1.0 / 10, atol=1e-12) + + +@requires_nss +def test__nssamples_log_evidence_error_property(): + model = af.Model(af.m.MockClassx2) + model.one = af.UniformPrior(lower_limit=-3.0, upper_limit=3.0) + model.two = af.UniformPrior(lower_limit=-3.0, upper_limit=3.0) + + search = af.NSS() + internal = _make_synthetic_internal(n_live=10, ndim=model.prior_count) + + samples = search.samples_via_internal_from(model=model, search_internal=internal) + + assert samples.log_evidence_error == pytest.approx( + float(internal.logZs.std()), abs=1e-12 + ) From 82484d7bdfaff95a0370f090b558f1173af2892b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 16:34:06 +0000 Subject: [PATCH 2/2] test: exempt af.NSS from the quick-update Fitness-wiring scan af.NSS samples through inline JAX closures, never through Fitness -- its Fitness exists only for AbstractNest.perform_update post-fit work, so forwarding iterations_per_quick_update there would wire the cadence to a dead path. Quick updates fire through the search's own outer loop (_fire_quick_update), so the announced cadence is real. The wiring test postdates the search's removal, which is why the restored stash tripped it. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_018dV9h5h76hZ6trGZKWJFHe --- .../non_linear/search/test_quick_update_wiring.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/test_autofit/non_linear/search/test_quick_update_wiring.py b/test_autofit/non_linear/search/test_quick_update_wiring.py index cc0e62286..f77617eb1 100644 --- a/test_autofit/non_linear/search/test_quick_update_wiring.py +++ b/test_autofit/non_linear/search/test_quick_update_wiring.py @@ -23,6 +23,14 @@ # path suffix -> why this construction site legitimately omits the kwarg. EXEMPT = { + "nest/nss/search.py": ( + "af.NSS samples through inline JAX closures, never through Fitness -- " + "its Fitness exists only for AbstractNest.perform_update post-fit " + "work (latent samples, batch_size), so forwarding the cadence there " + "would wire it to a dead path. Quick updates fire through the " + "search's own outer loop instead: `iterations_per_quick_update` is " + "consumed directly in `_fit` via `_fire_quick_update`." + ), "mle/multi_start_gradient/search.py": ( "MultiStartGradient differentiates `fitness.call` inside its own " "jit/vmap step loop, so the Python-side counter in `call_wrap` would run "