bci_essentials.classification.erp_rg_classifier_hyperparamgridsearch

  1"""**ERP RG Classifier**
  2
  3This classifier is used to classify ERPs using the Riemannian Geometry
  4approach.
  5
  6"""
  7
  8# Stock libraries
  9import numpy as np
 10import matplotlib.pyplot as plt
 11from sklearn.model_selection import StratifiedKFold, GridSearchCV
 12from sklearn.metrics import (
 13    confusion_matrix,
 14    ConfusionMatrixDisplay,
 15    precision_score,
 16    recall_score,
 17    roc_auc_score,
 18    make_scorer,
 19)
 20from sklearn.pipeline import Pipeline
 21from pyriemann.tangentspace import TangentSpace
 22from pyriemann.estimation import XdawnCovariances
 23from pyriemann.channelselection import FlatChannelRemover
 24from sklearn.discriminant_analysis import LinearDiscriminantAnalysis
 25
 26# Import bci_essentials modules and methods
 27from .generic_classifier import (
 28    GenericClassifier,
 29    Prediction,
 30)
 31from ..signal_processing import lico, random_oversampling, random_undersampling
 32from ..utils.logger import Logger  # Logger wrapper
 33
 34# Instantiate a logger for the module at the default level of logging.INFO
 35# Logs to bci_essentials.__module__) where __module__ is the name of the module
 36logger = Logger(name=__name__)
 37
 38
 39class ErpRgClassifierHyperparamGridSearch(GenericClassifier):
 40    """ERP RG Classifier with hyperparameter grid search
 41    class (*inherits from `GenericClassifier`*)."""
 42
 43    def set_p300_clf_settings(
 44        self,
 45        n_splits=3,
 46        resampling_method=None,
 47        lico_expansion_factor=1,
 48        oversample_ratio=0,
 49        undersample_ratio=0,
 50        random_seed=42,
 51        remove_flats=True,
 52    ):
 53        """Set P300 Classifier Settings.
 54
 55        Parameters
 56        ----------
 57        n_splits : int, *optional*
 58            Number of folds for cross-validation.
 59            - Default is `3`.
 60        resampling_method : str, *optional*, None
 61            Resampling method to use ["lico", "oversample", "undersample"].
 62            Default is None.
 63        lico_expansion_factor : int, *optional*
 64            Linear Combination Oversampling expansion factor, which is the
 65            factor by which the number of ERPs in the training set will be
 66            expanded.
 67            - Default is `1`.
 68        oversample_ratio : float, *optional*
 69            Traditional oversampling. Range is from from 0.1-1 resulting
 70            from the ratio of erp to non-erp class. 0 for no oversampling.
 71            - Default is `0`.
 72        undersample_ratio : float, *optional*
 73            Traditional undersampling. Range is from from 0.1-1 resulting
 74            from the ratio of erp to non-erp class. 0 for no undersampling.
 75            - Default is `0`.
 76        random_seed : int, *optional*
 77            Random seed.
 78            - Default is `42`.
 79        remove_flats : bool, *optional*
 80            Whether to remove flat channels.
 81            - Default is `True`.
 82
 83        Returns
 84        -------
 85        `None`
 86
 87        """
 88        self.n_splits = n_splits
 89        self.resampling_method = resampling_method
 90        self.lico_expansion_factor = lico_expansion_factor
 91        self.oversample_ratio = oversample_ratio
 92        self.undersample_ratio = undersample_ratio
 93        self.random_seed = random_seed
 94
 95        # # Create steps list with proper formatting
 96        steps = []
 97        if remove_flats:
 98            steps.append(("remove_flats", FlatChannelRemover()))
 99
