rename OperateDeflections to LensCalc, move fermat_potential_from into LensCalc

- Rename operate/deflections.py -> operate/lens_calc.py
- Rename class OperateDeflections -> LensCalc throughout
- LensCalc.__init__ now accepts optional potential_2d_from callable
- from_mass_obj and from_tracer capture potential_2d_from automatically
- fermat_potential_from moved into LensCalc (removed from MassProfile, Galaxy, Galaxies)
- Update all imports, call sites, tests, and docs in autogalaxy

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 2 people

autogalaxy/__init__.py

@@ -69,7 +69,7 @@
 from .operate.image import OperateImage
 from .operate.image import OperateImageList
 from .operate.image import OperateImageGalaxies
-from .operate.deflections import OperateDeflections
+from .operate.lens_calc import LensCalc
 from .gui.scribbler import Scribbler
 from .imaging.fit_imaging import FitImaging
 from .imaging.model.analysis import AnalysisImaging

autogalaxy/galaxy/galaxies.py

@@ -243,34 +243,6 @@ def potential_2d_from(
         """
         return sum(map(lambda g: g.potential_2d_from(grid=grid, xp=xp), self))
 
-    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
-        """
-        Returns the Fermat potential of the galaxies for a given grid of image-plane positions.
-
-        This is the sum of the geometric time delay term and the lensing potential, computed
-        using the combined deflections and potential of all galaxies:
-
-        .. math::
-            \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2
-            - \\psi(\\boldsymbol{\\theta})
-
-        Parameters
-        ----------
-        grid
-            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
-        xp
-            The array module (``numpy`` or ``jax.numpy``).
-        """
-        from autogalaxy.operate.deflections import OperateDeflections
-
-        od = OperateDeflections.from_mass_obj(self)
-        time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp)
-        potential = self.potential_2d_from(grid=grid, xp=xp)
-        fermat_potential = time_delay - potential
-        if isinstance(grid, aa.Grid2DIrregular):
-            return aa.ArrayIrregular(values=fermat_potential)
-        return aa.Array2D(values=fermat_potential, mask=grid.mask)
-
     def has(self, cls: Union[Type, Tuple[Type]]) -> bool:
         """
         Returns a bool specifying whether any of the galaxies has a certain class type.

autogalaxy/galaxy/galaxy.py

@@ -343,34 +343,6 @@ def potential_2d_from(
             )
         return xp.zeros((grid.shape[0],))
 
-    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
-        """
-        Returns the Fermat potential of the galaxy for a given grid of image-plane positions.
-
-        This is the sum of the geometric time delay term and the lensing potential, computed
-        using the galaxy's total deflections and total potential across all mass profiles:
-
-        .. math::
-            \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\alpha}|^2
-            - \\psi(\\boldsymbol{\\theta})
-
-        Parameters
-        ----------
-        grid
-            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
-        xp
-            The array module (``numpy`` or ``jax.numpy``).
-        """
-        from autogalaxy.operate.deflections import OperateDeflections
-
-        od = OperateDeflections.from_mass_obj(self)
-        time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp)
-        potential = self.potential_2d_from(grid=grid, xp=xp)
-        fermat_potential = time_delay - potential
-        if isinstance(grid, aa.Grid2DIrregular):
-            return aa.ArrayIrregular(values=fermat_potential)
-        return aa.Array2D(values=fermat_potential, mask=grid.mask)
-
     @property
     def half_light_radius(self):
         return None

autogalaxy/operate/deflections.py autogalaxy/operate/lens_calc.pyautogalaxy/operate/deflections.py renamed to autogalaxy/operate/lens_calc.py

