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