Efficient tensor contraction in python

Question

I have a list L of tensors (ndarray objects), with several indices each. I need to contract these indices according to a graph of connections.

The connections are encoded in a list of tuples in the form ((m,i),(n,j)) signifying "contract the i-th index of the tensor L[m] with the j-th index of the tensor L[n].

How can I handle non-trivial connectivity graphs? The first problem is that as soon as I contract a pair of indices, the result is a new tensor that does not belong to the list L. But even if I solved this (e.g. by giving a unique identifier to all the indices of all the tensors), there is the issue that one can pick any order to perform the contractions, and some choices yield unnecessarily enormous beasts in mid-computation (even if the final result is small). Suggestions?

Answer 1

Maybe helpful: Take a look into https://arxiv.org/abs/1402.0939 which is a description of an efficient framework for the problem of contracting so called tensor networks in a single function ncon(...). As far as I see implementations of it are directly available for Matlab (can be found within in the link) and for Python3 (https://github.com/mhauru/ncon).

Answer 2

Memory considerations aside, I believe you can do the contractions in a single call to einsum, although you'll need some preprocessing. I'm not entirely sure what you mean by "as I contract a pair of indices, the result is a new tensor that does not belong to the list L", but I think doing the contraction in a single step would exactly solve this problem.

I suggest using the alternative, numerically indexed syntax of einsum:

einsum(op0, sublist0, op1, sublist1, ..., [sublistout])

So what you need to do is encode the indices to contract in integer sequences. First you'll need to set up a range of unique indices initially, and keep another copy to be used as sublistout. Then, iterating over your connectivity graph, you need to set contracted indices to the same index where necessary, and at the same time remove the contracted index from sublistout.

import numpy as np

def contract_all(tensors,conns):
    '''
    Contract the tensors inside the list tensors
    according to the connectivities in conns

    Example input:
    tensors = [np.random.rand(2,3),np.random.rand(3,4,5),np.random.rand(3,4)]
    conns = [((0,1),(2,0)), ((1,1),(2,1))]
    returned shape in this case is (2,3,5)
    '''

    ndims = [t.ndim for t in tensors]
    totdims = sum(ndims)
    dims0 = np.arange(totdims)
    # keep track of sublistout throughout
    sublistout = set(dims0.tolist())
    # cut up the index array according to tensors
    # (throw away empty list at the end)
    inds = np.split(dims0,np.cumsum(ndims))[:-1]
    # we also need to convert to a list, otherwise einsum chokes
    inds = [ind.tolist() for ind in inds]

    # if there were no contractions, we'd call
    # np.einsum(*zip(tensors,inds),sublistout)

    # instead we need to loop over the connectivity graph
    # and manipulate the indices
    for (m,i),(n,j) in conns:
        # tensors[m][i] contracted with tensors[n][j]

        # remove the old indices from sublistout which is a set
        sublistout -= {inds[m][i],inds[n][j]}

        # contract the indices
        inds[n][j] = inds[m][i]

    # zip and flatten the tensors and indices
    args = [subarg for arg in zip(tensors,inds) for subarg in arg]

    # assuming there are no multiple contractions, we're done here
    return np.einsum(*args,sublistout)

A trivial example:

>>> tensors = [np.random.rand(2,3), np.random.rand(4,3)]
>>> conns = [((0,1),(1,1))]
>>> contract_all(tensors,conns)
array([[ 1.51970003,  1.06482209,  1.61478989,  1.86329518],
       [ 1.16334367,  0.60125945,  1.00275992,  1.43578448]])
>>> np.einsum('ij,kj',tensors[0],tensors[1])
array([[ 1.51970003,  1.06482209,  1.61478989,  1.86329518],
       [ 1.16334367,  0.60125945,  1.00275992,  1.43578448]])

In case there are multiple contractions, the logistics in the loop becomes a bit more complex, because we need to handle all the duplicates. The logic, however, is the same. Furthermore, the above is obviously missing checks to ensure that the corresponding indices can be contracted.

In hindsight I realized that the default sublistout doesn't have to be specified, einsum uses that order anyway. I decided to leave that variable in the code, because in case we want a non-trivial output index order, we'll have to handle that variable appropriately, and it might come handy.

---

As for optimization of the contraction order, you can effect internal optimization in np.einsum as of version 1.12 (as noted by @hpaulj in a now-deleted comment). This version introduced the optimize optional keyword argument to np.einsum, allowing to choose a contraction order that cuts down on computational time at the cost of memory. Passing 'greedy' or 'optimal' as the optimize keyword will make numpy choose a contraction order in roughly decreasing order of sizes of the dimensions.

The options available for the optimize keyword come from the apparently undocumented (as far as online documentation goes; help() fortunately works) function np.einsum_path:

einsum_path(subscripts, *operands, optimize='greedy')

Evaluates the lowest cost contraction order for an einsum expression by
considering the creation of intermediate arrays.

The output contraction path from np.einsum_path can also be used as an input for the optimize argument of np.einsum. In your question you were worried about too much memory being used, so I suspect that the default of no optimization (with potentially longer runtime and smaller memory footprint).