# [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()