W6: Nautilus n_batch scan — SEARCHES_NAUTILUS_N_BATCH override + A100 array submits (#163)

hpc/batch_gpu/submit_search_nautilus_imaging_delaunay_a100_hst_fp64_nbatch

@@ -0,0 +1,78 @@
+#!/bin/bash -l
+#
+# W6 (issue #163) — Nautilus n_batch scan, searches/nautilus/imaging/delaunay x hst x fp64.
+#
+# QUESTION: Phase 2 "Sample economy" — Nautilus needs ~15 likelihood evals per
+#      effective sample but pays 11.4 ms/eval on MGE vs af.NSS's 2.85 ms for
+#      the SAME likelihood: the proposal loop is host-side numpy/sklearn around
+#      n_batch-sized GPU batches (re-baseline 339070 ran n_batch=64; fork-era
+#      rows 100). If ms/eval falls toward ~3 ms as n_batch grows, the "JAX
+#      Nautilus" question is closed as a config knob and W2/W4 adopt the best
+#      value; if it does not, the host overhead is recorded as irreducible.
+#      n_batch changes WHEN the network retrains, so logZ / max logL are
+#      reported per arm to confirm the answer does not move.
+#
+# ARM: one n_batch per array task — BATCHES=(16 64 256); n_live and every other
+#      knob at the cell default. --config-name carries n_batch: it sets the
+#      results basename AND the af.Nautilus search name, and n_batch is NOT an
+#      identifier field, so this is the ONLY thing keeping the arms' output/
+#      directories and results JSONs disjoint (see the seeds submits).
+#
+# WALL-TIME: re-baseline 1,939 s total at n_batch=16 (VRAM-probed cap; 64/256 rely on the 80 GB A100). 3:00:00 is ~5.5x.
+#
+# WALL-BASIS: — one row per cell this submit runs. See scripts/misc/wall/README.md.
+#   `rates` cites a step rate measured on THAT cell (wall/rates.py); an MGE rate
+#   can never stand in for a pixelized cell — that is what killed RAL job 340576.
+#   Every arm of this submit has now RUN, so the basis is this cell's own
+#   completed totals rather than the re-baseline the WALL-TIME line cites:
+#   n_batch 16/64/256 came in at 2088.2 / 1750.8 / 1758.4 s. NOTE the slowest
+#   arm (2088 s) OVERRAN that 1,939 s re-baseline by 8% — the --time was ample
+#   either way, but the re-baseline was not itself a safe ceiling.
+#   cell: imaging/delaunay/hst  device: a100  precision: fp64
+#   wall: 2089  source: measured-wall  ref: nbatch-scan-arms-max
+#   compile: 0  headroom: 1.25
+
+#SBATCH -J search_nautilus_imaging_delaunay_hst_fp64_nbatch
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=64gb
+#SBATCH --time=3:00:00
+#SBATCH --array=0-2
+#SBATCH -o output/output.%A_%a.out
+#SBATCH -e error/error.%A_%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/nautilus/imaging/delaunay"
+echo "Instrument: hst"
+echo "Precision:  fp64"
+echo "Arm:        W6 n_batch scan — n_batch=$NB (SEARCHES_NAUTILUS_N_BATCH), default n_live, cores=1"
+
+BATCHES=(16 64 256)
+NB=${BATCHES[$SLURM_ARRAY_TASK_ID]}
+export SEARCHES_NAUTILUS_N_BATCH=$NB
+
+cd $AP_ROOT
+python3 scripts/imaging/searches/nautilus/delaunay.py \
+    --instrument hst \
+    --config-name hpc_a100_fp64_nbatch${NB} \
+    --output-dir $AP_ROOT/results/searches/nautilus/imaging/delaunay/hst
+
+echo "Finished."
+date

hpc/batch_gpu/submit_search_nautilus_imaging_mge_a100_hst_fp64_nbatch

@@ -0,0 +1,78 @@
+#!/bin/bash -l
+#
+# W6 (issue #163) — Nautilus n_batch scan, searches/nautilus/imaging/mge x hst x fp64.
+#
+# QUESTION: Phase 2 "Sample economy" — Nautilus needs ~15 likelihood evals per
+#      effective sample but pays 11.4 ms/eval on MGE vs af.NSS's 2.85 ms for
+#      the SAME likelihood: the proposal loop is host-side numpy/sklearn around
+#      n_batch-sized GPU batches (re-baseline 339070 ran n_batch=64; fork-era
+#      rows 100). If ms/eval falls toward ~3 ms as n_batch grows, the "JAX
+#      Nautilus" question is closed as a config knob and W2/W4 adopt the best
+#      value; if it does not, the host overhead is recorded as irreducible.
+#      n_batch changes WHEN the network retrains, so logZ / max logL are
+#      reported per arm to confirm the answer does not move.
+#
+# ARM: one n_batch per array task — BATCHES=(64 128 256 512 1000); n_live and every other
+#      knob at the cell default. --config-name carries n_batch: it sets the
+#      results basename AND the af.Nautilus search name, and n_batch is NOT an
+#      identifier field, so this is the ONLY thing keeping the arms' output/
+#      directories and results JSONs disjoint (see the seeds submits).
+#
+# WALL-TIME: re-baseline 775 s total at n_batch=64; larger batches should be faster. 1:30:00 is ~7x.
+#
+# WALL-BASIS: — one row per cell this submit runs. See scripts/misc/wall/README.md.
+#   `rates` cites a step rate measured on THAT cell (wall/rates.py); an MGE rate
+#   can never stand in for a pixelized cell — that is what killed RAL job 340576.
+#   Every arm of this submit has now RUN, so the basis is this cell's own
+#   completed totals rather than the re-baseline the WALL-TIME line cites:
+#   n_batch 64/128/256/512/1000 came in at 737.8 / 680.8 / 612.3 / 633.2 / 525.6 s.
+#   `wall` is the slowest of them, not the mean — the budget has to cover the
+#   worst arm, and n_batch=64 is the slowest by construction here.
+#   cell: imaging/mge/hst  device: a100  precision: fp64
+#   wall: 738  source: measured-wall  ref: nbatch-scan-arms-max
+#   compile: 0  headroom: 1.25
+
+#SBATCH -J search_nautilus_imaging_mge_hst_fp64_nbatch
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=64gb
+#SBATCH --time=1:30:00
+#SBATCH --array=0-4
+#SBATCH -o output/output.%A_%a.out
+#SBATCH -e error/error.%A_%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/nautilus/imaging/mge"
+echo "Instrument: hst"
+echo "Precision:  fp64"
+echo "Arm:        W6 n_batch scan — n_batch=$NB (SEARCHES_NAUTILUS_N_BATCH), default n_live, cores=1"
+
+BATCHES=(64 128 256 512 1000)
+NB=${BATCHES[$SLURM_ARRAY_TASK_ID]}
+export SEARCHES_NAUTILUS_N_BATCH=$NB
+
+cd $AP_ROOT
+python3 scripts/imaging/searches/nautilus/mge.py \
+    --instrument hst \
+    --config-name hpc_a100_fp64_nbatch${NB} \
+    --output-dir $AP_ROOT/results/searches/nautilus/imaging/mge/hst
+
+echo "Finished."
+date

results/notes/inference/PROGRAMME.md

@@ -964,7 +964,7 @@ ship as autolens_profiling issues/PRs, not Mind prompts).
 |---|---|---|---|---|
 | W1 | Phase 4 Stage 1 — PositionsLH plumbing in `_setup.py` + hazard transects | autolens_profiling#159 | laptop / RAL CPU | — (unblocks B pt 2, Phase 5) |
 | W3 | Phase 6 pre-req — PyAutoFit warm-start abstraction, vmapped multi-chain NUTS, metric/start injection | PyAutoFit#1521 | source, CPU | — |
