SHAP Support + Various Bug Fixes (#45)

* fix `ValueError: y_true and y_pred contain different number of classes`

* use number instead of str

* Set default LogisticRegression solver to saga to make it working for RFE

* Add a script to inspect and plot output db

* Add `y_true_collector` and `y_pred_proba_collector` metrics

To collect true binary labels and predicted probabilities for ROC curve generation

* Move getting the best replicate into a function

* Remove empty lines

* Return `best_models_dict` by `get_best_replicate`

* Plot ROC

* Keep only the models we are interested in

* Improve figure titles

* Plot the mean ± std ROC curve for each model across replicates.

* Unify titles across figures

* include metric (e.g., 'balanced_accuracy (test)' or 'balanced_accuracy (validate)') to ROC figure fnames

* Add max_iter parameter to log_reg configuration

* Add proper argparse

* Get unique targets dynamically and iterate over them

* Initial commit; added SHAP value calculation to the list of available metrics.

* Swapped SHAP values from list to dict (bound by feature name)

* Fixed error when a model cannot natively be parsed by SHAP.

* Added new "VarianceDrop" data hook, allowing low-variance features to be dropped as part of a trial's run.

* Updated iris testing dataset + config with new encoders.

* Added Jupyter Notebooks to the git ignore, as we occasionally use them for visual validation of tests.

* Added catch for homogeneity when running SHAP tests.

* Removed "inspect_output_db", as it is too specific to Jan's analysis.

* Pinned to pre-3.0 version of Pandas until `dtype` issues can be addressed.

---------

Co-authored-by: valosekj <jan.valosek@upol.cz>
Co-authored-by: Jan Valosek <39456460+valosekj@users.noreply.github.com>

Author: SomeoneInParticular

.gitignore

@@ -7,6 +7,7 @@
 *__pycache__/
 .vscode/
 .DS_Store
+/.ipynb_checkpoints/
 
 # Environments
 .env

data/hooks/__init__.py

@@ -30,7 +30,7 @@ def _decorator(cls: Type[DataHook]):
 # TODO: Find a more elegant way to do this
 from data.hooks.feature_selection import (
     SampleNullityDrop, FeatureNullityDrop, ExplicitDrop, ExplicitKeep, PrincipalComponentAnalysis,
-    RecursiveFeatureElimination
+    RecursiveFeatureElimination, VarianceDrop
 )
 from data.hooks.imputation import SimpleImputation
 from data.hooks.encoding import OneHotEncoding

data/hooks/feature_selection.py

@@ -6,7 +6,7 @@
 import pandas as pd
 from optuna import Trial
 from sklearn.decomposition import PCA
-from sklearn.feature_selection import RFE
+from sklearn.feature_selection import RFE, VarianceThreshold
 from sklearn.linear_model import LogisticRegression
 
 from config.utils import default_as, is_float, is_list, parse_data_config_entry
@@ -135,6 +135,81 @@ def run(self, x: BaseDataManager, y: Optional[BaseDataManager] = None) -> BaseDa
         return x.drop_features(drop_idx)
 
 
+## Feature selection by homogeneity
+@registered_data_hook("drop_low_variance")
+class VarianceDrop(FittedDataHook):
+    """
+    Thin wrapper for SciKit-Learn's VarianceThreshold class, for use as a data hook within MOOP.
+
+    Runs additional checks on top of the default implementation provided by SciKit-Learn:
+        * Ensures the resulting dataset always contains at least 1 feature.
+
+    Example usage:
+    {
+      "type": "drop_low_variance",
+      "threshold": 0.1
+    }
+    """
+    def __init__(self, config, **kwargs):
+        # TODO: make this tunable
+        super().__init__(config, **kwargs)
+
+        # Get the variance threshold
+        threshold = parse_data_config_entry(
+            "threshold", config,
+            default_as(0.0, self.logger), is_float(self.logger)
+        )
+
+        # Build the wrapped VarianceThreshold object
+        self.threshold = threshold
+        self.selected_features: list[str] | None = None
+
+    @classmethod
+    def from_config(cls, config: dict, logger: Logger = Logger.root) -> Self:
+        return cls(config=config, logger=logger)
+
+    def run(self, x: BaseDataManager, y: Optional[BaseDataManager] = None) -> BaseDataManager:
+        # If x contains only one feature already, just return that feature, as RFE has a stroke otherwise
+        if x.n_features() == 1:
+            self.logger.warning("Only one feature in the dataset was found; "
+                                "dropping any further would result in a null dataset."
+                                "Original (unmodified) dataset returned instead.")
+            return x
+
+        # Fit the model to the dataset
+        vt = VarianceThreshold(threshold=self.threshold)
+        vt.fit(x.as_array(), np.ravel(y.as_array()))  # Ravel prevents some warning spam
+
+        # Select only the features with variance less than the threshold
+        self.selected_features = vt.get_feature_names_out(x.features())
+
+        # Ensure that at least one feature was kept
+        if self.selected_features.shape[0] < 1:
+            # Find the
+            highest_var = np.max(vt.variances_)
+            highest_var_feature = list(x.features())[np.argmax(vt.variances_)]
+            self.selected_features = [highest_var_feature]
+            self.logger.warning(
+                f"Low-variance filter almost dropped all features; kept highest variance "
+                f"feature ({highest_var_feature}, variance {highest_var}) alone to prevent crash!"
+            )
+
+        # Return the copy of x containing only these features
+        x_out = x.get_features(self.selected_features)
+        return x_out
+
+    def run_fitted(self, x_train: BaseDataManager, x_test: Optional[BaseDataManager],
+                   y_train: Optional[BaseDataManager] = None, y_test: Optional[BaseDataManager] = None) -> \
+            tuple[BaseDataManager, BaseDataManager]:
+        # Run the fitted analysis first
+        train_out = self.run(x_train, y_train)
+
+        # Use the same set of features to filter the x_test set
+        test_out = x_test.get_features(self.selected_features)
+
+        return train_out, test_out
+
+
 ### Principal Component Analysis ###
 @registered_data_hook("principal_component_analysis")
 class PrincipalComponentAnalysis(Tunable, FittedDataHook):
@@ -263,8 +338,7 @@ def from_config(cls, config: dict, logger: Logger = Logger.root) -> Self:
     def tune(self, trial: Trial):
         self.prop_tuner.tune(trial)
         # Generate the new backing model based on this setup
-        # TODO: Generalize this to work with continuous targets as well
-        new_lor = LogisticRegression()
+        new_lor = LogisticRegression(solver='saga')
         self.backing_rfe = RFE(estimator=new_lor, n_features_to_select=self.prop_tuner.value)
 
     def tunable_params(self) -> list[TunableParam]:

environment.yml

@@ -7,5 +7,6 @@ dependencies:
   - ca-certificates
   - openssl
   - scikit-learn
-  - pandas
+  - pandas<3
   - pytest
+  - shap

study/__init__.py

@@ -17,6 +17,9 @@
     "sk_f1_weighted_avg": sk_f1_weighted_avg,
     "sk_f1_perclass": sk_f1_perclass,
     "importance_by_permutation": importance_by_permutation,
+    "shap_additive": shap_additive,
     "correct_samples": correct_samples,
-    "incorrect_samples": incorrect_samples
+    "incorrect_samples": incorrect_samples,
+    "y_true_collector": y_true_collector,
+    "y_pred_proba_collector": y_pred_proba_collector
 }

