diff --git a/doc/whats_new/v1.4.rst b/doc/whats_new/v1.4.rst index adb362d2802d5..953c4906d3fb2 100644 --- a/doc/whats_new/v1.4.rst +++ b/doc/whats_new/v1.4.rst @@ -320,6 +320,13 @@ Changelog with `np.int64` indices are not supported. :pr:`27240` by :user:`Yao Xiao `. +:mod:`sklearn.covariance` +......................... + +- |Enhancement| Allow :func:`covariance.shrunk_covariance` to process + multiple covariance matrices at once by handling nd-arrays. + :pr:`25275` by :user:`Quentin Barthélemy `. + - |API| |FIX| :class:`~compose.ColumnTransformer` now replaces `"passthrough"` with a corresponding :class:`~preprocessing.FunctionTransformer` in the fitted ``transformers_`` attribute. diff --git a/sklearn/covariance/_shrunk_covariance.py b/sklearn/covariance/_shrunk_covariance.py index 5a568192dd3c3..3a79afa30729f 100644 --- a/sklearn/covariance/_shrunk_covariance.py +++ b/sklearn/covariance/_shrunk_covariance.py @@ -109,14 +109,14 @@ def _oas(X, *, assume_centered=False): prefer_skip_nested_validation=True, ) def shrunk_covariance(emp_cov, shrinkage=0.1): - """Calculate a covariance matrix shrunk on the diagonal. + """Calculate covariance matrices shrunk on the diagonal. Read more in the :ref:`User Guide `. Parameters ---------- - emp_cov : array-like of shape (n_features, n_features) - Covariance matrix to be shrunk. + emp_cov : array-like of shape (..., n_features, n_features) + Covariance matrices to be shrunk, at least 2D ndarray. shrinkage : float, default=0.1 Coefficient in the convex combination used for the computation @@ -124,8 +124,8 @@ def shrunk_covariance(emp_cov, shrinkage=0.1): Returns ------- - shrunk_cov : ndarray of shape (n_features, n_features) - Shrunk covariance. + shrunk_cov : ndarray of shape (..., n_features, n_features) + Shrunk covariance matrices. Notes ----- @@ -135,12 +135,13 @@ def shrunk_covariance(emp_cov, shrinkage=0.1): where `mu = trace(cov) / n_features`. """ - emp_cov = check_array(emp_cov) - n_features = emp_cov.shape[0] + emp_cov = check_array(emp_cov, allow_nd=True) + n_features = emp_cov.shape[-1] - mu = np.trace(emp_cov) / n_features shrunk_cov = (1.0 - shrinkage) * emp_cov - shrunk_cov.flat[:: n_features + 1] += shrinkage * mu + mu = np.trace(emp_cov, axis1=-2, axis2=-1) / n_features + mu = np.expand_dims(mu, axis=tuple(range(mu.ndim, emp_cov.ndim))) + shrunk_cov += shrinkage * mu * np.eye(n_features) return shrunk_cov diff --git a/sklearn/covariance/tests/test_covariance.py b/sklearn/covariance/tests/test_covariance.py index 0866c209a10c3..ef4bd63149d60 100644 --- a/sklearn/covariance/tests/test_covariance.py +++ b/sklearn/covariance/tests/test_covariance.py @@ -81,7 +81,25 @@ def test_covariance(): assert_array_equal(cov.location_, np.zeros(X.shape[1])) +@pytest.mark.parametrize("n_matrices", [1, 3]) +def test_shrunk_covariance_func(n_matrices): + """Check `shrunk_covariance` function.""" + + n_features = 2 + cov = np.ones((n_features, n_features)) + cov_target = np.array([[1, 0.5], [0.5, 1]]) + + if n_matrices > 1: + cov = np.repeat(cov[np.newaxis, ...], n_matrices, axis=0) + cov_target = np.repeat(cov_target[np.newaxis, ...], n_matrices, axis=0) + + cov_shrunk = shrunk_covariance(cov, 0.5) + assert_allclose(cov_shrunk, cov_target) + + def test_shrunk_covariance(): + """Check consistency between `ShrunkCovariance` and `shrunk_covariance`.""" + # Tests ShrunkCovariance module on a simple dataset. # compare shrunk covariance obtained from data and from MLE estimate cov = ShrunkCovariance(shrinkage=0.5)