feat: scripts/model_composition/multi_galaxy_mge.py port (task 2/9)

[Jammy2211]

Overview

Port multi_galaxy_mge.py from autolens_workspace_test/scripts/model_composition/ to
autogalaxy_workspace_test/scripts/model_composition/, stripping it to autogalaxy semantics
(multiple galaxies sharing one plane with MGE light bases — no lens mass, no external shear,
no ray-tracing). This is task 2/9 of the autogalaxy_workspace_test parity epic
(#5).

Plan

• Port multi_galaxy_mge.py from autolens_workspace_test, stripping to autogalaxy semantics (multiple galaxies with MGE light bases, no mass / shear / ray-tracing).
• Two galaxies sharing one plane — galaxy_0 with gaussian_per_basis=2 (6 priors), galaxy_1 with gaussian_per_basis=1 (4 priors). Full model: 10 priors total.
• Keep the structural assertions (within-basis prior identity, cross-basis ell_comps independence, gaussian_per_basis variants, spherical, centre_per_basis, centre_fixed) — they exercise ag.model_util.mge_model_from directly and are profile-only.
• Drop autolens-specific keywords (Isothermal, ExternalShear) from the model.info keyword check.
• Generate a fresh identifier-stability regression anchor specific to the autogalaxy model.
• Wire into smoke_tests.txt so CI exercises it on every push.

Detailed implementation plan

Affected Repositories

• autogalaxy_workspace_test (primary)

Work Classification

Workspace

Branch Survey

Repository	Current Branch	Dirty?	Notes
./autogalaxy_workspace_test	main	dirty (3 modified dataset/*/jax_test/galaxies.json + untracked test_report.md)	Unrelated to this task. The worktree gets a clean checkout, so dirty main does not block.

Suggested branch: feature/autogalaxy-wst-model-composition
Worktree root: ~/Code/PyAutoLabs-wt/autogalaxy-wst-model-composition/ (created later by /start_workspace)

Implementation Steps

1. Create autogalaxy_workspace_test/scripts/model_composition/__init__.py (empty, mirrors autolens version).
2. Create autogalaxy_workspace_test/scripts/model_composition/multi_galaxy_mge.py:
   • import autofit as af + import autogalaxy as ag (drop al).
   • Build bulge_0 = ag.model_util.mge_model_from(..., gaussian_per_basis=2, centre_prior_is_uniform=True) → assert prior_count == 6.
   • Build galaxy_0 = af.Model(ag.Galaxy, redshift=0.5, bulge=bulge_0) → assert prior_count == 6.
   • Build bulge_1 = ag.model_util.mge_model_from(..., centre_prior_is_uniform=False) → assert prior_count == 4.
   • Build galaxy_1 = af.Model(ag.Galaxy, redshift=0.5, bulge=bulge_1) → assert prior_count == 4.
   • model = af.Collection(galaxies=af.Collection(galaxy_0=galaxy_0, galaxy_1=galaxy_1)) → assert prior_count == 10, 10 unique paths, top-level galaxies, second-level galaxy_0/galaxy_1.
   • Keep the within-basis identity / cross-basis independence checks verbatim against bulge_0.
   • Update model.info keywords to ["Galaxy", "Basis"] (drop Isothermal, ExternalShear).
   • Identifier-stability anchor: emit a placeholder, then run the script once locally, paste the resulting hash, re-run to confirm.
   • Keep gaussian_per_basis variants, spherical, centre_per_basis, centre_fixed sub-tests as-is — they call ag.model_util.mge_model_from only.
3. Append model_composition/multi_galaxy_mge.py to autogalaxy_workspace_test/smoke_tests.txt (placed between aggregator/ellipse.py and jax_likelihood_functions/imaging/lp.py so it runs before the JAX block).
4. Local verify: PYAUTO_TEST_MODE=2 PYAUTO_WORKSPACE_SMALL_DATASETS=1 python scripts/model_composition/multi_galaxy_mge.py.

Key Files

• autogalaxy_workspace_test/scripts/model_composition/__init__.py — new (empty, namespace marker).
• autogalaxy_workspace_test/scripts/model_composition/multi_galaxy_mge.py — new port from autolens.
• autogalaxy_workspace_test/smoke_tests.txt — one-line addition.

Risks & Notes

• Identifier hash unknown until first run. The autolens identifier 5a3c480de681f6958048b22b3db8ecf9 will not match — the model is structurally different (no mass, no shear, two light-only galaxies sharing a redshift). Emit on first run and lock it in.
• The dirty dataset/*/jax_test/galaxies.json files on main are unrelated and not part of this task. The worktree starts from origin/main so they don't follow.
• PR fix: rename stale PYAUTOFIT_TEST_MODE / PYAUTO_WORKSPACE_SMALL_DATASETS in build config #21 renamed PYAUTOFIT_TEST_MODE → PYAUTO_TEST_MODE. The verify command in the original prompt should be updated accordingly.

Original Prompt

Click to expand starting prompt

Create scripts/model_composition/ in @autogalaxy_workspace_test with an autogalaxy version of
the autolens multi_galaxy_mge.py script.

Reference

@autolens_workspace_test/scripts/model_composition/multi_galaxy_mge.py

That script exercises PyAutoFit's model-composition machinery on an autolens lens-plus-source
model with MGE bases. Strip it to an autogalaxy equivalent — multiple al.Galaxy objects in a
single plane (no lens/source split, no ray-tracing), each composed from an MGE light basis.

Deliverables

1. autogalaxy_workspace_test/scripts/model_composition/__init__.py
2. autogalaxy_workspace_test/scripts/model_composition/multi_galaxy_mge.py — ported script. Use
   ag.Galaxy, ag.Galaxies, ag.ImagingAnalysis (not al. / Tracer). Pick a small dataset
   from an existing autogalaxy workspace example and model two galaxies each with an MGE.
3. Append model_composition/multi_galaxy_mge.py to smoke_tests.txt.
4. Verify locally with PYAUTOFIT_TEST_MODE=2 PYAUTO_WORKSPACE_SMALL_DATASETS=1 python scripts/model_composition/multi_galaxy_mge.py.

Depends on

Task 1 (CI) must have merged so the new script runs in GitHub Actions. Review failure is cheaper
on CI than locally once CI is wired up.

Umbrella issue

Task 2/9. Track under the epic issue on PyAutoLabs/autogalaxy_workspace_test.

[Jammy2211]

Shipped — 2026-05-05

PRs

• feat: scripts/model_composition/multi_galaxy_mge.py port (task 2/9) #27 (merged)

Summary

Ported multi_galaxy_mge.py from autolens_workspace_test, stripped to autogalaxy semantics (two galaxies sharing one plane with MGE light bases, no mass / shear / ray-tracing). Full model has 10 priors; identifier regression anchor locked to a6eb928ed9a1fb92d0c18cf5443af4a6. Script passes on 3.12 (12.0s) and 3.13 (9.8s).

Session Notes

• The PYAUTO_SMALL_DATASETS=1 smoke default silently reduces total_gaussians inside ag.model_util.mge_model_from, which collapses gaussian_per_basis=2 to 1 and breaks the structural prior_count assertions. Required adding a model_composition/ override to config/build/env_vars.yaml (mirrors the existing autolens override). The original prompt's verify command used the legacy var name PYAUTO_WORKSPACE_SMALL_DATASETS which is no longer respected, so a naive local run masks the issue — the smoke runner is what surfaces it.
• The umbrella issue (epic: expand autogalaxy_workspace_test to mirror autolens_workspace_test coverage #5) checkboxes were stale: tasks 4 and 5 were marked unshipped despite PRs feat: jax_likelihood_functions/interferometer/ port #17 and feat: jax_likelihood_functions/multi/ port #19 already merging. Updated as part of this work.
• Remaining tasks in the parity epic, in numerical order: 6/9 (jax_grad/imaging), 7/9 (jax_grad/interferometer), 8/9 (jax_grad/multi).
• 3.13 has broader pre-existing JAX smoke failures than 3.12 (delaunay/multi suite) — not blocking for this task but worth tracking separately.