-| W6 | Nautilus `n_batch` scan (MGE + Delaunay) — bounds the "JAX Nautilus" ceiling | #163 | A100, ~30 min | — (submit now) |
+| W6 | Nautilus `n_batch` scan (MGE + Delaunay) — bounds the "JAX Nautilus" ceiling | #163 | A100, ~30 min | **COMPLETE 2026-08-25** (RAL 339842/339843, 8/8 arms): MGE 1.78x recoverable (10.56 -> 5.95 ms/eval at n_batch 64 -> 1000, logZ flat to 0.12 nats, not yet plateaued); Delaunay saturates by n_batch=64 at 1.26x. Partly a knob, not the full 4x — `methods/nautilus.md` "n_batch SCAN" |
 | W8 | slogdet default for GPU gradient-work cells in the searches framework | #165 | source | — |
 | W7 | CP-4 follow-up: NaN-under-both draws, transect gradients, marginal-band tier dependence | #164 | laptop / A100 replay | — |
 | W4 | Phase 1 targets registry, schema v2 (ESS + reject-inclusive evals), `slam_source_pix`, reference baselines — **HARVESTED 2026-08-26 (RAL 340210): 7 of 11 rows certified.** 2 lost to missing `_N_LIVE` presets (fixed, resubmit pending); 2 lost to a Sibson jit defect blocking `slam_source_pix_nn` (PyAutoArray, filed) — `targets/REFS_V1_HARVEST.md` | #161 | CPU + one A100 bake | rides behind W1/W6 |