100        steps.extend(
101            [
102                ("xdawn", XdawnCovariances()),
103                ("tangent", TangentSpace()),
104                ("lda", LinearDiscriminantAnalysis()),
105            ]
106        )
107
108        # Create pipeline
109        self.clf = Pipeline(steps)
110
111        # Hyperparameters to be optimized
112        # TODO: Implement an extended nfilter set, dynamically based on the number of channels
113        # Example of dynamic nfilter set
114        # n_channels = self.X.shape[1]
115        # nfilter_set = list(range(2, n_channels+1))  # Example range from 2 to n_channels inclusive
116        # Then set "xdawn__nfilter": nfilter_set in the param_grid below
117        self.param_grid = {
118            "xdawn__nfilter": [2, 3, 4],
119            "xdawn__estimator": ["oas", "lwf"],
120            "tangent__metric": ["riemann"],
121            "lda__solver": ["lsqr", "eigen"],
122            "lda__shrinkage": np.linspace(0.5, 0.9, 5),
123        }
124
125    def fit(
126        self,
127        plot_cm=False,
128        plot_roc=False,
129    ):
130        """Fit the model.
131
132        Parameters
133        ----------
134        plot_cm : bool, *optional*
135            Whether to plot the confusion matrix during training.
136            - Default is `False`.
137        plot_roc : bool, *optional*
138            Whether to plot the ROC curve during training.
139            - Default is `False`.
140
141        Returns
142        -------
143        `None`
144            Models created used in `predict()`.
145
146        """
147        logger.info("Fitting the model using RG")
148        logger.info("X shape: %s", self.X.shape)
149        logger.info("y shape: %s", self.y.shape)
150
151        # Resample data if needed
152        self.X, self.y = self.__resample_data()
153
154        # Optimize hyperparameters with cross-validation
155        self.__optimize_hyperparameters()
156
157        # Fit the model with the complete dataset and optimized hyperparameters
158        self.clf.fit(self.X, self.y)
159
160        # Get predictions for final model
161        y_pred_proba = self.clf.predict_proba(self.X)[:, 1]
162
163        # Calculate estimate of training metrics of final model
164        # TODO: Implement proper training metrics calculation, using cross validation.
165        # self.offline_accuracy = sum(y_pred == self.y) / len(self.y)
166        # self.offline_precision = precision_score(self.y, y_pred)
167        # self.offline_recall = recall_score(self.y, y_pred)
168
169        try:
170            roc_auc = roc_auc_score(self.y, y_pred_proba)
171            logger.info(f"ROC AUC Score: {roc_auc:0.3f}")
172        except Exception as e:
173            logger.warning(f"Could not calculate ROC AUC score: {e}")
174
175        # Display training confusion matrix
176        # self.offline_cm = confusion_matrix(self.y, y_pred)
177        if plot_cm:
178            disp = ConfusionMatrixDisplay(confusion_matrix=self.offline_cm)
179            disp.plot()
180            plt.title("Training confusion matrix")
181
182        if plot_roc:
183            # TODO Implementation missing
184            pass
185
186        # Log training metrics
187        logger.info("Final model training performance metrics:")
188        logger.info(f"Accuracy: {self.offline_accuracy:0.3f} - MAY NOT BE ACCURATE")
189        logger.info(f"Precision: {self.offline_precision:0.3f} - MAY NOT BE ACCURATE")
190        logger.info(f"Recall: {self.offline_recall:0.3f} - MAY NOT BE ACCURATE")
191        logger.info(f"Confusion Matrix:\n{self.offline_cm} ")
192        logger.warning(
193            "Note: Training metrics may not be accurate due to the use of "
194            "cross-validation and resampling methods. Use with caution."
195        )
196
197    def predict(self, X):
198        """Predict the class of the data
199
200        Parameters
201        ----------
202        X : numpy.ndarray
203            3D array where shape = (n_epochs, n_channels, n_samples)
204
205        Returns
206        -------
207        prediction : Prediction
208            Predict object. Contains the predicted labels and and the probability.
209            Because this classifier chooses the P300 object with the highest posterior probability,
210            the probability is only the posterior probability of the chosen object.
211
212        """
213
214        subset_X = self.get_subset(X, self.subset, self.channel_labels)
215
216        # Get posterior probability for each target
217        posterior_prob = self.clf.predict_proba(subset_X)[:, 1]
218
219        label = [int(np.argmax(posterior_prob))]
220        probability = [np.max(posterior_prob)]
221
222        return Prediction(label, probability)
223
224    # TODO implement additional resampling methods, JIRA ticket: B4K-342
225    def __resample_data(self):
226        """Resample data based on the selected method"""
227
228        X_resampled = self.X.copy()
229        y_resampled = self.y.copy()
230
231        try:
232            if (self.resampling_method == "lico") and (self.lico_expansion_factor > 1):
233                [X_resampled, y_resampled] = lico(
234                    self.X, self.y, self.lico_expansion_factor
235                )
236                pass
237
238            elif (self.resampling_method == "oversample") and (
239                self.oversample_ratio > 0
240            ):
241                [X_resampled, y_resampled] = random_oversampling(
242                    self.X, self.y, self.oversample_ratio
243                )
244                pass
245
246            elif (self.resampling_method == "undersample") and (
247                self.undersample_ratio > 0
248            ):
249                [X_resampled, y_resampled] = random_undersampling(
250                    self.X, self.y, self.undersample_ratio
251                )
252                pass
253
254            logger.info(f"Resampling  with {self.resampling_method} done")
255            logger.info(f"X_resampled shape: {X_resampled.shape}")
256            logger.info(f"y_resampled shape: {y_resampled.shape}")
257
258        except Exception as e:
259            logger.error(
260                f"{self.resampling_method.capitalize()} resampling method failed"
261            )
262            logger.error(e)
263
264        return X_resampled, y_resampled
265
266    def __optimize_hyperparameters(self):
267        """Optimize hyperparameters with cross-validation using brute force grid search
268
269        Returns
270        -------
271        `None`
272            Model with best hyperparameters to be used in `predict()`.
273
274        """
275
276        # Perform cross-validation
277        cv = StratifiedKFold(
278            n_splits=self.n_splits, shuffle=True, random_state=self.random_seed
279        )
280
281        # Create custom scorer function
282        custom_scorer = make_scorer(
283            self._valid_roc_auc, response_method="predict_proba", greater_is_better=True
284        )
285
286        # Create GridSearchCV object
287        grid_search = GridSearchCV(
288            estimator=self.clf,
289            param_grid=self.param_grid,
290            cv=cv,
291            n_jobs=-1,
292            verbose=1,
293            scoring=custom_scorer,
294            refit=True,
295            return_train_score=True,
296        )
297
298        # Start grid search optimization
299        logger.info("Starting grid search optimization...")
300        grid_search.fit(self.X, self.y)
301
302        # Get best parameters and score
303        logger.info("Grid search optimization completed.")
304        best_params = grid_search.best_params_
305        best_score = grid_search.best_score_
306
307        # Report training metrics: TODO: Verify this is the right way to calculate training metrics
308        self.offline_accuracy = grid_search.best_estimator_.score(self.X, self.y)
309        self.offline_cm = confusion_matrix(
310            self.y, grid_search.best_estimator_.predict(self.X)
311        )
312        self.offline_precision = precision_score(
313            self.y, grid_search.best_estimator_.predict(self.X)
314        )
315        self.offline_recall = recall_score(
316            self.y, grid_search.best_estimator_.predict(self.X)
317        )
318
319        # Update classifier with best parameters
320        self.clf.set_params(**best_params)
321        logger.info(f"Best parameters found: {best_params}")
322        logger.info(f"Best CV score: {best_score:0.3f}")
323
324    def _valid_roc_auc(self, y_true, y_pred, **kwargs):
325        """Calculate the ROC AUC score for the classifier.
326        This method is used because the stock `roc_auc_score` function
327        does not handle the case where one class is missing in the fold.
328        This method will return 0.5 in that case.
329
330        Parameters
331        ----------
332        y_true : numpy.ndarray
333            True labels.
334        y_pred : numpy.ndarray
335            Predicted labels.
336        **kwargs : dict
337            Additional keyword arguments passed by make_scorer.
338
339        Returns
340        -------
341        roc_auc : float
342            ROC AUC score.
343
344        """
345        try:
346            # Check if we have both classes in the fold
347            if len(np.unique(y_true)) < 2:
348                logger.warning("Fold contains only one class")
349                return 0.5
350
351            return roc_auc_score(y_true, y_pred)
352
353        except Exception as e:
354            logger.warning(f"ROC AUC calculation failed: {e}")
355            return 0.5