gitlabds 2.1.10__tar.gz → 2.1.11__tar.gz

{gitlabds-2.1.10 → gitlabds-2.1.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gitlabds
-Version: 2.1.10
+Version: 2.1.11
 Summary: GitLab Data Science Tools
 Author-email: Kevin Dietz <kdietz@gitlab.com>
 License-Expression: MIT

{gitlabds-2.1.10 → gitlabds-2.1.11}/gitlabds/model_evaluator.py

@@ -2044,6 +2044,70 @@ class ModelEvaluator:
                     ax.set_yticks([])
                     ax.set_title(f"Class {class_idx}")
 
+    def _expand_categorical_shap(self, shap_values, X):
+        """
+        Expand each categorical column's single SHAP column into one binary-
+        indicator column per level (e.g. ``industry_prev`` -> ``industry_prev=Manufacturing``,
+        ``industry_prev=Tech``, ...).
+
+        For each row, the per-level SHAP value equals the original column's SHAP
+        value when the row had that level, otherwise 0. The feature value is a
+        binary indicator (1 if the row has that level, else 0). This lets the
+        beeswarm color gradient (red/blue) carry direction-of-effect information
+        per level rather than collapsing every row to a single gray dot when
+        native categorical handling (``enable_categorical=True``) is used.
+
+        Works with both 2D (binary classification / regression) and 3D
+        (multiclass) SHAP value arrays. Non-categorical columns pass through
+        untouched. If no categorical columns are detected, the input is returned
+        unchanged.
+        """
+        try:
+            import shap
+        except ImportError:
+            return shap_values
+
+        cat_cols = [c for c in X.columns if isinstance(X[c].dtype, pd.CategoricalDtype)]
+        if not cat_cols:
+            return shap_values
+
+        feature_names = list(X.columns)
+        values = shap_values.values
+        is_3d = values.ndim == 3
+
+        new_feature_names = []
+        new_value_cols = []
+        new_data_cols = []
+
+        for j, col in enumerate(feature_names):
+            col_shap = values[:, j, :] if is_3d else values[:, j]
+            if col in cat_cols:
+                col_raw = X[col].astype(object).values
+                # Stable, sorted level ordering for reproducibility
+                levels = sorted([lv for lv in pd.unique(X[col]) if pd.notna(lv)],
+                                key=lambda v: str(v))
+                for level in levels:
+                    mask = (col_raw == level)
+                    if is_3d:
+                        masked = np.where(mask[:, None], col_shap, 0.0)
+                    else:
+                        masked = np.where(mask, col_shap, 0.0)
+                    new_value_cols.append(masked)
+                    new_data_cols.append(mask.astype(float))
+                    new_feature_names.append(f"{col}={level}")
+            else:
+                new_value_cols.append(col_shap)
+                new_data_cols.append(X[col].values)
+                new_feature_names.append(col)
+
+        new_values = np.stack(new_value_cols, axis=1)
+        new_data = np.stack(new_data_cols, axis=1)
+
+        kwargs = dict(values=new_values, data=new_data, feature_names=new_feature_names)
+        if hasattr(shap_values, "base_values") and shap_values.base_values is not None:
+            kwargs["base_values"] = shap_values.base_values
+        return shap.Explanation(**kwargs)
+
     def plot_shap_beeswarm(
         self,
         n_features: int = 20,
@@ -2055,17 +2119,29 @@ class ModelEvaluator:
         plot_name: str = "shap_beeswarm",
         plot_save_format: str = None,
         plot_save_dpi: str = None,
-        show_plot: bool = None
+        show_plot: bool = None,
+        expand_categoricals: bool = True,
     ) -> None:
         """
         Create a SHAP beeswarm plot for feature importance.
-        
+
         Parameters
         ----------
         class_index : int, optional
             For multi-class classification, specify which class to plot.
             If None, plots global importance (averaged across classes).
             Only applicable for multi-class models.
+        expand_categoricals : bool, default True
+            If True, each categorical column (pandas ``category`` dtype) is split
+            into per-level binary-indicator pseudo-features (e.g.
+            ``industry_prev=Manufacturing``) before plotting. This lets the
+            red/blue color gradient encode direction-of-effect per level when
+            using native categorical handling (``enable_categorical=True``)
+            instead of showing a single gray-dot row per categorical feature.
+            The ``n_features`` (max_display) cap then ranks numerical features
+            and per-level indicators on the same mean-|SHAP| scale, so only the
+            impactful levels surface. Set False to preserve the original
+            behavior.
         """
         # Set defaults if not provided
         save_plot = self.save_plots if save_plot is None else save_plot
@@ -2094,10 +2170,17 @@ class ModelEvaluator:
         try:
             # Create an explainer
             explainer = shap.Explainer(self.model)
-            
+
             # Get SHAP values
             shap_values = explainer(X)
-            
+
+            # Expand categorical columns into per-level binary indicators so
+            # the beeswarm color gradient encodes direction-of-effect per level
+            # (e.g. industry_prev=Manufacturing red dots vs blue dots) rather
+            # than collapsing each categorical feature to a single gray-dot row.
+            if expand_categoricals:
+                shap_values = self._expand_categorical_shap(shap_values, X)
+
             # For multi-class classification
             if self.classification and self.n_classes > 2:
                 if hasattr(shap_values, 'values') and len(shap_values.values.shape) == 3:
@@ -2112,8 +2195,8 @@ class ModelEvaluator:
                         class_values = shap_values.values[:, :, class_index]
                         class_exp = shap.Explanation(
                             values=class_values,
-                            data=X.values,
-                            feature_names=X.columns.tolist()
+                            data=shap_values.data,
+                            feature_names=shap_values.feature_names,
                         )
                         
                         # Create a figure first
@@ -2132,8 +2215,8 @@ class ModelEvaluator:
                         global_values = shap_values.values.mean(axis=2)
                         global_exp = shap.Explanation(
                             values=global_values,
-                            data=X.values,
-                            feature_names=X.columns.tolist()
+                            data=shap_values.data,
+                            feature_names=shap_values.feature_names,
                         )
                         
                         # Create a figure first
@@ -2489,8 +2572,8 @@ class ModelEvaluator:
             # Colored scatter plot
             # If specific pairs not provided, use top correlation pairs
             if feature_pairs is None:
-                # Calculate correlation matrix
-                corr_matrix = self.x_train.corr().abs()
+                # Calculate correlation matrix (skip categoricals — corr is undefined for them)
+                corr_matrix = self.x_train.corr(numeric_only=True).abs()
                 
                 # Set diagonal to 0
                 np.fill_diagonal(corr_matrix.values, 0)

{gitlabds-2.1.10 → gitlabds-2.1.11}/gitlabds.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gitlabds
-Version: 2.1.10
+Version: 2.1.11
 Summary: GitLab Data Science Tools
 Author-email: Kevin Dietz <kdietz@gitlab.com>
 License-Expression: MIT