Add mean_var (#46)

Co-authored-by: Ilan Gold <[email protected]>

Author: flying-sheep

src/fast_array_utils/stats/__init__.py

@@ -5,7 +5,8 @@
 
 from ._is_constant import is_constant
 from ._mean import mean
+from ._mean_var import mean_var
 from ._sum import sum
 
 
-__all__ = ["is_constant", "mean", "sum"]
+__all__ = ["is_constant", "mean", "mean_var", "sum"]

src/fast_array_utils/stats/_mean.py

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: MPL-2.0
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, overload
 
 import numpy as np
 
@@ -13,19 +13,34 @@
 
     from numpy._typing._array_like import _ArrayLikeFloat_co as ArrayLike
     from numpy.typing import DTypeLike, NDArray
+    from optype.numpy import ToDType
 
     from .. import types
 
-    # all supported types except OutOfCoreDataset (TODO)
-    Array = (
+    # all supported types except Dask and OutOfCoreDataset (TODO)
+    NonDaskArray = (
         NDArray[Any]
         | types.CSBase
         | types.H5Dataset
         | types.ZarrArray
         | types.CupyArray
         | types.CupySparseMatrix
-        | types.DaskArray
     )
+    Array = NonDaskArray | types.DaskArray
+
+
+@overload
+def mean(
+    x: ArrayLike | NonDaskArray, /, *, axis: Literal[None] = None, dtype: DTypeLike | None = None
+) -> np.number[Any]: ...
+@overload
+def mean(
+    x: ArrayLike | NonDaskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None
+) -> NDArray[np.number[Any]]: ...
+@overload
+def mean(
+    x: types.DaskArray, /, *, axis: Literal[0, 1], dtype: ToDType[Any] | None = None
+) -> types.DaskArray: ...
 
 
 def mean(
@@ -34,7 +49,7 @@ def mean(
     *,
     axis: Literal[0, 1, None] = None,
     dtype: DTypeLike | None = None,
-) -> NDArray[Any] | types.DaskArray:
+) -> NDArray[np.number[Any]] | np.number[Any] | types.DaskArray:
     """Mean over both or one axis.
 
     Returns
@@ -49,7 +64,7 @@ def mean(
     if not hasattr(x, "shape"):
         raise NotImplementedError  # TODO(flying-sheep): infer shape  # noqa: TD003
     if TYPE_CHECKING:
-        assert isinstance(x, Array)
+        assert isinstance(x, Array)  # type:ignore[unused-ignore]
     total = sum_(x, axis=axis, dtype=dtype)
     n = np.prod(x.shape) if axis is None else x.shape[axis]
     return total / n

src/fast_array_utils/stats/_mean_var.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, no_type_check, overload
+
+import numba
+import numpy as np
+
+from .. import types
+from ._mean import mean
+from ._power import power
+
+
+if TYPE_CHECKING:
+    from typing import Any, Literal
+
+    from numpy.typing import NDArray
+
+    MemArray = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix
+
+
+__all__ = ["mean_var"]
+
+
+@overload
+def mean_var(
+    x: MemArray, /, *, axis: Literal[None] = None, correction: int = 0
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ...
+@overload
+def mean_var(
+    x: MemArray, /, *, axis: Literal[0, 1], correction: int = 0
+) -> tuple[np.float64, np.float64]: ...
+@overload
+def mean_var(
+    x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None, correction: int = 0
+) -> tuple[types.DaskArray, types.DaskArray]: ...
+
+
+@no_type_check  # mypy is extremely confused
+def mean_var(
+    x: MemArray | types.DaskArray,
+    /,
+    *,
+    axis: Literal[0, 1, None] = None,
+    correction: int = 0,
+) -> (
+    tuple[NDArray[np.float64], NDArray[np.float64]]
+    | tuple[np.float64, np.float64]
+    | tuple[types.DaskArray, types.DaskArray]
+):
+    if axis is not None and isinstance(x, types.CSBase):
+        mean_, var = _sparse_mean_var(x, axis=axis)
+    else:
+        mean_ = mean(x, axis=axis, dtype=np.float64)
+        mean_sq = mean(power(x, 2), axis=axis, dtype=np.float64)
+        var = mean_sq - mean_**2
+    if correction:  # R convention == 1 (unbiased estimator)
+        n = np.prod(x.shape) if axis is None else x.shape[axis]
+        if n != 1:
+            var *= n / (n - correction)
+    return mean_, var
+
+
+def _sparse_mean_var(
+    mtx: types.CSBase, /, *, axis: Literal[0, 1]
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Calculate means and variances for each row or column of a sparse matrix.
+
+    This code and internal functions are based on sklearns `sparsefuncs.mean_variance_axis`.
+
+    Modifications:
+    - allow deciding on the output type,
+      which can increase accuracy when calculating the mean and variance of 32bit floats.
+    - Doesn't currently implement support for null values, but could.
+    - Uses numba instead of Cython
+    """
+    assert axis in (0, 1)
+    if mtx.format == "csr":
+        ax_minor = 1
+        shape = mtx.shape
+    elif mtx.format == "csc":
+        ax_minor = 0
+        shape = mtx.shape[::-1]
+    else:
+        msg = "This function only works on sparse csr and csc matrices"
+        raise TypeError(msg)
+    if len(shape) == 1:
+        msg = "array must have 2 dimensions"
+        raise TypeError(msg)
+    f = sparse_mean_var_major_axis if axis == ax_minor else sparse_mean_var_minor_axis
+    return f(
+        mtx.data,
+        mtx.indptr,
+        mtx.indices,
+        major_len=shape[0],
+        minor_len=shape[1],
+        n_threads=numba.get_num_threads(),
+    )
+
+
+@numba.njit
+def sparse_mean_var_minor_axis(
+    data: NDArray[np.number[Any]],
+    indptr: NDArray[np.integer[Any]],
+    indices: NDArray[np.integer[Any]],
+    *,
+    major_len: int,
+    minor_len: int,
+    n_threads: int,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Compute mean and variance along the minor axis of a compressed sparse matrix."""
+    rows = len(indptr) - 1
+    sums = np.zeros((n_threads, minor_len))
+    squared_sums = np.zeros((n_threads, minor_len))
+    means = np.zeros(minor_len)
+    variances = np.zeros(minor_len)
+    for i in numba.prange(n_threads):
+        for r in range(i, rows, n_threads):
+            for j in range(indptr[r], indptr[r + 1]):
+                minor_index = indices[j]
+                if minor_index >= minor_len:
+                    continue
+                value = data[j]
+                sums[i, minor_index] += value
+                squared_sums[i, minor_index] += value * value
+    for c in numba.prange(minor_len):
+        sum = sums[:, c].sum()
+        means[c] = sum / major_len
+        variances[c] = squared_sums[:, c].sum() / major_len - (sum / major_len) ** 2
+    return means, variances
+
+
+@numba.njit
+def sparse_mean_var_major_axis(
+    data: NDArray[np.number[Any]],
+    indptr: NDArray[np.integer[Any]],
+    indices: NDArray[np.integer[Any]],  # noqa: ARG001
+    *,
+    major_len: int,
+    minor_len: int,
+    n_threads: int,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Compute means and variances along the major axis of a compressed sparse matrix."""
+    rows = len(indptr) - 1
+    means = np.zeros(major_len)
+    variances = np.zeros_like(means)
+
+    for i in numba.prange(n_threads):
+        for r in range(i, rows, n_threads):
+            sum_major = np.float64(0.0)
+            squared_sum_minor = np.float64(0.0)
+            for j in range(indptr[r], indptr[r + 1]):
+                value = np.float64(data[j])
+                sum_major += value
+                squared_sum_minor += value * value
+            means[r] = sum_major
+            variances[r] = squared_sum_minor
+    for c in numba.prange(major_len):
+        mean = means[c] / minor_len
+        means[c] = mean
+        variances[c] = variances[c] / minor_len - mean * mean
+    return means, variances

src/fast_array_utils/stats/_power.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: MPL-2.0
+from __future__ import annotations
+
+from functools import singledispatch
+from typing import TYPE_CHECKING, cast
+
+from .. import types
+
+
+if TYPE_CHECKING:
+    from typing import Any, TypeVar
+
+    from numpy.typing import NDArray
+
+    # All supported array types except for disk ones and OutOfCoreDataset
+    Array = NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+
+    _Arr = TypeVar("_Arr", bound=Array)
+
+
+def power(x: _Arr, n: int, /) -> _Arr:
+    """Take array or matrix to a power."""
+    # This wrapper is necessary because TypeVars can’t be used in `singledispatch` functions
+    return _power(x, n)  # type: ignore[return-value]
+
+
+@singledispatch
+def _power(x: Array, n: int, /) -> Array:
+    if TYPE_CHECKING:
+        assert not isinstance(x, types.DaskArray | types.CSMatrix)
+    return x**n  # type: ignore[operator]
+
+
+@_power.register(types.CSMatrix)  # type: ignore[call-overload,misc]
+def _power_cs(x: types.CSMatrix, n: int, /) -> types.CSMatrix:
+    return x.power(n)
+
+
+@_power.register(types.DaskArray)
+def _power_dask(x: types.DaskArray, n: int, /) -> types.DaskArray:
+    import dask.array as da
+
+    return cast(types.DaskArray, da.map_blocks(power, x, n))  # type: ignore[no-untyped-call]

stubs/cupy.pyi

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
 import numpy as np
 from numpy.typing import ArrayLike, DTypeLike, NDArray
@@ -8,6 +8,7 @@ class ndarray:
     dtype: np.dtype[Any]
     shape: tuple[int, ...]
     def get(self) -> NDArray[Any]: ...
+    def __power__(self, other: int) -> Self: ...
 
 def asarray(
     a: ArrayLike,

stubs/cupyx/scipy/sparse.pyi

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: MPL-2.0
-from typing import Any, Literal
+from typing import Any, Literal, Self
 
 import cupy
 import numpy as np
@@ -8,3 +8,4 @@ class spmatrix:
     dtype: np.dtype[Any]
     shape: tuple[int, int]
     def toarray(self, order: Literal["C", "F", None] = None, out: None = None) -> cupy.ndarray: ...
+    def __power__(self, other: int) -> Self: ...

stubs/numba/__init__.pyi

@@ -31,3 +31,4 @@ def prange(stop: SupportsIndex, /) -> Iterable[int]: ...
 def prange(
     start: SupportsIndex, stop: SupportsIndex, step: SupportsIndex = ..., /
 ) -> Iterable[int]: ...
+def get_num_threads() -> int: ...

tests/test_stats.py

@@ -90,15 +90,41 @@ def test_sum(
 
 @pytest.mark.parametrize(("axis", "expected"), [(None, 3.5), (0, [2.5, 3.5, 4.5]), (1, [2.0, 5.0])])
 def test_mean(
-    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: list[float]
+    array_type: ArrayType[Array], axis: Literal[0, 1, None], expected: float | list[float]
 ) -> None:
-    arr = array_type(np.array([[1, 2, 3], [4, 5, 6]]))
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), expected)
+
+    arr = array_type(np_arr)
     result = stats.mean(arr, axis=axis)
     if isinstance(result, types.DaskArray):
-        result = result.compute()  # type: ignore[no-untyped-call]
+        result = result.compute()
     np.testing.assert_array_equal(result, expected)
 
 
+@pytest.mark.array_type(skip=Flags.Disk)
+@pytest.mark.parametrize(
+    ("axis", "mean_expected", "var_expected"),
+    [(None, 3.5, 3.5), (0, [2.5, 3.5, 4.5], [4.5, 4.5, 4.5]), (1, [2.0, 5.0], [1.0, 1.0])],
+)
+def test_mean_var(
+    array_type: ArrayType[
+        NDArray[Any] | types.CSBase | types.CupyArray | types.CupySparseMatrix | types.DaskArray
+    ],
+    axis: Literal[0, 1, None],
+    mean_expected: float | list[float],
+    var_expected: float | list[float],
+) -> None:
+    np_arr = np.array([[1, 2, 3], [4, 5, 6]])
+    np.testing.assert_array_equal(np.mean(np_arr, axis=axis), mean_expected)
+    np.testing.assert_array_equal(np.var(np_arr, axis=axis, correction=1), var_expected)
+
+    arr = array_type(np_arr)
+    mean, var = stats.mean_var(arr, axis=axis, correction=1)
+    np.testing.assert_array_equal(mean, mean_expected)
+    np.testing.assert_array_almost_equal_nulp(var, var_expected, nulp=8)
+
+
 # TODO(flying-sheep): enable for GPU  # noqa: TD003
 @pytest.mark.array_type(skip=Flags.Disk | Flags.Gpu)
 @pytest.mark.parametrize(
@@ -149,7 +175,7 @@ def test_dask_constant_blocks(
 
 @pytest.mark.benchmark
 @pytest.mark.array_type(skip=Flags.Matrix | Flags.Dask | Flags.Disk | Flags.Gpu)
-@pytest.mark.parametrize("func", [stats.sum, stats.mean, stats.is_constant])
+@pytest.mark.parametrize("func", [stats.sum, stats.mean, stats.mean_var, stats.is_constant])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])  # random only supports float
 def test_stats_benchmark(
     benchmark: BenchmarkFixture,