From 204f45b2f0941d17875bb3a9b7f9b612cb0eda11 Mon Sep 17 00:00:00 2001 From: Joseph Tindall <51231103+JoeyT1994@users.noreply.github.com> Date: Thu, 29 Jun 2023 07:13:10 -0400 Subject: [PATCH] Optimizations, `norm_network`, faster `dag`, gauging example (#94) --- examples/gauging/gauging_itns.jl | 116 +++++++++++++++++++++++++++++++ src/abstractitensornetwork.jl | 26 ++++++- src/apply.jl | 39 +++++------ src/gauging.jl | 31 +++++---- test/test_apply.jl | 5 +- 5 files changed, 181 insertions(+), 36 deletions(-) create mode 100644 examples/gauging/gauging_itns.jl diff --git a/examples/gauging/gauging_itns.jl b/examples/gauging/gauging_itns.jl new file mode 100644 index 00000000..6c30cfd4 --- /dev/null +++ b/examples/gauging/gauging_itns.jl @@ -0,0 +1,116 @@ +using Compat +using ITensors +using Metis +using ITensorNetworks +using Random +using SplitApplyCombine +using ProfileView + +using ITensorNetworks: + message_tensors, + nested_graph_leaf_vertices, + belief_propagation_iteration, + belief_propagation, + find_subgraph, + vidal_gauge, + symmetric_gauge, + vidal_itn_canonicalness, + vidal_to_symmetric_gauge, + initialize_bond_tensors, + vidal_itn_isometries, + norm_network + +using NamedGraphs +using NamedGraphs: add_edges!, rem_vertex!, hexagonal_lattice_graph +using Graphs + +"""Eager Gauging""" +function eager_gauging(ψ::ITensorNetwork, bond_tensors::DataGraph, mts::DataGraph) + isometries = vidal_itn_isometries(ψ, bond_tensors) + + ψ = copy(ψ) + mts = copy(mts) + for e in edges(ψ) + s1, s2 = find_subgraph((src(e), 1), mts), find_subgraph((dst(e), 1), mts) + normalize!(isometries[e]) + normalize!(isometries[reverse(e)]) + mts[s1 => s2], mts[s2 => s1] = ITensorNetwork(isometries[e]), + ITensorNetwork(isometries[reverse(e)]) + end + + ψ, bond_tensors = vidal_gauge(ψ, mts, bond_tensors) + + return ψ, bond_tensors, mts +end + +"""Bring an ITN into the Vidal gauge, various methods possible. Result is timed""" +function benchmark_state_gauging( + ψ::ITensorNetwork; mode="BeliefPropagation", no_iterations=50 +) + s = siteinds(ψ) + + C = zeros((no_iterations)) + times_iters = zeros((no_iterations)) + times_gauging = zeros((no_iterations)) + + ψψ = norm_network(ψ) + ψinit = copy(ψ) + vertex_groups = nested_graph_leaf_vertices(partition(ψψ, group(v -> v[1], vertices(ψψ)))) + mts = message_tensors(partition(ψψ; subgraph_vertices=vertex_groups)) + bond_tensors = initialize_bond_tensors(ψ) + for e in edges(mts) + mts[e] = ITensorNetwork(dense(delta(inds(ITensors.contract(ITensor(mts[e])))))) + end + + for i in 1:no_iterations + println("On Iteration " * string(i)) + + if mode == "BeliefPropagation" + times_iters[i] = @elapsed mts, _ = belief_propagation_iteration( + ψψ, mts; contract_kwargs=(; alg="exact") + ) + times_gauging[i] = @elapsed ψ, bond_tensors = vidal_gauge(ψinit, mts) + elseif mode == "Eager" + times_iters[i] = @elapsed ψ, bond_tensors, mts = eager_gauging(ψ, bond_tensors, mts) + else + times_iters[i] = @elapsed begin + for e in edges(ψ) + ψ, bond_tensors = apply(e, ψ, bond_tensors; normalize=true, cutoff=1e-16) + end + end + end + + C[i] = vidal_itn_canonicalness(ψ, bond_tensors) + end + + simulation_times = cumsum(times_iters) + times_gauging + + return simulation_times, C +end + +L, χ = 10, 10 +g = named_grid((L, L)) +s = siteinds("S=1/2", g) +ψ = randomITensorNetwork(s; link_space=χ) +no_iterations = 30 + +BPG_simulation_times, BPG_Cs = benchmark_state_gauging(ψ; no_iterations) +Eager_simulation_times, Eager_Cs = benchmark_state_gauging(ψ; mode="Eager", no_iterations) +SU_simulation_times, SU_Cs = benchmark_state_gauging(ψ; mode="SU", no_iterations) + +epsilon = 1e-6 +println( + "Time for BPG to reach C < epsilon was " * + string(BPG_simulation_times[findfirst(x -> x < 0, BPG_Cs .- epsilon)]) * + " seconds", +) +println( + "Time for Eager to reach C < epsilon was " * + string(Eager_simulation_times[findfirst(x -> x < 0, Eager_Cs .- epsilon)]) * + " seconds", +) +println( + "Time for SU to reach C < epsilon was " * + string(SU_simulation_times[findfirst(x -> x < 0, SU_Cs .- epsilon)]) * + " seconds", +) diff --git a/src/abstractitensornetwork.jl b/src/abstractitensornetwork.jl index cdd742d9..4fb56a27 100644 --- a/src/abstractitensornetwork.jl +++ b/src/abstractitensornetwork.jl @@ -330,7 +330,14 @@ end adjoint(tn::Union{IndsNetwork,AbstractITensorNetwork}) = prime(tn) -dag(tn::AbstractITensorNetwork) = map_vertex_data(dag, tn) +#dag(tn::AbstractITensorNetwork) = map_vertex_data(dag, tn) +function dag(tn::AbstractITensorNetwork) + tndag = copy(tn) + for v in vertices(tndag) + setindex_preserve_graph!(tndag, dag(tndag[v]), v) + end + return tndag +end # TODO: should this make sure that internal indices # don't clash? @@ -646,6 +653,23 @@ function flatten_networks( return flatten_networks(flatten_networks(tn1, tn2; kwargs...), tn3, tn_tail...; kwargs...) end +#Ideally this will dispatch to inner_network but this is a temporary fast version for now +function norm_network(tn::AbstractITensorNetwork) + tnbra = rename_vertices(v -> (v, 1), data_graph(tn)) + tndag = copy(tn) + for v in vertices(tndag) + setindex_preserve_graph!(tndag, dag(tndag[v]), v) + end + tnket = rename_vertices(v -> (v, 2), data_graph(prime(tndag; sites=[]))) + tntn = ITensorNetwork(union(tnbra, tnket)) + for v in vertices(tn) + if !isempty(commoninds(tntn[(v, 1)], tntn[(v, 2)])) + add_edge!(tntn, (v, 1) => (v, 2)) + end + end + return tntn +end + # TODO: Use or replace with `flatten_networks` function inner_network( tn1::AbstractITensorNetwork, diff --git a/src/apply.jl b/src/apply.jl index 908e21d1..9b0edd3d 100644 --- a/src/apply.jl +++ b/src/apply.jl @@ -19,7 +19,7 @@ function ITensors.apply( if normalize oψᵥ ./= norm(oψᵥ) end - ψ[v⃗[1]] = oψᵥ + setindex_preserve_graph!(ψ, oψᵥ, v⃗[1]) elseif length(v⃗) == 2 e = v⃗[1] => v⃗[2] if !has_edge(ψ, e) @@ -72,8 +72,8 @@ function ITensors.apply( ψᵥ₂ ./= norm(ψᵥ₂) end - ψ[v⃗[1]] = ψᵥ₁ - ψ[v⃗[2]] = ψᵥ₂ + setindex_preserve_graph!(ψ, ψᵥ₁, v⃗[1]) + setindex_preserve_graph!(ψ, ψᵥ₂, v⃗[2]) elseif length(v⃗) < 1 error("Gate being applied does not share indices with tensor network.") @@ -123,15 +123,18 @@ end _gate_vertices(o::ITensor, ψ) = neighbor_vertices(ψ, o) _gate_vertices(o::AbstractEdge, ψ) = [src(o), dst(o)] -function _contract_factorized_gate(o::ITensor, ψv1, ψv2) - G1, G2 = factorize(o, Index[commonind(ψv1, o), commonind(ψv1, o)']; cutoff=1e-16) - ψv1 = noprime(ψv1 * G1) - ψv2 = noprime(ψv2 * G2) - return ψv1, ψv2 +function _contract_gate(o::ITensor, ψv1, ψv2) + Qᵥ₁, Rᵥ₁ = qr(ψv1, setdiff(uniqueinds(ψv1, ψv2), commoninds(ψv1, o))) + Qᵥ₂, Rᵥ₂ = qr(ψv2, setdiff(uniqueinds(ψv2, ψv1), commoninds(ψv2, o))) + theta = noprime(Rᵥ₁ * Rᵥ₂ * o) + return Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta end -function _contract_factorized_gate(o::AbstractEdge, ψv1, ψv2) - return ψv1, ψv2 +function _contract_gate(o::AbstractEdge, ψv1, ψv2) + Qᵥ₁, Rᵥ₁ = qr(ψv1, uniqueinds(ψv1, ψv2)) + Qᵥ₂, Rᵥ₂ = qr(ψv2, uniqueinds(ψv2, ψv1)) + theta = Rᵥ₁ * Rᵥ₂ + return Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta end #In the future we will try to unify this into apply() above but currently leave it mostly as a separate function @@ -149,7 +152,7 @@ function ITensors.apply( v⃗ = _gate_vertices(o, ψ) if length(v⃗) == 2 e = NamedEdge(v⃗[1] => v⃗[2]) - ψv1, ψv2 = copy(ψ[src(e)]), copy(ψ[dst(e)]) + ψv1, ψv2 = ψ[src(e)], ψ[dst(e)] e_ind = commonind(ψv1, ψv2) for vn in neighbors(ψ, src(e)) @@ -164,14 +167,9 @@ function ITensors.apply( end end - ψv1, ψv2 = _contract_factorized_gate(o, ψv1, ψv2) - ψv2 = noprime(ψv2 * bond_tensors[e]) - Qᵥ₁, Rᵥ₁ = factorize(ψv1, uniqueinds(ψv1, ψv2); cutoff=1e-16) - Qᵥ₂, Rᵥ₂ = factorize(ψv2, uniqueinds(ψv2, ψv1); cutoff=1e-16) - - theta = Rᵥ₁ * Rᵥ₂ + Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta = _contract_gate(o, ψv1, ψv2) U, S, V = ITensors.svd( theta, @@ -201,12 +199,13 @@ function ITensors.apply( end if normalize - normalize!(ψv1) - normalize!(ψv2) + ψv1 /= norm(ψv1) + ψv2 /= norm(ψv2) normalize!(bond_tensors[e]) end - ψ[src(e)], ψ[dst(e)] = ψv1, ψv2 + setindex_preserve_graph!(ψ, ψv1, src(e)) + setindex_preserve_graph!(ψ, ψv2, dst(e)) return ψ, bond_tensors diff --git a/src/gauging.jl b/src/gauging.jl index f0616d49..194dce98 100644 --- a/src/gauging.jl +++ b/src/gauging.jl @@ -23,6 +23,7 @@ function vidal_gauge( for e in edges(ψ_vidal) vsrc, vdst = src(e), dst(e) + ψvsrc, ψvdst = ψ_vidal[vsrc], ψ_vidal[vdst] s1, s2 = find_subgraph((vsrc, 1), mts), find_subgraph((vdst, 1), mts) edge_ind = commoninds(ψ_vidal[vsrc], ψ_vidal[vdst]) @@ -44,8 +45,7 @@ function vidal_gauge( inv_rootX = X_U * inv_rootX_D * prime(dag(X_U)) inv_rootY = Y_U * inv_rootY_D * prime(dag(Y_U)) - ψ_vidal[vsrc] = noprime(ψ_vidal[vsrc] * inv_rootX) - ψ_vidal[vdst] = noprime(ψ_vidal[vdst] * inv_rootY) + ψvsrc, ψvdst = noprime(ψvsrc * inv_rootX), noprime(ψvdst * inv_rootY) Ce = rootX * prime(bond_tensors[e]) replaceinds!(Ce, edge_ind'', edge_ind') @@ -55,9 +55,12 @@ function vidal_gauge( new_edge_ind = Index[Index(dim(commoninds(S, U)), tags(first(edge_ind)))] - ψ_vidal[vsrc] = replaceinds(ψ_vidal[vsrc] * U, commoninds(S, U), new_edge_ind) - ψ_vidal[vdst] = replaceinds(ψ_vidal[vdst], edge_ind, edge_ind_sim) - ψ_vidal[vdst] = replaceinds(ψ_vidal[vdst] * V, commoninds(V, S), new_edge_ind) + ψvsrc = replaceinds(ψvsrc * U, commoninds(S, U), new_edge_ind) + ψvdst = replaceinds(ψvdst, edge_ind, edge_ind_sim) + ψvdst = replaceinds(ψvdst * V, commoninds(V, S), new_edge_ind) + + setindex_preserve_graph!(ψ_vidal, ψvsrc, vsrc) + setindex_preserve_graph!(ψ_vidal, ψvdst, vdst) S = replaceinds( S, @@ -93,7 +96,7 @@ function vidal_gauge( target_canonicalness::Union{Nothing,Float64}=nothing, svd_kwargs..., ) - ψψ = ψ ⊗ prime(dag(ψ); sites=[]) + ψψ = norm_network(ψ) Z = partition(ψψ; subgraph_vertices=collect(values(group(v -> v[1], vertices(ψψ))))) mts = message_tensors(Z) @@ -108,17 +111,18 @@ end """Transform from an ITensor in the Vidal Gauge (bond tensors) to the Symmetric Gauge (message tensors)""" function vidal_to_symmetric_gauge(ψ::ITensorNetwork, bond_tensors::DataGraph) ψsymm = copy(ψ) - ψψsymm = ψsymm ⊗ prime(dag(ψsymm); sites=[]) + ψψsymm = norm_network(ψsymm) Z = partition( ψψsymm; subgraph_vertices=collect(values(group(v -> v[1], vertices(ψψsymm)))) ) ψsymm_mts = message_tensors_skeleton(Z) for e in edges(ψsymm) - s1, s2 = find_subgraph((src(e), 1), ψsymm_mts), find_subgraph((dst(e), 1), ψsymm_mts) + vsrc, vdst = src(e), dst(e) + s1, s2 = find_subgraph((vsrc, 1), ψsymm_mts), find_subgraph((vdst, 1), ψsymm_mts) root_S = sqrt_diag(bond_tensors[e]) - ψsymm[src(e)] = noprime(ψsymm[src(e)] * root_S) - ψsymm[dst(e)] = noprime(ψsymm[dst(e)] * root_S) + setindex_preserve_graph!(ψsymm, noprime(root_S * ψsymm[vsrc]), vsrc) + setindex_preserve_graph!(ψsymm, noprime(root_S * ψsymm[vdst]), vdst) ψsymm_mts[s1 => s2], ψsymm_mts[s2 => s1] = ITensorNetwork(bond_tensors[e]), ITensorNetwork(bond_tensors[e]) @@ -157,11 +161,12 @@ function symmetric_to_vidal_gauge( ψ_vidal = copy(ψ) for e in edges(ψ) - s1, s2 = find_subgraph((src(e), 1), mts), find_subgraph((dst(e), 1), mts) + vsrc, vdst = src(e), dst(e) + s1, s2 = find_subgraph((vsrc, 1), mts), find_subgraph((vdst, 1), mts) bond_tensors[e] = ITensor(mts[s1 => s2]) invroot_S = invsqrt_diag(map_diag(x -> x + regularization, bond_tensors[e])) - ψ_vidal[src(e)] = noprime(invroot_S * ψ_vidal[src(e)]) - ψ_vidal[dst(e)] = noprime(invroot_S * ψ_vidal[dst(e)]) + setindex_preserve_graph!(ψ_vidal, noprime(invroot_S * ψ_vidal[vsrc]), vsrc) + setindex_preserve_graph!(ψ_vidal, noprime(invroot_S * ψ_vidal[vdst]), vdst) end return ψ_vidal, bond_tensors diff --git a/test/test_apply.jl b/test/test_apply.jl index 3579dca9..64909adf 100644 --- a/test/test_apply.jl +++ b/test/test_apply.jl @@ -6,7 +6,8 @@ using ITensorNetworks: message_tensors, nested_graph_leaf_vertices, vidal_gauge, - vidal_to_symmetric_gauge + vidal_to_symmetric_gauge, + norm_network using Test using Compat using ITensors @@ -26,7 +27,7 @@ using SplitApplyCombine ψ = randomITensorNetwork(s; link_space=χ) v1, v2 = (2, 2), (1, 2) - ψψ = ψ ⊗ prime(dag(ψ); sites=[]) + ψψ = norm_network(ψ) #Simple Belief Propagation Grouping vertex_groupsSBP = nested_graph_leaf_vertices(