Profiling and run-time tracking for PyAutoLens likelihood functions, simulators, and samplers across CPU, laptop GPU, and HPC GPU.
This repository is the single home for PyAutoLens performance measurement. It exists so that the run-times that matter for science — fitting a real lens, simulating an Euclid-resolution dataset, sampling a model with Nautilus — are visible, reproducible, and versioned across PyAutoLens releases.
What is profiled:
- 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.
Hardware tiers covered:
- CPU (single-machine, numpy backend).
- Laptop GPU (consumer-class, JAX backend).
- HPC GPU (A100 and similar, JAX backend).
Dataset framing:
Results are framed by astronomy instrument (HST, Euclid, JWST, …) rather than by raw pixel counts. Pixel counts are recorded too, but the headline numbers a reader sees first are the ones that map onto a real observing programme.
Likelihood runtime — full-pipeline per-call cost per cell × config:
No data yet — run likelihood_runtime/sweep.py then aggregate.py to populate.
Likelihood breakdown — latest per-step decompositions:
| Cell | Instrument | Inversion path | Step-sum total | PyAutoLens version |
|---|---|---|---|---|
imaging/delaunay |
hst | dense (mapping) | 11.29 s | v2026.5.29.4 |
imaging/delaunay |
hst | sparse (w-tilde) | 7.96 s | v2026.5.29.4 |
imaging/mge |
hst | dense (mapping) | 178.2 ms | v2026.5.29.4 |
imaging/pixelization |
hst | dense (mapping) | 7.79 s | v2026.5.29.4 |
imaging/pixelization |
hst | sparse (w-tilde) | 8.94 s | v2026.5.29.4 |
The tables above are auto-generated by scripts/build_readme.py from the artifacts under results/ — never edit them by hand; run python scripts/build_readme.py after a profiling run and commit the result (CI checks idempotence via --check). Narrative context — per-cell "where to optimize next" recommendations and the mp-vs-fp64 verdicts — lives in likelihood_runtime/OPTIMIZATION_NOTES.md.
PreOptimizationTimes is the named baseline the upcoming optimization work is measured against: a frozen snapshot of the full campaign (laptop CPU, HPC CPU, HPC A100 × fp64/mp) under results/baselines/PreOptimizationTimes/, rendered as a baseline column in the dashboard once populated. The convention is defined in results/notes/design_lock_in.md.
(Historical multi-config sweeps up to 2026-07 were committed under autolens_workspace_developer/jax_profiling/results/jit/; sweeps now write in-repo to results/runtime/ by default.)
Gradient profiling (jax.grad of the likelihood, autodiff-based optimisers) is not yet part of this repo. It is tracked in PyAutoLabs/autolens_workspace_developer/jax_profiling/gradient/ and will fold into this repo in a future phase once the gradient story stabilises.
Profiling scripts write two artifact shapes under results/ (full reference: results/README.md):
# Versioned summaries — standalone runs; history retained side-by-side
results/<section>/<subfolder>/<cell>_<purpose>_<instrument>_v<YYYY>.<M>.<D>.<PATCH>[_sparse].{json,png}
# Per-config sweeps — sweep.py + aggregate.py; latest sweep per cell
results/runtime/<class>/<model>[/<instrument>]/<config_name>[_sparse].{json,png,log} + comparison.{json,png}
The version string matches the PyAutoLens release that produced the numbers (e.g. v2026.5.29.4). The JSON carries structured timings; the PNG is the at-a-glance plot. Cross-release trend questions read the versioned summaries; cross-hardware comparison questions read comparison.json.
| Folder | Contents |
|---|---|
likelihood_runtime/ |
Full-pipeline JIT only, driven by sweep.py across CPU/GPU/A100 × fp64/mp. How long will this likelihood take on this hardware? |
likelihood_breakdown/ |
Per-step JIT decomposition. Single config. Where does time go inside the likelihood? |
vram/ |
GPU memory profiling + the per-cell vmap batch-size table for the A100. |
latent/ |
Latent-variable profiling. |
instruments/ |
Instrument presets (pixel scale, shape) that frame every result. |
simulators/ |
Run-time tracking for the PyAutoLens simulators. |
searches/ |
Sampler / search profiling, Nautilus first. |
quick_update/ |
Fast incremental re-profiling helpers. |
hpc/ |
SLURM submit scripts for the RAL HPC (A100 rows of the sweep matrix). |
results/ |
JSON + PNG artifacts written by the above scripts; named baselines. |
This repo is being built in phases (bootstrap history now archived in PyAutoMind).
| Phase | Title | Status |
|---|---|---|
| 0 | Repo bootstrap | ✓ shipped |
| 1 | Mirror JIT likelihood profiling scripts + per-section READMEs | ✓ shipped |
| 2 | Mirror simulator profiling scripts + run-time tracking | ✓ shipped |
| 3 | Nautilus profiling, design for sampler expansion | ✓ shipped |
| 4 | Top-level + per-section README dashboard with instrument framing | ✓ shipped |
| 5 | GitHub Actions for lint + profile re-runs + README refresh | ✓ shipped (lint.yml per-PR; profile.yml manual/on-release) |
| 6 | Design lock-in + results/dashboard groundwork (#52) | in progress |
| 7 | PreOptimizationTimes baseline campaign (vram-first, then runtime + breakdown) | queued |
Dashboards can grow in many directions. The list below captures candidate improvements that fit the "profiling and run-times" theme; none of them block the current dashboard from being useful.
- Regression-watch indicator — colour or arrow per cell showing whether the latest cost regressed (>5%) or improved versus the previous PyAutoLens release. Needs the second-latest version per axis kept alongside the latest. Trivial to add to
scripts/build_readme.py. - Per-axis version-history PNGs — small inline plot of run-time vs PyAutoLens release version, generated from the JSON artifacts (reusing the
_developer/jax_profiling/results/jit/.../*_v<version>.pnggenerator). Embeds nicely above each section table. - Plotly-rendered interactive timeline — hostable on GitHub Pages once the static dashboard stabilises; lets readers hover/filter across instrument × model × release.
- Flamegraph captures — alongside the headline timing numbers, store a flamegraph per instrument × model for the most recent release.
- Hardware-tier columns — extend
scripts/build_readme.pytable renderers to show CPU / laptop GPU / HPC GPU as separate columns once result artifacts encode the hardware label (filename suffix or JSON"hardware"field). - Archive old versions — once a script has >6 minor releases of artifacts, move the older ones to
results/archive/so the latest views stay uncluttered.
PyAutoLabs/PyAutoLens— the library being profiled.PyAutoLabs/autolens_workspace— user-facing science scripts and tutorials.PyAutoLabs/autolens_workspace_developer— the developer workspace this repo's scripts were migrated from; still hosts the pre-2026-07 sweep history and the gradient-profiling work.Jammy2211/autolens_colab_profiling— sibling repo, Colab-specific scope. Not yet migrated to PyAutoLabs.
This repo is a collection of standalone profiling scripts, not an installable Python package. There is no pyproject.toml. Run scripts from the repo root.
Scripts follow the JIT conventions documented in autolens_workspace_developer/CLAUDE.md:
- Extract
.arrayfrom autoarray types before crossing thejax.jitboundary (autoarray types are not JAX pytrees as inputs). - Pass
xp=jnpthrough PyAutoLens / PyAutoGalaxy / PyAutoArray functions to select the JAX backend.
- Slack — PyAutoLens workspace for questions.
- Issues — file profiling bugs and feature requests on this repo's issue tracker.