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"
@@ -41,13 +41,15 @@ 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"
 
 [extensions]
 ITensorNetworksAdaptExt = "Adapt"
 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,15 @@
+module ITensorNetworksTensorOperationsExt
+
+using ITensors: ITensors, ITensor, dim, inds
+using ITensorNetworks: ITensorNetworks
+using TensorOperations: optimaltree, TensorOperations
+using NDTensors.AlgorithmSelection: @Algorithm_str
+
+function ITensorNetworks.contraction_sequence(::Algorithm"optimal", tn::Vector{ITensor})
+  network = collect.(inds.(tn))
+  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,23 @@ function contraction_sequence(tn::Vector{ITensor}; alg="optimal", kwargs...)
   return contraction_sequence(Algorithm(alg), tn; kwargs...)
 end
 
+function contraction_sequence(alg::Algorithm, tn::Vector{ITensor})
+  return throw(
+    ArgumentError(
+      "Algorithm $alg isn't defined for contraction sequence finding. Try loading a backend package like 
+        TensorOperations.jl or OMEinsumContractionOrders.jl.",
+    ),
+  )
+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_additensornetworks.jl

@@ -7,6 +7,7 @@ using ITensors: ITensors, apply, op, scalar, inner
 using LinearAlgebra: norm_sqr
 using StableRNGs: StableRNG
 using Test: @test, @testset
+using TensorOperations: TensorOperations
 @testset "add_itensornetworks" begin
   g = named_grid((2, 2))
   s = siteinds("S=1/2", g)

test/test_apply.jl

@@ -16,6 +16,7 @@ using NamedGraphs.NamedGraphGenerators: named_grid
 using NamedGraphs.PartitionedGraphs: PartitionVertex
 using SplitApplyCombine: group
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 @testset "apply" begin
   g_dims = (2, 2)

test/test_belief_propagation.jl

@@ -34,6 +34,7 @@ using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid
 using NamedGraphs.PartitionedGraphs: PartitionVertex, partitionedges
 using SplitApplyCombine: group
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 
 @testset "belief_propagation (eltype=$elt)" for elt in (

test/test_binary_tree_partition.jl

@@ -24,6 +24,7 @@ using NamedGraphs.GraphsExtensions:
   is_binary_arborescence, post_order_dfs_vertices, root_vertex
 using OMEinsumContractionOrders: OMEinsumContractionOrders
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 
 @testset "test mincut functions on top of MPS" begin

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: TensorOperations
 using StableRNGs: StableRNG
 using Test: @test, @testset
 @testset "contraction_sequence" begin

test/test_contraction_sequence_to_graph.jl

@@ -10,6 +10,7 @@ using NamedGraphs.GraphsExtensions:
   is_leaf_vertex, leaf_vertices, non_leaf_edges, root_vertex
 using NamedGraphs.NamedGraphGenerators: named_grid
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 @testset "contraction_sequence_to_graph" begin
   n = 3

test/test_expect.jl

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

test/test_forms.jl

@@ -23,6 +23,7 @@ using ITensorNetworks:
 using ITensors: contract, dag, inds, prime, random_itensor
 using LinearAlgebra: norm
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @testset
 @testset "FormNetworks" begin
   g = named_grid((1, 4))

test/test_inner.jl

@@ -16,6 +16,7 @@ using Graphs: SimpleGraph, uniform_tree
 using NamedGraphs: NamedGraph
 using StableRNGs: StableRNG
 using Test: @test, @testset
+using TensorOperations: TensorOperations
 @testset "Inner products, BP vs exact comparison" begin
   L = 4
   χ = 2

test/test_itensornetwork.jl

@@ -60,6 +60,7 @@ using NamedGraphs.NamedGraphGenerators: named_comb_tree, named_grid
 using NDTensors: NDTensors, dim
 using Random: randn!
 using StableRNGs: StableRNG
+using TensorOperations: TensorOperations
 using Test: @test, @test_broken, @testset
 const elts = (Float32, Float64, Complex{Float32}, Complex{Float64})
 @testset "ITensorNetwork tests" begin