Merge pull request #277 from PyAutoLabs/feature/data-typing-simplify

refactor: simplify decorator type-dispatch system

Author: Jammy2211

autoarray/__init__.py

@@ -84,7 +84,8 @@
 from .structures.triangles.shape import Triangle
 from .structures.triangles.shape import Square
 from .structures.triangles.shape import Polygon
-from .structures import decorators as grid_dec
+from .structures import decorators as grid_dec  # deprecated alias
+from .structures import decorators
 from .structures.header import Header
 from .layout.region import Region1D
 from .layout.region import Region2D

autoarray/mock.py

@@ -20,3 +20,4 @@
 from autoarray.structures.mock.mock_grid import MockMeshGrid
 from autoarray.structures.mock.mock_decorators import MockGrid1DLikeObj
 from autoarray.structures.mock.mock_decorators import MockGrid2DLikeObj
+from autoarray.structures.mock.mock_decorators import MockTransformProfile

autoarray/structures/decorators/abstract.py

@@ -1,10 +1,6 @@
-from typing import Union
-
 import numpy as np
 
-from autoarray.mask.mask_1d import Mask1D
 from autoarray.mask.mask_2d import Mask2D
-from autoarray.structures.grids.uniform_1d import Grid1D
 from autoarray.structures.grids.irregular_2d import Grid2DIrregular
 from autoarray.structures.grids.uniform_2d import Grid2D
 
@@ -18,20 +14,16 @@ def __init__(self, func, obj, grid, xp=np, *args, **kwargs):
         This is used by the `to_array`, `to_grid` and `to_vector_yx` decorators to ensure that the input grid and output
         data structure are consistent.
 
-        There are three types of consistent data structures and therefore decorated function mappings:
+        There are two types of consistent data structures and therefore decorated function mappings:
 
         - Uniform: 2D structures defined on a uniform grid of data points, for example the `Array2D` and `Grid2D`
         objects. Both structures are defined according to a `Mask2D`, which the maker object ensures is passed through
         self consistently.
 
         - Irregular: 2D structures defined on an irregular grid of data points, for example an `ArrayIrregular`
-        and `Grid2DIrregular` objects. Neither structure is defined according to a mask and the maker sures the lack of
+        and `Grid2DIrregular` objects. Neither structure is defined according to a mask and the maker ensures the lack of
         a mask does not prevent the function from being evaluated.
 
-        - 1D: 1D structures defined on a 1D grid of data points, for example the `Array1D` and `Grid1D` objects.
-        These project the 1D grid to a 2D grid to ensure the function can be evaluated, and then deproject the 2D grid
-        back to a 1D grid to ensure the output data structure is consistent with the input grid.
-
         Parameters
         ----------
         func
@@ -66,7 +58,7 @@ def _xp(self):
         return np
 
     @property
-    def mask(self) -> Union[Mask1D, Mask2D]:
+    def mask(self) -> Mask2D:
         return self.grid.mask
 
     @property
@@ -79,30 +71,8 @@ def via_grid_2d(self, result):
     def via_grid_2d_irr(self, result):
         raise NotImplementedError
 
-    def via_grid_1d(self, result):
-        raise NotImplementedError
-
     @property
     def evaluate_func(self):
-        """
-        Evaluate the function that is being decorated, using the grid that is passed to the maker object when it is
-        initialized.
-
-        In normal usage, the input grid is 2D and it is simply passed to the decorated function.
-
-        However, if the input grid is 1D, the grid is projected to a 2D grid before being passed to the function. This
-        is because the function is expected to evaluate a 2D grid, and the maker object ensures that the function can
-        be evaluated by projecting the 1D grid to a 2D grid.
-
-        Returns
-        -------
-        The result of the function that is being decorated, which is the output data structure that is consistent with
-        the input grid.
-        """
-
-        if isinstance(self.grid, Grid1D):
-            grid = self.grid.grid_2d_radial_projected_from()
-            return self.func(self.obj, grid, self._xp, *self.args, **self.kwargs)
         return self.func(self.obj, self.grid, self._xp, *self.args, **self.kwargs)
 
     @property
