Re: [Numpy-discussion] Re: Dot product threading?

[Fernando P. <fpe...@gm...>]

On 5/18/06, Bill Baxter <wb...@gm...> wrote:
> In Einstein summation notation, what numpy.dot() does now is:
> c_riqk =3D a_rij * b_qjk
>
> And you want:
> c_[r]ik =3D a_[r]ij * b_[r]jk
>
> where the brackets indicate a 'no summation' index.
> Writing the ESN makes it clearer to me anyway.  :-)

I recently needed something similar to this, and being too lazy to
think up the proper numpy ax-swapping kung-fu, I just opened up weave
and was done with it in a hurry.   Here it is, in case anyone finds
the basic idea of any use.

Cheers,

f

###

class mt3_dot_factory(object):
    """Generate the mt3t contract function, holding necessary state.

    This class only needs to be instantiated once, though it doesn't try to
    enforce this via singleton/borg approaches at all."""

    def __init__(self):
        # The actual expression to contract indices, as a list of strings t=
o be
        # interpolated into the C++ source
        mat_ten =3D ['mat(i,m)*ten(m,j,k)', # first index
                   'mat(j,m)*ten(i,m,k)', # second
                   'mat(k,m)*ten(i,j,m)', # third
                   ]
        # Source template
        code_tpl =3D """
        for (int i=3D0;i<order;i++) {
            for (int j=3D0;j<order;j++) {
                for (int k=3D0;k<order;k++) {
                    double sum=3D0;
                    for (int m=3D0;m<order;m++) {
                        sum +=3D %s;
                    }
                    out(i,j,k) =3D sum;
                }
            }
        }
        """
        self.code =3D [code_tpl % mat_ten[idx] for idx in (0,1,2)]

    def __call__(self,mat,ten,idx):
        """mt3s_contract(mat,ten,idx) -> tensor.

        A special type of matrix-tensor contraction over a single index.  T=
he
        returned array has the following structure:

              out(i,j,k) =3D sum_m(mat(i,m)*ten(m,j,k)) if idx=3D=3D0
              out(i,j,k) =3D sum_m(mat(j,m)*ten(i,m,k)) if idx=3D=3D1
              out(i,j,k) =3D sum_m(mat(k,m)*ten(i,j,m)) if idx=3D=3D2

        Inputs:

         - mat: an NxN array.
         - ten: an NxNxN array.
         - idx: the position of the index to contract over, 0 1 or 2."""

        # Minimal input validation - we use asserts so they don't kick in
        # under a -O run of python.
        assert len(mat.shape)=3D=3D2,\
               "mat must be a rank 2 array, shape=3D%s" % mat.shape
        assert mat.shape[0]=3D=3Dmat.shape[1],\
               "Only square matrices are supported: mat shape=3D%s" % mat.s=
hape
        assert len(ten.shape)=3D=3D3,\
               "mat must be a rank 3 array, shape=3D%s" % ten.shape
        assert ten.shape[0]=3D=3Dten.shape[1]=3D=3Dten.shape[2],\
               "Only equal-dim tensors are supported: ten shape=3D%s" % ten=
.shape

        order =3D mat.shape[0]
        out =3D zeros_like(ten)
        inline(self.code[idx],('mat','ten','out','order'),
               type_converters =3D converters.blitz)
        return out

# Make actual instance
mt3_dot =3D mt3_dot_factory()