diff --git a/src/ITensorNetworks.jl b/src/ITensorNetworks.jl index fb4c6711..35f5702e 100644 --- a/src/ITensorNetworks.jl +++ b/src/ITensorNetworks.jl @@ -78,7 +78,6 @@ include("opsum.jl") include("sitetype.jl") include("abstractitensornetwork.jl") include("contraction_sequences.jl") -include("apply.jl") include("expect.jl") include("models.jl") include("tebd.jl") @@ -95,11 +94,12 @@ include("contract.jl") include("utility.jl") include("specialitensornetworks.jl") include("boundarymps.jl") -include(joinpath("beliefpropagation", "beliefpropagation.jl")) -include(joinpath("beliefpropagation", "beliefpropagation_schedule.jl")) +include("partitioneditensornetwork.jl") +include("edge_sequences.jl") include(joinpath("formnetworks", "abstractformnetwork.jl")) include(joinpath("formnetworks", "bilinearformnetwork.jl")) include(joinpath("formnetworks", "quadraticformnetwork.jl")) +include(joinpath("caches", "beliefpropagationcache.jl")) include("contraction_tree_to_graph.jl") include("gauging.jl") include("utils.jl") @@ -128,6 +128,7 @@ include(joinpath("treetensornetworks", "solvers", "dmrg_x.jl")) include(joinpath("treetensornetworks", "solvers", "contract.jl")) include(joinpath("treetensornetworks", "solvers", "linsolve.jl")) include(joinpath("treetensornetworks", "solvers", "tree_sweeping.jl")) +include("apply.jl") include("exports.jl") diff --git a/src/apply.jl b/src/apply.jl index c7745973..89edf656 100644 --- a/src/apply.jl +++ b/src/apply.jl @@ -169,7 +169,7 @@ function simple_update_bp(o, ψ, v⃗; envs, (observer!)=nothing, apply_kwargs.. end function ITensors.apply( - o::ITensor, + o, ψ::AbstractITensorNetwork; envs=ITensor[], normalize=false, @@ -297,15 +297,9 @@ end #In the future we will try to unify this into apply() above but currently leave it mostly as a separate function """Apply() function for an ITN in the Vidal Gauge. Hence the bond tensors are required. Gate does not necessarily need to be passed. Can supply an edge to do an identity update instead. Uses Simple Update procedure assuming gate is two-site""" -function vidal_apply( - o::Union{ITensor,NamedEdge}, - ψ::AbstractITensorNetwork, - bond_tensors::DataGraph; - normalize=false, - apply_kwargs..., -) - ψ = copy(ψ) - bond_tensors = copy(bond_tensors) +function ITensors.apply(o, ψ::VidalITensorNetwork; normalize=false, apply_kwargs...) + updated_ψ = copy(site_tensors(ψ)) + updated_bond_tensors = copy(bond_tensors(ψ)) v⃗ = _gate_vertices(o, ψ) if length(v⃗) == 2 e = NamedEdge(v⃗[1] => v⃗[2]) @@ -314,17 +308,17 @@ function vidal_apply( for vn in neighbors(ψ, src(e)) if (vn != dst(e)) - ψv1 = noprime(ψv1 * bond_tensors[vn => src(e)]) + ψv1 = noprime(ψv1 * bond_tensor(ψ, vn => src(e))) end end for vn in neighbors(ψ, dst(e)) if (vn != src(e)) - ψv2 = noprime(ψv2 * bond_tensors[vn => dst(e)]) + ψv2 = noprime(ψv2 * bond_tensor(ψ, vn => dst(e))) end end - Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta = _contract_gate(o, ψv1, bond_tensors[e], ψv2) + Qᵥ₁, Rᵥ₁, Qᵥ₂, Rᵥ₂, theta = _contract_gate(o, ψv1, bond_tensor(ψ, e), ψv2) U, S, V = ITensors.svd( theta, @@ -339,34 +333,34 @@ function vidal_apply( S = replaceind(S, ind_to_replace => ind_to_replace_with') V = replaceind(V, ind_to_replace => ind_to_replace_with) - ψv1, bond_tensors[e], ψv2 = U * Qᵥ₁, S, V * Qᵥ₂ + ψv1, updated_bond_tensors[e], ψv2 = U * Qᵥ₁, S, V * Qᵥ₂ for vn in neighbors(ψ, src(e)) if (vn != dst(e)) - ψv1 = noprime(ψv1 * inv_diag(bond_tensors[vn => src(e)])) + ψv1 = noprime(ψv1 * inv_diag(bond_tensor(ψ, vn => src(e)))) end end for vn in neighbors(ψ, dst(e)) if (vn != src(e)) - ψv2 = noprime(ψv2 * inv_diag(bond_tensors[vn => dst(e)])) + ψv2 = noprime(ψv2 * inv_diag(bond_tensor(ψ, vn => dst(e)))) end end if normalize ψv1 /= norm(ψv1) ψv2 /= norm(ψv2) - normalize!(bond_tensors[e]) + updated_bond_tensors[e] /= norm(updated_bond_tensors[e]) end - setindex_preserve_graph!(ψ, ψv1, src(e)) - setindex_preserve_graph!(ψ, ψv2, dst(e)) + setindex_preserve_graph!(updated_ψ, ψv1, src(e)) + setindex_preserve_graph!(updated_ψ, ψv2, dst(e)) - return ψ, bond_tensors + return VidalITensorNetwork(updated_ψ, updated_bond_tensors) else - ψ = ITensors.apply(o, ψ; normalize) - return ψ, bond_tensors + updated_ψ = ITensors.apply(o, updated_ψ; normalize) + return VidalITensorNetwork(ψ, updated_bond_tensors) end end diff --git a/src/beliefpropagation/beliefpropagation.jl b/src/beliefpropagation/beliefpropagation.jl deleted file mode 100644 index 3b16e3b3..00000000 --- a/src/beliefpropagation/beliefpropagation.jl +++ /dev/null @@ -1,143 +0,0 @@ -default_mt_constructor(inds_e) = ITensor[denseblocks(delta(inds_e))] -default_bp_cache(ptn::PartitionedGraph) = Dict() -function default_contractor(contract_list::Vector{ITensor}; kwargs...) - return contract_exact(contract_list; kwargs...) -end -default_contractor_kwargs() = (; normalize=true, contraction_sequence_alg="optimal") - -function message_tensor( - ptn::PartitionedGraph, edge::PartitionEdge; mt_constructor=default_mt_constructor -) - src_e_itn = subgraph(ptn, src(edge)) - dst_e_itn = subgraph(ptn, dst(edge)) - inds_e = commoninds(src_e_itn, dst_e_itn) - return mt_constructor(inds_e) -end - -""" -Do a single update of a message tensor using the current subgraph and the incoming mts -""" -function update_message_tensor( - ptn::PartitionedGraph, - edge::PartitionEdge, - mts; - contractor=default_contractor, - contractor_kwargs=default_contractor_kwargs(), - mt_constructor=default_mt_constructor, -) - pb_edges = partitionedges(ptn, boundary_edges(ptn, vertices(ptn, src(edge)); dir=:in)) - - incoming_messages = [ - e_in ∈ keys(mts) ? mts[e_in] : message_tensor(ptn, e_in; mt_constructor) for - e_in in setdiff(pb_edges, [reverse(edge)]) - ] - incoming_messages = reduce(vcat, incoming_messages; init=ITensor[]) - - contract_list = ITensor[ - incoming_messages - Vector{ITensor}(subgraph(ptn, src(edge))) - ] - - return contractor(contract_list; contractor_kwargs...) -end - -""" -Do a sequential update of message tensors on `edges` for a given ITensornetwork and its partition into sub graphs -""" -function belief_propagation_iteration( - ptn::PartitionedGraph, mts, edges::Vector{<:PartitionEdge}; compute_norm=false, kwargs... -) - new_mts = copy(mts) - c = 0 - for e in edges - new_mts[e] = update_message_tensor(ptn, e, new_mts; kwargs...) - - if compute_norm - LHS = e ∈ keys(mts) ? contract(mts[e]) : contract(message_tensor(ptn, e)) - RHS = contract(new_mts[e]) - #This line only makes sense if the message tensors are rank 2??? Should fix this. - LHS /= sum(diag(LHS)) - RHS /= sum(diag(RHS)) - c += 0.5 * norm(denseblocks(LHS) - denseblocks(RHS)) - end - end - return new_mts, c / (length(edges)) -end - -""" -Do parallel updates between groups of edges of all message tensors for a given ITensornetwork and its partition into sub graphs. -Currently we send the full message tensor data struct to belief_propagation_iteration for each subgraph. But really we only need the -mts relevant to that subgraph. -""" -function belief_propagation_iteration( - ptn::PartitionedGraph, mts, edge_groups::Vector{<:Vector{<:PartitionEdge}}; kwargs... -) - new_mts = copy(mts) - c = 0 - for edges in edge_groups - updated_mts, ct = belief_propagation_iteration(ptn, mts, edges; kwargs...) - for e in edges - new_mts[e] = updated_mts[e] - end - c += ct - end - return new_mts, c / (length(edge_groups)) -end - -function belief_propagation_iteration( - ptn::PartitionedGraph, mts; edges=default_edge_sequence(ptn), kwargs... -) - return belief_propagation_iteration(ptn, mts, edges; kwargs...) -end - -function belief_propagation( - ptn::PartitionedGraph, - mts; - niters=default_bp_niters(partitioned_graph(ptn)), - target_precision=nothing, - edges=default_edge_sequence(ptn), - verbose=false, - kwargs..., -) - compute_norm = !isnothing(target_precision) - if isnothing(niters) - error("You need to specify a number of iterations for BP!") - end - for i in 1:niters - mts, c = belief_propagation_iteration(ptn, mts, edges; compute_norm, kwargs...) - if compute_norm && c <= target_precision - if verbose - println("BP converged to desired precision after $i iterations.") - end - break - end - end - return mts -end - -function belief_propagation(ptn::PartitionedGraph; bp_cache=default_bp_cache, kwargs...) - mts = bp_cache(ptn) - return belief_propagation(ptn, mts; kwargs...) -end -""" -Given a subet of partitionvertices of a ptn get the incoming message tensors to that region -""" -function environment_tensors(ptn::PartitionedGraph, mts, verts::Vector) - partition_verts = partitionvertices(ptn, verts) - central_verts = vertices(ptn, partition_verts) - - pedges = partitionedges(ptn, boundary_edges(ptn, central_verts; dir=:in)) - env_tensors = [mts[e] for e in pedges] - env_tensors = reduce(vcat, env_tensors; init=ITensor[]) - central_tensors = ITensor[ - (unpartitioned_graph(ptn))[v] for v in setdiff(central_verts, verts) - ] - - return vcat(env_tensors, central_tensors) -end - -function environment_tensors( - ptn::PartitionedGraph, mts, partition_verts::Vector{<:PartitionVertex} -) - return environment_tensors(ptn, mts, vertices(ptn, partition_verts)) -end diff --git a/src/caches/beliefpropagationcache.jl b/src/caches/beliefpropagationcache.jl new file mode 100644 index 00000000..f5337784 --- /dev/null +++ b/src/caches/beliefpropagationcache.jl @@ -0,0 +1,234 @@ +default_message(inds_e) = ITensor[denseblocks(delta(inds_e))] +default_messages(ptn::PartitionedGraph) = Dictionary() +function default_message_update(contract_list::Vector{ITensor}; kwargs...) + return contract_exact(contract_list; kwargs...) +end +default_message_update_kwargs() = (; normalize=true, contraction_sequence_alg="optimal") +@traitfn default_bp_maxiter(g::::(!IsDirected)) = is_tree(g) ? 1 : nothing +@traitfn function default_bp_maxiter(g::::IsDirected) + return default_bp_maxiter(undirected_graph(underlying_graph(g))) +end +function message_diff(message_a::Vector{ITensor}, message_b::Vector{ITensor}) + lhs, rhs = contract(message_a), contract(message_b) + return 0.5 * + norm((denseblocks(lhs) / sum(diag(lhs))) - (denseblocks(rhs) / sum(diag(rhs)))) +end + +struct BeliefPropagationCache{PTN,MTS,DM} + partitioned_itensornetwork::PTN + messages::MTS + default_message::DM +end + +#Constructors... +function BeliefPropagationCache( + ptn::PartitionedGraph; messages=default_messages(ptn), default_message=default_message +) + return BeliefPropagationCache(ptn, messages, default_message) +end + +function BeliefPropagationCache(tn::ITensorNetwork, partitioned_vertices; kwargs...) + ptn = PartitionedGraph(tn, partitioned_vertices) + return BeliefPropagationCache(ptn; kwargs...) +end + +function partitioned_itensornetwork(bp_cache::BeliefPropagationCache) + return bp_cache.partitioned_itensornetwork +end +messages(bp_cache::BeliefPropagationCache) = bp_cache.messages +default_message(bp_cache::BeliefPropagationCache) = bp_cache.default_message +function tensornetwork(bp_cache::BeliefPropagationCache) + return unpartitioned_graph(partitioned_itensornetwork(bp_cache)) +end + +#Forward from partitioned graph +for f in [ + :(NamedGraphs.partitioned_graph), + :(NamedGraphs.partitionedge), + :(NamedGraphs.partitionvertices), + :(NamedGraphs.vertices), + :(NamedGraphs.boundary_partitionedges), + :linkinds, +] + @eval begin + function $f(bp_cache::BeliefPropagationCache, args...; kwargs...) + return $f(partitioned_itensornetwork(bp_cache), args...; kwargs...) + end + end +end + +function default_message(bp_cache::BeliefPropagationCache, edge::PartitionEdge) + return default_message(bp_cache)(linkinds(bp_cache, edge)) +end + +function message(bp_cache::BeliefPropagationCache, edge::PartitionEdge) + mts = messages(bp_cache) + return get(mts, edge, default_message(bp_cache, edge)) +end +function messages( + bp_cache::BeliefPropagationCache, edges::Vector{<:PartitionEdge}; kwargs... +) + return [message(bp_cache, edge; kwargs...) for edge in edges] +end + +function copy(bp_cache::BeliefPropagationCache) + return BeliefPropagationCache( + copy(partitioned_itensornetwork(bp_cache)), + copy(messages(bp_cache)), + default_message(bp_cache), + ) +end + +function default_bp_maxiter(bp_cache::BeliefPropagationCache) + return default_bp_maxiter(partitioned_graph(bp_cache)) +end +function default_edge_sequence(bp_cache::BeliefPropagationCache) + return default_edge_sequence(partitioned_itensornetwork(bp_cache)) +end + +function set_messages(cache::BeliefPropagationCache, messages) + return BeliefPropagationCache( + partitioned_itensornetwork(cache), messages, default_message(cache) + ) +end + +function incoming_messages( + bp_cache::BeliefPropagationCache, + partition_vertices::Vector{<:PartitionVertex}; + ignore_edges=PartitionEdge[], +) + bpes = boundary_partitionedges(bp_cache, partition_vertices; dir=:in) + ms = messages(bp_cache, setdiff(bpes, ignore_edges)) + return reduce(vcat, ms; init=[]) +end + +function incoming_messages( + bp_cache::BeliefPropagationCache, partition_vertex::PartitionVertex; kwargs... +) + return incoming_messages(bp_cache, [partition_vertex]; kwargs...) +end + +function incoming_messages(bp_cache::BeliefPropagationCache, verts::Vector) + partition_verts = partitionvertices(bp_cache, verts) + messages = incoming_messages(bp_cache, partition_verts) + central_tensors = ITensor[ + tensornetwork(bp_cache)[v] for v in setdiff(vertices(bp_cache, partition_verts), verts) + ] + return vcat(messages, central_tensors) +end + +function factor(bp_cache::BeliefPropagationCache, vertex::PartitionVertex) + ptn = partitioned_itensornetwork(bp_cache) + return Vector{ITensor}(subgraph(ptn, vertex)) +end + +""" +Compute message tensor as product of incoming mts and local state +""" +function update_message( + bp_cache::BeliefPropagationCache, + edge::PartitionEdge; + message_update=default_message_update, + message_update_kwargs=default_message_update_kwargs(), +) + vertex = src(edge) + messages = incoming_messages(bp_cache, vertex; ignore_edges=PartitionEdge[reverse(edge)]) + state = factor(bp_cache, vertex) + + return message_update(ITensor[messages; state]; message_update_kwargs...) +end + +""" +Do a sequential update of the message tensors on `edges` +""" +function update( + bp_cache::BeliefPropagationCache, + edges::Vector{<:PartitionEdge}; + (update_diff!)=nothing, + kwargs..., +) + bp_cache_updated = copy(bp_cache) + mts = messages(bp_cache_updated) + for e in edges + set!(mts, e, update_message(bp_cache_updated, e; kwargs...)) + if !isnothing(update_diff!) + update_diff![] += message_diff(message(bp_cache, e), mts[e]) + end + end + return bp_cache_updated +end + +""" +Update the message tensor on a single edge +""" +function update(bp_cache::BeliefPropagationCache, edge::PartitionEdge; kwargs...) + return update(bp_cache, [edge]; kwargs...) +end + +""" +Do parallel updates between groups of edges of all message tensors +Currently we send the full message tensor data struct to update for each edge_group. But really we only need the +mts relevant to that group. +""" +function update( + bp_cache::BeliefPropagationCache, + edge_groups::Vector{<:Vector{<:PartitionEdge}}; + kwargs..., +) + new_mts = copy(messages(bp_cache)) + for edges in edge_groups + bp_cache_t = update(bp_cache, edges; kwargs...) + for e in edges + new_mts[e] = message(bp_cache_t, e) + end + end + return set_messages(bp_cache, new_mts) +end + +""" +More generic interface for update, with default params +""" +function update( + bp_cache::BeliefPropagationCache; + edges=default_edge_sequence(bp_cache), + maxiter=default_bp_maxiter(bp_cache), + tol=nothing, + verbose=false, + kwargs..., +) + compute_error = !isnothing(tol) + diff = compute_error ? Ref(0.0) : nothing + if isnothing(maxiter) + error("You need to specify a number of iterations for BP!") + end + for i in 1:maxiter + bp_cache = update(bp_cache, edges; (update_diff!)=diff, kwargs...) + if compute_error && (diff.x / length(edges)) <= tol + if verbose + println("BP converged to desired precision after $i iterations.") + end + break + end + end + return bp_cache +end + +""" +Update the tensornetwork inside the cache +""" +function update_factors( + bp_cache::BeliefPropagationCache, vertices::Vector, factors::Vector{ITensor} +) + bp_cache = copy(bp_cache) + tn = tensornetwork(bp_cache) + + for (vertex, factor) in zip(vertices, factors) + # TODO: Add a check that this preserves the graph structure. + setindex_preserve_graph!(tn, factor, vertex) + end + return bp_cache +end + +function update_factor(bp_cache, vertex, factor) + return update_factors(bp_cache, [vertex], ITensor[factor]) +end diff --git a/src/beliefpropagation/beliefpropagation_schedule.jl b/src/edge_sequences.jl similarity index 86% rename from src/beliefpropagation/beliefpropagation_schedule.jl rename to src/edge_sequences.jl index a56814cd..786c1606 100644 --- a/src/beliefpropagation/beliefpropagation_schedule.jl +++ b/src/edge_sequences.jl @@ -3,11 +3,6 @@ function default_edge_sequence(pg::PartitionedGraph) return PartitionEdge.(edge_sequence(partitioned_graph(pg))) end -@traitfn default_bp_niters(g::::(!IsDirected)) = is_tree(g) ? 1 : nothing -@traitfn function default_bp_niters(g::::IsDirected) - return default_bp_niters(undirected_graph(underlying_graph(g))) -end - @traitfn function edge_sequence( g::::(!IsDirected); alg=default_edge_sequence_alg(), kwargs... ) diff --git a/src/gauging.jl b/src/gauging.jl index c6d3c3f6..41bd02f0 100644 --- a/src/gauging.jl +++ b/src/gauging.jl @@ -1,39 +1,88 @@ -"""initialize bond tensors of an ITN to identity matrices""" -function initialize_bond_tensors(ψ::ITensorNetwork; index_map=prime) - bond_tensors = DataGraph{vertextype(ψ),Nothing,ITensor}(underlying_graph(ψ)) +function default_bond_tensors(ψ::ITensorNetwork) + return DataGraph{vertextype(ψ),Nothing,ITensor}(underlying_graph(ψ)) +end + +struct VidalITensorNetwork{V,BTS} <: AbstractITensorNetwork{V} + itensornetwork::ITensorNetwork{V} + bond_tensors::BTS +end + +site_tensors(ψ::VidalITensorNetwork) = ψ.itensornetwork +bond_tensors(ψ::VidalITensorNetwork) = ψ.bond_tensors +bond_tensor(ψ::VidalITensorNetwork, e) = bond_tensors(ψ)[e] + +function data_graph_type(TN::Type{<:VidalITensorNetwork}) + return data_graph_type(fieldtype(TN, :itensornetwork)) +end +data_graph(ψ::VidalITensorNetwork) = data_graph(site_tensors(ψ)) +function copy(ψ::VidalITensorNetwork) + return VidalITensorNetwork(copy(site_tensors(ψ)), copy(bond_tensors(ψ))) +end + +function default_norm_cache(ψ::ITensorNetwork) + ψψ = norm_network(ψ) + return BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) +end +default_cache_update_kwargs(cache) = (; maxiter=20, tol=1e-5) + +function ITensorNetwork( + ψ_vidal::VidalITensorNetwork; (cache!)=nothing, update_gauge=false, update_kwargs... +) + if update_gauge + ψ_vidal = update(ψ_vidal; update_kwargs...) + end + + ψ = copy(site_tensors(ψ_vidal)) for e in edges(ψ) - index = commoninds(ψ[src(e)], ψ[dst(e)]) - bond_tensors[e] = denseblocks(delta(index, index_map(index))) + vsrc, vdst = src(e), dst(e) + root_S = sqrt_diag(bond_tensor(ψ_vidal, e)) + setindex_preserve_graph!(ψ, noprime(root_S * ψ[vsrc]), vsrc) + setindex_preserve_graph!(ψ, noprime(root_S * ψ[vdst]), vdst) end - return bond_tensors + if !isnothing(cache!) + bp_cache = default_norm_cache(ψ) + mts = messages(bp_cache) + + for e in edges(ψ) + vsrc, vdst = src(e), dst(e) + pe = partitionedge(bp_cache, (vsrc, 1) => (vdst, 1)) + set!(mts, pe, copy(ITensor[dense(bond_tensor(ψ_vidal, e))])) + set!(mts, reverse(pe), copy(ITensor[dense(bond_tensor(ψ_vidal, e))])) + end + + bp_cache = set_messages(bp_cache, mts) + cache![] = bp_cache + end + + return ψ end -"""Use an ITensorNetwork ψ, its bond tensors and gauging mts to put ψ into the vidal gauge, return the bond tensors and ψ_vidal.""" -function vidal_gauge( - ψ::ITensorNetwork, - pψψ::PartitionedGraph, - mts, - bond_tensors::DataGraph; - eigen_message_tensor_cutoff=10 * eps(real(scalartype(ψ))), +"""Use an ITensorNetwork ψ, its bond tensors and belief propagation cache to put ψ into the vidal gauge, return the bond tensors and updated_ψ.""" +function vidalitensornetwork_preserve_cache( + ψ::ITensorNetwork; + cache=default_norm_cache(ψ), + bond_tensors=default_bond_tensors, + message_cutoff=10 * eps(real(scalartype(ψ))), regularization=10 * eps(real(scalartype(ψ))), edges=NamedGraphs.edges(ψ), svd_kwargs..., ) - ψ_vidal = copy(ψ) + ψ_vidal_site_tensors = copy(ψ) + ψ_vidal_bond_tensors = bond_tensors(ψ) for e in edges vsrc, vdst = src(e), dst(e) - ψvsrc, ψvdst = ψ_vidal[vsrc], ψ_vidal[vdst] + ψvsrc, ψvdst = ψ_vidal_site_tensors[vsrc], ψ_vidal_site_tensors[vdst] - pe = partitionedge(pψψ, (vsrc, 1) => (vdst, 1)) + pe = partitionedge(cache, (vsrc, 1) => (vdst, 1)) edge_ind = commoninds(ψvsrc, ψvdst) edge_ind_sim = sim(edge_ind) - X_D, X_U = eigen(only(mts[pe]); ishermitian=true, cutoff=eigen_message_tensor_cutoff) + X_D, X_U = eigen(only(message(cache, pe)); ishermitian=true, cutoff=message_cutoff) Y_D, Y_U = eigen( - only(mts[reverse(pe)]); ishermitian=true, cutoff=eigen_message_tensor_cutoff + only(message(cache, reverse(pe))); ishermitian=true, cutoff=message_cutoff ) X_D, Y_D = map_diag(x -> x + regularization, X_D), map_diag(x -> x + regularization, Y_D) @@ -47,8 +96,7 @@ function vidal_gauge( ψvsrc, ψvdst = noprime(ψvsrc * inv_rootX), noprime(ψvdst * inv_rootY) - Ce = rootX * prime(bond_tensors[e]) - replaceinds!(Ce, edge_ind'', edge_ind') + Ce = rootX Ce = Ce * replaceinds(rootY, edge_ind, edge_ind_sim) U, S, V = svd(Ce, edge_ind; svd_kwargs...) @@ -59,172 +107,71 @@ function vidal_gauge( ψ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) + setindex_preserve_graph!(ψ_vidal_site_tensors, ψvsrc, vsrc) + setindex_preserve_graph!(ψ_vidal_site_tensors, ψvdst, vdst) S = replaceinds( S, [commoninds(S, U)..., commoninds(S, V)...] => [new_edge_ind..., prime(new_edge_ind)...], ) - bond_tensors[e] = S + ψ_vidal_bond_tensors[e] = S end - return ψ_vidal, bond_tensors -end - -"""Use an ITensorNetwork ψ in the symmetric gauge and its mts to put ψ into the vidal gauge. Return the bond tensors and ψ_vidal.""" -function vidal_gauge( - ψ::ITensorNetwork, - pψψ::PartitionedGraph, - mts; - eigen_message_tensor_cutoff=10 * eps(real(scalartype(ψ))), - regularization=10 * eps(real(scalartype(ψ))), - edges=NamedGraphs.edges(ψ), - svd_kwargs..., -) - bond_tensors = initialize_bond_tensors(ψ) - return vidal_gauge( - ψ, - pψψ, - mts, - bond_tensors; - eigen_message_tensor_cutoff, - regularization, - edges, - svd_kwargs..., - ) + return VidalITensorNetwork(ψ_vidal_site_tensors, ψ_vidal_bond_tensors) end -"""Put an ITensorNetwork into the vidal gauge (by computing the message tensors), return the network and the bond tensors. Will also return the mts that were constructed""" -function vidal_gauge( +function VidalITensorNetwork( ψ::ITensorNetwork; - eigen_message_tensor_cutoff=10 * eps(real(scalartype(ψ))), - regularization=10 * eps(real(scalartype(ψ))), - niters=30, - target_canonicalness::Union{Nothing,Float64}=nothing, - verbose=false, - svd_kwargs..., + (cache!)=nothing, + update_cache=isnothing(cache!), + cache_update_kwargs=default_cache_update_kwargs(cache!), + kwargs..., ) - ψψ = norm_network(ψ) - pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) - - mts = belief_propagation(pψψ; niters, target_precision=target_canonicalness, verbose) - return vidal_gauge( - ψ, pψψ, mts; eigen_message_tensor_cutoff, regularization, svd_kwargs... - ) -end - -"""Transform from an ITensor in the Vidal Gauge (bond tensors) to the Symmetric Gauge (partitionedgraph, message tensors)""" -function vidal_to_symmetric_gauge(ψ::ITensorNetwork, bond_tensors::DataGraph) - ψsymm = copy(ψ) - ψψsymm = norm_network(ψsymm) - pψψsymm = PartitionedGraph(ψψsymm, group(v -> v[1], vertices(ψψsymm))) - ψsymm_mts = default_bp_cache(pψψsymm) - - for e in edges(ψsymm) - vsrc, vdst = src(e), dst(e) - pe = partitionedge(pψψsymm, NamedEdge((vsrc, 1) => (vdst, 1))) - root_S = sqrt_diag(bond_tensors[e]) - setindex_preserve_graph!(ψsymm, noprime(root_S * ψsymm[vsrc]), vsrc) - setindex_preserve_graph!(ψsymm, noprime(root_S * ψsymm[vdst]), vdst) - - ψsymm_mts[pe], ψsymm_mts[reverse(pe)] = copy(ITensor[dense(bond_tensors[e])]), - copy(ITensor[dense(bond_tensors[e])]) + if isnothing(cache!) + cache! = Ref(default_norm_cache(ψ)) end - - ψψsymm = norm_network(ψsymm) - pψψsymm = PartitionedGraph(ψψsymm, group(v -> v[1], vertices(ψψsymm))) - - return ψsymm, pψψsymm, ψsymm_mts + cache![] = update(cache![]; cache_update_kwargs...) + return vidalitensornetwork_preserve_cache(ψ; cache=cache![], kwargs...) end -"""Put an ITensorNetwork into the symmetric gauge and also return the message tensors (which are the diagonal bond matrices from the Vidal Gauge)""" -function symmetric_gauge( - ψ::ITensorNetwork; - eigen_message_tensor_cutoff=10 * eps(real(scalartype(ψ))), - regularization=10 * eps(real(scalartype(ψ))), - niters=30, - target_canonicalness::Union{Nothing,Float64}=nothing, - svd_kwargs..., -) - ψ_vidal, bond_tensors = vidal_gauge( - ψ; - eigen_message_tensor_cutoff, - regularization, - niters, - target_canonicalness, - svd_kwargs..., - ) - - return vidal_to_symmetric_gauge(ψ_vidal, bond_tensors) +function update(ψ::VidalITensorNetwork; kwargs...) + return VidalITensorNetwork(ITensorNetwork(ψ; update_gauge=false); kwargs...) end -"""Transform from the Symmetric Gauge (message tensors) to the Vidal Gauge (bond tensors)""" -function symmetric_to_vidal_gauge( - ψ::ITensorNetwork, - pψψ::PartitionedGraph, - mts; - regularization=10 * eps(real(scalartype(ψ))), -) - bond_tensors = DataGraph{vertextype(ψ),Nothing,ITensor}(underlying_graph(ψ)) +"""Function to construct the 'isometry' of a state in the Vidal Gauge on the given edge""" +function vidal_gauge_isometry(ψ::VidalITensorNetwork, edge) + vsrc, vdst = src(edge), dst(edge) + ψ_vsrc = copy(ψ[vsrc]) - ψ_vidal = copy(ψ) - - for e in edges(ψ) - vsrc, vdst = src(e), dst(e) - pe = partitionedge(pψψ, NamedEdge((vsrc, 1) => (vdst, 1))) - bond_tensors[e], bond_tensors[reverse(e)] = only(mts[pe]), only(mts[pe]) - invroot_S = invsqrt_diag(map_diag(x -> x + regularization, bond_tensors[e])) - setindex_preserve_graph!(ψ_vidal, noprime(invroot_S * ψ_vidal[vsrc]), vsrc) - setindex_preserve_graph!(ψ_vidal, noprime(invroot_S * ψ_vidal[vdst]), vdst) + for vn in setdiff(neighbors(ψ, vsrc), [vdst]) + ψ_vsrc = noprime(ψ_vsrc * bond_tensor(ψ, vn => vsrc)) end - return ψ_vidal, bond_tensors -end - -"""Function to measure the 'isometries' of a state in the Vidal Gauge""" -function vidal_itn_isometries( - ψ::ITensorNetwork, - bond_tensors::DataGraph; - edges=vcat(NamedGraphs.edges(ψ), reverse.(NamedGraphs.edges(ψ))), -) - isometries = Dict() + ψ_vsrcdag = dag(ψ_vsrc) + ψ_vsrcdag = replaceind(ψ_vsrcdag, commonind(ψ_vsrc, ψ[vdst]), commonind(ψ_vsrc, ψ[vdst])') - for e in edges - vsrc, vdst = src(e), dst(e) - ψv = copy(ψ[vsrc]) - for vn in setdiff(neighbors(ψ, vsrc), [vdst]) - ψv = noprime(ψv * bond_tensors[vn => vsrc]) - end + return ψ_vsrcdag * ψ_vsrc +end - ψvdag = dag(ψv) - replaceind!(ψvdag, commonind(ψv, ψ[vdst]), commonind(ψv, ψ[vdst])') - isometries[e] = ψvdag * ψv - end +function vidal_gauge_isometries(ψ::VidalITensorNetwork, edges::Vector) + return Dict([e => vidal_gauge_isometry(ψ, e) for e in edges]) +end - return isometries +function vidal_gauge_isometries(ψ::VidalITensorNetwork) + return vidal_gauge_isometries( + ψ, vcat(NamedGraphs.edges(ψ), reverse.(NamedGraphs.edges(ψ))) + ) end -"""Function to measure the 'canonicalness' of a state in the Vidal Gauge""" -function vidal_itn_canonicalness(ψ::ITensorNetwork, bond_tensors::DataGraph) +"""Function to measure the 'distance' of a state from the Vidal Gauge""" +function gauge_error(ψ::VidalITensorNetwork) f = 0 - - isometries = vidal_itn_isometries(ψ, bond_tensors) - + isometries = vidal_gauge_isometries(ψ) for e in keys(isometries) - LHS = isometries[e] / sum(diag(isometries[e])) - id = dense(delta(inds(LHS))) - id /= sum(diag(id)) - f += 0.5 * norm(id - LHS) + lhs = isometries[e] + f += message_diff(ITensor[lhs], ITensor[denseblocks(delta(inds(lhs)))]) end return f / (length(keys(isometries))) end - -"""Function to measure the 'canonicalness' of a state in the Symmetric Gauge""" -function symmetric_itn_canonicalness(ψ::ITensorNetwork, pψψ::PartitionedGraph, mts) - ψ_vidal, bond_tensors = symmetric_to_vidal_gauge(ψ, pψψ, mts) - - return vidal_itn_canonicalness(ψ_vidal, bond_tensors) -end diff --git a/src/imports.jl b/src/imports.jl index e7071450..72a07910 100644 --- a/src/imports.jl +++ b/src/imports.jl @@ -22,7 +22,8 @@ import NamedGraphs: rename_vertices, disjoint_union, mincut_partitions, - incident_edges + incident_edges, + boundary_partitionedges import .DataGraphs: underlying_graph, diff --git a/src/partitioneditensornetwork.jl b/src/partitioneditensornetwork.jl new file mode 100644 index 00000000..8c00d3db --- /dev/null +++ b/src/partitioneditensornetwork.jl @@ -0,0 +1,5 @@ +function linkinds(pitn::PartitionedGraph, edge::PartitionEdge) + src_e_itn = subgraph(pitn, src(edge)) + dst_e_itn = subgraph(pitn, dst(edge)) + return commoninds(src_e_itn, dst_e_itn) +end diff --git a/test/test_apply.jl b/test/test_apply.jl index 2df2f491..0a815b9a 100644 --- a/test/test_apply.jl +++ b/test/test_apply.jl @@ -1,12 +1,11 @@ using ITensorNetworks using ITensorNetworks: - belief_propagation, - environment_tensors, + incoming_messages, + update, contract_inner, - vidal_gauge, - vidal_apply, - vidal_to_symmetric_gauge, - norm_network + norm_network, + BeliefPropagationCache, + VidalITensorNetwork using Test using Compat using ITensors @@ -25,20 +24,19 @@ using SplitApplyCombine χ = 2 ψ = randomITensorNetwork(s; link_space=χ) v1, v2 = (2, 2), (1, 2) - ψψ = norm_network(ψ) #Simple Belief Propagation Grouping - pψψ_SBP = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) - mtsSBP = belief_propagation(pψψ_SBP; niters=20) - envsSBP = environment_tensors(pψψ_SBP, mtsSBP, PartitionVertex.([v1, v2])) + bp_cache = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bp_cache = update(bp_cache; maxiter=20) + envsSBP = incoming_messages(bp_cache, PartitionVertex.([v1, v2])) - ψ_vidal, bond_tensors = vidal_gauge(ψ, pψψ_SBP, mtsSBP) + ψv = VidalITensorNetwork(ψ) #This grouping will correspond to calculating the environments exactly (each column of the grid is a partition) - pψψ_GBP = PartitionedGraph(ψψ, group(v -> v[1][1], vertices(ψψ))) - mtsGBP = belief_propagation(pψψ_GBP; niters=20) - envsGBP = environment_tensors(pψψ_GBP, mtsGBP, [(v1, 1), (v1, 2), (v2, 1), (v2, 2)]) + bp_cache = BeliefPropagationCache(ψψ, group(v -> v[1][1], vertices(ψψ))) + bp_cache = update(bp_cache; maxiter=20) + envsGBP = incoming_messages(bp_cache, [(v1, 1), (v1, 2), (v2, 1), (v2, 2)]) ngates = 5 @@ -55,10 +53,8 @@ using SplitApplyCombine print_fidelity_loss=true, envisposdef=true, ) - ψOVidal, bond_tensors_t = vidal_apply( - o, ψ_vidal, bond_tensors; maxdim=χ, normalize=true - ) - ψOVidal_symm, _ = vidal_to_symmetric_gauge(ψOVidal, bond_tensors_t) + ψOv = apply(o, ψv; maxdim=χ, normalize=true) + ψOVidal_symm = ITensorNetwork(ψOv) ψOGBP = apply( o, ψ; diff --git a/test/test_belief_propagation.jl b/test/test_belief_propagation.jl index cab105b9..a388e758 100644 --- a/test/test_belief_propagation.jl +++ b/test/test_belief_propagation.jl @@ -1,11 +1,14 @@ using ITensorNetworks using ITensorNetworks: ising_network, - belief_propagation, split_index, contract_inner, contract_boundary_mps, - environment_tensors + BeliefPropagationCache, + tensornetwork, + update, + update_factor, + incoming_messages using Test using Compat using ITensors @@ -35,14 +38,21 @@ ITensors.disable_warn_order() Oψ[v] = apply(op("Sz", s[v]), ψ[v]) exact_sz = contract_inner(Oψ, ψ) / contract_inner(ψ, ψ) - pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) - mts = belief_propagation(pψψ) - env_tensors = environment_tensors(pψψ, mts, [PartitionVertex(v)]) + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bpc = update(bpc) + env_tensors = incoming_messages(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 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) @@ -58,15 +68,15 @@ ITensors.disable_warn_order() Oψ[v] = apply(op("Sz", s[v]), ψ[v]) exact_sz = contract_inner(Oψ, ψ) / contract_inner(ψ, ψ) - pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) - mts = belief_propagation(pψψ) - env_tensors = environment_tensors(pψψ, mts, [PartitionVertex(v)]) + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bpc = update(bpc) + env_tensors = incoming_messages(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 - # #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) @@ -80,16 +90,16 @@ ITensors.disable_warn_order() ITensors.contract(ψOψ; sequence=contract_seq)[] / ITensors.contract(ψψ; sequence=contract_seq)[] - pψψ = PartitionedGraph(ψψ; nvertices_per_partition=2, backend="Metis") - mts = belief_propagation(pψψ; niters=20) + bpc = BeliefPropagationCache(ψψ, group(v -> v[1], vertices(ψψ))) + bpc = update(bpc; maxiter=20) - env_tensors = environment_tensors(pψψ, mts, vs) + env_tensors = incoming_messages(bpc, vs) numerator = contract(vcat(env_tensors, ITensor[ψOψ[v] for v in vs]))[] denominator = contract(vcat(env_tensors, ITensor[ψψ[v] for v in vs]))[] @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 + # # #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) @@ -98,11 +108,11 @@ ITensors.disable_warn_order() ψ = randomITensorNetwork(s; link_space=χ) ψψ = ψ ⊗ prime(dag(ψ); sites=[]) - pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) - mts = belief_propagation(pψψ; niters=20) + 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_tensors(pψψ, mts, [(v, 2) for v in vs]) + env_tensors = incoming_messages(bpc, [(v, 2) for v in vs]) rdm = ITensors.contract( vcat(env_tensors, ITensor[ψψsplit[vp] for vp in [(v, 2) for v in vs]]) ) @@ -114,7 +124,7 @@ ITensors.disable_warn_order() @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 + # # #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) @@ -130,14 +140,15 @@ ITensors.disable_warn_order() combiners = linkinds_combiners(ψψ) ψψ = combine_linkinds(ψψ, combiners) ψOψ = combine_linkinds(ψOψ, combiners) - pψψ = PartitionedGraph(ψψ, group(v -> v[1], vertices(ψψ))) - mts = belief_propagation( - pψψ; - contractor=ITensorNetworks.contract_density_matrix, - contractor_kwargs=(; cutoff=1e-6, maxdim=4), + + 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_tensors(pψψ, mts, [v]) + env_tensors = incoming_messages(bpc, [v]) numerator = contract(vcat(env_tensors, ITensor[ψOψ[v]]))[] denominator = contract(vcat(env_tensors, ITensor[ψψ[v]]))[] diff --git a/test/test_gauging.jl b/test/test_gauging.jl index 6441b766..f0a7d10b 100644 --- a/test/test_gauging.jl +++ b/test/test_gauging.jl @@ -1,13 +1,7 @@ using ITensors using ITensorNetworks using ITensorNetworks: - belief_propagation, - contract_inner, - symmetric_gauge, - symmetric_to_vidal_gauge, - vidal_itn_canonicalness, - vidal_gauge, - symmetric_itn_canonicalness + contract_inner, gauge_error, update, messages, BeliefPropagationCache, VidalITensorNetwork using NamedGraphs using Test using Compat @@ -23,26 +17,23 @@ using SplitApplyCombine Random.seed!(5467) ψ = randomITensorNetwork(s; link_space=χ) - ψ_symm, pψψ_symm, ψ_symm_mts = symmetric_gauge(ψ; niters=50) - @test symmetric_itn_canonicalness(ψ_symm, pψψ_symm, ψ_symm_mts) < 1e-5 + # Move directly to vidal gauge + ψ_vidal = VidalITensorNetwork(ψ) + @test gauge_error(ψ_vidal) < 1e-5 - #Test we just did a gauge transform and didn't change the overall network + # Move to symmetric gauge + cache_ref = Ref{BeliefPropagationCache}() + ψ_symm = ITensorNetwork(ψ_vidal; (cache!)=cache_ref) + bp_cache = cache_ref[] + + # Test we just did a gauge transform and didn't change the overall network @test contract_inner(ψ_symm, ψ) / sqrt(contract_inner(ψ_symm, ψ_symm) * contract_inner(ψ, ψ)) ≈ 1.0 - ψψ_symm_V2 = ψ_symm ⊗ prime(dag(ψ_symm); sites=[]) - pψψ_symm_V2 = PartitionedGraph(ψψ_symm_V2, group(v -> v[1], vertices(ψψ_symm_V2))) - ψ_symm_mts_V2 = belief_propagation(pψψ_symm_V2; niters=50) - - for m_e in values(ψ_symm_mts_V2) - #Test all message tensors are approximately diagonal + #Test all message tensors are approximately diagonal even when we keep running BP + bp_cache = update(bp_cache; maxiter=20) + for m_e in values(messages(bp_cache)) @test diagITensor(vector(diag(only(m_e))), inds(only(m_e))) ≈ only(m_e) atol = 1e-8 end - - ψ_vidal, bond_tensors = vidal_gauge(ψ; target_canonicalness=1e-6) - @test vidal_itn_canonicalness(ψ_vidal, bond_tensors) < 1e-5 - - ψ_vidal, bond_tensors = symmetric_to_vidal_gauge(ψ_symm, pψψ_symm, ψ_symm_mts) - @test vidal_itn_canonicalness(ψ_vidal, bond_tensors) < 1e-5 end