Optimize sparse COO matrix sum

[liuzhenqi77]

Summary

When trying to speed up MKDAChi2(), I noticed COO matrix sum takes a lot of time.

Additional details

Examples here and here. It seems to me that the implementation of sum in sparse only uses np.add.reduce without any optimization (I opened an issue there waiting for response).

Next steps

I tried some simple optimization like below using numba, which makes it much faster. (There are many better ways to implement this. Also maybe there are ways to avoid calculating the sum.)

from numba import njit

@njit
def _coo_sum_axis_0(shape, coords, data):
    res = np.zeros(shape)
    for i in range(len(data)):
        x, y, z = coords[1:, i]
        res[x, y, z] += data[i]
    return res

# _coo_sum_axis_0(ma_maps.shape[1:], ma_maps.coords, ma_maps.data)

[JulioAPeraza]

Thank you for noticing this issue, @liuzhenqi77. We think that other functions may be suffering from the same issue as well, for example

NiMARE/nimare/meta/cbma/ale.py

Lines 178 to 179 in c5f18ee

	# Need to convert to dense because np.ceil is too slow with sparse
	max_ma_values = ma_maps.max(axis=[1, 2, 3]).todense()

where max also uses np.maximum.reduce.

We will submit a PR with a solution to this soon.
Thanks!

[JulioAPeraza]

I think we can close this since #857 addressed this performance issue.