Tutorial

Example 1: Optimize contraction order without backend specified

OMEinsumContractionOrders can be used as a standalone package to optimize the contraction order of an einsum notation. The first step is to construct an OMEinsumContractionOrders.EinCode object, which is a data type to represent the einsum notation.

julia> using OMEinsumContractionOrders, Graphs, KaHyPar
julia> function random_regular_eincode(n, k; optimize=nothing, seed)
           g = Graphs.random_regular_graph(n, k; seed)
           ixs = [[minmax(e.src,e.dst)...] for e in Graphs.edges(g)]  # input indices
           iy = Int[]  # output indices (scalar output)
           return OMEinsumContractionOrders.EinCode(ixs, iy)
       endrandom_regular_eincode (generic function with 1 method)
julia> code = random_regular_eincode(100, 3; seed=42);

Here, we define an einsum notation with 3-regular graph topology. The vertices correspond to indices. On each edge, we specify a rank 2 tensor that associates with the vertices it connects. The output is a scalar.

One can use contraction_complexity function to get the time, space and rewrite cost for contracting this einsum notation. This function takes two arguments: the einsum notation and a dictionary to specify the size of the variables. Here, we use the uniformsize function to specify that all variables have the same size $2$. This function returns a dictionary that maps each variable to $2$: Dict(i=>2 for i in uniquelabels(code)), where OMEinsumContractionOrders.uniquelabels returns the unique labels in the einsum notation.

julia> size_dict = uniformsize(code, 2)Dict{Int64, Int64} with 100 entries:
  5  => 2
  56 => 2
  55 => 2
  35 => 2
  60 => 2
  30 => 2
  32 => 2
  6  => 2
  67 => 2
  45 => 2
  73 => 2
  64 => 2
  90 => 2
  4  => 2
  13 => 2
  54 => 2
  86 => 2
  63 => 2
  91 => 2
  ⋮  => ⋮
julia> contraction_complexity(code, size_dict)Time complexity: 2^100.0
Space complexity: 2^0.0
Read-write complexity: 2^9.231221180711184

Since we do not specify a contraction order, the direct contraction corresponds to brute force enumeration and costs $2^{\text{number of vertices}}$ operations. No space is required to store the intermediate contraction result and the space complexity is $0$. The read-write complexity corresponds to how many element-wise read and write operations are required to perform contraction.

The order of contraction is optimized by the optimize_code function. It takes three arguments: code, size_dict, and optimizer. The optimizer argument is the optimizer to be used. The available optimizers are listed in the optimizers page.

julia> score = ScoreFunction(sc_target=10, sc_weight=3.0)ScoreFunction(1.0, 3.0, 0.0, 10.0)
julia> optcode_tree = optimize_code(code, uniformsize(code, 2),
       	TreeSA(; βs=0.1:0.1:10, ntrials=2, niters=20, score);
       	slicer=TreeSASlicer(; score)
           );
julia> contraction_complexity(optcode_tree, size_dict)Time complexity: 2^17.099183410079284
Space complexity: 2^10.0
Read-write complexity: 2^16.148873745302843
julia> optcode_tree.slicing2-element Vector{Int64}:
 86
 85

The optimize_code function returns the optimized contraction order. The optimized contraction order is a OMEinsumContractionOrders.NestedEinsum object, which is a data type to represent the nested einsum notation.

Example 2: Use it in OMEinsum

OMEinsumContractionOrders is shipped with OMEinsum package, which is a powerful package for tensor network contraction (or einsum contraction). You can use it to optimize the contraction order of an OMEinsum notation.

julia> using OMEinsum
julia> code = ein"ij, jk, kl, il->"ij, jk, kl, il ->
julia> optimized_code = optimize_code(code, uniformsize(code, 2), TreeSA())ji, ij ->
├─ jl, il -> ji
│  ├─ jk, kl -> jl
│  │  ├─ jk
│  │  └─ kl
│  └─ il
└─ ij

Example 3: Visualization

The visualization is provided by the LuxorTensorPlot extension. To use it, just load the extension by using LuxorGraphPlot.

pkg> add OMEinsumContractionOrders, LuxorGraphPlot

julia> using OMEinsumContractionOrders, LuxorGraphPlot

The extension provides the following two functions, 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.