@@ -78,28 +78,40 @@ def wrapper(
     return wrapper
 
 
-class OperateDeflections:
+class LensCalc:
     """
-    Packages methods which manipulate the 2D deflection angle map returned from the `deflections_yx_2d_from` function
-    of a mass object (e.g. a `MassProfile`, `Galaxy`).
+    Computes lensing quantities from a deflection-angle callable and an optional potential callable.
 
-    The majority of methods are those which from the 2D deflection angle map compute lensing quantities like a 2D
-    shear field, magnification map or the Einstein Radius.
+    The deflection callable is used to compute the Hessian, Jacobian, convergence, shear,
+    magnification, critical curves, caustics, and Einstein radius/mass.  If a potential
+    callable is also supplied, ``fermat_potential_from`` is available as well.
 
     Parameters
     ----------
     deflections_yx_2d_from
-        A callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D deflection angles on the given
-        grid.  Typically a bound method of a ``MassProfile``, ``Galaxy``, or ``Galaxies`` instance.
+        A callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D deflection
+        angles on the given grid.  Typically a bound method of a ``MassProfile``, ``Galaxy``,
+        or ``Galaxies`` instance.
+    potential_2d_from
+        Optional callable with signature ``(grid, xp=np, **kwargs)`` that returns the 2D
+        lensing potential on the given grid.  Required only for ``fermat_potential_from``.
     """
 
-    def __init__(self, deflections_yx_2d_from):
+    def __init__(self, deflections_yx_2d_from, potential_2d_from=None):
         self.deflections_yx_2d_from = deflections_yx_2d_from
+        self.potential_2d_from = potential_2d_from
 
     @classmethod
     def from_mass_obj(cls, mass_obj):
-        """Construct from any object that has a ``deflections_yx_2d_from`` method."""
-        return cls(deflections_yx_2d_from=mass_obj.deflections_yx_2d_from)
+        """Construct from any object that has a ``deflections_yx_2d_from`` method.
+
+        If the object also exposes ``potential_2d_from``, it is captured so that
+        ``fermat_potential_from`` is available on the returned instance.
+        """
+        return cls(
+            deflections_yx_2d_from=mass_obj.deflections_yx_2d_from,
+            potential_2d_from=getattr(mass_obj, "potential_2d_from", None),
+        )
 
     @classmethod
     def from_tracer(
@@ -130,6 +142,8 @@ def from_tracer(
             Index of the second plane used by ``deflections_between_planes_from``.
             Ignored when ``use_multi_plane=False``.  Defaults to ``-1`` (source plane).
         """
+        potential_2d_from = getattr(tracer, "potential_2d_from", None)
+
         if use_multi_plane:
             from functools import partial
 
@@ -138,9 +152,13 @@ def from_tracer(
                     tracer.deflections_between_planes_from,
                     plane_i=plane_i,
                     plane_j=plane_j,
-                )
+                ),
+                potential_2d_from=potential_2d_from,
             )
-        return cls(deflections_yx_2d_from=tracer.deflections_yx_2d_from)
+        return cls(
+            deflections_yx_2d_from=tracer.deflections_yx_2d_from,
+            potential_2d_from=potential_2d_from,
+        )
 
     def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D:
         """
@@ -177,6 +195,39 @@ def time_delay_geometry_term_from(self, grid, xp=np) -> aa.Array2D:
             return aa.ArrayIrregular(values=delay)
         return aa.Array2D(values=delay, mask=grid.mask)
 
+    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
+        """
+        Returns the Fermat potential for a given grid of image-plane positions.
+
+        This is the sum of the geometric time delay term and the gravitational (Shapiro) delay
+        term (i.e. the lensing potential), and is given by:
+
+        .. math::
+            \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\theta} - \\boldsymbol{\\beta}|^2
+            - \\psi(\\boldsymbol{\\theta})
+
+        Requires that ``potential_2d_from`` was supplied at construction (e.g. via
+        ``LensCalc.from_mass_obj`` or ``LensCalc.from_tracer``).
+
+        Parameters
+        ----------
+        grid
+            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
+        xp
+            The array module (``numpy`` or ``jax.numpy``).
+        """
+        if self.potential_2d_from is None:
+            raise ValueError(
+                "fermat_potential_from requires a potential_2d_from callable. "
+                "Construct LensCalc with potential_2d_from, or use from_mass_obj / from_tracer."
+            )
+        time_delay = self.time_delay_geometry_term_from(grid=grid, xp=xp)
+        potential = self.potential_2d_from(grid=grid, xp=xp)
+        fermat_potential = time_delay - potential
+        if isinstance(grid, aa.Grid2DIrregular):
+            return aa.ArrayIrregular(values=fermat_potential)
+        return aa.Array2D(values=fermat_potential, mask=grid.mask)
+
     def tangential_eigen_value_from(self, grid, xp=np) -> aa.Array2D:
         """
         Returns the tangential eigen values of lensing jacobian, which are given by the expression:

autogalaxy/plot/mass_plotter.py

@@ -118,10 +118,10 @@ def figures_2d(
             )
 
         if magnification:
-            from autogalaxy.operate.deflections import OperateDeflections
+            from autogalaxy.operate.lens_calc import LensCalc
 
             self.mat_plot_2d.plot_array(
-                array=OperateDeflections.from_mass_obj(
+                array=LensCalc.from_mass_obj(
                     self.mass_obj
                 ).magnification_2d_from(grid=self.grid),
                 visuals_2d=self.visuals_2d_with_critical_curves,

autogalaxy/plot/visuals/one_d.py

@@ -170,12 +170,12 @@ def add_einstein_radius(
             The collection of attributes that can be plotted by a `Plotter` object.
         """
 
-        from autogalaxy.operate.deflections import OperateDeflections
+        from autogalaxy.operate.lens_calc import LensCalc
 
         einstein_radius = None
 
         try:
-            od = OperateDeflections.from_mass_obj(mass_obj)
+            od = LensCalc.from_mass_obj(mass_obj)
             einstein_radius = od.einstein_radius_from(grid=grid)
         except (TypeError, AttributeError):
             pass
@@ -217,13 +217,13 @@ def add_einstein_radius_errors(
             The mean value and errors of each attribute that are plotted in 1D by a `Plotter` object.
         """
 
-        from autogalaxy.operate.deflections import OperateDeflections
+        from autogalaxy.operate.lens_calc import LensCalc
 
         einstein_radius_list = []
 
         for mass_obj in mass_obj_list:
             try:
-                od = OperateDeflections.from_mass_obj(mass_obj)
+                od = LensCalc.from_mass_obj(mass_obj)
                 einstein_radius_list.append(od.einstein_radius_from(grid=grid))
             except TypeError:
                 einstein_radius_list.append(None)

autogalaxy/plot/visuals/two_d.py

@@ -184,9 +184,9 @@ def add_critical_curves(self, mass_obj, grid: aa.type.Grid2DLike):
         vis.Visuals2D
             A collection of attributes that can be plotted by a `Plotter` object.
         """
-        from autogalaxy.operate.deflections import OperateDeflections
+        from autogalaxy.operate.lens_calc import LensCalc
 
-        od = OperateDeflections.from_mass_obj(mass_obj)
+        od = LensCalc.from_mass_obj(mass_obj)
 
         tangential_critical_curves = od.tangential_critical_curve_list_from(grid=grid)
 
@@ -227,9 +227,9 @@ def add_caustics(self, mass_obj, grid: aa.type.Grid2DLike):
         vis.Visuals2D
             A collection of attributes that can be plotted by a `Plotter` object.
         """
-        from autogalaxy.operate.deflections import OperateDeflections
+        from autogalaxy.operate.lens_calc import LensCalc
 
-        od = OperateDeflections.from_mass_obj(mass_obj)
+        od = LensCalc.from_mass_obj(mass_obj)
 
         tangential_caustics = od.tangential_caustic_list_from(grid=grid)
         radial_caustics = od.radial_caustic_list_from(grid=grid)

autogalaxy/profiles/mass/abstract/abstract.py

@@ -27,34 +27,6 @@ def __init__(
     def deflections_yx_2d_from(self, grid):
         raise NotImplementedError
 
-    def fermat_potential_from(self, grid, xp=np) -> aa.Array2D:
-        """
-        Returns the Fermat potential for a given grid of image-plane positions.
-
-        This is the sum of the geometric time delay term and the gravitational (Shapiro) delay term (i.e. the lensing
-        potential), and is given by:
-
-        .. math::
-            \\phi(\\boldsymbol{\\theta}) = \\frac{1}{2} |\\boldsymbol{\\theta} - \\boldsymbol{\\beta}|^2
-            - \\psi(\\boldsymbol{\\theta})
-
-        Parameters
-        ----------
-        grid
-            The 2D grid of (y,x) arc-second coordinates the Fermat potential is computed on.
-        xp
-            The array module (``numpy`` or ``jax.numpy``).
-        """
-        from autogalaxy.operate.deflections import OperateDeflections
-
-        od = OperateDeflections.from_mass_obj(self)
-        time_delay = od.time_delay_geometry_term_from(grid=grid, xp=xp)
-        potential = self.potential_2d_from(grid=grid, xp=xp)
-        fermat_potential = time_delay - potential
-        if isinstance(grid, aa.Grid2DIrregular):
-            return aa.ArrayIrregular(values=fermat_potential)
-        return aa.Array2D(values=fermat_potential, mask=grid.mask)
-
     def deflections_2d_via_potential_2d_from(self, grid):
         potential = self.potential_2d_from(grid=grid)
 

docs/api/source.rst

@@ -31,7 +31,7 @@ Operators
    :recursive:
 
    OperateImage
-   OperateDeflections
+   LensCalc
 
 Total [ag.mp]
 -------------