diff --git a/src/array_api_extra/__init__.py b/src/array_api_extra/__init__.py index dd37953e..d901a7f9 100644 --- a/src/array_api_extra/__init__.py +++ b/src/array_api_extra/__init__.py @@ -2,6 +2,7 @@ from ._delegation import ( atleast_nd, + cov, expand_dims, isclose, nan_to_num, @@ -13,7 +14,6 @@ from ._lib._funcs import ( apply_where, broadcast_shapes, - cov, create_diagonal, default_dtype, kron, diff --git a/src/array_api_extra/_delegation.py b/src/array_api_extra/_delegation.py index d9e45838..75b6caeb 100644 --- a/src/array_api_extra/_delegation.py +++ b/src/array_api_extra/_delegation.py @@ -18,7 +18,95 @@ from ._lib._utils._helpers import asarrays from ._lib._utils._typing import Array, DType -__all__ = ["expand_dims", "isclose", "nan_to_num", "one_hot", "pad", "sinc"] +__all__ = [ + "cov", + "expand_dims", + "isclose", + "nan_to_num", + "one_hot", + "pad", + "sinc", +] + + +def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array: + """ + Estimate a covariance matrix. + + Covariance indicates the level to which two variables vary together. + If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`, + then the covariance matrix element :math:`C_{ij}` is the covariance of + :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance + of :math:`x_i`. + + This provides a subset of the functionality of ``numpy.cov``. + + Parameters + ---------- + m : array + A 1-D or 2-D array containing multiple variables and observations. + Each row of `m` represents a variable, and each column a single + observation of all those variables. + xp : array_namespace, optional + The standard-compatible namespace for `m`. Default: infer. + + Returns + ------- + array + The covariance matrix of the variables. + + Examples + -------- + >>> import array_api_strict as xp + >>> import array_api_extra as xpx + + Consider two variables, :math:`x_0` and :math:`x_1`, which + correlate perfectly, but in opposite directions: + + >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T + >>> x + Array([[0, 1, 2], + [2, 1, 0]], dtype=array_api_strict.int64) + + Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance + matrix shows this clearly: + + >>> xpx.cov(x, xp=xp) + Array([[ 1., -1.], + [-1., 1.]], dtype=array_api_strict.float64) + + Note that element :math:`C_{0,1}`, which shows the correlation between + :math:`x_0` and :math:`x_1`, is negative. + + Further, note how `x` and `y` are combined: + + >>> x = xp.asarray([-2.1, -1, 4.3]) + >>> y = xp.asarray([3, 1.1, 0.12]) + >>> X = xp.stack((x, y), axis=0) + >>> xpx.cov(X, xp=xp) + Array([[11.71 , -4.286 ], + [-4.286 , 2.14413333]], dtype=array_api_strict.float64) + + >>> xpx.cov(x, xp=xp) + Array(11.71, dtype=array_api_strict.float64) + + >>> xpx.cov(y, xp=xp) + Array(2.14413333, dtype=array_api_strict.float64) + """ + + if xp is None: + xp = array_namespace(m) + + if ( + is_numpy_namespace(xp) + or is_cupy_namespace(xp) + or is_torch_namespace(xp) + or is_dask_namespace(xp) + or is_jax_namespace(xp) + ): + return xp.cov(m) + + return _funcs.cov(m, xp=xp) def expand_dims( diff --git a/src/array_api_extra/_lib/_funcs.py b/src/array_api_extra/_lib/_funcs.py index 88703ecc..65c78b04 100644 --- a/src/array_api_extra/_lib/_funcs.py +++ b/src/array_api_extra/_lib/_funcs.py @@ -248,73 +248,8 @@ def broadcast_shapes(*shapes: tuple[float | None, ...]) -> tuple[int | None, ... return tuple(out) -def cov(m: Array, /, *, xp: ModuleType | None = None) -> Array: - """ - Estimate a covariance matrix. - - Covariance indicates the level to which two variables vary together. - If we examine N-dimensional samples, :math:`X = [x_1, x_2, ... x_N]^T`, - then the covariance matrix element :math:`C_{ij}` is the covariance of - :math:`x_i` and :math:`x_j`. The element :math:`C_{ii}` is the variance - of :math:`x_i`. - - This provides a subset of the functionality of ``numpy.cov``. - - Parameters - ---------- - m : array - A 1-D or 2-D array containing multiple variables and observations. - Each row of `m` represents a variable, and each column a single - observation of all those variables. - xp : array_namespace, optional - The standard-compatible namespace for `m`. Default: infer. - - Returns - ------- - array - The covariance matrix of the variables. - - Examples - -------- - >>> import array_api_strict as xp - >>> import array_api_extra as xpx - - Consider two variables, :math:`x_0` and :math:`x_1`, which - correlate perfectly, but in opposite directions: - - >>> x = xp.asarray([[0, 2], [1, 1], [2, 0]]).T - >>> x - Array([[0, 1, 2], - [2, 1, 0]], dtype=array_api_strict.int64) - - Note how :math:`x_0` increases while :math:`x_1` decreases. The covariance - matrix shows this clearly: - - >>> xpx.cov(x, xp=xp) - Array([[ 1., -1.], - [-1., 1.]], dtype=array_api_strict.float64) - - Note that element :math:`C_{0,1}`, which shows the correlation between - :math:`x_0` and :math:`x_1`, is negative. - - Further, note how `x` and `y` are combined: - - >>> x = xp.asarray([-2.1, -1, 4.3]) - >>> y = xp.asarray([3, 1.1, 0.12]) - >>> X = xp.stack((x, y), axis=0) - >>> xpx.cov(X, xp=xp) - Array([[11.71 , -4.286 ], - [-4.286 , 2.14413333]], dtype=array_api_strict.float64) - - >>> xpx.cov(x, xp=xp) - Array(11.71, dtype=array_api_strict.float64) - - >>> xpx.cov(y, xp=xp) - Array(2.14413333, dtype=array_api_strict.float64) - """ - if xp is None: - xp = array_namespace(m) - +def cov(m: Array, /, *, xp: ModuleType) -> Array: # numpydoc ignore=PR01,RT01 + """See docstring in array_api_extra._delegation.""" m = xp.asarray(m, copy=True) dtype = ( xp.float64 if xp.isdtype(m.dtype, "integral") else xp.result_type(m, xp.float64) diff --git a/tests/test_funcs.py b/tests/test_funcs.py index 7da01591..ef627eb1 100644 --- a/tests/test_funcs.py +++ b/tests/test_funcs.py @@ -419,38 +419,44 @@ def test_none(self, args: tuple[tuple[float | None, ...], ...]): class TestCov: def test_basic(self, xp: ModuleType): xp_assert_close( - cov(xp.asarray([[0, 2], [1, 1], [2, 0]]).T), + cov(xp.asarray([[0, 2], [1, 1], [2, 0]], dtype=xp.float64).T), xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64), ) def test_complex(self, xp: ModuleType): - actual = cov(xp.asarray([[1, 2, 3], [1j, 2j, 3j]])) + actual = cov(xp.asarray([[1, 2, 3], [1j, 2j, 3j]], dtype=xp.complex128)) expect = xp.asarray([[1.0, -1.0j], [1.0j, 1.0]], dtype=xp.complex128) xp_assert_close(actual, expect) + @pytest.mark.xfail_xp_backend(Backend.JAX, reason="jax#32296") @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="sparse#877") def test_empty(self, xp: ModuleType): with warnings.catch_warnings(record=True): warnings.simplefilter("always", RuntimeWarning) - xp_assert_equal(cov(xp.asarray([])), xp.asarray(xp.nan, dtype=xp.float64)) + warnings.simplefilter("always", UserWarning) xp_assert_equal( - cov(xp.reshape(xp.asarray([]), (0, 2))), + cov(xp.asarray([], dtype=xp.float64)), + xp.asarray(xp.nan, dtype=xp.float64), + ) + xp_assert_equal( + cov(xp.reshape(xp.asarray([], dtype=xp.float64), (0, 2))), xp.reshape(xp.asarray([], dtype=xp.float64), (0, 0)), ) xp_assert_equal( - cov(xp.reshape(xp.asarray([]), (2, 0))), + cov(xp.reshape(xp.asarray([], dtype=xp.float64), (2, 0))), xp.asarray([[xp.nan, xp.nan], [xp.nan, xp.nan]], dtype=xp.float64), ) def test_combination(self, xp: ModuleType): - x = xp.asarray([-2.1, -1, 4.3]) - y = xp.asarray([3, 1.1, 0.12]) + x = xp.asarray([-2.1, -1, 4.3], dtype=xp.float64) + y = xp.asarray([3, 1.1, 0.12], dtype=xp.float64) X = xp.stack((x, y), axis=0) desired = xp.asarray([[11.71, -4.286], [-4.286, 2.144133]], dtype=xp.float64) xp_assert_close(cov(X), desired, rtol=1e-6) xp_assert_close(cov(x), xp.asarray(11.71, dtype=xp.float64)) xp_assert_close(cov(y), xp.asarray(2.144133, dtype=xp.float64), rtol=1e-6) + @pytest.mark.xfail_xp_backend(Backend.TORCH, reason="array-api-extra#455") def test_device(self, xp: ModuleType, device: Device): x = xp.asarray([1, 2, 3], device=device) assert get_device(cov(x)) == device @@ -458,7 +464,10 @@ def test_device(self, xp: ModuleType, device: Device): @pytest.mark.skip_xp_backend(Backend.NUMPY_READONLY, reason="xp=xp") def test_xp(self, xp: ModuleType): xp_assert_close( - cov(xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]]).T, xp=xp), + cov( + xp.asarray([[0.0, 2.0], [1.0, 1.0], [2.0, 0.0]], dtype=xp.float64).T, + xp=xp, + ), xp.asarray([[1.0, -1.0], [-1.0, 1.0]], dtype=xp.float64), )