TST Extend tests for scipy.sparse.*array in test_polynomial.py (scikit-learn#27166)

Co-authored-by: Guillaume Lemaitre <g.lemaitre58@gmail.com>

Author: work-mohit

doc/whats_new/v1.4.rst

@@ -119,6 +119,8 @@ and classes are impacted:
   :user:`Lohit SundaramahaLingam <lohitslohit>`;
 - :class:`neural_network.BernoulliRBM` in :pr:`27252` by
   :user:`Yao Xiao <Charlie-XIAO>`.
+- :class:`preprocessing.PolynomialFeatures` in :pr:`27166` by
+  :user:`Mohit Joshi <work-mohit>`.
 
 Changelog
 ---------

sklearn/preprocessing/_polynomial.py

@@ -495,7 +495,7 @@ def transform(self, X):
                 if combi:
                     out_col = 1
                     for col_idx in combi:
-                        out_col = X[:, col_idx].multiply(out_col)
+                        out_col = X[:, [col_idx]].multiply(out_col)
                     columns.append(out_col)
                 else:
                     bias = sparse.csc_matrix(np.ones((X.shape[0], 1)))

sklearn/preprocessing/tests/test_polynomial.py

@@ -20,7 +20,12 @@
     _get_sizeof_LARGEST_INT_t,
 )
 from sklearn.utils._testing import assert_array_almost_equal
-from sklearn.utils.fixes import parse_version, sp_version
+from sklearn.utils.fixes import (
+    CSC_CONTAINERS,
+    CSR_CONTAINERS,
+    parse_version,
+    sp_version,
+)
 
 
 @pytest.mark.parametrize("est", (PolynomialFeatures, SplineTransformer))
@@ -522,27 +527,24 @@ def single_feature_degree3():
         ((2, 3), False, True, []),
     ],
 )
-@pytest.mark.parametrize(
-    "sparse_X",
-    [False, sparse.csr_matrix, sparse.csc_matrix],
-)
+@pytest.mark.parametrize("X_container", [None] + CSR_CONTAINERS + CSC_CONTAINERS)
 def test_polynomial_features_one_feature(
     single_feature_degree3,
     degree,
     include_bias,
     interaction_only,
     indices,
-    sparse_X,
+    X_container,
 ):
     """Test PolynomialFeatures on single feature up to degree 3."""
     X, P = single_feature_degree3
-    if sparse_X:
-        X = sparse_X(X)
+    if X_container is not None:
+        X = X_container(X)
     tf = PolynomialFeatures(
         degree=degree, include_bias=include_bias, interaction_only=interaction_only
     ).fit(X)
     out = tf.transform(X)
-    if sparse_X:
+    if X_container is not None:
         out = out.toarray()
     assert_allclose(out, P[:, indices])
     if tf.n_output_features_ > 0:
@@ -596,27 +598,24 @@ def two_features_degree3():
         ((3, 3), False, True, []),  # would need 3 input features
     ],
 )
-@pytest.mark.parametrize(
-    "sparse_X",
-    [False, sparse.csr_matrix, sparse.csc_matrix],
-)
+@pytest.mark.parametrize("X_container", [None] + CSR_CONTAINERS + CSC_CONTAINERS)
 def test_polynomial_features_two_features(
     two_features_degree3,
     degree,
     include_bias,
     interaction_only,
     indices,
-    sparse_X,
+    X_container,
 ):
     """Test PolynomialFeatures on 2 features up to degree 3."""
     X, P = two_features_degree3
-    if sparse_X:
-        X = sparse_X(X)
+    if X_container is not None:
+        X = X_container(X)
     tf = PolynomialFeatures(
         degree=degree, include_bias=include_bias, interaction_only=interaction_only
     ).fit(X)
     out = tf.transform(X)
-    if sparse_X:
+    if X_container is not None:
         out = out.toarray()
     assert_allclose(out, P[:, indices])
     if tf.n_output_features_ > 0:
@@ -712,10 +711,13 @@ def test_polynomial_feature_names():
         (4, False, True, np.float64),
     ],
 )
-def test_polynomial_features_csc_X(deg, include_bias, interaction_only, dtype):
+@pytest.mark.parametrize("csc_container", CSC_CONTAINERS)
+def test_polynomial_features_csc_X(
+    deg, include_bias, interaction_only, dtype, csc_container
+):
     rng = np.random.RandomState(0)
     X = rng.randint(0, 2, (100, 2))
