A100 re-measure: JAX Delaunay walk early exit (PyAutoArray#531)

hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64_control

@@ -0,0 +1,85 @@
+#!/bin/bash -l
+#
+# A100 likelihood-breakdown profile: imaging/delaunay x hst x fp64 — PyAutoArray#531
+# A/B leg: **control**.
+#
+# Verification run for the JAX Delaunay walk early-exit change
+# (PyAutoArray issue #530 / PR #531: lax.while_loop early exit, seed argmin
+# chunked, single concatenated locate). Both legs of the A/B run in the same
+# session on the same node class with the same everything except the
+# PyAutoArray checkout prepended to PYTHONPATH:
+#
+#   control  -> bcd15cd9 (the branch's merge base = PyAutoArray main)
+#   feature  -> 6152bdf0 (feature/delaunay-walk-early-exit)
+#
+# The shared /mnt/ral/jnightin/PyAuto/PyAutoArray install is NOT touched — long
+# science runs (subhalo validation, Euclid DR1) share it. The private checkouts
+# live under /mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/ and are
+# prepended to the PYTHONPATH that $AP_ROOT/activate.sh exports, so only
+# PyAutoArray is overridden and Nerves/Fit/Galaxy/Lens stay on the shared tree.
+#
+# WITNESS. The feature concatenates the data-grid and split-point walks into a
+# single locate call, so the split walk now runs *inside* the "Triangulation +
+# interpolation" prefix; before, XLA dead-code-eliminated it there and its cost
+# was charged to the "Regularization matrix (H)" row (a prefix difference).
+# Compare the **params->H prefix** (regularization_matrix_prefix_s, equivalently
+# Tri+interp + H) across the two legs — the H row alone can go ~0 or negative.
+#
+#   --split-setup    four-way split of the inversion-setup block
+#   --vmap-batch 16  batched re-timing at the Nautilus reference n_batch
+
+#SBATCH -J bd_delaunay_ctrl
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=64gb
+#SBATCH --time=12:00:00
+#SBATCH -o output/output.%A.out
+#SBATCH -e error/error.%A.err
+
+export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
+echo "AP_ROOT=$AP_ROOT"
+source $AP_ROOT/activate.sh
+
+# --- PyAutoArray A/B override (private checkout, shared install untouched) ---
+AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/PyAutoArray_control
+export PYTHONPATH=$AA_ROOT:$PYTHONPATH
+
+export JAX_PLATFORM_NAME=cuda
+export JAX_PLATFORMS=cuda,cpu
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export JAX_ENABLE_X64=True
+export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding"
+export NUMBA_CACHE_DIR=/tmp/numba_cache
+export MPLCONFIGDIR=/tmp/matplotlib
+
+nvidia-smi
+
+echo "=========================================="
+date
+echo "Cell:           likelihood_breakdown/imaging/delaunay"
+echo "Instrument:     hst"
+echo "Precision:      fp64"
+echo "A/B leg:        control"
+echo "Flags:          --split-setup --vmap-batch 16 --config-name hpc_a100_fp64_control_2026_09_07"
+
+echo "=== source revisions ==="
+printf '%-14s %s\n' "PyAutoArray(AB)" "$(git -C $AA_ROOT rev-parse HEAD 2>&1)"
+for r in PyAutoNerves PyAutoFit PyAutoGalaxy PyAutoLens; do
+    printf '%-14s %s\n' "$r" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)"
+done
+echo "AP_ROOT rev:   $(git -C $AP_ROOT rev-parse HEAD 2>&1)"
+
+cd $AP_ROOT
+
+echo "=== import provenance ==="
+python3 -c "import autoarray, autolens; print('autoarray.__file__ =', autoarray.__file__); print('autoarray.__version__ =', autoarray.__version__); print('autolens.__file__ =', autolens.__file__)"
+
+python3 -u scripts/imaging/likelihood_breakdown/delaunay.py \
+    --config-name hpc_a100_fp64_control_2026_09_07 \
+    --split-setup \
+    --vmap-batch 16
+
+echo "Finished."
+date

hpc/batch_gpu/submit_breakdown_imaging_delaunay_a100_hst_fp64_walk_early_exit

