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_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)))
def test_einsum_errors(self):
# number of inputs in specification must match number of inputs
with self.assertRaises(Exception):
einsum_sparse('abc,def->ad', tocoo(np.ones((1, 2, 3))))
with self.assertRaises(Exception):
einsum_sparse('abc->a', tocoo(np.ones((1, 2, 3))),
tocoo(np.ones((1, 2, 3))))
# must have an output
with self.assertRaises(Exception):
einsum_sparse('abc', tocoo(np.ones((1, 2, 3))))
# output indices must be unique
with self.assertRaises(Exception):
einsum_sparse('abc->bb', tocoo(np.ones((1, 2, 3))))
# output indices must be present in an input
with self.assertRaises(Exception):
einsum_sparse('abc->bd', tocoo(np.ones((1, 2, 3))))
# number of indices must match number of dimensions for each input
with self.assertRaises(Exception):
einsum_sparse('ab->a', tocoo(np.ones((1, 2, 3))))
with self.assertRaises(Exception):
einsum_sparse('abcd->a', tocoo(np.ones((1, 2, 3))),
tocoo(np.ones((1, 2, 3))))
# repeated indices must always have consistent sizes
with self.assertRaises(Exception):
einsum_sparse('aaa->a', tocoo(np.ones((1, 2, 3))))
with self.assertRaises(Exception):
einsum_sparse('abc,bac->a', tocoo(np.ones((1, 2, 3))),
tocoo(np.ones((1, 2, 3))))