From f9d70b1b3d2aeb9489816615b2810dcb4f096bf0 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Fri, 10 Jul 2026 16:21:13 +0100 Subject: [PATCH] =?UTF-8?q?perf:=20multi=20(imaging)=20shared=5Fpreloads?= =?UTF-8?q?=20runtime=20measurement=20=E2=80=94=201.14x=20at=20hst=20x=204?= =?UTF-8?q?=20exposures?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mesh-geometry-only sharing as designed (PyAutoLens#599 D1): the ~12% saving is the per-exposure image-mesh ray-trace + Delaunay triangulation; consistency (identical source-pixel grid across exposures) is the primary win. Includes the NaN adapt-cache re-run trap in OPTIMIZATION_NOTES. Co-Authored-By: Claude Fable 5 --- likelihood_runtime/OPTIMIZATION_NOTES.md | 28 +++ likelihood_runtime/multi/__init__.py | 0 likelihood_runtime/multi/shared_preloads.py | 181 ++++++++++++++++++ .../shared_preloads_hst_local_cpu_fp64.stdout | 8 + 4 files changed, 217 insertions(+) create mode 100644 likelihood_runtime/multi/__init__.py create mode 100644 likelihood_runtime/multi/shared_preloads.py create mode 100644 results/runtime/multi/shared_preloads_hst_local_cpu_fp64.stdout diff --git a/likelihood_runtime/OPTIMIZATION_NOTES.md b/likelihood_runtime/OPTIMIZATION_NOTES.md index f533083..d6d07a9 100644 --- a/likelihood_runtime/OPTIMIZATION_NOTES.md +++ b/likelihood_runtime/OPTIMIZATION_NOTES.md @@ -660,6 +660,34 @@ with raised `N_CHANNELS` on a quiet A100 (needs the RAL login node). SMA seconds are provisional; the ratio is the robust deliverable, and it should grow with per-channel inversion cost at ALMA visibility counts. +### multi (imaging) shared_preloads — measured (2026-07-10) + +The imaging generalisation (PyAutoArray#380 / PyAutoLens#600, design +PyAutoLens#599 D1–D6): exposures of one lens share the **source-plane +mesh geometry only** — per-exposure PSFs/offsets keep the mapping +matrix, blurred mapping matrix, curvature matrix and regularization +matrix per-exposure, so the design predicted a modest speed-up with +**consistency** (identical source-pixel grid across exposures) as the +primary win. Measured by `likelihood_runtime/multi/shared_preloads.py` +(vmap-honest, 4 identical exposures, Hilbert-1500 + Delaunay): + +| Config | unshared | shared | speed-up | +|--------------------------------------------|----------------|---------------|-----------| +| hst × 4 exp, local_cpu_fp64 (v2026.7.6.649) | 53279.5 ms/eval | 46747.4 ms/eval | **1.14×** | + +As predicted: the ~12% saved is the per-exposure image-mesh ray-trace + +Delaunay `pure_callback` triangulation; the dominant per-exposure +inversion work is untouched by design. Raw log: +`results/runtime/multi/shared_preloads_hst_local_cpu_fp64.stdout`. +Correctness (identical-exposure shared-vs-unshared bit-parity + g+r +shared-mesh jit) is asserted in +`autolens_workspace_test/scripts/jax_likelihood_functions/multi/shared_preloads.py`. + +Trap for re-runs: a crashed JAX run can poison the gitignored +`dataset/imaging//lensed_source.fits` adapt cache with in-mask +NaNs (qhull then fails with "Points cannot contain NaN") — delete the +cache and let `_adapt_image_util` regenerate it. + --- ## Cross-cell: known follow-ups diff --git a/likelihood_runtime/multi/__init__.py b/likelihood_runtime/multi/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/likelihood_runtime/multi/shared_preloads.py b/likelihood_runtime/multi/shared_preloads.py new file mode 100644 index 0000000..6af8f32 --- /dev/null +++ b/likelihood_runtime/multi/shared_preloads.py @@ -0,0 +1,181 @@ +""" +Multi Runtime: shared_preloads speed-up +======================================= + +Measures the realized end-to-end speed-up of the `FactorGraphModel` shared-state path for +multi-exposure imaging: the same identical exposure is fitted N times with `shared_preloads=False` +(every exposure computes its own image-mesh ray-trace and source-plane mesh) and with +`shared_preloads=True` (the lead exposure ray-traces the source-plane Delaunay mesh once and every +exposure maps its own grid onto it), and the ratio is reported. + +Unlike the datacube (`likelihood_runtime/datacube/shared_preloads.py`), only the mesh geometry is +shared — each exposure still builds its own mapping matrix, PSF-blurred mapping matrix, curvature +matrix and regularization matrix, which dominate the imaging inversion budget. The expected speed-up +is therefore modest; the primary win of the shared path is *consistency* (every exposure reconstructs +on the identical source-pixel grid — see PyAutoLens#599 D1). This script records the honest number. + +Timing uses `Fitness._vmap` over a batch rather than a single jit on concrete parameters, because a +single jit over fixed parameters can constant-fold the work and report a misleadingly fast time; the +vmapped per-evaluation time is honest (see PyAutoLabs notes on jit const-folding). + +Correctness of the imaging `shared_preloads` is asserted separately in +`autolens_workspace_test/scripts/jax_likelihood_functions/multi/shared_preloads.py`; this script is +purely a runtime measurement. + +Run from the repo root: + + python likelihood_runtime/multi/shared_preloads.py [--instrument hst|hst_up|euclid|...] +""" + +import os +import sys +import time +from pathlib import Path + +import autofit as af +import autolens as al +import jax.numpy as jnp +import numpy as np +from autofit.non_linear.fitness import Fitness + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +# CI lint smoke: verify the import graph + module setup without the heavy timing run. +if os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": + from simulators.imaging import INSTRUMENTS # noqa: F401 + + print(f"[smoke] {__file__}: imports + module setup OK; exiting.") + sys.exit(0) + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 +from _profile_cli import auto_simulate_if_missing, parse_profile_cli # noqa: E402 +from simulators.imaging import INSTRUMENTS # noqa: E402 + +_cli = parse_profile_cli() +instrument = _cli.instrument or "hst" + +N_EXPOSURES = 4 +BATCH_SIZE = 2 +N_REPEATS = 5 +N_MESH_VERTICES = 1500 + +_workspace_root = Path(__file__).resolve().parents[2] +pixel_scale = INSTRUMENTS[instrument]["pixel_scale"] +dataset_path = Path("dataset") / "imaging" / instrument + +auto_simulate_if_missing( + dataset_path, + dataset_type="imaging", + instrument=instrument, + workspace_root=_workspace_root, +) + + +def _dataset(): + dataset = al.Imaging.from_fits( + data_path=dataset_path / "data.fits", + psf_path=dataset_path / "psf.fits", + noise_map_path=dataset_path / "noise_map.fits", + pixel_scales=pixel_scale, + ) + mask = al.Mask2D.circular( + shape_native=dataset.shape_native, + pixel_scales=dataset.pixel_scales, + radius=3.5, + ) + dataset = dataset.apply_mask(mask=mask) + return dataset.apply_over_sampling(over_sample_size_lp=1, over_sample_size_pixelization=1) + + +def _adapt_images(dataset): + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) + image_mesh = al.image_mesh.Hilbert(pixels=N_MESH_VERTICES, weight_power=1.0, weight_floor=0.0) + image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( + mask=dataset.mask, adapt_data=adapt_image + ) + return al.AdaptImages( + galaxy_name_image_dict={"('galaxies', 'source')": adapt_image}, + galaxy_name_image_plane_mesh_grid_dict={"('galaxies', 'source')": image_plane_mesh_grid}, + ) + + +def _model(): + mass = af.Model(al.mp.Isothermal) + mass.centre.centre_0 = af.GaussianPrior(mean=0.0, sigma=0.005) + mass.centre.centre_1 = af.GaussianPrior(mean=0.0, sigma=0.005) + mass.einstein_radius = af.GaussianPrior(mean=1.6, sigma=0.05) + mass.ell_comps.ell_comps_0 = af.GaussianPrior(mean=0.0, sigma=0.01) + mass.ell_comps.ell_comps_1 = af.GaussianPrior(mean=0.05, sigma=0.01) + + lens = af.Model(al.Galaxy, redshift=0.5, mass=mass) + + pixelization = af.Model( + al.Pixelization, + mesh=al.mesh.Delaunay(pixels=N_MESH_VERTICES, zeroed_pixels=0), + regularization=al.reg.ConstantSplit, + ) + source = af.Model(al.Galaxy, redshift=1.0, pixelization=pixelization) + + return af.Collection(galaxies=af.Collection(lens=lens, source=source)) + + +def _factor_graph(shared_preloads): + analysis_list = [] + for _ in range(N_EXPOSURES): + dataset = _dataset() + analysis_list.append( + al.AnalysisImaging( + dataset=dataset, + adapt_images=_adapt_images(dataset), + use_jax=True, + shared_preloads=shared_preloads, + raise_inversion_positions_likelihood_exception=False, + ) + ) + + model = _model() + analysis_factor_list = [ + af.AnalysisFactor(prior_model=model.copy(), analysis=analysis) for analysis in analysis_list + ] + + return af.FactorGraphModel(*analysis_factor_list, use_jax=True) + + +def _time_vmap(shared_preloads): + factor_graph = _factor_graph(shared_preloads) + + fitness = Fitness( + model=factor_graph.global_prior_model, + analysis=factor_graph, + fom_is_log_likelihood=True, + resample_figure_of_merit=-1.0e99, + ) + + medians = factor_graph.global_prior_model.physical_values_from_prior_medians + parameters = jnp.array(np.tile(np.asarray(medians), (BATCH_SIZE, 1))) + + fitness._vmap(parameters).block_until_ready() # warm-up (compile excluded) + + start = time.perf_counter() + for _ in range(N_REPEATS): + fitness._vmap(parameters).block_until_ready() + elapsed = time.perf_counter() - start + + return elapsed / (N_REPEATS * BATCH_SIZE) + + +if __name__ == "__main__": + print( + f"multi shared_preloads runtime " + f"(instrument={instrument}, N_EXPOSURES={N_EXPOSURES}, Hilbert+Delaunay mesh)" + ) + + per_eval_unshared = _time_vmap(shared_preloads=False) + per_eval_shared = _time_vmap(shared_preloads=True) + + print(f" per-likelihood (unshared): {per_eval_unshared * 1e3:.1f} ms") + print(f" per-likelihood (shared): {per_eval_shared * 1e3:.1f} ms") + print(f" speed-up: {per_eval_unshared / per_eval_shared:.2f}x") + print(" (the mesh is the only shared quantity for imaging — per-exposure blurred mapping") + print(" matrix / curvature / regularization dominate, so expect a modest ratio; the") + print(" primary win is mesh consistency across exposures — see PyAutoLens#599 D1)") diff --git a/results/runtime/multi/shared_preloads_hst_local_cpu_fp64.stdout b/results/runtime/multi/shared_preloads_hst_local_cpu_fp64.stdout new file mode 100644 index 0000000..4c17720 --- /dev/null +++ b/results/runtime/multi/shared_preloads_hst_local_cpu_fp64.stdout @@ -0,0 +1,8 @@ +You can also set the environment variable PYAUTO_SKIP_WORKSPACE_VERSION_CHECK=1 to disable temporarily. +multi shared_preloads runtime (instrument=hst, N_EXPOSURES=4, Hilbert+Delaunay mesh) + per-likelihood (unshared): 53279.5 ms + per-likelihood (shared): 46747.4 ms + speed-up: 1.14x + (the mesh is the only shared quantity for imaging — per-exposure blurred mapping + matrix / curvature / regularization dominate, so expect a modest ratio; the + primary win is mesh consistency across exposures — see PyAutoLens#599 D1)