test_einsum.py

import sys
from decimal import Decimal

import numpy as np
from numpy.testing import *
from numpy.testing.utils import WarningManager
import warnings

myeinsum = np.einsum
from einsum_py import myeinsum
from unittest import skip

class TestEinSum(TestCase):
    @skip("skipping valueerrors")
    def test_einsum_errors(self):
        # Need enough arguments
        ##assert_raises(ValueError, myeinsum)
        assert_raises(ValueError, myeinsum, "")

        # subscripts must be a string
        ##assert_raises(TypeError, myeinsum, 0, 0)

        # out parameter must be an array
        ##assert_raises(TypeError, myeinsum, "", 0, out='test')

        # order parameter must be a valid order
        ##assert_raises(TypeError, myeinsum, "", 0, order='W')

        # casting parameter must be a valid casting
        ##assert_raises(ValueError, myeinsum, "", 0, casting='blah')

        # dtype parameter must be a valid dtype
        ##assert_raises(TypeError, myeinsum, "", 0, dtype='bad_data_type')

        # other keyword arguments are rejected
        ##assert_raises(TypeError, myeinsum, "", 0, bad_arg=0)

        # number of operands must match count in subscripts string
        ##assert_raises(ValueError, myeinsum, "", 0, 0)
        ##assert_raises(ValueError, myeinsum, ",", 0, [0], [0])
        ##assert_raises(ValueError, myeinsum, ",", [0])

        # can't have more subscripts than dimensions in the operand
        ##assert_raises(ValueError, myeinsum, "i", 0)
        assert_raises(ValueError, myeinsum, "ij", [0,0])
        ##assert_raises(ValueError, myeinsum, "...i", 0)
        assert_raises(ValueError, myeinsum, "i...j", [0,0])
        ##assert_raises(ValueError, myeinsum, "i...", 0)
        assert_raises(ValueError, myeinsum, "ij...", [0,0])

        # invalid ellipsis
        assert_raises(ValueError, myeinsum, "i..", [0,0])
        assert_raises(ValueError, myeinsum, ".i...", [0,0])
        assert_raises(ValueError, myeinsum, "j->..j", [0,0])
        assert_raises(ValueError, myeinsum, "j->.j...", [0,0])

        # invalid subscript character
        assert_raises(ValueError, myeinsum, "i%...", [0,0])
        assert_raises(ValueError, myeinsum, "...j$", [0,0])
        assert_raises(ValueError, myeinsum, "i->&", [0,0])

        # output subscripts must appear in input
        assert_raises(ValueError, myeinsum, "i->ij", [0,0])

        # output subscripts may only be specified once
        assert_raises(ValueError, myeinsum, "ij->jij", [[0,0],[0,0]])

        # dimensions much match when being collapsed
        assert_raises(ValueError, myeinsum, "ii", np.arange(6).reshape(2,3))
        assert_raises(ValueError, myeinsum, "ii->i", np.arange(6).reshape(2,3))

        # broadcasting to new dimensions must be enabled explicitly
        assert_raises(ValueError, myeinsum, "i", np.arange(6).reshape(2,3))
        assert_raises(ValueError, myeinsum, "i->i", [[0,1],[0,1]],
                                            out=np.arange(4).reshape(2,2))

    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))

        ##b = myeinsum(a, [0,1,2], [1,0,2])
        ##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)
    """
    casting issues with out=c
    def test_einsum_sums_int8(self):
        self.check_einsum_sums('i1');

    def test_einsum_sums_uint8(self):
        self.check_einsum_sums('u1');

    def test_einsum_sums_int16(self):
        self.check_einsum_sums('i2');

    def test_einsum_sums_uint16(self):
        self.check_einsum_sums('u2');
    """
    def test_einsum_sums_int32(self):
        self.check_einsum_sums('i4');
    """
    def test_einsum_sums_uint32(self):
        self.check_einsum_sums('u4');

    def test_einsum_sums_int64(self):
        self.check_einsum_sums('i8');

    def test_einsum_sums_uint64(self):
        self.check_einsum_sums('u8');

    def test_einsum_sums_float16(self):
        self.check_einsum_sums('f2');

    def test_einsum_sums_float32(self):
        self.check_einsum_sums('f4');
    """
    def test_einsum_sums_float64(self):
        self.check_einsum_sums('f8');

    def test_einsum_sums_longdouble(self):
        self.check_einsum_sums(np.longdouble);
    """
    def test_einsum_sums_cfloat64(self):
        self.check_einsum_sums('c8');

    def test_einsum_sums_cfloat128(self):
        self.check_einsum_sums('c16');

    def test_einsum_sums_clongdouble(self):
        self.check_einsum_sums(np.clongdouble);
    """
    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]]])

if __name__ == "__main__":
    run_module_suite()