Eager contraction sequence/order optimization (#21)

Author: mtfishman

Project.toml

@@ -1,13 +1,15 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.1.15"
+version = "0.2.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BackendSelection = "680c2d7c-f67a-4cc9-ae9c-da132b1447a5"
+Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
+DerivableInterfaces = "6c5e35bf-e59e-4898-b73c-732dcc4ba65f"
 DiagonalArrays = "74fd4be6-21e2-4f6f-823a-4360d37c7a77"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -31,7 +33,9 @@ ITensorNetworksNextTensorOperationsExt = "TensorOperations"
 AbstractTrees = "0.4.5"
 Adapt = "4.3"
 BackendSelection = "0.1.6"
+Combinatorics = "1"
 DataGraphs = "0.2.7"
+DerivableInterfaces = "0.5.5"
 DiagonalArrays = "0.3.23"
 Dictionaries = "0.4.5"
 Graphs = "1.13.1"

docs/Project.toml

@@ -6,4 +6,4 @@ ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 [compat]
 Documenter = "1"
 Literate = "2"
-ITensorNetworksNext = "0.1"
+ITensorNetworksNext = "0.2"

examples/Project.toml

@@ -2,4 +2,4 @@
 ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 
 [compat]
-ITensorNetworksNext = "0.1"
+ITensorNetworksNext = "0.2"

ext/ITensorNetworksNextTensorOperationsExt/ITensorNetworksNextTensorOperationsExt.jl

@@ -1,16 +1,28 @@
 module ITensorNetworksNextTensorOperationsExt
 
 using BackendSelection: @Algorithm_str, Algorithm
+using ITensorNetworksNext.LazyNamedDimsArrays: LazyNamedDimsArrays, ismul, symnameddims,
+    substitute
+using ITensorNetworksNext.LazyNamedDimsArrays.TermInterface: arguments
 using NamedDimsArrays: inds
-using ITensorNetworksNext: ITensorNetworksNext, contraction_sequence_to_expr
 using TensorOperations: TensorOperations, optimaltree
 
-function ITensorNetworksNext.contraction_sequence(::Algorithm"optimal", tn::Vector{<:AbstractArray})
-    network = collect.(inds.(tn))
-    #Converting dims to Float64 to minimize overflow issues
-    inds_to_dims = Dict(i => Float64(length(i)) for i in unique(reduce(vcat, network)))
-    seq, _ = optimaltree(network, inds_to_dims)
-    return contraction_sequence_to_expr(seq)
+function contraction_tree_to_expr(f, tree)
+    return if !(tree isa AbstractVector)
+        f(tree)
+    else
+        prod(Base.Fix1(contraction_tree_to_expr, f), tree)
+    end
+end
+
+function LazyNamedDimsArrays.optimize_contraction_order(alg::Algorithm"optimal", a)
+    @assert ismul(a)
+    ts = arguments(a)
+    inds_network = collect.(inds.(ts))
+    # Converting dims to Float64 to minimize overflow issues
+    inds_to_dims = Dict(i => Float64(length(i)) for i in reduce(∪, inds_network))
+    tree, _ = optimaltree(inds_network, inds_to_dims)
+    return contraction_tree_to_expr(i -> ts[i], tree)
 end
 
 end