ENH: a simple cov

Author: lucascolley

src/array_api_extra/__init__.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
-from ._funcs import atleast_nd
+from ._funcs import atleast_nd, cov
 
 __version__ = "0.1.2.dev0"
 
-__all__ = ["__version__", "atleast_nd"]
+__all__ = ["__version__", "atleast_nd", "cov"]

src/array_api_extra/_funcs.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
+import warnings
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from ._typing import Array, ModuleType
 
-__all__ = ["atleast_nd"]
+__all__ = ["atleast_nd", "cov"]
 
 
 def atleast_nd(x: Array, *, ndim: int, xp: ModuleType) -> Array:
@@ -46,3 +47,48 @@ def atleast_nd(x: Array, *, ndim: int, xp: ModuleType) -> Array:
         x = xp.expand_dims(x, axis=0)
         x = atleast_nd(x, ndim=ndim, xp=xp)
     return x
+
+
+def cov(x: Array, *, xp: ModuleType) -> Array:
+    """..."""
+    x = xp.asarray(x, copy=True)
+    dtype = (
+        xp.float64 if xp.isdtype(x.dtype, "integral") else xp.result_type(x, xp.float64)
+    )
+
+    x = atleast_nd(x, ndim=2, xp=xp)
+    x = xp.astype(x, dtype)
+
+    avg = mean(x, axis=1, xp=xp)
+    fact = x.shape[1] - 1
+
+    if fact <= 0:
+        warnings.warn("Degrees of freedom <= 0 for slice", RuntimeWarning, stacklevel=2)
+        fact = 0.0
+
+    x -= avg[:, None]
+    y_transpose = x.T
+    if xp.isdtype(y_transpose.dtype, "complex floating"):
+        y_transpose = xp.conj(y_transpose)
+    c = x @ y_transpose
+    c /= fact
+    axes = tuple(axis for axis, length in enumerate(c.shape) if length == 1)
+    return xp.squeeze(c, axis=axes)
+
+
+def mean(
+    x: Array,
+    /,
+    *,
+    axis: int | tuple[int, ...] | None = None,
+    keepdims: bool = False,
+    xp: ModuleType,
+) -> Array:
+    """..."""
+    if xp.isdtype(x.dtype, "complex floating"):
+        x_real = xp.real(x)
+        x_imag = xp.imag(x)
+        mean_real = xp.mean(x_real, axis=axis, keepdims=keepdims)
+        mean_imag = xp.mean(x_imag, axis=axis, keepdims=keepdims)
+        return mean_real + (mean_imag * xp.asarray(1j))
+    return xp.mean(x, axis=axis, keepdims=keepdims)

tests/test_funcs.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
+import warnings
+
 # array-api-strict#6
 import array_api_strict as xp  # type: ignore[import-untyped]
-from numpy.testing import assert_array_equal
+from numpy.testing import assert_allclose, assert_array_equal
 
-from array_api_extra import atleast_nd
+from array_api_extra import atleast_nd, cov
 
 
 class TestAtLeastND:
@@ -67,3 +69,29 @@ def test_5D(self):
 
         y = atleast_nd(x, ndim=9, xp=xp)
         assert_array_equal(y, xp.ones((1, 1, 1, 1, 1, 1, 1, 1, 1)))
+
+
+class TestCov:
+    def test_basic(self):
+        assert_allclose(
+            cov(xp.asarray([[0, 2], [1, 1], [2, 0]]).T, xp=xp),
+            xp.asarray([[1.0, -1.0], [-1.0, 1.0]]),
+        )
+
+    def test_complex(self):
+        x = xp.asarray([[1, 2, 3], [1j, 2j, 3j]])
+        res = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]])
+        assert_allclose(cov(x, xp=xp), res)
+
+    def test_empty(self):
+        with warnings.catch_warnings(record=True):
+            warnings.simplefilter("always", RuntimeWarning)
+            assert_array_equal(cov(xp.asarray([]), xp=xp), xp.nan)
+            assert_array_equal(
+                cov(xp.reshape(xp.asarray([]), (0, 2)), xp=xp),
+                xp.reshape(xp.asarray([]), (0, 0)),
+            )
+            assert_array_equal(
+                cov(xp.reshape(xp.asarray([]), (2, 0)), xp=xp),
+                xp.asarray([[xp.nan, xp.nan], [xp.nan, xp.nan]]),
+            )