ENH: add nanmin (#804)

Author: OmarManzoor

docs/api-assorted.md

@@ -21,6 +21,7 @@
     isin
     kron
     nan_to_num
+    nanmin
     nunique
     one_hot
     pad

src/array_api_extra/__init__.py

@@ -13,6 +13,7 @@
     isin,
     kron,
     nan_to_num,
+    nanmin,
     one_hot,
     pad,
     partition,
@@ -54,6 +55,7 @@
     "kron",
     "lazy_apply",
     "nan_to_num",
+    "nanmin",
     "nunique",
     "one_hot",
     "pad",

src/array_api_extra/_delegation.py

@@ -35,6 +35,7 @@
     "isin",
     "kron",
     "nan_to_num",
+    "nanmin",
     "one_hot",
     "pad",
     "partition",
@@ -1583,3 +1584,54 @@ def unravel_index(
         return xp.unravel_index(indices, shape)
 
     return _funcs.unravel_index(indices, shape)
+
+
+def nanmin(
+    a: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None = None,
+    xp: ModuleType | None = None,
+) -> Array:
+    """
+    Return the minimum 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 minimum is computed. The default is to compute
+        the minimum of the flattened array.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `a`. Default: infer.
+
+    Returns
+    -------
+    array
+        An array of minimum 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.nanmin(a)
+    Array(1., dtype=array_api_strict.float64)
+    >>> xpx.nanmin(a, axis=0)
+    Array([4., 3., 2., 1.], dtype=array_api_strict.float64)
+    >>> xpx.nanmin(a, axis=1)
+    Array([1., 2.], dtype=array_api_strict.float64)
+    """
+    if xp is None:
+        xp = array_namespace(a)
+
+    if (
+        is_numpy_namespace(xp)
+        or is_cupy_namespace(xp)
+        or is_dask_namespace(xp)
+        or is_jax_namespace(xp)
+    ):
+        return xp.nanmin(a, axis=axis)
+
+    return _funcs.nanmin(a, axis=axis, xp=xp)

src/array_api_extra/_lib/_funcs.py

@@ -38,6 +38,7 @@
     "isin",
     "kron",
     "nan_to_num",
+    "nanmin",
     "nunique",
     "one_hot",
     "pad",
@@ -823,3 +824,24 @@ def unravel_index(indices: Array, shape: tuple[int, ...], /) -> tuple[Array, ...
         coords.append(indices % dim)
         indices = indices // dim
     return tuple(reversed(coords))
+
+
+def nanmin(  # numpydoc ignore=PR01,RT01
+    a: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None,
+    xp: ModuleType,
+) -> Array:
+    """See docstring in `array_api_extra._delegation.py`."""
+    mask = xp.isnan(a)
+    device_a = _compat.device(a)
+    x = xp.min(
+        xp.where(mask, xp.asarray(+xp.inf, dtype=a.dtype, device=device_a), a),
+        axis=axis,
+    )
+    # Replace Infs from all NaN slices with NaN again
+    mask = xp.all(mask, axis=axis)
+    if xp.any(mask):
+        x = xp.where(mask, xp.asarray(xp.nan, dtype=x.dtype, device=device_a), x)
+    return x

tests/test_funcs.py

@@ -30,6 +30,7 @@
     isin,
     kron,
     nan_to_num,
+    nanmin,
     nunique,
     one_hot,
     pad,
@@ -2203,3 +2204,84 @@ def test_xp(self, xp: ModuleType):
         res = unravel_index(indices, shape, xp=xp)
         for res_arr, exp_arr in zip(res, expected, strict=True):
             assert_equal(res_arr, exp_arr)
+
+
+class TestNanMin:
+    def test_simple(self, xp: ModuleType):
+        a = xp.asarray([[1, 2], [3, xp.nan]])
+
+        # with the default `axis=None` a single scalar is returned
+        res = nanmin(a)
+        expected = 1.0
+        assert res == expected
+
+        res = nanmin(a, axis=0)
+        expected = xp.asarray([1.0, 2.0])
+        assert_equal(res, expected)
+
+        res = nanmin(a, axis=1)
+        expected = xp.asarray([1.0, 3.0])
+        assert_equal(res, expected)
+
+    def test_bigger(self, xp: ModuleType):
+        a = xp.asarray(
+            [
+                [1, xp.nan, 4, 5],
+                [xp.nan, -2, xp.nan, -4],
+                [2, 1, 3, xp.nan],
+            ]
+        )
+
+        res = nanmin(a, axis=0)
+        expected = xp.asarray([1.0, -2.0, 3.0, -4.0])
+        assert_equal(res, expected)
+
+        res = nanmin(a, axis=1)
+        expected = xp.asarray([1.0, -4.0, 1.0])
+        assert_equal(res, expected)
+
+    def test_with_infinity(self, xp: ModuleType):
+        a = xp.asarray([0.1, 1.0, xp.nan, xp.inf])
+        res = nanmin(a)
+        expected = 0.1
+        assert res == expected
+
+        a = xp.asarray([0.1, 1.0, xp.nan, -xp.inf])
+        res = nanmin(a)
+        expected = -xp.inf
+        assert res == expected
+
+    def test_scalar(self, xp: ModuleType):
+        a = xp.asarray(1.0)
+        assert nanmin(a) == 1.0
+
+    @pytest.mark.filterwarnings("ignore:.*All-NaN slice*.:RuntimeWarning")
+    def test_all_nan_slice_2d(self, xp: ModuleType):
+        a = xp.asarray(
+            [
+                [xp.nan, 5.0],
+                [xp.nan, 2.0],
+            ]
+        )
+
+        res = nanmin(a, axis=0, xp=xp)
+        expected = xp.asarray([xp.nan, 2.0])
+        assert_equal(res, expected)
+
+    @pytest.mark.skip_xp_backend(
+        Backend.TORCH, reason="torch.nanmin does not support tensors on meta device"
+    )
+    @pytest.mark.parametrize("axis", [None, 0, 1])
+    def test_device(self, axis: int | None, xp: ModuleType, device: Device):
+        a = xp.asarray([[4, xp.nan, 1], [2, 5, xp.nan]], device=device)
+        res = nanmin(a, axis=axis)
+        assert get_device(res) == device
+
+    @pytest.mark.parametrize(
+        ("axis", "expected_list"), [(0, [2.0, 3.0, 1.0]), (1, [1.0, 2.0])]
+    )
+    def test_xp(self, axis: int | None, expected_list: list[float], xp: ModuleType):
+        a = xp.asarray([[4, xp.nan, 1], [2, 3, xp.nan]])
+        res = nanmin(a, axis=axis, xp=xp)
+        expected = xp.asarray(expected_list)
+        assert_equal(res, expected)