feat: Phase 1 — mirror JIT likelihood profiling scripts

[Jammy2211]

Overview

Phase 1 of the autolens_profiling z_feature: mirror the JIT-profiling likelihood scripts currently in autolens_workspace_developer/jax_profiling/jit/ (10 files, ~7,000 LOC across imaging × 3 model compositions, interferometer × 3, point-source × 2, datacube × 1) into the new repo at likelihood/. _developer stays the source of truth — the originals do not move. This phase also writes the section READMEs (top-level + 4 per-dataset-type) that the Phase 4 dashboard will read from.

Plan

• Mirror 9 likelihood scripts + 1 __init__.py from _developer/jax_profiling/jit/ into autolens_profiling/likelihood/, preserving filenames and the step-by-step profile narrative each script prints.
• Mirror the ~900K of tracked input datasets under _developer/jax_profiling/dataset/{imaging,interferometer,point_source}/ into the new repo at top-level dataset/ (shared with Phase 2 simulators and Phase 3 searches — not under likelihood/).
• Rewrite path references inside the scripts: Path("jax_profiling") / "dataset" / ... → Path("dataset") / ..., and Path("jax_profiling") / "results" / "jit" / <type> → Path("results") / "likelihood" / <type>. The artifact filename convention (<script>_likelihood_summary_<instrument>_v<al.__version__>.{json,png}) is kept verbatim.
• Update the new repo's .gitignore so dataset/{cluster,group,multi}/, dataset/imaging/simple/, dataset/interferometer/simple/, and dataset/point_source/simple/*.fits stay regenerable, matching _developer's policy.
• Write likelihood/README.md (replacing the placeholder) and four per-section READMEs at likelihood/{imaging,interferometer,point_source,datacube}/README.md — narrative on what each likelihood path profiles, how to read script output, where results land, and a placeholder run-time table that Phase 4 will auto-populate.
• Smoke-check: run one script from each of the four subfolders end-to-end and confirm JSON+PNG summaries land under results/likelihood/<type>/.

Detailed implementation plan

Affected Repositories

• PyAutoLabs/autolens_profiling (primary, target) — Phase 0 just shipped at 0087d6a
• PyAutoLabs/autolens_workspace_developer (read-only source-of-truth, no changes)

Work Classification

Scripts mirroring — not a standard library nor a standard workspace. The new repo lacks the /start_workspace scaffolding (no smoke_tests.txt, no activate.sh, no config/). Treat operationally as workspace-style work (run scripts to validate), but the smoke step is the ad-hoc per-subfolder run described below rather than a configured smoke harness.

Branch Survey

Repository	Current Branch	Dirty?
./autolens_profiling	main	clean (just bootstrapped)

Suggested branch: feature/likelihood-jit-mirror
Worktree root: ~/Code/PyAutoLabs-wt/likelihood-jit-mirror/ (created later by the routing skill)

Implementation Steps

1. Mirror the dataset files — copy _developer/jax_profiling/dataset/{imaging,interferometer,point_source}/<instrument>/ directories into autolens_profiling/dataset/<type>/<instrument>/. ~900K, 14 tracked files. Verify checksums after copy.

2. Update .gitignore in the new repo to match _developer's policy on simulated/regenerable datasets but with the new layout's path prefixes (dataset/cluster/ instead of jax_profiling/dataset/cluster/, etc.).

3. Create the section folder skeleton — likelihood/{imaging,interferometer,point_source,datacube}/ directories.

4. Mirror the scripts — copy each of the 10 files preserving names:

   • likelihood/imaging/{mge,pixelization,delaunay}.py
   • likelihood/interferometer/{mge,pixelization,delaunay}.py
   • likelihood/point_source/{image_plane,source_plane}.py
   • likelihood/datacube/delaunay.py + __init__.py

