Repository: scipy/scipy

Issue:
Reshaping COO sparse matrices fails if the matrix has 2^31 or more elements
The [`coo_matrix.reshape` method](https://github.com/scipy/scipy/blob/master/scipy/sparse/coo.py#L194-L224) computes the linear index of each nonzero element (variable `flat_indexes` in the code) as an intermediate step. If there are more than 2^31 elements, the linear index overflows to negative and then the construction of the new matrix blows up. (This is the number of elements in the abstract matrix, not the number of nonzero elements in the sparse representation.)

This is particularly problematic because COO reshape is the default reshape method for sparse matrices so it affect CSC, CSR, and others.

What do you think the best solution is?

The obvious path is to do nothing and require that users supply the indexes as int64 if they want correct reshape behavior even if int32 is sufficient for the row and column indexes. However, this is very painful. Constructing CSR and CSC matrices with int64 indexes is not possible. If the user provides the rows and columns or indices and indptr as int64, the indexes and indptr get lowered to int32 anyway.

There are a couple of conservative solutions to the problem:

1. Raise the dtype to np.int64 for the linear index calculation.
2. Raise the dtype to np.int64 for the linear index calculation only when then number of elements is 2^31 or greater.

In the first two cases, we would have to decide under what circumstances we should lower the dtype back down to int32. 

It might be possible to rewrite the index calculation so that it simply does not overflow, but it is not obvious to me how to do that.
Reshaping COO sparse matrices fails if the matrix has 2^31 or more elements
The [`coo_matrix.reshape` method](https://github.com/scipy/scipy/blob/master/scipy/sparse/coo.py#L194-L224) computes the linear index of each nonzero element (variable `flat_indexes` in the code) as an intermediate step. If there are more than 2^31 elements, the linear index overflows to negative and then the construction of the new matrix blows up. (This is the number of elements in the abstract matrix, not the number of nonzero elements in the sparse representation.)

This is particularly problematic because COO reshape is the default reshape method for sparse matrices so it affect CSC, CSR, and others.

What do you think the best solution is?

The obvious path is to do nothing and require that users supply the indexes as int64 if they want correct reshape behavior even if int32 is sufficient for the row and column indexes. However, this is very painful. Constructing CSR and CSC matrices with int64 indexes is not possible. If the user provides the rows and columns or indices and indptr as int64, the indexes and indptr get lowered to int32 anyway.

There are a couple of conservative solutions to the problem:

1. Raise the dtype to np.int64 for the linear index calculation.
2. Raise the dtype to np.int64 for the linear index calculation only when then number of elements is 2^31 or greater.

In the first two cases, we would have to decide under what circumstances we should lower the dtype back down to int32. 

It might be possible to rewrite the index calculation so that it simply does not overflow, but it is not obvious to me how to do that.


Discussion:
If you want to experiment with alternative calculation methods, you have an expression of the form `divmod(a * b + c, d)`. If we rewrite `a = a0 * d + a1`, `b = b0 * d + b1` and `c = c0 * d + c1` we have `divmod((a0 * b0 * d + a1 * b0 + a0 * b1 + c0) * d + a1 * b1 + c1, d)`, which is very striaghtforward to calculate. So at the cost of replacing one division with 3, maybe 4:

```
rows, cols = self.shape
new_rows, new_cols = shape
cols_quot, cols_rem = divmod(cols, new_cols)
row_quot, row_rem = divmod(self.row, new_cols)
col_quot, col_rem = divmod(self.col, new_cols)

new_row = (cols_quot * row_quot * new_cols + cols_rem * row_quot +
           cols_quot * row_rem + col_quot)
new_col = cols_rem * row_rem + col_rem

if np.any(new_col >= new_cols):  # EDITED: s/>/>=/
    row_delta, new_col = divmod(new_col, new_cols)
    new_row += row_delta
```
I am trying @jaimefrio's solution. I think that 

> `if np.any(new_col > new_cols):`

needs to be 

`if np.any(new_col >= new_cols):`

Yup, I think you are right, have edited my answer to not trip future readers.
I am planning on making a pull request soon for this. I think this is the full chunk of code that needs to be added:

```
        if self.size == 0:
            # Special case an empty matrix to avoid divide by zero error
            new_row = []
            new_col = []
        elif order == 'C':
            ncols_quot, ncols_rem = divmod(ncols, new_ncols)
            row_quot, row_rem = divmod(self.row, new_ncols)
            col_quot, col_rem = divmod(self.col, new_ncols)

            new_row = (ncols_quot * row_quot * new_ncols + ncols_quot * row_rem +
                       ncols_rem * row_quot + col_quot)
            new_col = ncols_rem * row_rem + col_rem

            if np.any(new_col >= new_ncols):
                row_delta, new_col = divmod(new_col, new_ncols)
                new_row += row_delta
        elif order == 'F':
            nrows_quot, nrows_rem = divmod(nrows, new_nrows)
            col_quot, col_rem = divmod(self.col, new_nrows)
            row_quot, row_rem = divmod(self.row, new_nrows)

            new_col = (nrows_quot * col_quot * new_ncols + nrows_quot * col_rem +
                       nrows_rem * col_quot + row_quot)
            new_row = nrows_rem * col_rem + row_rem

            if np.any(new_row >= new_nrows):
                col_delta, new_row = divmod(new_row, new_nrows)
                new_col += col_delta
```

One question I have is: is that `if np.any(new_col >= new_ncols)` check worth it? As far as I can tell, this doesn't change the behavior of the method at all; it merely avoids an integer division when the quotient would be zero anyway. Is this to improve performance? This will trigger pretty much anytime the number of columns shrinks and this logical comparison will have been for naught.

Another question is should this new method only be triggered when the size will overflow, and it should follow the original path for small matrices?
If you want to experiment with alternative calculation methods, you have an expression of the form `divmod(a * b + c, d)`. If we rewrite `a = a0 * d + a1`, `b = b0 * d + b1` and `c = c0 * d + c1` we have `divmod((a0 * b0 * d + a1 * b0 + a0 * b1 + c0) * d + a1 * b1 + c1, d)`, which is very striaghtforward to calculate. So at the cost of replacing one division with 3, maybe 4:

```
rows, cols = self.shape
new_rows, new_cols = shape
cols_quot, cols_rem = divmod(cols, new_cols)
row_quot, row_rem = divmod(self.row, new_cols)
col_quot, col_rem = divmod(self.col, new_cols)

new_row = (cols_quot * row_quot * new_cols + cols_rem * row_quot +
           cols_quot * row_rem + col_quot)
new_col = cols_rem * row_rem + col_rem

if np.any(new_col >= new_cols):  # EDITED: s/>/>=/
    row_delta, new_col = divmod(new_col, new_cols)
    new_row += row_delta
```
I am trying @jaimefrio's solution. I think that 

> `if np.any(new_col > new_cols):`

needs to be 

`if np.any(new_col >= new_cols):`

Yup, I think you are right, have edited my answer to not trip future readers.
I am planning on making a pull request soon for this. I think this is the full chunk of code that needs to be added:

```
        if self.size == 0:
            # Special case an empty matrix to avoid divide by zero error
            new_row = []
            new_col = []
        elif order == 'C':
            ncols_quot, ncols_rem = divmod(ncols, new_ncols)
            row_quot, row_rem = divmod(self.row, new_ncols)
            col_quot, col_rem = divmod(self.col, new_ncols)

            new_row = (ncols_quot * row_quot * new_ncols + ncols_quot * row_rem +
                       ncols_rem * row_quot + col_quot)
            new_col = ncols_rem * row_rem + col_rem

            if np.any(new_col >= new_ncols):
                row_delta, new_col = divmod(new_col, new_ncols)
                new_row += row_delta
        elif order == 'F':
            nrows_quot, nrows_rem = divmod(nrows, new_nrows)
            col_quot, col_rem = divmod(self.col, new_nrows)
            row_quot, row_rem = divmod(self.row, new_nrows)

            new_col = (nrows_quot * col_quot * new_ncols + nrows_quot * col_rem +
                       nrows_rem * col_quot + row_quot)
            new_row = nrows_rem * col_rem + row_rem

            if np.any(new_row >= new_nrows):
                col_delta, new_row = divmod(new_row, new_nrows)
                new_col += col_delta
```

One question I have is: is that `if np.any(new_col >= new_ncols)` check worth it? As far as I can tell, this doesn't change the behavior of the method at all; it merely avoids an integer division when the quotient would be zero anyway. Is this to improve performance? This will trigger pretty much anytime the number of columns shrinks and this logical comparison will have been for naught.

Another question is should this new method only be triggered when the size will overflow, and it should follow the original path for small matrices?

Resolving patch:
diff --git a/scipy/sparse/coo.py b/scipy/sparse/coo.py
--- a/scipy/sparse/coo.py
+++ b/scipy/sparse/coo.py
@@ -153,7 +153,7 @@ def __init__(self, arg1, shape=None, dtype=None, copy=False):
                     M, N = shape
                     self._shape = check_shape((M, N))
 
-                idx_dtype = get_index_dtype(maxval=max(self.shape))
+                idx_dtype = get_index_dtype(maxval=self.shape[0] * self.shape[1])
                 self.row = np.array(row, copy=copy, dtype=idx_dtype)
                 self.col = np.array(col, copy=copy, dtype=idx_dtype)
                 self.data = np.array(obj, copy=copy)
@@ -262,7 +262,7 @@ def _check(self):
             warn("col index array has non-integer dtype (%s) "
                     % self.col.dtype.name)
 
-        idx_dtype = get_index_dtype(maxval=max(self.shape))
+        idx_dtype = get_index_dtype(maxval=self.shape[0] * self.shape[1])
         self.row = np.asarray(self.row, dtype=idx_dtype)
         self.col = np.asarray(self.col, dtype=idx_dtype)
         self.data = to_native(self.data)


Test patch:
diff --git a/scipy/sparse/tests/test_base.py b/scipy/sparse/tests/test_base.py
--- a/scipy/sparse/tests/test_base.py
+++ b/scipy/sparse/tests/test_base.py
@@ -4074,6 +4074,17 @@ def test_reshape_copy(self):
         y = x.reshape(new_shape, copy=True)
         assert_(not np.may_share_memory(y.data, x.data))
 
+    def test_large_dimensions(self):
+        # Test that reshape is immune to integer overflow when number of elements
+        # exceeds 2^31-1
+        mat1 = coo_matrix(([1], ([3000000], [1000])), (3000001, 1001))
+        mat2 = coo_matrix(([1], ([1000], [3000000])), (1001, 3000001))
+
+        # assert_array_equal is slow for big matrices because it expects dense
+        # Using __ne__ and nnz instead
+        assert_((mat1.reshape((1001, 3000001), order='C') != mat2).nnz == 0)
+        assert_((mat2.reshape((3000001, 1001), order='F') != mat1).nnz == 0)
+
 
 TestCOO.init_class()
 
@@ -4510,6 +4521,7 @@ def cases_64bit():
         'test_inv': 'linsolve for 64-bit indices not available',
         'test_solve': 'linsolve for 64-bit indices not available',
         'test_scalar_idx_dtype': 'test implemented in base class',
+        'test_large_dimensions': 'test actually requires 64-bit to work',
     }
 
     for cls in TEST_CLASSES: