Skip to content

feat: Phase 1 — mirror JIT likelihood profiling scripts#2

Merged
Jammy2211 merged 1 commit into
mainfrom
feature/likelihood-jit-mirror
May 16, 2026
Merged

feat: Phase 1 — mirror JIT likelihood profiling scripts#2
Jammy2211 merged 1 commit into
mainfrom
feature/likelihood-jit-mirror

Conversation

@Jammy2211
Copy link
Copy Markdown
Contributor

Summary

Phase 1 of the autolens_profiling z_feature roadmap — mirrors the JIT likelihood profiling scripts and their input datasets from autolens_workspace_developer/jax_profiling/jit/ into this repo at likelihood/ and dataset/. _developer remains the source of truth; nothing is moved or deleted upstream.

Closes #1.

What changed

  • 9 likelihood profiling scripts + 1 __init__.py mirrored at likelihood/{imaging,interferometer,point_source,datacube}/. Filenames preserved verbatim. The step-by-step profile narrative each script prints is unchanged.
  • 14 input dataset files mirrored at dataset/{imaging/hst, interferometer/sma, point_source/simple}/. Checksums identical to source.
  • 5 READMEs authored: top-level likelihood/README.md (replacing the placeholder from Phase 0) plus one per subfolder. Each per-section README covers what the likelihood path profiles, scripts available, default dataset, and a headline run-time table placeholder for Phase 4's dashboard.
  • .gitignore extended to cover regenerable dataset variants (dataset/{cluster,group,multi}/, */simple/ data files) and per-run profiling artifacts under results/likelihood/.

Path rewrites (mechanical, uniform across all 9 scripts)

Pattern Before After
Dataset path Path("jax_profiling") / "dataset" / ... Path("dataset") / ...
Workspace-root walk-up _script_dir.parents[2] _script_dir.parents[1]
Results path _workspace_root / "jax_profiling" / "results" / "jit" / <type> _workspace_root / "results" / "likelihood" / <type>
Docstring sibling refs ``jax_profiling/{jit,imaging,interferometer,…}/`` ``likelihood/...`` and `dataset/...`

The artifact filename convention is preserved verbatim: <script>_likelihood_summary_<instrument>_v<al.__version__>.{json,png}.

Dataset regeneration is documented, not bundled

The new repo does not include dataset_setup/ regeneration scripts (out of scope for Phase 1). The if should_simulate(...) block at the top of each script is replaced with a clear FileNotFoundError that points readers back at the source-of-truth regen scripts in autolens_workspace_developer/jax_profiling/dataset_setup/. Default smoke runs use the pre-mirrored HST / SMA / simple datasets and never trigger this path.

Decisions captured

  • Dataset folder is top-level dataset/, not likelihood/dataset/. Reason: Phase 2 (simulators) and Phase 3 (Nautilus) will consume the same input data and a shared folder avoids duplication.
  • No CI in this PR. Phase 5 lands the GitHub Actions for lint + profile re-runs.

Test plan — smoke results (all 4 ran end-to-end on CPU)

Script Artifacts at expected path? Regression assertion
likelihood/imaging/mge.py mge_likelihood_summary_hst_v2026.5.14.2.{json,png} ⚠️ pre-existing drift (got 27542.08, expected 27379.39; same drift in _developer/jax_profiling/jit/imaging/mge.py — script logic identical, constant EXPECTED_LOG_LIKELIHOOD_HST unchanged)
likelihood/interferometer/mge.py mge_likelihood_summary_sma_v2026.5.14.2.{json,png} ✓ PASSED
likelihood/point_source/image_plane.py image_plane_summary_v2026.5.14.2.{json,png} ⚠️ pre-existing drift (got -362.21, expected 0.07; large drift, but constant EXPECTED_LOG_LIKELIHOOD_IMAGE_PLANE unchanged from _developer — same script logic)
likelihood/datacube/delaunay.py delaunay_likelihood_summary_sma_v2026.5.14.2.{json,png} ✓ PASSED (both eager and full-pipeline cube assertions)

