Contraction path optimization with EinExprs

Project.toml

@@ -31,6 +31,12 @@ Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
+[weakdeps]
+EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
+
+[extensions]
+ITensorNetworksEinExprsExt = "EinExprs"
+
 [compat]
 AbstractTrees = "0.4.4"
 Combinatorics = "1"
@@ -40,6 +46,7 @@ DataStructures = "0.18"
 Dictionaries = "0.3.15"
 Distributions = "0.25.86"
 DocStringExtensions = "0.8, 0.9"
+EinExprs = "0.6.4"
 Graphs = "1.8"
 GraphsFlows = "0.1.1"
 ITensors = "0.3.23"
@@ -60,6 +67,7 @@ julia = "1.7"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 
 [targets]
 test = ["Test"]

ext/ITensorNetworksEinExprsExt/ITensorNetworksEinExprsExt.jl

@@ -0,0 +1,52 @@
+module ITensorNetworksEinExprsExt
+
+using ITensors: Index, ITensor, @Algorithm_str, inds, noncommoninds
+using ITensorNetworks:
+  ITensorNetworks, ITensorNetwork, vertextype, vertex_data, contraction_sequence
+using EinExprs: EinExprs, EinExpr, einexpr, SizedEinExpr
+
+function to_einexpr(ts::Vector{ITensor})
+  IndexType = Any
+
+  tensor_exprs = EinExpr{IndexType}[]
+  inds_dims = Dict{IndexType,Int}()
+
+  for tensor_v in ts
+    inds_v = collect(inds(tensor_v))
+    push!(tensor_exprs, EinExpr{IndexType}(; head=inds_v))
+    merge!(inds_dims, Dict(inds_v .=> size(tensor_v)))
+  end
+
+  externalinds_tn = reduce(noncommoninds, ts)
+  return SizedEinExpr(sum(tensor_exprs; skip=externalinds_tn), inds_dims)
+end
+
+function tensor_inds_to_vertex(ts::Vector{ITensor})
+  IndexType = Any
+
+  mapping = Dict{Set{IndexType},VertexType}()
+
+  for (v, tensor_v) in enumerate(ts)
+    inds_v = collect(inds(tensor_v))
+    mapping[Set(inds_v)] = v
+  end
+
+  return mapping
+end
+
+function ITensorNetworks.contraction_sequence(
+  ::Algorithm"einexpr", tn::Vector{ITensor}; optimizer=EinExprs.Exhaustive()
+)
+  expr = to_einexpr(tn)
+  path = einexpr(optimizer, expr)
+  return to_contraction_sequence(path, tensor_inds_to_vertex(tn))
+end
+
+function to_contraction_sequence(expr, tensor_inds_to_vertex)
+  EinExprs.nargs(expr) == 0 && return tensor_inds_to_vertex[Set(expr.head)]
+  return map(
+    expr -> to_contraction_sequence(expr, tensor_inds_to_vertex), EinExprs.args(expr)
+  )
+end
+
+end

test/Project.toml

@@ -4,6 +4,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 Glob = "c27321d9-0574-5035-807b-f59d2c89b15c"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 GraphsFlows = "06909019-6f44-4949-96fc-b9d9aaa02889"

test/test_contraction_sequence.jl

@@ -3,6 +3,7 @@ using ITensorNetworks
 using OMEinsumContractionOrders
 using Random
 using Test
+using EinExprs: Exhaustive, Greedy, HyPar
 
 Random.seed!(1234)
 
@@ -23,7 +24,15 @@ ITensors.disable_warn_order()
   res_tree_sa = contract(tn; sequence=seq_tree_sa)[]
   seq_sa_bipartite = contraction_sequence(tn; alg="sa_bipartite")
   res_sa_bipartite = contract(tn; sequence=seq_sa_bipartite)[]
-  @test res_optimal ≈ res_greedy ≈ res_tree_sa ≈ res_sa_bipartite
+  seq_einexprs_exhaustive = contraction_sequence(tn; alg="einexpr", optimizer=Exhaustive())
+  res_einexprs_exhaustive = contract(tn; sequence=seq_einexprs_exhaustive)[]
+  seq_einexprs_greedy = contraction_sequence(tn; alg="einexpr", optimizer=Greedy())
+  res_einexprs_greedy = contract(tn; sequence=seq_einexprs_exhaustive)[]
+  @test res_greedy ≈ res_optimal
+  @test res_tree_sa ≈ res_optimal
+  @test res_sa_bipartite ≈ res_optimal
+  @test res_einexprs_exhaustive ≈ res_optimal
+  @test res_einexprs_greedy ≈ res_optimal
 
   if !Sys.iswindows()
     # KaHyPar doesn't work on Windows
@@ -34,5 +43,8 @@ ITensors.disable_warn_order()
     seq_kahypar_bipartite = contraction_sequence(tn; alg="kahypar_bipartite", sc_target=200)
     res_kahypar_bipartite = contract(tn; sequence=seq_kahypar_bipartite)[]
     @test res_optimal ≈ res_kahypar_bipartite
+    seq_einexprs_kahypar = contraction_sequence(tn; alg="einexpr", optimizer=HyPar())
+    res_einexprs_kahypar = contract(tn; sequence=seq_einexprs_kahypar)[]
+    @test res_einexprs_kahypar ≈ res_optimal
   end
 end