feat: AnalysisEllipse.fit_from + JAX pytree registration

Wires AnalysisEllipse for end-to-end JAX traceability. Adds use_jax
flag (default True, matching AnalysisImaging), a fit_from method
returning a new FitEllipseSummed aggregate, and a lazy idempotent
_register_fit_ellipse_pytrees that registers FitEllipse,
FitEllipseSummed, DatasetModel, Ellipse, EllipseMultipole, and
EllipseMultipoleScaled as JAX pytrees on the first fit_from call.

log_likelihood_function now returns figure_of_merit (sum over the
per-ellipse fit_list) instead of log_likelihood, matching
AnalysisImaging exactly. The two differ by the noise_normalization
term per FitEllipse: figure_of_merit = -0.5 * (chi_squared +
noise_normalization), log_likelihood = -0.5 * chi_squared.

Step 7 of 7 in z_features/ellipse_fitting_jax.md — the keystone.
After this lands, ellipse modeling runs inside any JAX-compatible
search via AnalysisEllipse(dataset, use_jax=True).

Workspace follow-up flips prompt 2's workspace_test scripts to a
JIT round-trip with rtol=1e-4 parity against the locked-in numpy
reference numbers.

Issue PyAutoGalaxy#411.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: 2 people

autogalaxy/__init__.py

@@ -65,6 +65,7 @@
 from .ellipse.ellipse.ellipse_multipole import EllipseMultipole
 from .ellipse.ellipse.ellipse_multipole import EllipseMultipoleScaled
 from .ellipse.fit_ellipse import FitEllipse
+from .ellipse.fit_ellipse import FitEllipseSummed
 from .ellipse.model.analysis import AnalysisEllipse
 from .operate.image import OperateImage
 from .operate.image import OperateImageList

autogalaxy/ellipse/fit_ellipse.py

@@ -11,6 +11,10 @@
 This is the classical isophote-fitting approach used in tools such as IRAF/ELLIPSE and galfit, and is
 appropriate for measuring galaxy morphology, position angles, axis ratios, and multipole perturbations
 directly from imaging data without fitting a parametric light profile model.
+
+`FitEllipseSummed` aggregates multiple `FitEllipse` objects (one per ellipse in a model) and exposes
+sum-over-list versions of `figure_of_merit`, `log_likelihood`, and `chi_squared`. This is the return
+type of `AnalysisEllipse.fit_from`, mirroring the single-object return of `AnalysisImaging.fit_from`.
 """
 import numpy as np
 from typing import List, Optional
@@ -333,3 +337,54 @@ def figure_of_merit(self) -> float:
         The figure of merit of the fit.
         """
         return -0.5 * (self.chi_squared + self.noise_normalization)
+
+
+class FitEllipseSummed:
+    """
+    Aggregate of one or more :class:`FitEllipse` objects whose
+    ``figure_of_merit`` / ``log_likelihood`` / ``chi_squared`` properties
+    sum over the contained fits.
+
+    Used by :class:`AnalysisEllipse.fit_from` so the return type is a
+    single object (matching :class:`AnalysisImaging.fit_from`'s pattern)
+    even when a model contains multiple ellipses. Each contained
+    :class:`FitEllipse` carries the same shared ``dataset``; this class
+    exposes that dataset for JAX pytree-registration purposes
+    (``no_flatten=("dataset",)``).
+    """
+
+    def __init__(self, fit_list: List[FitEllipse]):
+        self.fit_list = fit_list
+
+    @property
+    def dataset(self):
+        """
+        All fits in the list share the same dataset; expose the first
+        for pytree no_flatten purposes.
+        """
+        return self.fit_list[0].dataset
+
+    @property
+    def figure_of_merit(self):
+        """
+        The sum of the ``figure_of_merit`` values of every contained
+        :class:`FitEllipse`, used by the non-linear search as the
+        overall objective.
+        """
+        return sum(f.figure_of_merit for f in self.fit_list)
+
+    @property
+    def log_likelihood(self):
+        """
+        The sum of the ``log_likelihood`` values of every contained
+        :class:`FitEllipse`.
+        """
+        return sum(f.log_likelihood for f in self.fit_list)
+
+    @property
+    def chi_squared(self):
+        """
+        The sum of the ``chi_squared`` values of every contained
+        :class:`FitEllipse`.
+        """
+        return sum(f.chi_squared for f in self.fit_list)