Element-wise maximum of two sparse matrices

This did the trick:

def maximum (A, B):
    BisBigger = A-B
    BisBigger.data = np.where(BisBigger.data < 0, 1, 0)
    return A - A.multiply(BisBigger) + B.multiply(BisBigger)

No, there's no built-in way to do this in scipy.sparse. The easy solution is

np.maximum(X.A, Y.A)

but this is obviously going to be very memory-intensive when the matrices have large dimensions and it might crash your machine. A memory-efficient (but by no means fast) solution is

# convert to COO, if necessary
X = X.tocoo()
Y = Y.tocoo()

Xdict = dict(((i, j), v) for i, j, v in zip(X.row, X.col, X.data))
Ydict = dict(((i, j), v) for i, j, v in zip(Y.row, Y.col, Y.data))

keys = list(set(Xdict.iterkeys()).union(Ydict.iterkeys()))

XmaxY = [max(Xdict.get((i, j), 0), Ydict.get((i, j), 0)) for i, j in keys]
XmaxY = coo_matrix((XmaxY, zip(*keys)))

Note that this uses pure Python instead of vectorized idioms. You can try shaving some of the running time off by vectorizing parts of it.


Here's another memory-efficient solution that should be a bit quicker than larsmans'. It's based on finding the set of unique indices for the nonzero elements in the two arrays using code from Jaime's excellent answer here.

import numpy as np
from scipy import sparse

def sparsemax(X, Y):

    # the indices of all non-zero elements in both arrays
    idx = np.hstack((X.nonzero(), Y.nonzero()))

    # find the set of unique non-zero indices
    idx = tuple(unique_rows(idx.T).T)

    # take the element-wise max over only these indices
    X[idx] = np.maximum(X[idx].A, Y[idx].A)

    return X

def unique_rows(a):
    void_type = np.dtype((np.void, a.dtype.itemsize * a.shape[1]))
    b = np.ascontiguousarray(a).view(void_type)
    idx = np.unique(b, return_index=True)[1]
    return a[idx]

Testing:

def setup(n=1000, fmt='csr'):
    return sparse.rand(n, n, format=fmt), sparse.rand(n, n, format=fmt)

X, Y = setup()
Z = sparsemax(X, Y)
print np.all(Z.A == np.maximum(X.A, Y.A))
# True

%%timeit X, Y = setup()
sparsemax(X, Y)
# 100 loops, best of 3: 4.92 ms per loop