study/metrics.py

@@ -1,7 +1,10 @@
 """
 Metric-reporting closures for use in this framework.
 """
+import sys
+
 import numpy as np
+import shap
 from sklearn.inspection import permutation_importance
 from sklearn.metrics import balanced_accuracy_score, log_loss, roc_auc_score, precision_score, recall_score, f1_score
 
@@ -18,7 +21,8 @@ def clean_val_for_db(val):
 def sk_log_loss(manager: OptunaModelManager, x: BaseDataManager, y: BaseDataManager):
     # Log Loss
     py = manager.predict_proba(x.as_array())
-    return log_loss(y.as_array(), py)
+    y_labels = [i for i in range(py.shape[1])]
+    return log_loss(y.as_array(), py, labels=y_labels)
 
 def sk_balanced_accuracy(manager: OptunaModelManager, x: BaseDataManager, y: BaseDataManager):
     # Balanced Accuracy
@@ -101,6 +105,95 @@ def importance_by_permutation(manager: OptunaModelManager, x: BaseDataManager, y
     importance_vals = clean_val_for_db(importance_vals)
     return importance_vals
 
+def shap_additive(manager: OptunaModelManager, x: BaseDataManager, _: BaseDataManager):
+    """
+    To restore the (raw) values for a given run, run the following snippet:
+
+    ```
+    from io import StringIO
+
+    # You can omit the Numpy import if you manually
+    # parse the inner string in the list comp
+    import numpy as np
+
+    # This is the SHAP value within the database you want to parse
+    val = ...
+
+    # Strip the brackets first
+    val = val.strip("{").strip("}")
+    # Split by commas
+    entry_strs = val.split(", ")
+    # Split the dataset by feature name
+    shap_map = dict()
+    for entry_str in entry_strs:
+        # Split the text along the colon to get the feature label back
+        feature_label, shap_value_str = entry_str.split(": ")
+        # Parse the shap value string back into numeric form
+        shap_vals = [list(np.fromstring(x, sep=" ")) for x in shap_value_str.split("\n")]
+        # Add it to the map
+        shap_map[feature_label] = shap_vals
+    ```
+
+    Each entry in `shap_map` will be Numpy array of with the following dimensions:
+        * n is the number of samples in the input dataset (train, validate, or test), and
+        * c is the number of categorical classes used during training;
+            If this is binary classification, or a continuous target, c=1.
+
+    For categorical targets with more than 2 classes, each class is treated as
+    unique feature by SHAP for the purpose of calculating SHAP values.
+
+    TODO: Save the `shap_values` directly via pickle into a SQLite blob
+    """
+    # Initialize the explainer, using the x data as both the mask and feature list
+    x_arr = x.as_array()
+    model = manager.get_model()
+
+    if np.unique(x_arr).shape[0] < 2:
+        # SHAP cannot run on a dataset which is entirely homogenous;
+        # return early to avoid an error
+        return "NULL"
+
+    try:
+        # Default to the "generic" explainer
+        explainer = shap.Explainer(
+            model, x_arr, feature_names=x.features()
+        )
+        # Calculate the Shapley values from this dataset
+        shap_values = explainer(x_arr)
+    except TypeError as err:
+        # If that failed, try to use the model's "predict" function instead
+        if hasattr(model, "predict"):
+            explainer = shap.Explainer(
+                model.predict, x_arr, feature_names=x.features()
+            )
+            # Calculate the Shapley values from this dataset
+            shap_values = explainer(x_arr)
+        else:
+            raise err
+
+    shap_list = list()
+    for i, v in enumerate(shap_values.feature_names):
+        # SHAP auto-reduces the shape of its features if it is targeting
+        # a binary classification OR a continuous metric
+        if len(shap_values.values.shape) < 3:
+            val_str = np.array2string(shap_values.values[:, i], max_line_width=sys.maxsize, threshold=sys.maxsize)
+        else:
+            val_str = np.array2string(shap_values.values[:, i, :], max_line_width=sys.maxsize, threshold=sys.maxsize)
+        # Remove the brackets; despite Numpy adding them, it cannot parse them after...
+        val_str = val_str.replace("[", "").replace("]", "")
+        val_str = f"{v}: {val_str}"
+        shap_list.append(val_str)
+
+    # This nonsense is required because Python maps
+    # "/n" to "//n" if you string convert a dict;
+    # why the hell does it do that?!?!?
+    full_str = "{"
+    full_str += ", ".join(shap_list)
+    full_str += "}"
+
+    # Return the result to be saved
+    return full_str
+
 
 """ Sample Reporting """
 def correct_samples(manager: OptunaModelManager, x: BaseDataManager, y: BaseDataManager):
@@ -133,4 +226,16 @@ def incorrect_samples(manager: OptunaModelManager, x: BaseDataManager, y: BaseDa
     # Strip quotation marks from the result so the DB backend doesn't explode
     bad_samples = clean_val_for_db(bad_samples)
 
-    return bad_samples
+    return bad_samples
+
+""" ROC Curve """
+def y_true_collector(_: OptunaModelManager, __: BaseDataManager, y: BaseDataManager):
+    """ Collects the true binary labels for ROC curve generation. """
+    return clean_val_for_db(list(y.as_array().flatten()))
+
+def y_pred_proba_collector(manager: OptunaModelManager, x: BaseDataManager, _: BaseDataManager):
+    """ Collects predicted probabilities for the positive class. """
+    py = manager.predict_proba(x.as_array())
+    if py.shape[1] != 2:
+        raise ValueError(f"Expected binary classification with two probability columns; found {py.shape[1]}.")
+    return clean_val_for_db(list(py[:, 1]))  # Probabilities for the positive class

testing/iris_data/iris_config.json

@@ -15,9 +15,14 @@
     }
   ],
   "post_split_hooks": [
+    {
+      "type": "imputation_simple",
+      "strategy": "most_frequent",
+      "features": ["color", "flower_category", "is_flower", "size"]
+    },
     {
       "type": "one_hot_encode",
-      "features": ["color", "flower_category"]
+      "features": ["color", "flower_category", "is_flower"]
     },
     {
       "type": "ladder_encode",
@@ -33,6 +38,19 @@
       "type": "standard_scaling",
       "run_per_cross": true
     },
+    {
+      "type": "drop_low_variance",
+      "threshold": 0.0
+    },
+    {
+      "type": "principal_component_analysis",
+      "proportion": {
+        "label": "pca_feature_proportion",
+        "type": "float",
+        "low": 0.1,
+        "high": 0.9
+      }
+    },
     {
       "type": "recursive_feature_elimination",
       "proportion": {