Working ContractNetwork

Project.toml

@@ -1,31 +1,40 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
-version = "0.1.4"
+version = "0.1.5"
 
 [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"
 
+[weakdeps]
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+
+[extensions]
+ITensorNetworksNextTensorOperationsExt = "TensorOperations"
+
 [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, 0.7"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
+TensorOperations = "5.3.1"
 julia = "1.10"

ext/ITensorNetworksNextTensorOperationsExt/ITensorNetworksNextTensorOperationsExt.jl

@@ -0,0 +1,16 @@
+module ITensorNetworksNextTensorOperationsExt
+
+using BackendSelection: @Algorithm_str, Algorithm
+using ITensorBase: inds, dim
+using ITensorNetworksNext: ITensorNetworksNext
+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(dim(i)) for i in unique(reduce(vcat, network)))
+    seq, _ = optimaltree(network, inds_to_dims)
+    return seq
+end
+
+end

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,42 @@
+using BackendSelection: @Algorithm_str, Algorithm
+
+default_contract_alg = nothing
+
+#Algorithmic defaults
+default_sequence(::Algorithm"exact") = "leftassociative"
+function set_default_kwargs(alg::Algorithm"exact")
+    sequence = get(alg, :sequence, default_sequence(alg))
+    return Algorithm("exact"; sequence)
+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)
+    return contraction_sequence(Algorithm(alg), tn)
+end
+
+# Internal recursive worker
+function recursive_contractnetwork(tn::Union{AbstractVector, AbstractNamedDimsArray})
+    tn isa AbstractVector && return prod(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(alg::Algorithm"exact", tn::Vector{<:AbstractArray})
+    contract_sequence = isa(alg.sequence, String) ? contraction_sequence(tn; alg = alg.sequence) : sequence
+    return recursive_contractnetwork(rearrange(tn, contract_sequence))
+end
+
+function contractnetwork(alg::Algorithm"exact", tn::AbstractTensorNetwork)
+    return contractnetwork(alg, [tn[v] for v in vertices(tn)])
+end
+
+function contractnetwork(tn::Union{AbstractTensorNetwork, Vector{<:AbstractArray}}; alg = default_contract_alg, kwargs...)
+    alg == nothing && error("Must specify an algorithm to contract the network with")
+    return contractnetwork(set_default_kwargs(Algorithm(alg; kwargs...)), tn)
+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, 0.7"
 SafeTestsets = "0.1"
 Suppressor = "0.2.8"
+TensorOperations = "5.3.1"
 Test = "1.10"

test/test_contractnetwork.jl

@@ -0,0 +1,39 @@
+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 TensorOperations: TensorOperations
+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]; alg = "exact", sequence = "leftassociative")
+        ABCD_2 = contractnetwork([A, B, C, D]; alg = "exact", sequence = "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; alg = "exact", sequence = "optimal")[]
+        z2 = contractnetwork(tn; alg = "exact", sequence = "leftassociative")[]
+
+        @test abs(z1 - z2) / abs(z1) <= 1.0e3 * eps(Float64)
+    end
+end