@@ -0,0 +1,85 @@
+#!/bin/bash -l
+#
+# A100 likelihood-breakdown profile: imaging/delaunay x hst x fp64 — PyAutoArray#531
+# A/B leg: **walk_early_exit**.
+#
+# Verification run for the JAX Delaunay walk early-exit change
+# (PyAutoArray issue #530 / PR #531: lax.while_loop early exit, seed argmin
+# chunked, single concatenated locate). Both legs of the A/B run in the same
+# session on the same node class with the same everything except the
+# PyAutoArray checkout prepended to PYTHONPATH:
+#
+#   control  -> bcd15cd9 (the branch's merge base = PyAutoArray main)
+#   feature  -> 6152bdf0 (feature/delaunay-walk-early-exit)
+#
+# The shared /mnt/ral/jnightin/PyAuto/PyAutoArray install is NOT touched — long
+# science runs (subhalo validation, Euclid DR1) share it. The private checkouts
+# live under /mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/ and are
+# prepended to the PYTHONPATH that $AP_ROOT/activate.sh exports, so only
+# PyAutoArray is overridden and Nerves/Fit/Galaxy/Lens stay on the shared tree.
+#
+# WITNESS. The feature concatenates the data-grid and split-point walks into a
+# single locate call, so the split walk now runs *inside* the "Triangulation +
+# interpolation" prefix; before, XLA dead-code-eliminated it there and its cost
+# was charged to the "Regularization matrix (H)" row (a prefix difference).
+# Compare the **params->H prefix** (regularization_matrix_prefix_s, equivalently
+# Tri+interp + H) across the two legs — the H row alone can go ~0 or negative.
+#
+#   --split-setup    four-way split of the inversion-setup block
+#   --vmap-batch 16  batched re-timing at the Nautilus reference n_batch
+
+#SBATCH -J bd_delaunay_walk
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=64gb
+#SBATCH --time=12:00:00
+#SBATCH -o output/output.%A.out
+#SBATCH -e error/error.%A.err
+
+export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
+echo "AP_ROOT=$AP_ROOT"
+source $AP_ROOT/activate.sh
+
+# --- PyAutoArray A/B override (private checkout, shared install untouched) ---
+AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/PyAutoArray_feature
+export PYTHONPATH=$AA_ROOT:$PYTHONPATH
+
+export JAX_PLATFORM_NAME=cuda
+export JAX_PLATFORMS=cuda,cpu
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export JAX_ENABLE_X64=True
+export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding"
+export NUMBA_CACHE_DIR=/tmp/numba_cache
+export MPLCONFIGDIR=/tmp/matplotlib
+
+nvidia-smi
+
+echo "=========================================="
+date
+echo "Cell:           likelihood_breakdown/imaging/delaunay"
+echo "Instrument:     hst"
+echo "Precision:      fp64"
+echo "A/B leg:        walk_early_exit"
+echo "Flags:          --split-setup --vmap-batch 16 --config-name hpc_a100_fp64_walk_early_exit"
+
+echo "=== source revisions ==="
+printf '%-14s %s\n' "PyAutoArray(AB)" "$(git -C $AA_ROOT rev-parse HEAD 2>&1)"
+for r in PyAutoNerves PyAutoFit PyAutoGalaxy PyAutoLens; do
+    printf '%-14s %s\n' "$r" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)"
+done
+echo "AP_ROOT rev:   $(git -C $AP_ROOT rev-parse HEAD 2>&1)"
+
+cd $AP_ROOT
+
+echo "=== import provenance ==="
+python3 -c "import autoarray, autolens; print('autoarray.__file__ =', autoarray.__file__); print('autoarray.__version__ =', autoarray.__version__); print('autolens.__file__ =', autolens.__file__)"
+
+python3 -u scripts/imaging/likelihood_breakdown/delaunay.py \
+    --config-name hpc_a100_fp64_walk_early_exit \
+    --split-setup \
+    --vmap-batch 16
+
+echo "Finished."
+date

hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_control

