feat: pytree registration for FitEllipse + FitQuantity (Phase 0c)

[Jammy2211]

Overview

Phase 0c of the JAX visualization roadmap. PyAutoGalaxy's FitImaging and FitInterferometer are JAX pytree-registered (PRs #364 and #376 respectively). The remaining gap is FitEllipse and FitQuantity — neither has pytree registration today, so use_jax_for_visualization=True on ag.AnalysisEllipse or ag.AnalysisQuantity would crash when fit_for_visualization tries to lift the fit across the JIT boundary.

This task ports the imaging / interferometer pytree pattern to those two remaining dataset types and adds round-trip JIT tests.

Plan

• Add _register_fit_ellipse_pytrees to ag.AnalysisEllipse mirroring _register_fit_interferometer_pytrees from PR feat: register pytrees for AnalysisInterferometer #376 (autogalaxy/interferometer/model/analysis.py:165-184).
• Add _register_fit_quantity_pytrees to ag.AnalysisQuantity mirroring the same pattern.
• Reuse the shared autogalaxy/analysis/jax_pytrees.py::register_galaxies_pytree() helper.
• For each fit class, identify dynamic (JAX-traced arrays / sub-pytrees) vs static (no_flatten aux: masks, pixel scales, redshifts, settings) attributes by reading the populated Fit* instance.
• Wire self._register_fit_*_pytrees() into each Analysis.__init__ under if use_jax:.
• If ag.AnalysisEllipse or ag.AnalysisQuantity lacks **kwargs passthrough on __init__, fix it inline (2-line addition, same pattern as fix: ag.AnalysisInterferometer.__init__ kwargs passthrough #399).
• Add round-trip JIT tests under test_autogalaxy/ellipse/jax/ and test_autogalaxy/quantity/jax/ following the three-step pattern from autolens_workspace_test/scripts/hessian_jax.py (build → tree_flatten/unflatten round-trip → jax.jit returning the fit).

Detailed implementation plan

Affected Repositories

• PyAutoGalaxy (primary, only)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoGalaxy	main	clean (recent merge from jax-interp-2d #398)

Suggested branch: feature/ag-ellipse-quantity-pytree
Worktree root: ~/Code/PyAutoLabs-wt/ag-ellipse-quantity-pytree/ (parallel-safe with in-flight nfw-jax-port mass-profile worktree — file-disjoint)

Implementation Steps

1. autogalaxy/ellipse/model/analysis.py: add _register_fit_ellipse_pytrees static method + invoke under if use_jax: in __init__. Register FitEllipse with no_flatten=("dataset", "settings", ...), plus DatasetModel if applicable, plus register_galaxies_pytree().

2. autogalaxy/quantity/model/analysis.py: same pattern for _register_fit_quantity_pytrees. Register FitQuantity with appropriate no_flatten aux, plus DatasetQuantity (its dataset analogue), plus register_galaxies_pytree().

3. If either analysis __init__ lacks **kwargs passthrough (parallel pattern to PR fix: ag.AnalysisInterferometer.__init__ kwargs passthrough #399), add 2-line fix.

4. test_autogalaxy/ellipse/jax/test_fit_ellipse_pytree.py (NEW) — three-step JIT round-trip test.

5. test_autogalaxy/quantity/jax/test_fit_quantity_pytree.py (NEW) — same.

6. pytest test_autogalaxy/ellipse/ test_autogalaxy/quantity/ — all existing tests must still pass.

Key Files

• autogalaxy/ellipse/model/analysis.py — add registration method + invoke under if use_jax
• autogalaxy/quantity/model/analysis.py — add registration method + invoke under if use_jax
• test_autogalaxy/ellipse/jax/test_fit_ellipse_pytree.py (NEW)
• test_autogalaxy/quantity/jax/test_fit_quantity_pytree.py (NEW)

Reference patterns

• autogalaxy/interferometer/model/analysis.py:165-184 — _register_fit_interferometer_pytrees (closest sibling)
• autogalaxy/imaging/model/analysis.py:169-187 — _register_fit_imaging_pytrees
• autogalaxy/analysis/jax_pytrees.py — shared register_galaxies_pytree() helper

Parallel worktree safety

Coexisting with nfw-jax-port (PyAutoGalaxy mass profiles) and interferometer-nufftax-updates (workspaces). This task touches autogalaxy/ellipse/ + autogalaxy/quantity/ only — file-disjoint from all in-flight work. User confirmed file-level safety.

Original Prompt

Click to expand starting prompt

(elided — full text in PyAutoPrompt/issued/fit_pytree_registration_other_datasets.md)

[Jammy2211]

Shipped — 2026-05-14

PR

• feat: pytree registration for FitEllipse + FitQuantity (Phase 0c) #401 (merged)

Summary

Phase 0c shipped. ag.AnalysisEllipse and ag.AnalysisQuantity now register FitEllipse + FitQuantity (plus Ellipse, EllipseMultipole, and reused Galaxies registration) as JAX pytrees when constructed with use_jax=True. Plus the **kwargs passthrough fix on both __init__ methods (parity with #399). 154/154 unit tests across test_autogalaxy/{ellipse,quantity,imaging,interferometer}/ pass; pytree round-trip verified interactively (FitEllipse: 8 dynamic leaves; FitQuantity: 3 dynamic leaves with Galaxies correctly reconstructed).

With this PR, all four PyAutoGalaxy Fit* types now have JAX pytree registration.

Session Notes

• Scope narrowed mid-task after audit. The original prompt assumed registering pytrees would unblock use_jax_for_visualization=True for these analyses. Reading the code revealed both visualizers bypass analysis.fit_for_visualization entirely: VisualizerEllipse.visualize calls analysis.fit_list_from(instance) (returns a list); VisualizerQuantity.visualize calls analysis.fit_quantity_for_instance(instance) (singular). So use_jax_for_visualization=True remains a no-op for these two analyses despite the pytree work — separate dispatch follow-ups needed.

• Two deferred follow-ups:

  • Quantity dispatch swap — small. Add fit_from = fit_quantity_for_instance (or wrapper) on AnalysisQuantity, switch VisualizerQuantity.visualize to use analysis.fit_for_visualization(instance). Mirrors the autolens / autogalaxy imaging + interferometer pattern. Unlocks use_jax_for_visualization=True end-to-end on quantity.
  • Ellipse dispatch swap — needs design. AnalysisEllipse.fit_list_from returns List[FitEllipse], not a single fit. Either generalize the autofit fit_for_visualization contract to accept list returns, or have ellipse-side compose its list into a wrapper. Likely a separate small design pass.

• Pytree no_flatten decisions:

  • Ellipse: no no_flatten (centre, ell_comps, major_axis are all fittable params).
  • EllipseMultipole: no_flatten=("m",) — m is the multipole order (integer discriminator, not fittable).
  • FitEllipse: no_flatten=("dataset",) — dataset is per-analysis-constant.
  • FitQuantity: no_flatten=("dataset", "func_str", "use_mask_in_fit") — all three are per-analysis-constant or static config.

• Parallel-worktree safety confirmed. Coexisted with in-flight nfw-jax-port (PyAutoGalaxy mass profiles) — file-disjoint from autogalaxy/ellipse/ and autogalaxy/quantity/. Worktree-conflict guard bypassed per user-confirmed file-level safety, same pattern as PR fix: ag.AnalysisInterferometer.__init__ kwargs passthrough #399.

Library kwargs-gap series — COMPLETE

• al.AnalysisImaging — always had **kwargs ✓
• al.AnalysisInterferometer — fixed in fix: JAX-trace Sérsic stellar-mass CSE deflection path (#499) #500 ✓
• al.AnalysisPoint — fixed in dPIE: make Lenstool-native parameterization the default profile #506 ✓
• ag.AnalysisImaging — always had **kwargs ✓
• ag.AnalysisInterferometer — fixed in fix: ag.AnalysisInterferometer.__init__ kwargs passthrough #399 ✓
• ag.AnalysisEllipse + ag.AnalysisQuantity — fixed in this PR (feat: pytree registration for FitEllipse + FitQuantity (Phase 0c) #401) ✓

PyAutoGalaxy pytree series — COMPLETE

• ag.FitImaging — PR feat: register FitImaging, DatasetModel, Galaxies as pytrees in AnalysisImaging #364 ✓
• ag.FitInterferometer — PR feat: register pytrees for AnalysisInterferometer #376 ✓
• ag.FitEllipse + ag.FitQuantity — this PR (feat: pytree registration for FitEllipse + FitQuantity (Phase 0c) #401) ✓