OMEinsumContractionOrders

This is the documentation for OMEinsumContractionOrders, a Julia package for the optimization of the contraction order of tensor networks.

Installation guide is available in README.md. You can also access its features in OMEinsum, which uses it as the default contraction order optimizer.

Example 1: Use it directly

The contraction order optimizer is implemented in the optimize_code function. It takes three arguments: code, size, and optimizer. The code argument is the einsum notation to be optimized. The size argument is the size of the variables in the einsum notation. The optimizer argument is the optimizer to be used. The optimize_code function returns the optimized contraction order. One can use contraction_complexity function to get the time, space and rewrite complexity of returned contraction order. Supported solvers include:

The KaHyParBipartite is implemented as an extension. If you have issues in installing KaHyPar, please check these issues: #12 and #19. Additionally, code simplifiers can be used to preprocess the tensor network to reduce the optimization time:

More details could be found in this blog post.

julia> using OMEinsumContractionOrders, Graphs, KaHyPar
julia> function random_regular_eincode(n, k; optimize=nothing)
       	g = Graphs.random_regular_graph(n, k)
       	ixs = [[minmax(e.src,e.dst)...] for e in Graphs.edges(g)]
       	return OMEinsumContractionOrders.EinCode([ixs..., [[i] for i in Graphs.vertices(g)]...], Int[])
       endrandom_regular_eincode (generic function with 1 method)
julia> code = random_regular_eincode(100, 3);
julia> contraction_complexity(code, uniformsize(code, 2))Time complexity: 2^100.0
Space complexity: 2^0.0
Read-write complexity: 2^9.645658432408709
julia> optcode_tree = optimize_code(code, uniformsize(code, 2),
       	TreeSA(sc_target=28, βs=0.1:0.1:10, ntrials=2, niters=100, sc_weight=3.0));
julia> contraction_complexity(optcode_tree, uniformsize(code, 2))Time complexity: 2^18.15720875469124
Space complexity: 2^14.0
Read-write complexity: 2^16.795139164888298
julia> optcode_tree_with_slice = optimize_code(code, uniformsize(code, 2),
       	TreeSA(sc_target=28, βs=0.1:0.1:10, ntrials=2, niters=20, sc_weight=3.0, nslices=5));
julia> optcode_kahypar = optimize_code(code, uniformsize(code, 2),
       	KaHyParBipartite(sc_target=30, max_group_size=50));
julia> contraction_complexity(optcode_kahypar, uniformsize(code, 2))Time complexity: 2^21.783443113504337
Space complexity: 2^16.0
Read-write complexity: 2^18.043983136945002
julia> optcode_sa = optimize_code(code, uniformsize(code, 2),
       	SABipartite(sc_target=30, max_group_size=50));
julia> contraction_complexity(optcode_sa, uniformsize(code, 2))Time complexity: 2^20.30878007163327
Space complexity: 2^14.0
Read-write complexity: 2^17.630686617500075

Example 2: Use it in OMEinsum

OMEinsumContractionOrders is shipped with OMEinsum package. You can use it to optimize the contraction order of an OMEinsum expression.

julia> using OMEinsum
julia> code = ein"ij, jk, kl, il->"ij, jk, kl, il ->
julia> optimized_code = optimize_code(code, uniformsize(code, 2), TreeSA())OMEinsum.SlicedEinsum{Char, OMEinsum.DynamicNestedEinsum{Char}}(Char[], il, il ->
├─ ik, kl -> il
│  ├─ ij, jk -> ik
│  │  ├─ ij
│  │  └─ jk
│  └─ kl
└─ il
)

For multi-GPU contraction of tensor networks, please check this Gist.

Example 3: Visualization

LuxorTensorPlot

LuxorTensorPlot is an extension of the OMEinsumContractionOrders package that provides a visualization of the contraction order. It is designed to work with the OMEinsumContractionOrders package. To use LuxorTensorPlot, please follow these steps:

pkg> add OMEinsumContractionOrders, LuxorGraphPlot

julia> using OMEinsumContractionOrders, LuxorGraphPlot

and then the extension will be loaded automatically.

The extension provides the following to function, viz_eins and viz_contraction, where the former will plot the tensor network as a graph, and the latter will generate a video or gif of the contraction process. Here is an example:

julia> using OMEinsumContractionOrders, LuxorGraphPlot

julia> eincode = OMEinsumContractionOrders.EinCode([['a', 'b'], ['a', 'c', 'd'], ['b', 'c', 'e', 'f'], ['e'], ['d', 'f']], ['a'])
ab, acd, bcef, e, df -> a

julia> viz_eins(eincode, filename = "eins.png")

julia> nested_eins = optimize_code(eincode, uniformsize(eincode, 2), GreedyMethod())
ab, ab -> a
├─ ab
└─ acf, bcf -> ab
   ├─ acd, df -> acf
   │  ├─ acd
   │  └─ df
   └─ bcef, e -> bcf
      ├─ bcef
      └─ e


julia> viz_contraction(nested_code)
[ Info: Generating frames, 7 frames in total
[ Info: Creating video at: /var/folders/3y/xl2h1bxj4ql27p01nl5hrrnc0000gn/T/jl_SiSvrH/contraction.mp4
"/var/folders/3y/xl2h1bxj4ql27p01nl5hrrnc0000gn/T/jl_SiSvrH/contraction.mp4"

The resulting image and video will be saved in the current working directory, and the image is shown below:

The large white nodes represent the tensors, and the small colored nodes represent the indices, red for closed indices and green for open indices.