Working contract_network

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"
+authors = ["ITensor developers <[email protected]> and contributors"]
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -20,6 +20,13 @@ TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 
+[weakdeps]
+ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
+TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
+
+[extensions]
+ITensorNetworksNextTensorOperationsExt = "TensorOperations"
+
 [compat]
 AbstractTrees = "0.4.5"
 Adapt = "4.3"
@@ -33,7 +40,5 @@ NamedDimsArrays = "0.8"
 NamedGraphs = "0.6.9, 0.7"
 SimpleTraits = "0.9.5"
 SplitApplyCombine = "1.2.3"
-TermInterface = "2"
-TypeParameterAccessors = "0.4.4"
-WrappedUnions = "0.3"
+TensorOperations = "5.3.1"
 julia = "1.10"

ext/ITensorNetworksNextTensorOperationsExt/ITensorNetworksNextTensorOperationsExt.jl

@@ -0,0 +1,17 @@
+module ITensorNetworksNextTensorOperationsExt
+
+using BackendSelection: @Algorithm_str, Algorithm
+using NamedDimsArrays: inds
+using ITensorNetworksNext: ITensorNetworksNext
+using ITensorNetworksNext.LazyNamedDimsArrays: nested_array_to_lazy_multiply
+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 nested_array_to_lazy_multiply(seq)
+end
+
+end

src/ITensorNetworksNext.jl

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

src/contractnetwork.jl

@@ -0,0 +1,35 @@
+using BackendSelection: @Algorithm_str, Algorithm
+using ITensorNetworksNext.LazyNamedDimsArrays: nested_array_to_lazy_multiply, substitute_lazy, materialize, lazy,
+    symnameddims
+
+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 nested_array_to_lazy_multiply(collect.(1:length(tn)))
+end
+
+function contraction_sequence(tn::Vector{<:AbstractArray}; alg = default_sequence_alg)
+    return contraction_sequence(Algorithm(alg), tn)
+end
+
+function contractnetwork(alg::Algorithm"exact", tn::Vector{<:AbstractArray})
+    contract_sequence = isa(alg.sequence, String) ? contraction_sequence(tn; alg = alg.sequence) : sequence
+    contract_sequence = substitute_lazy(contract_sequence, Dict(symnameddims(i) => lazy(tn[i]) for i in 1:length(tn)))
+    return materialize(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

src/lazynameddimsarrays.jl

@@ -417,4 +417,12 @@ Base.:*(a::SymbolicNamedDimsArray, b::SymbolicNamedDimsArray) = lazy(a) * lazy(b
 Base.:*(a::SymbolicNamedDimsArray, b::LazyNamedDimsArray) = lazy(a) * b
 Base.:*(a::LazyNamedDimsArray, b::SymbolicNamedDimsArray) = a * lazy(b)
 
+function contraction_sequence_to_expr(seq)
+    if seq isa AbstractVector
+        return prod(contraction_sequence_to_expr, seq)
+    else
+        return symnameddims(seq)
+    end
+end
+
 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_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, symnameddims, lazy, substitute_lazy, materialize
+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