-    X_csc = sparse.csc_matrix(X)
+    X_csc = csc_container(X)
 
     est = PolynomialFeatures(
         deg, include_bias=include_bias, interaction_only=interaction_only
@@ -739,10 +741,13 @@ def test_polynomial_features_csc_X(deg, include_bias, interaction_only, dtype):
         (3, False, True, np.float64),
     ],
 )
-def test_polynomial_features_csr_X(deg, include_bias, interaction_only, dtype):
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_polynomial_features_csr_X(
+    deg, include_bias, interaction_only, dtype, csr_container
+):
     rng = np.random.RandomState(0)
     X = rng.randint(0, 2, (100, 2))
-    X_csr = sparse.csr_matrix(X)
+    X_csr = csr_container(X)
 
     est = PolynomialFeatures(
         deg, include_bias=include_bias, interaction_only=interaction_only
@@ -761,17 +766,14 @@ def test_polynomial_features_csr_X(deg, include_bias, interaction_only, dtype):
 )
 @pytest.mark.parametrize("interaction_only", [True, False])
 @pytest.mark.parametrize("include_bias", [True, False])
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
 def test_num_combinations(
-    n_features,
-    min_degree,
-    max_degree,
-    interaction_only,
-    include_bias,
+    n_features, min_degree, max_degree, interaction_only, include_bias, csr_container
 ):
     """
     Test that n_output_features_ is calculated correctly.
     """
-    x = sparse.csr_matrix(([1], ([0], [n_features - 1])))
+    x = csr_container(([1], ([0], [n_features - 1])))
     est = PolynomialFeatures(
         degree=max_degree,
         interaction_only=interaction_only,
@@ -799,8 +801,11 @@ def test_num_combinations(
         (3, False, True, np.float64),
     ],
 )
-def test_polynomial_features_csr_X_floats(deg, include_bias, interaction_only, dtype):
-    X_csr = sparse_random(1000, 10, 0.5, random_state=0).tocsr()
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_polynomial_features_csr_X_floats(
+    deg, include_bias, interaction_only, dtype, csr_container
+):
+    X_csr = csr_container(sparse_random(1000, 10, 0.5, random_state=0))
     X = X_csr.toarray()
 
     est = PolynomialFeatures(
@@ -831,8 +836,11 @@ def test_polynomial_features_csr_X_floats(deg, include_bias, interaction_only, d
         (2, 3, False),
     ],
 )
-def test_polynomial_features_csr_X_zero_row(zero_row_index, deg, interaction_only):
-    X_csr = sparse_random(3, 10, 1.0, random_state=0).tocsr()
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_polynomial_features_csr_X_zero_row(
+    zero_row_index, deg, interaction_only, csr_container
+):
+    X_csr = csr_container(sparse_random(3, 10, 1.0, random_state=0))
     X_csr[zero_row_index, :] = 0.0
     X = X_csr.toarray()
 
@@ -851,8 +859,11 @@ def test_polynomial_features_csr_X_zero_row(zero_row_index, deg, interaction_onl
     ["include_bias", "interaction_only"],
     [(True, True), (True, False), (False, True), (False, False)],
 )
-def test_polynomial_features_csr_X_degree_4(include_bias, interaction_only):
-    X_csr = sparse_random(1000, 10, 0.5, random_state=0).tocsr()
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_polynomial_features_csr_X_degree_4(
+    include_bias, interaction_only, csr_container
+):
+    X_csr = csr_container(sparse_random(1000, 10, 0.5, random_state=0))
     X = X_csr.toarray()
 
     est = PolynomialFeatures(
@@ -881,8 +892,9 @@ def test_polynomial_features_csr_X_degree_4(include_bias, interaction_only):
         (3, 3, False),
     ],
 )
-def test_polynomial_features_csr_X_dim_edges(deg, dim, interaction_only):
-    X_csr = sparse_random(1000, dim, 0.5, random_state=0).tocsr()
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_polynomial_features_csr_X_dim_edges(deg, dim, interaction_only, csr_container):
+    X_csr = csr_container(sparse_random(1000, dim, 0.5, random_state=0))
     X = X_csr.toarray()
 
     est = PolynomialFeatures(deg, interaction_only=interaction_only)