@@ -0,0 +1,85 @@
+#!/bin/bash -l
+#
+# A100 likelihood-breakdown profile: imaging/delaunay_nn x hst x fp64 — PyAutoArray#531
+# A/B leg: **control**.
+#
+# Verification run for the JAX Delaunay walk early-exit change
+# (PyAutoArray issue #530 / PR #531: lax.while_loop early exit, seed argmin
+# chunked, single concatenated locate). Both legs of the A/B run in the same
+# session on the same node class with the same everything except the
+# PyAutoArray checkout prepended to PYTHONPATH:
+#
+#   control  -> bcd15cd9 (the branch's merge base = PyAutoArray main)
+#   feature  -> 6152bdf0 (feature/delaunay-walk-early-exit)
+#
+# The shared /mnt/ral/jnightin/PyAuto/PyAutoArray install is NOT touched — long
+# science runs (subhalo validation, Euclid DR1) share it. The private checkouts
+# live under /mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/ and are
+# prepended to the PYTHONPATH that $AP_ROOT/activate.sh exports, so only
+# PyAutoArray is overridden and Nerves/Fit/Galaxy/Lens stay on the shared tree.
+#
+# WITNESS. The feature concatenates the data-grid and split-point walks into a
+# single locate call, so the split walk now runs *inside* the "Triangulation +
+# interpolation" prefix; before, XLA dead-code-eliminated it there and its cost
+# was charged to the "Regularization matrix (H)" row (a prefix difference).
+# Compare the **params->H prefix** (regularization_matrix_prefix_s, equivalently
+# Tri+interp + H) across the two legs — the H row alone can go ~0 or negative.
+#
+#   --split-setup    four-way split of the inversion-setup block
+#   --vmap-batch 16  batched re-timing at the Nautilus reference n_batch
+
+#SBATCH -J bd_delnn_ctrl
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=64gb
+#SBATCH --time=12:00:00
+#SBATCH -o output/output.%A.out
+#SBATCH -e error/error.%A.err
+
+export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
+echo "AP_ROOT=$AP_ROOT"
+source $AP_ROOT/activate.sh
+
+# --- PyAutoArray A/B override (private checkout, shared install untouched) ---
+AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/PyAutoArray_control
+export PYTHONPATH=$AA_ROOT:$PYTHONPATH
+
+export JAX_PLATFORM_NAME=cuda
+export JAX_PLATFORMS=cuda,cpu
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export JAX_ENABLE_X64=True
+export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding"
+export NUMBA_CACHE_DIR=/tmp/numba_cache
+export MPLCONFIGDIR=/tmp/matplotlib
+
+nvidia-smi
+
+echo "=========================================="
+date
+echo "Cell:           likelihood_breakdown/imaging/delaunay_nn"
+echo "Instrument:     hst"
+echo "Precision:      fp64"
+echo "A/B leg:        control"
+echo "Flags:          --split-setup --vmap-batch 16 --config-name hpc_a100_fp64_control_2026_09_07"
+
+echo "=== source revisions ==="
+printf '%-14s %s\n' "PyAutoArray(AB)" "$(git -C $AA_ROOT rev-parse HEAD 2>&1)"
+for r in PyAutoNerves PyAutoFit PyAutoGalaxy PyAutoLens; do
+    printf '%-14s %s\n' "$r" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)"
+done
+echo "AP_ROOT rev:   $(git -C $AP_ROOT rev-parse HEAD 2>&1)"
+
+cd $AP_ROOT
+
+echo "=== import provenance ==="
+python3 -c "import autoarray, autolens; print('autoarray.__file__ =', autoarray.__file__); print('autoarray.__version__ =', autoarray.__version__); print('autolens.__file__ =', autolens.__file__)"
+
+python3 -u scripts/imaging/likelihood_breakdown/delaunay_nn.py \
+    --config-name hpc_a100_fp64_control_2026_09_07 \
+    --split-setup \
+    --vmap-batch 16
+
+echo "Finished."
+date

hpc/batch_gpu/submit_breakdown_imaging_delaunay_nn_a100_hst_fp64_walk_early_exit

