Implementation of Block Covariance estimation

doc/whatsnew.rst

@@ -36,6 +36,8 @@ v0.2.8.dev
 
 - Add masked and NaN means with Riemannian metric: :func:`pyriemann.utils.mean.maskedmean_riemann` and :func:`pyriemann.utils.mean.nanmean_riemann`
 
+- Add block covariance matrix: :func:`pyriemann.utils.covariance.block_covariances`
+
 - Add ``corr`` option in :func:`pyriemann.utils.covariance.normalize`, to normalize covariance into correlation matrices
 
 v0.2.7 (June 2021)

examples/SSVEP/plot_classify_ssvep_mdm.py

@@ -23,7 +23,7 @@
 from mne.datasets import fetch_dataset
 
 # pyriemann import
-from pyriemann.estimation import Covariances
+from pyriemann.estimation import BlockCovariances
 from pyriemann.utils.mean import mean_riemann
 from pyriemann.classification import MDM
 
@@ -190,7 +190,8 @@ def _bandpass_filter(signal, lowcut, highcut):
 # The covariance matrices will be estimated using the Ledoit-Wolf shrinkage
 # estimator on the extended signal.
 
-cov_ext_trials = Covariances(estimator='lwf').transform(epochs.get_data())
+cov_ext_trials = BlockCovariances(estimator='lwf',
+                                  block_size=8).transform(epochs.get_data())
 
 # This plot shows an example of a covariance matrix observed for each class:
 

pyriemann/estimation.py

@@ -3,7 +3,7 @@
 
 from .spatialfilters import Xdawn
 from .utils.covariance import (covariances, covariances_EP, cospectrum,
-                               coherence)
+                               coherence, block_covariances)
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.covariance import shrunk_covariance
 
@@ -651,3 +651,89 @@ def transform(self, X):
             covmats[ii] = shrunk_covariance(x, self.shrinkage)
 
         return covmats
+
+
+class BlockCovariances(BaseEstimator, TransformerMixin):
+    """Estimation of block covariance matrix.
+    Perform a block covariance matrix estimation for each given trial. The
+    resulting matrices are block diagonal matrices.
+
+    The blocks on the diagonal are calculated as individual covariance
+    matrices for a subset of channels using the given the estimator.
+    Varying block sized possible by passing a list to allow incorporation
+    of different modalities with different number of channels (e.g. EEG,
+    ECoG, LFP, EMG) with their own respective covariance matrices.
+
+    Parameters
+    ----------
+    estimator : string (default: 'scm')
+        covariance matrix estimator. For regularization consider 'lwf' or 'oas'
+        For a complete list of estimator, see `utils.covariance`.
+    block_size : list or int
+        Sizes of individual blocks given as int for same-size block or list for
+        varying block sizes.
+
+    See Also
+    --------
+    ERPCovariances
+    XdawnCovariances
+    CospCovariances
+    HankelCovariances
+    """
+
+    def __init__(self, block_size, estimator='scm'):
+        """Init."""
+        self.estimator = estimator
+        self.block_size = block_size
+
+    def fit(self, X, y=None):
+        """Fit.
+
+        Do nothing. For compatibility purpose.
+
+        Parameters
+        ----------
+        X : ndarray, shape (n_matrices, n_channels, n_samples)
+            ndarray of matrices.
+        y : ndarray shape (n_matrices,)
+            labels corresponding to each trial, not used.
+
+        Returns
+        -------
+        self : Covariances instance
+            The Covariances instance.
+        """
+        return self
+
+    def transform(self, X):
+        """Estimate block covariance matrices.
+
+        Parameters
+        ----------
+        X : ndarray, shape (n_matrices, n_channels, n_samples)
+            ndarray of matrices.
+
+        Returns
+        -------
+        covmats : ndarray, shape (n_matrices, n_channels, n_channels)
+            ndarray of covariance matrices for each trials.
+        """
+        n_matrices, n_channels, n_times = X.shape
+
+        if isinstance(self.block_size, int):
+            n_blocks = n_channels // self.block_size
+
+            if n_blocks*self.block_size != n_channels:
+                raise ValueError('block_size must be divisor '
+                                 'of number of channels of X.')
+
+            blocks = [self.block_size for b in range(n_blocks)]
+
+        else:
+            blocks = self.block_size
+
+        if np.sum(blocks) != n_channels:
+            raise ValueError('Sum of individual block sizes '
+                             'must match number of channels of X.')
+
+        return block_covariances(X, blocks, self.estimator)

