add uniform feature subsampling to gam_shap

csinva · Nov 6, 2024 · e555c01 · e555c01
1 parent d47ead5
commit e555c01
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/imodels/algebraic/gam_shap.py b/imodels/algebraic/gam_shap.py
@@ -30,9 +30,13 @@ def fit(self, X, y):
 
         for _ in range(self.n_estimators):
             # Randomly select a subset of features
+            if isinstance(self.feature_fraction, float):
+                n_features = max(1, int(self.feature_fraction * n_features))
+            elif self.feature_fraction == 'uniform':
+                n_features = rng.integers(1, n_features + 1)
             feature_subset = rng.choice(
                 n_features,
-                size=int(self.feature_fraction * n_features),
+                size=n_features,
                 replace=False
             )
             self.feature_subsets.append(feature_subset)