@@ -111,21 +81,17 @@ def result(self):
         The result of the function that is being decorated, which this function converts to the output data structure
         that is consistent with the input grid.
 
-        This function called one of three methods, depending on the type of the input grid:
+        This function calls one of two methods, depending on the type of the input grid:
 
         - `via_grid_2d`: If the input grid is a `Grid2D` object.
         - `via_grid_2d_irr`: If the input grid is a `Grid2DIrregular` object.
-        - `via_grid_1d`: If the input grid is a `Grid1D` object.
 
-        These functions are over written depending on whether the decorated function returns an array, grid or vector.
-        The over written functions are in the child classes `ArrayMaker`, `GridMaker` and `VectorYXMaker`.
+        If the input is a raw ndarray (e.g. numpy or JAX), the function result is returned unchanged.
         """
 
         if isinstance(self.grid, Grid2D):
             return self.via_grid_2d(self.evaluate_func)
         elif isinstance(self.grid, Grid2DIrregular):
             return self.via_grid_2d_irr(self.evaluate_func)
-        elif isinstance(self.grid, Grid1D):
-            return self.via_grid_1d(self.evaluate_func)
 
         return self.evaluate_func

autoarray/structures/decorators/to_array.py

@@ -1,125 +1,49 @@
-import numpy as np
-from functools import wraps
-
-
-from typing import List, Union
-
-from autoarray.structures.arrays.irregular import ArrayIrregular
-from autoarray.structures.arrays.uniform_1d import Array1D
-from autoarray.structures.arrays.uniform_2d import Array2D
-from autoarray.structures.decorators.abstract import AbstractMaker
-from autoarray.structures.grids.uniform_1d import Grid1D
-from autoarray.structures.grids.irregular_2d import Grid2DIrregular
-from autoarray.structures.grids.uniform_2d import Grid2D
-
-
-class ArrayMaker(AbstractMaker):
-    def via_grid_2d(self, result) -> Union[Array2D, List[Array2D]]:
-        """
-        Convert the result of a decorated function which receives as input a `Grid2D` object to an `Array2D` object.
-
-        If the result returns a list, a list of `Array2D` objects is returned.
-
-        Parameters
-        ----------
-        result
-            The input result (e.g. of a decorated function) that is converted to an Array2D or list of Array2D objects.
-        """
-
-        if not isinstance(result, list):
-            return Array2D(values=result, mask=self.mask)
-        return [Array2D(values=res, mask=self.mask) for res in result]
-
-    def via_grid_2d_irr(self, result) -> Union[ArrayIrregular, List[ArrayIrregular]]:
-        """
-        Convert the result of a decorated function which receives as input a `Grid2DIrregular` object to an `ArrayIrregular`
-        object.
-
-        If the result returns a list, a list of `ArrayIrregular` objects is returned.
-
-        Parameters
-        ----------
-        result
-            The input result (e.g. of a decorated function) that is converted to an ArrayIrregular or list of
-            ArrayIrregular objects.
-        """
-        if not isinstance(result, list):
-            return ArrayIrregular(values=result)
-        return [ArrayIrregular(values=res) for res in result]
-
-    def via_grid_1d(self, result) -> Union[Array1D, List[Array1D]]:
-        """
-        Convert the result of a decorated function which receives as input a `Grid1D` object to an `Array1D` object.
-
-        If the result returns a list, a list of `Array1D` objects is returned.
-
-        Parameters
-        ----------
-        result
-            The input result (e.g. of a decorated function) that is converted to an Array1D or list of Array1D objects.
-        """
-        if not isinstance(result, list):
-            return Array1D(values=result, mask=self.mask)
-        return [Array1D(values=res, mask=self.mask) for res in result]
-
-
-def to_array(func):
-    """
-    Homogenize the inputs and outputs of functions that take 1D or 2D grids of coordinates and return a 1D ndarray
-    which is converted to an `Array2D`, `ArrayIrregular` or `Array1D` object.
-
-    Parameters
-    ----------
-    func
-        A function which computes a set of values from a 1D or 2D grid of coordinates.
-
-    Returns
-    -------
-        A function that has its outputs homogenized to `Array2D`, `ArrayIrregular` or `Array1D` objects.
-    """
-
-    @wraps(func)
-    def wrapper(
-        obj: object,
-        grid: Union[np.ndarray, Grid2D, Grid2DIrregular, Grid1D],
-        xp=np,
-        *args,
-        **kwargs,
-    ) -> Union[np.ndarray, Array1D, Array2D, ArrayIrregular, List]:
-        """
-        This decorator homogenizes the input of a "grid_like" 2D structure (`Grid2D`, `Grid2DIrregular` or `Grid1D`)
-        into a function which outputs an array-like structure (`Array2D`, `ArrayIrregular` or `Array1D`).
-
-        It allows these classes to be interchangeably input into a function, such that the grid is used to evaluate
-        the function at every (y,x) coordinates of the grid using specific functionality of the input grid.
-
-        The grid_like objects `Grid2D` and `Grid2DIrregular` are input into the function as a slimmed 2D ndarray array
-        of shape [total_coordinates, 2] where the second dimension stores the (y,x)
-
-        There are three types of consistent data structures and therefore decorated function mappings:
-
-        - Uniform (`Grid2D` -> `Array`): 2D structures defined on a uniform grid of data points. Both structures are
-        defined according to a `Mask2D`, which the maker object ensures is passed through self consistently.
-
-        - Irregular (`Grid2DIrregular` -> `ArrayIrregular`: 2D structures defined on an irregular grid of data points,
-        Neither structure is defined according to a mask and the maker sures the lack of a mask does not prevent the
-        function from being evaluated.
-
-        - 1D (`Grid1D` -> `Array1D`): 1D structures defined on a 1D grid of data points. These project the 1D grid
-        to a 2D grid to ensure the function can be evaluated, and then deproject the 2D grid back to a 1D grid to
-        ensure the output data structure is consistent with the input grid.
-
-        Parameters
-        ----------
-        obj
-            An object whose function uses grid_like inputs to compute quantities at every coordinate on the grid.
-        grid
-            A grid_like object of coordinates on which the function values are evaluated.
-
-        Returns
-        -------
-            The function values evaluated on the grid with the same structure as the input grid_like object.
-        """
-        return ArrayMaker(func=func, obj=obj, grid=grid, xp=xp, *args, **kwargs).result
-
-    return wrapper
+import numpy as np
+from functools import wraps
+from typing import List, Union
+
+from autoarray.structures.arrays.irregular import ArrayIrregular
+from autoarray.structures.arrays.uniform_2d import Array2D
+from autoarray.structures.decorators.abstract import AbstractMaker
+from autoarray.structures.grids.irregular_2d import Grid2DIrregular
+from autoarray.structures.grids.uniform_2d import Grid2D
+
+
+class ArrayMaker(AbstractMaker):
+    def via_grid_2d(self, result) -> Union[Array2D, List[Array2D]]:
+        if not isinstance(result, list):
+            return Array2D(values=result, mask=self.mask)
+        return [Array2D(values=res, mask=self.mask) for res in result]
+
+    def via_grid_2d_irr(self, result) -> Union[ArrayIrregular, List[ArrayIrregular]]:
+        if not isinstance(result, list):
+            return ArrayIrregular(values=result)
+        return [ArrayIrregular(values=res) for res in result]
+
+
+def to_array(func):
+    """
+    Homogenize the inputs and outputs of functions that take 2D grids of coordinates and return a 1D ndarray
+    which is converted to an `Array2D` or `ArrayIrregular` object.
+
+    Parameters
+    ----------
+    func
+        A function which computes a set of values from a 2D grid of coordinates.
+
+    Returns
+    -------
+        A function that has its outputs homogenized to `Array2D` or `ArrayIrregular` objects.
+    """
+
+    @wraps(func)
+    def wrapper(
+        obj: object,
+        grid: Union[np.ndarray, Grid2D, Grid2DIrregular],
+        xp=np,
+        *args,
+        **kwargs,
+    ) -> Union[np.ndarray, Array2D, ArrayIrregular, List]:
+        return ArrayMaker(func=func, obj=obj, grid=grid, xp=xp, *args, **kwargs).result
+
+    return wrapper