pyriemann/utils/covariance.py

@@ -1,6 +1,7 @@
 import warnings
 
 import numpy as np
+from scipy.linalg import block_diag
 
 from sklearn.covariance import oas, ledoit_wolf, fast_mcd, empirical_covariance
 from .test import is_square
@@ -444,6 +445,49 @@ def coherence(X, window=128, overlap=0.75, fmin=None, fmax=None, fs=None,
     return C, freqs
 
 
+def block_covariances(X, blocks, estimator='cov'):
+    """Compute block diagonal covariance.
+
+    Calculates block diagonal matrices where each block is a covariance
+    matrix of a subset of channels.
+    Block sizes are passed as a list of integers and can vary. The sum
+    of block sizes must equal the number of channels in X.
+
+    Parameters
+    ----------
+    X : ndarray, shape (n_matrices, n_channels, n_times)
+        Multi-channel time-series.
+    blocks: list of int
+        list of block sizes.
+    estimator : {'cov', 'scm', 'lwf', 'oas', 'mcd', 'sch', 'corr'} \
+        (default: 'scm')
+        Covariance matrix estimator, see
+            :func:`pyriemann.utils.covariance.covariances`.
+
+    Returns
+    -------
+    C : ndarray, shape (n_matrices, n_channels, n_channels)
+        Block diagonal covariance matrices.
+
+    Notes
+    -----
+    .. versionadded:: 0.2.8
+    """
+    est = _check_est(estimator)
+    n_matrices, n_channels, n_times = X.shape
+
+    covmats = np.empty((n_matrices, n_channels, n_channels))
+    for i in range(n_matrices):
+        blockcov = []
+        idx_start = 0
+        for j in blocks:
+            blockcov.append(est(X[i, idx_start:idx_start+j, :]))
+            idx_start += j
+        covmats[i, :, :] = block_diag(*tuple(blockcov))
+
+    return covmats
+
+
 ###############################################################################
 
 

tests/test_estimation.py

@@ -6,6 +6,7 @@
     HankelCovariances,
     Coherences,
     Shrinkage,
+    BlockCovariances
 )
 import pytest
 
@@ -206,3 +207,47 @@ def test_shrinkage(shrinkage, rndstate):
     assert sh.get_params() == dict(shrinkage=shrinkage)
     assert covmats.shape == (n_trials, n_channels, n_channels)
     assert is_spd(covmats)
