Contract Network Stuff

Author: JoeyT1994

src/contractnetwork.jl

@@ -1,15 +1,6 @@
-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
+default_sequence_alg = "leftassociative"
 
 function contraction_sequence(::Algorithm"leftassociative", tn::Vector{<:AbstractArray})
   return Any[i for i in 1:length(tn)]
@@ -20,20 +11,20 @@ function contraction_sequence(tn::Vector{<:AbstractArray}; alg=default_sequence_
 end
 
 # Internal recursive worker
-function recursive_contractnetwork(tn::Union{AbstractVector,AbstractArray})
-  tn isa AbstractVector && return reduce(*, map(recursive_contractnetwork, tn))
+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(tn::Vector{<:AbstractArray}; sequence_alg=default_sequence_alg)
-  sequence = contraction_sequence(tn; alg=sequence_alg)
-  return recursive_contractnetwork(rearrange(tn, sequence))
+function contractnetwork(tn::Vector{<:AbstractArray}; sequence=default_sequence_alg)
+  contract_sequence = isa(sequence, String) ? contraction_sequence(tn; alg=sequence) : sequence
+  return recursive_contractnetwork(rearrange(tn, contract_sequence))
 end
 
-function contractnetwork(tn::AbstractTensorNetwork; sequence_alg=default_sequence_alg)
-  return contractnetwork([tn[v] for v in vertices(tn)]; sequence_alg)
+function contractnetwork(tn::AbstractTensorNetwork; sequence=default_sequence_alg)
+  return contractnetwork([tn[v] for v in vertices(tn)]; sequence)
 end

test/test_contractnetwork.jl

@@ -4,6 +4,7 @@ using NamedGraphs.NamedGraphGenerators: named_grid
 using ITensorBase: Index, ITensor
 using ITensorNetworksNext:
   TensorNetwork, linkinds, siteinds, contractnetwork, contraction_sequence
+using TensorOperations
 using Test: @test, @testset
 
 @testset "ContractNetwork" begin
@@ -14,8 +15,8 @@ using Test: @test, @testset
     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")
+    ABCD_1 = contractnetwork([A, B, C, D]; sequence="leftassociative")
+    ABCD_2 = contractnetwork([A, B, C, D]; sequence="optimal")
 
     @test ABCD_1 == ABCD_2
   end
@@ -30,8 +31,8 @@ using Test: @test, @testset
       return randn(Tuple(is))
     end
 
-    z1 = contractnetwork(tn; sequence_alg="optimal")[]
-    z2 = contractnetwork(tn; sequence_alg="leftassociative")[]
+    z1 = contractnetwork(tn; sequence="optimal")[]
+    z2 = contractnetwork(tn; sequence="leftassociative")[]
 
     @test abs(z1 - z2) / abs(z1) <= 1e3*eps(Float64)
   end