DataContainer agnostic on backend, second attempt

Wrappers/Python/cil/framework/acquisition_data.py

@@ -16,12 +16,14 @@
 # Authors:
 # CIL Developers, listed at: https://github.com/TomographicImaging/CIL/blob/master/NOTICE.txt
 # Joshua DM Hellier (University of Manchester) [refactorer]
-import numpy
+import numpy as np
 import warnings
 
 from .labels import AcquisitionDimension, Backend, AcquisitionType
 from .data_container import DataContainer
 from .partitioner import Partitioner
+import array_api_compat
+from array_api_compat import array_namespace # https://data-apis.org/array-api-compat/
 
 
 class AcquisitionData(DataContainer, Partitioner):
@@ -83,17 +85,15 @@ def __init__(self,
         if array is None:
             if dtype is None:
                 dtype = geometry.dtype
-            array = numpy.empty(geometry.shape, dtype)
-
+            xp = np
+            array = xp.empty(geometry.shape, dtype=dtype)
+
         elif issubclass(type(array) , DataContainer):
             array = array.as_array()
 
-        elif issubclass(type(array) , numpy.ndarray):
-            # remove singleton dimensions
-            array = numpy.squeeze(array)
-
         else:
-            raise TypeError('array must be a CIL type DataContainer or numpy.ndarray got {}'.format(type(array)))
+            # remove singleton dimensions
+            array = array.squeeze()
 
         if array.shape != geometry.shape:
             raise ValueError('Shape mismatch got {} expected {}'.format(array.shape, geometry.shape))
@@ -118,13 +118,13 @@ def __eq__(self, other):
         bool
             True if the two objects are equal, False otherwise.
         '''
-
+        # FIXME: address this
         if isinstance(other, AcquisitionData):
-            if numpy.array_equal(self.as_array(), other.as_array()) \
+            if np.array_equal(self.as_array(), other.as_array()) \
                 and self.geometry == other.geometry \
                 and self.dtype == other.dtype:
-                return True
-        elif numpy.array_equal(self.as_array(), other) and self.dtype==other.dtype:
+                return True 
+        elif np.array_equal(self.as_array(), other) and self.dtype==other.dtype:
             return True
         else:
             return False
@@ -149,7 +149,7 @@ def _get_slice(self, **kwargs):
             dim = self.geometry.dimension_labels.index('vertical')
 
             centre_slice_pos = (self.geometry.shape[dim]-1) / 2.
-            ind0 = int(numpy.floor(centre_slice_pos))
+            ind0 = int(np.floor(centre_slice_pos))
             w2 = centre_slice_pos - ind0
             kwargs['vertical'] = ind0
             out = DataContainer.get_slice(self, **kwargs)

Wrappers/Python/cil/framework/array_api_compat.py

@@ -0,0 +1,127 @@
+# Updates to array API compatibility layer for CIL
+from array_api_compat import array_namespace
+import sys
+
+def expand_dims(array, axis):
+    '''Expand dimensions of an array along specified axes.
+
+    Parameters
+    ----------
+    array : array-like
+        The input array to expand.
+    axis : int or tuple of int
+        The axis or axes along which to expand the dimensions.
+
+    Returns
+    -------
+    array-like
+        The array with expanded dimensions. It may be a new array or the same array with expanded dimensions.
+
+    Raises:
+    --------
+    IndexError If provided an invalid axis position, an IndexError should be raised.
+
+    Notes:
+    This function recursively expands the dimensions of the input array along the specified axes if a list or tuple of ints is provided.
+    '''
+    xp = array_namespace(array)
+
+    if isinstance(axis, int):
+        return xp.expand_dims(array, axis=axis)
+    if len(axis) == 1:
+        return xp.expand_dims(array, axis=axis[0])
+    axis = list(axis)
+    ax = axis.pop(0)
+    if len(axis) == 1:
+        axis = axis[0]
+    return expand_dims(xp.expand_dims(array, axis=ax), axis=axis)
+
+def squeeze(array, axis=None):
+    '''squeezes the array, removing all singleton dimensions recursively
+
+    Parameters
+    ----------
+    array : array-like
+        The array to squeeze
+    axis : int or tuple of int, optional
+        The axis or axes to squeeze. If None, all singleton dimensions are removed.
+
+    Returns
+    -------
+    array-like
+        The squeezed array with all singleton dimensions removed. If the input array has no singleton dimensions, it is returned unchanged.
+    '''
+    xp = array_namespace(array)
+    # find and remove singleton dimensions
+    if axis is None:
+        s = xp.nonzero(xp.asarray(array.shape) == 1)[0]
+        axis = s.tolist()
+        if len(axis) == 1:
+            axis = axis[0]
+        elif len(axis) == 0:
+            # nothing to do
+            return array
+    if isinstance(axis, int):
+        return xp.squeeze(array, axis=axis)
+    if len(axis) == 1:
+        return xp.squeeze(array, axis=axis[0])
+
+    # process from the largest axis to the smallest
+    axis = list(axis)
+    axis.sort(reverse=True)
+    ax = axis.pop(0)
+    if len(axis) == 1:
+        axis = axis[0]
+    return squeeze(xp.squeeze(array, axis=ax), axis=axis)
+
+def allclose(a, b, rtol=1e-5, atol=1e-6):
+    """
+    Check if two arrays are element-wise equal within a tolerance.allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False)[source]
+    parameters:
+    a, b: DataContainer or array_like
+Returns True if two arrays are element-wise equal within a tolerance.
+
+The tolerance values are positive, typically very small numbers. 
+The relative difference (rtol * abs(b)) and the absolute difference atol are added together to compare against the absolute difference between a and b.
+    """
+    try:
+        aarr = a.as_array()
+    except AttributeError:
+        aarr = a
+    try:
+        barr = b.as_array()
+    except AttributeError:
+        barr = b
+
+    xp = array_namespace(aarr)
+    if array_namespace(barr) != xp:
+        raise TypeError('Can only compare arrays ' \
+        'with same namespace. Got {} and {}'.format(array_namespace(aarr), array_namespace(barr)))
+
+    diff = rtol * xp.abs(barr) + atol
+    if xp.any(diff < xp.abs(aarr - barr)):
+        print(f"Max difference: {diff.max()}")
+        return False
+    return True
+
+def dtype_namespace(dtype):
+    """
+    Get the namespace of a given dtype.
+
+    Parameters
+    ----------
+    dtype : data-type
+        The data type to check.
+
+    Returns
+    -------
+    str
+        The namespace of the dtype.
+
+    Raises
+    ------
+    TypeError
+        If the dtype is not recognized.
+    """
+    xp = sys.modules[dtype.__module__]
+    return xp