+
+
+@pytest.mark.parametrize("estimator", estim)
+def test_block_covariances_est(estimator, rndstate):
+    """Test Covariances"""
+    n_trials, n_channels, n_times = 2, 12, 100
+    x = rndstate.randn(n_trials, n_channels, n_times)
+    cov = BlockCovariances(block_size=6, estimator=estimator)
+    cov.fit(x)
+    covmats = cov.fit_transform(x)
+    assert cov.get_params() == dict(block_size=6, estimator=estimator)
+    assert covmats.shape == (n_trials, n_channels, n_channels)
+    assert is_spd(covmats)
+
+
+@pytest.mark.parametrize("block_size", [1, 6, [4, 8]])
+def test_block_covariances_blocks(block_size, rndstate):
+    """Test Covariances"""
+    n_trials, n_channels, n_times = 2, 12, 100
+    x = rndstate.randn(n_trials, n_channels, n_times)
+    cov = BlockCovariances(block_size=block_size)
+    cov.fit(x)
+    covmats = cov.fit_transform(x)
+    assert cov.get_params() == dict(block_size=block_size, estimator='scm')
+    assert covmats.shape == (n_trials, n_channels, n_channels)
+    assert is_spd(covmats)
+
+
+def test_block_covariances_int_value_error(rndstate):
+    """Test Covariances"""
+    n_trials, n_channels, n_times = 2, 12, 100
+    x = rndstate.randn(n_trials, n_channels, n_times)
+    cov = BlockCovariances(block_size=5)
+    with pytest.raises(ValueError):
+        cov.fit_transform(x)
+
+
+def test_block_covariances_array_value_error(rndstate):
+    """Test Covariances"""
+    n_trials, n_channels, n_times = 2, 12, 100
+    x = rndstate.randn(n_trials, n_channels, n_times)
+    cov = BlockCovariances(block_size=[4, 4, 5])
+    with pytest.raises(ValueError):
+        cov.fit_transform(x)

tests/test_utils_covariance.py

@@ -1,12 +1,13 @@
 from numpy.testing import assert_array_almost_equal
 import numpy as np
+from scipy.linalg import block_diag
 from scipy.signal import welch, csd, coherence as coherence_sp
 import pytest
 
 from pyriemann.utils.covariance import (
     covariances, covariances_EP, covariances_X, eegtocov,
     cross_spectrum, cospectrum, coherence,
-    normalize, get_nondiag_weight
+    normalize, get_nondiag_weight, block_covariances
 )
 from pyriemann.utils.test import is_real, is_hermitian
 
@@ -275,6 +276,47 @@ def test_covariances_coherence(coh, rndstate):
             assert c[0, 3, foi] == pytest.approx(0.0)
 
 
+@pytest.mark.parametrize(
+    'estimator', ['oas', 'lwf', 'scm', 'corr', 'mcd',
+                  'sch', np.cov, 'truc', None]
+)
+def test_block_covariances_est(estimator, rndstate):
+    """Test covariance for multiple estimators"""
+    n_matrices, n_channels, n_times = 2, 12, 100
+    x = rndstate.randn(n_matrices, n_channels, n_times)
+
+    if estimator is None:
+        cov = block_covariances(x, [4, 4, 4])
+        assert cov.shape == (n_matrices, n_channels, n_channels)
+    elif estimator == 'truc':
+        with pytest.raises(ValueError):
+            block_covariances(x, [4, 4, 4], estimator=estimator)
+    else:
+        cov = block_covariances(x, [4, 4, 4], estimator=estimator)
+        assert cov.shape == (n_matrices, n_channels, n_channels)
+
+
+def test_block_covariances(rndstate):
+    """Test block covariance"""
+    n_matrices, n_channels, n_times = 2, 12, 100
+    x = rndstate.randn(n_matrices, n_channels, n_times)
+
+    cov = block_covariances(x, [12], estimator='cov')
+    assert_array_almost_equal(cov, covariances(x, estimator='cov'))
+
+    cov = block_covariances(x, [6, 6], estimator='cov')
+    cov2 = covariances(x, estimator='cov')
+    covcomp = block_diag(*(cov2[0, :6, :6], cov2[0, 6:12, 6:12]))
+    assert_array_almost_equal(cov[0], covcomp)
+
+    cov = block_covariances(x, [3, 5, 4], estimator='cov')
+    cov2 = covariances(x, estimator='cov')
+    covcomp = block_diag(*(cov2[0, :3, :3],
+                           cov2[0, 3:8, 3:8],
+                           cov2[0, 8:12, 8:12]))
+    assert_array_almost_equal(cov[0], covcomp)
+
+
 @pytest.mark.parametrize('norm', ['corr', 'trace', 'determinant'])
 def test_normalize_shapes(norm, rndstate):
     """Test normalize shapes"""