def test_einsum_random(self):
for _ in range(10): # do 10 random tests
num_arrs = random.randint(3,
5) # use between 3 and 5 arrays as input
arrs = [
sp.random((20, ) * random.randint(1, 5), 0.05)
for _ in range(num_arrs)
]
# pick indices at random with replacement from the first 7 letters of the alphabet
dims = [
''.join(np.random.choice(list("abcdefg"), arr.ndim))
for arr in arrs
]
all_inds = set.union(*(set(inds) for inds in dims))
# of all of the distinct indices that appear in any input,
# pick a random subset of them (of size at most 5) to appear in the output
output = ''.join(
random.sample(all_inds, random.randint(0,
min(len(all_inds), 5))))
specification = ','.join(dims) + '->' + output
with self.subTest(spec=specification):
print(specification)
start = time.perf_counter()
spr = einsum_sparse(specification, *arrs)
mid = time.perf_counter()
der = np.einsum(specification, *[todense(arr) for arr in arrs])
end = time.perf_counter()
print(" sparse: {0}".format(mid - start))
print(" dense: {0}".format(end - mid))
self.assertTrue(np.allclose(todense(spr), der))
def test_transpose_compatible(self):
"""Test that the results of `transpose` are compatible for sparse and dense arrays."""
arr = tocoo(np.arange(27).reshape(3, 3, 3))
np.testing.assert_array_equal(todense(transpose(arr, (1, 2, 0))),
transpose(todense(arr), (1, 2, 0)))
for _ in range(5):
ndims = np.random.randint(2, 6)
dims = tuple(np.random.randint(5, 20, size=ndims))
axes = np.random.permutation(range(ndims))
arr = sp.random(dims, density=0.1)
out1 = todense(transpose(arr, axes))
out2 = transpose(todense(arr), axes)
np.testing.assert_allclose(out1, out2, verbose=True)
def test_einsum_basic(self):
# transpose
arr = sp.random((20, 30, 40), 0.1)
self.assertEqual(
(einsum_sparse('abc->cba', arr) != arr.transpose()).nnz, 0)
# tensordot
arr1 = sp.random((20, 30, 40), 0.1)
arr2 = sp.random((40, 30), 0.1)
arr3 = sp.random((40, 20, 10), 0.1)
self.assertTrue(
np.allclose(
todense(einsum_sparse('abc,cb->a', arr1, arr2)),
todense(sp.tensordot(arr1, arr2, axes=([1, 2], [1, 0])))))
self.assertTrue(
np.allclose(todense(einsum_sparse('ab,acd->bcd', arr2, arr3)),
todense(sp.tensordot(arr2, arr3, axes=(0, 0)))))
# trace
arr = sp.random((100, 100), 0.1)
self.assertAlmostEqual(
einsum_sparse('aa->', arr)[()], np.trace(todense(arr)))