def test_einsum_misc(self):
# This call used to crash because of a bug in
# PyArray_AssignZero
a = np.ones((1,2))
b = np.ones((2,2,1))
assert_equal(myeinsum('ij...,j...->i...',a,b), [[[2],[2]]])
# The iterator had an issue with buffering this reduction
a = np.ones((5, 12, 4, 2, 3), np.int64)
b = np.ones((5, 12, 11), np.int64)
assert_equal(myeinsum('ijklm,ijn,ijn->',a,b,b),
myeinsum('ijklm,ijn->',a,b))
# Issue #2027, was a problem in the contiguous 3-argument
# inner loop implementation
a = np.arange(1, 3)
b = np.arange(1, 5).reshape(2, 2)
c = np.arange(1, 9).reshape(4, 2)
assert_equal(myeinsum('x,yx,zx->xzy', a, b, c),
[[[1, 3], [3, 9], [5, 15], [7, 21]],
[[8, 16], [16, 32], [24, 48], [32, 64]]])
def test_einsum_views(self):
# pass-through
a = np.arange(6)
a.shape = (2,3)
b = myeinsum("...", a)
#assert_(b.base is a)
assert_equal(b,a)
##b = myeinsum(a, [Ellipsis])
##assert_(b.base is a)
b = myeinsum("ij", a)
#assert_(b.base is a)
assert_equal(b, a)
##b = myeinsum(a, [0,1])
##assert_(b.base is a)
##assert_equal(b, a)
# transpose
a = np.arange(6)
a.shape = (2,3)
b = myeinsum("ji", a)
#assert_(b.base is a)
assert_equal(b, a.T)
##b = myeinsum(a, [1,0])
##assert_(b.base is a)
##assert_equal(b, a.T)
# diagonal
a = np.arange(9)
a.shape = (3,3)
b = myeinsum("ii->i", a)
##assert_(b.base is a)
assert_equal(b, [a[i,i] for i in range(3)])
##b = myeinsum(a, [0,0], [0])
##assert_(b.base is a)
##assert_equal(b, [a[i,i] for i in range(3)])
# diagonal with various ways of broadcasting an additional dimension
a = np.arange(27)
a.shape = (3,3,3)
b = myeinsum("...ii->...i", a)
#assert_(b.base is a)
assert_equal(b, [[x[i,i] for i in range(3)] for x in a])
##b = myeinsum(a, [Ellipsis,0,0], [Ellipsis,0])
##assert_(b.base is a)
##assert_equal(b, [[x[i,i] for i in range(3)] for x in a])
b = myeinsum("ii...->...i", a)
#assert_(b.base is a)
assert_equal(b, [[x[i,i] for i in range(3)]
for x in a.transpose(2,0,1)])
##b = myeinsum(a, [0,0,Ellipsis], [Ellipsis,0])
##assert_(b.base is a)
##assert_equal(b, [[x[i,i] for i in range(3)]
## for x in a.transpose(2,0,1)])
b = myeinsum("...ii->i...", a)
#assert_(b.base is a)
assert_equal(b, [a[:,i,i] for i in range(3)])
##b = myeinsum(a, [Ellipsis,0,0], [0,Ellipsis])
##assert_(b.base is a)
##assert_equal(b, [a[:,i,i] for i in range(3)])
b = myeinsum("jii->ij", a)
#assert_(b.base is a)
assert_equal(b, [a[:,i,i] for i in range(3)])
##b = myeinsum(a, [1,0,0], [0,1])
##assert_(b.base is a)
##assert_equal(b, [a[:,i,i] for i in range(3)])
b = myeinsum("ii...->i...", a)
#assert_(b.base is a)
assert_equal(b, [a.transpose(2,0,1)[:,i,i] for i in range(3)])
##b = myeinsum(a, [0,0,Ellipsis], [0,Ellipsis])
##assert_(b.base is a)
##assert_equal(b, [a.transpose(2,0,1)[:,i,i] for i in range(3)])
b = myeinsum("i...i->i...", a)
#assert_(b.base is a)
assert_equal(b, [a.transpose(1,0,2)[:,i,i] for i in range(3)])
##b = myeinsum(a, [0,Ellipsis,0], [0,Ellipsis])
##assert_(b.base is a)
##assert_equal(b, [a.transpose(1,0,2)[:,i,i] for i in range(3)])
b = myeinsum("i...i->...i", a)
#assert_(b.base is a)
assert_equal(b, [[x[i,i] for i in range(3)]
for x in a.transpose(1,0,2)])
##b = myeinsum(a, [0,Ellipsis,0], [Ellipsis,0])
##assert_(b.base is a)
##assert_equal(b, [[x[i,i] for i in range(3)]
## for x in a.transpose(1,0,2)])
# triple diagonal
a = np.arange(27)
a.shape = (3,3,3)
b = myeinsum("iii->i", a)
#assert_(b.base is a)
assert_equal(b, [a[i,i,i] for i in range(3)])
##b = myeinsum(a, [0,0,0], [0])
##assert_(b.base is a)
##assert_equal(b, [a[i,i,i] for i in range(3)])
# swap axes
a = np.arange(24)
a.shape = (2,3,4)
b = myeinsum("ijk->jik", a)
#assert_(b.base is a)
assert_equal(b, a.swapaxes(0,1))
def check_einsum_sums(self, dtype):
# Check various sums. Does many sizes to exercise unrolled loops.
# sum(a, axis=-1)
for n in range(1,17):
a = np.arange(n, dtype=dtype)
assert_equal(myeinsum("i->", a), np.sum(a, axis=-1).astype(dtype))
##assert_equal(myeinsum(a, [0], []),
## np.sum(a, axis=-1).astype(dtype))
for n in range(1,17):
a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
assert_equal(myeinsum("...i->...", a),
np.sum(a, axis=-1).astype(dtype))
##assert_equal(myeinsum(a, [Ellipsis,0], [Ellipsis]),
## np.sum(a, axis=-1).astype(dtype))
# sum(a, axis=0)
for n in range(1,17):
a = np.arange(2*n, dtype=dtype).reshape(2,n)
assert_equal(myeinsum("i...->...", a),
np.sum(a, axis=0).astype(dtype))
##assert_equal(myeinsum(a, [0,Ellipsis], [Ellipsis]),
## np.sum(a, axis=0).astype(dtype))
for n in range(1,17):
a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
assert_equal(myeinsum("i...->...", a),
np.sum(a, axis=0).astype(dtype))
##assert_equal(myeinsum(a, [0,Ellipsis], [Ellipsis]),
## np.sum(a, axis=0).astype(dtype))
# trace(a)
for n in range(1,17):
a = np.arange(n*n, dtype=dtype).reshape(n,n)
assert_equal(myeinsum("ii", a), np.trace(a).astype(dtype))
##assert_equal(myeinsum(a, [0,0]), np.trace(a).astype(dtype))
# multiply(a, b)
for n in range(1,17):
a = np.arange(3*n, dtype=dtype).reshape(3,n)
b = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
assert_equal(myeinsum("..., ...", a, b), np.multiply(a, b))
##assert_equal(myeinsum(a, [Ellipsis], b, [Ellipsis]),
## np.multiply(a, b))
# inner(a,b)
for n in range(1,17):
a = np.arange(2*3*n, dtype=dtype).reshape(2,3,n)
b = np.arange(n, dtype=dtype)
assert_equal(myeinsum("...i, ...i", a, b), np.inner(a, b))
##assert_equal(myeinsum(a, [Ellipsis,0], b, [Ellipsis,0]),
## np.inner(a, b))
for n in range(1,11):
a = np.arange(n*3*2, dtype=dtype).reshape(n,3,2)
b = np.arange(n, dtype=dtype)
assert_equal(myeinsum("i..., i...", a, b), np.inner(a.T, b.T).T)
##assert_equal(myeinsum(a, [0,Ellipsis], b, [0,Ellipsis]),
## np.inner(a.T, b.T).T)
# outer(a,b)
for n in range(1,17):
a = np.arange(3, dtype=dtype)+1
b = np.arange(n, dtype=dtype)+1
assert_equal(myeinsum("i,j", a, b), np.outer(a, b))
##assert_equal(myeinsum(a, [0], b, [1]), np.outer(a, b))
# Suppress the complex warnings for the 'as f8' tests
ctx = WarningManager()
ctx.__enter__()
try:
warnings.simplefilter('ignore', np.ComplexWarning)
# matvec(a,b) / a.dot(b) where a is matrix, b is vector
for n in range(1,17):
a = np.arange(4*n, dtype=dtype).reshape(4,n)
b = np.arange(n, dtype=dtype)
assert_equal(myeinsum("ij, j", a, b), np.dot(a, b))
##assert_equal(myeinsum(a, [0,1], b, [1]), np.dot(a, b))
c = np.arange(4, dtype=dtype)
myeinsum("ij,j", a, b, out=c,
dtype='f8', casting='unsafe')
assert_equal(c,
np.dot(a.astype('f8'),
b.astype('f8')).astype(dtype))
##c[...] = 0
##myeinsum(a, [0,1], b, [1], out=c,
## dtype='f8', casting='unsafe')
##assert_equal(c,
## np.dot(a.astype('f8'),
## b.astype('f8')).astype(dtype))
for n in range(1,17):
a = np.arange(4*n, dtype=dtype).reshape(4,n)
b = np.arange(n, dtype=dtype)
assert_equal(myeinsum("ji,j", a.T, b.T), np.dot(b.T, a.T))
##assert_equal(myeinsum(a.T, [1,0], b.T, [1]), np.dot(b.T, a.T))
c = np.arange(4, dtype=dtype)
myeinsum("ji,j", a.T, b.T, out=c, dtype='f8', casting='unsafe')
assert_equal(c,
np.dot(b.T.astype('f8'),
a.T.astype('f8')).astype(dtype))
##c[...] = 0
##myeinsum(a.T, [1,0], b.T, [1], out=c,
## dtype='f8', casting='unsafe')
##assert_equal(c,
## np.dot(b.T.astype('f8'),
## a.T.astype('f8')).astype(dtype))
# matmat(a,b) / a.dot(b) where a is matrix, b is matrix
for n in range(1,17):
if n < 8 or dtype != 'f2':
a = np.arange(4*n, dtype=dtype).reshape(4,n)
b = np.arange(n*6, dtype=dtype).reshape(n,6)
assert_equal(myeinsum("ij,jk", a, b), np.dot(a, b))
##assert_equal(myeinsum(a, [0,1], b, [1,2]), np.dot(a, b))
for n in range(1,17):
a = np.arange(4*n, dtype=dtype).reshape(4,n)
b = np.arange(n*6, dtype=dtype).reshape(n,6)
c = np.arange(24, dtype=dtype).reshape(4,6)
myeinsum("ij,jk", a, b, out=c, dtype='f8', casting='unsafe')
assert_equal(c,
np.dot(a.astype('f8'),
b.astype('f8')).astype(dtype))
##c[...] = 0
##myeinsum(a, [0,1], b, [1,2], out=c,
## dtype='f8', casting='unsafe')
##assert_equal(c,
## np.dot(a.astype('f8'),
## b.astype('f8')).astype(dtype))
# matrix triple product (note this is not currently an efficient
# way to multiply 3 matrices)
a = np.arange(12, dtype=dtype).reshape(3,4)
b = np.arange(20, dtype=dtype).reshape(4,5)
c = np.arange(30, dtype=dtype).reshape(5,6)
if dtype != 'f2':
assert_equal(myeinsum("ij,jk,kl", a, b, c),
a.dot(b).dot(c))
##assert_equal(myeinsum(a, [0,1], b, [1,2], c, [2,3]),
## a.dot(b).dot(c))
d = np.arange(18, dtype=dtype).reshape(3,6)
myeinsum("ij,jk,kl", a, b, c, out=d,
dtype='f8', casting='unsafe')
assert_equal(d, a.astype('f8').dot(b.astype('f8')
).dot(c.astype('f8')).astype(dtype))
##d[...] = 0
##myeinsum(a, [0,1], b, [1,2], c, [2,3], out=d,
## dtype='f8', casting='unsafe')
##assert_equal(d, a.astype('f8').dot(b.astype('f8')
## ).dot(c.astype('f8')).astype(dtype))
# tensordot(a, b)
if np.dtype(dtype) != np.dtype('f2'):
a = np.arange(60, dtype=dtype).reshape(3,4,5)
b = np.arange(24, dtype=dtype).reshape(4,3,2)
assert_equal(myeinsum("ijk, jil -> kl", a, b),
np.tensordot(a,b, axes=([1,0],[0,1])))
##assert_equal(myeinsum(a, [0,1,2], b, [1,0,3], [2,3]),
## np.tensordot(a,b, axes=([1,0],[0,1])))
c = np.arange(10, dtype=dtype).reshape(5,2)
myeinsum("ijk,jil->kl", a, b, out=c,
dtype='f8', casting='unsafe')
assert_equal(c, np.tensordot(a.astype('f8'), b.astype('f8'),
axes=([1,0],[0,1])).astype(dtype))
##c[...] = 0
##myeinsum(a, [0,1,2], b, [1,0,3], [2,3], out=c,
## dtype='f8', casting='unsafe')
##assert_equal(c, np.tensordot(a.astype('f8'), b.astype('f8'),
## axes=([1,0],[0,1])).astype(dtype))
finally:
ctx.__exit__()
"""
# logical_and(logical_and(a!=0, b!=0), c!=0)
a = np.array([1, 3, -2, 0, 12, 13, 0, 1], dtype=dtype)
b = np.array([0, 3.5, 0., -2, 0, 1, 3, 12], dtype=dtype)
c = np.array([True,True,False,True,True,False,True,True])
assert_equal(myeinsum("i,i,i->i", a, b, c,
dtype='?', casting='unsafe'),
np.logical_and(np.logical_and(a!=0, b!=0), c!=0))
##assert_equal(myeinsum(a, [0], b, [0], c, [0], [0],
## dtype='?', casting='unsafe'),
## np.logical_and(np.logical_and(a!=0, b!=0), c!=0))
"""
a = np.arange(9, dtype=dtype)
assert_equal(myeinsum(",i->", 3, a), 3*np.sum(a))
##assert_equal(myeinsum(3, [], a, [0], []), 3*np.sum(a))
assert_equal(myeinsum("i,->", a, 3), 3*np.sum(a))
##assert_equal(myeinsum(a, [0], 3, [], []), 3*np.sum(a))
# Various stride0, contiguous, and SSE aligned variants
for n in range(1,25):
a = np.arange(n, dtype=dtype)
if np.dtype(dtype).itemsize > 1:
assert_equal(myeinsum("...,...",a,a), np.multiply(a,a))
assert_equal(myeinsum("i,i", a, a), np.dot(a,a))
assert_equal(myeinsum("i,->i", a, 2), 2*a)
assert_equal(myeinsum(",i->i", 2, a), 2*a)
assert_equal(myeinsum("i,->", a, 2), 2*np.sum(a))
assert_equal(myeinsum(",i->", 2, a), 2*np.sum(a))
assert_equal(myeinsum("...,...",a[1:],a[:-1]),
np.multiply(a[1:],a[:-1]))
assert_equal(myeinsum("i,i", a[1:], a[:-1]),
np.dot(a[1:],a[:-1]))
assert_equal(myeinsum("i,->i", a[1:], 2), 2*a[1:])
assert_equal(myeinsum(",i->i", 2, a[1:]), 2*a[1:])
assert_equal(myeinsum("i,->", a[1:], 2), 2*np.sum(a[1:]))
assert_equal(myeinsum(",i->", 2, a[1:]), 2*np.sum(a[1:]))
"""
# An object array, summed as the data type
a = np.arange(9, dtype=object)
b = myeinsum("i->", a, dtype=dtype, casting='unsafe')
assert_equal(b, np.sum(a))
assert_equal(b.dtype, np.dtype(dtype))
##b = myeinsum(a, [0], [], dtype=dtype, casting='unsafe')
##assert_equal(b, np.sum(a))
##assert_equal(b.dtype, np.dtype(dtype))
"""
# A case which was failing (ticket #1885)
p = np.arange(2) + 1
q = np.arange(4).reshape(2,2) + 3
r = np.arange(4).reshape(2,2) + 7
assert_equal(myeinsum('z,mz,zm->', p, q, r), 253)