@@ -0,0 +1,85 @@
+#!/bin/bash -l
+#
+# A100 likelihood-breakdown profile: imaging/delaunay_nn x hst x fp64 — PyAutoArray#531
+# A/B leg: **walk_early_exit**.
+#
+# Verification run for the JAX Delaunay walk early-exit change
+# (PyAutoArray issue #530 / PR #531: lax.while_loop early exit, seed argmin
+# chunked, single concatenated locate). Both legs of the A/B run in the same
+# session on the same node class with the same everything except the
+# PyAutoArray checkout prepended to PYTHONPATH:
+#
+#   control  -> bcd15cd9 (the branch's merge base = PyAutoArray main)
+#   feature  -> 6152bdf0 (feature/delaunay-walk-early-exit)
+#
+# The shared /mnt/ral/jnightin/PyAuto/PyAutoArray install is NOT touched — long
+# science runs (subhalo validation, Euclid DR1) share it. The private checkouts
+# live under /mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/ and are
+# prepended to the PYTHONPATH that $AP_ROOT/activate.sh exports, so only
+# PyAutoArray is overridden and Nerves/Fit/Galaxy/Lens stay on the shared tree.
+#
+# WITNESS. The feature concatenates the data-grid and split-point walks into a
+# single locate call, so the split walk now runs *inside* the "Triangulation +
+# interpolation" prefix; before, XLA dead-code-eliminated it there and its cost
+# was charged to the "Regularization matrix (H)" row (a prefix difference).
+# Compare the **params->H prefix** (regularization_matrix_prefix_s, equivalently
+# Tri+interp + H) across the two legs — the H row alone can go ~0 or negative.
+#
+#   --split-setup    four-way split of the inversion-setup block
+#   --vmap-batch 16  batched re-timing at the Nautilus reference n_batch
+
+#SBATCH -J bd_delnn_walk
+#SBATCH --partition=gpu
+#SBATCH --gres=gpu:1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=64gb
+#SBATCH --time=12:00:00
+#SBATCH -o output/output.%A.out
+#SBATCH -e error/error.%A.err
+
+export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
+echo "AP_ROOT=$AP_ROOT"
+source $AP_ROOT/activate.sh
+
+# --- PyAutoArray A/B override (private checkout, shared install untouched) ---
+AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/PyAutoArray_feature
+export PYTHONPATH=$AA_ROOT:$PYTHONPATH
+
+export JAX_PLATFORM_NAME=cuda
+export JAX_PLATFORMS=cuda,cpu
+export XLA_PYTHON_CLIENT_PREALLOCATE=false
+export JAX_ENABLE_X64=True
+export XLA_FLAGS="--xla_disable_hlo_passes=constant_folding"
+export NUMBA_CACHE_DIR=/tmp/numba_cache
+export MPLCONFIGDIR=/tmp/matplotlib
+
+nvidia-smi
+
+echo "=========================================="
+date
+echo "Cell:           likelihood_breakdown/imaging/delaunay_nn"
+echo "Instrument:     hst"
+echo "Precision:      fp64"
+echo "A/B leg:        walk_early_exit"
+echo "Flags:          --split-setup --vmap-batch 16 --config-name hpc_a100_fp64_walk_early_exit"
+
+echo "=== source revisions ==="
+printf '%-14s %s\n' "PyAutoArray(AB)" "$(git -C $AA_ROOT rev-parse HEAD 2>&1)"
+for r in PyAutoNerves PyAutoFit PyAutoGalaxy PyAutoLens; do
+    printf '%-14s %s\n' "$r" "$(git -C /mnt/ral/jnightin/PyAuto/$r rev-parse HEAD 2>&1)"
+done
+echo "AP_ROOT rev:   $(git -C $AP_ROOT rev-parse HEAD 2>&1)"
+
+cd $AP_ROOT
+
+echo "=== import provenance ==="
+python3 -c "import autoarray, autolens; print('autoarray.__file__ =', autoarray.__file__); print('autoarray.__version__ =', autoarray.__version__); print('autolens.__file__ =', autolens.__file__)"
+
+python3 -u scripts/imaging/likelihood_breakdown/delaunay_nn.py \
+    --config-name hpc_a100_fp64_walk_early_exit \
+    --split-setup \
+    --vmap-batch 16
+
+echo "Finished."
+date

hpc/batch_gpu/submit_runtime_imaging_delaunay_a100_hst_fp64_control

@@ -0,0 +1,50 @@
+#!/bin/bash -l
+#
+# A100 likelihood-runtime profile: imaging/delaunay x hst x fp64 — PyAutoArray#531
+# A/B leg: **control**. Single-JIT whole-likelihood cost, the end-to-end companion
+# to submit_breakdown_imaging_delaunay_a100_hst_fp64_control.
+#
+# Same private-checkout PYTHONPATH override as the breakdown A/B legs; the
+# shared /mnt/ral/jnightin/PyAuto install is untouched. Phase A (--vmap-probe)
+# is skipped: only the full single-JIT + vmap timing is wanted here.
+
+#SBATCH -J rt_delaunay_control
+#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
+
+export AP_ROOT="$(cd "${SLURM_SUBMIT_DIR:-$(dirname "$0")}/../.." && pwd)"
+echo "AP_ROOT=$AP_ROOT"
+source $AP_ROOT/activate.sh
+
+AA_ROOT=/mnt/ral/jnightin/PyAuto_wt/delaunay-walk-early-exit/PyAutoArray_control
+export PYTHONPATH=$AA_ROOT:$PYTHONPATH
+
+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:           likelihood_runtime/imaging/delaunay"
+echo "A/B leg:        control"
+printf 'PyAutoArray(AB) %s\n' "$(git -C $AA_ROOT rev-parse HEAD 2>&1)"
+
+cd $AP_ROOT
+python3 -c "import autoarray; print('autoarray.__file__ =', autoarray.__file__)"
+
+python3 -u scripts/imaging/likelihood_runtime/delaunay.py \
+    --instrument hst \
+    --config-name hpc_a100_fp64_control_2026_09_07
+
+echo "Finished."
+date