Working ContractNetwork

Author: JoeyT1994

Project.toml

@@ -1,31 +1,35 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.2"
+version = "0.1.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 BackendSelection = "680c2d7c-f67a-4cc9-ae9c-da132b1447a5"
 DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
+ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [compat]
 Adapt = "4.3.0"
 BackendSelection = "0.1.6"
 DataGraphs = "0.2.7"
 Dictionaries = "0.4.5"
 Graphs = "1.13.1"
+ITensorBase = "0.2.13"
 LinearAlgebra = "1.10"
 MacroTools = "0.5.16"
 NamedDimsArrays = "0.7.13"
 NamedGraphs = "0.6.9"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
+TensorOperations = "5.3.1"
 julia = "1.10"

src/ITensorNetworksNext.jl

@@ -2,5 +2,6 @@ module ITensorNetworksNext
 
 include("abstracttensornetwork.jl")
 include("tensornetwork.jl")
+include("contractnetwork.jl")
 
 end

src/contractnetwork.jl

@@ -0,0 +1,39 @@
+using TensorOperations: TensorOperations, optimaltree
+using ITensorBase: inds, dim
+
+default_sequence_alg = "optimal"
+
+function 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(dim(i)) for i in unique(reduce(vcat, network)))
+  seq, _ = optimaltree(network, inds_to_dims)
+  return seq
+end
+
+function contraction_sequence(::Algorithm"leftassociative", tn::Vector{<:AbstractArray})
+  return Any[i for i in 1:length(tn)]
+end
+
+function contraction_sequence(tn::Vector{<:AbstractArray}; alg=default_sequence_alg)
+  contraction_sequence(Algorithm(alg), tn)
+end
+
+# Internal recursive worker
+function recursive_contractnetwork(tn::Union{AbstractVector,AbstractArray})
+  tn isa AbstractVector && return reduce(*, map(recursive_contractnetwork, tn))
+  return tn
+end
+
+# Recursive worker for ordering the tensors according to the sequence
+rearrange(tn::Vector{<:AbstractArray}, i::Integer) = tn[i]
+rearrange(tn::Vector{<:AbstractArray}, v::AbstractVector) = [rearrange(tn, s) for s in v]
+
+function contractnetwork(tn::Vector{<:AbstractArray}; sequence_alg=default_sequence_alg)
+  sequence = contraction_sequence(tn; alg=sequence_alg)
+  return recursive_contractnetwork(rearrange(tn, sequence))
+end
+
+function contractnetwork(tn::AbstractTensorNetwork; sequence_alg=default_sequence_alg)
+  return contractnetwork([tn[v] for v in vertices(tn)]; sequence_alg)
+end

test/Project.toml

@@ -8,6 +8,7 @@ NamedDimsArrays = "60cbd0c0-df58-4cb7-918c-6f5607b73fde"
 NamedGraphs = "678767b0-92e7-4007-89e4-4527a8725b19"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
@@ -20,4 +21,5 @@ NamedDimsArrays = "0.7.14"
 NamedGraphs = "0.6.8"
 SafeTestsets = "0.1"
 Suppressor = "0.2.8"
+TensorOperations = "5.3.1"
 Test = "1.10"

test/test_contractnetwork.jl

@@ -0,0 +1,38 @@
+using Graphs: edges
+using NamedGraphs.GraphsExtensions: arranged_edges, incident_edges
+using NamedGraphs.NamedGraphGenerators: named_grid
+using ITensorBase: Index, ITensor
+using ITensorNetworksNext:
+  TensorNetwork, linkinds, siteinds, contractnetwork, contraction_sequence
+using Test: @test, @testset
+
+@testset "ContractNetwork" begin
+  @testset "Contract Vectors of ITensors" begin
+    i, j, k = Index(2), Index(2), Index(5)
+    A = ITensor([1.0 1.0; 0.5 1.0], i, j)
+    B = ITensor([2.0, 1.0], i)
+    C = ITensor([5.0, 1.0], j)
+    D = ITensor([-2.0, 3.0, 4.0, 5.0, 1.0], k)
+
+    ABCD_1 = contractnetwork([A, B, C, D]; sequence_alg="leftassociative")
+    ABCD_2 = contractnetwork([A, B, C, D]; sequence_alg="optimal")
+
+    @test ABCD_1 == ABCD_2
+  end
+
+  @testset "Contract One Dimensional Network" begin
+    dims = (4, 4)
+    g = named_grid(dims)
+    l = Dict(e => Index(2) for e in edges(g))
+    l = merge(l, Dict(reverse(e) => l[e] for e in edges(g)))
+    tn = TensorNetwork(g) do v
+      is = map(e -> l[e], incident_edges(g, v))
+      return randn(Tuple(is))
+    end
+
+    z1 = contractnetwork(tn; sequence_alg="optimal")[]
+    z2 = contractnetwork(tn; sequence_alg="leftassociative")[]
+
+    @test abs(z1 - z2) / abs(z1) <= 1e-14
+  end
+end