@@ -896,8 +908,9 @@ def test_polynomial_features_csr_X_dim_edges(deg, dim, interaction_only):
 
 @pytest.mark.parametrize("interaction_only", [True, False])
 @pytest.mark.parametrize("include_bias", [True, False])
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
 def test_csr_polynomial_expansion_index_overflow_non_regression(
-    interaction_only, include_bias
+    interaction_only, include_bias, csr_container
 ):
     """Check the automatic index dtype promotion to `np.int64` when needed.
 
@@ -925,7 +938,7 @@ def degree_2_calc(d, i, j):
     col = np.array(
         [n_features - 2, n_features - 1, n_features - 2, n_features - 1], dtype=np.int64
     )
-    X = sparse.csr_matrix(
+    X = csr_container(
         (data, (row, col)),
         shape=(n_samples, n_features),
         dtype=data_dtype,
@@ -1028,8 +1041,9 @@ def degree_2_calc(d, i, j):
 )
 @pytest.mark.parametrize("interaction_only", [True, False])
 @pytest.mark.parametrize("include_bias", [True, False])
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
 def test_csr_polynomial_expansion_index_overflow(
-    degree, n_features, interaction_only, include_bias
+    degree, n_features, interaction_only, include_bias, csr_container
 ):
     """Tests known edge-cases to the dtype promotion strategy and custom
     Cython code, including a current bug in the upstream
@@ -1050,7 +1064,7 @@ def test_csr_polynomial_expansion_index_overflow(
         n_features * (n_features + 1) * (n_features + 2) // 6 + expected_indices[1]
     )
 
-    X = sparse.csr_matrix((data, (row, col)))
+    X = csr_container((data, (row, col)))
     pf = PolynomialFeatures(
         interaction_only=interaction_only, include_bias=include_bias, degree=degree
     )
@@ -1131,12 +1145,15 @@ def test_csr_polynomial_expansion_index_overflow(
 
 @pytest.mark.parametrize("interaction_only", [True, False])
 @pytest.mark.parametrize("include_bias", [True, False])
-def test_csr_polynomial_expansion_too_large_to_index(interaction_only, include_bias):
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_csr_polynomial_expansion_too_large_to_index(
+    interaction_only, include_bias, csr_container
+):
     n_features = np.iinfo(np.int64).max // 2
     data = [1.0]
     row = [0]
     col = [n_features - 1]
-    X = sparse.csr_matrix((data, (row, col)))
+    X = csr_container((data, (row, col)))
     pf = PolynomialFeatures(
         interaction_only=interaction_only, include_bias=include_bias, degree=(2, 2)
     )
@@ -1150,7 +1167,8 @@ def test_csr_polynomial_expansion_too_large_to_index(interaction_only, include_b
         pf.fit_transform(X)
 
 
-def test_polynomial_features_behaviour_on_zero_degree():
+@pytest.mark.parametrize("sparse_container", CSR_CONTAINERS + CSC_CONTAINERS)
+def test_polynomial_features_behaviour_on_zero_degree(sparse_container):
     """Check that PolynomialFeatures raises error when degree=0 and include_bias=False,
     and output a single constant column when include_bias=True
     """
@@ -1171,7 +1189,7 @@ def test_polynomial_features_behaviour_on_zero_degree():
     with pytest.raises(ValueError, match=err_msg):
         poly.fit_transform(X)
 
-    for _X in [X, sparse.csr_matrix(X), sparse.csc_matrix(X)]:
+    for _X in [X, sparse_container(X)]:
         poly = PolynomialFeatures(degree=0, include_bias=True)
         output = poly.fit_transform(_X)
         # convert to dense array if needed
@@ -1202,7 +1220,8 @@ def test_sizeof_LARGEST_INT_t():
     ),
     run=True,
 )
-def test_csr_polynomial_expansion_windows_fail():
+@pytest.mark.parametrize("csr_container", CSR_CONTAINERS)
+def test_csr_polynomial_expansion_windows_fail(csr_container):
     # Minimum needed to ensure integer overflow occurs while guaranteeing an
     # int64-indexable output.
     n_features = int(np.iinfo(np.int64).max ** (1 / 3) + 3)
@@ -1223,7 +1242,7 @@ def test_csr_polynomial_expansion_windows_fail():
         int(n_features * (n_features + 1) * (n_features + 2) // 6 + expected_indices[1])
     )
 
-    X = sparse.csr_matrix((data, (row, col)))
+    X = csr_container((data, (row, col)))
     pf = PolynomialFeatures(interaction_only=False, include_bias=False, degree=3)
     if sys.maxsize <= 2**32:
         msg = (