diff --git a/AGENTS.md b/AGENTS.md index bee232c..743afbd 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -10,18 +10,19 @@ overview (vision, latest run-times, roadmap); this file is the operational guide ## Repository Structure ``` -likelihood/ Per-instrument likelihood profile scripts (imaging, interferometer, - datacube, point_source) — import the shared _profile_cli helper likelihood_runtime/ Full-pipeline JIT runtime, driven by sweep.py across CPU/GPU/A100 × fp64/mp + (per-cell scripts under /.py import the shared _profile_cli helper) likelihood_breakdown/ Per-step JIT decomposition of a single likelihood config -vram/ GPU memory-usage profiling +vram/ GPU memory-usage profiling + the per-cell A100 vmap batch-size table instruments/ Instrument definitions (pixel scale, shape) used to frame results searches/ Sampler / search profiling (Nautilus first) simulators/ Run-time tracking for the PyAutoLens simulators latent/ Latent-variable profiling -quick_update/ Fast incremental re-profiling helpers -hpc/ SLURM submit scripts for the RAL HPC -results/ Versioned JSON + PNG artifacts (`*_v....{json,png}`) +quick_update/ Fast incremental re-profiling helpers (unversioned scratch tier) +hpc/ SLURM submit scripts for the RAL HPC (A100 rows of the sweep matrix) +scripts/ Repo tooling (build_readme.py dashboard generator) +results/ JSON + PNG artifacts: versioned summaries, sweep comparisons, named + baselines (results/README.md defines the shapes; sweeps default here) config/ dataset/ output/ Config, input data, runtime output ``` diff --git a/README.md b/README.md index be3cc05..7997cb7 100644 --- a/README.md +++ b/README.md @@ -24,9 +24,28 @@ Results are framed by **astronomy instrument** (HST, Euclid, JWST, …) rather t ## Latest run-times -Cell-level full-pipeline numbers live in [`likelihood_runtime/OPTIMIZATION_NOTES.md`](./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///`](https://github.com/PyAutoLabs/autolens_workspace_developer/tree/main/jax_profiling/results/jit) 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.) +**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/`](./results/README.md) — 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`](./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`](./results/notes/design_lock_in.md). + +(Historical multi-config sweeps up to 2026-07 were committed under [`autolens_workspace_developer/jax_profiling/results/jit/`](https://github.com/PyAutoLabs/autolens_workspace_developer/tree/main/jax_profiling/results/jit); sweeps now write in-repo to `results/runtime/` by default.) ## JAX gradients — currently out of scope @@ -34,34 +53,36 @@ Gradient profiling (`jax.grad` of the likelihood, autodiff-based optimisers) is ## How to read this repo -Each profiling script writes a **versioned artifact pair** under `results/`: +Profiling scripts write two artifact shapes under `results/` (full reference: [`results/README.md`](./results/README.md)): ``` -results/
//_summary_v....json -results/
//_summary_v....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. +# Versioned summaries — standalone runs; history retained side-by-side +results/
//___v...[_sparse].{json,png} -Examples that already exist in the source-of-truth repo: +# Per-config sweeps — sweep.py + aggregate.py; latest sweep per cell +results/runtime//[/]/[_sparse].{json,png,log} + comparison.{json,png} +``` -- `imaging_summary_v2026.5.1.4.json` / `.png` -- `point_source_summary_v2026.5.1.4.json` -- `delaunay_sparse_cpu_likelihood_summary_hst_v2026.5.1.4.json` +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`. ## Section index | Folder | Contents | |--------|----------| -| [`likelihood_breakdown/`](./likelihood_breakdown/README.md) | Per-step JIT decomposition. Single config. *Where does time go inside the likelihood?* | | [`likelihood_runtime/`](./likelihood_runtime/README.md) | 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/`](./likelihood_breakdown/README.md) | Per-step JIT decomposition. Single config. *Where does time go inside the likelihood?* | +| [`vram/`](./vram/README.md) | GPU memory profiling + the per-cell vmap batch-size table for the A100. | +| [`latent/`](./latent/README.md) | Latent-variable profiling. | +| [`instruments/`](./instruments/README.md) | Instrument presets (pixel scale, shape) that frame every result. | | [`simulators/`](./simulators/README.md) | Run-time tracking for the PyAutoLens simulators. | | [`searches/`](./searches/README.md) | Sampler / search profiling, Nautilus first. | -| [`results/`](./results/README.md) | Versioned JSON + PNG artifacts written by the above scripts. | +| [`quick_update/`](./quick_update/README.md) | Fast incremental re-profiling helpers. | +| [`hpc/`](./hpc/README.md) | SLURM submit scripts for the RAL HPC (A100 rows of the sweep matrix). | +| [`results/`](./results/README.md) | JSON + PNG artifacts written by the above scripts; named baselines. | ## Roadmap -This repo is being built in phases. Phase numbers correspond to internal sub-prompts under `PyAutoLabs/PyAutoPrompt/z_features/autolens_profiling.md`. +This repo is being built in phases (bootstrap history now archived in `PyAutoMind`). | Phase | Title | Status | |-------|-------|--------| @@ -70,7 +91,9 @@ This repo is being built in phases. Phase numbers correspond to internal sub-pro | 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 | +| 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](https://github.com/PyAutoLabs/autolens_profiling/issues/52)) | in progress | +| 7 | **PreOptimizationTimes** baseline campaign (vram-first, then runtime + breakdown) | queued | ### Future enhancements (Phase 4 follow-ups) @@ -87,7 +110,7 @@ Dashboards can grow in many directions. The list below captures candidate improv - [`PyAutoLabs/PyAutoLens`](https://github.com/PyAutoLabs/PyAutoLens) — the library being profiled. - [`PyAutoLabs/autolens_workspace`](https://github.com/PyAutoLabs/autolens_workspace) — user-facing science scripts and tutorials. -- [`PyAutoLabs/autolens_workspace_developer`](https://github.com/PyAutoLabs/autolens_workspace_developer) — the developer workspace; **source of truth during the migration**. Each phase mirrors the relevant subdirectories from here into this repo. +- [`PyAutoLabs/autolens_workspace_developer`](https://github.com/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`](https://github.com/Jammy2211/autolens_colab_profiling) — sibling repo, Colab-specific scope. Not yet migrated to PyAutoLabs. ## Package vs scripts diff --git a/_adapt_image_util.py b/_adapt_image_util.py index 00a0258..5607c55 100644 --- a/_adapt_image_util.py +++ b/_adapt_image_util.py @@ -25,9 +25,9 @@ def adapt_image_for_dataset( *, - dataset_path: Union[str, Path], + dataset_path: str | Path, dataset, - tracer: Optional[al.Tracer] = None, + tracer: al.Tracer | None = None, ) -> aa.Array2D: """Return the lensed-source image masked to ``dataset.mask``. diff --git a/_profile_cli.py b/_profile_cli.py index 4216b05..e23e0bb 100644 --- a/_profile_cli.py +++ b/_profile_cli.py @@ -1,7 +1,7 @@ """Shared CLI / JSON / auto-simulate helpers for the likelihood scripts. -Used by every script under ``likelihood/{imaging,interferometer, -datacube,point_source}/`` so the per-script boilerplate stays minimal +Used by every per-cell script under ``likelihood_runtime/`` and +``likelihood_breakdown/`` so the per-script boilerplate stays minimal and the sweep-driver flags (``--config-name``, ``--output-dir``, ``--use-mixed-precision``) and dataset auto-simulate hook are defined in one place. @@ -30,15 +30,15 @@ @dataclass(frozen=True) class ProfileCLI: - config_name: Optional[str] - output_dir: Optional[Path] + config_name: str | None + output_dir: Path | None use_mixed_precision: bool - instrument: Optional[str] + instrument: str | None vmap_probe: bool use_sparse_operator: bool -def parse_profile_cli(default_config_name: Optional[str] = None) -> ProfileCLI: +def parse_profile_cli(default_config_name: str | None = None) -> ProfileCLI: """Parse the sweep CLI flags accepted by every per-cell profile script. Returns ``ProfileCLI(config_name, output_dir, use_mixed_precision, @@ -69,8 +69,8 @@ def parse_profile_cli(default_config_name: Optional[str] = None) -> ProfileCLI: "--output-dir", default=None, help=( - "Override results dir. Defaults to " - "/results/likelihood//." + "Override results dir. Each per-cell script defaults to its " + "package's section under /results/." ), ) parser.add_argument( @@ -149,15 +149,19 @@ def device_info_dict() -> dict: } if info["backend"] == "gpu": try: - out = subprocess.check_output( - [ - "nvidia-smi", - "--query-gpu=name,memory.used,memory.total", - "--format=csv,noheader", - ], - stderr=subprocess.DEVNULL, - timeout=3, - ).decode().strip() + out = ( + subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=name,memory.used,memory.total", + "--format=csv,noheader", + ], + stderr=subprocess.DEVNULL, + timeout=3, + ) + .decode() + .strip() + ) info["nvidia_smi"] = out.replace("\n", "; ") except Exception: pass @@ -246,8 +250,10 @@ def auto_simulate_if_missing( [ sys.executable, str(simulator_script), - "--instrument", instrument, - "--output-root", str(workspace_root), + "--instrument", + instrument, + "--output-root", + str(workspace_root), ], check=True, ) diff --git a/hpc/README.md b/hpc/README.md new file mode 100644 index 0000000..1ac1b3e --- /dev/null +++ b/hpc/README.md @@ -0,0 +1,33 @@ +# hpc + +SLURM submit scripts for the RAL HPC — the `hpc_a100_fp64` / `hpc_a100_mp` +rows of the sweep matrix (and HPC-CPU rows) that the local +`likelihood_runtime/sweep.py` driver cannot run itself. + +## Layout + +``` +batch_gpu/ + submit____a100__[_sparse] # one submit per cell/config + output/ error/ # SLURM stdout/stderr (gitignored) +``` + +Submit names follow the same `/` cell grid as the rest of the +repo; `runtime_` prefixed submits drive `likelihood_runtime/` cells, `nss_` +prefixed ones drive `searches/` cells. + +## Running + +On the HPC login node: + +```bash +source activate.sh # repo-root helper: venv + PYTHONPATH at the canonical checkouts +sbatch hpc/batch_gpu/submit_runtime_imaging_mge_a100_hst_fp64 +``` + +Each job writes its per-config JSON into the same `results/` layout as a local +sweep (`--config-name hpc_a100_fp64` etc.), so `likelihood_runtime/aggregate.py` +merges local and A100 rows into one `comparison.json`. Copy/commit the result +JSONs from the HPC checkout back via the normal git flow. The PyAuto* +libraries resolve from sibling source checkouts on `PYTHONPATH` — never +pip-install them into the venv (`HPCPullPyAuto` is the update story). diff --git a/instruments/imaging.py b/instruments/imaging.py index ab5dd71..02653d8 100644 --- a/instruments/imaging.py +++ b/instruments/imaging.py @@ -26,7 +26,6 @@ from __future__ import annotations - INSTRUMENTS: dict[str, dict] = { "euclid": { "pixel_scale": 0.1, diff --git a/instruments/interferometer.py b/instruments/interferometer.py index eda02cc..3bf8e61 100644 --- a/instruments/interferometer.py +++ b/instruments/interferometer.py @@ -31,7 +31,6 @@ from typing import Optional - INSTRUMENTS: dict[str, dict] = { "sma": { "pixel_scale": 0.1, @@ -92,6 +91,6 @@ def mask_radius_pixels(instrument: str) -> int: return int(round(cfg["mask_radius"] / cfg["pixel_scale"])) -def transformer_chunk_size_for(instrument: str) -> Optional[int]: +def transformer_chunk_size_for(instrument: str) -> int | None: """Per-instrument NUFFT chunk_size (None for one-shot).""" return INSTRUMENTS[instrument].get("transformer_chunk_size") diff --git a/latent/aggregate.py b/latent/aggregate.py index 425ae3c..83f883b 100644 --- a/latent/aggregate.py +++ b/latent/aggregate.py @@ -32,10 +32,8 @@ import matplotlib.pyplot as plt import numpy as np - _REPO_ROOT = Path(__file__).resolve().parents[1] -_WT_ROOT = _REPO_ROOT.parent -_DEFAULT_OUTPUT_ROOT = _WT_ROOT / "autolens_workspace_developer" / "jax_profiling" / "results" / "latent" +_DEFAULT_OUTPUT_ROOT = _REPO_ROOT / "results" / "latent" # Canonical ordering of cells — drives auto-discovery sort order. _CELLS: list[tuple[str, str]] = [ @@ -134,8 +132,15 @@ def _render_table(comparison: dict, cell_id: str) -> str: lines = [f"=== {cell_id} ==="] is_einstein = "effective_einstein_radius" in cell_id if is_einstein: - header = ("config", "eager_time", "jit_first", "jit_steady", "vmap_per_call", - "cache_first", "cache_second") + header = ( + "config", + "eager_time", + "jit_first", + "jit_steady", + "vmap_per_call", + "cache_first", + "cache_second", + ) else: header = ("config", "eager_time", "jit_first", "jit_steady", "vmap_per_call") rows = [header] @@ -147,23 +152,27 @@ def _render_table(comparison: dict, cell_id: str) -> str: if is_einstein: cc_first = cfg.get("closure_cache_first_call_s") cc_second = cfg.get("closure_cache_second_call_s") - rows.append(( - name, - _format_seconds(eager_t), - _format_seconds(jit_first), - _format_seconds(jit_steady), - _format_seconds(vmap), - _format_seconds(cc_first), - _format_seconds(cc_second), - )) + rows.append( + ( + name, + _format_seconds(eager_t), + _format_seconds(jit_first), + _format_seconds(jit_steady), + _format_seconds(vmap), + _format_seconds(cc_first), + _format_seconds(cc_second), + ) + ) else: - rows.append(( - name, - _format_seconds(eager_t), - _format_seconds(jit_first), - _format_seconds(jit_steady), - _format_seconds(vmap), - )) + rows.append( + ( + name, + _format_seconds(eager_t), + _format_seconds(jit_first), + _format_seconds(jit_steady), + _format_seconds(vmap), + ) + ) col_w = [max(len(r[i]) for r in rows) for i in range(len(rows[0]))] for r in rows: lines.append(" " + " ".join(s.ljust(w) for s, w in zip(r, col_w))) @@ -199,18 +208,11 @@ def _render_png(comparison: dict, cell_id: str, png_path: Path) -> None: series["jit_steady_state"].append(cfg.get("jit_steady_state_s", np.nan)) series["vmap_per_call"].append(cfg.get("vmap_per_call_s", np.nan)) if is_einstein: - series["closure_cache_first"].append( - cfg.get("closure_cache_first_call_s", np.nan) - ) - series["closure_cache_second"].append( - cfg.get("closure_cache_second_call_s", np.nan) - ) + series["closure_cache_first"].append(cfg.get("closure_cache_first_call_s", np.nan)) + series["closure_cache_second"].append(cfg.get("closure_cache_second_call_s", np.nan)) # Drop series that are entirely nan/zero — nothing to plot. - series = { - k: v for k, v in series.items() - if any(np.isfinite(x) and x > 0 for x in v) - } + series = {k: v for k, v in series.items() if any(np.isfinite(x) and x > 0 for x in v)} if not series: return @@ -262,9 +264,7 @@ def main() -> int: for spec in args.cell: parts = spec.split("/") if len(parts) != 2: - sys.stderr.write( - f"bad --cell argument: {spec!r} (expected class/latent)\n" - ) + sys.stderr.write(f"bad --cell argument: {spec!r} (expected class/latent)\n") return 2 cells.append((parts[0], parts[1])) else: @@ -274,7 +274,7 @@ def main() -> int: sys.stderr.write(f"no cells found under {args.output_root}\n") return 1 - for (cls, latent) in cells: + for cls, latent in cells: cell_id = f"{cls}/{latent}" cell_dir = args.output_root / cls / latent if not cell_dir.exists(): diff --git a/latent/imaging/effective_einstein_radius.py b/latent/imaging/effective_einstein_radius.py index c0d28f6..14af503 100644 --- a/latent/imaging/effective_einstein_radius.py +++ b/latent/imaging/effective_einstein_radius.py @@ -23,23 +23,23 @@ import argparse import json import os +import os as _os + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit. +import sys as _sys +import tempfile import time from pathlib import Path +import autofit as af +import autolens as al import jax import jax.numpy as jnp import numpy as np - -import tempfile - -import autofit as af -import autolens as al -from autolens import fixtures from autoconf import conf +from autolens import fixtures from autolens.analysis.latent import LATENT_FUNCTIONS -# AUTOLENS_PROFILING_SMOKE=1 short-circuit. -import sys as _sys, os as _os if _os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports OK; exiting.") _sys.exit(0) @@ -50,10 +50,7 @@ def _push_single_latent_config(latent_key: str) -> Path: """Write a temp config dir with only latent_key enabled and push it.""" tmpdir = Path(tempfile.mkdtemp(prefix="latent_cfg_")) - yaml_lines = [ - f"{k}: {'true' if k == latent_key else 'false'}" - for k in LATENT_FUNCTIONS - ] + yaml_lines = [f"{k}: {'true' if k == latent_key else 'false'}" for k in LATENT_FUNCTIONS] (tmpdir / "latent.yaml").write_text("\n".join(yaml_lines) + "\n") conf.instance.push(str(tmpdir)) return tmpdir @@ -75,8 +72,6 @@ def _time_closure_cache(tracer, dataset, xp=np): try: # Config already pushed by main(); the pushed config has this latent enabled. - analysis = al.AnalysisImaging(dataset=dataset, use_jax=(xp is not np), magzero=25.0) - lens = af.Model(al.Galaxy, redshift=0.5, mass=al.mp.Isothermal, bulge=al.lp.Sersic) source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.Sersic) model = af.Collection(galaxies=af.Collection(lens=lens, source=source)) @@ -97,6 +92,7 @@ def _time_closure_cache(tracer, dataset, xp=np): # Second call — LensCalc on `fit.tracer` still has its cache populated # from above. Construct a new LensCalc explicitly to test warm-cache path. from autogalaxy.operate.lens_calc import LensCalc + lc2 = LensCalc.from_mass_obj(fit.tracer) init_guess = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]]) # Warm the ZeroSolver cache on lc2 first (mirrors the cold first call). diff --git a/latent/imaging/magnification.py b/latent/imaging/magnification.py index 9db9c6d..09999a7 100644 --- a/latent/imaging/magnification.py +++ b/latent/imaging/magnification.py @@ -17,23 +17,23 @@ import argparse import json import os +import os as _os + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit. +import sys as _sys +import tempfile import time from pathlib import Path +import autofit as af +import autolens as al import jax import jax.numpy as jnp import numpy as np - -import tempfile - -import autofit as af -import autolens as al -from autolens import fixtures from autoconf import conf +from autolens import fixtures from autolens.analysis.latent import LATENT_FUNCTIONS -# AUTOLENS_PROFILING_SMOKE=1 short-circuit. -import sys as _sys, os as _os if _os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports OK; exiting.") _sys.exit(0) @@ -44,10 +44,7 @@ def _push_single_latent_config(latent_key: str) -> Path: """Write a temp config dir with only latent_key enabled and push it.""" tmpdir = Path(tempfile.mkdtemp(prefix="latent_cfg_")) - yaml_lines = [ - f"{k}: {'true' if k == latent_key else 'false'}" - for k in LATENT_FUNCTIONS - ] + yaml_lines = [f"{k}: {'true' if k == latent_key else 'false'}" for k in LATENT_FUNCTIONS] (tmpdir / "latent.yaml").write_text("\n".join(yaml_lines) + "\n") conf.instance.push(str(tmpdir)) return tmpdir diff --git a/latent/imaging/total_lens_flux_mujy.py b/latent/imaging/total_lens_flux_mujy.py index 5e42338..e31306d 100644 --- a/latent/imaging/total_lens_flux_mujy.py +++ b/latent/imaging/total_lens_flux_mujy.py @@ -16,23 +16,23 @@ import argparse import json import os +import os as _os + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit. +import sys as _sys +import tempfile import time from pathlib import Path +import autofit as af +import autolens as al import jax import jax.numpy as jnp import numpy as np - -import tempfile - -import autofit as af -import autolens as al -from autolens import fixtures from autoconf import conf +from autolens import fixtures from autolens.analysis.latent import LATENT_FUNCTIONS -# AUTOLENS_PROFILING_SMOKE=1 short-circuit. -import sys as _sys, os as _os if _os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports OK; exiting.") _sys.exit(0) @@ -43,10 +43,7 @@ def _push_single_latent_config(latent_key: str) -> Path: """Write a temp config dir with only latent_key enabled and push it.""" tmpdir = Path(tempfile.mkdtemp(prefix="latent_cfg_")) - yaml_lines = [ - f"{k}: {'true' if k == latent_key else 'false'}" - for k in LATENT_FUNCTIONS - ] + yaml_lines = [f"{k}: {'true' if k == latent_key else 'false'}" for k in LATENT_FUNCTIONS] (tmpdir / "latent.yaml").write_text("\n".join(yaml_lines) + "\n") conf.instance.push(str(tmpdir)) return tmpdir diff --git a/latent/imaging/total_lensed_source_flux_mujy.py b/latent/imaging/total_lensed_source_flux_mujy.py index a40eae8..c657f31 100644 --- a/latent/imaging/total_lensed_source_flux_mujy.py +++ b/latent/imaging/total_lensed_source_flux_mujy.py @@ -17,23 +17,23 @@ import argparse import json import os +import os as _os + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit. +import sys as _sys +import tempfile import time from pathlib import Path +import autofit as af +import autolens as al import jax import jax.numpy as jnp import numpy as np - -import tempfile - -import autofit as af -import autolens as al -from autolens import fixtures from autoconf import conf +from autolens import fixtures from autolens.analysis.latent import LATENT_FUNCTIONS -# AUTOLENS_PROFILING_SMOKE=1 short-circuit. -import sys as _sys, os as _os if _os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports OK; exiting.") _sys.exit(0) @@ -44,10 +44,7 @@ def _push_single_latent_config(latent_key: str) -> Path: """Write a temp config dir with only latent_key enabled and push it.""" tmpdir = Path(tempfile.mkdtemp(prefix="latent_cfg_")) - yaml_lines = [ - f"{k}: {'true' if k == latent_key else 'false'}" - for k in LATENT_FUNCTIONS - ] + yaml_lines = [f"{k}: {'true' if k == latent_key else 'false'}" for k in LATENT_FUNCTIONS] (tmpdir / "latent.yaml").write_text("\n".join(yaml_lines) + "\n") conf.instance.push(str(tmpdir)) return tmpdir diff --git a/latent/imaging/total_source_flux_mujy.py b/latent/imaging/total_source_flux_mujy.py index bef7bfd..7ef4656 100644 --- a/latent/imaging/total_source_flux_mujy.py +++ b/latent/imaging/total_source_flux_mujy.py @@ -19,23 +19,23 @@ import argparse import json import os +import os as _os + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit. +import sys as _sys +import tempfile import time from pathlib import Path +import autofit as af +import autolens as al import jax import jax.numpy as jnp import numpy as np - -import tempfile - -import autofit as af -import autolens as al -from autolens import fixtures from autoconf import conf +from autolens import fixtures from autolens.analysis.latent import LATENT_FUNCTIONS -# AUTOLENS_PROFILING_SMOKE=1 short-circuit. -import sys as _sys, os as _os if _os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports OK; exiting.") _sys.exit(0) @@ -46,10 +46,7 @@ def _push_single_latent_config(latent_key: str) -> Path: """Write a temp config dir with only latent_key enabled and push it.""" tmpdir = Path(tempfile.mkdtemp(prefix="latent_cfg_")) - yaml_lines = [ - f"{k}: {'true' if k == latent_key else 'false'}" - for k in LATENT_FUNCTIONS - ] + yaml_lines = [f"{k}: {'true' if k == latent_key else 'false'}" for k in LATENT_FUNCTIONS] (tmpdir / "latent.yaml").write_text("\n".join(yaml_lines) + "\n") conf.instance.push(str(tmpdir)) return tmpdir diff --git a/latent/sweep.py b/latent/sweep.py index 3865aab..88ac142 100644 --- a/latent/sweep.py +++ b/latent/sweep.py @@ -13,7 +13,7 @@ ///.log (captured stdout/stderr) Default ``--output-root`` is -``autolens_workspace_developer/jax_profiling/results/latent`` — mirrors the +``results/latent/`` in this repo — mirrors the ``jit/`` convention used by ``likelihood_runtime/sweep.py`` and is read by ``aggregate.py`` to produce ``comparison.json`` / ``comparison.png``. @@ -39,10 +39,8 @@ from dataclasses import dataclass from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[1] # autolens_profiling/ -_WT_ROOT = _REPO_ROOT.parent # PyAutoLabs-wt// (or PyAutoLabs/) -_DEFAULT_OUTPUT_ROOT = _WT_ROOT / "autolens_workspace_developer" / "jax_profiling" / "results" / "latent" +_REPO_ROOT = Path(__file__).resolve().parents[1] # autolens_profiling/ +_DEFAULT_OUTPUT_ROOT = _REPO_ROOT / "results" / "latent" _DEFAULT_PYTHON = "/home/jammy/venv/PyAutoGPU/bin/python" @@ -176,8 +174,10 @@ def _run_one( cmd = [ python, str(script_path), - "--config-name", config.name, - "--output-dir", str(out_dir), + "--config-name", + config.name, + "--output-dir", + str(out_dir), *config.extra_args, ] @@ -206,7 +206,9 @@ def _run_one( ) elapsed = time.time() - t0 ok = proc.returncode == 0 - print(f" {'OK ' if ok else 'FAIL'} ({elapsed:.1f}s, exit={proc.returncode}) -> {log_path.name}") + print( + f" {'OK ' if ok else 'FAIL'} ({elapsed:.1f}s, exit={proc.returncode}) -> {log_path.name}" + ) return ok, elapsed, str(log_path) except KeyboardInterrupt: elapsed = time.time() - t0 @@ -219,8 +221,10 @@ def main() -> int: cells = _resolve_cells(args) configs = _resolve_configs(args) - print(f"sweep_latent: {len(cells)} cells x {len(configs)} configs " - f"= {len(cells) * len(configs)} runs") + print( + f"sweep_latent: {len(cells)} cells x {len(configs)} configs " + f"= {len(cells) * len(configs)} runs" + ) print(f" cells: {[f'{c}/{m}' for (c, m) in cells]}") print(f" configs: {[c.name for c in configs]}") print(f" output: {args.output_root}") @@ -231,7 +235,7 @@ def main() -> int: summary: list[tuple[str, str, bool, float]] = [] overall_t0 = time.time() - for (cls, latent) in cells: + for cls, latent in cells: script_path = _REPO_ROOT / "latent" / cls / f"{latent}.py" if not script_path.exists(): print(f"\n!!! missing script: {script_path}") @@ -243,9 +247,7 @@ def main() -> int: for cfg in configs: try: - ok, elapsed, _log = _run_one( - args.python, script_path, cfg, out_dir, args.dry_run - ) + ok, elapsed, _log = _run_one(args.python, script_path, cfg, out_dir, args.dry_run) except KeyboardInterrupt: print("\n\nsweep interrupted by user") return 130 @@ -256,7 +258,7 @@ def main() -> int: print(f"sweep_latent summary ({total:.1f}s total)") print("=" * 70) print(f" {'cell':<40}{'config':<22}{'ok':<6}{'elapsed':>10}") - print(f" {'-'*40}{'-'*22}{'-'*6}{'-'*10}") + print(f" {'-' * 40}{'-' * 22}{'-' * 6}{'-' * 10}") failures = 0 for cell, cfg, ok, t in summary: flag = "OK" if ok else "FAIL" diff --git a/likelihood_breakdown/README.md b/likelihood_breakdown/README.md index 13a15cf..39e0922 100644 --- a/likelihood_breakdown/README.md +++ b/likelihood_breakdown/README.md @@ -8,6 +8,20 @@ Run one of these scripts when you want to find the **next source-code-level opti For *how long the likelihood actually takes* on production hardware — i.e. a single end-to-end number per (hardware, precision) config — use the sibling package [`likelihood_runtime/`](../likelihood_runtime/) instead. The two packages are deliberately disjoint so neither has to pay the other's cost. +## Latest results + + +| 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 | + + +Auto-generated by `scripts/build_readme.py` from the versioned artifacts under `results/breakdown/`. + ## Methodology For each pipeline step (e.g. *ray-trace grids* → *blurred mapping matrix* → *curvature matrix F* → *NNLS reconstruction* → *log-evidence*), the script: diff --git a/likelihood_breakdown/datacube/delaunay.py b/likelihood_breakdown/datacube/delaunay.py index 6443e9c..5617365 100644 --- a/likelihood_breakdown/datacube/delaunay.py +++ b/likelihood_breakdown/datacube/delaunay.py @@ -76,34 +76,33 @@ same schema as the interferometer sibling. """ -import numpy as np -import jax -import jax.numpy as jnp import os -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -111,12 +110,13 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 + _cli = parse_profile_cli() instrument = _cli.instrument or "sma" @@ -132,6 +132,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -253,18 +254,14 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset_list[0] - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset_list[0]) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") print("\n--- Image mesh construction (Hilbert) ---") with timer.section("image_mesh_hilbert"): - image_mesh = al.image_mesh.Hilbert( - pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0 - ) + image_mesh = al.image_mesh.Hilbert(pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0) image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( mask=dataset_list[0].real_space_mask, adapt_data=adapt_image ) @@ -389,9 +386,9 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n" + "=" * 70) print("PER-STEP JIT PROFILING — CUBE") print("=" * 70) -print(f" Channel-invariant steps are timed once.") -print(f" Channel-variant steps are JIT-compiled on channel 0; the reported") -print(f" cube cost is N × the per-channel steady-state per-call.") +print(" Channel-invariant steps are timed once.") +print(" Channel-variant steps are JIT-compiled on channel 0; the reported") +print(" cube cost is N × the per-channel steady-state per-call.") # Reference single-channel context (channel 0) fit = fit_list[0] @@ -409,9 +406,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Step 1: Ray-trace data grid (shared) ---") with timer.section("ray_trace_data_eager"): - traced_grids = tracer.traced_grid_2d_list_from( - grid=dataset.grids.pixelization, xp=jnp - ) + traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp) for tg in traced_grids: block(tg) @@ -424,9 +419,7 @@ def ray_trace_data_raw(grid_raw): _, _ = jit_profile(ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw) ray_trace_data_per_call = timer.records[-1][1] / 10 -likelihood_steps.append( - ("Ray-trace data grid (shared)", ray_trace_data_per_call) -) +likelihood_steps.append(("Ray-trace data grid (shared)", ray_trace_data_per_call)) # --------------------------------------------------------------------------- # Step 2: Ray-trace mesh grid (channel-invariant) @@ -443,9 +436,7 @@ def ray_trace_mesh_raw(mesh_raw): _, _ = jit_profile(ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw) ray_trace_mesh_per_call = timer.records[-1][1] / 10 -likelihood_steps.append( - ("Ray-trace mesh grid (shared)", ray_trace_mesh_per_call) -) +likelihood_steps.append(("Ray-trace mesh grid (shared)", ray_trace_mesh_per_call)) # --------------------------------------------------------------------------- # Always use the sparse-operator breakdown path @@ -552,26 +543,39 @@ def transformed_mm_from_params(params_tree): print("\n--- Step 4: Data vector D (per channel) ---") def compute_data_vector( - transformed_mm_real, transformed_mm_imag, data_real, data_imag, - noise_real, noise_imag, + transformed_mm_real, + transformed_mm_imag, + data_real, + data_imag, + noise_real, + noise_imag, ): - weighted_data_real = data_real / (noise_real ** 2) - weighted_data_imag = data_imag / (noise_imag ** 2) + weighted_data_real = data_real / (noise_real**2) + weighted_data_imag = data_imag / (noise_imag**2) return jnp.matmul(transformed_mm_real.T, weighted_data_real) + jnp.matmul( transformed_mm_imag.T, weighted_data_imag ) with timer.section("data_vector_eager"): data_vector = compute_data_vector( - transformed_mm_real_jnp, transformed_mm_imag_jnp, - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) block(data_vector) _, data_vector = jit_profile( - compute_data_vector, "data_vector_jit", - transformed_mm_real_jnp, transformed_mm_imag_jnp, - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, + compute_data_vector, + "data_vector_jit", + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) data_vector_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( @@ -590,7 +594,10 @@ def compute_data_vector( no_reg_list = list(inversion.no_regularization_index_list) def compute_curvature_matrix( - transformed_mm_real, transformed_mm_imag, noise_real, noise_imag, + transformed_mm_real, + transformed_mm_imag, + noise_real, + noise_imag, ): real_curv = al.util.inversion.curvature_matrix_via_mapping_matrix_from( mapping_matrix=transformed_mm_real, @@ -612,13 +619,20 @@ def compute_curvature_matrix( with timer.section("curvature_matrix_eager"): curvature_matrix = compute_curvature_matrix( - transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) block(curvature_matrix) _, curvature_matrix = jit_profile( - compute_curvature_matrix, "curvature_matrix_jit", - transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp, + compute_curvature_matrix, + "curvature_matrix_jit", + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) curvature_matrix_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( @@ -638,9 +652,7 @@ def compute_curvature_matrix( regularization_matrix = jnp.array(inversion.regularization_matrix) block(regularization_matrix) - likelihood_steps.append( - ("Regularization matrix H (shared)", timer.records[-1][1]) - ) + likelihood_steps.append(("Regularization matrix H (shared)", timer.records[-1][1])) # --------------------------------------------------------------- # Step 7: Reconstruction NNLS (per channel) @@ -665,8 +677,11 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) block(reconstruction) _, reconstruction = jit_profile( - compute_reconstruction, "reconstruction_jit", - jnp.array(data_vector), jnp.array(curvature_matrix), jnp.array(regularization_matrix), + compute_reconstruction, + "reconstruction_jit", + jnp.array(data_vector), + jnp.array(curvature_matrix), + jnp.array(regularization_matrix), ) reconstruction_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( @@ -683,9 +698,16 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) print("\n--- Step 8: Mapped recon + log evidence (per channel) ---") def compute_log_evidence( - data_real, data_imag, noise_real, noise_imag, - transformed_mm_real, transformed_mm_imag, - reconstruction, curvature_matrix, regularization_matrix, mapper_indices, + data_real, + data_imag, + noise_real, + noise_imag, + transformed_mm_real, + transformed_mm_imag, + reconstruction, + curvature_matrix, + regularization_matrix, + mapper_indices, ): mapped_real = jnp.matmul(transformed_mm_real, reconstruction) mapped_imag = jnp.matmul(transformed_mm_imag, reconstruction) @@ -701,21 +723,19 @@ def compute_log_evidence( curvature_reg_matrix = curvature_matrix + regularization_matrix creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices] reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices] - log_det_curvature_reg = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced))) - ) - log_det_regularization = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced))) - ) + log_det_curvature_reg = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))) + log_det_regularization = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))) - noise_normalization = ( - jnp.sum(jnp.log(2 * jnp.pi * noise_real ** 2)) - + jnp.sum(jnp.log(2 * jnp.pi * noise_imag ** 2)) + noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_real**2)) + jnp.sum( + jnp.log(2 * jnp.pi * noise_imag**2) ) return -0.5 * ( - chi_squared + regularization_term + log_det_curvature_reg - - log_det_regularization + noise_normalization + chi_squared + + regularization_term + + log_det_curvature_reg + - log_det_regularization + + noise_normalization ) mapper_indices_jnp = jnp.array(np.asarray(inversion.mapper_indices)) @@ -725,17 +745,32 @@ def compute_log_evidence( with timer.section("log_evidence_eager"): log_evidence_one_channel = compute_log_evidence( - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - transformed_mm_real_jnp, transformed_mm_imag_jnp, - reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + reconstruction, + curvature_matrix, + reg_jnp, + mapper_indices_jnp, ) block(log_evidence_one_channel) _, log_evidence_one_channel = jit_profile( - compute_log_evidence, "log_evidence_jit", - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - transformed_mm_real_jnp, transformed_mm_imag_jnp, - reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp, + compute_log_evidence, + "log_evidence_jit", + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + reconstruction, + curvature_matrix, + reg_jnp, + mapper_indices_jnp, ) log_evidence_per_channel_cost = timer.records[-1][1] / 10 likelihood_steps.append( @@ -759,7 +794,8 @@ def compute_log_evidence( jnp.array(dataset_list[c].data.imag), jnp.array(dataset_list[c].noise_map.real), jnp.array(dataset_list[c].noise_map.imag), - tm_real, tm_imag, + tm_real, + tm_imag, jnp.asarray(inv_c.reconstruction), jnp.asarray(inv_c.curvature_matrix), jnp.array(inv_c.regularization_matrix), @@ -768,19 +804,14 @@ def compute_log_evidence( log_evidence_check_per_channel.append(float(le_c)) cube_log_evidence_check = float(sum(log_evidence_check_per_channel)) - print( - f"\n cube log_evidence (per-step recompute) = {cube_log_evidence_check:.6f}" - ) + print(f"\n cube log_evidence (per-step recompute) = {cube_log_evidence_check:.6f}") print(f" cube log_evidence (reference) = {cube_log_evidence_ref:.6f}") np.testing.assert_allclose( cube_log_evidence_check, cube_log_evidence_ref, rtol=1e-4, - err_msg=( - "Per-step cube log_evidence does not match summed " - "FitInterferometer.log_evidence" - ), + err_msg=("Per-step cube log_evidence does not match summed FitInterferometer.log_evidence"), ) print( " Assertion PASSED: per-step cube log_evidence matches summed " @@ -797,8 +828,8 @@ def compute_log_evidence( # sparse_operator W~ kernel (~41 MB) — all fit. Extract them and # profile each sparse step individually, mirroring the dense breakdown. - from autoarray.inversion.mappers.mapper_util import sparse_triplets_from from autoarray.inversion.mappers.abstract import Mapper as _MapperCls + from autoarray.inversion.mappers.mapper_util import sparse_triplets_from # --------------------------------------------------------------- # Extract sparse references from channel 0's inversion @@ -866,9 +897,7 @@ def inversion_setup_from_params(params_tree): ) return jnp.asarray(fit_jax.inversion.data_vector) - _, _ = jit_profile( - inversion_setup_from_params, "inversion_setup_jit", params_tree - ) + _, _ = jit_profile(inversion_setup_from_params, "inversion_setup_jit", params_tree) inversion_setup_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( ( @@ -942,8 +971,10 @@ def compute_data_vector_sparse(mapping_matrix, dirty_image): block(data_vector) _, data_vector = jit_profile( - compute_data_vector_sparse, "data_vector_sparse_jit", - mapping_matrix_jnp, dirty_image_jnp, + compute_data_vector_sparse, + "data_vector_sparse_jit", + mapping_matrix_jnp, + dirty_image_jnp, ) data_vector_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( @@ -960,17 +991,18 @@ def compute_data_vector_sparse(mapping_matrix, dirty_image): print("\n--- Step 5: Sparse curvature matrix F via FFT W~ (per channel) ---") def compute_curvature_matrix_sparse(rows, cols, vals): - return sparse_operator.curvature_matrix_diag_from( - rows=rows, cols=cols, vals=vals, S=S - ) + return sparse_operator.curvature_matrix_diag_from(rows=rows, cols=cols, vals=vals, S=S) with timer.section("curvature_matrix_sparse_eager"): curvature_matrix = compute_curvature_matrix_sparse(rows_jnp, cols_jnp, vals_jnp) block(curvature_matrix) _, curvature_matrix = jit_profile( - compute_curvature_matrix_sparse, "curvature_matrix_sparse_jit", - rows_jnp, cols_jnp, vals_jnp, + compute_curvature_matrix_sparse, + "curvature_matrix_sparse_jit", + rows_jnp, + cols_jnp, + vals_jnp, ) curvature_matrix_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( @@ -986,9 +1018,7 @@ def compute_curvature_matrix_sparse(rows, cols, vals): print("\n--- Step 6: Regularization matrix H (shared) ---") - likelihood_steps.append( - ("Regularization matrix H (shared)", 0.0) - ) + likelihood_steps.append(("Regularization matrix H (shared)", 0.0)) # --------------------------------------------------------------- # Step 7: Reconstruction NNLS (per channel) — same as dense path @@ -1011,8 +1041,11 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) block(reconstruction) _, reconstruction = jit_profile( - compute_reconstruction, "reconstruction_jit", - data_vector, curvature_matrix, regularization_matrix, + compute_reconstruction, + "reconstruction_jit", + data_vector, + curvature_matrix, + regularization_matrix, ) reconstruction_per_channel = timer.records[-1][1] / 10 likelihood_steps.append( @@ -1029,17 +1062,20 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) print("\n--- Step 8: fast chi-squared + log evidence (per channel) ---") def compute_log_evidence_sparse( - data_real, data_imag, noise_real, noise_imag, - reconstruction, data_vector, curvature_matrix, - regularization_matrix, mapper_indices, + data_real, + data_imag, + noise_real, + noise_imag, + reconstruction, + data_vector, + curvature_matrix, + regularization_matrix, + mapper_indices, ): # fast_chi_squared (avoids forming residual visibilities) chi_sq_t1 = jnp.dot(reconstruction, jnp.dot(curvature_matrix, reconstruction)) chi_sq_t2 = -2.0 * jnp.dot(reconstruction, data_vector) - chi_sq_t3 = ( - jnp.sum(data_real ** 2 / noise_real ** 2) - + jnp.sum(data_imag ** 2 / noise_imag ** 2) - ) + chi_sq_t3 = jnp.sum(data_real**2 / noise_real**2) + jnp.sum(data_imag**2 / noise_imag**2) chi_squared = chi_sq_t1 + chi_sq_t2 + chi_sq_t3 regularization_term = jnp.dot( @@ -1049,36 +1085,47 @@ def compute_log_evidence_sparse( curvature_reg_matrix = curvature_matrix + regularization_matrix creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices] reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices] - log_det_curvature_reg = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced))) - ) - log_det_regularization = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced))) - ) + log_det_curvature_reg = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))) + log_det_regularization = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))) - noise_normalization = ( - jnp.sum(jnp.log(2 * jnp.pi * noise_real ** 2)) - + jnp.sum(jnp.log(2 * jnp.pi * noise_imag ** 2)) + noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_real**2)) + jnp.sum( + jnp.log(2 * jnp.pi * noise_imag**2) ) return -0.5 * ( - chi_squared + regularization_term + log_det_curvature_reg - - log_det_regularization + noise_normalization + chi_squared + + regularization_term + + log_det_curvature_reg + - log_det_regularization + + noise_normalization ) with timer.section("log_evidence_sparse_eager"): log_evidence_one_channel = compute_log_evidence_sparse( - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - reconstruction, data_vector, curvature_matrix, - regularization_matrix, mapper_indices_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + reconstruction, + data_vector, + curvature_matrix, + regularization_matrix, + mapper_indices_jnp, ) block(log_evidence_one_channel) _, log_evidence_one_channel = jit_profile( - compute_log_evidence_sparse, "log_evidence_sparse_jit", - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - reconstruction, data_vector, curvature_matrix, - regularization_matrix, mapper_indices_jnp, + compute_log_evidence_sparse, + "log_evidence_sparse_jit", + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + reconstruction, + data_vector, + curvature_matrix, + regularization_matrix, + mapper_indices_jnp, ) log_evidence_per_channel_cost = timer.records[-1][1] / 10 likelihood_steps.append( @@ -1112,14 +1159,16 @@ def compute_log_evidence_sparse( jnp.array(dataset_list[c].data.imag), jnp.array(dataset_list[c].noise_map.real), jnp.array(dataset_list[c].noise_map.imag), - rc_c, dv_c, cm_c, rm_c, mi_c, + rc_c, + dv_c, + cm_c, + rm_c, + mi_c, ) log_evidence_check_per_channel.append(float(le_c)) cube_log_evidence_check = float(sum(log_evidence_check_per_channel)) - print( - f"\n cube log_evidence (per-step recompute) = {cube_log_evidence_check:.6f}" - ) + print(f"\n cube log_evidence (per-step recompute) = {cube_log_evidence_check:.6f}") print(f" cube log_evidence (reference) = {cube_log_evidence_ref:.6f}") np.testing.assert_allclose( @@ -1127,8 +1176,7 @@ def compute_log_evidence_sparse( cube_log_evidence_ref, rtol=1e-4, err_msg=( - "Per-step sparse cube log_evidence does not match summed " - "FitInterferometer.log_evidence" + "Per-step sparse cube log_evidence does not match summed FitInterferometer.log_evidence" ), ) print( @@ -1142,14 +1190,18 @@ def compute_log_evidence_sparse( # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt al_version = al.__version__ print("\n" + "=" * 70) -print(f"JAX LIKELIHOOD BREAKDOWN SUMMARY — CUBE {instrument.upper()} × {n_channels} — v{al_version}") +print( + f"JAX LIKELIHOOD BREAKDOWN SUMMARY — CUBE {instrument.upper()} × {n_channels} — v{al_version}" +) print("=" * 70) print(f" Instrument: {instrument}") print(f" Channels: {n_channels}") @@ -1185,19 +1237,21 @@ def compute_log_evidence_sparse( # the sparse path where the sub-decomposition ran. shared_invariant_savings = None if "inversion_setup_invariant_per_channel" in dir() and curvature_matrix_per_channel is not None: - per_channel_invariant = ( - inversion_setup_invariant_per_channel + curvature_matrix_per_channel - ) + per_channel_invariant = inversion_setup_invariant_per_channel + curvature_matrix_per_channel shared_invariant_savings = (n_channels - 1) * per_channel_invariant print("-" * 70) print(f" {'TOTAL (step-by-step cube cost)':<{max_label}} {step_total:>12.6f} s") if shared_lwl_savings is not None: - print(f" {f'Shared-Lᵀ W̃ L savings (curvature only, est.)':<{max_label}} {shared_lwl_savings:>12.6f} s") + print( + f" {'Shared-Lᵀ W̃ L savings (curvature only, est.)':<{max_label}} {shared_lwl_savings:>12.6f} s" + ) else: - print(f" Shared-Lᵀ W̃ L savings: N/A (run at SMA scale for per-step granularity)") + print(" Shared-Lᵀ W̃ L savings: N/A (run at SMA scale for per-step granularity)") if shared_invariant_savings is not None: - print(f" {f'Shared-state savings (mapper+L+F, est.)':<{max_label}} {shared_invariant_savings:>12.6f} s") + print( + f" {'Shared-state savings (mapper+L+F, est.)':<{max_label}} {shared_invariant_savings:>12.6f} s" + ) print("=" * 70) # --- Save results dictionary --- @@ -1290,7 +1344,7 @@ def compute_log_evidence_sparse( fontweight="bold", ) ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | " + f'AutoLens v{al_version} | {pixel_scale}"/px | ' f"{real_space_shape[0]}x{real_space_shape[1]} real-space | " f"{n_visibilities} visibilities/chan | {n_mesh_vertices} Delaunay verts | " f"step-by-step total: {step_total:.6f} s", @@ -1320,9 +1374,7 @@ def compute_log_evidence_sparse( } _per_channel = EXPECTED_LOG_EVIDENCE_PER_CHANNEL.get(instrument) -expected_cube_log_evidence = ( - n_channels * _per_channel if _per_channel is not None else None -) +expected_cube_log_evidence = n_channels * _per_channel if _per_channel is not None else None if expected_cube_log_evidence is None: print( diff --git a/likelihood_breakdown/datacube/inversion_setup_decompose.py b/likelihood_breakdown/datacube/inversion_setup_decompose.py index 00e12e7..f6578b7 100644 --- a/likelihood_breakdown/datacube/inversion_setup_decompose.py +++ b/likelihood_breakdown/datacube/inversion_setup_decompose.py @@ -64,36 +64,38 @@ Python imports it fine.) """ -import numpy as np -import jax -import jax.numpy as jnp import os -import time import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 + _cli = parse_profile_cli() instrument = _cli.instrument or "sma" @@ -350,10 +352,14 @@ def cut_data_vector(params_tree): print(f" {'channel-INVARIANT (shareable 1x)':<34}{channel_invariant:>12.6f} s") print(f" {'channel-VARIANT (stays Nx)':<34}{channel_variant:>12.6f} s") if per_channel_total > 0: - print(f" shareable fraction of inversion total: {100*channel_invariant/per_channel_total:.1f}%") + print( + f" shareable fraction of inversion total: {100 * channel_invariant / per_channel_total:.1f}%" + ) print("-" * 64) print(" raw cutpoint leaves (NOT a linear chain — C,D are parallel children of B):") -print(f" A trace={cost_A:.6f} B mapper+L={cost_B:.6f} C +curv={cost_C:.6f} D +datavec={cost_D:.6f}") +print( + f" A trace={cost_A:.6f} B mapper+L={cost_B:.6f} C +curv={cost_C:.6f} D +datavec={cost_D:.6f}" +) print("=" * 64) # --------------------------------------------------------------------------- diff --git a/likelihood_breakdown/imaging/delaunay.py b/likelihood_breakdown/imaging/delaunay.py index b489c54..8e4d5bb 100644 --- a/likelihood_breakdown/imaging/delaunay.py +++ b/likelihood_breakdown/imaging/delaunay.py @@ -44,33 +44,32 @@ the basename ``delaunay_breakdown_{instrument}_v{al_version}``. """ -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -78,12 +77,13 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.imaging import INSTRUMENTS # noqa: E402 + _cli = parse_profile_cli() instrument = "hst" # <-- change this to profile a different instrument @@ -93,6 +93,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -228,9 +229,7 @@ def jit_profile(func, label, *args, n_repeats=10): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") @@ -239,9 +238,7 @@ def jit_profile(func, label, *args, n_repeats=10): n_mesh_vertices = 1500 # 1500-tier production fiducial with timer.section("image_mesh_hilbert"): - image_mesh = al.image_mesh.Hilbert( - pixels=n_mesh_vertices, weight_power=1.0, weight_floor=0.0 - ) + 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 ) @@ -281,9 +278,7 @@ def jit_profile(func, label, *args, n_repeats=10): shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005) shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005) - lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear - ) + lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) mesh = al.mesh.Delaunay( pixels=n_mesh_vertices, @@ -402,23 +397,21 @@ def jit_profile(func, label, *args, n_repeats=10): print("\n--- Step 1: Ray-trace data grid ---") with timer.section("ray_trace_data_eager"): - traced_grids = tracer.traced_grid_2d_list_from( - grid=dataset.grids.pixelization, xp=jnp - ) + traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp) for tg in traced_grids: block(tg) print(f" Number of planes traced: {len(traced_grids)}") + def ray_trace_data_raw(grid_raw): """Wraps ray-tracing so inputs/outputs are raw arrays.""" grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp) traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp) return jnp.stack([tg.array for tg in traced]) -_, traced_data_grids_raw = jit_profile( - ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw -) + +_, traced_data_grids_raw = jit_profile(ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw) likelihood_steps.append(("Ray-trace data grid", timer.records[-1][1] / 10)) print(f" traced_data_grids shape: {traced_data_grids_raw.shape}") @@ -436,15 +429,15 @@ def ray_trace_data_raw(grid_raw): for tg in traced_mesh: block(tg) + def ray_trace_mesh_raw(mesh_raw): """Ray-trace image-plane mesh vertices to source plane.""" grid = aa.Grid2DIrregular(values=mesh_raw, xp=jnp) traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp) return jnp.stack([tg.array for tg in traced]) -_, traced_mesh_grids_raw = jit_profile( - ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw -) + +_, traced_mesh_grids_raw = jit_profile(ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw) likelihood_steps.append(("Ray-trace mesh grid", timer.records[-1][1] / 10)) print(f" traced_mesh_grids shape: {traced_mesh_grids_raw.shape}") @@ -470,13 +463,12 @@ def ray_trace_mesh_raw(mesh_raw): def lens_image_raw(grid_raw, blurring_grid_raw): """Compute lens light images on masked + blurring grids (no PSF).""" grid = aa.Grid2D(values=grid_raw, mask=_grid_lp_mask, xp=jnp) - blurring_grid = aa.Grid2D( - values=blurring_grid_raw, mask=_grid_blurring_mask, xp=jnp - ) + blurring_grid = aa.Grid2D(values=blurring_grid_raw, mask=_grid_blurring_mask, xp=jnp) image = tracer.image_2d_from(grid=grid, xp=jnp) blurring_image = tracer.image_2d_from(grid=blurring_grid, xp=jnp) return image.array, blurring_image.array + with timer.section("lens_image_eager"): img_eager, blur_img_eager = lens_image_raw(grid_lp_raw, grid_blurring_raw) block(img_eager) @@ -499,17 +491,20 @@ def lens_image_raw(grid_raw, blurring_grid_raw): print(f" blurred_image shape: {blurred_image.array.shape}") + def blurred_image_from_params(params_tree): """Compute blurred image directly from a pytree ModelInstance — fully JIT-traceable.""" t = al.Tracer(galaxies=list(params_tree.galaxies)) result = t.blurred_image_2d_from( - grid=grid_lp, psf=dataset.psf, blurring_grid=grid_blurring, xp=jnp, + grid=grid_lp, + psf=dataset.psf, + blurring_grid=grid_blurring, + xp=jnp, ) return result.array -_, blurred_img_jit = jit_profile( - blurred_image_from_params, "blurred_image_jit", params_tree -) + +_, blurred_img_jit = jit_profile(blurred_image_from_params, "blurred_image_jit", params_tree) likelihood_steps.append(("Blurred image (PSF convolution)", timer.records[-1][1] / 10)) # --------------------------------------------------------------------------- @@ -518,9 +513,11 @@ def blurred_image_from_params(params_tree): print("\n--- Step 4: Profile-subtracted image ---") + def profile_subtract(data, blurred_image): return data - blurred_image + with timer.section("profile_subtract_eager"): blurred_img_jnp = jnp.array(blurred_image.array) profile_subtracted = profile_subtract(data_array, blurred_img_jnp) @@ -544,9 +541,7 @@ def profile_subtract(data, blurred_image): with timer.section("border_relocator_setup"): border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1) -traced_source_grid = tracer.traced_grid_2d_list_from( - grid=dataset.grids.pixelization, xp=jnp -)[-1] +traced_source_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp)[-1] traced_mesh_source = tracer.traced_grid_2d_list_from( grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp )[-1] @@ -554,7 +549,8 @@ def profile_subtract(data, blurred_image): with timer.section("border_relocation_eager"): relocated_grid = border_relocator.relocated_grid_from(grid=traced_source_grid) relocated_mesh_grid = border_relocator.relocated_mesh_grid_from( - grid=traced_source_grid, mesh_grid=traced_mesh_source, + grid=traced_source_grid, + mesh_grid=traced_mesh_source, ) block(relocated_grid) block(relocated_mesh_grid) @@ -638,6 +634,7 @@ def profile_subtract(data, blurred_image): # border relocation → Delaunay triangulation → interpolation → mapper → mapping matrix → PSF convolution. # These steps are tightly sequential; the full pipeline JIT-compiles them all together. + def blurred_mm_from_params(params_tree): """Compute blurred mapping matrix via full inversion setup from a pytree ModelInstance.""" t = al.Tracer(galaxies=list(params_tree.galaxies)) @@ -651,14 +648,18 @@ def blurred_mm_from_params(params_tree): }, ) fit_jax = al.FitImaging( - dataset=dataset, tracer=t, adapt_images=adapt_images_jax, + dataset=dataset, + tracer=t, + adapt_images=adapt_images_jax, settings=al.Settings( use_border_relocator=True, use_mixed_precision=_cli.use_mixed_precision, - ), xp=jnp, + ), + xp=jnp, ) return jnp.array(fit_jax.inversion.operated_mapping_matrix) + _, bmm_jit = jit_profile(blurred_mm_from_params, "inversion_setup_jit", params_tree) likelihood_steps.append(("Inversion setup (steps 5-8 combined)", timer.records[-1][1] / 10)) @@ -673,6 +674,7 @@ def blurred_mm_from_params(params_tree): print("\n--- Step 9: Data vector ---") + def compute_data_vector(blurred_mapping_matrix, image, noise_map): return al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from( blurred_mapping_matrix=blurred_mapping_matrix, @@ -680,6 +682,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): noise_map=noise_map, ) + profile_sub_jnp = jnp.array(fit.profile_subtracted_image.array) noise_jnp = jnp.array(dataset.noise_map.array) @@ -702,6 +705,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): no_reg_list = list(inversion.no_regularization_index_list) + def compute_curvature_matrix(blurred_mapping_matrix, noise_map): return al.util.inversion.curvature_matrix_via_mapping_matrix_from( mapping_matrix=blurred_mapping_matrix, @@ -712,6 +716,7 @@ def compute_curvature_matrix(blurred_mapping_matrix, noise_map): xp=jnp, ) + with timer.section("curvature_matrix_eager"): curvature_matrix = compute_curvature_matrix(bmm_jnp, noise_jnp) block(curvature_matrix) @@ -756,6 +761,7 @@ def compute_curvature_matrix(blurred_mapping_matrix, noise_map): print("\n--- Step 12: Regularized reconstruction ---") + def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix): curvature_reg_matrix = curvature_matrix + regularization_matrix return al.util.inversion.reconstruction_positive_only_from( @@ -764,6 +770,7 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) xp=jnp, ) + with timer.section("reconstruction_eager"): reconstruction = compute_reconstruction( jnp.array(data_vector), @@ -773,7 +780,8 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) block(reconstruction) _, reconstruction = jit_profile( - compute_reconstruction, "reconstruction_jit", + compute_reconstruction, + "reconstruction_jit", jnp.array(data_vector), jnp.array(curvature_matrix), jnp.array(regularization_matrix), @@ -788,9 +796,16 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) print("\n--- Step 13: Mapped reconstruction + log evidence ---") + def compute_log_evidence( - data, noise_map, blurred_image, blurred_mapping_matrix, reconstruction, - reduced_indices, reg_reduced, curv_reg_reduced, + data, + noise_map, + blurred_image, + blurred_mapping_matrix, + reconstruction, + reduced_indices, + reg_reduced, + curv_reg_reduced, ): """Compute the full log evidence including all five terms: @@ -845,7 +860,7 @@ def compute_log_evidence( log_det_regularization = 2.0 * jnp.sum(jnp.log(jnp.diag(L_r))) # Noise normalization (over full masked image) - noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_map ** 2)) + noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_map**2)) return -0.5 * ( chi_squared @@ -855,6 +870,7 @@ def compute_log_evidence( + noise_normalization ) + # For the JIT profiling we use the step-by-step reconstruction for timing. # For the correctness assertion we use the inversion's own reconstruction, # because cumulative floating-point differences between JIT-compiled and @@ -874,15 +890,28 @@ def compute_log_evidence( with timer.section("log_evidence_eager"): log_evidence = compute_log_evidence( - data_array, noise_jnp, blurred_img_jnp, bmm_jnp, - recon_jnp, reduced_indices_jnp, reg_reduced_jnp, curv_reg_reduced_jnp, + data_array, + noise_jnp, + blurred_img_jnp, + bmm_jnp, + recon_jnp, + reduced_indices_jnp, + reg_reduced_jnp, + curv_reg_reduced_jnp, ) block(log_evidence) _, log_evidence = jit_profile( - compute_log_evidence, "log_evidence_jit", - data_array, noise_jnp, blurred_img_jnp, bmm_jnp, - recon_jnp, reduced_indices_jnp, reg_reduced_jnp, curv_reg_reduced_jnp, + compute_log_evidence, + "log_evidence_jit", + data_array, + noise_jnp, + blurred_img_jnp, + bmm_jnp, + recon_jnp, + reduced_indices_jnp, + reg_reduced_jnp, + curv_reg_reduced_jnp, ) likelihood_steps.append(("Mapped recon + log evidence", timer.records[-1][1] / 10)) @@ -894,8 +923,14 @@ def compute_log_evidence( inv_recon_jnp = jnp.array(inversion.reconstruction) log_evidence_check = compute_log_evidence( - data_array, noise_jnp, blurred_img_jnp, bmm_jnp, - inv_recon_jnp, reduced_indices_jnp, reg_reduced_jnp, curv_reg_reduced_jnp, + data_array, + noise_jnp, + blurred_img_jnp, + bmm_jnp, + inv_recon_jnp, + reduced_indices_jnp, + reg_reduced_jnp, + curv_reg_reduced_jnp, ) print(f" log_evidence (inv matrices) = {log_evidence_check}") print(f" log_evidence (reference) = {log_evidence_ref}") @@ -913,7 +948,9 @@ def compute_log_evidence( # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -996,7 +1033,7 @@ def compute_log_evidence( fontweight="bold", ) ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | {n_image_pixels} pixels | " + f'AutoLens v{al_version} | {pixel_scale}"/px | {n_image_pixels} pixels | ' f"{n_over_sampled_pixels} over-sampled | {n_source_pixels} Delaunay vertices | " f"total: {step_total:.6f} s", fontsize=9, @@ -1013,7 +1050,9 @@ def compute_log_evidence( # Regression assertion — eager log_evidence only # =================================================================== -EXPECTED_LOG_EVIDENCE_HST = 29110.92085793 # 1500-pixel Hilbert/Delaunay, MGE-60 lens, adapt_image=lensed_source +EXPECTED_LOG_EVIDENCE_HST = ( + 29110.92085793 # 1500-pixel Hilbert/Delaunay, MGE-60 lens, adapt_image=lensed_source +) np.testing.assert_allclose( log_evidence_ref, @@ -1024,7 +1063,4 @@ def compute_log_evidence( f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})" ), ) -print( - f" Eager regression assertion PASSED: log_evidence matches " - f"{EXPECTED_LOG_EVIDENCE_HST:.6f}" -) +print(f" Eager regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") diff --git a/likelihood_breakdown/imaging/mge.py b/likelihood_breakdown/imaging/mge.py index 5d5dab6..9dffbb1 100644 --- a/likelihood_breakdown/imaging/mge.py +++ b/likelihood_breakdown/imaging/mge.py @@ -42,30 +42,28 @@ the basename ``mge_breakdown_{instrument}_v{al_version}``. """ -import numpy as np -import jax -import jax.numpy as jnp -import time -import subprocess -import sys -from pathlib import Path -from contextlib import contextmanager - -import autofit as af -import autolens as al -import autoarray as aa -from autofit.jax import register_model as _register_model_pytrees - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os +import subprocess +import sys import sys as _smoke_sys +import time +from contextlib import contextmanager +from pathlib import Path + +import autoarray as aa +import autofit as af +import autolens as al +import jax +import jax.numpy as jnp +import numpy as np +from autofit.jax import register_model as _register_model_pytrees + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -74,12 +72,13 @@ # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.imaging import INSTRUMENTS # noqa: E402 + _cli = parse_profile_cli() instrument = "hst" # <-- change this to profile a different instrument @@ -89,6 +88,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -246,9 +246,7 @@ def jit_profile(func, label, *args, n_repeats=10): shear.gamma_1 = 0.05 shear.gamma_2 = 0.05 - lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear - ) + lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) source_bulge = al.model_util.mge_model_from( mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False @@ -353,12 +351,14 @@ def jit_profile(func, label, *args, n_repeats=10): print(f" Number of planes traced: {len(traced_grids)}") + def ray_trace_raw(grid_raw): """Wraps ray-tracing so inputs/outputs are raw arrays.""" grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp) traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp) return jnp.stack([tg.array for tg in traced]) + _, traced_grids_raw = jit_profile(ray_trace_raw, "ray_trace_jit", grid_lp_raw) likelihood_steps.append(("Ray-trace grids", timer.records[-1][1] / 10)) @@ -390,10 +390,13 @@ def ray_trace_raw(grid_raw): # mapping_matrix and operated_mapping_matrix_override already return raw arrays. with timer.section("mapping_matrix"): mapping_matrices = [func.mapping_matrix for func in lp_linear_funcs] - mapping_matrix = np.hstack(mapping_matrices) if len(mapping_matrices) > 1 else mapping_matrices[0] + mapping_matrix = ( + np.hstack(mapping_matrices) if len(mapping_matrices) > 1 else mapping_matrices[0] + ) print(f" mapping_matrix shape: {mapping_matrix.shape}") + def mapping_matrix_from_params(params_tree): """Compute mapping matrix from a pytree-shaped ``ModelInstance``. @@ -417,6 +420,7 @@ def mapping_matrix_from_params(params_tree): matrices = [f.mapping_matrix for f in funcs] return jnp.hstack(matrices) if len(matrices) > 1 else matrices[0] + _, mm_jit = jit_profile(mapping_matrix_from_params, "mapping_matrix_jit", params_tree) likelihood_steps.append(("Mapping matrix", timer.records[-1][1] / 10)) @@ -430,10 +434,13 @@ def mapping_matrix_from_params(params_tree): with timer.section("blurred_mapping_matrix"): blurred_matrices = [func.operated_mapping_matrix_override for func in lp_linear_funcs] - blurred_mapping_matrix = np.hstack(blurred_matrices) if len(blurred_matrices) > 1 else blurred_matrices[0] + blurred_mapping_matrix = ( + np.hstack(blurred_matrices) if len(blurred_matrices) > 1 else blurred_matrices[0] + ) print(f" blurred_mapping_matrix shape: {blurred_mapping_matrix.shape}") + def blurred_mm_from_params(params_tree): """Compute blurred mapping matrix from a pytree-shaped ``ModelInstance``.""" t = al.Tracer(galaxies=list(params_tree.galaxies)) @@ -453,6 +460,7 @@ def blurred_mm_from_params(params_tree): matrices = [f.operated_mapping_matrix_override for f in funcs] return jnp.hstack(matrices) if len(matrices) > 1 else matrices[0] + _, bmm_jit = jit_profile(blurred_mm_from_params, "blurred_mm_jit", params_tree) likelihood_steps.append(("Blurred mapping matrix", timer.records[-1][1] / 10)) @@ -464,6 +472,7 @@ def blurred_mm_from_params(params_tree): print("\n--- Step 4: Data vector ---") + def compute_data_vector(blurred_mapping_matrix, image, noise_map): return al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from( blurred_mapping_matrix=blurred_mapping_matrix, @@ -471,6 +480,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): noise_map=noise_map, ) + bmm_jnp = jnp.array(blurred_mapping_matrix) noise_jnp = jnp.array(dataset.noise_map.array) @@ -478,9 +488,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): data_vector = compute_data_vector(bmm_jnp, data_array, noise_jnp) block(data_vector) -_, data_vector = jit_profile( - compute_data_vector, "data_vector_jit", bmm_jnp, data_array, noise_jnp -) +_, data_vector = jit_profile(compute_data_vector, "data_vector_jit", bmm_jnp, data_array, noise_jnp) likelihood_steps.append(("Data vector (D)", timer.records[-1][1] / 10)) print(f" data_vector shape: {data_vector.shape}") @@ -493,6 +501,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): n_linear = bmm_jnp.shape[1] + def compute_curvature_matrix(blurred_mapping_matrix, noise_map): return al.util.inversion.curvature_matrix_via_mapping_matrix_from( mapping_matrix=blurred_mapping_matrix, @@ -502,6 +511,7 @@ def compute_curvature_matrix(blurred_mapping_matrix, noise_map): xp=jnp, ) + with timer.section("curvature_matrix_eager"): curvature_matrix = compute_curvature_matrix(bmm_jnp, noise_jnp) block(curvature_matrix) @@ -519,6 +529,7 @@ def compute_curvature_matrix(blurred_mapping_matrix, noise_map): print("\n--- Step 6: Reconstruction (NNLS) ---") + def compute_reconstruction(data_vector, curvature_matrix): return al.util.inversion.reconstruction_positive_only_from( data_vector=data_vector, @@ -526,15 +537,16 @@ def compute_reconstruction(data_vector, curvature_matrix): xp=jnp, ) + with timer.section("reconstruction_eager"): - reconstruction = compute_reconstruction( - jnp.array(data_vector), jnp.array(curvature_matrix) - ) + reconstruction = compute_reconstruction(jnp.array(data_vector), jnp.array(curvature_matrix)) block(reconstruction) _, reconstruction = jit_profile( - compute_reconstruction, "reconstruction_jit", - jnp.array(data_vector), jnp.array(curvature_matrix) + compute_reconstruction, + "reconstruction_jit", + jnp.array(data_vector), + jnp.array(curvature_matrix), ) likelihood_steps.append(("Reconstruction (NNLS)", timer.records[-1][1] / 10)) @@ -546,6 +558,7 @@ def compute_reconstruction(data_vector, curvature_matrix): print("\n--- Step 7: Mapped reconstructed image ---") + def compute_mapped_recon(blurred_mapping_matrix, reconstruction): return al.util.inversion.mapped_reconstructed_data_via_mapping_matrix_from( mapping_matrix=blurred_mapping_matrix, @@ -553,6 +566,7 @@ def compute_mapped_recon(blurred_mapping_matrix, reconstruction): xp=jnp, ) + with timer.section("mapped_recon_eager"): mapped_recon = compute_mapped_recon(bmm_jnp, jnp.array(reconstruction)) block(mapped_recon) @@ -570,23 +584,22 @@ def compute_mapped_recon(blurred_mapping_matrix, reconstruction): print("\n--- Step 8: Chi-squared & log likelihood ---") + def compute_log_likelihood(data, noise_map, mapped_recon): residual = data - mapped_recon chi_squared = jnp.sum((residual / noise_map) ** 2) - noise_norm = jnp.sum(jnp.log(2 * jnp.pi * noise_map ** 2)) + noise_norm = jnp.sum(jnp.log(2 * jnp.pi * noise_map**2)) return -0.5 * (chi_squared + noise_norm) + mapped_recon_jnp = jnp.array(mapped_recon) with timer.section("log_likelihood_eager"): - log_like = compute_log_likelihood( - data_array, noise_jnp, mapped_recon_jnp - ) + log_like = compute_log_likelihood(data_array, noise_jnp, mapped_recon_jnp) block(log_like) _, log_like = jit_profile( - compute_log_likelihood, "log_likelihood_jit", - data_array, noise_jnp, mapped_recon_jnp + compute_log_likelihood, "log_likelihood_jit", data_array, noise_jnp, mapped_recon_jnp ) likelihood_steps.append(("Chi-squared & log likelihood", timer.records[-1][1] / 10)) @@ -607,7 +620,9 @@ def compute_log_likelihood(data, noise_map, mapped_recon): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -688,7 +703,7 @@ def compute_log_likelihood(data, noise_map, mapped_recon): fontweight="bold", ) ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | {n_image_pixels} pixels | " + f'AutoLens v{al_version} | {pixel_scale}"/px | {n_image_pixels} pixels | ' f"{n_over_sampled_pixels} over-sampled | {n_linear_gaussians} Gaussians | " f"total: {step_total:.6f} s", fontsize=9, @@ -717,6 +732,5 @@ def compute_log_likelihood(data, noise_map, mapped_recon): ), ) print( - f" Eager regression assertion PASSED: log_likelihood matches " - f"{EXPECTED_LOG_LIKELIHOOD_HST:.6f}" + f" Eager regression assertion PASSED: log_likelihood matches {EXPECTED_LOG_LIKELIHOOD_HST:.6f}" ) diff --git a/likelihood_breakdown/imaging/pixelization.py b/likelihood_breakdown/imaging/pixelization.py index 0fffb9d..7508f81 100644 --- a/likelihood_breakdown/imaging/pixelization.py +++ b/likelihood_breakdown/imaging/pixelization.py @@ -37,35 +37,34 @@ the basename ``pixelization_breakdown_{instrument}_v{al_version}``. """ -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees # Shared adapt-image loader: load or compute+cache `lensed_source.fits` # next to the dataset, then return the masked ``aa.Array2D``. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -73,12 +72,13 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.imaging import INSTRUMENTS # noqa: E402 + _cli = parse_profile_cli() instrument = "hst" # <-- change this to profile a different instrument @@ -88,6 +88,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -250,18 +251,14 @@ def jit_profile(func, label, *args, n_repeats=10): shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005) shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005) - lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear - ) + lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) # ``RectangularAdaptImage`` weights mesh pixels by the lensed-source # adapt image — the production-grade alternative to the coordinate- # density-only ``RectangularAdaptDensity``. Adapt image is loaded / # cached below; the same shape and regularization are kept. pixelization = al.Pixelization( - mesh=al.mesh.RectangularAdaptImage( - shape=mesh_shape, weight_power=1.0, weight_floor=0.0 - ), + mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape, weight_power=1.0, weight_floor=0.0), regularization=al.reg.Constant(coefficient=1.0), ) @@ -316,9 +313,7 @@ def jit_profile(func, label, *args, n_repeats=10): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) adapt_images = al.AdaptImages( galaxy_image_dict={instance.galaxies.source: adapt_image}, galaxy_name_image_dict={"('galaxies', 'source')": adapt_image}, @@ -384,12 +379,14 @@ def jit_profile(func, label, *args, n_repeats=10): print(f" Number of planes traced: {len(traced_grids)}") + def ray_trace_raw(grid_raw): """Wraps ray-tracing so inputs/outputs are raw arrays.""" grid = aa.Grid2DIrregular(values=grid_raw, xp=jnp) traced = tracer.traced_grid_2d_list_from(grid=grid, xp=jnp) return jnp.stack([tg.array for tg in traced]) + _, traced_grids_raw = jit_profile(ray_trace_raw, "ray_trace_jit", grid_pix_raw) likelihood_steps.append(("Ray-trace grids", timer.records[-1][1] / 10)) @@ -417,13 +414,12 @@ def ray_trace_raw(grid_raw): def lens_image_raw(grid_raw, blurring_grid_raw): """Compute lens light images on masked + blurring grids (no PSF).""" grid = aa.Grid2D(values=grid_raw, mask=_grid_lp_mask, xp=jnp) - blurring_grid = aa.Grid2D( - values=blurring_grid_raw, mask=_grid_blurring_mask, xp=jnp - ) + blurring_grid = aa.Grid2D(values=blurring_grid_raw, mask=_grid_blurring_mask, xp=jnp) image = tracer.image_2d_from(grid=grid, xp=jnp) blurring_image = tracer.image_2d_from(grid=blurring_grid, xp=jnp) return image.array, blurring_image.array + with timer.section("lens_image_eager"): img_eager, blur_img_eager = lens_image_raw(grid_lp_raw, grid_blurring_raw) block(img_eager) @@ -446,17 +442,20 @@ def lens_image_raw(grid_raw, blurring_grid_raw): print(f" blurred_image shape: {blurred_image.array.shape}") + def blurred_image_from_params(params_tree): """Compute blurred image directly from a pytree ModelInstance — fully JIT-traceable.""" t = al.Tracer(galaxies=list(params_tree.galaxies)) result = t.blurred_image_2d_from( - grid=grid_lp, psf=dataset.psf, blurring_grid=grid_blurring, xp=jnp, + grid=grid_lp, + psf=dataset.psf, + blurring_grid=grid_blurring, + xp=jnp, ) return result.array -_, blurred_img_jit = jit_profile( - blurred_image_from_params, "blurred_image_jit", params_tree -) + +_, blurred_img_jit = jit_profile(blurred_image_from_params, "blurred_image_jit", params_tree) likelihood_steps.append(("Blurred image (PSF convolution)", timer.records[-1][1] / 10)) # --------------------------------------------------------------------------- @@ -465,9 +464,11 @@ def blurred_image_from_params(params_tree): print("\n--- Step 3: Profile-subtracted image ---") + def profile_subtract(data, blurred_image): return data - blurred_image + with timer.section("profile_subtract_eager"): blurred_img_jnp = jnp.array(blurred_image.array) profile_subtracted = profile_subtract(data_array, blurred_img_jnp) @@ -492,9 +493,7 @@ def profile_subtract(data, blurred_image): border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1) # The source plane grid is the last entry (index -1) of the traced grids list. -traced_source_grid = tracer.traced_grid_2d_list_from( - grid=dataset.grids.pixelization, xp=jnp -)[-1] +traced_source_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp)[-1] with timer.section("border_relocation_eager"): relocated_grid = border_relocator.relocated_grid_from(grid=traced_source_grid) @@ -521,13 +520,13 @@ def profile_subtract(data, blurred_image): ) block(mesh_grid) + def overlay_grid_raw_fn(relocated_grid_raw): grid = al.Grid2DIrregular(values=relocated_grid_raw, xp=jnp) return overlay_grid_from(shape_native=mesh_shape, grid=grid, xp=jnp) -_, mesh_grid_raw = jit_profile( - overlay_grid_raw_fn, "overlay_grid_jit", relocated_grid_raw -) + +_, mesh_grid_raw = jit_profile(overlay_grid_raw_fn, "overlay_grid_jit", relocated_grid_raw) likelihood_steps.append(("Overlay grid (source pixel centres)", timer.records[-1][1] / 10)) print(f" mesh_grid shape: {mesh_grid_raw.shape}") @@ -610,6 +609,7 @@ def overlay_grid_raw_fn(relocated_grid_raw): # border relocation → overlay grid → interpolation → mapper → mapping matrix → PSF convolution. # These steps are tightly sequential; the full pipeline JIT-compiles them all together. + def blurred_mm_from_params(params_tree): """Compute blurred mapping matrix via full inversion setup from a pytree ModelInstance. @@ -623,14 +623,18 @@ def blurred_mm_from_params(params_tree): galaxy_name_image_dict={"('galaxies', 'source')": adapt_image}, ) fit_jax = al.FitImaging( - dataset=dataset, tracer=t, adapt_images=adapt_images_jax, + dataset=dataset, + tracer=t, + adapt_images=adapt_images_jax, settings=al.Settings( use_border_relocator=True, use_mixed_precision=_cli.use_mixed_precision, - ), xp=jnp, + ), + xp=jnp, ) return jnp.array(fit_jax.inversion.operated_mapping_matrix) + _, bmm_jit = jit_profile(blurred_mm_from_params, "inversion_setup_jit", params_tree) likelihood_steps.append(("Inversion setup (steps 4-8 combined)", timer.records[-1][1] / 10)) @@ -645,6 +649,7 @@ def blurred_mm_from_params(params_tree): print("\n--- Step 9: Data vector ---") + def compute_data_vector(blurred_mapping_matrix, image, noise_map): return al.util.inversion_imaging.data_vector_via_blurred_mapping_matrix_from( blurred_mapping_matrix=blurred_mapping_matrix, @@ -652,6 +657,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): noise_map=noise_map, ) + profile_sub_jnp = jnp.array(fit.profile_subtracted_image.array) noise_jnp = jnp.array(dataset.noise_map.array) @@ -675,6 +681,7 @@ def compute_data_vector(blurred_mapping_matrix, image, noise_map): # Match the FitImaging inversion: add_to_curvature_diag=True, with settings no_reg_list = list(inversion.no_regularization_index_list) + def compute_curvature_matrix(blurred_mapping_matrix, noise_map): return al.util.inversion.curvature_matrix_via_mapping_matrix_from( mapping_matrix=blurred_mapping_matrix, @@ -685,6 +692,7 @@ def compute_curvature_matrix(blurred_mapping_matrix, noise_map): xp=jnp, ) + with timer.section("curvature_matrix_eager"): curvature_matrix = compute_curvature_matrix(bmm_jnp, noise_jnp) block(curvature_matrix) @@ -702,6 +710,7 @@ def compute_curvature_matrix(blurred_mapping_matrix, noise_map): print("\n--- Step 11: Regularization matrix ---") + def compute_regularization_matrix(neighbors_array, neighbors_sizes): return al.util.regularization.constant_regularization_matrix_from( coefficient=reg_coefficient, @@ -710,15 +719,13 @@ def compute_regularization_matrix(neighbors_array, neighbors_sizes): xp=jnp, ) + with timer.section("regularization_matrix_eager"): - regularization_matrix = compute_regularization_matrix( - neighbors_array, neighbors_sizes - ) + regularization_matrix = compute_regularization_matrix(neighbors_array, neighbors_sizes) block(regularization_matrix) _, regularization_matrix = jit_profile( - compute_regularization_matrix, "regularization_matrix_jit", - neighbors_array, neighbors_sizes + compute_regularization_matrix, "regularization_matrix_jit", neighbors_array, neighbors_sizes ) likelihood_steps.append(("Regularization matrix (H)", timer.records[-1][1] / 10)) @@ -737,6 +744,7 @@ def compute_regularization_matrix(neighbors_array, neighbors_sizes): print("\n--- Step 12: Regularized reconstruction ---") + def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix): curvature_reg_matrix = curvature_matrix + regularization_matrix return al.util.inversion.reconstruction_positive_only_from( @@ -745,6 +753,7 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) xp=jnp, ) + with timer.section("reconstruction_eager"): reconstruction = compute_reconstruction( jnp.array(data_vector), @@ -754,7 +763,8 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) block(reconstruction) _, reconstruction = jit_profile( - compute_reconstruction, "reconstruction_jit", + compute_reconstruction, + "reconstruction_jit", jnp.array(data_vector), jnp.array(curvature_matrix), regularization_matrix_full, @@ -769,9 +779,16 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) print("\n--- Step 13: Mapped reconstruction + log evidence ---") + def compute_log_evidence( - data, noise_map, blurred_image, blurred_mapping_matrix, reconstruction, - curvature_matrix, regularization_matrix, mapper_indices, + data, + noise_map, + blurred_image, + blurred_mapping_matrix, + reconstruction, + curvature_matrix, + regularization_matrix, + mapper_indices, ): """Compute the full log evidence including all five terms: @@ -800,9 +817,7 @@ def compute_log_evidence( chi_squared = jnp.sum((residual / noise_map) ** 2) # Regularization term: s^T H s - regularization_term = jnp.dot( - reconstruction, jnp.dot(regularization_matrix, reconstruction) - ) + regularization_term = jnp.dot(reconstruction, jnp.dot(regularization_matrix, reconstruction)) # Curvature + regularization matrix curvature_reg_matrix = curvature_matrix + regularization_matrix @@ -813,15 +828,11 @@ def compute_log_evidence( creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices] reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices] - log_det_curvature_reg = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced))) - ) - log_det_regularization = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced))) - ) + log_det_curvature_reg = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))) + log_det_regularization = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))) # Noise normalization - noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_map ** 2)) + noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_map**2)) return -0.5 * ( chi_squared @@ -831,6 +842,7 @@ def compute_log_evidence( + noise_normalization ) + # For the JIT profiling we use the step-by-step matrices for timing. # For the correctness assertion we use the inversion's own matrices, because # cumulative floating-point differences between JIT-compiled and eager paths @@ -844,15 +856,28 @@ def compute_log_evidence( with timer.section("log_evidence_eager"): log_evidence = compute_log_evidence( - data_array, noise_jnp, blurred_img_jnp, bmm_jnp, - recon_jnp, curv_jnp, reg_jnp, mapper_indices_jnp, + data_array, + noise_jnp, + blurred_img_jnp, + bmm_jnp, + recon_jnp, + curv_jnp, + reg_jnp, + mapper_indices_jnp, ) block(log_evidence) _, log_evidence = jit_profile( - compute_log_evidence, "log_evidence_jit", - data_array, noise_jnp, blurred_img_jnp, bmm_jnp, - recon_jnp, curv_jnp, reg_jnp, mapper_indices_jnp, + compute_log_evidence, + "log_evidence_jit", + data_array, + noise_jnp, + blurred_img_jnp, + bmm_jnp, + recon_jnp, + curv_jnp, + reg_jnp, + mapper_indices_jnp, ) likelihood_steps.append(("Mapped recon + log evidence", timer.records[-1][1] / 10)) @@ -864,8 +889,14 @@ def compute_log_evidence( inv_curv_jnp = jnp.array(inversion.curvature_matrix) log_evidence_check = compute_log_evidence( - data_array, noise_jnp, blurred_img_jnp, bmm_jnp, - inv_recon_jnp, inv_curv_jnp, reg_jnp, mapper_indices_jnp, + data_array, + noise_jnp, + blurred_img_jnp, + bmm_jnp, + inv_recon_jnp, + inv_curv_jnp, + reg_jnp, + mapper_indices_jnp, ) print(f" log_evidence (inv matrices) = {log_evidence_check}") print(f" log_evidence (reference) = {log_evidence_ref}") @@ -883,7 +914,9 @@ def compute_log_evidence( # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -966,7 +999,7 @@ def compute_log_evidence( fontweight="bold", ) ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | {n_image_pixels} pixels | " + f'AutoLens v{al_version} | {pixel_scale}"/px | {n_image_pixels} pixels | ' f"{n_over_sampled_pixels} over-sampled | {mesh_shape[0]}x{mesh_shape[1]} mesh | " f"total: {step_total:.6f} s", fontsize=9, @@ -983,7 +1016,9 @@ def compute_log_evidence( # Regression assertion — eager log_evidence only # =================================================================== -EXPECTED_LOG_EVIDENCE_HST = 28370.27770182 # 39x39 = 1521 source pixels, MGE-60 lens light, adapt_image=lensed_source +EXPECTED_LOG_EVIDENCE_HST = ( + 28370.27770182 # 39x39 = 1521 source pixels, MGE-60 lens light, adapt_image=lensed_source +) np.testing.assert_allclose( log_evidence_ref, @@ -994,7 +1029,4 @@ def compute_log_evidence( f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})" ), ) -print( - f" Eager regression assertion PASSED: log_evidence matches " - f"{EXPECTED_LOG_EVIDENCE_HST:.6f}" -) +print(f" Eager regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") diff --git a/likelihood_breakdown/interferometer/delaunay.py b/likelihood_breakdown/interferometer/delaunay.py index 2f3309a..5e1b8ba 100644 --- a/likelihood_breakdown/interferometer/delaunay.py +++ b/likelihood_breakdown/interferometer/delaunay.py @@ -41,32 +41,31 @@ """ import os -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path import autofit as af import autolens as al +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -74,12 +73,13 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 + _cli = parse_profile_cli() instrument = "sma" # <-- change this to profile a different instrument @@ -92,6 +92,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -212,18 +213,14 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") print("\n--- Image mesh construction (Hilbert) ---") with timer.section("image_mesh_hilbert"): - image_mesh = al.image_mesh.Hilbert( - pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0 - ) + image_mesh = al.image_mesh.Hilbert(pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0) image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( mask=dataset.real_space_mask, adapt_data=adapt_image ) @@ -368,9 +365,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Step 1: Ray-trace data grid ---") with timer.section("ray_trace_data_eager"): - traced_grids = tracer.traced_grid_2d_list_from( - grid=dataset.grids.pixelization, xp=jnp - ) + traced_grids = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp) for tg in traced_grids: block(tg) @@ -383,9 +378,7 @@ def ray_trace_data_raw(grid_raw): return jnp.stack([tg.array for tg in traced]) -_, traced_data_grids_raw = jit_profile( - ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw -) +_, traced_data_grids_raw = jit_profile(ray_trace_data_raw, "ray_trace_data_jit", grid_pix_raw) likelihood_steps.append(("Ray-trace data grid", timer.records[-1][1] / 10)) print(f" traced_data_grids shape: {traced_data_grids_raw.shape}") @@ -414,9 +407,7 @@ def ray_trace_mesh_raw(mesh_raw): return jnp.stack([tg.array for tg in traced]) -_, traced_mesh_grids_raw = jit_profile( - ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw -) +_, traced_mesh_grids_raw = jit_profile(ray_trace_mesh_raw, "ray_trace_mesh_jit", mesh_grid_raw) likelihood_steps.append(("Ray-trace mesh grid", timer.records[-1][1] / 10)) print(f" traced_mesh_grids shape: {traced_mesh_grids_raw.shape}") @@ -435,9 +426,7 @@ def ray_trace_mesh_raw(mesh_raw): with timer.section("border_relocator_setup"): border_relocator = BorderRelocator(mask=dataset.mask, sub_size=1) -traced_source_grid = tracer.traced_grid_2d_list_from( - grid=dataset.grids.pixelization, xp=jnp -)[-1] +traced_source_grid = tracer.traced_grid_2d_list_from(grid=dataset.grids.pixelization, xp=jnp)[-1] traced_mesh_source = tracer.traced_grid_2d_list_from( grid=al.Grid2DIrregular(image_plane_mesh_grid), xp=jnp )[-1] @@ -553,9 +542,7 @@ def transformed_mm_from_params(params_tree): return jnp.asarray(fit_jax.inversion.operated_mapping_matrix) -_, transformed_mm_jit = jit_profile( - transformed_mm_from_params, "inversion_setup_jit", params_tree -) +_, transformed_mm_jit = jit_profile(transformed_mm_from_params, "inversion_setup_jit", params_tree) likelihood_steps.append( ("Inversion setup (steps 5-8 combined, incl. NUFFT)", timer.records[-1][1] / 10) ) @@ -575,12 +562,16 @@ def transformed_mm_from_params(params_tree): def compute_data_vector( - transformed_mm_real, transformed_mm_imag, data_real, data_imag, - noise_real, noise_imag, + transformed_mm_real, + transformed_mm_imag, + data_real, + data_imag, + noise_real, + noise_imag, ): # Visibility-space data vector: D_i = sum_j f_ij d_j / sigma_j^2 (real + imag). - weighted_data_real = data_real / (noise_real ** 2) - weighted_data_imag = data_imag / (noise_imag ** 2) + weighted_data_real = data_real / (noise_real**2) + weighted_data_imag = data_imag / (noise_imag**2) return jnp.matmul(transformed_mm_real.T, weighted_data_real) + jnp.matmul( transformed_mm_imag.T, weighted_data_imag ) @@ -588,15 +579,24 @@ def compute_data_vector( with timer.section("data_vector_eager"): data_vector = compute_data_vector( - transformed_mm_real_jnp, transformed_mm_imag_jnp, - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) block(data_vector) _, data_vector = jit_profile( - compute_data_vector, "data_vector_jit", - transformed_mm_real_jnp, transformed_mm_imag_jnp, - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, + compute_data_vector, + "data_vector_jit", + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) likelihood_steps.append(("Data vector (D)", timer.records[-1][1] / 10)) @@ -612,7 +612,10 @@ def compute_data_vector( def compute_curvature_matrix( - transformed_mm_real, transformed_mm_imag, noise_real, noise_imag, + transformed_mm_real, + transformed_mm_imag, + noise_real, + noise_imag, ): real_curv = al.util.inversion.curvature_matrix_via_mapping_matrix_from( mapping_matrix=transformed_mm_real, @@ -635,13 +638,20 @@ def compute_curvature_matrix( with timer.section("curvature_matrix_eager"): curvature_matrix = compute_curvature_matrix( - transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) block(curvature_matrix) _, curvature_matrix = jit_profile( - compute_curvature_matrix, "curvature_matrix_jit", - transformed_mm_real_jnp, transformed_mm_imag_jnp, noise_real_jnp, noise_imag_jnp, + compute_curvature_matrix, + "curvature_matrix_jit", + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + noise_real_jnp, + noise_imag_jnp, ) likelihood_steps.append(("Curvature matrix (F)", timer.records[-1][1] / 10)) @@ -690,8 +700,11 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) block(reconstruction) _, reconstruction = jit_profile( - compute_reconstruction, "reconstruction_jit", - jnp.array(data_vector), jnp.array(curvature_matrix), jnp.array(regularization_matrix), + compute_reconstruction, + "reconstruction_jit", + jnp.array(data_vector), + jnp.array(curvature_matrix), + jnp.array(regularization_matrix), ) likelihood_steps.append(("Regularized reconstruction", timer.records[-1][1] / 10)) @@ -705,9 +718,16 @@ def compute_reconstruction(data_vector, curvature_matrix, regularization_matrix) def compute_log_evidence( - data_real, data_imag, noise_real, noise_imag, - transformed_mm_real, transformed_mm_imag, - reconstruction, curvature_matrix, regularization_matrix, mapper_indices, + data_real, + data_imag, + noise_real, + noise_imag, + transformed_mm_real, + transformed_mm_imag, + reconstruction, + curvature_matrix, + regularization_matrix, + mapper_indices, ): """Visibility-space log-evidence — five-term formula matching the production ``FitInterferometer.log_evidence``. @@ -724,30 +744,26 @@ def compute_log_evidence( chi_squared = chi_real + chi_imag # s^T H s - regularization_term = jnp.dot( - reconstruction, jnp.dot(regularization_matrix, reconstruction) - ) + regularization_term = jnp.dot(reconstruction, jnp.dot(regularization_matrix, reconstruction)) # Complexity terms (Cholesky log-det matching production) curvature_reg_matrix = curvature_matrix + regularization_matrix creg_reduced = curvature_reg_matrix[mapper_indices][:, mapper_indices] reg_reduced = regularization_matrix[mapper_indices][:, mapper_indices] - log_det_curvature_reg = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced))) - ) - log_det_regularization = 2.0 * jnp.sum( - jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced))) - ) + log_det_curvature_reg = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(creg_reduced)))) + log_det_regularization = 2.0 * jnp.sum(jnp.log(jnp.diag(jnp.linalg.cholesky(reg_reduced)))) # Noise normalisation (real + imag) - noise_normalization = ( - jnp.sum(jnp.log(2 * jnp.pi * noise_real ** 2)) - + jnp.sum(jnp.log(2 * jnp.pi * noise_imag ** 2)) + noise_normalization = jnp.sum(jnp.log(2 * jnp.pi * noise_real**2)) + jnp.sum( + jnp.log(2 * jnp.pi * noise_imag**2) ) return -0.5 * ( - chi_squared + regularization_term + log_det_curvature_reg - - log_det_regularization + noise_normalization + chi_squared + + regularization_term + + log_det_curvature_reg + - log_det_regularization + + noise_normalization ) @@ -761,17 +777,32 @@ def compute_log_evidence( with timer.section("log_evidence_eager"): log_evidence = compute_log_evidence( - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - transformed_mm_real_jnp, transformed_mm_imag_jnp, - reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + reconstruction, + curvature_matrix, + reg_jnp, + mapper_indices_jnp, ) block(log_evidence) _, log_evidence = jit_profile( - compute_log_evidence, "log_evidence_jit", - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - transformed_mm_real_jnp, transformed_mm_imag_jnp, - reconstruction, curvature_matrix, reg_jnp, mapper_indices_jnp, + compute_log_evidence, + "log_evidence_jit", + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + reconstruction, + curvature_matrix, + reg_jnp, + mapper_indices_jnp, ) likelihood_steps.append(("Mapped recon + log evidence", timer.records[-1][1] / 10)) @@ -779,9 +810,16 @@ def compute_log_evidence( # Correctness check: use the inversion's own reconstruction and curvature matrix log_evidence_check = compute_log_evidence( - data_real_jnp, data_imag_jnp, noise_real_jnp, noise_imag_jnp, - transformed_mm_real_jnp, transformed_mm_imag_jnp, - inv_recon_jnp, inv_curv_jnp, reg_jnp, mapper_indices_jnp, + data_real_jnp, + data_imag_jnp, + noise_real_jnp, + noise_imag_jnp, + transformed_mm_real_jnp, + transformed_mm_imag_jnp, + inv_recon_jnp, + inv_curv_jnp, + reg_jnp, + mapper_indices_jnp, ) print(f" log_evidence (inv matrices) = {log_evidence_check}") print(f" log_evidence (reference) = {figure_of_merit_ref}") @@ -795,10 +833,7 @@ def compute_log_evidence( "FitInterferometer.log_evidence" ), ) -print( - " Assertion PASSED: inversion-matrix log_evidence matches " - "FitInterferometer.log_evidence" -) +print(" Assertion PASSED: inversion-matrix log_evidence matches FitInterferometer.log_evidence") # =================================================================== @@ -806,7 +841,9 @@ def compute_log_evidence( # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -893,7 +930,7 @@ def compute_log_evidence( fontweight="bold", ) ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | " + f'AutoLens v{al_version} | {pixel_scale}"/px | ' f"{real_space_shape[0]}x{real_space_shape[1]} real-space | " f"{n_visibilities} visibilities | {n_mesh_vertices} Delaunay verts | " f"total: {step_total:.6f} s", @@ -935,7 +972,4 @@ def compute_log_evidence( f"drifted (got {figure_of_merit_ref}, expected {expected_log_evidence})" ), ) - print( - f" Eager regression assertion PASSED: log_evidence matches " - f"{expected_log_evidence:.6f}" - ) + print(f" Eager regression assertion PASSED: log_evidence matches {expected_log_evidence:.6f}") diff --git a/likelihood_runtime/README.md b/likelihood_runtime/README.md index fc2456f..92e1c7b 100644 --- a/likelihood_runtime/README.md +++ b/likelihood_runtime/README.md @@ -10,6 +10,14 @@ For *where the time goes inside the likelihood*, use the sibling package [`likel The empirical findings from previous sweeps — per-cell timings, mp verdicts, the GPU-NUFFT regression, the upstream blockers — live in [`OPTIMIZATION_NOTES.md`](OPTIMIZATION_NOTES.md) in this directory. +## Latest results + + +_No data yet — run `likelihood_runtime/sweep.py` then `aggregate.py` to populate._ + + +Auto-generated by `scripts/build_readme.py` from `results/runtime/**/comparison.json`; baseline columns appear once `results/baselines/` is populated. + ## Methodology Each script measures **one** quantity per run: the steady-state cost of the entire likelihood as a single JIT-compiled JAX program. There is no per-step decomposition. The measurement is: @@ -30,7 +38,7 @@ The sweep harness drives every in-scope cell through this matrix: | `local_cpu_mp` | CPU | mixed (fp32 inversion) | same + `--use-mixed-precision` | | `local_gpu_fp64` | RTX 2060 (consumer) | fp64 | `JAX_PLATFORM_NAME=cuda JAX_PLATFORMS=cuda,cpu` | | `local_gpu_mp` | RTX 2060 | mixed | same + `--use-mixed-precision` | -| `hpc_a100_fp64` | A100 (80 GB) | fp64 | SLURM-dispatched via `z_projects/profiling/hpc/sync` | +| `hpc_a100_fp64` | A100 (80 GB) | fp64 | SLURM-dispatched via the submit scripts under [`hpc/`](../hpc/README.md) | | `hpc_a100_mp` | A100 | mixed | same + `--use-mixed-precision` | The `cuda,cpu` listing on GPU configs is load-bearing: the Delaunay + datacube paths use `jax.pure_callback` for Hilbert-curve mesh generation, which needs a CPU device available even when the primary platform is CUDA. Without the trailing `cpu` the callback raises `pure_callback failed to find a local CPU device`. diff --git a/likelihood_runtime/aggregate.py b/likelihood_runtime/aggregate.py index f71a837..70a105f 100644 --- a/likelihood_runtime/aggregate.py +++ b/likelihood_runtime/aggregate.py @@ -1,11 +1,11 @@ """Aggregate per-config JSONs for a swept likelihood cell into comparison.{json,png}. -Reads every ``.json`` under a cell's output dir (see -``sweep.py``) and produces a single ``comparison.json`` whose schema -mirrors the existing -``autolens_workspace_developer/jax_profiling/results/jit/imaging/{mge, -pixelization,delaunay}/comparison.json`` artifacts so the existing -readers (and the OPTIMIZATION_NOTES doc) continue to work. +Reads every ``[_sparse].json`` under a cell's output dir (see +``sweep.py``; default root is ``results/runtime/`` in this repo) and +produces a single ``comparison.json`` whose schema mirrors the historical +``autolens_workspace_developer/jax_profiling/results/jit`` artifacts so the +existing readers (and the OPTIMIZATION_NOTES doc) continue to work. +``_sparse`` rows order after the canonical configs. The ``comparison.png`` is a log-scale grouped bar chart: one bar per (step, config), sorted by step cost on the slowest config. The @@ -34,10 +34,8 @@ import matplotlib.pyplot as plt import numpy as np - _REPO_ROOT = Path(__file__).resolve().parents[1] -_WT_ROOT = _REPO_ROOT.parent -_DEFAULT_OUTPUT_ROOT = _WT_ROOT / "autolens_workspace_developer" / "jax_profiling" / "results" / "jit" +_DEFAULT_OUTPUT_ROOT = _REPO_ROOT / "results" / "runtime" # Stable ordering — keep the same row order as sweep_likelihood + the prior @@ -85,10 +83,7 @@ def _discover_cells(output_root: Path) -> list[tuple[str, ...]]: return cells def _has_config_json(d: Path) -> bool: - return any( - p.stem in _CONFIG_ORDER or p.stem.endswith("_pre_fix") - for p in d.glob("*.json") - ) + return any(_is_config_stem(p.stem) for p in d.glob("*.json")) for cls_dir in sorted(output_root.iterdir()): if not cls_dir.is_dir(): @@ -106,6 +101,20 @@ def _has_config_json(d: Path) -> bool: return cells +def _is_config_stem(stem: str) -> bool: + """True for sweep-config JSONs (``local_gpu_mp``, ``hpc_a100_fp64_sparse``, …). + + Cell dirs may also hold versioned standalone summaries + (``mge_likelihood_summary_hst_v….json``) — those are not sweep rows and + must not enter ``comparison.json``. + """ + return ( + stem in _CONFIG_ORDER + or stem.removesuffix("_sparse") in _CONFIG_ORDER + or stem.endswith("_pre_fix") + ) + + def _read_config(json_path: Path) -> dict: data = json.loads(json_path.read_text()) # Normalise field names so downstream rendering is uniform whether the @@ -130,7 +139,7 @@ def _read_config(json_path: Path) -> dict: def _aggregate_cell(cell_dir: Path) -> dict: configs: dict[str, dict] = {} for json_path in sorted(cell_dir.glob("*.json")): - if json_path.name == "comparison.json": + if not _is_config_stem(json_path.stem): continue try: configs[json_path.stem] = _read_config(json_path) @@ -204,7 +213,8 @@ def _render_png(comparison: dict, cell_id: str, png_path: Path) -> None: n_steps, n_cfgs = len(step_union), len(config_names) # Drop configs that have no positive step data — log scale can't render NaN/0. config_names = [ - c for c in config_names + c + for c in config_names if any( isinstance(v, (int, float)) and np.isfinite(v) and v > 0 for v in configs[c].get("steps", {}).values() diff --git a/likelihood_runtime/datacube/delaunay.py b/likelihood_runtime/datacube/delaunay.py index 67d0d30..b3061c1 100644 --- a/likelihood_runtime/datacube/delaunay.py +++ b/likelihood_runtime/datacube/delaunay.py @@ -87,34 +87,33 @@ the shared-``Lᵀ W̃ L`` optimisation. """ -import numpy as np -import jax -import jax.numpy as jnp import os -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -122,16 +121,19 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 -from vram import vmap_batch_for, write_probe_json, ProbeResult # noqa: E402 +from vram import ProbeResult, vmap_batch_for, write_probe_json # noqa: E402 + _cli = parse_profile_cli() -instrument = _cli.instrument or "sma" # default; override via --instrument (cube is N copies of the per-instrument dataset) +instrument = ( + _cli.instrument or "sma" +) # default; override via --instrument (cube is N copies of the per-instrument dataset) # n_channels = 34 matches the prior Hannah ALMA cube fiducial. For quick # iteration on the smaller sma dataset, drop this to 4. @@ -144,6 +146,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -267,18 +270,14 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset_list[0] - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset_list[0]) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") print("\n--- Image mesh construction (Hilbert) ---") with timer.section("image_mesh_hilbert"): - image_mesh = al.image_mesh.Hilbert( - pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0 - ) + image_mesh = al.image_mesh.Hilbert(pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0) image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( mask=dataset_list[0].real_space_mask, adapt_data=adapt_image ) @@ -405,18 +404,23 @@ def _build_transformer(uv_wavelengths, real_space_mask): if _cli.vmap_probe: probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "datacube")) - / "vmap_probe.json" - ) + _cli.output_dir or (_workspace_root / "results" / "runtime" / "datacube" / "delaunay") + ) / "vmap_probe.json" probe_path.parent.mkdir(parents=True, exist_ok=True) import json - probe_path.write_text(json.dumps({ - "dataset": "datacube", - "model": "delaunay", - "instrument": instrument, - "recommended_batch_size": None, - "note": "datacube vmap intentionally skipped — natural batching axis is channels, not parameters", - }, indent=2)) + + probe_path.write_text( + json.dumps( + { + "dataset": "datacube", + "model": "delaunay", + "instrument": instrument, + "recommended_batch_size": None, + "note": "datacube vmap intentionally skipped — natural batching axis is channels, not parameters", + }, + indent=2, + ) + ) print(f" vmap_probe: cell intentionally skipped — wrote {probe_path}") sys.exit(0) @@ -507,7 +511,9 @@ def full_cube_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -529,7 +535,7 @@ def full_cube_pipeline_from_params(params_tree): if full_cube_result is not None: print(f" Cube JIT log_evidence: {float(full_cube_result)}") else: - print(f" Cube JIT log_evidence: SKIPPED (CUBE_FULL_JIT=1 to enable)") + print(" Cube JIT log_evidence: SKIPPED (CUBE_FULL_JIT=1 to enable)") print("-" * 70) # Shared-Lᵀ W̃ L savings estimate was a per-step-breakdown deliverable — @@ -565,9 +571,7 @@ def full_cube_pipeline_from_params(params_tree): "regularization_coefficient": regularization_coefficient, }, "cube_log_evidence_eager": cube_log_evidence_ref, - "cube_log_evidence_jit": ( - float(full_cube_result) if full_cube_result is not None else None - ), + "cube_log_evidence_jit": (float(full_cube_result) if full_cube_result is not None else None), "log_evidence_per_channel_eager": [float(le) for le in log_evidence_per_channel], "full_pipeline_cube_single_jit": full_pipeline_per_call, "shared_lwl_savings_estimate": shared_lwl_savings, @@ -576,7 +580,7 @@ def full_cube_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "datacube", + default_dir=_workspace_root / "results" / "runtime" / "datacube" / "delaunay", default_basename=f"delaunay_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -598,9 +602,7 @@ def full_cube_pipeline_from_params(params_tree): } _per_channel = EXPECTED_LOG_EVIDENCE_PER_CHANNEL.get(instrument) -expected_cube_log_evidence = ( - n_channels * _per_channel if _per_channel is not None else None -) +expected_cube_log_evidence = n_channels * _per_channel if _per_channel is not None else None if expected_cube_log_evidence is None: print( @@ -631,4 +633,4 @@ def full_cube_pipeline_from_params(params_tree): rtol=1e-3, err_msg=f"datacube/delaunay[{instrument}]: regression — full cube log_evidence drifted", ) - print(f" Full-pipeline cube regression assertion PASSED") + print(" Full-pipeline cube regression assertion PASSED") diff --git a/likelihood_runtime/datacube/shared_preloads.py b/likelihood_runtime/datacube/shared_preloads.py index e9c440c..fb552dc 100644 --- a/likelihood_runtime/datacube/shared_preloads.py +++ b/likelihood_runtime/datacube/shared_preloads.py @@ -34,11 +34,10 @@ import time from pathlib import Path -import numpy as np -import jax.numpy as jnp - 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])) diff --git a/likelihood_runtime/imaging/delaunay.py b/likelihood_runtime/imaging/delaunay.py index 008df4b..9d78abd 100644 --- a/likelihood_runtime/imaging/delaunay.py +++ b/likelihood_runtime/imaging/delaunay.py @@ -41,33 +41,32 @@ """ import json -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -75,10 +74,10 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.imaging import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -87,6 +86,7 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() instrument = _cli.instrument or "hst" # default; override via --instrument @@ -96,6 +96,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -241,9 +242,7 @@ def jit_profile(func, label, *args, n_repeats=10): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") @@ -252,9 +251,7 @@ def jit_profile(func, label, *args, n_repeats=10): n_mesh_vertices = 1500 # 1500-tier production fiducial with timer.section("image_mesh_hilbert"): - image_mesh = al.image_mesh.Hilbert( - pixels=n_mesh_vertices, weight_power=1.0, weight_floor=0.0 - ) + 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 ) @@ -294,9 +291,7 @@ def jit_profile(func, label, *args, n_repeats=10): shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005) shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005) - lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear - ) + lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) mesh = al.mesh.Delaunay( pixels=n_mesh_vertices, @@ -396,9 +391,11 @@ def jit_profile(func, label, *args, n_repeats=10): analysis = al.AnalysisImaging(dataset=dataset, adapt_images=adapt_images, use_jax=True) + def full_pipeline_from_params(params_tree): return analysis.log_likelihood_function(instance=params_tree) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 @@ -426,7 +423,7 @@ def full_pipeline_from_params(params_tree): } _early_dict_path, _ = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "imaging", + default_dir=_workspace_root / "results" / "runtime" / "imaging" / "delaunay", default_basename=f"delaunay_likelihood_summary_{instrument}_v{al.__version__}", ) _early_dict_path.write_text(json.dumps(_early_summary, indent=2)) @@ -449,14 +446,11 @@ def full_pipeline_from_params(params_tree): recommended = recommend_batch_size(probe) _inversion_path = "sparse" if _cli.use_sparse_operator else "dense" _probe_basename = ( - "vmap_probe_delaunay_sparse" - if _cli.use_sparse_operator - else "vmap_probe_delaunay" + "vmap_probe_delaunay_sparse" if _cli.use_sparse_operator else "vmap_probe_delaunay" ) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "imaging")) - / f"{_probe_basename}.json" - ) + _cli.output_dir or (_workspace_root / "results" / "runtime" / "imaging" / "delaunay") + ) / f"{_probe_basename}.json" write_probe_json( probe, recommended, @@ -546,7 +540,9 @@ def full_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -589,7 +585,9 @@ def full_pipeline_from_params(params_tree): "inversion_path": "sparse" if _cli.use_sparse_operator else "dense", }, "full_pipeline_single_jit": full_pipeline_per_call, - "vmap": "SKIPPED — vmap_batch_for() returned None for this (cell, instrument)" if _vmap_skipped else { + "vmap": "SKIPPED — vmap_batch_for() returned None for this (cell, instrument)" + if _vmap_skipped + else { "batch_size": batch_size, "batch_time": vmap_batch_time, "per_call": vmap_per_call, @@ -599,7 +597,7 @@ def full_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "imaging", + default_dir=_workspace_root / "results" / "runtime" / "imaging" / "delaunay", default_basename=f"delaunay_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -617,7 +615,9 @@ def full_pipeline_from_params(params_tree): # matches imaging/pixelization (adaptive meshes amplify fp drift through # Cholesky / log_det). vmap result asserted only when DELAUNAY_VMAP=1 # (vmap compile takes 20+ min). -EXPECTED_LOG_EVIDENCE_HST = 29110.92085793 # 1500-pixel Hilbert/Delaunay, MGE-60 lens, adapt_image=lensed_source +EXPECTED_LOG_EVIDENCE_HST = ( + 29110.92085793 # 1500-pixel Hilbert/Delaunay, MGE-60 lens, adapt_image=lensed_source +) np.testing.assert_allclose( log_evidence_ref, @@ -628,10 +628,7 @@ def full_pipeline_from_params(params_tree): f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})" ), ) -print( - f" Eager regression assertion PASSED: log_evidence matches " - f"{EXPECTED_LOG_EVIDENCE_HST:.6f}" -) +print(f" Eager regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") np.testing.assert_allclose( float(full_result), EXPECTED_LOG_EVIDENCE_HST, diff --git a/likelihood_runtime/imaging/mge.py b/likelihood_runtime/imaging/mge.py index add62fe..9d6c519 100644 --- a/likelihood_runtime/imaging/mge.py +++ b/likelihood_runtime/imaging/mge.py @@ -50,30 +50,29 @@ """ import json -import numpy as np -import jax -import jax.numpy as jnp -import time -import subprocess -import sys -from pathlib import Path -from contextlib import contextmanager - -import autofit as af -import autolens as al -import autoarray as aa -from autofit.jax import register_model as _register_model_pytrees # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os +import subprocess +import sys import sys as _smoke_sys +import time +from contextlib import contextmanager +from pathlib import Path + +import autoarray as aa +import autofit as af +import autolens as al +import jax +import jax.numpy as jnp +import numpy as np +from autofit.jax import register_model as _register_model_pytrees + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -82,10 +81,10 @@ # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.imaging import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -94,6 +93,7 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() instrument = _cli.instrument or "hst" # default; override via --instrument @@ -103,6 +103,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -264,9 +265,7 @@ def jit_profile(func, label, *args, n_repeats=10): shear.gamma_1 = 0.05 shear.gamma_2 = 0.05 - lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear - ) + lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) source_bulge = al.model_util.mge_model_from( mask_radius=mask_radius, total_gaussians=20, centre_prior_is_uniform=False @@ -354,6 +353,7 @@ def jit_profile(func, label, *args, n_repeats=10): # instead of going through ``model.instance_from_vector(parameters, xp=jnp)``. analysis = al.AnalysisImaging(dataset=dataset, use_jax=True) + def full_pipeline_from_params(params_tree): """Full likelihood from a pytree-shaped ``ModelInstance``. @@ -363,6 +363,7 @@ def full_pipeline_from_params(params_tree): """ return analysis.log_likelihood_function(instance=params_tree) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 @@ -389,7 +390,7 @@ def full_pipeline_from_params(params_tree): } _early_dict_path, _ = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "imaging", + default_dir=_workspace_root / "results" / "runtime" / "imaging" / "mge", default_basename=f"mge_likelihood_summary_{instrument}_v{al.__version__}", ) _early_dict_path.write_text(json.dumps(_early_summary, indent=2)) @@ -413,13 +414,10 @@ def full_pipeline_from_params(params_tree): ) recommended = recommend_batch_size(probe) _inversion_path = "sparse" if _cli.use_sparse_operator else "dense" - _probe_basename = ( - "vmap_probe_mge_sparse" if _cli.use_sparse_operator else "vmap_probe_mge" - ) + _probe_basename = "vmap_probe_mge_sparse" if _cli.use_sparse_operator else "vmap_probe_mge" probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "imaging")) - / f"{_probe_basename}.json" - ) + _cli.output_dir or (_workspace_root / "results" / "runtime" / "imaging" / "mge") + ) / f"{_probe_basename}.json" write_probe_json( probe, recommended, @@ -502,7 +500,9 @@ def full_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -520,7 +520,7 @@ def full_pipeline_from_params(params_tree): print("-" * 70) print(f" {'Full pipeline (single JIT)':<30} {full_pipeline_per_call:>12.6f} s") print(f" {f'vmap batch={batch_size} (per call)':<30} {vmap_per_call:>12.6f} s") -print(f" {f'vmap speedup vs single JIT':<30} {vmap_speedup:>11.1f}x") +print(f" {'vmap speedup vs single JIT':<30} {vmap_speedup:>11.1f}x") print("=" * 70) # --- Save results dictionary --- @@ -552,7 +552,7 @@ def full_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "imaging", + default_dir=_workspace_root / "results" / "runtime" / "imaging" / "mge", default_basename=f"mge_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -580,8 +580,7 @@ def full_pipeline_from_params(params_tree): ), ) print( - f" Eager regression assertion PASSED: log_likelihood matches " - f"{EXPECTED_LOG_LIKELIHOOD_HST:.6f}" + f" Eager regression assertion PASSED: log_likelihood matches {EXPECTED_LOG_LIKELIHOOD_HST:.6f}" ) np.testing.assert_allclose( float(full_result), diff --git a/likelihood_runtime/imaging/pixelization.py b/likelihood_runtime/imaging/pixelization.py index 631e455..57715b4 100644 --- a/likelihood_runtime/imaging/pixelization.py +++ b/likelihood_runtime/imaging/pixelization.py @@ -32,35 +32,34 @@ """ import json -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees # Shared adapt-image loader: load or compute+cache `lensed_source.fits` # next to the dataset, then return the masked ``aa.Array2D``. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -68,10 +67,10 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.imaging import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -80,6 +79,7 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() instrument = _cli.instrument or "hst" # default; override via --instrument @@ -89,6 +89,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -253,18 +254,14 @@ def jit_profile(func, label, *args, n_repeats=10): shear.gamma_1 = af.GaussianPrior(mean=0.05, sigma=0.005) shear.gamma_2 = af.GaussianPrior(mean=0.05, sigma=0.005) - lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear - ) + lens = af.Model(al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear) # ``RectangularAdaptImage`` weights mesh pixels by the lensed-source # adapt image — the production-grade alternative to the coordinate- # density-only ``RectangularAdaptDensity``. Adapt image is loaded / # cached below; the same shape and regularization are kept. pixelization = al.Pixelization( - mesh=al.mesh.RectangularAdaptImage( - shape=mesh_shape, weight_power=1.0, weight_floor=0.0 - ), + mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape, weight_power=1.0, weight_floor=0.0), regularization=al.reg.Constant(coefficient=1.0), ) @@ -322,9 +319,7 @@ def jit_profile(func, label, *args, n_repeats=10): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) # ``galaxy_image_dict`` (Galaxy-object-keyed) feeds the eager-path # ``image_for_galaxy`` lookup; ``galaxy_name_image_dict`` (path-tuple # str-keyed) is rebuilt inside JIT closures where the Galaxy objects @@ -378,9 +373,11 @@ def jit_profile(func, label, *args, n_repeats=10): use_jax=True, ) + def full_pipeline_from_params(params_tree): return analysis.log_likelihood_function(instance=params_tree) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 @@ -413,7 +410,7 @@ def full_pipeline_from_params(params_tree): } _early_dict_path, _ = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "imaging", + default_dir=_workspace_root / "results" / "runtime" / "imaging" / "pixelization", default_basename=f"pixelization_likelihood_summary_{instrument}_v{al.__version__}", ) _early_dict_path.write_text(json.dumps(_early_summary, indent=2)) @@ -441,14 +438,11 @@ def full_pipeline_from_params(params_tree): recommended = recommend_batch_size(probe) _inversion_path = "sparse" if _cli.use_sparse_operator else "dense" _probe_basename = ( - "vmap_probe_pixelization_sparse" - if _cli.use_sparse_operator - else "vmap_probe_pixelization" + "vmap_probe_pixelization_sparse" if _cli.use_sparse_operator else "vmap_probe_pixelization" ) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "imaging")) - / f"{_probe_basename}.json" - ) + _cli.output_dir or (_workspace_root / "results" / "runtime" / "imaging" / "pixelization") + ) / f"{_probe_basename}.json" write_probe_json( probe, recommended, @@ -482,8 +476,7 @@ def full_pipeline_from_params(params_tree): _vmap_skipped_reason = None if _n_leaves == 0: _vmap_skipped_reason = ( - "model has 0 free parameters (all fixed to truth); vmap " - "requires at least one array leaf." + "model has 0 free parameters (all fixed to truth); vmap requires at least one array leaf." ) else: parameters = jax.tree_util.tree_map( @@ -563,7 +556,9 @@ def full_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -606,7 +601,9 @@ def full_pipeline_from_params(params_tree): "inversion_path": "sparse" if _cli.use_sparse_operator else "dense", }, "full_pipeline_single_jit": full_pipeline_per_call, - "vmap": "SKIPPED — model has 0 free parameters (all fixed to truth)" if vmap_per_call is None else { + "vmap": "SKIPPED — model has 0 free parameters (all fixed to truth)" + if vmap_per_call is None + else { "batch_size": batch_size, "batch_time": vmap_batch_time, "per_call": vmap_per_call, @@ -616,7 +613,7 @@ def full_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "imaging", + default_dir=_workspace_root / "results" / "runtime" / "imaging" / "pixelization", default_basename=f"pixelization_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -635,7 +632,9 @@ def full_pipeline_from_params(params_tree): # 1581x1581 mapping matrix relative to the non-adaptive baseline (which # previously matched at 1e-4). The 1e-3 envelope is still tight enough to # catch real numerical regressions while accommodating the adaptive path. -EXPECTED_LOG_EVIDENCE_HST = 28370.27770182 # 39x39 = 1521 source pixels, MGE-60 lens light, adapt_image=lensed_source +EXPECTED_LOG_EVIDENCE_HST = ( + 28370.27770182 # 39x39 = 1521 source pixels, MGE-60 lens light, adapt_image=lensed_source +) np.testing.assert_allclose( log_evidence_ref, @@ -646,10 +645,7 @@ def full_pipeline_from_params(params_tree): f"(got {log_evidence_ref}, expected {EXPECTED_LOG_EVIDENCE_HST})" ), ) -print( - f" Eager regression assertion PASSED: log_evidence matches " - f"{EXPECTED_LOG_EVIDENCE_HST:.6f}" -) +print(f" Eager regression assertion PASSED: log_evidence matches {EXPECTED_LOG_EVIDENCE_HST:.6f}") np.testing.assert_allclose( float(full_result), EXPECTED_LOG_EVIDENCE_HST, diff --git a/likelihood_runtime/interferometer/delaunay.py b/likelihood_runtime/interferometer/delaunay.py index c30579b..fc108e4 100644 --- a/likelihood_runtime/interferometer/delaunay.py +++ b/likelihood_runtime/interferometer/delaunay.py @@ -83,32 +83,31 @@ pytree support landed in PyAutoFit#1222. """ -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path import autofit as af import autolens as al +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -116,10 +115,10 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -128,6 +127,7 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() instrument = _cli.instrument or "sma" # default; override via --instrument @@ -140,6 +140,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -264,18 +265,14 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") print("\n--- Image mesh construction (Hilbert) ---") with timer.section("image_mesh_hilbert"): - image_mesh = al.image_mesh.Hilbert( - pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0 - ) + image_mesh = al.image_mesh.Hilbert(pixels=hilbert_pixels, weight_power=1.0, weight_floor=0.0) image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( mask=dataset.real_space_mask, adapt_data=adapt_image ) @@ -409,6 +406,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): use_jax=True, ) + def full_pipeline_from_params(params_tree): """Full interferometer likelihood from a pytree-shaped ``ModelInstance``. @@ -418,6 +416,7 @@ def full_pipeline_from_params(params_tree): """ return analysis.log_likelihood_function(instance=params_tree) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 @@ -455,9 +454,8 @@ def full_pipeline_from_params(params_tree): ) recommended = recommend_batch_size(probe) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "interferometer")) - / "vmap_probe.json" - ) + _cli.output_dir or (_workspace_root / "results" / "runtime" / "interferometer" / "delaunay") + ) / "vmap_probe.json" write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -487,8 +485,10 @@ def full_pipeline_from_params(params_tree): if _vmap_skipped: print(" SKIPPED: vmap_batch_for() returned None for this (cell, instrument).") elif _n_leaves == 0: - print(f" SKIPPED: model has 0 free parameters (all fixed to truth); " - f"vmap requires at least one array leaf.") + print( + " SKIPPED: model has 0 free parameters (all fixed to truth); " + "vmap requires at least one array leaf." + ) else: parameters = jax.tree_util.tree_map( lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)), @@ -551,7 +551,9 @@ def full_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -617,7 +619,9 @@ def full_pipeline_from_params(params_tree): "log_evidence_jit": float(full_result), "full_pipeline_single_jit": full_pipeline_per_call, "vmap": vmap_payload, - "memory_mb": None if memory_analysis is None else { + "memory_mb": None + if memory_analysis is None + else { "output": memory_analysis.output_size_in_bytes / 1024**2, "temp": memory_analysis.temp_size_in_bytes / 1024**2, }, @@ -625,7 +629,7 @@ def full_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "interferometer", + default_dir=_workspace_root / "results" / "runtime" / "interferometer" / "delaunay", default_basename=f"delaunay_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -667,17 +671,14 @@ def full_pipeline_from_params(params_tree): f"drifted (got {figure_of_merit_ref}, expected {expected_log_evidence})" ), ) - print( - f" Eager regression assertion PASSED: log_evidence matches " - f"{expected_log_evidence:.6f}" - ) + print(f" Eager regression assertion PASSED: log_evidence matches {expected_log_evidence:.6f}") np.testing.assert_allclose( float(full_result), expected_log_evidence, rtol=1e-3, err_msg=f"interferometer/delaunay[{instrument}]: regression — full log_evidence drifted", ) - print(f" Full-pipeline regression assertion PASSED") + print(" Full-pipeline regression assertion PASSED") if result_vmap is not None: np.testing.assert_allclose( np.array(result_vmap), @@ -685,4 +686,4 @@ def full_pipeline_from_params(params_tree): rtol=1e-3, err_msg=f"interferometer/delaunay[{instrument}]: regression — vmap log_evidence drifted", ) - print(f" vmap regression assertion PASSED") + print(" vmap regression assertion PASSED") diff --git a/likelihood_runtime/interferometer/mge.py b/likelihood_runtime/interferometer/mge.py index 3dd9b2d..d55e5f7 100644 --- a/likelihood_runtime/interferometer/mge.py +++ b/likelihood_runtime/interferometer/mge.py @@ -48,29 +48,27 @@ support landed in PyAutoFit#1222. """ -import numpy as np -import jax -import jax.numpy as jnp -import time -import subprocess -import sys -from pathlib import Path -from contextlib import contextmanager - -import autofit as af -import autolens as al -from autofit.jax import register_model as _register_model_pytrees - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os +import subprocess +import sys import sys as _smoke_sys +import time +from contextlib import contextmanager +from pathlib import Path + +import autofit as af +import autolens as al +import jax +import jax.numpy as jnp +import numpy as np +from autofit.jax import register_model as _register_model_pytrees + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -80,10 +78,10 @@ # historical DFT baseline on SMA. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -92,9 +90,11 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() import argparse as _argparse # noqa: E402 + _local_parser = _argparse.ArgumentParser(add_help=False, allow_abbrev=False) _local_parser.add_argument("--use-dft", action="store_true") _local_args, _ = _local_parser.parse_known_args() @@ -107,6 +107,7 @@ # Profiling helpers (copied verbatim from imaging/mge.py) # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -268,6 +269,7 @@ def jit_profile(func, label, *args, n_repeats=10): # --------------------------------------------------------------------------- from autogalaxy.profiles.basis import Basis as _Basis + _basis_list = [b for g in instance.galaxies for b in g.cls_list_from(cls=_Basis)] n_linear_gaussians = sum(len(b.profile_list) for b in _basis_list) @@ -313,6 +315,7 @@ def jit_profile(func, label, *args, n_repeats=10): use_jax=True, ) + def full_pipeline_from_params(params_tree): """Full interferometer likelihood from a pytree-shaped ``ModelInstance``. @@ -322,6 +325,7 @@ def full_pipeline_from_params(params_tree): """ return analysis.log_likelihood_function(instance=params_tree) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 @@ -347,9 +351,8 @@ def full_pipeline_from_params(params_tree): ) recommended = recommend_batch_size(probe) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "interferometer")) - / "vmap_probe.json" - ) + _cli.output_dir or (_workspace_root / "results" / "runtime" / "interferometer" / "mge") + ) / "vmap_probe.json" write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -434,7 +437,9 @@ def full_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -491,7 +496,7 @@ def full_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "interferometer", + default_dir=_workspace_root / "results" / "runtime" / "interferometer" / "mge", default_basename=f"mge_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -500,7 +505,7 @@ def full_pipeline_from_params(params_tree): # --- Save bar chart --- labels = [ - f"Full pipeline (single JIT)", + "Full pipeline (single JIT)", f"vmap batch={batch_size} (per call)", ] times = [full_pipeline_per_call, vmap_per_call] @@ -528,7 +533,7 @@ def full_pipeline_from_params(params_tree): fontweight="bold", ) ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | " + f'AutoLens v{al_version} | {pixel_scale}"/px | ' f"{real_space_shape[0]}x{real_space_shape[1]} real-space | " f"{n_visibilities} visibilities | {n_linear_gaussians} Gaussians | " f"vmap speedup: {vmap_speedup:.1f}x", @@ -593,7 +598,4 @@ def full_pipeline_from_params(params_tree): rtol=_regression_rtol, err_msg=f"interferometer/mge[{instrument}]: regression — vmap log_likelihood drifted", ) - print( - f" Regression assertion PASSED: log_likelihood matches " - f"{expected_log_likelihood:.6f}" - ) + print(f" Regression assertion PASSED: log_likelihood matches {expected_log_likelihood:.6f}") diff --git a/likelihood_runtime/interferometer/pixelization.py b/likelihood_runtime/interferometer/pixelization.py index a23c6d7..4827bd9 100644 --- a/likelihood_runtime/interferometer/pixelization.py +++ b/likelihood_runtime/interferometer/pixelization.py @@ -43,33 +43,32 @@ pytree support landed in PyAutoFit#1222. """ -import numpy as np -import jax -import jax.numpy as jnp -import time import subprocess import sys -from pathlib import Path +import time from contextlib import contextmanager +from pathlib import Path +import autoarray as aa import autofit as af import autolens as al -import autoarray as aa +import jax +import jax.numpy as jnp +import numpy as np from autofit.jax import register_model as _register_model_pytrees sys.path.insert(0, str(Path(__file__).resolve().parents[2])) -from _adapt_image_util import adapt_image_for_dataset # noqa: E402 - # --------------------------------------------------------------------------- # Instrument configuration # --------------------------------------------------------------------------- - - # AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). # Verifies the import graph + module-level setup succeeded without running # the full profiling pipeline. Skipped entirely when the env var is unset. import os as _smoke_os import sys as _smoke_sys + +from _adapt_image_util import adapt_image_for_dataset # noqa: E402 + if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -77,10 +76,10 @@ # Sweep-driver CLI args (--config-name / --output-dir / --use-mixed-precision). # Tolerates extra/unknown args via parse_known_args inside the helper. from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.interferometer import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -89,6 +88,7 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() instrument = _cli.instrument or "sma" # default; override via --instrument @@ -102,6 +102,7 @@ # Profiling helpers # --------------------------------------------------------------------------- + class Timer: """Accumulates named timing measurements and prints a summary.""" @@ -240,9 +241,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): pixelization = af.Model( al.Pixelization, - mesh=al.mesh.RectangularAdaptImage( - shape=mesh_shape, weight_power=1.0, weight_floor=0.0 - ), + mesh=al.mesh.RectangularAdaptImage(shape=mesh_shape, weight_power=1.0, weight_floor=0.0), regularization=al.reg.Constant(coefficient=regularization_coefficient), ) @@ -297,9 +296,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): print("\n--- Adapt image (lensed source) ---") with timer.section("adapt_image_build"): - adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset - ) + adapt_image = adapt_image_for_dataset(dataset_path=dataset_path, dataset=dataset) adapt_images = al.AdaptImages( galaxy_image_dict={instance.galaxies.source: adapt_image}, galaxy_name_image_dict={"('galaxies', 'source')": adapt_image}, @@ -343,6 +340,7 @@ def _build_transformer(uv_wavelengths, real_space_mask): use_jax=True, ) + def full_pipeline_from_params(params_tree): """Full interferometer likelihood from a pytree-shaped ``ModelInstance``. @@ -352,6 +350,7 @@ def full_pipeline_from_params(params_tree): """ return analysis.log_likelihood_function(instance=params_tree) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 @@ -377,9 +376,9 @@ def full_pipeline_from_params(params_tree): ) recommended = recommend_batch_size(probe) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "interferometer")) - / "vmap_probe.json" - ) + _cli.output_dir + or (_workspace_root / "results" / "runtime" / "interferometer" / "pixelization") + ) / "vmap_probe.json" write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -405,8 +404,10 @@ def full_pipeline_from_params(params_tree): _n_leaves = len(jax.tree_util.tree_leaves(params_tree)) if _n_leaves == 0: - print(f" SKIPPED: model has 0 free parameters (all fixed to truth); " - f"vmap requires at least one array leaf.") + print( + " SKIPPED: model has 0 free parameters (all fixed to truth); " + "vmap requires at least one array leaf." + ) else: parameters = jax.tree_util.tree_map( lambda leaf: jnp.broadcast_to(leaf, (batch_size, *leaf.shape)), @@ -483,7 +484,9 @@ def full_pipeline_from_params(params_tree): # =================================================================== import json + import matplotlib + matplotlib.use("Agg") import matplotlib.pyplot as plt @@ -509,7 +512,7 @@ def full_pipeline_from_params(params_tree): print(f" vmap batch={batch_size} per call: {vmap_per_call:.6f} s") print(f" vmap speedup: {vmap_speedup:.1f}x") else: - print(f" vmap: SKIPPED (0 free params)") + print(" vmap: SKIPPED (0 free params)") print("=" * 70) # --- Save results dictionary --- @@ -532,13 +535,17 @@ def full_pipeline_from_params(params_tree): "figure_of_merit_eager": float(figure_of_merit_ref), "log_evidence_jit": float(full_result), "full_pipeline_single_jit": full_pipeline_per_call, - "vmap": "SKIPPED — model has 0 free parameters (all fixed to truth)" if vmap_per_call is None else { + "vmap": "SKIPPED — model has 0 free parameters (all fixed to truth)" + if vmap_per_call is None + else { "batch_size": batch_size, "batch_time": vmap_batch_time, "per_call": vmap_per_call, "speedup_vs_single_jit": round(vmap_speedup, 1), }, - "memory_mb": None if memory_analysis is None else { + "memory_mb": None + if memory_analysis is None + else { "output": memory_analysis.output_size_in_bytes / 1024**2, "temp": memory_analysis.temp_size_in_bytes / 1024**2, }, @@ -546,7 +553,7 @@ def full_pipeline_from_params(params_tree): dict_path, chart_path = resolve_output_paths( _cli, - default_dir=_workspace_root / "results" / "likelihood" / "interferometer", + default_dir=_workspace_root / "results" / "runtime" / "interferometer" / "pixelization", default_basename=f"pixelization_likelihood_summary_{instrument}_v{al_version}", ) dict_path.write_text(json.dumps(likelihood_summary, indent=2)) @@ -554,7 +561,7 @@ def full_pipeline_from_params(params_tree): # --- Save bar chart --- -labels = [f"Full pipeline (single JIT)"] +labels = ["Full pipeline (single JIT)"] times = [full_pipeline_per_call] bar_colors = ["#4C72B0"] if vmap_per_call is not None: @@ -586,7 +593,7 @@ def full_pipeline_from_params(params_tree): ) _vmap_title = f"vmap speedup: {vmap_speedup:.1f}x" if vmap_speedup is not None else "vmap: SKIPPED" ax.set_title( - f"AutoLens v{al_version} | {pixel_scale}\"/px | " + f'AutoLens v{al_version} | {pixel_scale}"/px | ' f"{real_space_shape[0]}x{real_space_shape[1]} real-space | " f"{n_visibilities} visibilities | {mesh_shape[0]}x{mesh_shape[1]} mesh | " f"{_vmap_title}", @@ -647,7 +654,4 @@ def full_pipeline_from_params(params_tree): rtol=1e-3, err_msg=f"interferometer/pixelization[{instrument}]: regression — vmap log_evidence drifted", ) - print( - f" Regression assertion PASSED: log_evidence matches " - f"{expected_log_evidence:.6f}" - ) + print(f" Regression assertion PASSED: log_evidence matches {expected_log_evidence:.6f}") diff --git a/likelihood_runtime/point_source/image_plane.py b/likelihood_runtime/point_source/image_plane.py index a0ed1ce..ef63bd4 100644 --- a/likelihood_runtime/point_source/image_plane.py +++ b/likelihood_runtime/point_source/image_plane.py @@ -38,23 +38,23 @@ """ import json -import time + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). +# Verifies the import graph + module-level setup succeeded without running +# the full profiling pipeline. Skipped entirely when the env var is unset. +import os as _smoke_os import subprocess import sys +import sys as _smoke_sys +import time from contextlib import contextmanager from pathlib import Path -import numpy as np import jax import jax.numpy as jnp import matplotlib +import numpy as np - -# AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). -# Verifies the import graph + module-level setup succeeded without running -# the full profiling pipeline. Skipped entirely when the env var is unset. -import os as _smoke_os -import sys as _smoke_sys if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -63,10 +63,10 @@ # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.point_source import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -75,16 +75,15 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() matplotlib.use("Agg") -import matplotlib.pyplot as plt - import autofit as af import autolens as al +import matplotlib.pyplot as plt from autofit.jax import register_model as _register_model_pytrees - # --------------------------------------------------------------------------- # Profiling helpers (mirrors imaging/mge.py and source_plane.py) # --------------------------------------------------------------------------- @@ -157,9 +156,7 @@ def jit_profile(func, label, *args, n_repeats=10): _script_dir = Path(__file__).resolve().parent _workspace_root = _script_dir.parents[1] -dataset_path = ( - Path("dataset") / "point_source" / instrument -) +dataset_path = Path("dataset") / "point_source" / instrument auto_simulate_if_missing( dataset_path, @@ -303,9 +300,9 @@ def full_pipeline_from_params(params_tree): ) recommended = recommend_batch_size(probe) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "point_source")) - / "vmap_probe.json" - ) + _cli.output_dir + or (_workspace_root / "results" / "runtime" / "point_source" / "image_plane") + ) / "vmap_probe.json" write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -388,10 +385,7 @@ def full_pipeline_from_params(params_tree): mem = compiled_batched.memory_analysis() print(f" Output size: {mem.output_size_in_bytes / 1024**2:.3f} MB") print(f" Temp size: {mem.temp_size_in_bytes / 1024**2:.3f} MB") -print( - f" Total: " - f"{(mem.output_size_in_bytes + mem.temp_size_in_bytes) / 1024**2:.3f} MB" -) +print(f" Total: {(mem.output_size_in_bytes + mem.temp_size_in_bytes) / 1024**2:.3f} MB") # =================================================================== @@ -407,7 +401,7 @@ def full_pipeline_from_params(params_tree): print(f" Observed image positions: {n_observed_positions}") print(f" Position noise sigma: {positions_noise_sigma}") print(f" Free parameters: {model.total_free_parameters}") -print(f" fit_positions_cls: FitPositionsImagePairAll (image-plane chi-squared)") +print(" fit_positions_cls: FitPositionsImagePairAll (image-plane chi-squared)") print("-" * 70) print(f" Eager full likelihood: {eager_per_call:.6f} s/call ({log_likelihood_ref:.6f})") print(f" Full pipeline (JIT): {full_pipeline_per_call:.6f} s/call") @@ -436,7 +430,7 @@ def full_pipeline_from_params(params_tree): }, } -results_dir = _workspace_root / "results" / "likelihood" / "point_source" +results_dir = _workspace_root / "results" / "runtime" / "point_source" / "image_plane" results_dir.mkdir(parents=True, exist_ok=True) dict_path = results_dir / f"image_plane_summary_v{al_version}.json" diff --git a/likelihood_runtime/point_source/source_plane.py b/likelihood_runtime/point_source/source_plane.py index 0073548..f833627 100644 --- a/likelihood_runtime/point_source/source_plane.py +++ b/likelihood_runtime/point_source/source_plane.py @@ -20,23 +20,23 @@ """ import json -import time + +# AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). +# Verifies the import graph + module-level setup succeeded without running +# the full profiling pipeline. Skipped entirely when the env var is unset. +import os as _smoke_os import subprocess import sys +import sys as _smoke_sys +import time from contextlib import contextmanager from pathlib import Path -import numpy as np import jax import jax.numpy as jnp import matplotlib +import numpy as np - -# AUTOLENS_PROFILING_SMOKE=1 short-circuit (Phase 5 / CI lint smoke). -# Verifies the import graph + module-level setup succeeded without running -# the full profiling pipeline. Skipped entirely when the env var is unset. -import os as _smoke_os -import sys as _smoke_sys if _smoke_os.environ.get("AUTOLENS_PROFILING_SMOKE") == "1": print(f"[smoke] {__file__}: imports + module setup OK; exiting.") _smoke_sys.exit(0) @@ -45,10 +45,10 @@ # Tolerates extra/unknown args via parse_known_args inside the helper. sys.path.insert(0, str(Path(__file__).resolve().parents[2])) from _profile_cli import ( # noqa: E402 - parse_profile_cli, + auto_simulate_if_missing, device_info_dict, + parse_profile_cli, resolve_output_paths, - auto_simulate_if_missing, ) from simulators.point_source import INSTRUMENTS # noqa: E402 from vram import ( # noqa: E402 @@ -57,17 +57,16 @@ vmap_batch_for, write_probe_json, ) + _cli = parse_profile_cli() matplotlib.use("Agg") -import matplotlib.pyplot as plt - -import autofit as af import autoarray as aa +import autofit as af import autolens as al +import matplotlib.pyplot as plt from autofit.jax import register_model as _register_model_pytrees - # --------------------------------------------------------------------------- # Profiling helpers (mirrors imaging/mge.py) # --------------------------------------------------------------------------- @@ -140,9 +139,7 @@ def jit_profile(func, label, *args, n_repeats=10): _script_dir = Path(__file__).resolve().parent _workspace_root = _script_dir.parents[1] -dataset_path = ( - Path("dataset") / "point_source" / instrument -) +dataset_path = Path("dataset") / "point_source" / instrument auto_simulate_if_missing( dataset_path, @@ -269,13 +266,11 @@ def full_pipeline_from_params(params_tree): full_pipeline_blocker = None try: - _, full_result = jit_profile( - full_pipeline_from_params, "full_pipeline", params_tree - ) + _, full_result = jit_profile(full_pipeline_from_params, "full_pipeline", params_tree) full_pipeline_per_call = timer.records[-1][1] / 10 full_pipeline_jits = True print(f" full log_likelihood = {full_result}") -except jax.errors.TracerArrayConversionError as e: +except jax.errors.TracerArrayConversionError: full_pipeline_blocker = ( "Grid2DIrregular.grid_2d_via_deflection_grid_from does not propagate xp; " "model_data ends up with _xp=np while holding JAX tracers, so " @@ -307,9 +302,9 @@ def full_pipeline_from_params(params_tree): ) recommended = recommend_batch_size(probe) probe_path = ( - (_cli.output_dir or (_workspace_root / "results" / "likelihood" / "point_source")) - / "vmap_probe.json" - ) + _cli.output_dir + or (_workspace_root / "results" / "runtime" / "point_source" / "source_plane") + ) / "vmap_probe.json" write_probe_json(probe, recommended, probe_path) print(f"\n vmap_probe samples: {probe.samples}") print(f" per_replica: {probe.per_replica_mb:.1f} MB / replica") @@ -425,10 +420,7 @@ def ray_trace_to_source_plane(params_tree, positions_raw): mem = compiled_batched.memory_analysis() print(f" Output size: {mem.output_size_in_bytes / 1024**2:.3f} MB") print(f" Temp size: {mem.temp_size_in_bytes / 1024**2:.3f} MB") -print( - f" Total: " - f"{(mem.output_size_in_bytes + mem.temp_size_in_bytes) / 1024**2:.3f} MB" -) +print(f" Total: {(mem.output_size_in_bytes + mem.temp_size_in_bytes) / 1024**2:.3f} MB") # =================================================================== @@ -444,7 +436,7 @@ def ray_trace_to_source_plane(params_tree, positions_raw): print(f" Observed image positions: {n_observed_positions}") print(f" Position noise sigma: {positions_noise_sigma}") print(f" Free parameters: {model.total_free_parameters}") -print(f" fit_positions_cls: FitPositionsSource (source-plane chi-squared)") +print(" fit_positions_cls: FitPositionsSource (source-plane chi-squared)") print("-" * 70) print(f" Eager full likelihood: {eager_per_call:.6f} s/call ({log_likelihood_ref:.6f})") if full_pipeline_jits: @@ -482,7 +474,7 @@ def ray_trace_to_source_plane(params_tree, positions_raw): }, } -results_dir = _workspace_root / "results" / "likelihood" / "point_source" +results_dir = _workspace_root / "results" / "runtime" / "point_source" / "source_plane" results_dir.mkdir(parents=True, exist_ok=True) dict_path = results_dir / f"source_plane_summary_v{al_version}.json" @@ -525,9 +517,7 @@ def ray_trace_to_source_plane(params_tree, positions_raw): fontsize=12, fontweight="bold", ) -title_extra = ( - " | full pipeline JIT BLOCKED" if not full_pipeline_jits else "" -) +title_extra = " | full pipeline JIT BLOCKED" if not full_pipeline_jits else "" ax.set_title( f"AutoLens v{al_version} | {n_observed_positions} positions | " f"{model.total_free_parameters} free params{title_extra}", diff --git a/likelihood_runtime/sweep.py b/likelihood_runtime/sweep.py index 1d23edc..8ab8a09 100644 --- a/likelihood_runtime/sweep.py +++ b/likelihood_runtime/sweep.py @@ -1,8 +1,8 @@ """Multi-config likelihood profiling driver. Runs each in-scope cell across the CPU/GPU x fp64/mp matrix (4 configs per -cell locally; HPC A100 configs are dispatched separately via -`z_projects/profiling/hpc/sync`). +cell locally; HPC A100 configs are dispatched separately via the SLURM +submit scripts under ``hpc/``). Each subprocess invokes the existing per-cell likelihood script under ``autolens_profiling/likelihood_runtime//.py`` with the @@ -13,10 +13,11 @@ ///.png ///.log (captured stdout/stderr) -Default ``--output-root`` is -``autolens_workspace_developer/jax_profiling/results/jit`` — matches the -existing imaging precedent and is read by ``aggregate.py`` to produce -``comparison.json`` / ``comparison.png``. +Default ``--output-root`` is ``results/runtime/`` in this repo, read by +``aggregate.py`` to produce ``comparison.json`` / ``comparison.png``. +(Earlier sweeps wrote to ``autolens_workspace_developer/jax_profiling/ +results/jit`` — that tree remains readable history but is no longer the +default; see ``results/notes/design_lock_in.md``.) Usage:: @@ -28,6 +29,9 @@ # Single cell, single backend python likelihood_runtime/sweep.py --only interferometer/mge --skip-cpu + + # Imaging sparse-operator (w-tilde) rows — filenames gain a _sparse suffix + python likelihood_runtime/sweep.py --only imaging/pixelization imaging/delaunay --sparse """ from __future__ import annotations @@ -40,22 +44,21 @@ from dataclasses import dataclass from pathlib import Path - -_REPO_ROOT = Path(__file__).resolve().parents[1] # autolens_profiling/ -_WT_ROOT = _REPO_ROOT.parent # PyAutoLabs-wt// (or PyAutoLabs/) -_DEFAULT_OUTPUT_ROOT = _WT_ROOT / "autolens_workspace_developer" / "jax_profiling" / "results" / "jit" -_DEFAULT_PYTHON = "/home/jammy/venv/PyAutoGPU/bin/python" +_REPO_ROOT = Path(__file__).resolve().parents[1] # autolens_profiling/ +_DEFAULT_OUTPUT_ROOT = _REPO_ROOT / "results" / "runtime" +_DEFAULT_PYTHON = sys.executable # (dataset_class, model). Order is roughly cheapest -> heaviest so failures # surface quickly during iteration. CELLS: list[tuple[str, str]] = [ - ("imaging", "pixelization"), - ("imaging", "delaunay"), - ("interferometer", "mge"), - ("interferometer", "pixelization"), - ("interferometer", "delaunay"), - ("datacube", "delaunay"), + ("imaging", "mge"), + ("imaging", "pixelization"), + ("imaging", "delaunay"), + ("interferometer", "mge"), + ("interferometer", "pixelization"), + ("interferometer", "delaunay"), + ("datacube", "delaunay"), ] @@ -121,6 +124,16 @@ def _parse_args() -> argparse.Namespace: action="store_true", help="Skip the use_mixed_precision rows (just fp64).", ) + p.add_argument( + "--sparse", + action="store_true", + help=( + "Pass --sparse to every selected cell (w-tilde sparse-operator " + "inversion path; imaging cells only — combine with --only). " + "Result filenames gain a _sparse suffix so dense and sparse " + "rows coexist." + ), + ) p.add_argument( "--output-root", type=Path, @@ -170,17 +183,22 @@ def _run_one( config: SweepConfig, out_dir: Path, dry_run: bool, + sparse: bool = False, ) -> tuple[bool, float, str]: """Run one (cell, config) pair as a subprocess. Returns (ok, elapsed, log_path).""" out_dir.mkdir(parents=True, exist_ok=True) - log_path = out_dir / f"{config.name}.log" + log_suffix = "_sparse" if sparse else "" + log_path = out_dir / f"{config.name}{log_suffix}.log" cmd = [ python, str(script_path), - "--config-name", config.name, - "--output-dir", str(out_dir), + "--config-name", + config.name, + "--output-dir", + str(out_dir), *config.extra_args, + *(("--sparse",) if sparse else ()), ] env = dict(os.environ) @@ -208,12 +226,15 @@ def _run_one( ) elapsed = time.time() - t0 ok = proc.returncode == 0 - print(f" {'OK ' if ok else 'FAIL'} ({elapsed:.1f}s, exit={proc.returncode}) -> {log_path.name}") + print( + f" {'OK ' if ok else 'FAIL'} ({elapsed:.1f}s, exit={proc.returncode}) -> {log_path.name}" + ) # Verify the device.backend in the JSON matches expectations. if ok: import json - json_path = out_dir / f"{config.name}.json" + + json_path = out_dir / f"{config.name}{log_suffix}.json" if json_path.exists(): try: data = json.loads(json_path.read_text()) @@ -238,8 +259,10 @@ def main() -> int: cells = _resolve_cells(args) configs = _resolve_configs(args) - print(f"sweep_likelihood: {len(cells)} cells x {len(configs)} configs " - f"= {len(cells) * len(configs)} runs") + print( + f"sweep_likelihood: {len(cells)} cells x {len(configs)} configs " + f"= {len(cells) * len(configs)} runs" + ) print(f" cells: {[f'{c}/{m}' for (c, m) in cells]}") print(f" configs: {[c.name for c in configs]}") print(f" output: {args.output_root}") @@ -250,7 +273,7 @@ def main() -> int: summary: list[tuple[str, str, bool, float]] = [] overall_t0 = time.time() - for (cls, model) in cells: + for cls, model in cells: script_path = _REPO_ROOT / "likelihood_runtime" / cls / f"{model}.py" if not script_path.exists(): print(f"\n!!! missing script: {script_path}") @@ -263,7 +286,12 @@ def main() -> int: for cfg in configs: try: ok, elapsed, _log = _run_one( - args.python, script_path, cfg, out_dir, args.dry_run + args.python, + script_path, + cfg, + out_dir, + args.dry_run, + sparse=args.sparse, ) except KeyboardInterrupt: print("\n\nsweep interrupted by user") @@ -275,7 +303,7 @@ def main() -> int: print(f"sweep_likelihood summary ({total:.1f}s total)") print("=" * 70) print(f" {'cell':<32}{'config':<22}{'ok':<6}{'elapsed':>10}") - print(f" {'-'*32}{'-'*22}{'-'*6}{'-'*10}") + print(f" {'-' * 32}{'-' * 22}{'-' * 6}{'-' * 10}") failures = 0 for cell, cfg, ok, t in summary: flag = "OK" if ok else "FAIL" diff --git a/quick_update/README.md b/quick_update/README.md new file mode 100644 index 0000000..f5c2758 --- /dev/null +++ b/quick_update/README.md @@ -0,0 +1,22 @@ +# quick_update + +Fast incremental re-profiling helpers: one thin script per dataset class that +re-runs the headline likelihood measurement in seconds, for tight +edit-profile-edit loops while iterating on the source libraries. + +These are the **scratch tier** of the repo: + +- Outputs land in `results/quick_update/` as **unversioned** JSONs + (`_quick_update_.json`) that are overwritten on every run — + no history, no dashboard row. +- Numbers here are for steering a source-code change, never for citing: when a + result matters, re-run the real cell under + [`likelihood_runtime/`](../likelihood_runtime/README.md) or + [`likelihood_breakdown/`](../likelihood_breakdown/README.md). + +## Running + +```bash +python quick_update/imaging.py +python quick_update/interferometer_delaunay.py +``` diff --git a/quick_update/imaging.py b/quick_update/imaging.py index 8418429..dc7a393 100644 --- a/quick_update/imaging.py +++ b/quick_update/imaging.py @@ -26,11 +26,10 @@ os.environ.setdefault("MPLBACKEND", "Agg") -from autoconf import jax_wrapper # noqa: E402 - import autofit as af # noqa: E402 import autolens as al # noqa: E402 import numpy as np # noqa: E402 +from autoconf import jax_wrapper # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -39,6 +38,7 @@ try: from _profile_cli import device_info_dict except ImportError: + def device_info_dict(): return {"backend": "unknown"} @@ -49,6 +49,7 @@ def device_info_dict(): import argparse # noqa: E402 + def _parse_args(): p = argparse.ArgumentParser(prog="quick_update/imaging.py") p.add_argument("--instrument", default="hst", choices=list(INSTRUMENTS)) @@ -63,6 +64,7 @@ def _parse_args(): # Timer # --------------------------------------------------------------------------- + class Timer: def __init__(self): self.records: dict[str, list[float]] = {} @@ -151,7 +153,11 @@ def first(self, label: str) -> float: shear = af.Model(al.mp.ExternalShear) lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear, + al.Galaxy, + redshift=0.5, + bulge=lens_bulge, + mass=mass, + shear=shear, ) source_bulge = al.model_util.mge_model_from( @@ -219,9 +225,7 @@ def first(self, label: str) -> float: for i in range(n_repeats): with timer.section("critical_curves"): - ip_lines, ip_colors, sp_lines, sp_colors = ( - _compute_critical_curves_from_fit(fit) - ) + ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curves_from_fit(fit) # --------------------------------------------------------------------------- # Phase 3: Render comparison (12-panel vs 6-panel) @@ -311,15 +315,11 @@ def first(self, label: str) -> float: "n_repeats": n_repeats, "device": device_info_dict(), "phases": summary, - "all_timings": { - k: [round(v, 4) for v in vals] - for k, vals in timer.records.items() - }, + "all_timings": {k: [round(v, 4) for v in vals] for k, vals in timer.records.items()}, } output_dir = ( - Path(args.output_dir) if args.output_dir - else workspace_root / "results" / "quick_update" + Path(args.output_dir) if args.output_dir else workspace_root / "results" / "quick_update" ) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"imaging_mge_quick_update_{instrument}.json" diff --git a/quick_update/imaging_delaunay.py b/quick_update/imaging_delaunay.py index a0de559..bb41cbc 100644 --- a/quick_update/imaging_delaunay.py +++ b/quick_update/imaging_delaunay.py @@ -30,21 +30,21 @@ os.environ.setdefault("MPLBACKEND", "Agg") -from autoconf import jax_wrapper # noqa: E402 - +import autoarray as aa # noqa: E402 import autofit as af # noqa: E402 import autolens as al # noqa: E402 -import autoarray as aa # noqa: E402 import numpy as np # noqa: E402 +from autoconf import jax_wrapper # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from instruments.imaging import INSTRUMENTS # noqa: E402 from _adapt_image_util import adapt_image_for_dataset # noqa: E402 +from instruments.imaging import INSTRUMENTS # noqa: E402 try: from _profile_cli import device_info_dict except ImportError: + def device_info_dict(): return {"backend": "unknown"} @@ -55,6 +55,7 @@ def device_info_dict(): import argparse # noqa: E402 + def _parse_args(): p = argparse.ArgumentParser(prog="quick_update/imaging_delaunay.py") p.add_argument("--instrument", default="hst", choices=list(INSTRUMENTS)) @@ -70,6 +71,7 @@ def _parse_args(): # Timer # --------------------------------------------------------------------------- + class Timer: def __init__(self): self.records: dict[str, list[float]] = {} @@ -153,7 +155,8 @@ def first(self, label: str) -> float: print("\n Loading adapt image...") adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset, + dataset_path=dataset_path, + dataset=dataset, ) print(f" adapt_image shape (slim): {adapt_image.shape_slim}") @@ -161,7 +164,7 @@ def first(self, label: str) -> float: # Model: MGE lens light + Isothermal mass + Delaunay source # --------------------------------------------------------------------------- -print(f"\n Building Delaunay model...") +print("\n Building Delaunay model...") lens_bulge = al.model_util.mge_model_from( mask_radius=mask_radius, @@ -173,7 +176,11 @@ def first(self, label: str) -> float: shear = af.Model(al.mp.ExternalShear) lens = af.Model( - al.Galaxy, redshift=0.5, bulge=lens_bulge, mass=mass, shear=shear, + al.Galaxy, + redshift=0.5, + bulge=lens_bulge, + mass=mass, + shear=shear, ) mesh = al.mesh.Delaunay(pixels=mesh_pixels, zeroed_pixels=0) @@ -190,10 +197,13 @@ def first(self, label: str) -> float: print("\n Building Hilbert image-plane mesh grid...") image_mesh = al.image_mesh.Hilbert( - pixels=mesh_pixels, weight_power=1.0, weight_floor=0.0, + pixels=mesh_pixels, + 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, + mask=dataset.mask, + adapt_data=adapt_image, ) print(f" mesh vertices: {image_plane_mesh_grid.shape[0]}") @@ -257,9 +267,7 @@ def first(self, label: str) -> float: for i in range(n_repeats): with timer.section("critical_curves"): - ip_lines, ip_colors, sp_lines, sp_colors = ( - _compute_critical_curves_from_fit(fit) - ) + ip_lines, ip_colors, sp_lines, sp_colors = _compute_critical_curves_from_fit(fit) # --------------------------------------------------------------------------- # Phase 3: Render comparison @@ -351,15 +359,11 @@ def first(self, label: str) -> float: "n_repeats": n_repeats, "device": device_info_dict(), "phases": summary, - "all_timings": { - k: [round(v, 4) for v in vals] - for k, vals in timer.records.items() - }, + "all_timings": {k: [round(v, 4) for v in vals] for k, vals in timer.records.items()}, } output_dir = ( - Path(args.output_dir) if args.output_dir - else workspace_root / "results" / "quick_update" + Path(args.output_dir) if args.output_dir else workspace_root / "results" / "quick_update" ) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"imaging_delaunay_quick_update_{instrument}.json" diff --git a/quick_update/interferometer.py b/quick_update/interferometer.py index 707ad63..ca5b34e 100644 --- a/quick_update/interferometer.py +++ b/quick_update/interferometer.py @@ -23,11 +23,10 @@ os.environ.setdefault("MPLBACKEND", "Agg") -from autoconf import jax_wrapper # noqa: E402 - import autofit as af # noqa: E402 import autolens as al # noqa: E402 import numpy as np # noqa: E402 +from autoconf import jax_wrapper # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parents[1])) @@ -36,12 +35,14 @@ try: from _profile_cli import device_info_dict except ImportError: + def device_info_dict(): return {"backend": "unknown"} import argparse # noqa: E402 + def _parse_args(): p = argparse.ArgumentParser(prog="quick_update/interferometer.py") p.add_argument("--instrument", default="sma", choices=list(INSTRUMENTS)) @@ -149,7 +150,9 @@ def first(self, label: str) -> float: # Phase 2: Critical curves print("\n=== Phase 2: Critical curves ===") -from autolens.interferometer.plot.fit_interferometer_plots import _compute_critical_curve_lines # noqa: E402 +from autolens.interferometer.plot.fit_interferometer_plots import ( + _compute_critical_curve_lines, # noqa: E402 +) for i in range(n_repeats): tracer = fit.tracer_linear_light_profiles_to_light_profiles @@ -213,15 +216,11 @@ def first(self, label: str) -> float: "n_repeats": n_repeats, "device": device_info_dict(), "phases": summary, - "all_timings": { - k: [round(v, 4) for v in vals] - for k, vals in timer.records.items() - }, + "all_timings": {k: [round(v, 4) for v in vals] for k, vals in timer.records.items()}, } output_dir = ( - Path(args.output_dir) if args.output_dir - else workspace_root / "results" / "quick_update" + Path(args.output_dir) if args.output_dir else workspace_root / "results" / "quick_update" ) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"interferometer_mge_quick_update_{instrument}.json" diff --git a/quick_update/interferometer_delaunay.py b/quick_update/interferometer_delaunay.py index 3f0eedb..aaad7c0 100644 --- a/quick_update/interferometer_delaunay.py +++ b/quick_update/interferometer_delaunay.py @@ -22,27 +22,28 @@ os.environ.setdefault("MPLBACKEND", "Agg") -from autoconf import jax_wrapper # noqa: E402 - +import autoarray as aa # noqa: E402 import autofit as af # noqa: E402 import autolens as al # noqa: E402 -import autoarray as aa # noqa: E402 import numpy as np # noqa: E402 +from autoconf import jax_wrapper # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parents[1])) -from instruments.interferometer import INSTRUMENTS # noqa: E402 from _adapt_image_util import adapt_image_for_dataset # noqa: E402 +from instruments.interferometer import INSTRUMENTS # noqa: E402 try: from _profile_cli import device_info_dict except ImportError: + def device_info_dict(): return {"backend": "unknown"} import argparse # noqa: E402 + def _parse_args(): p = argparse.ArgumentParser(prog="quick_update/interferometer_delaunay.py") p.add_argument("--instrument", default="sma", choices=list(INSTRUMENTS)) @@ -112,14 +113,16 @@ def first(self, label: str) -> float: # Adapt image print("\n Loading adapt image...") adapt_image = adapt_image_for_dataset( - dataset_path=dataset_path, dataset=dataset, + dataset_path=dataset_path, + dataset=dataset, ) print(f" adapt_image shape: {adapt_image.shape_slim}") print(" Building Hilbert mesh grid...") image_mesh = al.image_mesh.Hilbert(pixels=mesh_pixels, weight_power=1.0, weight_floor=0.0) image_plane_mesh_grid = image_mesh.image_plane_mesh_grid_from( - mask=real_space_mask, adapt_data=adapt_image, + mask=real_space_mask, + adapt_data=adapt_image, ) print(f" mesh vertices: {image_plane_mesh_grid.shape[0]}") @@ -227,15 +230,11 @@ def first(self, label: str) -> float: "n_repeats": n_repeats, "device": device_info_dict(), "phases": summary, - "all_timings": { - k: [round(v, 4) for v in vals] - for k, vals in timer.records.items() - }, + "all_timings": {k: [round(v, 4) for v in vals] for k, vals in timer.records.items()}, } output_dir = ( - Path(args.output_dir) if args.output_dir - else workspace_root / "results" / "quick_update" + Path(args.output_dir) if args.output_dir else workspace_root / "results" / "quick_update" ) output_dir.mkdir(parents=True, exist_ok=True) output_path = output_dir / f"interferometer_delaunay_quick_update_{instrument}.json" diff --git a/quick_update/point_source.py b/quick_update/point_source.py index 4062a89..9ec7997 100644 --- a/quick_update/point_source.py +++ b/quick_update/point_source.py @@ -23,18 +23,18 @@ os.environ.setdefault("MPLBACKEND", "Agg") -from autoconf import jax_wrapper # noqa: E402 -from autoconf.dictable import from_dict # noqa: E402 - import autofit as af # noqa: E402 import autolens as al # noqa: E402 import numpy as np # noqa: E402 +from autoconf import jax_wrapper # noqa: E402 +from autoconf.dictable import from_dict # noqa: E402 sys.path.insert(0, str(Path(__file__).resolve().parents[1])) try: from _profile_cli import device_info_dict except ImportError: + def device_info_dict(): return {"backend": "unknown"} @@ -162,10 +162,7 @@ def first(self, label: str) -> float: "n_repeats": n_repeats, "device": device_info_dict(), "phases": summary, - "all_timings": { - k: [round(v, 4) for v in vals] - for k, vals in timer.records.items() - }, + "all_timings": {k: [round(v, 4) for v in vals] for k, vals in timer.records.items()}, } output_dir = workspace_root / "results" / "quick_update" diff --git a/results/README.md b/results/README.md index f244ee6..ef3a0e6 100644 --- a/results/README.md +++ b/results/README.md @@ -1,22 +1,50 @@ # results -Versioned profiling artifacts written by scripts under [`likelihood/`](../likelihood/README.md), [`simulators/`](../simulators/README.md), and [`searches/`](../searches/README.md). Layout mirrors the source folders. +Profiling artifacts written by the packages above. Layout mirrors the source +packages; the dashboard tables in every README are rendered from this tree by +`scripts/build_readme.py`. -## Filename convention +## Sections + +| Folder | Written by | Contents | +|--------|-----------|----------| +| `runtime/` | [`likelihood_runtime/`](../likelihood_runtime/README.md) sweeps | Per-config sweep outputs + `comparison.{json,png}` per cell; A100 logs/probes | +| `breakdown/` | [`likelihood_breakdown/`](../likelihood_breakdown/README.md) | Versioned per-step decompositions | +| `simulators/` | [`simulators/`](../simulators/README.md) | Versioned simulator run-time summaries | +| `searches/` | [`searches/`](../searches/README.md) | Versioned sampler profiling summaries | +| `quick_update/` | [`quick_update/`](../quick_update/README.md) | Unversioned fast re-profiling snapshots (scratch tier) | +| `notes/` | humans + agents | Narrative findings and design notes (e.g. [`design_lock_in.md`](./notes/design_lock_in.md)) | +| `baselines/` | campaign snapshots | Named, frozen baselines (e.g. `PreOptimizationTimes/`) — see below | + +## The two artifact shapes + +**Versioned summaries** — written by per-cell scripts run standalone; history +is retained side-by-side so cross-release trends stay inspectable: ``` -_summary_v....json -_summary_v....png +___v...[_sparse].json # purpose = summary | breakdown +___v...[_sparse].png ``` -The version string is the PyAutoLens release that produced the numbers (e.g. `v2026.5.1.4`). Older versions are kept alongside newer ones so cross-release trends stay inspectable. +The version string is the PyAutoLens release that produced the numbers +(e.g. `v2026.5.29.4`). -Example: +**Per-config sweeps** — written by `likelihood_runtime/sweep.py` and +aggregated by `aggregate.py`; each cell dir holds the *latest* sweep: ``` -results/likelihood/imaging/imaging_summary_v2026.5.1.4.json -results/likelihood/imaging/imaging_summary_v2026.5.1.4.png -results/likelihood/imaging/mge/delaunay_sparse_cpu_likelihood_summary_hst_v2026.5.1.4.json +runtime//[/]/[_sparse].{json,png,log} +runtime//[/]/comparison.{json,png} ``` -Populated as Phases 1–3 land. See the top-level [README](../README.md) for the full phase plan. +Config names: `local_cpu_fp64 | local_cpu_mp | local_gpu_fp64 | local_gpu_mp | +hpc_a100_fp64 | hpc_a100_mp`, with `_sparse` as a filename suffix. + +## Named baselines + +A **baseline** is a frozen snapshot of a full campaign under +`baselines//`, mirroring the `runtime/` layout plus a rendered +`.md` with every headline number on one page. Baselines are +append-only — never edited after the campaign closes. The first baseline is +**`PreOptimizationTimes`** (the pre-optimization reference). Full convention: +[`notes/design_lock_in.md`](./notes/design_lock_in.md). diff --git a/results/notes/design_lock_in.md b/results/notes/design_lock_in.md new file mode 100644 index 0000000..cbbf580 --- /dev/null +++ b/results/notes/design_lock_in.md @@ -0,0 +1,101 @@ +# Design lock-in — pre-PreOptimizationTimes review (2026-07-08) + +One final holistic review of the repo before the **PreOptimizationTimes** +baseline campaign — the last full profiling sweep before the optimization +push. This note records what is locked in, what changed, and why, so future +extensions (more datasets, instruments, packages) build on a settled core. +Tracked by [autolens_profiling#52](https://github.com/PyAutoLabs/autolens_profiling/issues/52); +parent intent: PyAutoMind `maintenance/autolens_profiling/polish.md`. + +## What is locked in (unchanged — the design is right) + +- **The package split.** `likelihood_runtime/` (how long per call?) vs + `likelihood_breakdown/` (where does the time go?) vs `vram/` (does it fit, + and at what vmap batch size?) are deliberately disjoint questions with + deliberately disjoint packages. `latent/`, `searches/`, `simulators/`, + `quick_update/` follow the same one-question-per-package rule. New profiling + concerns get a **new package**, not a flag on an existing one. +- **The cell grid.** Work is addressed as `/` cells + (e.g. `imaging/mge`, `datacube/delaunay`), optionally deepened by + `/`. Sweep drivers, aggregators, output dirs and HPC submit + scripts all speak this grid. New datasets/models slot in as new cells. +- **`_profile_cli.py` as the single CLI surface.** Every per-cell script takes + `--config-name / --output-dir / --use-mixed-precision / --instrument / + --vmap-probe / --sparse` through the shared helper. New flags go here, once, + never per-script. +- **Instrument framing** via `instruments/` presets; headline numbers are + named for observing programmes (HST, ALMA, JWST…), not pixel counts. +- **Config axis names**: `local_cpu_fp64 | local_cpu_mp | local_gpu_fp64 | + local_gpu_mp | hpc_a100_fp64 | hpc_a100_mp` (the 6-config matrix), with + `_sparse` as a filename suffix, not a seventh config. +- **Correctness gates inside every timing script** (eager ≡ JIT ≡ vmap at + documented rtol) stay mandatory for new cells. + +## What changed in this review + +1. **Results live wholly in this repo.** `sweep.py` / `aggregate.py` + previously defaulted `--output-root` to + `../autolens_workspace_developer/jax_profiling/results/jit` — a migration + leftover. The default is now **`results/runtime/`** in-repo. The + workspace_developer tree remains readable history; nothing new is written + there. (`--output-root` still overrides for scratch runs.) + The same fix propagated to every default that still pointed at a legacy + home: the per-cell `likelihood_runtime` scripts (standalone/HPC runs + defaulted to the retired `results/likelihood//`; now + `results/runtime///`, so SLURM jobs land exactly where a + local sweep writes) and `latent/sweep.py` / `latent/aggregate.py` + (workspace_developer → `results/latent/`). Because a cell dir can now hold + both sweep configs and versioned standalone summaries, `aggregate.py` + filters to config-shaped stems when building `comparison.json`. +2. **No machine-specific defaults.** `sweep.py`'s hard-coded + `/home/jammy/venv/PyAutoGPU/bin/python` default became `sys.executable` + (violated the repo's own "no machine-specific absolute paths" rule). +3. **`imaging/mge` joined the sweep `CELLS` grid** so the runtime campaign + covers the same imaging cells as the breakdown package; `sweep.py` also + gained `--sparse` passthrough so the imaging sparse-vs-mapping comparison + runs through the same driver as everything else. +4. **The README dashboard was revived and retargeted.** + `scripts/build_readme.py` still pointed at the retired `likelihood/` + package (deleted in the runtime/breakdown split) and scanned 0 artifacts. + It now scans the real `results/` sections (`breakdown`, `runtime`, + `simulators`, `searches`), understands both artifact shapes (see below), + and renders auto-tables into the top-level README + per-package READMEs. + CI's `--check` idempotence gate is unchanged. +5. **Stale docs fixed**: retired-`likelihood/` references in `AGENTS.md`, + `README.md`, `results/README.md` and `_profile_cli.py`; the + `z_projects/profiling/hpc/sync` references (HPC dispatch lives in-repo + under `hpc/`); the roadmap's retired `PyAutoPrompt` registry pointer. + `quick_update/`, `hpc/` and `scripts/` gained READMEs. + +## The two artifact shapes (both canonical) + +| Shape | Written by | Pattern | Versioning | +|-------|-----------|---------|------------| +| **Versioned summary** | per-cell scripts run standalone | `___v[_sparse].{json,png}` | in the filename; history retained side-by-side | +| **Per-config sweep** | `sweep.py` → `aggregate.py` | `/[/]/[_sparse].{json,png,log}` + `comparison.{json,png}` | in each JSON's metadata; dirs hold the *latest* sweep | + +Rule of thumb: **cross-release trend** questions read versioned summaries; +**cross-hardware comparison** questions read `comparison.json`. + +## The PreOptimizationTimes convention (for phases 2–4) + +- A **baseline** is a named, frozen snapshot of campaign results: + `results/baselines//` containing (a) the `comparison.json` + per swept cell, (b) the versioned summary JSONs the campaign produced, and + (c) a rendered **`.md`** — every headline number in one + browsable page (cells × configs × instruments). +- `PreOptimizationTimes` is the first such baseline: laptop CPU, HPC CPU and + HPC A100 (+ mp where supported); vmap-only runtime numbers using the + `vram/` batch-size table; laptop GPU appended later by hand. +- The top-level README dashboard grows a baseline column once + `results/baselines/PreOptimizationTimes/` exists — the "compare against + this" anchor for all optimization work that follows. +- Baselines are **append-only**: never edited after the campaign closes; a + post-optimization campaign snapshots a new name next to it. + +## Deliberately out of scope here + +Profiling runs (phases 2–4); searches profiling; `point_source` cells (in the +grid, excluded from this campaign); laptop-GPU rows (human-run follow-up); +gradient profiling (still in `autolens_workspace_developer`, folds in later); +the future PyAutoBrain profiling agent (separate `feature/pyautobrain/` task). diff --git a/ruff.toml b/ruff.toml index 6f152d5..331eb79 100644 --- a/ruff.toml +++ b/ruff.toml @@ -41,6 +41,10 @@ ignore = [ "F401", # unused imports — some scripts deliberately import-for-side-effect # (e.g. `from autoconf import jax_wrapper # noqa: F401`) "B008", # function call in default argument — common autofit pattern + "B905", # zip() without strict= — scripts zip equal-length lists they built + # themselves; threading strict= through adds noise, not safety + "B007", # unused loop-control variable — warm-up/repeat loops (`for _i in + # range(n)`) are the dominant pattern in timing code ] [lint.per-file-ignores] diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 0000000..b52c487 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,8 @@ +# scripts + +Repo tooling (not profiling scripts). + +- **`build_readme.py`** — renders the auto-generated dashboard tables in the + top-level and per-package READMEs from the artifacts under `results/`. + Run `python scripts/build_readme.py` after a profiling run and commit the + result; CI's `lint.yml` runs `--check` to enforce idempotence. diff --git a/scripts/build_readme.py b/scripts/build_readme.py index 9ce91d1..d393c1c 100644 --- a/scripts/build_readme.py +++ b/scripts/build_readme.py @@ -1,6 +1,6 @@ """ build_readme.py — refresh auto-generated tables in every README from the -latest versioned artifacts under `results/`. +latest artifacts under `results/`. Run from the repo root: @@ -10,50 +10,51 @@ Each table region in a README is delimited by sentinel comments, e.g. - + | ... | - + This script: - 1. Scans `results/**/*_summary_v.json`. - 2. Parses filenames into (section, sub-folder, script, instrument, version). - 3. Picks the latest version per group via PEP 440-ish dotted-version sort. - 4. Generates a markdown table per known region type and replaces the - content inside the matching sentinel block. - -Sections covered today: - - - top-level README.md - ... - - likelihood/README.md (section overview) - - likelihood/imaging/README.md | likelihood-imaging - - likelihood/interferometer/README.md | likelihood-interferometer - - likelihood/point_source/README.md | likelihood-point_source - - likelihood/datacube/README.md | likelihood-datacube - - simulators/README.md | simulators - - searches/nautilus/README.md | searches-nautilus - -Hardware-tier columns (CPU / laptop GPU / HPC GPU) are deferred — every -artifact today is implicitly CPU and the table shows a single "Latest" -column. Once future artifacts encode hardware in the filename or JSON -(`*_summary_v_.json` or `{"hardware": "a100"}`), the -column logic in `_render_*_table` will be extended without touching the -sentinel layout. + 1. Scans `results/{breakdown,simulators,searches}/**` for **versioned + artifacts** (`