Merge pull request #272 from PyAutoLabs/feature/transform-decorator

refactor: reduce magic in transform decorator

Author: Jammy2211

autoarray/abstract_ndarray.py

@@ -76,6 +76,14 @@ def __init__(self, array, xp=np):
 
         self.use_jax = xp is not np
 
+    @property
+    def is_transformed(self) -> bool:
+        return self._is_transformed
+
+    @is_transformed.setter
+    def is_transformed(self, value: bool):
+        self._is_transformed = value
+
     @property
     def _xp(self):
         if self.use_jax:

autoarray/structures/decorators/transform.py

@@ -1,67 +1,117 @@
-from functools import wraps
-import numpy as np
-from typing import Union
-
-from autoarray.structures.grids.uniform_1d import Grid1D
-from autoarray.structures.grids.irregular_2d import Grid2DIrregular
-from autoarray.structures.grids.uniform_2d import Grid2D
-
-
-def transform(func):
-    """
-    Checks whether the input Grid2D of (y,x) coordinates have previously been transformed. If they have not
-    been transformed then they are transformed.
-
-    Parameters
-    ----------
-    func
-        A function where the input grid is the grid whose coordinates are transformed.
-
-    Returns
-    -------
-        A function that can accept cartesian or transformed coordinates
-    """
-
-    @wraps(func)
-    def wrapper(
-        obj: object,
-        grid: Union[np.ndarray, Grid2D, Grid2DIrregular, Grid1D],
-        xp=np,
-        *args,
-        **kwargs,
-    ) -> Union[np.ndarray, Grid2D, Grid2DIrregular]:
-        """
-        This decorator checks whether the input grid has been transformed to the reference frame of the class
-        that owns the function. If it has not been transformed, it is transformed.
-
-        A function call which uses this decorator often has many subsequent function calls which also use the
-        decorator. To ensure the grid is only transformed once, the `is_transformed` keyword is used to track
-        whether the grid has been transformed.
-
-        Parameters
-        ----------
-        obj
-            An object whose function uses grid_like inputs to compute quantities at every coordinate on the grid.
-        grid
-            The (y, x) coordinates in the original reference frame of the grid.
-
-        Returns
-        -------
-            A grid_like object whose coordinates may be transformed.
-        """
-
-        if not kwargs.get("is_transformed"):
-            kwargs["is_transformed"] = True
-
-            transformed_grid = obj.transformed_to_reference_frame_grid_from(
-                grid, xp, **kwargs
-            )
-
-            result = func(obj, transformed_grid, xp, *args, **kwargs)
-
-        else:
-            result = func(obj, grid, xp, *args, **kwargs)
-
-        return result
-
-    return wrapper
+from functools import wraps
+import numpy as np
+from typing import Union
+
+from autoarray.structures.grids.uniform_1d import Grid1D
+from autoarray.structures.grids.irregular_2d import Grid2DIrregular
+from autoarray.structures.grids.uniform_2d import Grid2D
+
+
+def transform(func=None, *, rotate_back=False):
+    """
+    Checks whether the input Grid2D of (y,x) coordinates have previously been transformed. If they have not
+    been transformed then they are transformed.
+
+    Can be used with or without arguments::
+
+        @transform
+        def convergence_2d_from(self, grid, xp=np, **kwargs): ...
+
+        @transform(rotate_back=True)
+        def deflections_yx_2d_from(self, grid, xp=np, **kwargs): ...
+
+    **Frame conventions and rotate_back**
+
+    This decorator transforms the input grid into the profile's reference frame (centred on the
+    profile centre and rotated by its position angle) before calling the decorated function.
+
+    For **scalar** quantities (convergence, potential), the returned value is frame-invariant — no
+    back-rotation is needed, so use ``@transform`` without ``rotate_back``.
+
+    For **vector** quantities (e.g. deflection angles), whether back-rotation is needed depends on
+    which frame the returned components are expressed in:
+
+    - If the function computes vector components using the rotated grid coordinates (i.e. the
+      components are expressed in the profile's frame), they must be rotated back to the observer
+      frame before use in ray-tracing. Set ``rotate_back=True`` for this case.
+
+    - If the function reconstructs observer-frame components from scalar quantities (e.g. computing
+      a radial deflection magnitude and converting to Cartesian using observer-frame geometry), the
+      result is already in the observer frame and ``rotate_back`` should remain ``False``.
+
+    When ``rotate_back=True``, the decorator calls ``obj.rotated_grid_from_reference_frame_from``
+    on the result after evaluation, applying the inverse rotation by the profile's position angle.
+
+    For **spin-2** quantities (shear), the transformation law uses twice the profile angle. This
+    is not handled by ``rotate_back`` — shear methods must apply the 2-theta rotation manually.
+
+    Parameters
+    ----------
+    func
+        A function where the input grid is the grid whose coordinates are transformed.
+    rotate_back
+        If ``True``, the result is rotated back from the profile's reference frame after
+        evaluation. Use this when the decorated function returns vector components that were
+        computed in the profile's rotated coordinate basis and need to be expressed in the
+        original observer frame.
+
+    Returns
+    -------
+        A function that can accept cartesian or transformed coordinates
+    """
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(
+            obj: object,
+            grid: Union[np.ndarray, Grid2D, Grid2DIrregular, Grid1D],
+            xp=np,
+            *args,
+            **kwargs,
+        ) -> Union[np.ndarray, Grid2D, Grid2DIrregular]:
+            """
+            This decorator checks whether the input grid has been transformed to the reference frame of the class
+            that owns the function. If it has not been transformed, it is transformed.
+
+            The transform state is tracked via the ``is_transformed`` property on the grid object itself.
+            When a decorated function calls another decorated function with the same (already-transformed)
+            grid, the flag prevents the grid from being transformed a second time.
+
+            Parameters
+            ----------
+            obj
+                An object whose function uses grid_like inputs to compute quantities at every coordinate on the grid.
+            grid
+                The (y, x) coordinates in the original reference frame of the grid.
+
+            Returns
+            -------
+                A grid_like object whose coordinates may be transformed.
+            """
+
+            if not getattr(grid, "is_transformed", False):
+                transformed_grid = obj.transformed_to_reference_frame_grid_from(
+                    grid, xp, **kwargs
+                )
+                transformed_grid.is_transformed = True
+
+                result = func(obj, transformed_grid, xp, *args, **kwargs)
+
+            else:
+                result = func(obj, grid, xp, *args, **kwargs)
+
+            if rotate_back:
+                result = obj.rotated_grid_from_reference_frame_from(
+                    grid=result, xp=xp
+                )
+
+            return result
+
+        return wrapper
+
+    if func is not None:
+        # Called without arguments: @transform
+        return decorator(func)
+
+    # Called with arguments: @transform(rotate_back=True)
+    return decorator