results/notes/inference/methods/nautilus.md

@@ -82,6 +82,50 @@ current truth-bar source.
 - Full matrix: the searches dashboard
   (`scripts/misc/searches/README.md`, auto-table `searches`).
 
+## n_batch SCAN (W6, issue #163 — A100, 2026-08-25)
+
+Question: is the JAX-Nautilus per-eval overhead a config knob, or structural?
+Answer: **partly a knob on MGE, essentially structural on Delaunay.**
+
+MGE/hst, `n_live=200`, fp64, one seed per arm
+(`results/searches/nautilus/imaging/mge/hst/hpc_hpc_a100_fp64_nbatch*.json`):
+
+| n_batch | ms/eval | sampler wall | evals | Kish ESS | logZ |
+|--------:|--------:|-------------:|------:|---------:|-----:|
+| 64      | 10.56   | 670 s        | 63,424 | 4,304   | 31690.45 |
+| 128     |  9.52   | 617 s        | 64,768 | 4,248   | 31690.45 |
+| 256     |  8.15   | 540 s        | 66,304 | 4,156   | 31690.48 |
+| 512     |  8.37   | 565 s        | 67,584 | 4,575   | 31690.47 |
+| 1000    |  **5.95** | **458 s**  | 77,000 | 4,694   | 31690.36 |
+
+1.78x of the per-eval cost is recovered for free: logZ spans 0.12 nats across
+the whole scan, max logL 31786.73-31787.04, best-fit r_E 1.5996-1.5998, and
+Kish ESS is flat. The scan has **not** plateaued at n_batch=1000, but the
+trade is visible — evals rise 21% (63,424 -> 77,000) as larger batches
+overshoot the shrinking live set, while wall still falls. There is an optimum
+past 1000 that this scan does not bracket.
+
+Delaunay/hst, `n_live=150`
+(`.../imaging/delaunay/hst/hpc_hpc_a100_fp64_nbatch*.json`):
+
+| n_batch | ms/eval | sampler wall | evals | Kish ESS | logZ |
+|--------:|--------:|-------------:|------:|---------:|-----:|
+| 16      | 66.73   | 2,031 s      | 30,432 | 2,464   | 30562.10 |
+| 64      | 52.95   | 1,698 s      | 32,064 | 2,335   | 30562.17 |
+| 256     | 50.80   | 1,704 s      | 33,536 | 2,387   | 30562.17 |
+
+Saturates by n_batch=64 (1.26x, then flat). The pixelized cell's cost is
+dominated by the per-eval inversion, not by batch occupancy, so batching
+cannot buy back what MGE's cheaper likelihood gives up to launch overhead.
+
+READING: raise `n_batch` on parametric cells; leave it at the default on
+pixelized ones. This does **not** close the ~4x JAX-vs-NumPy per-eval gap —
+it recovers under half of it on the cell where it is recoverable at all.
+
+CONFIDENCE: **single-seed per arm.** logZ agreement across five independent
+MGE arms is itself weak evidence of seed-stability, but no arm was repeated;
+treat the wall numbers as one draw each.
+
 ## RECOMMENDED
 
 - All SLaM stages today (baseline pipeline); the global engine wherever no

results/searches/nautilus/imaging/delaunay/hst/hpc_hpc_a100_fp64_nbatch16.json

