[Numpy-discussion] Tensor contraction

Pauli Virtanen pav@iki...
Sat Jun 12 18:12:58 CDT 2010

Sat, 12 Jun 2010 16:30:14 -0400, Alan Bromborsky wrote:
> If I have a single numpy array, for example with 3 indices T_{ijk} and I
> want to sum over two them in the sense of tensor contraction -
> T_{k} = \sum_{i=0}^{n-1} T_{iik}.  Is there an easy way to do this with
> numpy?

HTH, (not really tested, so you get what you pay for) 

def tensor_contraction_single(tensor, dimensions):
    """Perform a single tensor contraction over the dimensions given"""
    swap = [x for x in range(tensor.ndim) 
            if x not in dimensions] + list(dimensions)
    x = tensor.transpose(swap)
    for k in range(len(dimensions) - 1):
        x = np.diagonal(x, axis1=-2, axis2=-1)
    return x.sum(axis=-1)

def _preserve_indices(indices, removed):
    """Adjust values of indices after some items are removed"""
    for r in reversed(sorted(removed)):
        indices = [j if j <= r else j-1 for j in indices]
    return indices

def tensor_contraction(tensor, contractions):
    """Perform several tensor contractions"""
    while contractions:
        dimensions = contractions.pop(0)
        tensor = tensor_contraction_single(tensor, dimensions)
        contractions = [_preserve_indices(c, dimensions) 
                        for c in contractions]
    return tensor

# b_{km} = sum_{ij} a_{i,j,k,j,i,m}

a = np.random.rand(3,2,4,2,3,5)
b = tensor_contraction(a, [(0,4), (1,3)])

b2 = np.zeros((4, 5))
for i in range(3):
    for j in range(2):
        b2 += a[i,j,:,j,i,:]

assert (abs(b - b2) < 1e-12).all()

More information about the NumPy-Discussion mailing list