diff --git a/AGENTS.md b/AGENTS.md index e44f5d7a8..5c9d5a6e5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -63,7 +63,7 @@ import autofit as af ``` Key subsystems: `non_linear/search/` (MCMC: emcee/zeus; nested: dynesty, -nautilus, NSS; MLE: LBFGS/BFGS/drawer), `mapper/` (model + priors), +nautilus; MLE: LBFGS/BFGS/drawer), `mapper/` (model + priors), `non_linear/analysis/` (`af.Analysis.log_likelihood_function`), `aggregator/`, `database/` (SQLAlchemy), `graphical/` (expectation propagation), `interpolator/`. @@ -81,10 +81,6 @@ nautilus, NSS; MLE: LBFGS/BFGS/drawer), `mapper/` (model + priors), (`test_autofit/graphical/test_declarative_deterministic.py` is the pattern). Capabilities that exist below but are silently absent above are the seam's known failure mode (see PyAutoFit#1336/#1337). -- The `[nss]` extra (for `af.NSS`) needs a pinned **handley-lab/blackjax** fork - installed manually *after* the extras step; that fork conflicts with the - mainline `blackjax` pinned in `[optional]`. Do not naively combine or bump - them — see the install notes in `pyproject.toml`. - All files use Unix line endings (LF, `\n`) — never `\r\n`. ## Working on issues diff --git a/autofit/__init__.py b/autofit/__init__.py index b80d6edd5..2edf6a1fe 100644 --- a/autofit/__init__.py +++ b/autofit/__init__.py @@ -83,7 +83,6 @@ 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 deleted file mode 100644 index f40825df6..000000000 --- a/autofit/non_linear/search/nest/nss/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from autofit.non_linear.search.nest.nss.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 deleted file mode 100644 index d23ee7a34..000000000 --- a/autofit/non_linear/search/nest/nss/_chunked_nss.py +++ /dev/null @@ -1,115 +0,0 @@ -"""Local replica of ``blackjax.ns.nss.as_top_level_api`` with chunked init. - -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`` -(``blackjax/ns/nss.py:223-230`` in the handley-lab fork) — 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 goes through ``make_chunked_update_strategy`` - (PyAutoFit#1303). Behaviour unchanged from that PR. -- 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 gated by the - # ``[nss]`` extra; see PyAutoFit pyproject.toml). - 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.nss import build_kernel, update_inner_kernel_params - - from autofit.non_linear.search.nest.nss._chunked_update import ( - make_chunked_update_strategy, - ) - - init_state_fn = partial( - init_state_strategy, - logprior_fn=logprior_fn, - loglikelihood_fn=loglikelihood_fn, - ) - - kernel = build_kernel( - init_state_fn, - num_inner_steps, - num_delete, - update_strategy=make_chunked_update_strategy(chunk_size), - ) - - def init_fn(position, rng_key=None): - # Mirror the upstream signature (rng_key is unused but accepted for - # API parity with ``blackjax.ns.nss.as_top_level_api``). - 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=update_inner_kernel_params, - ) - - 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 deleted file mode 100644 index 5a226a82e..000000000 --- a/autofit/non_linear/search/nest/nss/_chunked_update.py +++ /dev/null @@ -1,119 +0,0 @@ -"""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. - -``blackjax.nss(...)`` already exposes ``update_strategy`` as a kwarg -(see ``blackjax/ns/nss.py:157``), so ``af.NSS._fit`` only needs to pass -this builder to opt in: - - algo = _blackjax.nss( - ..., - update_strategy=make_chunked_update_strategy(chunk_size), - ) -""" - -from __future__ import annotations - -from functools import partial -from typing import Callable, Optional - - -def make_chunked_update_strategy(chunk_size: Optional[int]) -> Callable: - """Return an ``update_strategy`` callable for ``blackjax.nss(...)``. - - Signature matches ``blackjax.ns.from_mcmc.update_with_mcmc_take_last`` - so it can be passed through the ``update_strategy=`` kwarg 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 deleted file mode 100644 index b63b78376..000000000 --- a/autofit/non_linear/search/nest/nss/samples.py +++ /dev/null @@ -1,25 +0,0 @@ -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 deleted file mode 100644 index 744d78ba7..000000000 --- a/autofit/non_linear/search/nest/nss/search.py +++ /dev/null @@ -1,603 +0,0 @@ -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 autofit.non_linear.search.nest.nss.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 nss.ns import ( - log_weights as _nss_log_weights, - finalise as _nss_finalise, - ) - - _HAS_NSS = True -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 ``nss.ns.run_nested_sampling``. - - 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, ``yallup/nss``). - - ``af.NSS`` wraps ``nss.ns.run_nested_sampling`` 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. - - Phase 4 (``pip install autofit[nss]`` extra) is still pending — for now - ``af.NSS`` is an optional requirement and must be installed manually - via ``pip install git+https://github.com/yallup/nss.git``. - - 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 the optional `nss` package and the matching " - "`handley-lab/blackjax` fork. Install via:\n" - " pip install autofit[nss]\n" - "The extra pins specific upstream commits — see PyAutoFit's " - "pyproject.toml `[project.optional-dependencies] nss` entry." - ) - - 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 the upstream - ``nss.ns.run_nested_sampling`` 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 - # ``blackjax.ns.nss.as_top_level_api`` (~30 lines) with both vmap - # sites chunked: the inner MCMC step (matches 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 before this PR). ``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 autofit.non_linear.search.nest.nss._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))) - evals = int( - final_state.update_info.num_steps.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 1fa53dc55..4bed414bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,35 +75,6 @@ optional = [ "nautilus-sampler==1.0.5", "zeus-mcmc==2.5.4", ] -# Phase 4 of nss_first_class_sampler — support package for `af.NSS`. -# -# The `handley-lab/blackjax` fork (BSD-3-Clause) carries the -# `blackjax.ns.adaptive.init` entrypoint that mainline blackjax lacks, and -# `yallup/nss` (BSD-3-Clause) sits on top. Both are pinned to specific SHAs -# below — but installed manually rather than declared here, because PyPI and -# TestPyPI reject direct `git+https://` URLs in uploaded wheels with: -# -# 400 Can't have direct dependency: blackjax @ git+... -# -# After `pip install autofit[nss]`, complete the install with: -# -# pip install \ -# "blackjax @ git+https://github.com/handley-lab/blackjax.git@ef45acd2f2fa0cca15adbdcd3ff7cb3a98987cb5" \ -# "nss @ git+https://github.com/yallup/nss.git@69159b0f4a3a53123b9eec7df91e4ed3885e4dc4" -# -# blackjax SHA `ef45acd2` is the May 2026 "Merge PR #60 — double_compile" -# revision, locally validated against the Phase 1-3 work. blackjax HEAD -# (5dbf89a9 at pin time) introduced `numpy>=1.25`. autofit's anesthetic -# pin has since been bumped to `>=2.9.0` (numpy>=2 compatible) so the -# previous "bump only after both pins move together" constraint is now -# resolved. See PyAutoPrompt/issued/nss_install_simplification.md notes -# #3-4 for the historic context. -# -# `fastprogress<1.1` stays here as a regular PyPI dep — it keeps fastprogress -# from pulling `python-fasthtml` (1.1.5 transitively requires it). -nss = [ - "fastprogress<1.1", -] docs=[ "sphinx", "furo", diff --git a/test_autofit/non_linear/search/nest/nss/__init__.py b/test_autofit/non_linear/search/nest/nss/__init__.py deleted file mode 100644 index e69de29bb..000000000 diff --git a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py b/test_autofit/non_linear/search/nest/nss/test_checkpoint.py deleted file mode 100644 index cca204016..000000000 --- a/test_autofit/non_linear/search/nest/nss/test_checkpoint.py +++ /dev/null @@ -1,245 +0,0 @@ -""" -Unit tests for the ``af.NSS`` checkpoint/resume helpers -(``_save_checkpoint`` / ``_load_checkpoint``). - -No real ``nss.ns.run_nested_sampling`` 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 optional `nss` extra", -) - - -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 deleted file mode 100644 index 04b75e2ee..000000000 --- a/test_autofit/non_linear/search/nest/nss/test_search.py +++ /dev/null @@ -1,250 +0,0 @@ -""" -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 optional `nss` extra", -) - - -@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 ``af.NSS._fit`` drop it into - ``blackjax.nss(update_strategy=...)`` 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 the optional `nss` package"): - 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 - )