5. Rewrite paths in each script — search for Path(\"jax_profiling\") and rewrite. Two patterns to handle:

   • Dataset input: Path(\"jax_profiling\") / \"dataset\" / <type> / <instrument> → Path(\"dataset\") / <type> / <instrument>
   • Results output: Path(\"jax_profiling\") / \"results\" / \"jit\" / <type> → Path(\"results\") / \"likelihood\" / <type>
   • Keep the al.__version__ derivation for the artifact version string.
   • Keep all xp parameter usage, .array extraction, and step-by-step print() narrative unchanged.

6. Replace likelihood/README.md placeholder with the real section README — what JIT likelihood profiling means, how to read the per-script step-by-step output, the versioned artifact convention, and the 4-section navigation.

7. Write per-section READMEs at likelihood/{imaging,interferometer,point_source,datacube}/README.md — each:

   • 2-3 paragraphs on what the likelihood path looks like for that data type / model composition.
   • Table of script files with one-line summaries.
   • Run-time table placeholder ("populated by Phase 4 auto-generation"), with the column headers (instrument × hardware) already in place.

8. Smoke check — from the repo root, run one representative script per subfolder:

   • python likelihood/imaging/mge.py
   • python likelihood/interferometer/mge.py
   • python likelihood/point_source/image_plane.py
   • python likelihood/datacube/delaunay.py

   Confirm JSON+PNG summaries land at results/likelihood/<type>/<script>_likelihood_summary_<instrument>_v<al.__version__>.{json,png} with the expected step-by-step stdout.

9. Open PR against main of PyAutoLabs/autolens_profiling. PR body: link this issue + link to Phase 0 commit + summary of which scripts moved.

Key Files

• Source (read-only): autolens_workspace_developer/jax_profiling/jit/{imaging,interferometer,point_source,datacube}/*.py — copied verbatim then path-rewritten.
• Source (read-only): autolens_workspace_developer/jax_profiling/dataset/{imaging,interferometer,point_source}/<instrument>/* — copied verbatim.
• New: autolens_profiling/likelihood/{imaging,interferometer,point_source,datacube}/*.py — mirrored scripts.
• New: autolens_profiling/dataset/{imaging,interferometer,point_source}/<instrument>/* — mirrored datasets.
• Edit: autolens_profiling/.gitignore — add dataset/cluster/, dataset/group/, dataset/multi/, dataset/imaging/simple/, dataset/interferometer/simple/, dataset/point_source/simple/*.fits.
• Edit: autolens_profiling/likelihood/README.md — replace placeholder with real section README.
• New: autolens_profiling/likelihood/{imaging,interferometer,point_source,datacube}/README.md — 4 per-section READMEs.

Documented decision: dataset folder is top-level

Datasets live at autolens_profiling/dataset/<type>/<instrument>/ (top-level), not at likelihood/dataset/. Phase 2 (simulators) and Phase 3 (Nautilus) will consume the same input data; a shared top-level folder avoids duplication. Locked in at issue-filing time.

Original Prompt

Click to expand starting prompt

Phase 1 of the autolens_profiling z_feature
(see z_features/autolens_profiling.md for the full roadmap).

Mirror the JIT-profiling likelihood scripts currently in
autolens_workspace_developer/jax_profiling/jit/ into the new
autolens_profiling/likelihood/ directory. _developer remains the
source of truth — do not delete the originals.

Scripts to mirror (one-to-one, preserving file names):

• jit/imaging/{mge,pixelization,delaunay}.py
• jit/interferometer/{mge,pixelization,delaunay}.py
• jit/point_source/{image_plane,source_plane}.py
• jit/datacube/delaunay.py (+ its __init__.py)

Adjustments per-script:

• Update any path references that assume the _developer workspace layout
  (e.g. Path(\"jax_profiling\") / \"imaging\" / \"dataset\" / ...) to the new
  repo's layout (e.g. Path(\"likelihood\") / \"imaging\" / \"dataset\" / ...).
• Make sure results/ writes use the same versioned-artifact naming the
  _developer repo already uses (*_summary_v<version>.{json,png}) so the
  Phase 4 README dashboard can pick them up.
• Keep the printed step-by-step profile narrative each script already
  produces — it is the user-facing documentation. The user has been
  explicit: "make sure you keep that".

Per-section READMEs the implementer should author at the same time:

• likelihood/README.md — overview of the section, what "JIT likelihood
  profiling" means, how to read the per-script output, and how the
  versioned result artifacts work.
• likelihood/imaging/README.md, likelihood/interferometer/README.md,
  likelihood/point_source/README.md, likelihood/datacube/README.md —
  each with a short narrative on what the profiled likelihood path looks
  like and a table of the most-recent run-times for the scripts in that
  folder. Table population can be a placeholder filled in by Phase 4's
  auto-generation, but the section structure should already be there.

Dataset handling: _developer/jax_profiling/dataset/ provides the input
datasets these scripts depend on. Decide whether to mirror those files
into likelihood/dataset/, symlink, or document a one-time make dataset step. Lean toward mirroring the small simulated datasets so the
new repo is fully self-contained — but check sizes first before
committing binary blobs.

Smoke check before declaring done: run one script from each subfolder
end-to-end from the new repo (cd autolens_profiling && python likelihood/imaging/mge.py etc.) and confirm the JSON / PNG summaries
land under autolens_profiling/results/.

Out of scope: anything from _developer/jax_profiling/gradient/ (still
gated on the JAX gradient story stabilising) or _developer/jax_profiling/misc/
(exploratory work, not part of the stable profiling surface).