refactor code and factors out the int64 casting

wdevazelhes · wdevazelhes · commit 90e12b827d37 · 2021-04-19T12:01:29.000+02:00
diff --git a/sklearn/preprocessing/_csr_polynomial_expansion.pyx b/sklearn/preprocessing/_csr_polynomial_expansion.pyx
@@ -8,8 +8,6 @@ from scipy.sparse import csr_matrix
 from numpy cimport ndarray
 import numpy as np
 cimport numpy as np
-cdef int MAX_INT32 = 2147483647
-# TODO: use finfo instead
 
 np.import_array()
 ctypedef fused INDEX_T:
@@ -56,7 +54,9 @@ cdef inline INDEX_T _deg3_column(INDEX_T d, INDEX_T i, INDEX_T j, INDEX_T k,
 def _csr_polynomial_expansion(ndarray[DATA_T, ndim=1] data,
                               ndarray[INDEX_T, ndim=1] indices,
                               ndarray[INDEX_T, ndim=1] indptr,
-                              INDEX_T d, INDEX_T interaction_only,
+                              INDEX_T d,
+                              INDEX_T expanded_dimensionality,
+                              INDEX_T interaction_only,
                               INDEX_T degree):
     """
     Perform a second-degree polynomial or interaction expansion on a scipy
@@ -91,13 +91,6 @@ def _csr_polynomial_expansion(ndarray[DATA_T, ndim=1] data,
     Matrices Using K-Simplex Numbers" by Andrew Nystrom and John Hughes.
     """
 
-    assert degree in (2, 3)
-
-    if degree == 2:
-        expanded_dimensionality = int((d**2 + d) / 2 - interaction_only*d)
-    else:
-        expanded_dimensionality = int((d**3 + 3*d**2 + 2*d) / 6
-                                      - interaction_only*d**2)
     if expanded_dimensionality == 0:
         return None
     assert expanded_dimensionality > 0
@@ -106,12 +99,7 @@ def _csr_polynomial_expansion(ndarray[DATA_T, ndim=1] data,
 
     # Count how many nonzero elements the expanded matrix will contain.
     for row_i in range(indptr.shape[0]-1):
-        # nnz is the number of nonzero elements in this row.
-        # The number of nonzero elements can explode
-        # in the expanded space, so we cast to int64
-        # before the expansion computation to
-        # avoid overflow:
-        nnz = np.int64(indptr[row_i + 1] - indptr[row_i])
+        nnz = indptr[row_i + 1] - indptr[row_i]
         # TODO: check that the casting is indeed done
         if degree == 2:
             total_nnz += (nnz ** 2 + nnz) / 2 - interaction_only * nnz
@@ -128,16 +116,6 @@ def _csr_polynomial_expansion(ndarray[DATA_T, ndim=1] data,
     cdef ndarray[INDEX_T, ndim=1] expanded_indptr = ndarray(
         shape=num_rows + 1, dtype=indptr.dtype)
 
-    # if the expanded_dimensionality is too big, we need to cast
-    # indices to int64. If total_nnz is too big, we need to cast
-    # expanded_indptr to int64 too.
-    # TODO: check that the casting is indeed done
-    if expanded_dimensionality > MAX_INT32:
-        indices = np.array(indices, dtype=np.int64)
-        expanded_indices = np.array(expanded_indices, dtype=np.int64)
-    if total_nnz > MAX_INT32:
-        expanded_indptr = np.array(expanded_indptr, dtype=np.int64)
-
     cdef INDEX_T expanded_index = 0, row_starts, row_ends, i, j, k, \
                  i_ptr, j_ptr, k_ptr, num_cols_in_row, col
 
diff --git a/sklearn/preprocessing/_polynomial.py b/sklearn/preprocessing/_polynomial.py
@@ -23,6 +23,32 @@
 ]
 
 
+def _cast_to_int64_if_needed(X, degree, interaction_only):
+    """Computes the expanded dimensionality and casts X to int64 if the
+    expanded dim is too big"""
+    d = int(X.shape[1])
+    assert degree in (2, 3)
+    if degree == 2:
+        expanded_dimensionality = (d ** 2 + d) / 2 - interaction_only * d
+    else:
+        expanded_dimensionality = ((d ** 3 + 3 * d ** 2 + 2 * d) / 6 -
+                                   interaction_only * d ** 2)
+    if expanded_dimensionality > np.iinfo(np.int64).max:
+        raise ValueError("The expanded dimensionality will be too big to "
+                         "be contained in an int64.")
+    if expanded_dimensionality > np.iinfo(np.int32).max:
+        # if the expansion needs int64 ints, we cast every index value in X
+        # to int64
+        X = X.copy()
+        X.data = X.data.astype(np.int64)
+        X.indices = X.indices.astype(np.int64)
+        X.indptr = X.indptr.astype(np.int64)
+        return X, np.int64(d), np.int64(expanded_dimensionality)
+    else:
+        # otherwise we keep X as is
+        return X, np.int32(d), np.int32(expanded_dimensionality)
+
+
 class PolynomialFeatures(TransformerMixin, BaseEstimator):
     """Generate polynomial and interaction features.
 
@@ -249,12 +275,13 @@ def transform(self, X):
                 to_stack.append(np.ones(shape=(n_samples, 1), dtype=X.dtype))
             to_stack.append(X)
             for deg in range(2, self.degree+1):
-                # Dimensionality of the expanded space can become very
-                # large so we cast X.shape[1] to int64 to avoid overflow
-                # when computing expanded_dimensionality
+                (X, d, expanded_dim
+                 ) = _cast_to_int64_if_needed(X, deg,
+                                              self.interaction_only)
                 Xp_next = _csr_polynomial_expansion(X.data, X.indices,
                                                     X.indptr,
-                                                    np.int64(X.shape[1]),
+                                                    d,
+                                                    expanded_dim,
                                                     self.interaction_only,
                                                     deg)
                 if Xp_next is None:
