diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c15e379..96bd127 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -54,16 +54,18 @@ jobs: steps: - uses: actions/checkout@v7 with: { fetch-depth: 0, filter: "blob:none" } + # https://github.com/pypa/hatch/pull/2382 + - run: uv tool install 'hatch @ git+https://github.com/pypa/hatch.git@ed8e30bebf98f2fe4d70c18a32a50a8160c391cb' - uses: astral-sh/setup-uv@v9.0.0 with: python-version: ${{ matrix.env.python }} - name: create environment - run: uvx hatch -v env create ${{ matrix.env.name }} + run: hatch -v env create ${{ matrix.env.name }} - name: run tests with coverage run: | - uvx hatch run ${{ matrix.env.name }}:run-cov + hatch run ${{ matrix.env.name }}:run-cov # https://github.com/codecov/codecov-cli/issues/648 - uvx hatch run ${{ matrix.env.name }}:coverage xml + hatch run ${{ matrix.env.name }}:coverage xml rm test-data/.coverage - uses: codecov/codecov-action@v7 with: diff --git a/.vscode/settings.json b/.vscode/settings.json index f1e4981..fa0dd07 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,10 +1,7 @@ { - "[toml][json][jsonc][python]": { + "[json][jsonc][python]": { "editor.formatOnSave": true, }, - "[toml]": { - "editor.defaultFormatter": "tamasfe.even-better-toml", - }, "[json][jsonc]": { "editor.defaultFormatter": "biomejs.biome", }, diff --git a/pyproject.toml b/pyproject.toml index 8127f5b..f2814b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ dynamic = [ "version" ] dependencies = [ "array-api-compat", "numpy>=2" ] optional-dependencies.accel = [ "numba>=0.57" ] -optional-dependencies.dask = [ "dask>=2023.6.1" ] +optional-dependencies.dask = [ "dask>=2024.5.1" ] optional-dependencies.full = [ "fast-array-utils[accel,dask,sparse]", "h5py", "zarr" ] optional-dependencies.sparse = [ "scipy>=1.13" ] optional-dependencies.testing = [ "packaging" ] @@ -86,7 +86,7 @@ envs.hatch-test.overrides.matrix.extras.dependency-groups = [ ] envs.hatch-test.overrides.matrix.resolution.dependencies = [ { value = "numpy==2", if = [ "lowest" ] }, - { value = "dask==2023.6.1", if = [ "lowest" ] }, + { value = "dask==2024.5.1", if = [ "lowest" ] }, { value = "scipy==1.13.0", if = [ "lowest" ] }, ] envs.hatch-test.default-args = [] diff --git a/src/fast_array_utils/stats/__init__.py b/src/fast_array_utils/stats/__init__.py index 5cfd439..0d2ac8d 100644 --- a/src/fast_array_utils/stats/__init__.py +++ b/src/fast_array_utils/stats/__init__.py @@ -91,7 +91,7 @@ def mean(x: CpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | @overload def mean(x: GpuArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None) -> types.CupyArray: ... @overload -def mean(x: types.DaskArray, /, *, axis: Literal[0, 1], dtype: ToDType[Any] | None = None) -> types.DaskArray: ... +def mean(x: types.DaskArray, /, *, axis: Literal[0, 1] | None = None, dtype: ToDType[Any] | None = None) -> types.DaskArray: ... @overload def mean[A: types.HasArrayNamespace](x: A, /, *, axis: Literal[0, 1] | None = None, dtype: DTypeLike | None = None) -> A: ... diff --git a/src/fast_array_utils/stats/_mean_var.py b/src/fast_array_utils/stats/_mean_var.py index c704bc9..4cbb07b 100644 --- a/src/fast_array_utils/stats/_mean_var.py +++ b/src/fast_array_utils/stats/_mean_var.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: MPL-2.0 from __future__ import annotations -from typing import TYPE_CHECKING, no_type_check +from collections.abc import Sequence +from typing import TYPE_CHECKING, cast, no_type_check import numba import numpy as np @@ -9,14 +10,32 @@ from .. import types from ..numba import njit from ._power import power +from ._utils import _get_shape, _normalize_axis if TYPE_CHECKING: - from typing import Literal + from collections.abc import Iterator + from typing import Any, Literal, TypedDict + from dask.array.reductions import _Chunk from numpy.typing import NDArray from ..typing import CpuArray, GpuArray + from ._utils import ComplexAxis + + class _Moments(TypedDict): + """A (count, mean, M2) triple as tracked by Chan's parallel-variance algorithm. + + `mean`/`m2` are shaped like `_get_shape`'s `keepdims=True` convention. + """ + + n: int + mean: NDArray[np.float64] + m2: NDArray[np.float64] + + # what dask passes to `combine`/`aggregate`: our own chunk output, a (nested) list + # thereof (`concatenate=False`), or a plain array while it’s computing `meta` + type MomentsIn = _Moments | CpuArray | GpuArray | Sequence[Any] @no_type_check # mypy is extremely confused @@ -32,6 +51,9 @@ def mean_var_( | tuple[np.float64, np.float64] | tuple[types.DaskArray, types.DaskArray] ): + if isinstance(x, types.DaskArray): + return _dask_mean_var(x, axis=axis, correction=correction) + from . import mean if isinstance(x, np.ndarray | types.CSBase) or not isinstance(x, types.HasArrayNamespace): @@ -45,7 +67,7 @@ def mean_var_( mean_, var = _sparse_mean_var(x, axis=axis) else: mean_ = mean(x, axis=axis, dtype=xp.float64) - mean_sq = mean(power(x, 2, dtype=xp.float64), axis=axis) if isinstance(x, types.DaskArray) else mean(power(x, 2), axis=axis, dtype=xp.float64) + mean_sq = mean(power(x, 2, dtype=xp.float64), axis=axis, dtype=xp.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] @@ -54,6 +76,125 @@ def mean_var_( return mean_, var +def _dask_mean_var(x: types.DaskArray, /, *, axis: Literal[0, 1] | None, correction: int) -> tuple[types.DaskArray, types.DaskArray]: + """Mean and variance of a dask array. + + ``mean`` is a normal (associative) sum-based dask reduction. + ``var`` is instead derived from per-chunk ``(count, mean, M2)`` triples + (``M2 = sum((x - mean)**2)``, computed per chunk by recursing into :func:`mean_var_`) + that get merged pairwise across chunks using Chan's parallel-variance algorithm. + Separately dask-reducing ``mean(x)``/``mean(x**2)`` and subtracting at the end + (the naive two-pass formula used before) loses precision once many chunks are + combined, especially for float32 data on GPUs. + https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Parallel_algorithm + """ + import dask.array as da + + from . import mean + + n = np.prod(x.shape) if axis is None else x.shape[axis] + mean_ = mean(x, axis=axis, dtype=np.float64) + # mypy can’t infer `reduction`’s type parameter from the callbacks, so pin it here + chunk: _Chunk[_Moments] = _moments_chunk + m2 = da.reduction( + x, chunk, _moments_aggregate, axis=axis, combine=_moments_combine, concatenate=False, dtype=np.float64, meta=np.array([], dtype=np.float64) + ) + if axis is None: # match `mean`/`sum`’s convention of reducing to a true scalar + m2 = m2.map_blocks(lambda a: a.reshape(())[()], meta=m2.dtype.type(0)) + denom = n - correction if correction and n != 1 else n + return mean_, m2 / denom + + +def _moments_chunk( + a: CpuArray | GpuArray, + /, + *, + axis: ComplexAxis = None, + keepdims: bool = False, + computing_meta: bool = False, + **kwargs: object, # noqa: ARG001 # `_Chunk`/`_CB` let dask pass arbitrary keywords +) -> _Moments | NDArray[np.float64]: + if computing_meta: # pragma: no cover + return np.ndarray((), dtype=np.float64) + axis_ = _normalize_axis(axis, a.ndim) + mean_, var_ = mean_var_(a, axis=axis_, correction=0) + n = int(np.prod(a.shape)) if axis_ is None else a.shape[axis_] + shape = _get_shape(mean_, axis=axis_, keepdims=keepdims) + moments: _Moments = {"n": n, "mean": mean_.reshape(shape), "m2": (var_ * n).reshape(shape)} + return moments + + +def _moments_combine( + pairs: MomentsIn, + /, + *, + axis: ComplexAxis = None, # noqa: ARG001 + keepdims: bool = False, # noqa: ARG001 + computing_meta: bool = False, + **kwargs: object, # noqa: ARG001 # `_Chunk`/`_CB` let dask pass arbitrary keywords +) -> _Moments | NDArray[np.float64]: + if computing_meta: # pragma: no cover + return np.ndarray((), dtype=np.float64) + return _combine_all(pairs) + + +def _moments_aggregate( + pairs: MomentsIn, + /, + *, + axis: ComplexAxis = None, + keepdims: bool = False, + computing_meta: bool = False, + **kwargs: object, # noqa: ARG001 # `_Chunk`/`_CB` let dask pass arbitrary keywords +) -> NDArray[np.float64]: + if computing_meta: # pragma: no cover + return np.ndarray((), dtype=np.float64) + m2 = _combine_all(pairs)["m2"] + axis_ = _normalize_axis(axis, 2) + return m2.reshape(_final_moments_shape(m2.size, axis=axis_, keepdims=keepdims)) + + +def _combine_all(pairs: MomentsIn) -> _Moments: + """Merge every moment triple in a (possibly nested) list of `_moments_chunk` outputs.""" + combined = None + for pair in _flatten_moments(pairs): + combined = pair if combined is None else _chan_combine(combined, pair) + assert combined is not None + return combined + + +def _flatten_moments(pairs: MomentsIn) -> Iterator[_Moments]: + match pairs: + case {"n": _, "mean": _, "m2": _}: + yield cast("_Moments", pairs) + case Sequence(): # `concatenate=False` means dask hands us nested lists + for pair in pairs: + yield from _flatten_moments(pair) + case _: # pragma: no cover + msg = f"Unexpected moments input: {type(pairs)}" + raise TypeError(msg) + + +def _chan_combine(a: _Moments, b: _Moments) -> _Moments: + """Pairwise-merge two ``(count, mean, M2)`` moment triples.""" + n_a, mean_a, m2_a = a["n"], a["mean"], a["m2"] + n_b, mean_b, m2_b = b["n"], b["mean"], b["m2"] + n = n_a + n_b + delta = mean_b - mean_a + mean_ = mean_a + delta * (n_b / n) + m2 = m2_a + m2_b + delta**2 * (n_a * n_b / n) + return {"n": n, "mean": mean_, "m2": m2} + + +def _final_moments_shape(size: int, *, axis: Literal[0, 1] | None, keepdims: bool) -> tuple[int, ...]: + """Shape for a fully-combined moment array, mirroring `_get_shape`'s convention.""" + if axis is None: + return (1, 1) if keepdims else (1,) + if not keepdims: + return (size,) + return (1, size) if axis == 0 else (size, 1) + + 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. diff --git a/src/fast_array_utils/stats/_utils.py b/src/fast_array_utils/stats/_utils.py index f468508..1198265 100644 --- a/src/fast_array_utils/stats/_utils.py +++ b/src/fast_array_utils/stats/_utils.py @@ -19,7 +19,11 @@ from ..typing import CpuArray, GpuArray from ._typing import DTypeKw, Ops - type ComplexAxis = tuple[Literal[0], Literal[1]] | tuple[Literal[0, 1]] | Literal[0, 1] | None + # `axis` as dask passes it into reduction callbacks: + # normalized to a tuple of non-negative ints, + # or not passed at all (hence `None`) when dask only computes a callback’s `meta`. + # `Literal[0, 1]` is what `_normalize_axis` narrows that down to. + type ComplexAxis = Literal[0, 1] | tuple[int, ...] | None __all__ = ["_dask_inner"] @@ -85,12 +89,15 @@ def _normalize_axis(axis: ComplexAxis, ndim: int) -> Literal[0, 1] | None: match axis: case int() | None: # pragma: no cover pass - case (0 | 1,): - axis = axis[0] + case (0,): + axis = 0 + case (1,): + axis = 1 case (0, 1) | (1, 0): axis = None case _: # pragma: no cover - raise AxisError(axis, ndim) + msg = f"axis {axis!r} invalid for {ndim}-dimensional array" + raise AxisError(msg) if axis == 0 and ndim == 1: return None # dask’s aggregate doesn’t know we don’t accept `axis=0` for 1D arrays return axis diff --git a/src/testing/fast_array_utils/_array_type.py b/src/testing/fast_array_utils/_array_type.py index 74bd09c..7105c7d 100644 --- a/src/testing/fast_array_utils/_array_type.py +++ b/src/testing/fast_array_utils/_array_type.py @@ -311,7 +311,7 @@ def _to_scipy_sparse( x = to_dense(x, to_cpu_memory=True) # type: ignore[arg-type] # doesn’t officially handle ArrayLike cls = cast("type[types.CSBase]", cls or self.cls) - return cls(x, dtype=dtype) # type: ignore[arg-type,misc] + return cls(x, dtype=dtype) # type: ignore[arg-type] def _to_cupy_array(self, x: ArrayLike | Array, /, *, dtype: DTypeLike | None = None) -> types.CupyArray: import cupy as cu diff --git a/tests/test_stats.py b/tests/test_stats.py index 993020a..1565191 100644 --- a/tests/test_stats.py +++ b/tests/test_stats.py @@ -316,10 +316,7 @@ def test_mean_var_sparse_32(array_type: ArrayType[types.CSArray], subtests: pyte @pytest.mark.array_type({at for at in SUPPORTED_TYPES if at.flags & Flags.Sparse and at.flags & Flags.Dask}) def test_mean_var_pbmc_dask(array_type: ArrayType[types.DaskArray], pbmc64k_reduced_raw: sps.csr_array[np.float32]) -> None: - """Test float32 precision for bigger data. - - This test is flaky for sparse-in-dask for some reason. - """ + """Test float32 precision for bigger data.""" mat = pbmc64k_reduced_raw arr = array_type(mat) diff --git a/typings/dask/array/core.pyi b/typings/dask/array/core.pyi index 0a19cc0..2e9c8f5 100644 --- a/typings/dask/array/core.pyi +++ b/typings/dask/array/core.pyi @@ -39,6 +39,7 @@ class Array: def __eq__(self, value: object, /) -> Array: ... # type: ignore[override] def __getitem__(self, index: object) -> Array: ... def all(self) -> Array: ... + def __truediv__(self, other: float | np.number[Any] | Array, /) -> Array: ... # dask methods and attrs _meta: _Array diff --git a/typings/dask/array/reductions.pyi b/typings/dask/array/reductions.pyi index 0007725..1290712 100644 --- a/typings/dask/array/reductions.pyi +++ b/typings/dask/array/reductions.pyi @@ -7,25 +7,40 @@ from numpy.typing import ArrayLike, DTypeLike, NDArray from .core import Array, _Array -class _Chunk(Protocol): +# dask only passes `computing_meta` to callbacks that accept it, so it’s optional here. +class _Chunk[T](Protocol): @overload - def __call__(self, x_chunk: _Array, /, *, weights_chunk: NDArray[Any] | None = None, axis: tuple[int, ...], keepdims: bool, **kwargs: object) -> _Array: ... + def __call__( + self, + x_chunk: _Array, + /, + *, + weights_chunk: NDArray[Any] | None = None, + axis: tuple[int, ...], + keepdims: bool, + computing_meta: bool = ..., + **kwargs: object, + ) -> _Array | T: ... @overload - def __call__(self, x_chunk: _Array, /, *, axis: tuple[int, ...], keepdims: bool, **kwargs: object) -> _Array: ... + def __call__(self, x_chunk: _Array, /, *, axis: tuple[int, ...], keepdims: bool, computing_meta: bool = ..., **kwargs: object) -> _Array | T: ... -class _CB(Protocol): - def __call__(self, x_chunk: _Array, /, *, axis: tuple[int, ...], keepdims: bool, **kwargs: object) -> _Array: ... +class _CB[T](Protocol): + # When `concatenate=False`, dask passes a (possibly nested) list of the previous + # step’s raw outputs instead of concatenating them into a single `_Array`. + def __call__( + self, x_chunk: _Array | T | Sequence[Any], /, *, axis: tuple[int, ...], keepdims: bool, computing_meta: bool = ..., **kwargs: object + ) -> _Array | T: ... -def reduction( +def reduction[T]( x: Array, - chunk: _Chunk, - aggregate: _CB, + chunk: _Chunk[T], + aggregate: _CB[T], *, axis: int | Sequence[int] | None = None, keepdims: bool = False, dtype: DTypeLike | None = None, split_every: int | Mapping[int, int] | None = None, - combine: _CB | None = None, + combine: _CB[T] | None = None, name: str | None = None, out: Array | None = None, concatenate: bool = True,