Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ repos:
- id: trailing-whitespace
- id: no-commit-to-branch
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.12.12
rev: v0.13.0
hooks:
- id: ruff-check
args: [--fix, --exit-non-zero-on-fix]
Expand Down
3 changes: 1 addition & 2 deletions docs/conf.py
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,6 @@
source_directory="docs/",
)

_np_aliases = {"bool_": "bool"}
_np_nocls = {"float64": "attr"}
_optional_types = {
"CupyArray": "cupy.ndarray",
Expand All @@ -111,7 +110,7 @@ def find_type_alias(name: str) -> tuple[str, str, str | None] | tuple[None, None
return "class", path, None
return "data", f"fast_array_utils.{name}", None
if name.startswith("np."):
name = _np_aliases.get(name[3:], name[3:])
name = name.removeprefix("np.")
return _np_nocls.get(name, "class"), f"numpy.{name}", f"np.{name}"
if name in npt.__all__:
return "data", f"numpy.typing.{name}", None
Expand Down
10 changes: 5 additions & 5 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,9 @@ classifiers = [
"Programming Language :: Python :: 3.13",
]
dynamic = [ "description", "readme", "version" ]
dependencies = [ "numpy>=1.25.2" ]
dependencies = [ "numpy>=2" ]
optional-dependencies.accel = [ "numba>=0.57" ]
optional-dependencies.dask = [ "dask>=2023.6.1" ]
optional-dependencies.doc = [
"furo>=2024.8.6",
"pytest>=8.4",
Expand All @@ -34,7 +35,7 @@ optional-dependencies.doc = [
"sphinx-autodoc-typehints>=3.2",
"sphinx-autofixture>=0.4.1",
]
optional-dependencies.full = [ "dask", "fast-array-utils[accel,sparse]", "h5py", "zarr" ]
optional-dependencies.full = [ "fast-array-utils[accel,dask,sparse]", "h5py", "zarr" ]
optional-dependencies.sparse = [ "scipy>=1.11" ]
optional-dependencies.test = [
"anndata",
Expand Down Expand Up @@ -99,7 +100,7 @@ overrides.matrix.resolution.features = [
]
overrides.matrix.resolution.dependencies = [
# TODO: move to min dep once this is fixed: https://github.com/tlambert03/hatch-min-requirements/issues/5
{ if = [ "lowest" ], value = "dask==2023.5.1" },
{ if = [ "lowest" ], value = "dask==2023.6.1" },
]

[[tool.hatch.envs.hatch-test.matrix]]
Expand Down Expand Up @@ -144,13 +145,12 @@ lint.per-file-ignores."typings/**/*.pyi" = [ "A002", "F403", "F405", "N801" ] #
lint.allowed-confusables = [ "×", "’" ]
lint.flake8-bugbear.extend-immutable-calls = [ "testing.fast_array_utils.Flags" ]
lint.flake8-copyright.notice-rgx = "SPDX-License-Identifier: MPL-2\\.0"
lint.flake8-tidy-imports.banned-api."numpy.bool".msg = "Use `np.bool_` instead for numpy>=1.24<2 compatibility"
lint.flake8-type-checking.exempt-modules = [ ]
lint.flake8-type-checking.strict = true
lint.isort.known-first-party = [ "fast_array_utils" ]
lint.isort.lines-after-imports = 2
lint.isort.required-imports = [ "from __future__ import annotations" ]
lint.pydocstyle.convention = "numpy"
lint.future-annotations = true

[tool.pytest.ini_options]
addopts = [
Expand Down
24 changes: 12 additions & 12 deletions src/fast_array_utils/stats/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,19 +29,19 @@
@overload
def is_constant(x: NDArray[Any] | types.CSBase | types.CupyArray, /, *, axis: None = None) -> bool: ...
@overload
def is_constant(x: NDArray[Any] | types.CSBase, /, *, axis: Literal[0, 1]) -> NDArray[np.bool_]: ...
def is_constant(x: NDArray[Any] | types.CSBase, /, *, axis: Literal[0, 1]) -> NDArray[np.bool]: ...
@overload
def is_constant(x: types.CupyArray, /, *, axis: Literal[0, 1]) -> types.CupyArray: ...
@overload
def is_constant(x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None) -> types.DaskArray: ...
def is_constant(x: types.DaskArray, /, *, axis: Literal[0, 1] | None = None) -> types.DaskArray: ...


def is_constant(
x: NDArray[Any] | types.CSBase | types.CupyArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
) -> bool | NDArray[np.bool_] | types.CupyArray | types.DaskArray:
axis: Literal[0, 1] | None = None,
) -> bool | NDArray[np.bool] | types.CupyArray | types.DaskArray:
"""Check whether values in array are constant.

Parameters
Expand Down Expand Up @@ -80,7 +80,7 @@ def is_constant(
# TODO(flying-sheep): support CSDataset (TODO)
# https://github.com/scverse/fast-array-utils/issues/52
@overload
def mean(x: CpuArray | GpuArray | DiskArray, /, *, axis: Literal[None] = None, dtype: DTypeLike | None = None) -> np.number[Any]: ...
def mean(x: CpuArray | GpuArray | DiskArray, /, *, axis: None = None, dtype: DTypeLike | None = None) -> np.number[Any]: ...
@overload
def mean(x: CpuArray | DiskArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None) -> NDArray[np.number[Any]]: ...
@overload
Expand All @@ -93,7 +93,7 @@ def mean(
x: CpuArray | GpuArray | DiskArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
) -> NDArray[np.number[Any]] | types.CupyArray | np.number[Any] | types.DaskArray:
"""Mean over both or one axis.
Expand Down Expand Up @@ -131,24 +131,24 @@ def mean(
from ._mean import mean_

validate_axis(x.ndim, axis)
return mean_(x, axis=axis, dtype=dtype) # type: ignore[no-any-return] # literally the same type, wtf mypy
return mean_(x, axis=axis, dtype=dtype)


@overload
def mean_var(x: CpuArray | GpuArray, /, *, axis: Literal[None] = None, correction: int = 0) -> tuple[np.float64, np.float64]: ...
def mean_var(x: CpuArray | GpuArray, /, *, axis: None = None, correction: int = 0) -> tuple[np.float64, np.float64]: ...
@overload
def mean_var(x: CpuArray, /, *, axis: Literal[0, 1], correction: int = 0) -> tuple[NDArray[np.float64], NDArray[np.float64]]: ...
@overload
def mean_var(x: GpuArray, /, *, axis: Literal[0, 1], correction: int = 0) -> tuple[types.CupyArray, types.CupyArray]: ...
@overload
def mean_var(x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None, correction: int = 0) -> tuple[types.DaskArray, types.DaskArray]: ...
def mean_var(x: types.DaskArray, /, *, axis: Literal[0, 1] | None = None, correction: int = 0) -> tuple[types.DaskArray, types.DaskArray]: ...


def mean_var(
x: CpuArray | GpuArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
correction: int = 0,
) -> (
tuple[np.float64, np.float64]
Expand Down Expand Up @@ -218,14 +218,14 @@ def sum(x: GpuArray, /, *, axis: Literal[0, 1], dtype: DTypeLike | None = None,


@overload
def sum(x: types.DaskArray, /, *, axis: Literal[0, 1, None] = None, dtype: DTypeLike | None = None, keep_cupy_as_array: bool = False) -> types.DaskArray: ...
def sum(x: types.DaskArray, /, *, axis: Literal[0, 1] | None = None, dtype: DTypeLike | None = None, keep_cupy_as_array: bool = False) -> types.DaskArray: ...


def sum(
x: CpuArray | GpuArray | DiskArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
keep_cupy_as_array: bool = False,
) -> NDArray[Any] | types.CupyArray | np.number[Any] | types.DaskArray:
Expand Down
20 changes: 10 additions & 10 deletions src/fast_array_utils/stats/_is_constant.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,13 @@ def is_constant_(
a: NDArray[Any] | types.CSBase | types.CupyArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
) -> bool | NDArray[np.bool_] | types.CupyArray | types.DaskArray: # pragma: no cover
axis: Literal[0, 1] | None = None,
) -> bool | NDArray[np.bool] | types.CupyArray | types.DaskArray: # pragma: no cover
raise NotImplementedError


@is_constant_.register(np.ndarray | types.CupyArray)
def _is_constant_ndarray(a: NDArray[Any] | types.CupyArray, /, *, axis: Literal[0, 1, None] = None) -> bool | NDArray[np.bool_] | types.CupyArray:
def _is_constant_ndarray(a: NDArray[Any] | types.CupyArray, /, *, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool] | types.CupyArray:
# Should eventually support nd, not now.
match axis:
case None:
Expand All @@ -41,13 +41,13 @@ def _is_constant_ndarray(a: NDArray[Any] | types.CupyArray, /, *, axis: Literal[
return _is_constant_rows(a)


def _is_constant_rows(a: NDArray[Any] | types.CupyArray) -> NDArray[np.bool_] | types.CupyArray:
def _is_constant_rows(a: NDArray[Any] | types.CupyArray) -> NDArray[np.bool] | types.CupyArray:
b = np.broadcast_to(a[:, 0][:, np.newaxis], a.shape)
return cast("NDArray[np.bool_]", (a == b).all(axis=1))
return cast("NDArray[np.bool]", (a == b).all(axis=1))


@is_constant_.register(types.CSBase)
def _is_constant_cs(a: types.CSBase, /, *, axis: Literal[0, 1, None] = None) -> bool | NDArray[np.bool_]:
def _is_constant_cs(a: types.CSBase, /, *, axis: Literal[0, 1] | None = None) -> bool | NDArray[np.bool]:
from . import is_constant

if len(a.shape) == 1: # pragma: no cover
Expand All @@ -68,9 +68,9 @@ def _is_constant_cs(a: types.CSBase, /, *, axis: Literal[0, 1, None] = None) ->


@numba.njit(cache=True)
def _is_constant_cs_major(a: types.CSBase, shape: tuple[int, int]) -> NDArray[np.bool_]:
def _is_constant_cs_major(a: types.CSBase, shape: tuple[int, int]) -> NDArray[np.bool]:
n = len(a.indptr) - 1
result = np.ones(n, dtype=np.bool_)
result = np.ones(n, dtype=np.bool)
for i in numba.prange(n):
start = a.indptr[i]
stop = a.indptr[i + 1]
Expand All @@ -83,13 +83,13 @@ def _is_constant_cs_major(a: types.CSBase, shape: tuple[int, int]) -> NDArray[np


@is_constant_.register(types.DaskArray)
def _is_constant_dask(a: types.DaskArray, /, *, axis: Literal[0, 1, None] = None) -> types.DaskArray:
def _is_constant_dask(a: types.DaskArray, /, *, axis: Literal[0, 1] | None = None) -> types.DaskArray:
import dask.array as da

from . import is_constant

if axis is not None:
return da.map_blocks(partial(is_constant, axis=axis), a, drop_axis=axis, meta=np.array([], dtype=np.bool_))
return da.map_blocks(partial(is_constant, axis=axis), a, drop_axis=axis, meta=np.array([], dtype=np.bool))

rv = (
(a == a[0, 0].compute()).all()
Expand Down
7 changes: 3 additions & 4 deletions src/fast_array_utils/stats/_mean.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: MPL-2.0
from __future__ import annotations

from typing import TYPE_CHECKING, no_type_check
from typing import TYPE_CHECKING

import numpy as np

Expand All @@ -17,14 +17,13 @@
from ..typing import CpuArray, DiskArray, GpuArray


@no_type_check # mypy is very confused
def mean_(
x: CpuArray | GpuArray | DiskArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
) -> NDArray[np.number[Any]] | np.number[Any] | types.DaskArray:
total = sum_(x, axis=axis, dtype=dtype)
n = np.prod(x.shape) if axis is None else x.shape[axis]
return total / n
return total / n # type: ignore[call-overload,operator,return-value]
2 changes: 1 addition & 1 deletion src/fast_array_utils/stats/_mean_var.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ def mean_var_(
x: CpuArray | GpuArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
correction: int = 0,
) -> (
tuple[NDArray[np.float64], NDArray[np.float64]]
Expand Down
14 changes: 7 additions & 7 deletions src/fast_array_utils/stats/_sum.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,15 @@

from ..typing import CpuArray, DiskArray, GpuArray

ComplexAxis: TypeAlias = tuple[Literal[0], Literal[1]] | tuple[Literal[0, 1]] | Literal[0, 1, None]
ComplexAxis: TypeAlias = tuple[Literal[0], Literal[1]] | tuple[Literal[0, 1]] | Literal[0, 1] | None


@singledispatch
def sum_(
x: CpuArray | GpuArray | DiskArray | types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
keep_cupy_as_array: bool = False,
) -> NDArray[Any] | np.number[Any] | types.CupyArray | types.DaskArray:
Expand All @@ -43,7 +43,7 @@ def _sum_cupy(
x: GpuArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
keep_cupy_as_array: bool = False,
) -> types.CupyArray | np.number[Any]:
Expand All @@ -56,7 +56,7 @@ def _sum_cs(
x: types.CSBase,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
keep_cupy_as_array: bool = False,
) -> NDArray[Any] | np.number[Any]:
Expand All @@ -76,7 +76,7 @@ def _sum_dask(
x: types.DaskArray,
/,
*,
axis: Literal[0, 1, None] = None,
axis: Literal[0, 1] | None = None,
dtype: DTypeLike | None = None,
keep_cupy_as_array: bool = False,
) -> types.DaskArray:
Expand Down Expand Up @@ -129,7 +129,7 @@ def sum_dask_inner(
return cast("NDArray[Any] | types.CupyArray", rv.reshape(shape))


def normalize_axis(axis: ComplexAxis, ndim: int) -> Literal[0, 1, None]:
def normalize_axis(axis: ComplexAxis, ndim: int) -> Literal[0, 1] | None:
"""Adapt `axis` parameter passed by Dask to what we support."""
match axis:
case int() | None:
Expand All @@ -145,7 +145,7 @@ def normalize_axis(axis: ComplexAxis, ndim: int) -> Literal[0, 1, None]:
return axis


def get_shape(a: NDArray[Any] | np.number[Any] | types.CupyArray, *, axis: Literal[0, 1, None], keepdims: bool) -> tuple[int] | tuple[int, int]:
def get_shape(a: NDArray[Any] | np.number[Any] | types.CupyArray, *, axis: Literal[0, 1] | None, keepdims: bool) -> tuple[int] | tuple[int, int]:
"""Get the output shape of an axis-flattening operation."""
match keepdims, a.ndim:
case False, 0:
Expand Down
Loading
Loading