Add support for more numpy functions to dense and sparse DelayedArrays (#62)

Update tests and documentation

Author: jkanche

src/delayedarray/DelayedArray.py

@@ -1,6 +1,6 @@
 from typing import Sequence, Tuple, Union, Optional, List, Callable
 import numpy
-from numpy import array, dtype, integer, issubdtype, ndarray, prod, array2string
+from numpy import dtype, ndarray, array2string
 from collections import namedtuple
 
 from .SparseNdarray import SparseNdarray
@@ -24,7 +24,7 @@
 
 from ._subset import _getitem_subset_preserves_dimensions, _getitem_subset_discards_dimensions, _repr_subset
 from ._isometric import translate_ufunc_to_op_simple, translate_ufunc_to_op_with_args
-from ._statistics import array_mean, array_var, array_sum, _create_offset_multipliers
+from ._statistics import array_mean, array_var, array_sum, _create_offset_multipliers, array_any, array_all
 
 __author__ = "ltla"
 __copyright__ = "ltla"
@@ -255,6 +255,21 @@ def __array_function__(self, func, types, args, kwargs) -> "DelayedArray":
                 decimals = 0
             return DelayedArray(Round(seed, decimals=decimals))
 
+        if func == numpy.mean:
+            return self.mean(**kwargs)
+
+        if func == numpy.sum:
+            return self.sum(**kwargs)
+
+        if func == numpy.var:
+            return self.var(**kwargs)
+
+        if func == numpy.any:
+            return self.any(**kwargs)
+
+        if func == numpy.all:
+            return self.all(**kwargs)
+
         if func == numpy.shape:
             return self.shape 
 
@@ -691,6 +706,66 @@ def __abs__(self) -> "DelayedArray":
         """
         return DelayedArray(UnaryIsometricOpSimple(self._seed, operation="abs"))
 
+    def __or__(self, other) -> "DelayedArray":
+        """Element-wise OR with something.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            A ``DelayedArray`` containing the delayed OR operation.
+        """
+        return _wrap_isometric_with_args(self, other, operation="logical_or", right=True)
+
+    def __ror__(self, other) -> "DelayedArray":
+        """Element-wise OR with the right-hand-side of a ``DelayedArray``.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            A ``DelayedArray`` containing the delayed OR operation.
+        """
+        return _wrap_isometric_with_args(self, other, operation="logical_or", right=False)
+
+    def __and__(self, other) -> "DelayedArray":
+        """Element-wise AND with something.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            A ``DelayedArray`` containing the delayed AND operation.
+        """
+        return _wrap_isometric_with_args(self, other, operation="logical_and", right=True)
+
+    def __rand__(self, other) -> "DelayedArray":
+        """Element-wise AND with the right-hand-side of a ``DelayedArray``.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            A ``DelayedArray`` containing the delayed AND operation.
+        """
+        return _wrap_isometric_with_args(self, other, operation="logical_and", right=False)
+
     # Subsetting.
     def __getitem__(self, subset: Tuple[Union[slice, Sequence], ...]) -> Union["DelayedArray", ndarray]:
         """Take a subset of this ``DelayedArray``. This follows the same logic as NumPy slicing and will generate a
@@ -832,6 +907,79 @@ def var(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
                 masked=is_masked(self),
             )
 
+    def any(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: int = 1e8) -> numpy.ndarray:
+        """Test whether any array element along a given axis evaluates to True.
+
+        Compute this test across the ``DelayedArray``, possibly over a
+        given axis or set of axes. If the seed has a ``any()`` method, that
+        method is called directly with the supplied arguments.
+
+        Args:
+            axis: 
+                A single integer specifying the axis over which to test
+                for any. Alternatively, a tuple (multiple axes) or None (no
+                axes), see :py:func:`~numpy.any` for details.
+
+            dtype:
+                NumPy type for the output array. If None, this is automatically
+                chosen based on the type of the ``DelayedArray``, see
+                :py:func:`~numpy.any` for details.
+
+            buffer_size:
+                Buffer size in bytes to use for block processing. Larger values
+                generally improve speed at the cost of memory.
+
+        Returns:
+            A NumPy array containing the boolean values. If ``axis = None``, this will
+            be a NumPy scalar instead.
+        """
+        if hasattr(self._seed, "any"):
+            return self._seed.any(axis=axis).astype(dtype)
+        else:
+            return array_any(
+                self, 
+                axis=axis, 
+                dtype=dtype, 
+                reduce_over_x=lambda x, axes, op : _reduce(x, axes, op, buffer_size),
+                masked=is_masked(self),
+            )
+
+    def all(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None, buffer_size: int = 1e8) -> numpy.ndarray:
+        """Test whether all array elements along a given axis evaluate to True.
+
+        Compute this test across the ``DelayedArray``, possibly over a
+        given axis or set of axes. If the seed has a ``all()`` method, that
+        method is called directly with the supplied arguments.
+
+        Args:
+            axis: 
+                A single integer specifying the axis over which to test 
+                for all. Alternatively, a tuple (multiple axes) or None (no
+                axes), see :py:func:`~numpy.all` for details.
+
+            dtype:
+                NumPy type for the output array. If None, this is automatically
+                chosen based on the type of the ``DelayedArray``, see
+                :py:func:`~numpy.all` for details.
+
+            buffer_size:
+                Buffer size in bytes to use for block processing. Larger values
+                generally improve speed at the cost of memory.
+
+        Returns:
+            A NumPy array containing the boolean values. If ``axis = None``, this will
+            be a NumPy scalar instead.
+        """
+        if hasattr(self._seed, "all"):
+            return self._seed.all(axis=axis).astype(dtype)
+        else:
+            return array_all(
+                self, 
+                axis=axis, 
+                dtype=dtype, 
+                reduce_over_x=lambda x, axes, op : _reduce(x, axes, op, buffer_size),
+                masked=is_masked(self),
+            )
 
 @extract_dense_array.register
 def extract_dense_array_DelayedArray(x: DelayedArray, subset: Tuple[Sequence[int], ...]) -> numpy.ndarray:

src/delayedarray/SparseNdarray.py

@@ -14,7 +14,7 @@
     _concatenate_unmasked_ndarrays,
     _concatenate_maybe_masked_ndarrays
 )
-from ._statistics import array_mean, array_var, array_sum, _create_offset_multipliers
+from ._statistics import array_mean, array_var, array_sum, _create_offset_multipliers, array_all, array_any
 
 __author__ = "ltla"
 __copyright__ = "ltla"
@@ -665,6 +665,71 @@ def __abs__(self):
             A ``SparseNdarray`` containing the delayed absolute value operation.
         """
         return _transform_sparse_array_from_SparseNdarray(self, lambda l, i, v : (i, abs(v)), self._dtype)
+    
+    def __or__(self, other) -> Union["SparseNdarray", numpy.ndarray]:
+        """Element-wise OR with something.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            Array containing the result of the check.
+            This may or may not be sparse depending on ``other``.
+        """
+        return _operate_with_args_on_SparseNdarray(self, other, operation="logical_or", right=True)
+
+    def __ror__(self, other) -> Union["SparseNdarray", numpy.ndarray]:
+        """Element-wise OR with the right-hand-side of a ``DelayedArray``.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            Array containing the result of the check.
+            This may or may not be sparse depending on ``other``.
+        """
+        return _operate_with_args_on_SparseNdarray(self, other, operation="logical_or", right=False)
+
+    def __and__(self, other) -> Union["SparseNdarray", numpy.ndarray]:
+        """Element-wise AND with something.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            Array containing the result of the check.
+            This may or may not be sparse depending on ``other``.
+        """
+        return _operate_with_args_on_SparseNdarray(self, other, operation="logical_and", right=True)
+
+    def __rand__(self, other) -> Union["SparseNdarray", numpy.ndarray]:
+        """Element-wise AND with the right-hand-side of a ``DelayedArray``.
+
+        Args:
+            other:
+                A numeric scalar;
+                or a NumPy array with dimensions as described in
+                :py:class:`~delayedarray.UnaryIsometricOpWithArgs.UnaryIsometricOpWithArgs`;
+                or a ``DelayedArray`` of the same dimensions as :py:attr:`~shape`.
+
+        Returns:
+            Array containing the result of the check.
+            This may or may not be sparse depending on ``other``.
+        """
+        return _operate_with_args_on_SparseNdarray(self, other, operation="logical_and", right=False)
+
 
     # Subsetting.
     def __getitem__(self, subset: Tuple[Union[slice, Sequence], ...]) -> Union["SparseNdarray", numpy.ndarray]:
@@ -760,6 +825,21 @@ def __array_function__(self, func, types, args, kwargs) -> "SparseNdarray":
         if func == numpy.round:
             return _transform_sparse_array_from_SparseNdarray(self, lambda l, i, v : (i, func(v, **kwargs)), self._dtype)
 
+        if func == numpy.mean:
+            return self.mean(**kwargs)
+
+        if func == numpy.sum:
+            return self.sum(**kwargs)
+
+        if func == numpy.var:
+            return self.var(**kwargs)
+
+        if func == numpy.any:
+            return self.any(**kwargs)
+
+        if func == numpy.all:
+            return self.all(**kwargs)
+
         raise NotImplementedError(f"'{func.__name__}' is not implemented!")
 
 
@@ -872,6 +952,64 @@ def var(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optiona
             masked=self._is_masked,
         )
 
+    def any(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None) -> numpy.ndarray:
+        """Test whether any array element along a given axis evaluates to True.
+
+        Compute this test across the ``SparseNdarray``, possibly over a
+        given axis or set of axes. If the seed has a ``any()`` method, that
+        method is called directly with the supplied arguments.
+
+        Args:
+            axis: 
+                A single integer specifying the axis over which to test
+                for any. Alternatively, a tuple (multiple axes) or None
+                (no axes), see :py:func:`~numpy.any` for details.
+
+            dtype:
+                NumPy type for the output array. If None, this is automatically
+                chosen based on the type of the ``SparseNdarray``, see
+                :py:func:`~numpy.any` for details.
+
+        Returns:
+            A NumPy array containing the variances. If ``axis = None``,
+            this will be a NumPy scalar instead.
+        """
+        return array_any(
+            self, 
+            axis=axis, 
+            dtype=dtype, 
+            reduce_over_x=_reduce_SparseNdarray,
+            masked=self._is_masked,
+        )
+
+    def all(self, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[numpy.dtype] = None) -> numpy.ndarray:
+        """Test whether all array elements along a given axis evaluate to True.
+
+        Compute this test across the ``SparseNdarray``, possibly over a
+        given axis or set of axes. If the seed has a ``all()`` method, that
+        method is called directly with the supplied arguments.
+        Args:
+            axis: 
+                A single integer specifying the axis over which to test
+                for any. Alternatively, a tuple (multiple axes) or None
+                (no axes), see :py:func:`~numpy.any` for details.
+
+            dtype:
+                NumPy type for the output array. If None, this is automatically
+                chosen based on the type of the ``SparseNdarray``, see
+                :py:func:`~numpy.any` for details.
+
+        Returns:
+            A NumPy array containing the variances. If ``axis = None``,
+            this will be a NumPy scalar instead.
+        """
+        return array_all(
+            self, 
+            axis=axis, 
+            dtype=dtype, 
+            reduce_over_x=_reduce_SparseNdarray,
+            masked=self._is_masked,
+        )
 
     # Other stuff
     def __copy__(self) -> "SparseNdarray":

src/delayedarray/_isometric.py

@@ -106,6 +106,7 @@ def _execute(left, right, operation):
         "floor",
         "trunc",
         "sign",
+        "isnan"
     ]
 )