@@ -0,0 +1,150 @@
+{
+  "sampler": "nautilus",
+  "dataset_class": "imaging",
+  "model": "delaunay",
+  "instrument": "hst",
+  "config_name": "hpc_a100_fp64_nbatch16",
+  "version": "2026.8.17.1",
+  "device": {
+    "backend": "gpu",
+    "device": "cuda:0",
+    "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0",
+    "omp_num_threads": null,
+    "cpu_count": 124,
+    "nvidia_smi": "NVIDIA A100 80GB PCIe, 5149 MiB, 81920 MiB"
+  },
+  "use_mixed_precision": false,
+  "sampler_config": {
+    "n_live": 150,
+    "n_batch": 16,
+    "number_of_cores": 1,
+    "use_jax_vmap": true,
+    "force_x1_cpu": true,
+    "iterations_per_update": 450,
+    "seed": null,
+    "positions": {
+      "enabled": false
+    }
+  },
+  "positions": {
+    "enabled": false
+  },
+  "log_det_method": null,
+  "model_summary": {
+    "free_parameters": 12,
+    "best_fit": "lens.mass.einstein_radius=1.0276  lens.mass.centre=(0.013, 0.053)  shear=(-0.1322, 0.0317)"
+  },
+  "results": {
+    "log_evidence": 30562.10070250484,
+    "max_log_likelihood": 30623.80545400236,
+    "posterior_samples": 30394,
+    "truth_log_likelihood": 31521.743789581255,
+    "delta_max_ll_vs_truth": -897.9383355788959,
+    "bar_source": "truth_tracer",
+    "posterior_stats": {
+      "centre_0": {
+        "mean": 0.01061524325203249,
+        "std": 0.010397783201260884
+      },
+      "centre_1": {
+        "mean": 0.05449800358496196,
+        "std": 0.004480491447109771
+      },
+      "ell_comps_0": {
+        "mean": 0.002812239829831496,
+        "std": 0.013176746148423025
+      },
+      "ell_comps_1": {
+        "mean": -0.4817502077803952,
+        "std": 0.01065440426431391
+      },
+      "einstein_radius": {
+        "mean": 1.0387916772156465,
+        "std": 0.011542989579443238
+      },
+      "gamma_1": {
+        "mean": -0.13398776703271736,
+        "std": 0.005890816681999582
+      },
+      "gamma_2": {
+        "mean": 0.02912851119625473,
+        "std": 0.006006543814750087
+      },
+      "coefficient": {
+        "mean": 1.422216308752476,
+        "std": 0.0424122028153029
+      }
+    }
+  },
+  "performance": {
+    "total_wall_s": 2088.1522002220154,
+    "viz_wall_s": 57.415554126841016,
+    "viz_n_calls": 8,
+    "viz_disabled": false,
+    "sampler_wall_s": 2030.7366460951744,
+    "likelihood_evals": 30432,
+    "time_per_eval_ms": 66.73030514245447,
+    "stored_samples": 30432,
+    "gradient_evals": null,
+    "kish_ess": 2464.346615985722,
+    "evals_per_ess": 12.348912203581152,
+    "ess_per_min": 72.81140922111156
+  },
+  "schema_version": 2,
+  "target": {
+    "target_id": "sha256:a45841ff08e6",
+    "cell": "imaging/delaunay/hst",
+    "model_dim": 12,
+    "priors_ref": "_targets.py@bf2c8742c334",
+    "likelihood": {
+      "log_det_method": null,
+      "positive_only": true,
+      "border_relocator": true,
+      "curvature_floor": 0.001,
+      "regularization_term_method": "matmul",
+      "precision": "fp64"
+    },
+    "positions": {
+      "enabled": false
+    },
+    "tolerances": {
+      "delta_max_ll_nats": 2.0,
+      "mean_shift_sigma": 0.2,
+      "sigma_ratio": [
+        0.8,
+        1.25
+      ],
+      "per_parameter_recovery": {}
+    },
+    "target_class_vs_v1": null
+  },
+  "algorithm": {
+    "name": "nautilus",
+    "config_id": "hpc_a100_fp64_nbatch16",
+    "settings": {
+      "n_live": 150,
+      "n_batch": 16,
+      "number_of_cores": 1,
+      "use_jax_vmap": true,
+      "force_x1_cpu": true,
+      "iterations_per_update": 450,
+      "seed": null,
+      "positions": {
+        "enabled": false
+      }
+    },
+    "seed": null
+  },
+  "hardware": {
+    "tier": "hpc_a100_fp64_nbatch16",
+    "precision": "fp64",
+    "device": {
+      "backend": "gpu",
+      "device": "cuda:0",
+      "xla_flags": "--xla_disable_hlo_passes=constant_folding --xla_gpu_autotune_level=0",
+      "omp_num_threads": null,
+      "cpu_count": 124,
+      "nvidia_smi": "NVIDIA A100 80GB PCIe, 5149 MiB, 81920 MiB"
+    }
+  }
+}