feature/jax merge

.gitignore

@@ -1,3 +1,4 @@
+files/
 test_autoarray/dataset/files/array/output_test/uv_wavelengths.fits
 test_autoarray/dataset/files/array/output_test/visibilities.fits
 test_autoarray/dataset/plot/files/

autoarray/abstract_ndarray.py

@@ -1,49 +1,272 @@
 from __future__ import annotations
+
+from copy import copy
+
 from abc import ABC
 from abc import abstractmethod
 import numpy as np
-from pathlib import Path
-from typing import Union, TYPE_CHECKING
 
-from autoconf import conf
+from autoarray.numpy_wrapper import numpy as npw, register_pytree_node, Array
+
+from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from autoarray.structures.abstract_structure import Structure
 
 from autoarray.structures.arrays import array_2d_util
+from autoconf import conf
 
 
-class AbstractNDArray(np.ndarray, ABC):
-    def __reduce__(self):
-        pickled_state = super().__reduce__()
+def to_new_array(func):
+    """
+    Decorator for functions that returns an array. The array is wrapped in a new instance of the class.
 
-        class_dict = {}
-        for key, value in self.__dict__.items():
-            class_dict[key] = value
-        new_state = pickled_state[2] + (class_dict,)
+    Parameters
+    ----------
+    func
+        The function to be decorated.
 
-        return pickled_state[0], pickled_state[1], new_state
+    Returns
+    -------
+    The decorated function.
+    """
 
-    # noinspection PyMethodOverriding
-    def __setstate__(self, state):
-        for key, value in state[-1].items():
-            setattr(self, key, value)
-        super().__setstate__(state[0:-1])
+    def wrapper(self, *args, **kwargs) -> "AbstractNDArray":
+        return self.with_new_array(func(self, *args, **kwargs))
 
-    @property
-    @abstractmethod
-    def native(self) -> Structure:
+    return wrapper
+
+
+def unwrap_array(func):
+    """
+    Decorator for functions that take an array as an argument. If the argument is an AbstractNDArray, the underlying
+    array is used instead.
+
+    Parameters
+    ----------
+    func
+        The function to be decorated.
+
+    Returns
+    -------
+    The decorated function.
+    """
+
+    def wrapper(self, other):
+        try:
+            return func(self, other.array)
+        except AttributeError:
+            return func(self, other)
+
+    return wrapper
+
+
+class AbstractNDArray(ABC):
+    def __init__(self, array):
+        while isinstance(array, AbstractNDArray):
+            array = array.array
+        self._array = array
+        try:
+            register_pytree_node(
+                type(self),
+                self.instance_flatten,
+                self.instance_unflatten,
+            )
+        except ValueError:
+            pass
+
+    __no_flatten__ = ()
+
+    def invert(self):
+        new = self.copy()
+        new._array = np.invert(new._array)
+        return new
+
+    @classmethod
+    def instance_flatten(cls, instance):
         """
-        Returns the data structure in its `native` format which contains all unmaksed values to the native dimensions.
+        Flatten an instance of an autoarray class into a tuple of its attributes (i.e.. a pytree)
         """
+        keys, values = zip(
+            *sorted(
+                {
+                    key: value
+                    for key, value in instance.__dict__.items()
+                    if key not in cls.__no_flatten__
+                }.items()
+            )
+        )
+        return values, keys
 
     @staticmethod
     def flip_hdu_for_ds9(values):
         if conf.instance["general"]["fits"]["flip_for_ds9"]:
             return np.flipud(values)
         return values
 
