Add regularization of convex mean

pyriemann_qiskit/utils/mean.py

@@ -2,10 +2,12 @@
 from pyriemann.utils.mean import mean_methods
 from pyriemann_qiskit.utils.docplex import (ClassicalOptimizer,
                                             get_global_optimizer)
+from pyriemann.estimation import Shrinkage
 
 
 def fro_mean_convex(covmats, sample_weight=None,
-                    optimizer=ClassicalOptimizer()):
+                    optimizer=ClassicalOptimizer(),
+                    shrink=True):
     """Convex formulation of the mean
     with frobenius distance.
     Parameters
@@ -16,7 +18,10 @@ def fro_mean_convex(covmats, sample_weight=None,
         Weights for each matrix. Never used in practice.
         It is kept only for standardization with pyRiemann.
     optimizer: pyQiskitOptimizer
-      An instance of pyQiskitOptimizer.
+        An instance of pyQiskitOptimizer.
+    shrink: boolean (default: true)
+        If True, it applies shrinkage regularization [2]_
+        of the resulting covariance matrix.
 
     Returns
     -------
@@ -26,11 +31,15 @@ def fro_mean_convex(covmats, sample_weight=None,
     Notes
     -----
     .. versionadded:: 0.0.3
+    .. versionchanged:: 0.0.4
+        Add regularization of the results.
 
     References
     ----------
     .. [1] \
         http://ibmdecisionoptimization.github.io/docplex-doc/mp/_modules/docplex/mp/model.html#Model
+    .. [2] \
+        https://pyriemann.readthedocs.io/en/v0.4/generated/pyriemann.estimation.Shrinkage.html
     """
 
     optimizer = get_global_optimizer(optimizer)
@@ -55,6 +64,8 @@ def _fro_dist(A, B):
 
     result = optimizer.solve(prob)
 
+    if shrink:
+        return Shrinkage().transform([result])[0]
     return result
 
 

tests/test_utils_mean.py

@@ -30,7 +30,7 @@ def test_mean_convex_vs_euclid(get_covmats):
     """Test that euclidian and convex mean returns close results"""
     n_trials, n_channels = 5, 3
     covmats = get_covmats(n_trials, n_channels)
-    C = fro_mean_convex(covmats)
+    C = fro_mean_convex(covmats, shrink=False)
     C_euclid = mean_euclid(covmats)
     assert np.allclose(C, C_euclid, atol=0.001)
 
@@ -51,7 +51,7 @@ def test_mean_convex_all_zeros(optimizer):
     is a matrix filled with zeros"""
     n_trials, n_channels = 5, 2
     covmats = np.zeros((n_trials, n_channels, n_channels))
-    C = fro_mean_convex(covmats, optimizer=optimizer)
+    C = fro_mean_convex(covmats, optimizer=optimizer, shrink=False)
     assert np.allclose(covmats[0], C, atol=0.001)
 
 
@@ -60,7 +60,7 @@ def test_mean_convex_all_ones():
     is a matrix filled with ones"""
     n_trials, n_channels = 5, 2
     covmats = np.ones((n_trials, n_channels, n_channels))
-    C = fro_mean_convex(covmats)
+    C = fro_mean_convex(covmats, shrink=False)
     assert np.allclose(covmats[0], C, atol=0.001)
 
 
@@ -69,7 +69,7 @@ def test_mean_convex_all_equals():
     is a matrix identical to the input"""
     n_trials, n_channels, value = 5, 2, 2.5
     covmats = np.full((n_trials, n_channels, n_channels), value)
-    C = fro_mean_convex(covmats)
+    C = fro_mean_convex(covmats, shrink=False)
     assert np.allclose(covmats[0], C, atol=0.001)
 
 
@@ -80,5 +80,6 @@ def test_mean_convex_mixed():
     covmats_0 = np.zeros((n_trials, n_channels, n_channels))
     covmats_1 = np.ones((n_trials, n_channels, n_channels))
     expected_mean = np.full((n_channels, n_channels), 0.5)
-    C = fro_mean_convex(np.concatenate((covmats_0, covmats_1), axis=0))
+    C = fro_mean_convex(np.concatenate((covmats_0, covmats_1), axis=0),
+                        shrink=False)
     assert np.allclose(expected_mean, C, atol=0.001)