refactor: retire the NSS profiling runs

README.md

@@ -10,7 +10,7 @@ This repository is the single home for PyAutoLens performance measurement. It ex
 
 - **Likelihood functions** — imaging, interferometer, point-source, and datacube paths, across the MGE, pixelization, and Delaunay model compositions used in real science cases.
 - **Simulators** — run-time tracking for the imaging, interferometer, point-source, cluster, group, and multi-plane simulators.
-- **Searches / samplers** — sampler-level profiling, starting with Nautilus. Other samplers (Dynesty, Emcee, BlackJAX, NumPyro, NSS, LBFGS, PocoMC) follow in later sweeps.
+- **Searches / samplers** — sampler-level profiling, starting with Nautilus. Other samplers (Dynesty, Emcee, BlackJAX, NumPyro, LBFGS, PocoMC) follow in later sweeps.
 
 **Hardware tiers covered:**
 

hpc/README.md

@@ -13,8 +13,7 @@ batch_gpu/
 ```
 
 Submit names follow the same `<class>/<model>` cell grid as the rest of the
-repo; `runtime_` prefixed submits drive `likelihood_runtime/` cells, `nss_`
-prefixed ones drive `searches/` cells.
+repo; `runtime_` prefixed submits drive `likelihood_runtime/` cells.
 
 ## Running
 

searches/_runner.py

@@ -38,7 +38,6 @@
 )
 from searches._metrics import attach_viz_timer, collect_metrics  # noqa: E402
 from searches._samplers import (  # noqa: E402
-    _NSS_DEFAULTS,
     SAMPLER_BUILDERS,
     n_live_for,
     vmap_batch_for_cell,
@@ -187,16 +186,6 @@ def _sampler_config_dict(
             "force_x1_cpu": use_jax,
             "iterations_per_update": 3 * n_live,
         }
-    if sampler == "nss":
-        return {
-            "n_live": n_live,
-            "num_mcmc_steps": int(_NSS_DEFAULTS["num_mcmc_steps"]),
-            "num_delete": int(_NSS_DEFAULTS["num_delete"]),
-            "chunk_size": None,
-            "termination": float(_NSS_DEFAULTS["termination"]),
-            "seed": int(_NSS_DEFAULTS["seed"]),
-            "jax_native": True,
-        }
     return {"n_live": n_live, "_note": f"unknown sampler {sampler!r}"}
 
 

searches/_samplers.py

@@ -49,22 +49,6 @@
 }
 
 
-# NSS-specific defaults. ``af.NSS`` exposes three knobs above ``n_live``:
-# - ``num_mcmc_steps``: inner slice-MCMC steps per dead-point batch.
-# - ``num_delete``: particles killed per outer iteration. Larger reduces
-#   JIT overhead per iteration at the cost of slightly worse posterior
-#   coverage.
-# - ``termination``: stop when ``logZ_live - logZ < termination``. Default
-#   ``-3.0`` corresponds to delta-logZ < 1e-3 — matched to PyAutoFit's
-#   own production-tested default.
-_NSS_DEFAULTS: dict[str, int | float] = {
-    "num_mcmc_steps": 5,
-    "num_delete": 50,
-    "termination": -3.0,
-    "seed": 42,
-}
-
-
 def n_live_for(dataset_class: str, model_type: str) -> int:
     """Look up the canonical n_live for a (dataset_class, model_type) cell."""
     try:
@@ -138,66 +122,7 @@ def build_nautilus(
     )
 
 
-def build_nss(
-    *,
-    sampler: str,
-    dataset_class: str,
-    model_type: str,
-    instrument: str,
-    config_name: str,
-    use_jax: bool,
-) -> af.NSS:
-    """Construct a first-class ``af.NSS`` (JAX nested slice sampler) for one cell.
-
-    ``af.NSS`` is JAX-native: the likelihood + prior closures both run inside
-    ``jax.jit`` and execute on whatever device JAX is configured for. There
-    is no NumPy / multiprocessing fallback, so:
-
-    - ``use_jax`` must be True (the analysis must be JAX-traceable).
-    - ``number_of_cores`` is accepted by ``af.NSS`` for API parity but
-      ignored. We leave it at the default ``1``.
-    - There is no ``force_x1_cpu`` / ``use_jax_vmap`` toggle; NSS always
-      batches across live particles natively.
-
-    ``n_live`` matches the per-(ds, model) SLaM choice used for Nautilus
-    so timing comparisons line up. ``num_delete=50``, ``num_mcmc_steps=5``,
-    ``termination=-3.0`` are PyAutoFit's production-tested defaults.
-    """
-    if not use_jax:
-        raise ValueError(
-            "af.NSS is JAX-native; running with use_jax=False is not "
-            "supported. Unset PYAUTO_DISABLE_JAX or use sampler='nautilus' "
-            "for the NumPy-front profile."
-        )
-    n_live = n_live_for(dataset_class, model_type)
-    # Memory-budget plumbing:
-    # - ``num_delete`` stays at the sampler's preferred default (50 particles
-    #   per outer iteration) so convergence isn't compromised.
-    # - ``chunk_size = None`` lets blackjax use its native ``jax.vmap`` over
-    #   ``num_delete`` particles (no chunking, no inner-step serialisation).
-    #   This is the best-case sampler shape and is only viable because
-    #   ``searches/_setup.py:_build_imaging`` now calls
-    #   ``dataset.apply_sparse_operator()``, which makes the inversion factory
-    #   pick ``InversionImagingSparse``. A100 fp64 vmap-probe results
-    #   (autolens_profiling#44):
-    #     dense  per-replica VRAM:  931 MB  ->  n_live=150 needs 140 GB (OOM)
-    #     sparse per-replica VRAM:   95 MB  ->  n_live=150 needs  14 GB
-    #   With sparse the full ``num_delete=50`` vmap fits comfortably and the
-    #   chunk_size workaround from PyAutoFit#1303/#1305 is no longer needed.
-    return af.NSS(
-        name=config_name,
-        path_prefix=f"searches/{sampler}/{dataset_class}/{model_type}/{instrument}",
-        n_live=n_live,
-        num_mcmc_steps=int(_NSS_DEFAULTS["num_mcmc_steps"]),
-        num_delete=int(_NSS_DEFAULTS["num_delete"]),
-        chunk_size=None,
-        termination=float(_NSS_DEFAULTS["termination"]),
-        seed=int(_NSS_DEFAULTS["seed"]),
-    )
-
-
 SamplerBuilder = Callable[..., af.NonLinearSearch]
 SAMPLER_BUILDERS: dict[str, SamplerBuilder] = {
     "nautilus": build_nautilus,
-    "nss": build_nss,
 }