diff --git a/docs/sphinx/api-statistical.md b/docs/sphinx/api-statistical.md index 37cf3a9f..ee54a90f 100644 --- a/docs/sphinx/api-statistical.md +++ b/docs/sphinx/api-statistical.md @@ -8,6 +8,7 @@ cov nanmax + nanmean nanmin nansum ``` diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index 5b4b6cfc..8459f9df 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -13,7 +13,7 @@ from ._searching import searchsorted from ._set import isin, nunique, setdiff1d, union1d from ._sorting import argpartition, partition -from ._statistical import cov, nanmax, nanmin, nansum +from ._statistical import cov, nanmax, nanmean, nanmin, nansum __version__ = "0.11.2.dev0" @@ -37,6 +37,7 @@ "lazy_apply", "nan_to_num", "nanmax", + "nanmean", "nanmin", "nansum", "nunique", diff --git a/src/array_api_extra/_agnostic/_statistical.py b/src/array_api_extra/_agnostic/_statistical.py index a8dfbb84..fe7da0be 100644 --- a/src/array_api_extra/_agnostic/_statistical.py +++ b/src/array_api_extra/_agnostic/_statistical.py @@ -8,7 +8,7 @@ from .._lib._typing import Array, ArrayNamespace from . import _manipulation -__all__ = ["cov", "nanmax", "nanmin", "nansum"] +__all__ = ["cov", "nanmax", "nanmean", "nanmin", "nansum"] def cov( @@ -154,3 +154,29 @@ def nansum( # numpydoc ignore=PR01,RT01 """See docstring in `array_api_extra._statistical`.""" mask = xp.isnan(a) return xp.sum(xp.where(mask, xp.zeros_like(a), a), axis=axis) + + +def nanmean( # numpydoc ignore=PR01,RT01 + a: Array, + /, + *, + axis: int | tuple[int, ...] | None, + xp: ArrayNamespace, +) -> Array: + """See docstring in `array_api_extra._statistical`.""" + mask = xp.isnan(a) + sum_ = nansum(a, axis=axis, xp=xp) + count = xp.sum(xp.where(mask, xp.zeros_like(a), xp.ones_like(a)), axis=axis) + safe_count = xp.astype( + xp.where(count == 0, xp.ones_like(count), count), + sum_.dtype, + copy=False, + ) + result = sum_ / safe_count + if xp.any(count == 0): + result = xp.where( + count == 0, + xp.full_like(result, xp.nan), + result, + ) + return result diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py index 492e65eb..752860fb 100644 --- a/src/array_api_extra/_statistical.py +++ b/src/array_api_extra/_statistical.py @@ -4,7 +4,7 @@ from ._lib import _compat from ._lib._typing import Array, ArrayNamespace -__all__ = ["cov", "nanmax", "nanmin", "nansum"] +__all__ = ["cov", "nanmax", "nanmean", "nanmin", "nansum"] def cov( @@ -321,6 +321,58 @@ def nanmax( return _agnostic._statistical.nanmax(a, axis=axis, xp=xp) +def nanmean( + a: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + xp: ArrayNamespace | None = None, +) -> Array: + """ + Return the mean of the array elements along a given axis, ignoring NaNs. + + Parameters + ---------- + a : Array + Input array. + axis : int or tuple of ints or None, optional + Axis or axes along which the mean is computed. The default is to compute + the mean of the flattened array. + xp : array_namespace, optional + The standard-compatible namespace for `a`. Default: infer. + + Returns + ------- + array + An array of mean values along the given axis, ignoring NaNs. + + Examples + -------- + >>> import array_api_extra as xpx + >>> import array_api_strict as xp + >>> a = xp.asarray([[5, 3, xp.nan, 1], [4, xp.nan, 2, xp.nan]]) + >>> xpx.nanmean(a) + Array(3., dtype=array_api_strict.float64) + >>> xpx.nanmean(a, axis=0) + Array([4.5, 3., 2., 1.], dtype=array_api_strict.float64) + >>> xpx.nanmean(a, axis=1) + Array([3., 3.], dtype=array_api_strict.float64) + """ + if xp is None: + xp = _compat.array_namespace(a) + + if ( + _compat.is_numpy_namespace(xp) + or _compat.is_cupy_namespace(xp) + or _compat.is_dask_namespace(xp) + or _compat.is_jax_namespace(xp) + or _compat.is_torch_namespace(xp) + ): + return xp.nanmean(a, axis=axis) + + return _agnostic._statistical.nanmean(a, axis=axis, xp=xp) + + def nansum( a: Array, /, diff --git a/tests/main/test_statistical.py b/tests/main/test_statistical.py index 91319645..6732cdd9 100644 --- a/tests/main/test_statistical.py +++ b/tests/main/test_statistical.py @@ -5,7 +5,7 @@ import numpy as np import pytest -from array_api_extra import cov, nanmax, nanmin, nansum +from array_api_extra import cov, nanmax, nanmean, nanmin, nansum from array_api_extra._lib._backends import Backend from array_api_extra._lib._compat import array_namespace from array_api_extra._lib._compat import device as get_device @@ -13,6 +13,7 @@ from array_api_extra.testing import assert_close, assert_equal, lazy_xp_function lazy_xp_function(cov) +lazy_xp_function(nanmean) lazy_xp_function(nansum) @@ -506,3 +507,66 @@ def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespa res = nansum(a, axis=axis, xp=xp) expected = xp.asarray(expected_list) assert_equal(res, expected) + + +class TestNanMean: + def test_simple(self, xp: ArrayNamespace): + a = xp.asarray([[1.0, 2.0], [3.0, xp.nan]]) + + res = nanmean(a) + assert res == 2.0 + + res = nanmean(a, axis=0) + expected = xp.asarray([2.0, 2.0]) + assert_equal(res, expected) + + res = nanmean(a, axis=1) + expected = xp.asarray([1.5, 3.0]) + assert_equal(res, expected) + + def test_bigger(self, xp: ArrayNamespace): + a = xp.asarray( + [ + [1.0, xp.nan, 4.0, 5.0], + [xp.nan, -2.0, xp.nan, -4.0], + [2.0, 1.0, 3.0, xp.nan], + ] + ) + + res = nanmean(a, axis=0) + expected = xp.asarray([1.5, -0.5, 3.5, 0.5]) + assert_equal(res, expected) + + res = nanmean(a, axis=1) + expected = xp.asarray([3.3333333, -3.0, 2.0]) + assert_close(res, expected) + + @pytest.mark.filterwarnings("ignore:.*Mean of empty slice.*:RuntimeWarning") + def test_all_nan_slice(self, xp: ArrayNamespace): + a = xp.asarray([[xp.nan, 1.0], [xp.nan, xp.nan]]) + + res = nanmean(a, axis=0, xp=xp) + expected = xp.asarray([xp.nan, 1.0]) + assert_equal(res, expected) + + def test_scalar(self, xp: ArrayNamespace): + a = xp.asarray(1.0) + assert nanmean(a) == 1.0 + + @pytest.mark.skip_xp_backend( + Backend.TORCH, reason="torch.nanmean does not support tensors on meta device" + ) + @pytest.mark.parametrize("axis", [None, 0, 1]) + def test_device(self, axis: int | None, xp: ArrayNamespace, device: Device): + a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 5.0, xp.nan]], device=device) + res = nanmean(a, axis=axis) + assert get_device(res) == device + + @pytest.mark.parametrize( + ("axis", "expected_list"), [(0, [3.0, 5.0, 1.0]), (1, [2.5, 3.5])] + ) + def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespace): + a = xp.asarray([[4.0, xp.nan, 1.0], [2.0, 5.0, xp.nan]]) + res = nanmean(a, axis=axis, xp=xp) + expected = xp.asarray(expected_list) + assert_equal(res, expected)