Contraction sequence finding

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorNetworks"
 uuid = "2919e153-833c-4bdc-8836-1ea460a35fc7"
 authors = ["Matthew Fishman <[email protected]>, Joseph Tindall <[email protected]> and contributors"]
-version = "0.11.27"
+version = "0.12.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -32,6 +32,7 @@ SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StructWalk = "31cdf514-beb7-4750-89db-dda9d2eb8d3d"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 
@@ -48,6 +49,7 @@ ITensorNetworksEinExprsExt = "EinExprs"
 ITensorNetworksGraphsFlowsExt = "GraphsFlows"
 ITensorNetworksOMEinsumContractionOrdersExt = "OMEinsumContractionOrders"
 ITensorNetworksObserversExt = "Observers"
+ITensorNetworksTensorOperationsExt = "TensorOperations"
 
 [compat]
 AbstractTrees = "0.4.4"
@@ -80,6 +82,7 @@ SplitApplyCombine = "1.2"
 StaticArrays = "1.5.12"
 StructWalk = "0.2"
 Suppressor = "0.2"
+TensorOperations = "5.1.4"
 TimerOutputs = "0.5.22"
 TupleTools = "1.4"
 julia = "1.10"
@@ -90,6 +93,7 @@ EinExprs = "b1794770-133b-4de1-afb4-526377e9f4c5"
 GraphsFlows = "06909019-6f44-4949-96fc-b9d9aaa02889"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Observers = "338f10d5-c7f1-4033-a7d1-f9dec39bcaa0"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]

ext/ITensorNetworksTensorOperationsExt/ITensorNetworksTensorOperationsExt.jl

@@ -0,0 +1,19 @@
+module ITensorNetworksTensorOperationsExt
+
+using ITensors: ITensors, ITensor, dim, inds
+using ITensorNetworks: ITensorNetworks
+using TensorOperations: optimaltree
+using NDTensors.AlgorithmSelection: @Algorithm_str
+
+function ITensorNetworks.contraction_sequence(::Algorithm"optimal", tn::Vector{ITensor})
+  return optimal_contraction_sequence(tn)
+end
+
+function optimal_contraction_sequence(tensors::Vector{<:ITensor})
+  network = collect.(inds.(tensors))
+  inds_to_dims = Dict(i => dim(i) for i in unique(reduce(vcat, network)))
+  seq, _ = optimaltree(network, inds_to_dims)
+  return seq
+end
+
+end

src/contract.jl

@@ -1,5 +1,4 @@
 using ITensors: ITensor, scalar
-using ITensors.ContractionSequenceOptimization: deepmap
 using ITensors.NDTensors: NDTensors, Algorithm, @Algorithm_str, contract
 using LinearAlgebra: normalize!
 using NamedGraphs: NamedGraphs

src/contraction_sequences.jl

@@ -1,6 +1,5 @@
 using Graphs: vertices
-using ITensors: ITensor, contract
-using ITensors.ContractionSequenceOptimization: deepmap, optimal_contraction_sequence
+using ITensors: ITensor
 using ITensors.NDTensors: Algorithm, @Algorithm_str
 using NamedGraphs.Keys: Key
 using NamedGraphs.OrdinalIndexing: th
@@ -9,14 +8,14 @@ function contraction_sequence(tn::Vector{ITensor}; alg="optimal", kwargs...)
   return contraction_sequence(Algorithm(alg), tn; kwargs...)
 end
 
+function deepmap(f, tree; filter=(x -> x isa AbstractArray))
+  return filter(tree) ? map(t -> deepmap(f, t; filter=filter), tree) : f(tree)
+end
+
 function contraction_sequence(tn::AbstractITensorNetwork; kwargs...)
   # TODO: Use `token_vertex` and/or `token_vertices` here.
   ts = map(v -> tn[v], (1:nv(tn))th)
   seq_linear_index = contraction_sequence(ts; kwargs...)
   # TODO: Use `Functors.fmap` or `StructWalk`?
   return deepmap(n -> Key(vertices(tn)[n * th]), seq_linear_index)
 end
-
-function contraction_sequence(::Algorithm"optimal", tn::Vector{ITensor})
-  return optimal_contraction_sequence(tn)
-end

test/test_contraction_sequence.jl

@@ -5,6 +5,7 @@ using ITensorNetworks:
 using ITensors: ITensors, contract
 using NamedGraphs.NamedGraphGenerators: named_grid
 using OMEinsumContractionOrders: OMEinsumContractionOrders
+using TensorOperations
 using StableRNGs: StableRNG
 using Test: @test, @testset
 @testset "contraction_sequence" begin

test/test_expect.jl

@@ -12,6 +12,7 @@ using ITensorNetworks:
 using SplitApplyCombine: group
 using StableRNGs: StableRNG
 using Test: @test, @testset
+using TensorOperations
 @testset "Test Expect" begin
   #Test on a tree
   L, χ = 4, 2