Compute outer product of arrays with arbitrary dimensions

arrays, numpy, python, scipy, slice

Solution

After discovering the use of ellipsis in numpy/scipy arrays I ended up implementing it as a recursive function:

def array_outer_product(A, B, result=None):
    ''' Compute the outer-product in the final two dimensions of the given arrays.
    If the result array is provided, the results are written into it.
    '''
    assert(A.shape[:-1] == B.shape[:-1])
    if result is None:
        result=scipy.zeros(A.shape+B.shape[-1:], dtype=A.dtype)
    if A.ndim==1:
        result[:,:]=scipy.outer(A, B)
    else:
        for idx in xrange(A.shape[0]):
            array_outer_product(A[idx,...], B[idx,...], result[idx,...])
    return result

Problem

I have two arrays `A,B` and want to take the outer product on their last dimension, e.g. `result[:,i,j]=A[:,i]*B[:,j]` when `A,B` are 2-dimensional. How can I do this if I don't know whether they will be 2 or 3 dimensional? In my specific problem `A,B` are slices out of a bigger 3-dimensional array `Z`, Sometimes this may be called with integer indices `A=Z[:,1,:], B=Z[:,2,:]` and other times with slices `A=Z[:,1:3,:],B=Z[:,4:6,:]`. Since scipy "squeezes" singleton dimensions, I won't know what dimensions my inputs will be. The array-outer-product I'm trying to define should satisfy ``` array_outer_product( Y[a,b,:], Z[i,j,:] ) == scipy.outer( Y[a,b,:], Z[i,j,:] ) array_outer_product( Y[a:a+N,b,:], Z[i:i+N,j,:])[n,:,:] == scipy.outer( Y[a+n,b,:], Z[i+n,j,:] ) array_outer_product( Y[a:a+N,b:b+M,:], Z[i:i+N, j:j+M,:] )[n,m,:,:]==scipy.outer( Y[a+n,b+m,:] , Z[i+n,j+m,:] ) ``` for any rank-3 arrays `Y,Z` and integers `a,b,...i,j,k...n,N,...` The kind of problem I'm dealing with involves a 2-D spatial grid, with a vector-valued function at each grid point. I want to be able to calculate the covariance matrix (outer product) of these vectors, over regions defined by slices in the first two axes.

Original source