FIX bug in SplineTransformer.n_features_out_ (#19577)

mlondschien · web-flow · commit f0a6f054e03b · 2021-02-27T13:09:43.000+01:00
diff --git a/sklearn/preprocessing/_polynomial.py b/sklearn/preprocessing/_polynomial.py
@@ -307,7 +307,7 @@ def fit(self, X, y=None):
         ]
         self.bsplines_ = bsplines
 
-        self.n_features_out_ = n_out - n_features * self.include_bias
+        self.n_features_out_ = n_out - n_features * (1 - self.include_bias)
         return self
 
     def transform(self, X):
@@ -336,7 +336,7 @@ def transform(self, X):
 
         # Note that scipy BSpline returns float64 arrays and converts input
         # x=X[:, i] to c-contiguous float64.
-        n_out = self.n_features_out_ + n_features * self.include_bias
+        n_out = self.n_features_out_ + n_features * (1 - self.include_bias)
         if X.dtype in FLOAT_DTYPES:
             dtype = X.dtype
         else:
diff --git a/sklearn/preprocessing/tests/test_polynomial.py b/sklearn/preprocessing/tests/test_polynomial.py
@@ -243,3 +243,19 @@ def test_spline_transformer_kbindiscretizer():
     # Though they should be exactly equal, we test approximately with high
     # accuracy.
     assert_allclose(splines, kbins, rtol=1e-13)
+
+
+@pytest.mark.parametrize("n_knots", [5, 10])
+@pytest.mark.parametrize("include_bias", [True, False])
+@pytest.mark.parametrize("degree", [3, 5])
+def test_spline_transformer_n_features_out(n_knots, include_bias, degree):
+    """Test that transform results in n_features_out_ features."""
+    splt = SplineTransformer(
+        n_knots=n_knots,
+        degree=degree,
+        include_bias=include_bias
+    )
+    X = np.linspace(0, 1, 10)[:, None]
+    splt.fit(X)
+
+    assert splt.transform(X).shape[1] == splt.n_features_out_
