From fa0f8a77a2b5ece1caf1cba527cdf6e0ba0e3c54 Mon Sep 17 00:00:00 2001 From: Joey Date: Wed, 10 Apr 2024 09:44:10 -0400 Subject: [PATCH] Removed support for Vector{ITensor} in contract.jl functions for now --- src/ITensorsExtensions/ITensorsExtensions.jl | 82 ++++++++ src/abstractitensornetwork.jl | 3 +- src/caches/beliefpropagationcache.jl | 4 +- src/contract.jl | 33 +--- test/test_belief_propagation.jl | 197 ++++++++++--------- test/test_itensornetwork.jl | 4 +- 6 files changed, 191 insertions(+), 132 deletions(-) create mode 100644 src/ITensorsExtensions/ITensorsExtensions.jl diff --git a/src/ITensorsExtensions/ITensorsExtensions.jl b/src/ITensorsExtensions/ITensorsExtensions.jl new file mode 100644 index 00000000..66350c8f --- /dev/null +++ b/src/ITensorsExtensions/ITensorsExtensions.jl @@ -0,0 +1,82 @@ +module ITensorsExtensions +using LinearAlgebra: LinearAlgebra, eigen, pinv +using ITensors: + ITensor, + Index, + commonind, + dag, + hasqns, + inds, + isdiag, + itensor, + map_diag, + noncommonind, + noprime, + replaceinds, + space, + sqrt_decomp +using ITensors.NDTensors: + NDTensors, + Block, + Tensor, + blockdim, + blockoffsets, + denseblocks, + diaglength, + getdiagindex, + nzblocks, + setdiagindex!, + svd, + tensor, + DiagBlockSparseTensor, + DenseTensor, + BlockOffsets +using Observers: update!, insert_function! + +function NDTensors.blockoffsets(dense::DenseTensor) + return BlockOffsets{ndims(dense)}([Block(ntuple(Returns(1), ndims(dense)))], [0]) +end +function NDTensors.nzblocks(dense::DenseTensor) + return nzblocks(blockoffsets(dense)) +end +NDTensors.blockdim(ind::Int, ::Block{1}) = ind +NDTensors.blockdim(i::Index{Int}, b::Integer) = blockdim(i, Block(b)) +NDTensors.blockdim(i::Index{Int}, b::Block) = blockdim(space(i), b) + +LinearAlgebra.isdiag(it::ITensor) = isdiag(tensor(it)) + +# Convenience functions +sqrt_diag(it::ITensor) = map_diag(sqrt, it) +inv_diag(it::ITensor) = map_diag(inv, it) +invsqrt_diag(it::ITensor) = map_diag(inv ∘ sqrt, it) +pinv_diag(it::ITensor) = map_diag(pinv, it) +pinvsqrt_diag(it::ITensor) = map_diag(pinv ∘ sqrt, it) + +function map_itensor( + f::Function, A::ITensor, lind=first(inds(A)); regularization=nothing, kwargs... +) + USV = svd(A, lind; kwargs...) + U, S, V, spec, u, v = USV + S = map_diag(s -> f(s + regularization), S) + sqrtDL, δᵤᵥ, sqrtDR = sqrt_decomp(S, u, v) + sqrtDR = denseblocks(sqrtDR) * denseblocks(δᵤᵥ) + L, R = U * sqrtDL, V * sqrtDR + return L * R +end + +# Analagous to `denseblocks`. +# Extract the diagonal entries into a diagonal tensor. +function diagblocks(D::Tensor) + nzblocksD = nzblocks(D) + T = DiagBlockSparseTensor(eltype(D), nzblocksD, inds(D)) + for b in nzblocksD + for n in 1:diaglength(D) + setdiagindex!(T, getdiagindex(D, n), n) + end + end + return T +end + +diagblocks(it::ITensor) = itensor(diagblocks(tensor(it))) + +end diff --git a/src/abstractitensornetwork.jl b/src/abstractitensornetwork.jl index 77fb73b8..bfcb9397 100644 --- a/src/abstractitensornetwork.jl +++ b/src/abstractitensornetwork.jl @@ -738,9 +738,10 @@ function inner_network( return BilinearFormNetwork(A, x, y; kwargs...) end -# TODO: We should make this pass to inner_network and then to BiLinearForm. +# TODO: We should make this use the QuadraticFormNetwork constructor here. # Parts of the code (tests relying on norm_sqr being two layer and the gauging code # which relies on specific message tensors) currently would break in that case so we need to resolve +# We could have the option in the Form constructors to pre-contract the operator into the bra or ket function norm_sqr_network(ψ::AbstractITensorNetwork) return disjoint_union("bra" => dag(prime(ψ; sites=[])), "ket" => ψ) end diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl index 0f06effd..35769012 100644 --- a/src/caches/beliefpropagationcache.jl +++ b/src/caches/beliefpropagationcache.jl @@ -270,11 +270,11 @@ end function region_scalar(bp_cache::BeliefPropagationCache, pv::PartitionVertex) incoming_mts = environment(bp_cache, [pv]) local_state = factor(bp_cache, pv) - return scalar(vcat(incoming_mts, local_state)) + return contract(vcat(incoming_mts, local_state))[] end function region_scalar(bp_cache::BeliefPropagationCache, pe::PartitionEdge) - return scalar(vcat(message(bp_cache, pe), message(bp_cache, reverse(pe)))) + return contract(vcat(message(bp_cache, pe), message(bp_cache, reverse(pe))))[] end function vertex_scalars( diff --git a/src/contract.jl b/src/contract.jl index c24432c4..a5f3fdd7 100644 --- a/src/contract.jl +++ b/src/contract.jl @@ -15,10 +15,6 @@ function NDTensors.contract( return contract(Vector{ITensor}(tn); sequence=sequence_linear_index, kwargs...) end -function NDTensors.contract(alg::Algorithm"exact", tensors::Vector{ITensor}; kwargs...) - return contract(tensors; kwargs...) -end - function NDTensors.contract( alg::Union{Algorithm"density_matrix",Algorithm"ttn_svd"}, tn::AbstractITensorNetwork; @@ -28,40 +24,19 @@ function NDTensors.contract( return approx_tensornetwork(alg, tn, output_structure; kwargs...) end -function contract_density_matrix( - contract_list::Vector{ITensor}; normalize=true, contractor_kwargs... -) - tn, _ = contract( - ITensorNetwork(contract_list); alg="density_matrix", contractor_kwargs... - ) - out = Vector{ITensor}(tn) - if normalize - out .= normalize!.(copy.(out)) - end - return out -end - -function ITensors.scalar( - alg::Algorithm, tn::Union{AbstractITensorNetwork,Vector{ITensor}}; kwargs... -) +function ITensors.scalar(alg::Algorithm, tn::AbstractITensorNetwork; kwargs...) return contract(alg, tn; kwargs...)[] end -function ITensors.scalar( - tn::Union{AbstractITensorNetwork,Vector{ITensor}}; alg="exact", kwargs... -) +function ITensors.scalar(tn::AbstractITensorNetwork; alg="exact", kwargs...) return scalar(Algorithm(alg), tn; kwargs...) end -function logscalar( - tn::Union{AbstractITensorNetwork,Vector{ITensor}}; alg="exact", kwargs... -) +function logscalar(tn::AbstractITensorNetwork; alg="exact", kwargs...) return logscalar(Algorithm(alg), tn; kwargs...) end -function logscalar( - alg::Algorithm"exact", tn::Union{AbstractITensorNetwork,Vector{ITensor}}; kwargs... -) +function logscalar(alg::Algorithm"exact", tn::AbstractITensorNetwork; kwargs...) s = scalar(alg, tn; kwargs...) s = real(s) < 0 ? complex(s) : s return log(s) diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index f26f3f2e..6c380852 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -5,11 +5,13 @@ using ITensorNetworks: ITensorNetworks, BeliefPropagationCache, IndsNetwork, + ITensorNetwork, ⊗, apply, combine_linkinds, contract, contract_boundary_mps, + contract_density_matrix, contraction_sequence, environment, flatten_networks, @@ -32,61 +34,61 @@ using Test: @test, @testset @testset "belief_propagation" begin ITensors.disable_warn_order() - # #First test on an MPS, should be exact - # g_dims = (1, 6) - # g = named_grid(g_dims) - # s = siteinds("S=1/2", g) - # χ = 4 - # Random.seed!(1234) - # ψ = random_tensornetwork(s; link_space=χ) + #First test on an MPS, should be exact + g_dims = (1, 6) + g = named_grid(g_dims) + s = siteinds("S=1/2", g) + χ = 4 + Random.seed!(1234) + ψ = random_tensornetwork(s; link_space=χ) - # ψψ = ψ ⊗ prime(dag(ψ); sites=[]) + ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - # v = (1, 3) + v = (1, 3) - # Oψ = copy(ψ) - # Oψ[v] = apply(op("Sz", s[v]), ψ[v]) - # exact_sz = inner(Oψ, ψ) / inner(ψ, ψ) + Oψ = copy(ψ) + Oψ[v] = apply(op("Sz", s[v]), ψ[v]) + exact_sz = inner(Oψ, ψ) / inner(ψ, ψ) - # bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) - # bpc = update(bpc) - # env_tensors = environment(bpc, [PartitionVertex(v)]) - # numerator = contract(vcat(env_tensors, ITensor[ψ[v], op("Sz", s[v]), dag(prime(ψ[v]))]))[] - # denominator = contract(vcat(env_tensors, ITensor[ψ[v], op("I", s[v]), dag(prime(ψ[v]))]))[] + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bpc = update(bpc) + env_tensors = environment(bpc, [PartitionVertex(v)]) + numerator = contract(vcat(env_tensors, ITensor[ψ[v], op("Sz", s[v]), dag(prime(ψ[v]))]))[] + denominator = contract(vcat(env_tensors, ITensor[ψ[v], op("I", s[v]), dag(prime(ψ[v]))]))[] - # @test abs.((numerator / denominator) - exact_sz) <= 1e-14 + @test abs.((numerator / denominator) - exact_sz) <= 1e-14 - # #Test updating the underlying tensornetwork in the cache - # v = first(vertices(ψψ)) - # new_tensor = randomITensor(inds(ψψ[v])) - # bpc = update_factor(bpc, v, new_tensor) - # ψψ_updated = tensornetwork(bpc) - # @test ψψ_updated[v] == new_tensor + #Test updating the underlying tensornetwork in the cache + v = first(vertices(ψψ)) + new_tensor = randomITensor(inds(ψψ[v])) + bpc = update_factor(bpc, v, new_tensor) + ψψ_updated = tensornetwork(bpc) + @test ψψ_updated[v] == new_tensor - # #Now test on a tree, should also be exact - # g = named_comb_tree((4, 4)) - # s = siteinds("S=1/2", g) - # χ = 2 - # Random.seed!(1564) - # ψ = random_tensornetwork(s; link_space=χ) + #Now test on a tree, should also be exact + g = named_comb_tree((4, 4)) + s = siteinds("S=1/2", g) + χ = 2 + Random.seed!(1564) + ψ = random_tensornetwork(s; link_space=χ) - # ψψ = ψ ⊗ prime(dag(ψ); sites=[]) + ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - # v = (1, 3) + v = (1, 3) - # Oψ = copy(ψ) - # Oψ[v] = apply(op("Sz", s[v]), ψ[v]) - # exact_sz = inner(Oψ, ψ) / inner(ψ, ψ) + Oψ = copy(ψ) + Oψ[v] = apply(op("Sz", s[v]), ψ[v]) + exact_sz = inner(Oψ, ψ) / inner(ψ, ψ) - # bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) - # bpc = update(bpc) - # env_tensors = environment(bpc, [PartitionVertex(v)]) - # numerator = contract(vcat(env_tensors, ITensor[ψ[v], op("Sz", s[v]), dag(prime(ψ[v]))]))[] - # denominator = contract(vcat(env_tensors, ITensor[ψ[v], op("I", s[v]), dag(prime(ψ[v]))]))[] + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bpc = update(bpc) + env_tensors = environment(bpc, [PartitionVertex(v)]) + numerator = contract(vcat(env_tensors, ITensor[ψ[v], op("Sz", s[v]), dag(prime(ψ[v]))]))[] + denominator = contract(vcat(env_tensors, ITensor[ψ[v], op("I", s[v]), dag(prime(ψ[v]))]))[] - # @test abs.((numerator / denominator) - exact_sz) <= 1e-14 + @test abs.((numerator / denominator) - exact_sz) <= 1e-14 - # # #Now test two-site expec taking on the partition function of the Ising model. Not exact, but close + #Now test two-site expec taking on the partition function of the Ising model. Not exact, but close g_dims = (3, 4) g = named_grid(g_dims) s = IndsNetwork(g; link_space=2) @@ -106,63 +108,62 @@ using Test: @test, @testset numerator = contract(vcat(env_tensors, ITensor[ψOψ[v] for v in vs]))[] denominator = contract(vcat(env_tensors, ITensor[ψψ[v] for v in vs]))[] - @show abs.((numerator / denominator) - actual_szsz) @test abs.((numerator / denominator) - actual_szsz) <= 0.05 - # # #Test forming a two-site RDM. Check it has the correct size, trace 1 and is PSD - # g_dims = (3, 3) - # g = named_grid(g_dims) - # s = siteinds("S=1/2", g) - # vs = [(2, 2), (2, 3)] - # χ = 3 - # ψ = random_tensornetwork(s; link_space=χ) - # ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - - # bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) - # bpc = update(bpc; maxiter=20) - - # ψψsplit = split_index(ψψ, NamedEdge.([(v, 1) => (v, 2) for v in vs])) - # env_tensors = environment(bpc, [(v, 2) for v in vs]) - # rdm = contract(vcat(env_tensors, ITensor[ψψsplit[vp] for vp in [(v, 2) for v in vs]])) - - # rdm = array((rdm * combiner(inds(rdm; plev=0)...)) * combiner(inds(rdm; plev=1)...)) - # rdm /= tr(rdm) - - # eigs = eigvals(rdm) - # @test size(rdm) == (2^length(vs), 2^length(vs)) - # @test all(>=(0), real(eigs)) && all(==(0), imag(eigs)) - - # # # #Test more advanced block BP with MPS message tensors on a grid - # g_dims = (4, 3) - # g = named_grid(g_dims) - # s = siteinds("S=1/2", g) - # χ = 2 - # ψ = random_tensornetwork(s; link_space=χ) - # v = (2, 2) - - # ψψ = flatten_networks(ψ, dag(ψ); combine_linkinds=false, map_bra_linkinds=prime) - # Oψ = copy(ψ) - # Oψ[v] = apply(op("Sz", s[v]), ψ[v]) - # ψOψ = flatten_networks(ψ, dag(Oψ); combine_linkinds=false, map_bra_linkinds=prime) - - # combiners = linkinds_combiners(ψψ) - # ψψ = combine_linkinds(ψψ, combiners) - # ψOψ = combine_linkinds(ψOψ, combiners) - - # bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) - # bpc = update( - # bpc; - # message_update=ITensorNetworks.contract_density_matrix, - # message_update_kwargs=(; cutoff=1e-6, maxdim=4), - # ) - - # env_tensors = environment(bpc, [v]) - # numerator = contract(vcat(env_tensors, ITensor[ψOψ[v]]))[] - # denominator = contract(vcat(env_tensors, ITensor[ψψ[v]]))[] - - # exact_sz = - # contract_boundary_mps(ψOψ; cutoff=1e-16) / contract_boundary_mps(ψψ; cutoff=1e-16) - - # @test abs.((numerator / denominator) - exact_sz) <= 1e-5 + #Test forming a two-site RDM. Check it has the correct size, trace 1 and is PSD + g_dims = (3, 3) + g = named_grid(g_dims) + s = siteinds("S=1/2", g) + vs = [(2, 2), (2, 3)] + χ = 3 + ψ = random_tensornetwork(s; link_space=χ) + ψψ = ψ ⊗ prime(dag(ψ); sites=[]) + + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bpc = update(bpc; maxiter=20) + + ψψsplit = split_index(ψψ, NamedEdge.([(v, 1) => (v, 2) for v in vs])) + env_tensors = environment(bpc, [(v, 2) for v in vs]) + rdm = contract(vcat(env_tensors, ITensor[ψψsplit[vp] for vp in [(v, 2) for v in vs]])) + + rdm = array((rdm * combiner(inds(rdm; plev=0)...)) * combiner(inds(rdm; plev=1)...)) + rdm /= tr(rdm) + + eigs = eigvals(rdm) + @test size(rdm) == (2^length(vs), 2^length(vs)) + @test all(>=(0), real(eigs)) && all(==(0), imag(eigs)) + + #Test more advanced block BP with MPS message tensors on a grid + g_dims = (4, 3) + g = named_grid(g_dims) + s = siteinds("S=1/2", g) + χ = 2 + ψ = random_tensornetwork(s; link_space=χ) + v = (2, 2) + + ψψ = flatten_networks(ψ, dag(ψ); combine_linkinds=false, map_bra_linkinds=prime) + Oψ = copy(ψ) + Oψ[v] = apply(op("Sz", s[v]), ψ[v]) + ψOψ = flatten_networks(ψ, dag(Oψ); combine_linkinds=false, map_bra_linkinds=prime) + + combiners = linkinds_combiners(ψψ) + ψψ = combine_linkinds(ψψ, combiners) + ψOψ = combine_linkinds(ψOψ, combiners) + + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + message_update_func(tns; kwargs...) = + Vector{ITensor}(first(contract(ITensorNetwork(tns); alg="density_matrix", kwargs...))) + bpc = update( + bpc; message_update=message_update_func, message_update_kwargs=(; cutoff=1e-6, maxdim=4) + ) + + env_tensors = environment(bpc, [v]) + numerator = contract(vcat(env_tensors, ITensor[ψOψ[v]]))[] + denominator = contract(vcat(env_tensors, ITensor[ψψ[v]]))[] + + exact_sz = + contract_boundary_mps(ψOψ; cutoff=1e-16) / contract_boundary_mps(ψψ; cutoff=1e-16) + + @test abs.((numerator / denominator) - exact_sz) <= 1e-5 end end diff --git a/test/test_itensornetwork.jl b/test/test_itensornetwork.jl index ab67d9b7..93189137 100644 --- a/test/test_itensornetwork.jl +++ b/test/test_itensornetwork.jl @@ -41,6 +41,8 @@ using ITensorNetworks: inner_network, internalinds, linkinds, + neighbor_itensors, + norm_sqr, norm_sqr_network, orthogonalize, random_tensornetwork, @@ -50,8 +52,6 @@ using NamedGraphs: NamedEdge, incident_edges, named_comb_tree, named_grid using Random: Random, randn! using Test: @test, @test_broken, @testset -using ITensorNetworks: norm_sqr, neighbor_itensors - @testset "ITensorNetwork tests" begin @testset "ITensorNetwork Basics" begin Random.seed!(1234)