def test_large_sum():
n = 500000
x = np.random.randint(0, 10000, size=(n,))
y = np.random.randint(0, 1000, size=(n,))
z = np.random.randint(0, 3, size=(n,))
data = np.random.random(n)
a = COO((x, y, z), data)
assert a.shape == (10000, 1000, 3)
b = a.sum(axis=2)
assert b.nnz > 100000
def normalize_vecs(mat):
"""Normalizes probability vectors so they
sum to one."""
#convert to array for operation
order = len(mat.shape) - 1
mat = COO(mat)
row_sums = mat.sum(axis=order)
mat = DOK(mat)
for point in mat.data:
divisor = row_sums[point[:-1]]
mat[point] = mat[point] / divisor
mat = COO(mat)
return mat
