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.
Cell-level full-pipeline numbers live in likelihood_runtime/OPTIMIZATION_NOTES.md, which carries the latest CPU + local-GPU per-call costs together with per-cell "where to optimize next" recommendations and the mp-vs-fp64 verdicts. The detailed multi-config comparison.json artifacts are committed under autolens_workspace_developer/jax_profiling/results/jit/<class>/<model>/ and are re-aggregated by likelihood_runtime/aggregate.py whenever a new sweep finishes.
(The previous auto-generated table in this README was retired when the likelihood profiling was split into likelihood_breakdown/ + likelihood_runtime/; scripts/build_readme.py is queued for a path-update follow-up.)
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.
Each profiling script writes a versioned artifact pair under results/:
results/<section>/<subfolder>/<profile_name>_summary_v<YYYY>.<M>.<D>.<PATCH>.json
results/<section>/<subfolder>/<profile_name>_summary_v<YYYY>.<M>.<D>.<PATCH>.png
The version string matches the PyAutoLens release that produced the numbers (e.g. v2026.5.1.4). Older versions are retained alongside newer ones, so trends across releases stay inspectable. The JSON carries structured timings; the PNG is the at-a-glance plot.
Examples that already exist in the source-of-truth repo:
imaging_summary_v2026.5.1.4.json/.pngpoint_source_summary_v2026.5.1.4.jsondelaunay_sparse_cpu_likelihood_summary_hst_v2026.5.1.4.json
| Folder | Contents |
|---|---|
likelihood_breakdown/ |
Per-step JIT decomposition. Single config. Where does time go inside the likelihood? |
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? |
simulators/ |
Run-time tracking for the PyAutoLens simulators. |
searches/ |
Sampler / search profiling, Nautilus first. |
results/ |
Versioned JSON + PNG artifacts written by the above scripts. |
This repo is being built in phases. Phase numbers correspond to internal sub-prompts under PyAutoLabs/PyAutoPrompt/z_features/autolens_profiling.md.
| 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 | 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; source of truth during the migration. Each phase mirrors the relevant subdirectories from here into this repo.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.