ENH allow shrunk_covariance to handle multiple matrices at once

[qbarthelemy]

Reference Issues/PRs

None.

What does this implement/fix? Explain your changes.

Current version of shrunk_covariance processes only one covariance matrix.

This PR aims to handle ndarrays with n>=2.

Any other comments?

Example with a 3D array.

import timeit
import numpy as np
from sklearn.utils import check_array
from sklearn.covariance import empirical_covariance, shrunk_covariance


def shrunk_covariance_old(emp_covs, shrinkage=0.1):
    return np.array([shrunk_covariance(e, shrinkage) for e in emp_covs])

def shrunk_covariance_new(emp_cov, shrinkage=0.1):
    emp_cov = check_array(emp_cov, allow_nd=True)
    n_features = emp_cov.shape[-1]

    shrunk_cov = (1.0 - shrinkage) * emp_cov
    mu = np.trace(emp_cov, axis1=-2, axis2=-1) / n_features
    while mu.ndim != emp_cov.ndim:
        mu = mu[..., np.newaxis]
    shrunk_cov += shrinkage * mu * np.eye(n_features)
    return shrunk_cov


def compute_time_3D(func, n_reps=10, n_matrices=200, n_samples=1000, n_features=50):
    rs = np.random.RandomState(42)
    times = np.zeros((n_reps))
    for r in range(n_reps):
        X = rs.randn(n_matrices, n_samples, n_features)
        C = np.array([empirical_covariance(x) for x in X])
        t0 = timeit.default_timer()
        func(C)
        times[r] = timeit.default_timer() - t0
    print('Comput time = {0:.4f} +/- {1:.4f}'.format(times.mean(), times.std()))

compute_time_3D(shrunk_covariance_old)
compute_time_3D(shrunk_covariance_new)

Comput time = 0.0105 +/- 0.0004
Comput time = 0.0032 +/- 0.0001

[qbarthelemy]

Thx @agramfort if you can help the process!

[agramfort]

@qbarthelemy can you add a changelog entry?

here is a just for a bug

can you see why CIs are red?

thx @qbarthelemy

[qbarthelemy]

Thx for your review @agramfort ! Everything seems OK.

[agramfort]

Can you see why CIs fail?

[qbarthelemy]

Errors do not seem to be related to my modifications.

/bin/bash: ./build_tools/circle/build_test_arm.sh: No such file or directory

Failed to fetch http://azure.archive.ubuntu.com/ubuntu/pool/universe/m/matplotlib/python-matplotlib-data_3.1.2-1ubuntu4_all.deb Connection failed [IP: 40.119.46.219 80]

/home/vsts/work/1/s/build_tools/azure/test_docs.sh: line 8: testvenv/bin/activate: No such file or directory

[glemaitre]

Instead of adding this validation, we should instead use the parameter validation framework: #24862

Before to be able to do that, we will need a small refactoring to have the function calling the class as we did for other estimator/function. I will fix that in another PR and we can build upon it here.

[glemaitre]

Suggested change

	emp_cov : array-like of shape (..., n_features, n_features)
	emp_cov : array-like of shape (n_features, n_features) or \
	(n_matrices, n_features, n_features)

[qbarthelemy]

Suggestions about docstring seem to restrict usage to 2D and 3D arrays, whereas function is now able to process any nd array.

[glemaitre]

What is the use case of n-D shrinkage? Since the covariance matrix is always a 2-D matrix, I don't really see when you will pass a 4-D array, for instance.

[qbarthelemy]

Most of NumPy and SciPy functions use (..., n, n) to indicate that they can process ndarrays,
like numpy.linalg.eig and scipy.linalg.expm for example.

Use cases belong to users. But, for an example with a 4D array, shape (k, m, n, n), one might want to shrunk k sets, each set containing of m covariance matrices. I have tested, and code is ok.

For me, description should not restrain actual usage. But, as you wish.

[glemaitre]

I was thinking more of restraining on purpose the usage that we have in scikit-learn.
But we can go this road and see what other reviewers think.

[glemaitre]

Here, we materialize the eye matrix. If we restrain to a 3-D matrix then, we can efficiently do the in place addition with flattening as previously done.

[glemaitre]

Then we need to mention somehow that we handle nd-array.

[glemaitre]

I was thinking more of restraining on purpose the usage that we have in scikit-learn.
But we can go this road and see what other reviewers think.

[qbarthelemy]

Tests fail, due to np.expand_dims.
It seems that NumPy version used in tests does not accept tuple for axis argument (tuple added in NumPy 1.18.0).

  File "/home/vsts/work/1/s/sklearn/covariance/_shrunk_covariance.py", line 114, in shrunk_covariance
    mu = np.expand_dims(mu, axis=tuple(range(mu.ndim, emp_cov.ndim)))
  File "<__array_function__ internals>", line 5, in expand_dims
  File "/usr/lib/python3/dist-packages/numpy/lib/shape_base.py", line 577, in expand_dims
    if axis > a.ndim or axis < -a.ndim - 1:
TypeError: '>' not supported between instances of 'tuple' and 'int'

[glemaitre]

Indeed, it was introduced in NumPy 1.18 and we have minimal dependencies set to 1.17.3.
But I think that we should be fine here. NEP29 tells us that we will drop Python 3.8 support and the oldest-support-numpy (https://github.com/scipy/oldest-supported-numpy/blob/main/setup.cfg) tell us that we should bump to 1.19.3.

So for scikit-learn 1.3, we will be able to get with the current implementation. We might need to delay the merge of the feature depending on the bugfix release, but I will put this feature in the milestone.

[jeremiedbb]

Actually we're not bumping the numpy min version for 1.3. Let's target this for 1.4 and we'll merge this when we bump the min version.

[github-actions]

✔️ Linting Passed

All linting checks passed. Your pull request is in excellent shape! ☀️

_{Generated for commit: 4d88273. Link to the linter CI: here}

[glemaitre]

I open #27910 to bump the version of NumPy. Once merged, this PR should be mergeable.

[glemaitre]

Merging since everything is green. Thanks @qbarthelemy