ConfoundRegressor with GeneralizingEstimator

Hello there!

I would like to check whether, once the influence of some confound variables is ruled out, my decoding still works (EEG data).
To do so, I thought to try the cross-validated "confound regression" described in this paper.

I've made a simple and naive adaptation of the code provided with this paper to use it inside MNE sliders:

# Giulia Gennari November 2020
"""
Adaptation of the confounds module retrievable at https://github.com/lukassnoek/MVCA/blob/master/analyses/confounds.py
which contains code to handle/account for confounds in pattern analyses.
The main change relative to the original code consists in providing separate train
and test confound vectors at inizialization while avoiding to specify X. This should enable the
employment of ConfoundRegressor inside MNE sliders 
""" 

import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin


class ConfoundRegressorG(BaseEstimator, TransformerMixin):
    """ Fits a confound onto each feature in X and returns their residuals."""

    def __init__(self, confound_train, confound_test, skip_cross_val=False, stack_intercept=True):
        """ Regresses out a variable (confound) from each feature in X.
        Parameters
        ----------
        confound_train, confound_test: numpy arrays
            Array of length (n_samples, n_confounds) to regress out of each
            feature. May have multiple columns for multiple confounds.
        skip_cross_val : bool
            If True do not transform the test set 
        stack_intercept : bool
            Whether to stack an intercept to the confound (default is True in original module)

        Attributes
        ----------
        weights_ : numpy array
            Array with weights for the confound(s).
        """

        self.confound_train = confound_train
        self.confound_test = confound_test
        self.skip_cross_val = skip_cross_val
        self.stack_intercept = stack_intercept
        self.weights_ = None

    def fit(self, X, y=None):
        """ Fits the confound-regressor to X.
        Parameters
        ----------
        X : numpy array
            An array of shape (n_samples, n_features)
        y : None
            Included for compatibility; does nothing.
        """

        if self.stack_intercept:
            icept = np.ones(self.confound_train.shape[0])
            self.confound_train = np.c_[icept, self.confound_train]
            
            icept_bis = np.ones(self.confound_test.shape[0])
            self.confound_test = np.c_[icept_bis, self.confound_test]

        confound_fit = self.confound_train

        # Vectorized implementation estimating weights for all features
        self.weights_ = np.linalg.lstsq(confound_fit, X, rcond=None)[0]
        return self

    def transform(self, X):
        """ Regresses out confound from X.
        Parameters
        ----------
        X : numpy array
            An array of shape (n_samples, n_features)
        Returns
        -------
        X_new : ndarray
            ndarray with confound-regressed features
        """
        
        # no transformation when skip_cross_val is True and tranform is called upon 
        # test data (otherwise select the proper confounds and proceed)
        
        if self.confound_test.shape[-1] != self.confound_train.shape[-1]:
            icept_bis = np.ones(self.confound_test.shape[0])
            self.confound_test = np.c_[icept_bis, self.confound_test]
        
        if self.skip_cross_val and len(X)== len(self.confound_test):
            X_corr = X
        else:
            if len(X)== len(self.confound_train):
                confound_transform = self.confound_train
            if len(X)== len(self.confound_test):
                confound_transform = self.confound_test

            X_corr = X - confound_transform.dot(self.weights_)
            
        return X_corr

This works perfectly with SlidingEstimator. However, with GeneralizingEstimatos I found a behavior that needs to be explained to me (please!). Namely: when I call .score everything seems to work, BUT when I call either .predict/ .predict_proba or .decision_function it doesn't. Turns out that X_test.shape in the latter cases becomes (n_epochs*n_time_points, n_features) at some point.
My one-million question at the moment: how come no error comes up with .score ??
I thought that e.g. .predict_proba is called internally with .score and thus I would expect an error in this case as well.

By saving the lines above as confound_regressor_customized.py you can reproduce everything with the following:

import os
import numpy as np
import pandas as pd
import mne
from sklearn.model_selection import train_test_split
from confound_regressor_customized import ConfoundRegressorG
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import make_pipeline
from mne.decoding import GeneralizingEstimator

kiloword_data_folder = mne.datasets.kiloword.data_path()
kiloword_data_file = os.path.join(kiloword_data_folder,'kword_metadata-epo.fif')
epochs = mne.read_epochs(kiloword_data_file)
epochs.crop(0.,.076)

train_indx, test_indx = train_test_split(range(len(epochs)))
epochs_train, epochs_test = epochs[train_indx], epochs[test_indx]

y_train = epochs_train.metadata['NumberOfLetters'].to_numpy(dtype=int)
y_test = epochs_test.metadata['NumberOfLetters'].to_numpy(dtype=int)

c_train = epochs_train.metadata[['ConsonantVowelProportion', 'VisualComplexity']].to_numpy()
c_test = epochs_test.metadata[['ConsonantVowelProportion', 'VisualComplexity']].to_numpy()

cfr = ConfoundRegressorG(confound_train=c_train, confound_test=c_test)
pipeline = make_pipeline(cfr, StandardScaler(), LogisticRegression(max_iter=1000))
pip_gen = GeneralizingEstimator(pipeline, scoring='roc_auc_ovr_weighted')

pip_gen.fit(epochs_train.get_data(), y=y_train)
scores_gen =  pip_gen.score(epochs_test.get_data(), y=y_test)

# the error arrives here! (same for .predict or .decision_function) 
probas_gen =  pip_gen.predict_proba(epochs_test.get_data())

Thank you very much for looking at this issue!!!!

Giulia

Reproduce the behavior with confound_regressor_customized.py and the provided GeneralizingEstimator pipeline. Read the GeneralizingEstimator score, predict_proba, predict, and decision_function entry points, then trace the input shapes passed to the pipeline. Done means explaining why score differs from the other methods and documenting or correcting the inconsistent behavior.