From c6a05527f6ffc0d9eedd834c6e502c5eec377b18 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 08:37:17 +0100 Subject: [PATCH 1/2] chore: delete the superseded searches/nss archive copy NSS was re-mainlined into PyAutoFit (af.NSS, autofit/non_linear/search/ nest/nss) and the library copy has since diverged, so this archive copy duplicates live code and its README ("parked ... removed 2026-07-11") is wrong. Deleted via Gut per the human ruling of 2026-10-08. Recoverable at the pre-delete SHA: 7cb97c6d4380324beb6fc724e3edd5ad0592b549 (git show 7cb97c6:searches/nss/...) Refs PyAutoLabs/PyAutoFit#1668 Co-Authored-By: Claude Fable 5.1 --- searches/nss/README.md | 88 ---- searches/nss/__init__.py | 1 - searches/nss/_chunked_nss.py | 115 ----- searches/nss/_chunked_update.py | 119 ----- searches/nss/example.py | 112 ----- searches/nss/samples.py | 25 -- searches/nss/search.py | 602 -------------------------- searches/nss/tests/__init__.py | 0 searches/nss/tests/test_checkpoint.py | 243 ----------- searches/nss/tests/test_search.py | 252 ----------- 10 files changed, 1557 deletions(-) delete mode 100644 searches/nss/README.md delete mode 100644 searches/nss/__init__.py delete mode 100644 searches/nss/_chunked_nss.py delete mode 100644 searches/nss/_chunked_update.py delete mode 100644 searches/nss/example.py delete mode 100644 searches/nss/samples.py delete mode 100644 searches/nss/search.py delete mode 100644 searches/nss/tests/__init__.py delete mode 100644 searches/nss/tests/test_checkpoint.py delete mode 100644 searches/nss/tests/test_search.py diff --git a/searches/nss/README.md b/searches/nss/README.md deleted file mode 100644 index deaa02d..0000000 --- a/searches/nss/README.md +++ /dev/null @@ -1,88 +0,0 @@ -# NSS — Nested Slice Sampling (parked implementation) - -This directory preserves the **`af.NSS`** nested slice sampler that was -previously bundled with PyAutoFit at -`autofit/non_linear/search/nest/nss/`. It was removed from the library on -2026-07-11 (PyAutoFit#1356) because the bespoke install / CI / build machinery -it required was not justified by its measured performance — but the sampler -works, and this is the full, working implementation kept here so it can be -re-mainlined cheaply once `nss` ships as a genuine PyPI package. - -Like `../ultranest/` and `../pyswarms/`, `search.py` imports only from -`autofit`'s base classes, so it is a drop-in `af.NonLinearSearch`. The -intra-package imports were converted to **relative** imports (`from .samples -import NSSamples`, …), so the module works both from here (`from searches.nss -import NSS`) and verbatim if the folder is dropped back into the library. - -## What NSS is - -A JAX-native nested sampler whose inner sampling loop runs end-to-end inside -`jax.jit`. When the log-likelihood is itself JAX-traceable it avoids the -Python ↔ JAX boundary that Nautilus and Dynesty cross on every call. On the -production lensing likelihoods that motivated it the per-evaluation cost was -roughly **30× lower than Nautilus** (MGE). The catch — and the reason it was -parked rather than kept — is that it is **OOM-prone on pixelization / Delaunay -inversions** via the vmap fan-out, so it was never a clear across-the-board -win. - -## Why it couldn't be a normal `pip install autofit[nss]` - -Its real dependencies are two **pinned git forks**, and PyPI/TestPyPI reject -direct `git+https://` URLs inside an uploaded wheel -(`400 Can't have direct dependency: blackjax @ git+...`). So the `[nss]` extra -could only carry `fastprogress<1.1`, and the two forks had to be installed -manually afterwards. - -### The exact install recipe (needed to run this again) - -```bash -pip install "fastprogress<1.1" # keeps fastprogress from pulling python-fasthtml (1.1.5+) -pip install \ - "blackjax @ git+https://github.com/handley-lab/blackjax.git@ef45acd2f2fa0cca15adbdcd3ff7cb3a98987cb5" \ - "nss @ git+https://github.com/yallup/nss.git@69159b0f4a3a53123b9eec7df91e4ed3885e4dc4" -``` - -- The **`handley-lab/blackjax`** fork carries the `blackjax.ns.adaptive.init` - entrypoint that mainline blackjax lacks. SHA `ef45acd2` is the May 2026 - "Merge PR #60 — double_compile" revision, locally validated against the - Phase 1–3 work. -- **`yallup/nss`** (the sampler itself) sits on top of that fork. -- This fork **conflicts with the mainline `blackjax`** pinned in autofit's - `[optional]` extra — install NSS in its own environment, do not naively - combine or bump the two. - -## Requirement: a JAX-traceable analysis - -NSS runs the likelihood inside `jax.jit`, so the analysis must be built with -`use_jax=True` (its arithmetic then dispatches through `jax.numpy`). Handing it -a plain NumPy analysis raises `TracerArrayConversionError`. See `example.py`. - -## Key kwargs - -| kwarg | production default | meaning | -|-------|--------------------|---------| -| `n_live` | 200 | live particles maintained throughout the run | -| `num_mcmc_steps` | 5 | slice-MCMC inner steps per dead-point batch | -| `num_delete` | 50 | particles removed per outer iteration (larger ⇒ less JIT overhead, coarser coverage) | -| `termination` | -3.0 | stop on `logZ_live - logZ` (≈ remaining evidence < 1e-3) | -| `checkpoint_interval` | — | outer iterations between resumable disk checkpoints (SLURM-friendly) | -| `iterations_per_quick_update` | None | if set, calls `analysis.visualize(...)` every N outer iterations | - -## Re-mainlining checklist (when `nss` becomes a real pip package) - -1. Move this folder back to `PyAutoFit/autofit/non_linear/search/nest/nss/` - (relative imports mean the module files need no edits). -2. Restore `tests/` to `PyAutoFit/test_autofit/non_linear/search/nest/nss/` - — the tests were kept **verbatim** against the library import paths - (`autofit.non_linear.search.nest.nss.*`, including the `mock.patch` target - strings), so they run unchanged from that location. -3. Re-add to `autofit/__init__.py`: - `from .non_linear.search.nest.nss.search import NSS`. -4. Add the dependency to `pyproject.toml`. If `nss` is a real PyPI package by - then, this is finally a plain `nss = ["nss>=X"]` extra — the whole reason - this was removed. Otherwise restore the git+ recipe above. -5. Restore the `Search: NSS` tutorial section in - `autofit_workspace/scripts/searches/nest.py` (see PyAutoFit#1356 diff). - -Provenance: removed in PyAutoFit#1356 (2026-07-11); the removal diff is the -authoritative record of every touch-point across the six repos. diff --git a/searches/nss/__init__.py b/searches/nss/__init__.py deleted file mode 100644 index 76f18ac..0000000 --- a/searches/nss/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from .search import NSS diff --git a/searches/nss/_chunked_nss.py b/searches/nss/_chunked_nss.py deleted file mode 100644 index 9cbb94b..0000000 --- a/searches/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 ._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/searches/nss/_chunked_update.py b/searches/nss/_chunked_update.py deleted file mode 100644 index 5a226a8..0000000 --- a/searches/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/searches/nss/example.py b/searches/nss/example.py deleted file mode 100644 index 94c3891..0000000 --- a/searches/nss/example.py +++ /dev/null @@ -1,112 +0,0 @@ -""" -NSS Example ------------ - -This script demonstrates the NSS (Nested Slice Sampling) sampler that was -previously bundled with PyAutoFit as ``af.NSS``. It was removed from the -library on 2026-07-11 (PyAutoFit#1356); the full implementation is preserved in -this directory (``search.py`` and friends) as a drop-in ``af.NonLinearSearch``. - -NSS must be installed separately — see ``README.md`` in this directory for the -exact pinned install recipe (a ``handley-lab/blackjax`` fork + ``yallup/nss``). -The ``NSS`` class here imports from autofit's base classes, so once its -dependencies are installed it works as a drop-in replacement for any boundary -sampler (``af.Nautilus``, ``af.DynestyStatic``, …). - -__Analysis Must Be JAX-Traceable__ - -NSS runs the log-likelihood inside ``jax.jit``. The boundary-based samplers are -happy with a plain NumPy ``af.ex.Analysis(data, noise_map)`` — but when we hand -the same analysis to NSS the JIT trace hits the NumPy paths and raises -``TracerArrayConversionError``. The fix is to build the analysis with -``use_jax=True``, which makes its arithmetic dispatch through ``jax.numpy``. -""" - -import matplotlib.pyplot as plt -import numpy as np -from os import path - -import autofit as af - -# The parked NSS implementation living alongside this example. -from searches.nss import NSS - -""" -__Data__ - -This example fits a single 1D Gaussian, we therefore load and plot data -containing one Gaussian. -""" -dataset_path = path.join("dataset", "example_1d", "gaussian_x1") - -data = af.util.numpy_array_from_json(file_path=path.join(dataset_path, "data.json")) -noise_map = af.util.numpy_array_from_json( - file_path=path.join(dataset_path, "noise_map.json") -) - -""" -__Model + Analysis__ - -A single ``Gaussian`` (dimensionality N=3). The analysis is built with -``use_jax=True`` so its likelihood is JAX-traceable — the precondition for NSS. -""" -model = af.Model(af.ex.Gaussian) - -model.centre = af.UniformPrior(lower_limit=0.0, upper_limit=100.0) -model.normalization = af.LogUniformPrior(lower_limit=1e-2, upper_limit=1e2) -model.sigma = af.UniformPrior(lower_limit=0.0, upper_limit=30.0) - -analysis = af.ex.Analysis(data=data, noise_map=noise_map, use_jax=True) - -""" -__Search: NSS__ - -Swapping ``af.Nautilus(...)`` for ``NSS(...)`` in an existing script is a -one-line change — NSS exposes the same ``result.samples`` interface. See the -kwarg table in ``README.md`` for what each setting does. -""" -search = NSS( - path_prefix=path.join("searches"), - name="NSS", - n_live=200, # live particles maintained throughout the run - num_mcmc_steps=5, # slice-MCMC inner steps per dead-point batch - num_delete=50, # particles removed per outer iteration - termination=-3.0, # delta-logZ stopping criterion - seed=42, # JAX PRNG seed for reproducible runs - checkpoint_interval=100, # SLURM-friendly resume — see README -) - -result = search.fit(model=model, analysis=analysis) - -""" -__Result__ - -Confirm NSS converged to roughly the same posterior as the boundary samplers. -On this trivial 1D likelihood there is no speed advantage (the per-call cost is -dominated by Python overhead, and the first run pays a one-off ~25-30s JIT -compile) — the per-eval win only shows on a real JAX-traceable autolens / -autogalaxy MGE likelihood. -""" -model_data = result.max_log_likelihood_instance.model_data_from( - xvalues=np.arange(data.shape[0]) -) - -plt.errorbar( - x=range(data.shape[0]), - y=data, - yerr=noise_map, - linestyle="", - color="k", - ecolor="k", - elinewidth=1, - capsize=2, -) -plt.plot(range(data.shape[0]), model_data, color="r") -plt.title("NSS model fit to 1D Gaussian dataset.") -plt.xlabel("x values of profile") -plt.ylabel("Profile normalization") -plt.show() -plt.close() - -print(f"NSS log evidence: {result.samples.samples_info['log_evidence']:.4f}") -print(f"NSS max log L: {max(result.samples.log_likelihood_list):.4f}") diff --git a/searches/nss/samples.py b/searches/nss/samples.py deleted file mode 100644 index b63b783..0000000 --- a/searches/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/searches/nss/search.py b/searches/nss/search.py deleted file mode 100644 index 2483e3e..0000000 --- a/searches/nss/search.py +++ /dev/null @@ -1,602 +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 .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 ._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/searches/nss/tests/__init__.py b/searches/nss/tests/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/searches/nss/tests/test_checkpoint.py b/searches/nss/tests/test_checkpoint.py deleted file mode 100644 index 53eb2dc..0000000 --- a/searches/nss/tests/test_checkpoint.py +++ /dev/null @@ -1,243 +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/searches/nss/tests/test_search.py b/searches/nss/tests/test_search.py deleted file mode 100644 index d962552..0000000 --- a/searches/nss/tests/test_search.py +++ /dev/null @@ -1,252 +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 - ) From ae30c84e0cc7056a90046cbb286aae104219373a Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 8 Oct 2026 08:37:17 +0100 Subject: [PATCH 2/2] docs: define the searches/ archive in the root README The README described the repo only as an archive of UltraNest and PySwarms. It now says what the archive is for (searches PyAutoFit does not ship), that a copy is deleted once its search is mainlined, and that searches_minimal/ is a prototype tier, not part of the archive. Also fixes the PyAutoFit link (rhayes777 -> PyAutoLabs) and the matching AGENTS.md structure line. Refs PyAutoLabs/PyAutoFit#1668 Co-Authored-By: Claude Fable 5.1 --- AGENTS.md | 3 ++- README.md | 24 ++++++++++++++++++++---- 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 9dd2e26..1451457 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -16,7 +16,8 @@ Dependencies: `autofit`, plus optional sampler backends (`nautilus-sampler`, `bl searches_minimal/ Minimal direct-sampler examples (Nautilus, Dynesty, Emcee, LBFGS) that bypass the NonLinearSearch wrapper. Outputs land in searches_minimal/output/. -searches/ Search-interface prototypes (pyswarms, ultranest). +searches/ Archive of searches not in PyAutoFit (pyswarms, ultranest); + a copy is deleted once its search is mainlined. projects/ Example PyAutoFit projects (cosmology, ...). scripts/ Tutorial-style developer scripts (howtofit). config/ YAML configuration files (non_linear/...). diff --git a/README.md b/README.md index ea76be6..934f9c0 100644 --- a/README.md +++ b/README.md @@ -1,10 +1,26 @@ # autofit_workspace_developer -Archived non-linear search implementations removed from -[PyAutoFit](https://github.com/rhayes777/PyAutoFit). These searches are preserved -here so they are not lost and can still be used by advanced users. +Developer workspace for [PyAutoFit](https://github.com/PyAutoLabs/PyAutoFit): +search prototypes, minimal sampler examples, expectation-propagation and graphical +model experiments. -## Searches +## The `searches/` archive + +`searches/` holds non-linear search implementations that are **not** in PyAutoFit, +either because they were removed from the library or because they have not been +promoted into it. They are preserved here so they are not lost and can still be +used by advanced users. + +A search lives in the archive only while PyAutoFit does not ship it. When a search +is (re-)mainlined into PyAutoFit, its archive copy is deleted rather than kept in +parallel, so the library is the one maintained implementation. NSS is an example: +it was parked here and is now `af.NSS` in PyAutoFit. + +`searches_minimal/` is not part of the archive: it holds minimal scripts that call +samplers directly, bypassing the `NonLinearSearch` wrapper, for prototyping and +benchmarking. + +## Archived searches ### UltraNest