-    def output_to_fits(self, file_path: Union[Path, str], overwrite: bool = False):
+    @classmethod
+    def instance_unflatten(cls, aux_data, children):
+        """
+        Unflatten a tuple of attributes (i.e. a pytree) into an instance of an autoarray class
+        """
+        instance = cls.__new__(cls)
+        for key, value in zip(aux_data, children[1:]):
+            setattr(instance, key, value)
+        return instance
+
+    def with_new_array(self, array: np.ndarray) -> "AbstractNDArray":
+        """
+        Copy this object but give it a new array.
+
+        This is used to ensure that when an array is modified, associated
+        attributes such as pixel size are retained.
+
+        Parameters
+        ----------
+        array
+            The new array that is given to the copied object.
+
+        Returns
+        -------
+
+        """
+        new_array = self.copy()
+        new_array._array = array
+        return new_array
+
+    def copy(self):
+        new = copy(self)
+        return new
+
+    def __copy__(self):
+        """
+        When copying an autoarray also copy its underlying array.
+        """
+        new = self.__new__(self.__class__)
+        new.__dict__.update(self.__dict__)
+        new._array = self._array.copy()
+        return new
+
+    def __deepcopy__(self, memo):
+        """
+        When copying an autoarray also copy its underlying array.
+        """
+        new = self.__new__(self.__class__)
+        new.__dict__.update(self.__dict__)
+        new._array = self._array.copy()
+        return new
+
+    def __iter__(self):
+        return iter(self._array)
+
+    @to_new_array
+    def sqrt(self):
+        return np.sqrt(self._array)
+
+    @property
+    def array(self):
+        return self._array
+
+    @unwrap_array
+    def __lt__(self, other):
+        return self._array < other
+
+    @unwrap_array
+    def __le__(self, other):
+        return self._array <= other
+
+    @unwrap_array
+    def __gt__(self, other):
+        return self._array > other
+
+    @unwrap_array
+    def __ge__(self, other):
+        return self._array >= other
+
+    @unwrap_array
+    def __eq__(self, other):
+        return self._array == other
+
+    @to_new_array
+    @unwrap_array
+    def __pow__(self, other):
+        return self._array**other
+
+    @to_new_array
+    @unwrap_array
+    def __add__(self, other):
+        return self._array + other
+
+    @to_new_array
+    @unwrap_array
+    def __radd__(self, other):
+        return other + self._array
+
+    @to_new_array
+    @unwrap_array
+    def __sub__(self, other):
+        return self._array - other
+
+    @to_new_array
+    @unwrap_array
+    def __rsub__(self, other):
+        return other - self._array
+
+    @unwrap_array
+    def __ne__(self, other):
+        return self._array != other
+
+    @to_new_array
+    @unwrap_array
+    def __mul__(self, other):
+        return self._array * other
+
+    @to_new_array
+    @unwrap_array
+    def __rmul__(self, other):
+        return other * self._array
+
+    @to_new_array
+    def __neg__(self):
+        return -self._array
+
+    def __invert__(self):
+        return ~self._array
+
+    def __divmod__(self, other):
+        return divmod(self._array, other)
+
+    def __rdivmod__(self, other):
+        return divmod(other, self._array)
+
+    @to_new_array
+    @unwrap_array
+    def __truediv__(self, other):
+        return self._array / other
+
+    @to_new_array
+    @unwrap_array
+    def __rtruediv__(self, other):
+        return other / self._array
+
+    @to_new_array
+    def __abs__(self):
+        return abs(self._array)
+
+    def sum(self, *args, **kwargs):
+        return self._array.sum(*args, **kwargs)
+
+    def __float__(self):
+        return float(self._array)
+
+    @property
+    @abstractmethod
+    def native(self) -> Structure:
+        """
+        Returns the data structure in its `native` format which contains all unmaksed values to the native dimensions.
+        """
+
+    def output_to_fits(self, file_path: str, overwrite: bool = False):
         """
         Output the grid to a .fits file.
 
@@ -55,5 +278,86 @@ def output_to_fits(self, file_path: Union[Path, str], overwrite: bool = False):
             If a file already exists at the path, if overwrite=True it is overwritten else an error is raised.
         """
         array_2d_util.numpy_array_2d_to_fits(
-            array_2d=self.native, file_path=file_path, overwrite=overwrite
+            array_2d=self.native.array, file_path=file_path, overwrite=overwrite
         )
+
+    @property
+    def shape(self):
+        try:
+            return self._array.shape
+        except AttributeError:
+            return ()
+
+    @property
+    def size(self):
+        return self._array.size
+
+    @property
+    def dtype(self):
+        return self._array.dtype
+
+    @property
+    def ndim(self):
+        return self._array.ndim
+
+    def max(self, *args, **kwargs):
+        return self._array.max(*args, **kwargs)
+
+    def min(self, *args, **kwargs):
+        return self._array.min(*args, **kwargs)
+
+    @to_new_array
+    def reshape(self, *args, **kwargs):
+        return self._array.reshape(*args, **kwargs)
+
+    def __getattr__(self, item):
+        if item != "__setstate__":
+            try:
+                return getattr(self._array, item)
+            except AttributeError:
+                pass
+        raise AttributeError(
+            f"{self.__class__.__name__} does not have attribute {item}"
+        )
+
+    def __getitem__(self, item):
+        result = self._array[item]
+        if isinstance(item, slice):
+            result = self.with_new_array(result)
+        if isinstance(result, np.ndarray):
+            result = self.with_new_array(result)
+        return result
+
+    def __setitem__(self, key, value):
+        if isinstance(key, (np.ndarray, AbstractNDArray, Array)):
+            self._array = npw.where(key, value, self._array)
+        else:
+            self._array[key] = value
+
+    def __repr__(self):
+        return f"{self.__class__.__name__} {self.shape}"
+
+    def __array__(self, dtype=None):
+        if dtype:
+            return self._array.astype(dtype)
+        return self._array
+
+    def __len__(self):
+        return len(self._array)
+
+    @to_new_array
+    def astype(self, dtype):
+        return self._array.astype(dtype)
+
+    @property
+    @to_new_array
+    def real(self):
+        return self._array.real
+
+    @property
+    @to_new_array
+    def imag(self):
+        return self._array.imag
+
+    def all(self):
+        return self._array.all()