Skip to content

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

New issue
New issue
Closed
#401
Closed
feat: pytree registration for FitEllipse + FitQuantity (Phase 0c)#400
#401

Description

@Jammy2211
Jammy2211
opened on May 14, 2026

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)

Activity

  1. Jammy2211
    closed this as completedin #401on May 14, 2026

  2. Jammy2211 commented on May 14, 2026

    @Jammy2211
    Jammy2211
    on May 14, 2026
    CollaboratorAuthor

    Shipped — 2026-05-14

    PR

    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

    PyAutoGalaxy pytree series — COMPLETE

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions