WIP is-constant

Author: flying-sheep

src/fast_array_utils/stats/__init__.py

@@ -3,7 +3,8 @@
 
 from __future__ import annotations
 
+from ._is_constant import is_constant
 from ._sum import sum
 
 
-__all__ = ["sum"]
+__all__ = ["is_constant", "sum"]

src/fast_array_utils/stats/_is_constant.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from functools import partial, singledispatch
+from numbers import Integral
+from typing import TYPE_CHECKING, overload
+
+import numba
+import numpy as np
+
+from ..types import CSBase, DaskArray, H5Dataset, ZarrArray
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+    from typing import Any, Literal, TypeAlias, TypeVar
+
+    from numpy.typing import NDArray
+
+    _Array: TypeAlias = NDArray[Any] | CSBase | H5Dataset | ZarrArray | DaskArray
+
+    C = TypeVar("C", bound=Callable[..., Any])
+
+
+@overload
+def is_constant(a: _Array, axis: None = None) -> bool: ...
+@overload
+def is_constant(a: _Array, axis: Literal[0, 1]) -> NDArray[np.bool_]: ...
+
+
+def is_constant(a: _Array, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool_]:
+    """Check whether values in array are constant.
+
+    Params
+    ------
+    a
+        Array to check
+    axis
+        Axis to reduce over.
+
+
+    Returns:
+    -------
+    Boolean array, True values were constant.
+
+    Example:
+    -------
+    >>> a = np.array([[0, 1], [0, 0]])
+    >>> a
+    array([[0, 1],
+           [0, 0]])
+    >>> is_constant(a)
+    False
+    >>> is_constant(a, axis=0)
+    array([ True, False])
+    >>> is_constant(a, axis=1)
+    array([False,  True])
+
+    """
+    if axis is not None:
+        if not isinstance(axis, Integral):
+            msg = "axis must be integer or None."
+            raise TypeError(msg)
+        if axis not in (0, 1):
+            msg = "We only support axis 0 and 1 at the moment"
+            raise NotImplementedError(msg)
+
+    return _is_constant(a, axis)
+
+
+@singledispatch
+def _is_constant(a: _Array, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool_]:
+    raise NotImplementedError
+
+
+@_is_constant.register(np.ndarray)
+@_is_constant.register(H5Dataset)
+@_is_constant.register(ZarrArray)
+def _(a: NDArray[Any], axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool_]:
+    # Should eventually support nd, not now.
+    match axis:
+        case None:
+            return bool((a == a.flat[0]).all())
+        case 0:
+            return _is_constant_rows(a.T)
+        case 1:
+            return _is_constant_rows(a)
+
+
+def _is_constant_rows(a: NDArray[Any]) -> NDArray[np.bool_]:
+    b = np.broadcast_to(a[:, 0][:, np.newaxis], a.shape)
+    return (a == b).all(axis=1)  # type: ignore[no-any-return]
+
+
+@_is_constant.register(CSBase)  # type: ignore[call-overload,misc]
+def _(a: CSBase, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool_]:
+    n_row, n_col = a.shape
+    if axis is None:
+        if len(a.data) == n_row * n_col:
+            return is_constant(a.data)
+        return (a.data == 0).all()  # type: ignore[no-any-return]
+    shape = (n_row, n_col) if axis == 1 else (n_col, n_row)
+    match axis, a.format:
+        case 0, "csr":
+            a = a.T.tocsr()
+        case 1, "csc":
+            a = a.T.tocsc()
+    return _is_constant_csr_rows(a.data, a.indptr, shape)
+
+
+@numba.njit(cache=True)
+def _is_constant_csr_rows(
+    data: NDArray[np.number[Any]],
+    indptr: NDArray[np.integer[Any]],
+    shape: tuple[int, int],
+) -> NDArray[np.bool_]:
+    n = len(indptr) - 1
+    result = np.ones(n, dtype=np.bool_)
+    for i in numba.prange(n):
+        start = indptr[i]
+        stop = indptr[i + 1]
+        val = data[start] if stop - start == shape[1] else 0
+        for j in range(start, stop):
+            if data[j] != val:
+                result[i] = False
+                break
+    return result
+
+
+@_is_constant.register(DaskArray)
+def _(a: DaskArray, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool_]:
+    if TYPE_CHECKING:
+        from dask.array.core import map_blocks
+    else:
+        from dask.array import map_blocks
+
+    if axis is None:
+        v = a[tuple(0 for _ in range(a.ndim))].compute()
+        return (a == v).all()  # type: ignore[no-any-return]
+    # TODO(flying-sheep): use overlapping blocks and reduction instead of `drop_axis`  # noqa: TD003
+    return map_blocks(  # type: ignore[no-any-return,no-untyped-call]
+        partial(is_constant, axis=axis),
+        a,
+        drop_axis=axis,
+        meta=np.array([], dtype=a.dtype),
+    )

tests/test_stats.py

@@ -71,6 +71,32 @@ def test_sum(
     np.testing.assert_array_equal(sum_, np.sum(np_arr, axis=axis, dtype=dtype_arg))  # type: ignore[arg-type]
 
 
+@pytest.mark.parametrize(
+    ("axis", "expected"),
+    [
+        pytest.param(None, False, id="None"),
+        pytest.param(0, [True, True, False, False], id="0"),
+        pytest.param(1, [False, False, True, True, False, True], id="1"),
+    ],
+)
+def test_is_constant_dask(
+    to_array: ToArray, axis: Literal[0, 1, None], expected: bool | list[bool]
+) -> None:
+    x_data = [
+        [0, 0, 1, 1],
+        [0, 0, 1, 1],
+        [0, 0, 0, 0],
+        [0, 0, 0, 0],
+        [0, 0, 1, 0],
+        [0, 0, 0, 0],
+    ]
+    x = to_array(x_data)
+    result = stats.is_constant(x, axis=axis)
+    if isinstance(result, types.DaskArray):
+        result = result.compute()
+    np.testing.assert_array_equal(expected, result)
+
+
 @pytest.mark.benchmark
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])  # random only supports float
 def test_sum_benchmark(