From df5d7ac508596d21174b4c82aa0f9c67ca11a1c6 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Sat, 11 Jul 2026 15:56:46 +0100 Subject: [PATCH] refactor: retire the NSS profiling runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Removes NSS from the profiling framework now that af.NSS is gone (PyAutoFit#1356, Phase 5 — the final piece): the searches/nss/ cell subtree (9 files), the 9 nss sweep cells, build_nss + _NSS_DEFAULTS + the "nss" SAMPLER_BUILDERS registration, the _runner nss settings branch, the 3 nss A100 HPC submit scripts, and the nss mentions in README / hpc README. Kept as historical evidence: results/notes/sparse_vs_dense_inversion_path.md (A100 fp64 NSS validation table) and the HPC JAX venv reference (used for all GPU profiling; NSS-derived name only). Issue: PyAutoFit#1356 Co-Authored-By: Claude Opus 4.8 --- README.md | 2 +- hpc/README.md | 3 +- .../submit_nss_imaging_delaunay_a100_hst_fp64 | 45 ----------- .../submit_nss_imaging_mge_a100_hst_fp64 | 50 ------------- ...mit_nss_imaging_pixelization_a100_hst_fp64 | 48 ------------ searches/_runner.py | 11 --- searches/_samplers.py | 75 ------------------- searches/nss/datacube/delaunay.py | 24 ------ searches/nss/imaging/delaunay.py | 19 ----- searches/nss/imaging/mge.py | 24 ------ searches/nss/imaging/pixelization.py | 19 ----- searches/nss/interferometer/delaunay.py | 19 ----- searches/nss/interferometer/mge.py | 19 ----- searches/nss/interferometer/pixelization.py | 19 ----- searches/nss/point_source/image_plane.py | 19 ----- searches/nss/point_source/source_plane.py | 19 ----- searches/sweep.py | 9 --- 17 files changed, 2 insertions(+), 422 deletions(-) delete mode 100755 hpc/batch_gpu/submit_nss_imaging_delaunay_a100_hst_fp64 delete mode 100755 hpc/batch_gpu/submit_nss_imaging_mge_a100_hst_fp64 delete mode 100755 hpc/batch_gpu/submit_nss_imaging_pixelization_a100_hst_fp64 delete mode 100644 searches/nss/datacube/delaunay.py delete mode 100644 searches/nss/imaging/delaunay.py delete mode 100644 searches/nss/imaging/mge.py delete mode 100644 searches/nss/imaging/pixelization.py delete mode 100644 searches/nss/interferometer/delaunay.py delete mode 100644 searches/nss/interferometer/mge.py delete mode 100644 searches/nss/interferometer/pixelization.py delete mode 100644 searches/nss/point_source/image_plane.py delete mode 100644 searches/nss/point_source/source_plane.py diff --git a/README.md b/README.md index cc15b4d..8128f93 100644 --- a/README.md +++ b/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:** diff --git a/hpc/README.md b/hpc/README.md index 1ac1b3e..a0f340d 100644 --- a/hpc/README.md +++ b/hpc/README.md @@ -13,8 +13,7 @@ batch_gpu/ ``` Submit names follow the same `/` 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 diff --git a/hpc/batch_gpu/submit_nss_imaging_delaunay_a100_hst_fp64 b/hpc/batch_gpu/submit_nss_imaging_delaunay_a100_hst_fp64 deleted file mode 100755 index c1bd42d..0000000 --- a/hpc/batch_gpu/submit_nss_imaging_delaunay_a100_hst_fp64 +++ /dev/null @@ -1,45 +0,0 @@ -#!/bin/bash -l -# -# A100 first-class search profiling: searches/nss/imaging/delaunay × hst × fp64. -# -# Pair to submit_imaging_delaunay_a100_hst_fp64 (Nautilus); same dataset / -# model / analysis so timing and convergence numbers compare directly. - -#SBATCH -J search_nss_imaging_delaunay_hst_fp64 -#SBATCH --partition=gpu -#SBATCH --gres=gpu:1 -#SBATCH --ntasks=1 -#SBATCH --cpus-per-task=4 -#SBATCH --mem=64gb -#SBATCH --time=3:00:00 -#SBATCH -o output/output.%A.out -#SBATCH -e error/error.%A.err -#SBATCH --mail-type=END,FAIL -#SBATCH --mail-user=james.w.nightingale@durham.ac.uk - -export AP_ROOT=/mnt/ral/jnightin/autolens_profiling -source $AP_ROOT/activate.sh - -export JAX_PLATFORM_NAME=cuda -export JAX_PLATFORMS=cuda,cpu -export XLA_PYTHON_CLIENT_PREALLOCATE=false -export JAX_ENABLE_X64=True -export NUMBA_CACHE_DIR=/tmp/numba_cache -export MPLCONFIGDIR=/tmp/matplotlib - -nvidia-smi - -echo "==========================================" -date -echo "Cell: searches/nss/imaging/delaunay" -echo "Instrument: hst" -echo "Precision: fp64" - -cd $AP_ROOT -python3 searches/nss/imaging/delaunay.py \ - --instrument hst \ - --config-name hpc_a100_fp64 \ - --output-dir $AP_ROOT/results/searches/nss/imaging/delaunay/hst - -echo "Finished." -date diff --git a/hpc/batch_gpu/submit_nss_imaging_mge_a100_hst_fp64 b/hpc/batch_gpu/submit_nss_imaging_mge_a100_hst_fp64 deleted file mode 100755 index 83ff461..0000000 --- a/hpc/batch_gpu/submit_nss_imaging_mge_a100_hst_fp64 +++ /dev/null @@ -1,50 +0,0 @@ -#!/bin/bash -l -# -# A100 first-class search profiling: searches/nss/imaging/mge × hst × fp64. -# -# Drives af.NSS end-to-end (JAX-native nested slice sampling on the -# handley-lab blackjax fork + yallup/nss) on the same HST MGE model the -# Nautilus baseline submit (submit_imaging_mge_a100_hst_fp64) profiles, so -# log_Z / per-eval / wall-time / convergence numbers compare directly. -# -# af.NSS does not use multiprocessing — JAX handles batching internally. -# Same A100 device + memory budget as the Nautilus submit. - -#SBATCH -J search_nss_imaging_mge_hst_fp64 -#SBATCH --partition=gpu -#SBATCH --gres=gpu:1 -#SBATCH --ntasks=1 -#SBATCH --cpus-per-task=4 -#SBATCH --mem=64gb -#SBATCH --time=2:00:00 -#SBATCH -o output/output.%A.out -#SBATCH -e error/error.%A.err -#SBATCH --mail-type=END,FAIL -#SBATCH --mail-user=james.w.nightingale@durham.ac.uk - -export AP_ROOT=/mnt/ral/jnightin/autolens_profiling -source $AP_ROOT/activate.sh - -export JAX_PLATFORM_NAME=cuda -export JAX_PLATFORMS=cuda,cpu -export XLA_PYTHON_CLIENT_PREALLOCATE=false -export JAX_ENABLE_X64=True -export NUMBA_CACHE_DIR=/tmp/numba_cache -export MPLCONFIGDIR=/tmp/matplotlib - -nvidia-smi - -echo "==========================================" -date -echo "Cell: searches/nss/imaging/mge" -echo "Instrument: hst" -echo "Precision: fp64" - -cd $AP_ROOT -python3 searches/nss/imaging/mge.py \ - --instrument hst \ - --config-name hpc_a100_fp64 \ - --output-dir $AP_ROOT/results/searches/nss/imaging/mge/hst - -echo "Finished." -date diff --git a/hpc/batch_gpu/submit_nss_imaging_pixelization_a100_hst_fp64 b/hpc/batch_gpu/submit_nss_imaging_pixelization_a100_hst_fp64 deleted file mode 100755 index ad0db68..0000000 --- a/hpc/batch_gpu/submit_nss_imaging_pixelization_a100_hst_fp64 +++ /dev/null @@ -1,48 +0,0 @@ -#!/bin/bash -l -# -# A100 first-class search profiling: searches/nss/imaging/pixelization × hst × fp64. -# -# NSS control-cell paired with Nautilus' submit_imaging_pixelization_a100_hst_fp64 -# and the existing Delaunay results. Same model shape + memory budget as Delaunay -# (931 MB / replica) but no pure_callback in the mesh path — the comparison -# isolates whether NSS's Delaunay OOM is the pure_callback × vmap pytree -# blow-up or a generic inversion-matrix issue. - -#SBATCH -J search_nss_imaging_pixelization_hst_fp64 -#SBATCH --partition=gpu -#SBATCH --gres=gpu:1 -#SBATCH --ntasks=1 -#SBATCH --cpus-per-task=4 -#SBATCH --mem=64gb -#SBATCH --time=3:00:00 -#SBATCH -o output/output.%A.out -#SBATCH -e error/error.%A.err -#SBATCH --mail-type=END,FAIL -#SBATCH --mail-user=james.w.nightingale@durham.ac.uk - -export AP_ROOT=/mnt/ral/jnightin/autolens_profiling -source $AP_ROOT/activate.sh - -export JAX_PLATFORM_NAME=cuda -export JAX_PLATFORMS=cuda,cpu -export XLA_PYTHON_CLIENT_PREALLOCATE=false -export JAX_ENABLE_X64=True -export NUMBA_CACHE_DIR=/tmp/numba_cache -export MPLCONFIGDIR=/tmp/matplotlib - -nvidia-smi - -echo "==========================================" -date -echo "Cell: searches/nss/imaging/pixelization" -echo "Instrument: hst" -echo "Precision: fp64" - -cd $AP_ROOT -python3 searches/nss/imaging/pixelization.py \ - --instrument hst \ - --config-name hpc_a100_fp64 \ - --output-dir $AP_ROOT/results/searches/nss/imaging/pixelization/hst - -echo "Finished." -date diff --git a/searches/_runner.py b/searches/_runner.py index c8507fe..fca4f28 100644 --- a/searches/_runner.py +++ b/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}"} diff --git a/searches/_samplers.py b/searches/_samplers.py index 5748535..6360f76 100644 --- a/searches/_samplers.py +++ b/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, } diff --git a/searches/nss/datacube/delaunay.py b/searches/nss/datacube/delaunay.py deleted file mode 100644 index 3ba612f..0000000 --- a/searches/nss/datacube/delaunay.py +++ /dev/null @@ -1,24 +0,0 @@ -"""First-class af.NSS search profiling — datacube Delaunay. - -Multi-channel cube fit via ``af.FactorGraphModel`` (same wiring as the -Nautilus datacube cell). ``_DATACUBE_N_CHANNELS`` in ``_setup.py`` -controls the channel count. -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="datacube", - model_type="delaunay", - default_instrument="sma", -) diff --git a/searches/nss/imaging/delaunay.py b/searches/nss/imaging/delaunay.py deleted file mode 100644 index 3e1b8e0..0000000 --- a/searches/nss/imaging/delaunay.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — imaging Delaunay.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="imaging", - model_type="delaunay", - default_instrument="hst", -) diff --git a/searches/nss/imaging/mge.py b/searches/nss/imaging/mge.py deleted file mode 100644 index 5639276..0000000 --- a/searches/nss/imaging/mge.py +++ /dev/null @@ -1,24 +0,0 @@ -"""First-class af.NSS search profiling — imaging MGE. - -JAX-native nested slice sampling against the HST-MGE-style imaging model. -Same dataset / model / analysis as ``searches/nautilus/imaging/mge.py`` so -timing and convergence numbers compare directly. -""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="imaging", - model_type="mge", - default_instrument="hst", -) diff --git a/searches/nss/imaging/pixelization.py b/searches/nss/imaging/pixelization.py deleted file mode 100644 index 33c021f..0000000 --- a/searches/nss/imaging/pixelization.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — imaging pixelization.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="imaging", - model_type="pixelization", - default_instrument="hst", -) diff --git a/searches/nss/interferometer/delaunay.py b/searches/nss/interferometer/delaunay.py deleted file mode 100644 index ca99ecc..0000000 --- a/searches/nss/interferometer/delaunay.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — interferometer Delaunay.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="interferometer", - model_type="delaunay", - default_instrument="sma", -) diff --git a/searches/nss/interferometer/mge.py b/searches/nss/interferometer/mge.py deleted file mode 100644 index 030ddcf..0000000 --- a/searches/nss/interferometer/mge.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — interferometer MGE.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="interferometer", - model_type="mge", - default_instrument="sma", -) diff --git a/searches/nss/interferometer/pixelization.py b/searches/nss/interferometer/pixelization.py deleted file mode 100644 index cd25aca..0000000 --- a/searches/nss/interferometer/pixelization.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — interferometer pixelization.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="interferometer", - model_type="pixelization", - default_instrument="sma", -) diff --git a/searches/nss/point_source/image_plane.py b/searches/nss/point_source/image_plane.py deleted file mode 100644 index b3ffdfb..0000000 --- a/searches/nss/point_source/image_plane.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — point-source image-plane fit.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="point_source", - model_type="image_plane", - default_instrument="simple", -) diff --git a/searches/nss/point_source/source_plane.py b/searches/nss/point_source/source_plane.py deleted file mode 100644 index a203da4..0000000 --- a/searches/nss/point_source/source_plane.py +++ /dev/null @@ -1,19 +0,0 @@ -"""First-class af.NSS search profiling — point-source source-plane fit.""" - -from __future__ import annotations - -import sys -from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[3] -if str(_REPO_ROOT) not in sys.path: - sys.path.insert(0, str(_REPO_ROOT)) - -from searches._runner import run_search # noqa: E402 - -run_search( - sampler="nss", - dataset_class="point_source", - model_type="source_plane", - default_instrument="simple", -) diff --git a/searches/sweep.py b/searches/sweep.py index 9755e80..3081382 100644 --- a/searches/sweep.py +++ b/searches/sweep.py @@ -71,15 +71,6 @@ ("nautilus", "interferometer", "pixelization"), ("nautilus", "interferometer", "delaunay"), ("nautilus", "datacube", "delaunay"), - ("nss", "point_source", "image_plane"), - ("nss", "point_source", "source_plane"), - ("nss", "imaging", "mge"), - ("nss", "imaging", "pixelization"), - ("nss", "imaging", "delaunay"), - ("nss", "interferometer", "mge"), - ("nss", "interferometer", "pixelization"), - ("nss", "interferometer", "delaunay"), - ("nss", "datacube", "delaunay"), ]