ENH improved CSC matrix handling in VarianceThreshold

Thomas Unterthiner · Thomas Unterthiner · commit 50e20a0002d4 · 2014-09-03T10:51:33.000+02:00
diff --git a/sklearn/feature_selection/variance_threshold.py b/sklearn/feature_selection/variance_threshold.py
@@ -5,7 +5,7 @@
 from ..base import BaseEstimator
 from .base import SelectorMixin
 from ..utils import check_array
-from ..utils.sparsefuncs_fast import csr_mean_variance_axis0
+from ..utils.sparsefuncs import mean_variance_axis
 
 
 class VarianceThreshold(BaseEstimator, SelectorMixin):
@@ -58,10 +58,10 @@ def fit(self, X, y=None):
         -------
         self
         """
-        X = check_array(X, 'csr', dtype=np.float64)
+        X = check_array(X, ('csr', 'csc'), dtype=np.float64)
 
         if hasattr(X, "toarray"):   # sparse matrix
-            _, self.variances_ = csr_mean_variance_axis0(X)
+            _, self.variances_ = mean_variance_axis(X, axis=0)
         else:
             self.variances_ = np.var(X, axis=0)
 