Mirror-vs-source diff verified clean: only the documented path rewrites and the should_simulate-block replacement changed. The regression-constant drifts pre-date this PR and would manifest if the corresponding scripts in autolens_workspace_developer/jax_profiling/jit/ were rerun on PyAutoLens 2026.5.14.2. They are independent of the mirror and should be triaged in their own issue against _developer (or the constants refreshed and re-committed if the new values are correct).

Phase 0 bootstrap commit (for context): 0087d6a

…se 1)

Phase 1 of the autolens_profiling z_feature roadmap. Mirrors the JIT
likelihood profiling scripts and their pre-existing tracked input datasets
from PyAutoLabs/autolens_workspace_developer/jax_profiling/jit/ into this
repo. _developer remains the source of truth — no scripts or data are
moved or deleted upstream.

Scripts mirrored (9 .py + 1 __init__.py, verbatim filenames):
  likelihood/imaging/{mge,pixelization,delaunay}.py
  likelihood/interferometer/{mge,pixelization,delaunay}.py
  likelihood/point_source/{image_plane,source_plane}.py
  likelihood/datacube/{__init__.py,delaunay.py}

Datasets mirrored (14 files, ~900K, checksums verified identical to source):
  dataset/imaging/hst/, dataset/imaging/source_complex/
  dataset/interferometer/sma/
  dataset/point_source/simple/

Path rewrites applied uniformly across all 9 scripts:
- Path("jax_profiling") / "dataset" / ... -> Path("dataset") / ...
- _script_dir.parents[2] -> _script_dir.parents[1] (one fewer level in
  the new layout: likelihood/<type>/<script>.py vs jax_profiling/jit/...)
- _workspace_root / "jax_profiling" / "results" / "jit" / <type>
  -> _workspace_root / "results" / "likelihood" / <type>
- Docstring sibling-script and dataset refs rewritten to the new layout.

The autolens_profiling repo does not include dataset_setup/ regeneration
scripts (out of scope for Phase 1). The if should_simulate(...) block at
the top of each script is replaced with a clear FileNotFoundError that
points readers back at autolens_workspace_developer/jax_profiling/dataset_setup/
for regeneration. Default smoke runs use the pre-mirrored HST / SMA /
simple datasets and never trigger this path.

READMEs authored:
- likelihood/README.md (replaces Phase 0 placeholder) — section overview,
  what JIT likelihood profiling means, how to read per-script output,
  versioned-artifact convention, conventions inherited from _developer.
- likelihood/{imaging,interferometer,point_source,datacube}/README.md —
  per-section narrative, script tables, default datasets, headline
  run-time table placeholder for Phase 4's auto-generated dashboard.

.gitignore extended:
- dataset/{cluster,group,multi}/, dataset/{imaging,interferometer}/simple/,
  dataset/point_source/simple/*.fits + *.csv — regenerable variants.
- dataset/datacube/sim_*/ — datacube simulator outputs.
- results/{likelihood,simulators,searches}/ — per-run profiling artifacts;
  Phase 4's dashboard will explicitly select which to commit.

Smoke runs (all 4 ran end-to-end on CPU; artifacts landed at expected
paths):
- imaging/mge.py [hst] — artifacts ✓; regression assertion pre-existing
  drift (constant unchanged from _developer; same drift would manifest
  upstream on PyAutoLens 2026.5.14.2).
- interferometer/mge.py [sma] — artifacts ✓; all assertions PASSED.
- point_source/image_plane.py [simple] — artifacts ✓; regression
  assertion pre-existing drift (same as above).
- datacube/delaunay.py [sma × 4] — artifacts ✓; both eager and
  full-pipeline cube regression assertions PASSED.

Closes #1.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: Phase 1 — mirror JIT likelihood profiling scripts

1 participant