Scipy sparse memory explosion with simple matrix multiplication

I noted that Scipy must be storing some intermediate arrays when doing matrix multiplication. I assume this can be helpful in some cases, but it is a pain sometimes. Consider the following example:

from scipy.sparse import coo_matrix
n = 100000000000
row  = np.array([0, 0])
col  = np.array([0, n-1])
data = np.array([1, 1])
A = coo_matrix((data, (row, col)), shape=(2, n))

Yes, this is a very large matrix. However it has only two nonzero values. The result of B = A.dot(A.T) can be evaluated by hand since it has only one nonzero value. The matrix B is such that B[0, 0] = 2 and zero elsewhere. When I try to run this I get Memory Error, you can see the Traceback below:

    ---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-32-3e0d3e3c3f13> in <module>
----> 1 A.dot(A.T)

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/base.py in dot(self, other)
    361 
    362         """
--> 363         return self * other
    364 
    365     def power(self, n, dtype=None):

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/base.py in __mul__(self, other)
    479             if self.shape[1] != other.shape[0]:
    480                 raise ValueError('dimension mismatch')
--> 481             return self._mul_sparse_matrix(other)
    482 
    483         # If it's a list or whatever, treat it like a matrix

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/base.py in _mul_sparse_matrix(self, other)
    538 
    539     def _mul_sparse_matrix(self, other):
--> 540         return self.tocsr()._mul_sparse_matrix(other)
    541 
    542     def __rmul__(self, other):  # other * self

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/compressed.py in _mul_sparse_matrix(self, other)
    494 
    495         major_axis = self._swap((M, N))[0]
--> 496         other = self.__class__(other)  # convert to this format
    497 
    498         idx_dtype = get_index_dtype((self.indptr, self.indices,

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/compressed.py in __init__(self, arg1, shape, dtype, copy)
     35                 arg1 = arg1.copy()
     36             else:
---> 37                 arg1 = arg1.asformat(self.format)
     38             self._set_self(arg1)
     39 

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/base.py in asformat(self, format, copy)
    324             # Forward the copy kwarg, if it's accepted.
    325             try:
--> 326                 return convert_method(copy=copy)
    327             except TypeError:
    328                 return convert_method()

~/anaconda3/envs/tfx/lib/python3.7/site-packages/scipy/sparse/coo.py in tocsr(self, copy)
    402             col = self.col.astype(idx_dtype, copy=False)
    403 
--> 404             indptr = np.empty(M + 1, dtype=idx_dtype)
    405             indices = np.empty_like(col, dtype=idx_dtype)
    406             data = np.empty_like(self.data, dtype=upcast(self.dtype))

MemoryError:

The output is a 2 x 2 matrix, so it doesn't matter if it is dense or not. What the program is trying to do to cause this failure at such simple problem? How can I fix this behavior?

Thank you.

Solution

COO is a format that does not lend itself well to math for reasons I will leave to the reader.

In this edge condition case, I think you would be best-served by solving the problem directly:

from scipy.sparse import coo_matrix
import numpy as np

n = 100000000000
row  = np.array([0, 0])
col  = np.array([0, n-1])
data = np.array([1, 1])
A = coo_matrix((data, (row, col)), shape=(2, n))

B = A.tocsr()
C = A.tocsr().T

n, m = B.shape[0], C.shape[1]
out_arr = np.zeros((n, m), dtype=A.dtype)

for i in range(n):
    for j in range(m):
        out_arr[i, j] = B[i, :].multiply(C[:, j].T).data.sum()

For any problem with a reasonably small n & m this workaround will be sufficient.