Working contract_network (#7)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Matt Fishman <[email protected]>

Author: 3 people

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorNetworksNext"
 uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.10"
+version = "0.1.11"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -20,6 +20,12 @@ TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 
+[weakdeps]
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+
+[extensions]
+ITensorNetworksNextTensorOperationsExt = "TensorOperations"
+
 [compat]
 AbstractTrees = "0.4.5"
 Adapt = "4.3"
@@ -33,6 +39,7 @@ NamedDimsArrays = "0.8"
 NamedGraphs = "0.6.9, 0.7"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
+TensorOperations = "5.3.1"
 TermInterface = "2"
 TypeParameterAccessors = "0.4.4"
 WrappedUnions = "0.3"

ext/ITensorNetworksNextTensorOperationsExt/ITensorNetworksNextTensorOperationsExt.jl

@@ -0,0 +1,16 @@
+module ITensorNetworksNextTensorOperationsExt
+
+using BackendSelection: @Algorithm_str, Algorithm
+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)
+end
+
+end

src/ITensorNetworksNext.jl

@@ -3,5 +3,6 @@ module ITensorNetworksNext
 include("lazynameddimsarrays.jl")
 include("abstracttensornetwork.jl")
 include("tensornetwork.jl")
+include("contract_network.jl")
 
 end

src/contract_network.jl

@@ -0,0 +1,47 @@
+using BackendSelection: @Algorithm_str, Algorithm
+using ITensorNetworksNext.LazyNamedDimsArrays: substitute, materialize, lazy,
+    symnameddims
+
+#Algorithmic defaults
+default_sequence_alg(::Algorithm"exact") = "leftassociative"
+default_sequence(::Algorithm"exact") = nothing
+function set_default_kwargs(alg::Algorithm"exact")
+    sequence = get(alg, :sequence, nothing)
+    sequence_alg = get(alg, :sequence_alg, default_sequence_alg(alg))
+    return Algorithm("exact"; sequence, sequence_alg)
+end
+
+function contraction_sequence_to_expr(seq)
+    if seq isa AbstractVector
+        return prod(contraction_sequence_to_expr, seq)
+    else
+        return symnameddims(seq)
+    end
+end
+
+function contraction_sequence(::Algorithm"leftassociative", tn::Vector{<:AbstractArray})
+    return prod(symnameddims, 1:length(tn))
+end
+
+function contraction_sequence(tn::Vector{<:AbstractArray}; sequence_alg = default_sequence_alg(Algorithm("exact")))
+    return contraction_sequence(Algorithm(sequence_alg), tn)
+end
+
+function contract_network(alg::Algorithm"exact", tn::Vector{<:AbstractArray})
+    if !isnothing(alg.sequence)
+        sequence = alg.sequence
+    else
+        sequence = contraction_sequence(tn; sequence_alg = alg.sequence_alg)
+    end
+
+    sequence = substitute(sequence, Dict(symnameddims(i) => lazy(tn[i]) for i in 1:length(tn)))
+    return materialize(sequence)
+end
+
+function contract_network(alg::Algorithm"exact", tn::AbstractTensorNetwork)
+    return contract_network(alg, [tn[v] for v in vertices(tn)])
+end
+
+function contract_network(tn; alg, kwargs...)
+    return contract_network(set_default_kwargs(Algorithm(alg; kwargs...)), tn)
+end

test/Project.toml

@@ -9,6 +9,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"
 TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
@@ -25,5 +26,6 @@ NamedGraphs = "0.6.8, 0.7"
 SafeTestsets = "0.1"
 Suppressor = "0.2.8"
 TermInterface = "2"
+TensorOperations = "5.3.1"
 Test = "1.10"
 WrappedUnions = "0.3"

test/test_contract_network.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, contract_network
+using TensorOperations: TensorOperations
+using Test: @test, @testset
+
+@testset "contract_network" 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 = contract_network([A, B, C, D]; alg = "exact", sequence_alg = "leftassociative")
+        ABCD_2 = contract_network([A, B, C, D]; alg = "exact", 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 = contract_network(tn; alg = "exact", sequence_alg = "optimal")[]
+        z2 = contract_network(tn; alg = "exact", sequence_alg = "leftassociative")[]
+
+        @test abs(z1 - z2) / abs(z1) <= 1.0e3 * eps(Float64)
+    end
+end