From 456c9f9e33e661f86c73cba02466bdbf9fd01442 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 22 Jun 2026 11:19:10 -0400 Subject: [PATCH 1/4] Add two-site DMRG for tree tensor networks Adds `dmrg`, a two-site DMRG ground-state solver for trees: the operator is a `TensorNetworkOperator`, the exact projected-Hamiltonian environments live in a `MessageCache` of `NamedDimsArrays` operators, and the driver follows the layered `AlgorithmsInterface` structure used by `beliefpropagation` and `apply_operators`. Co-Authored-By: Claude Opus 4.8 --- Project.toml | 4 +- src/ITensorNetworksNext.jl | 4 +- src/dmrg.jl | 678 +++++++++++++++++++++++++++++++++++++ test/test_aqua.jl | 6 +- test/test_dmrg.jl | 216 ++++++++++++ 5 files changed, 905 insertions(+), 3 deletions(-) create mode 100644 src/dmrg.jl create mode 100644 test/test_dmrg.jl diff --git a/Project.toml b/Project.toml index b5c5900a..cb7aaf65 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" -version = "0.7.1" +version = "0.7.2" authors = ["ITensor developers and contributors"] [workspace] @@ -14,6 +14,7 @@ DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a" Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" ITensorBase = "4795dd04-0d67-49bb-8f44-b89c448a1dc7" +KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4" @@ -31,6 +32,7 @@ DataGraphs = "0.4" Dictionaries = "0.4.5" Graphs = "1.13.1" ITensorBase = "0.8" +KrylovKit = "0.10" LinearAlgebra = "1.10" MacroTools = "0.5.16" MatrixAlgebraKit = "0.6" diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index 92a541f1..869f9142 100644 --- a/src/ITensorNetworksNext.jl +++ b/src/ITensorNetworksNext.jl @@ -3,7 +3,7 @@ module ITensorNetworksNext if VERSION >= v"1.11.0-DEV.469" eval( Meta.parse( - "public apply_operator, apply_operators, beliefpropagation_normnetwork, identity_norm_message_env, normnetwork, norm_message_env, ones_norm_message_env, rand_norm_message_env, randn_norm_message_env, similar_norm_message_env" + "public apply_operator, apply_operators, beliefpropagation_normnetwork, dmrg, identity_norm_message_env, normnetwork, norm_message_env, ones_norm_message_env, rand_norm_message_env, randn_norm_message_env, similar_norm_message_env, StopWhenEnergyConverged, TensorNetworkOperator" ) ) end @@ -21,4 +21,6 @@ include("beliefpropagation/normnetwork.jl") include("apply/apply_operators.jl") +include("dmrg.jl") + end diff --git a/src/dmrg.jl b/src/dmrg.jl new file mode 100644 index 00000000..b774222b --- /dev/null +++ b/src/dmrg.jl @@ -0,0 +1,678 @@ +import .AlgorithmsInterfaceExtensions as AIE +import AlgorithmsInterface as AI +using .LazyNamedDimsArrays: lazy +using Base: @kwdef +using DataGraphs: DataGraphs, underlying_graph +using Graphs: Graphs, dst, edges, edgetype, neighbors, src, vertices +using KrylovKit: eigsolve +using NamedDimsArrays: + NamedDimsArrays as NDA, AbstractNamedDimsArray, dimnames, replacedimnames +using NamedGraphs.GraphsExtensions: edge_path, post_order_dfs_edges, vertextype +using NamedGraphs: NamedGraphs +using TensorAlgebra: TensorAlgebra as TA +using VectorInterface: VectorInterface as VI + +# ============================ TensorNetworkOperator ===================================== + +struct TensorNetworkOperator{V, VD, State, SiteMap} <: AbstractTensorNetwork{V, VD} + state::State + site_index_map::SiteMap +end + +# `state` is the underlying operator `TensorNetwork`, mirroring `NamedDimsArrays.state` on a +# `NamedDimsOperator`. The rest forwards the tensor-network interface to it. +NDA.state(o::TensorNetworkOperator) = o.state +site_index_map(o::TensorNetworkOperator) = o.site_index_map + +DataGraphs.underlying_graph(o::TensorNetworkOperator) = underlying_graph(NDA.state(o)) +function DataGraphs.is_vertex_assigned(o::TensorNetworkOperator, v) + return DataGraphs.is_vertex_assigned(NDA.state(o), v) +end +function DataGraphs.is_edge_assigned(o::TensorNetworkOperator, e) + return DataGraphs.is_edge_assigned(NDA.state(o), e) +end +function DataGraphs.get_vertex_data(o::TensorNetworkOperator, v) + return DataGraphs.get_vertex_data(NDA.state(o), v) +end + +function Base.copy(o::TensorNetworkOperator) + return TensorNetworkOperator(copy(NDA.state(o)), site_index_map(o)) +end + +function TensorNetworkOperator(state, site_index_map) + return TensorNetworkOperator{ + vertextype(state), eltype(state), typeof(state), typeof(site_index_map), + }( + state, site_index_map + ) +end + +# ============================ QuadraticFormNetwork ====================================== + +struct QuadraticFormNetwork{V, VD, Ket, Operator, LinkMap} <: + AbstractTensorNetwork{V, VD} + ket::Ket + operator::Operator + link_index_map::LinkMap +end + +function bra_name_map(qf::QuadraticFormNetwork) + return merge(site_index_map(qf.operator), qf.link_index_map) +end + +ket_tensor(qf::QuadraticFormNetwork, v) = qf.ket[v] +operator_tensor(qf::QuadraticFormNetwork, v) = qf.operator[v] +function bra_tensor(qf::QuadraticFormNetwork, v) + m = bra_name_map(qf) + return replacedimnames(n -> get(m, n, n), conj(qf.ket[v])) +end + +# === AbstractTensorNetwork / DataGraphs interface === + +DataGraphs.underlying_graph(qf::QuadraticFormNetwork) = underlying_graph(qf.ket) + +function DataGraphs.is_vertex_assigned(qf::QuadraticFormNetwork, v) + return DataGraphs.is_vertex_assigned(qf.ket, v) +end +DataGraphs.is_edge_assigned(::QuadraticFormNetwork, _e) = false + +function DataGraphs.get_vertex_data(qf::QuadraticFormNetwork, v) + return lazy(ket_tensor(qf, v)) * lazy(operator_tensor(qf, v)) * lazy(bra_tensor(qf, v)) +end + +function Base.copy(qf::QuadraticFormNetwork) + return QuadraticFormNetwork(copy(qf.ket), copy(qf.operator), qf.link_index_map) +end + +# === constructor === + +function QuadraticFormNetwork(ket, operator::TensorNetworkOperator, link_index_map) + V = vertextype(ket) + tmp = QuadraticFormNetwork{ + V, Any, typeof(ket), typeof(operator), typeof(link_index_map), + }( + ket, + operator, + link_index_map + ) + VD = typeof(DataGraphs.get_vertex_data(tmp, first(vertices(ket)))) + return QuadraticFormNetwork{ + V, VD, typeof(ket), typeof(operator), typeof(link_index_map), + }( + ket, + operator, + link_index_map + ) +end + +# ============================ Environments ============================================== +# +# On a tree, the projected-Hamiltonian environment of `⟨ψ|H|ψ⟩` is exact: the message on +# directed edge `v → w` is the contraction of the entire `⟨ψ|H|ψ⟩` subtree on `v`'s side +# of the cut `(v, w)`. It carries the three bonds crossing the cut (ket, operator, bra). +# +# The messages are computed by a single dependency-ordered pass over every directed edge +# (`forest_cover_edge_sequence` visits each tree edge in both directions, sources before +# targets), with each message an exact contraction of `qf[v]` against the already-computed +# incoming messages from `v`'s other neighbors. No fixed-point iteration is needed. + +function incoming_subtree_messages(messages, graph, v, w) + return [ + NDA.state(messages[edgetype(graph)(u, v)]) for + u in neighbors(graph, v) if u != w + ] +end + +function environment_operator(message, link_index_map) + ketnames = [n for n in dimnames(message) if haskey(link_index_map, n)] + branames = [link_index_map[n] for n in ketnames] + return NDA.operator(message, branames, ketnames) +end + +""" + quadratic_form_environments(qf::QuadraticFormNetwork; root) -> MessageCache + +Exact projected-Hamiltonian environments of `⟨ψ|H|ψ⟩` on a tree, as a `MessageCache` of +`NamedDimsArrays` operators keyed by directed edges. The message on `v → w` is the +contraction of the `⟨ψ|H|ψ⟩` subtree on `v`'s side of `(v, w)`, wrapped as an operator +recording the bra ↔ ket link correspondence (see [`environment_operator`](@ref)). +""" +function quadratic_form_environments(qf::QuadraticFormNetwork; kwargs...) + sequence = forest_cover_edge_sequence(qf; kwargs...) + messages = Dict{edgetype(qf), Any}() + for e in sequence + v, w = src(e), dst(e) + incoming = incoming_subtree_messages(messages, qf, v, w) + message = contract_network([[qf[v]]; incoming]) + messages[e] = environment_operator(message, qf.link_index_map) + end + return messagecache(messages) +end + +# ============================ Effective Hamiltonian ===================================== +# +# The effective Hamiltonian for a `region` (a vector of vertices) is the projected operator +# obtained by contracting the operator tensors on the region with the incoming environment +# messages on the region's boundary. Contracting those gives a single tensor whose codomain +# (output) names are the region's bra names and whose domain (input) names are its ket names +# — i.e. exactly a `NamedDimsArrays` operator. So `effective_hamiltonian` returns that +# operator, and its action on a region ket tensor `T` is `NamedDimsArrays.apply(H, T)` (which +# contracts the ket names and renames the resulting bra names back to ket names). Returning +# an operator object (rather than a bare closure) follows the effective-Hamiltonian designs +# in ITensorMPS (`ProjMPO`), MPSKit (`AC_hamiltonian`), and TeNPy (`EffectiveH`). + +""" + effective_hamiltonian(qf::QuadraticFormNetwork, env, region) -> NamedDimsOperator + +Effective (projected) Hamiltonian for `region` (a vector of vertices) as a `NamedDimsArrays` +operator: its domain (input) names are the region's ket names and its codomain (output) +names are the matching bra names. Apply it to a region ket tensor `T` with +`NamedDimsArrays.apply`. The environment `env` is a `MessageCache` of +[`quadratic_form_environments`](@ref). +""" +function effective_hamiltonian(qf::QuadraticFormNetwork, env, region) + operators = [operator_tensor(qf, v) for v in region] + boundary = [NDA.state(m) for m in incoming_edge_data(env, region)] + h = contract_network([operators; boundary]) + sitemap = site_index_map(qf.operator) + ketnames = [ + n for n in dimnames(h) if haskey(sitemap, n) || haskey(qf.link_index_map, n) + ] + branames = [bra_name_map(qf)[n] for n in ketnames] + return NDA.operator(h, branames, ketnames) +end + +# ============================ Orthogonalization ========================================= +# +# Isometric (QR) gauge on a tree. The orthogonality center is *not* stored on the network +# — it is tracked by the caller. `orthogonalize(state, center)` canonicalizes the whole +# tree so every tensor but `center` is an isometry pointing toward `center`; +# `orthogonalize(state, source, dest)` moves the center along the tree path, returning the +# walked edges so environments can be refreshed incrementally. +# +# Each gauge step keeps the link name on its edge stable (the fresh QR bond is renamed back +# to the original link name), so name maps built against the ket — e.g. a +# `QuadraticFormNetwork`'s `link_index_map` — stay valid across gauging. + +# Make `state[v]` an isometry whose only non-isometric leg points toward `w`, pushing the +# `R` factor into `state[w]`. Mutates `state`. +function gauge_move!(state, v, w) + ln = only(linknames(state, v => w)) + tv = state[v] + rows = collect(setdiff(dimnames(tv), [ln])) + Q, R = TA.qr(tv, rows) + r = only(setdiff(dimnames(Q), rows)) + new_w = R * state[w] + setindex_preserve_graph!(state, replacedimnames(Q, r => ln), v) + setindex_preserve_graph!(state, replacedimnames(new_w, r => ln), w) + return state +end + +""" + orthogonalize(state, center) -> state + +Canonicalize the tree tensor network `state` so that every tensor except `center` (a +vertex) is an isometry pointing toward `center`; the orthogonality center is then `center`. +""" +function orthogonalize(state, center) + state = copy(state) + for e in post_order_dfs_edges(state, center) + gauge_move!(state, src(e), dst(e)) + end + return state +end + +""" + orthogonalize(state, source, dest) -> (state, walked_edges) + +Move the orthogonality center of the tree tensor network `state` from vertex `source` to +vertex `dest` by QR gauge steps along the tree path, assuming `state` is already canonical +with center `source`. Returns the re-gauged `state` and the directed edges walked (for +incremental environment refresh). +""" +function orthogonalize(state, source, dest) + state = copy(state) + walked = collect(edge_path(state, source, dest)) + for e in walked + gauge_move!(state, src(e), dst(e)) + end + return state, walked +end + +# ============================ Local update (TwoSiteEigsolve) ============================ +# +# Per-region update strategy, resolved through `select_algorithm`/`default_algorithm` like +# the belief-propagation and apply-operator strategies. The default `TwoSiteEigsolve` +# extracts the region's two-site ket tensor, finds the lowest eigenpair of the effective +# Hamiltonian with `KrylovKit.eigsolve`, and SVD-splits the result back onto the two +# vertices (truncating with `trunc`). Assumes the orthogonality center sits on `region`, so +# the effective Hamiltonian is the exact projected Hamiltonian and the eigenvalue is the +# energy. Mutates the ket of `qf`. + +abstract type DMRGUpdateAlgorithm <: AbstractAlgorithm end + +function region_update! end + +@kwdef struct TwoSiteEigsolve{Trunc} <: DMRGUpdateAlgorithm + trunc::Trunc = nothing + which::Symbol = :SR +end + +function default_algorithm(::typeof(region_update!), ::Type{<:Tuple}; kwargs...) + return TwoSiteEigsolve(; kwargs...) +end + +""" + region_update!(qf::QuadraticFormNetwork, env, region; alg=nothing, kwargs...) -> (qf, energy) + +Optimize the ket tensors on `region` (a vector of vertices) against the effective +Hamiltonian of `qf` with environments `env`, returning the updated `qf` and the energy. +The orthogonality center must sit on `region`. +""" +function region_update!(qf, env, region; alg = nothing, kwargs...) + algorithm = select_algorithm(region_update!, alg, (qf, env, region); kwargs...) + return region_update!(algorithm, qf, env, region) +end + +function region_update!(alg::TwoSiteEigsolve, qf, env, region) + return region_update_nsite!(Val(length(region)), alg, qf, env, region) +end + +function region_update_nsite!(::Val{N}, alg, qf, env, region) where {N} + return throw(ArgumentError("$N-site region update not implemented")) +end + +function svd_split(T, rows; trunc) + return isnothing(trunc) ? TA.svd(T, rows) : TA.svd(T, rows; trunc) +end + +# --- Temporary `VectorInterface` overloads for named arrays --------------------------- +# +# `KrylovKit.eigsolve` drives its Krylov vectors through `VectorInterface`. The generic +# `AbstractArray` fallbacks broadcast in a way that fails on a named array (e.g. +# `zerovector(x, S)`), so we provide name-aware methods here. This is type piracy on +# `AbstractNamedDimsArray` and is intended to move into `NamedDimsArrays` before this work +# merges. +VI.scalartype(::Type{<:AbstractNamedDimsArray{T}}) where {T} = T +function VI.zerovector(x::AbstractNamedDimsArray, ::Type{S}) where {S <: Number} + return fill!(similar(x, S), zero(S)) +end +VI.scale(x::AbstractNamedDimsArray, α::Number) = x * α +VI.scale!!(x::AbstractNamedDimsArray, α::Number) = x * α +VI.scale!!(::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number) = x * α +function VI.add!!( + y::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number, β::Number + ) + return x * α + y * β +end +VI.inner(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) = (conj(x) * y)[] +# -------------------------------------------------------------------------------------- + +# Lowest eigenpair of the effective Hamiltonian operator `H_eff` acting on the named tensor +# `T` (via `NamedDimsArrays.apply`), using `KrylovKit.eigsolve`. +function eigsolve_named(H_eff, T, which) + vals, vecs = eigsolve(x -> NDA.apply(H_eff, x), T, 1, which; ishermitian = true) + return real(vals[1]), vecs[1] +end + +function region_update_nsite!(::Val{2}, alg::TwoSiteEigsolve, qf, env, region) + v1, v2 = region + ln = only(linknames(qf.ket, v1 => v2)) + rows = collect(setdiff(dimnames(qf.ket[v1]), [ln])) + + T = qf.ket[v1] * qf.ket[v2] + H_eff = effective_hamiltonian(qf, env, region) + energy, T_opt = eigsolve_named(H_eff, T, alg.which) + + U, S, Vt = svd_split(T_opt, rows; alg.trunc) + bond = only(setdiff(dimnames(U), rows)) + # Center on `v2` (the direction the sweep moves): `v1` is left isometric. + setindex_preserve_graph!(qf.ket, replacedimnames(U, bond => ln), v1) + setindex_preserve_graph!(qf.ket, replacedimnames(S * Vt, bond => ln), v2) + return qf, energy +end + +# ============================ Layered DMRG algorithm ==================================== +# +# Three layers mirroring `beliefpropagation` / `apply_operators`: the outer sweep loop +# (`DMRGProblem` / `DMRGAlgorithm`), one sweep over regions (`DMRGSweepProblem` / +# `DMRGSweepAlgorithm`), and the per-region update strategy (`DMRGUpdateAlgorithm`, default +# `TwoSiteEigsolve`). The iterate is the ket tensor network; the operator and ket→bra name +# maps are fixed problem data. +# +# Environments are maintained incrementally across a sweep. Each sweep starts by gauging the +# ket onto the first region and building the exact tree environment once (O(N)); each region +# step then moves the orthogonality center along the tree path to the region (refreshing the +# forward message on each walked edge) and refreshes the message leaving the region after the +# local solve. So a sweep costs O(N), not O(N^2). + +# Default sweep plan: every tree edge as a 2-site region, both directions (a back-and-forth +# sweep), via `forest_cover_edge_sequence`. +function default_dmrg_regions(ket; kwargs...) + return [[src(e), dst(e)] for e in forest_cover_edge_sequence(ket; kwargs...)] +end + +# Recompute the environment message on `v → w` from the (current) subtree on `v`'s side: +# `qf[v]` contracted with the incoming messages from `v`'s other neighbors. Used to refresh +# a message after the tensor at `v` has been re-gauged or updated. +function set_environment_message!(env, qf, v, w) + incoming = incoming_subtree_messages(env, qf, v, w) + message = contract_network([[qf[v]]; incoming]) + env[edgetype(qf)(v, w)] = environment_operator(message, qf.link_index_map) + return env +end + +# Walk the orthogonality center from its current vertex to `target` along the tree path, +# QR-gauging each edge and refreshing the forward environment message so the env stays +# exact. No-op when the center is already at `target`. +function move_center!(state, qf, target) + for e in edge_path(state.iterate, state.center, target) + gauge_move!(state.iterate, src(e), dst(e)) + set_environment_message!(state.env, qf, src(e), dst(e)) + end + state.center = target + return state +end + +# === Layer 2: one sweep over regions === + +struct DMRGSweepProblem{Operator, Regions, LinkMap} <: AI.Problem + operator::Operator + regions::Regions + link_index_map::LinkMap +end + +@kwdef struct DMRGSweepAlgorithm{ + UpdateAlgorithm, StoppingCriterion <: AI.StoppingCriterion, + } <: AI.Algorithm + update_algorithm::UpdateAlgorithm = TwoSiteEigsolve() + stopping_criterion::StoppingCriterion +end + +@kwdef mutable struct DMRGSweepState{ + Iterate, Env, V, StoppingCriterionState <: AI.StoppingCriterionState, + } <: AI.State + iterate::Iterate + env::Env + center::V + energy::Float64 = 0.0 + iteration::Int = 0 + stopping_criterion_state::StoppingCriterionState +end + +function quadraticformnetwork(problem::DMRGSweepProblem, ket) + return QuadraticFormNetwork(ket, problem.operator, problem.link_index_map) +end + +# Gauge the ket onto the first region's center and build the exact tree environment from +# scratch. Run once at the start of each sweep; per-step updates keep it current. +function prepare_sweep(problem::DMRGSweepProblem, iterate) + center = first(first(problem.regions)) + ket = orthogonalize(iterate, center) + qf = quadraticformnetwork(problem, ket) + env = quadratic_form_environments(qf; root_vertex = _ -> center) + return ket, env, center +end + +function AI.initialize_state( + problem::DMRGSweepProblem, algorithm::DMRGSweepAlgorithm; + iterate, iteration::Int = 0 + ) + stopping_criterion_state = AI.initialize_state( + problem, algorithm, algorithm.stopping_criterion; iterate + ) + ket, env, center = prepare_sweep(problem, iterate) + return DMRGSweepState(; iterate = ket, env, center, iteration, stopping_criterion_state) +end + +# The iterate, environment, and center are carried from the previous sweep (built once in +# `initialize_state`). Reset only the per-sweep counters and move the center back to the +# first region's start — a no-op for the default round-trip plan, an incremental walk for a +# custom plan that ends elsewhere. The environment is never rebuilt from scratch. +function AI.initialize_state!( + problem::DMRGSweepProblem, algorithm::DMRGSweepAlgorithm, + state::DMRGSweepState; iteration::Int = 0 + ) + state.iteration = iteration + AI.initialize_state!( + problem, algorithm, algorithm.stopping_criterion, state.stopping_criterion_state + ) + target = first(first(problem.regions)) + if state.center != target + move_center!(state, quadraticformnetwork(problem, state.iterate), target) + end + return state +end + +function AI.step!( + problem::DMRGSweepProblem, algorithm::DMRGSweepAlgorithm, state::DMRGSweepState + ) + region = problem.regions[state.iteration] + v1, v2 = region[1], region[2] + qf = quadraticformnetwork(problem, state.iterate) + move_center!(state, qf, v1) + qf, energy = region_update!(algorithm.update_algorithm, qf, state.env, region) + # The split left the center on `v2` and finalized `v1`; refresh the message it emits. + set_environment_message!(state.env, qf, v1, v2) + state.center = v2 + state.energy = energy + return state +end + +# === Layer 1: outer sweep loop === + +struct DMRGProblem{Operator, LinkMap} <: AI.Problem + operator::Operator + link_index_map::LinkMap +end + +# `sweep_update_algorithm(n)` returns the per-region update algorithm for sweep `n`, so the +# truncation (or any update-algorithm parameter) can vary by sweep. +@kwdef struct DMRGAlgorithm{ + Regions, SweepUpdateAlgorithm, StoppingCriterion <: AI.StoppingCriterion, + } <: AIE.NestedAlgorithm + regions::Regions + sweep_update_algorithm::SweepUpdateAlgorithm + stopping_criterion::StoppingCriterion +end + +function sweep_algorithm(algorithm::DMRGAlgorithm, n::Int) + return DMRGSweepAlgorithm(; + update_algorithm = algorithm.sweep_update_algorithm(n), + stopping_criterion = AI.StopAfterIteration(length(algorithm.regions)) + ) +end + +@kwdef mutable struct DMRGState{ + Substate <: AI.State, StoppingCriterionState <: AI.StoppingCriterionState, + } <: AIE.NestedState + substate::Substate + iteration::Int = 0 + stopping_criterion_state::StoppingCriterionState +end + +# The current energy lives on the substate; expose it on the outer state too. +energy(state::DMRGState) = state.substate.energy + +function AI.initialize_state( + problem::DMRGProblem, algorithm::DMRGAlgorithm; iterate, iteration::Int = 0 + ) + subproblem = DMRGSweepProblem( + problem.operator, algorithm.regions, problem.link_index_map + ) + substate = AI.initialize_state(subproblem, sweep_algorithm(algorithm, 1); iterate) + stopping_criterion_state = AI.initialize_state( + problem, algorithm, algorithm.stopping_criterion; iterate + ) + return DMRGState(; iteration, stopping_criterion_state, substate) +end + +function AI.initialize_state!( + problem::DMRGProblem, algorithm::DMRGAlgorithm, state::DMRGState; + iteration::Int = 0 + ) + state.iteration = iteration + AI.initialize_state!( + problem, algorithm, algorithm.stopping_criterion, state.stopping_criterion_state + ) + return state +end + +function AIE.initialize_subsolve( + problem::DMRGProblem, algorithm::DMRGAlgorithm, state::DMRGState + ) + subproblem = DMRGSweepProblem( + problem.operator, algorithm.regions, problem.link_index_map + ) + return subproblem, sweep_algorithm(algorithm, state.iteration), state.substate +end + +function AI.finalize_state!(::DMRGProblem, ::DMRGAlgorithm, state::DMRGState) + return state.iterate, energy(state) +end + +# === Energy-based convergence === + +@kwdef struct StopWhenEnergyConverged <: AI.StoppingCriterion + tol::Float64 +end + +@kwdef mutable struct StopWhenEnergyConvergedState <: AI.StoppingCriterionState + delta::Float64 = Inf + at_iteration::Int = -1 + previous_energy::Float64 = NaN +end + +function AI.initialize_state( + ::AI.Problem, ::AI.Algorithm, ::StopWhenEnergyConverged; iterate + ) + return StopWhenEnergyConvergedState() +end + +function AI.initialize_state!( + ::AI.Problem, ::AI.Algorithm, ::StopWhenEnergyConverged, + st::StopWhenEnergyConvergedState + ) + st.delta = Inf + st.previous_energy = NaN + return st +end + +function AI.is_finished!( + problem::AI.Problem, algorithm::AI.Algorithm, state::AI.State, + c::StopWhenEnergyConverged, st::StopWhenEnergyConvergedState + ) + current_energy = energy(state) + previous_energy = st.previous_energy + st.previous_energy = current_energy + # No previous-sweep energy to compare against before the first sweep completes. + state.iteration == 0 && return false + st.delta = abs(current_energy - previous_energy) + if AI.is_finished(problem, algorithm, state, c, st) + st.at_iteration = state.iteration + return true + end + return false +end + +function AI.is_finished( + ::AI.Problem, ::AI.Algorithm, ::AI.State, + c::StopWhenEnergyConverged, st::StopWhenEnergyConvergedState + ) + return st.delta < c.tol +end + +# === Top-level entry point === + +select_dmrg_stopping_criterion(c::AI.StoppingCriterion) = c +function select_dmrg_stopping_criterion(::Nothing) + return throw( + ArgumentError( + "`stopping_criterion` must be specified, e.g.\n" * + " `stopping_criterion = (; maxsweeps = 10)`,\n" * + " `stopping_criterion = (; maxsweeps = 10, tol = 1.0e-10)`, or\n" * + " `stopping_criterion = AI.StopAfterIteration(10) | StopWhenEnergyConverged(1.0e-10)`." + ) + ) +end +function select_dmrg_stopping_criterion(kwargs::NamedTuple) + return select_dmrg_stopping_criterion(; kwargs...) +end +function select_dmrg_stopping_criterion(; maxsweeps = nothing, tol = nothing, kwargs...) + if !isempty(kwargs) + throw( + ArgumentError( + "Unrecognized `stopping_criterion` kwargs: $(keys(kwargs)). " * + "Supported: `maxsweeps`, `tol`." + ) + ) + end + if isnothing(maxsweeps) && isnothing(tol) + throw(ArgumentError("At least one of `maxsweeps` or `tol` must be specified.")) + end + criterion = nothing + if !isnothing(maxsweeps) + criterion = AI.StopAfterIteration(maxsweeps) + end + if !isnothing(tol) + converged = StopWhenEnergyConverged(; tol) + criterion = isnothing(criterion) ? converged : criterion | converged + end + return criterion +end + +""" + dmrg(operator, ket; stopping_criterion, kwargs...) -> (ket, energy) + +Two-site DMRG for the ground state of `operator` (a [`TensorNetworkOperator`](@ref) on the +same tree as `ket`) starting from the ket tensor network `ket`. Returns the optimized ket +and its energy. + +Keyword arguments: + + - `stopping_criterion` (required): an `AI.StoppingCriterion`, or a `NamedTuple` accepting + `maxsweeps` (cap on back-and-forth passes) and/or `tol` (stop once the energy change + between consecutive sweeps drops below `tol`), e.g. `(; maxsweeps = 10, tol = 1.0e-10)`. + - `trunc`: a function `n -> trunc` giving the SVD-split truncation for sweep `n`. Cannot + be combined with an explicit `alg`. + - `regions`: sweep plan, a vector of vertex-vector regions (default: every tree edge in + both directions). + - `link_index_map`: ket→bra link name map (default: a fresh name per ket link). + - `alg`: per-region update algorithm (fixed across sweeps). Other keyword arguments + (e.g. `which`) are forwarded to the default `TwoSiteEigsolve`. +""" +function dmrg( + operator::TensorNetworkOperator, ket; + stopping_criterion = nothing, + trunc = nothing, + regions = default_dmrg_regions(ket), + link_index_map = Dict( + ln => randname(ln) for e in edges(ket) for ln in linknames(ket, e) + ), + alg = nothing, + kwargs... + ) + sweep_update_algorithm = if isnothing(alg) + sweep_trunc = isnothing(trunc) ? Returns(nothing) : trunc + n -> select_algorithm( + region_update!, + nothing, + (ket,); + trunc = sweep_trunc(n), + kwargs... + ) + else + isnothing(trunc) || throw( + ArgumentError( + "`trunc` cannot be combined with an explicit `alg`; set it on `alg`." + ) + ) + Returns(select_algorithm(region_update!, alg, (ket,))) + end + algorithm = DMRGAlgorithm(; + regions, sweep_update_algorithm, + stopping_criterion = select_dmrg_stopping_criterion(stopping_criterion) + ) + problem = DMRGProblem(operator, link_index_map) + return AI.solve(problem, algorithm; iterate = copy(ket)) +end diff --git a/test/test_aqua.jl b/test/test_aqua.jl index 8eb46122..6d3276da 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -3,5 +3,9 @@ using ITensorNetworksNext: ITensorNetworksNext using Test: @testset @testset "Code quality (Aqua.jl)" begin - Aqua.test_all(ITensorNetworksNext; persistent_tasks = false) + # Piracy is checked separately as `broken`: `dmrg.jl` temporarily pirates a few + # `VectorInterface` methods on `AbstractNamedDimsArray` (needed by `KrylovKit.eigsolve`). + # These are slated to move into `NamedDimsArrays`; drop the `broken` marker once they do. + Aqua.test_all(ITensorNetworksNext; persistent_tasks = false, piracies = false) + Aqua.test_piracies(ITensorNetworksNext; broken = true) end diff --git a/test/test_dmrg.jl b/test/test_dmrg.jl new file mode 100644 index 00000000..24e99a17 --- /dev/null +++ b/test/test_dmrg.jl @@ -0,0 +1,216 @@ +using Graphs: add_edge!, edges, vertices +using ITensorBase: Index +using ITensorNetworksNext: TensorNetwork, TensorNetworkOperator, dmrg, insertlink! +using LinearAlgebra: eigen +using MatrixAlgebraKit: truncrank +using NamedDimsArrays: name, nameddims, randname, setname +using NamedGraphs.NamedGraphGenerators: named_path_graph +using NamedGraphs: NamedGraph +using StableRNGs: StableRNG +using Test: @test, @test_throws, @testset + +# Transverse-field Ising model `H = -J Σ Z Z - h Σ X`, built by hand as operator networks +# and as dense matrices for exact-diagonalization references. + +const X = [0.0 1.0; 1.0 0.0] +const Z = [1.0 0.0; 0.0 -1.0] +const Id = [1.0 0.0; 0.0 1.0] + +kron_ops(ops) = foldl(kron, ops) + +# Bulk Ising MPO tensor `W[a, b, ket, bra]` (bond dim 3 finite-state machine). +function ising_mpo_tensor(; J, h) + W = zeros(3, 3, 2, 2) + W[1, 1, :, :] = Id + W[2, 1, :, :] = Z + W[3, 1, :, :] = -h * X + W[3, 2, :, :] = -J * Z + W[3, 3, :, :] = Id + return W +end + +# Path-graph TFIM as a `TensorNetworkOperator` on vertices `1:N`. +function tfim_path_operator(N, sites, sitemap; J, h) + verts = collect(1:N) + bond_edges = [(verts[i], verts[i + 1]) for i in 1:(N - 1)] + bonds = Dict( + e => setname(Index(Base.OneTo(3)), randname(name(sites[verts[1]]))) + for e in bond_edges + ) + W = ising_mpo_tensor(; J, h) + g = NamedGraph(verts) + for (v1, v2) in bond_edges + add_edge!(g, v1, v2) + end + function tensor(v) + k = name(sites[v]) + b = sitemap[k] + left = findfirst(e -> e[2] == v, bond_edges) + right = findfirst(e -> e[1] == v, bond_edges) + if isnothing(left) # left boundary: row 3 of W + data = W[3, :, :, :] + return nameddims(data, (name(bonds[bond_edges[right]]), k, b)) + elseif isnothing(right) # right boundary: column 1 of W + data = W[:, 1, :, :] + return nameddims(data, (name(bonds[bond_edges[left]]), k, b)) + end + return nameddims( + W, (name(bonds[bond_edges[left]]), name(bonds[bond_edges[right]]), k, b) + ) + end + tn = TensorNetwork(g) do v + return tensor(v) + end + return TensorNetworkOperator(tn, sitemap) +end + +function tfim_path_dense(N; J, h) + H = zeros(2^N, 2^N) + for i in 1:(N - 1) + H += -J * kron_ops([(j == i || j == i + 1) ? Z : Id for j in 1:N]) + end + for i in 1:N + H += -h * kron_ops([j == i ? X : Id for j in 1:N]) + end + return H +end + +# Star tree: center vertex 1 bonded to leaves 2, 3, 4. +function tfim_star_operator(sites, sitemap; J, h) + leaves = [2, 3, 4] + g = NamedGraph(collect(1:4)) + for v in leaves + add_edge!(g, 1, v) + end + bonds = Dict( + v => setname(Index(Base.OneTo(3)), randname(name(sites[1]))) for v in leaves + ) + # Center tensor: one bond per leaf, with `-hX` once and `-JZ`/`I` on each leaf channel. + center = zeros(3, 3, 3, 2, 2) + center[1, 1, 1, :, :] = -h * X + for (i, _) in enumerate(leaves) + center[ntuple(j -> j == i ? 2 : 1, 3)..., :, :] = -J * Z + center[ntuple(j -> j == i ? 3 : 1, 3)..., :, :] = Id + end + center_names = ( + name(bonds[2]), name(bonds[3]), name(bonds[4]), name(sites[1]), + sitemap[name(sites[1])], + ) + function leaf_tensor(v) + data = zeros(3, 2, 2) + data[1, :, :] = Id + data[2, :, :] = Z + data[3, :, :] = -h * X + return nameddims(data, (name(bonds[v]), name(sites[v]), sitemap[name(sites[v])])) + end + tensors = + Dict(1 => nameddims(center, center_names), (v => leaf_tensor(v) for v in leaves)...) + tn = TensorNetwork(g) do v + return tensors[v] + end + return TensorNetworkOperator(tn, sitemap) +end + +function tfim_star_dense(; J, h) + H = zeros(16, 16) + for v in [2, 3, 4] + H += -J * kron_ops([(j == 1 || j == v) ? Z : Id for j in 1:4]) + end + for i in 1:4 + H += -h * kron_ops([j == i ? X : Id for j in 1:4]) + end + return H +end + +function random_ket(rng, g) + sites = Dict(v => Index(Base.OneTo(2)) for v in vertices(g)) + ket = TensorNetwork(NamedGraph(collect(vertices(g)))) do v + return randn(rng, Float64, (sites[v],)) + end + for e in edges(g) + insertlink!(ket, e) + end + return ket, sites +end + +@testset "dmrg" begin + @testset "path TFIM matches ED (N=$N)" for N in (4, 6) + rng = StableRNG(8) + g = named_path_graph(N) + ket0, sites = random_ket(rng, g) + sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) + exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) + + _, energy = dmrg(operator, ket0; stopping_criterion = (; maxsweeps = 8)) + @test energy ≈ exact rtol = 1.0e-8 + end + + @testset "energy-tolerance early stop" begin + N = 6 + rng = StableRNG(8) + g = named_path_graph(N) + ket0, sites = random_ket(rng, g) + sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) + exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) + + _, energy = + dmrg(operator, ket0; stopping_criterion = (; maxsweeps = 50, tol = 1.0e-10)) + @test energy ≈ exact rtol = 1.0e-8 + end + + @testset "per-sweep truncation schedule" begin + N = 6 + rng = StableRNG(8) + g = named_path_graph(N) + ket0, sites = random_ket(rng, g) + sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) + exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) + + # Ramp the bond dimension up over sweeps; the final rank is exact for N=6. + trunc = n -> truncrank(min(2^n, 8)) + _, energy = dmrg(operator, ket0; stopping_criterion = (; maxsweeps = 8), trunc) + @test energy ≈ exact rtol = 1.0e-8 + end + + @testset "energy decreases with more sweeps" begin + N = 6 + rng = StableRNG(8) + g = named_path_graph(N) + ket0, sites = random_ket(rng, g) + sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) + exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) + + _, energy1 = dmrg(operator, ket0; stopping_criterion = (; maxsweeps = 1)) + _, energy4 = dmrg(operator, ket0; stopping_criterion = (; maxsweeps = 4)) + @test energy4 ≤ energy1 + 1.0e-10 + @test energy4 ≥ exact - 1.0e-10 + end + + @testset "star-tree TFIM matches ED" begin + rng = StableRNG(11) + g = NamedGraph(collect(1:4)) + for v in 2:4 + add_edge!(g, 1, v) + end + ket0, sites = random_ket(rng, g) + sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + operator = tfim_star_operator(sites, sitemap; J = 1.0, h = 0.6) + exact = minimum(eigen(tfim_star_dense(; J = 1.0, h = 0.6)).values) + + _, energy = dmrg(operator, ket0; stopping_criterion = (; maxsweeps = 8)) + @test energy ≈ exact rtol = 1.0e-8 + end + + @testset "stopping_criterion is required" begin + rng = StableRNG(8) + g = named_path_graph(4) + ket0, sites = random_ket(rng, g) + sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + operator = tfim_path_operator(4, sites, sitemap; J = 1.0, h = 0.7) + @test_throws ArgumentError dmrg(operator, ket0) + end +end From 7809b5d0da089a5a593fbd909b8792674a9efaae Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 22 Jun 2026 11:40:10 -0400 Subject: [PATCH 2/4] Reorganize tree-DMRG sources top-down Move `dmrg` to the top of `dmrg.jl` so the file reads entry-point first, matching `apply_operators.jl`. Split the generic pieces into their own files: `tensornetworkoperator.jl`, `quadraticformnetwork.jl`, `orthogonalize.jl`, and `nameddimsarrays_extensions.jl` (the temporary `VectorInterface` methods on `AbstractNamedDimsArray`). Co-Authored-By: Claude Opus 4.8 --- src/ITensorNetworksNext.jl | 4 + src/dmrg.jl | 739 +++++++++++------------------- src/nameddimsarrays_extensions.jl | 22 + src/orthogonalize.jl | 59 +++ src/quadraticformnetwork.jl | 111 +++++ src/tensornetworkoperator.jl | 38 ++ 6 files changed, 491 insertions(+), 482 deletions(-) create mode 100644 src/nameddimsarrays_extensions.jl create mode 100644 src/orthogonalize.jl create mode 100644 src/quadraticformnetwork.jl create mode 100644 src/tensornetworkoperator.jl diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index 869f9142..48192b44 100644 --- a/src/ITensorNetworksNext.jl +++ b/src/ITensorNetworksNext.jl @@ -21,6 +21,10 @@ include("beliefpropagation/normnetwork.jl") include("apply/apply_operators.jl") +include("nameddimsarrays_extensions.jl") +include("tensornetworkoperator.jl") +include("quadraticformnetwork.jl") +include("orthogonalize.jl") include("dmrg.jl") end diff --git a/src/dmrg.jl b/src/dmrg.jl index b774222b..b779baf1 100644 --- a/src/dmrg.jl +++ b/src/dmrg.jl @@ -1,335 +1,111 @@ import .AlgorithmsInterfaceExtensions as AIE import AlgorithmsInterface as AI -using .LazyNamedDimsArrays: lazy using Base: @kwdef -using DataGraphs: DataGraphs, underlying_graph -using Graphs: Graphs, dst, edges, edgetype, neighbors, src, vertices +using Graphs: dst, edges, edgetype, src using KrylovKit: eigsolve -using NamedDimsArrays: - NamedDimsArrays as NDA, AbstractNamedDimsArray, dimnames, replacedimnames -using NamedGraphs.GraphsExtensions: edge_path, post_order_dfs_edges, vertextype -using NamedGraphs: NamedGraphs +using NamedDimsArrays: NamedDimsArrays as NDA, dimnames, replacedimnames +using NamedGraphs.GraphsExtensions: edge_path using TensorAlgebra: TensorAlgebra as TA -using VectorInterface: VectorInterface as VI -# ============================ TensorNetworkOperator ===================================== +# ============================ Top-level entry point ===================================== -struct TensorNetworkOperator{V, VD, State, SiteMap} <: AbstractTensorNetwork{V, VD} - state::State - site_index_map::SiteMap -end - -# `state` is the underlying operator `TensorNetwork`, mirroring `NamedDimsArrays.state` on a -# `NamedDimsOperator`. The rest forwards the tensor-network interface to it. -NDA.state(o::TensorNetworkOperator) = o.state -site_index_map(o::TensorNetworkOperator) = o.site_index_map - -DataGraphs.underlying_graph(o::TensorNetworkOperator) = underlying_graph(NDA.state(o)) -function DataGraphs.is_vertex_assigned(o::TensorNetworkOperator, v) - return DataGraphs.is_vertex_assigned(NDA.state(o), v) -end -function DataGraphs.is_edge_assigned(o::TensorNetworkOperator, e) - return DataGraphs.is_edge_assigned(NDA.state(o), e) -end -function DataGraphs.get_vertex_data(o::TensorNetworkOperator, v) - return DataGraphs.get_vertex_data(NDA.state(o), v) -end - -function Base.copy(o::TensorNetworkOperator) - return TensorNetworkOperator(copy(NDA.state(o)), site_index_map(o)) -end - -function TensorNetworkOperator(state, site_index_map) - return TensorNetworkOperator{ - vertextype(state), eltype(state), typeof(state), typeof(site_index_map), - }( - state, site_index_map - ) -end - -# ============================ QuadraticFormNetwork ====================================== - -struct QuadraticFormNetwork{V, VD, Ket, Operator, LinkMap} <: - AbstractTensorNetwork{V, VD} - ket::Ket - operator::Operator - link_index_map::LinkMap -end - -function bra_name_map(qf::QuadraticFormNetwork) - return merge(site_index_map(qf.operator), qf.link_index_map) -end - -ket_tensor(qf::QuadraticFormNetwork, v) = qf.ket[v] -operator_tensor(qf::QuadraticFormNetwork, v) = qf.operator[v] -function bra_tensor(qf::QuadraticFormNetwork, v) - m = bra_name_map(qf) - return replacedimnames(n -> get(m, n, n), conj(qf.ket[v])) -end - -# === AbstractTensorNetwork / DataGraphs interface === - -DataGraphs.underlying_graph(qf::QuadraticFormNetwork) = underlying_graph(qf.ket) - -function DataGraphs.is_vertex_assigned(qf::QuadraticFormNetwork, v) - return DataGraphs.is_vertex_assigned(qf.ket, v) -end -DataGraphs.is_edge_assigned(::QuadraticFormNetwork, _e) = false - -function DataGraphs.get_vertex_data(qf::QuadraticFormNetwork, v) - return lazy(ket_tensor(qf, v)) * lazy(operator_tensor(qf, v)) * lazy(bra_tensor(qf, v)) -end +""" + dmrg(operator, ket; stopping_criterion, kwargs...) -> (ket, energy) -function Base.copy(qf::QuadraticFormNetwork) - return QuadraticFormNetwork(copy(qf.ket), copy(qf.operator), qf.link_index_map) -end +Two-site DMRG for the ground state of `operator` (a [`TensorNetworkOperator`](@ref) on the +same tree as `ket`) starting from the ket tensor network `ket`. Returns the optimized ket +and its energy. -# === constructor === +Keyword arguments: -function QuadraticFormNetwork(ket, operator::TensorNetworkOperator, link_index_map) - V = vertextype(ket) - tmp = QuadraticFormNetwork{ - V, Any, typeof(ket), typeof(operator), typeof(link_index_map), - }( - ket, - operator, - link_index_map + - `stopping_criterion` (required): an `AI.StoppingCriterion`, or a `NamedTuple` accepting + `maxsweeps` (cap on back-and-forth passes) and/or `tol` (stop once the energy change + between consecutive sweeps drops below `tol`), e.g. `(; maxsweeps = 10, tol = 1.0e-10)`. + - `trunc`: a function `n -> trunc` giving the SVD-split truncation for sweep `n`. Cannot + be combined with an explicit `alg`. + - `regions`: sweep plan, a vector of vertex-vector regions (default: every tree edge in + both directions). + - `link_index_map`: ket→bra link name map (default: a fresh name per ket link). + - `alg`: per-region update algorithm (fixed across sweeps). Other keyword arguments + (e.g. `which`) are forwarded to the default `TwoSiteEigsolve`. +""" +function dmrg( + operator::TensorNetworkOperator, ket; + stopping_criterion = nothing, + trunc = nothing, + regions = default_dmrg_regions(ket), + link_index_map = Dict( + ln => randname(ln) for e in edges(ket) for ln in linknames(ket, e) + ), + alg = nothing, + kwargs... ) - VD = typeof(DataGraphs.get_vertex_data(tmp, first(vertices(ket)))) - return QuadraticFormNetwork{ - V, VD, typeof(ket), typeof(operator), typeof(link_index_map), - }( - ket, - operator, - link_index_map + sweep_update_algorithm = if isnothing(alg) + sweep_trunc = isnothing(trunc) ? Returns(nothing) : trunc + n -> select_algorithm( + region_update!, + nothing, + (ket,); + trunc = sweep_trunc(n), + kwargs... + ) + else + isnothing(trunc) || throw( + ArgumentError( + "`trunc` cannot be combined with an explicit `alg`; set it on `alg`." + ) + ) + Returns(select_algorithm(region_update!, alg, (ket,))) + end + algorithm = DMRGAlgorithm(; + regions, sweep_update_algorithm, + stopping_criterion = select_dmrg_stopping_criterion(stopping_criterion) ) + problem = DMRGProblem(operator, link_index_map) + return AI.solve(problem, algorithm; iterate = copy(ket)) end -# ============================ Environments ============================================== -# -# On a tree, the projected-Hamiltonian environment of `⟨ψ|H|ψ⟩` is exact: the message on -# directed edge `v → w` is the contraction of the entire `⟨ψ|H|ψ⟩` subtree on `v`'s side -# of the cut `(v, w)`. It carries the three bonds crossing the cut (ket, operator, bra). -# -# The messages are computed by a single dependency-ordered pass over every directed edge -# (`forest_cover_edge_sequence` visits each tree edge in both directions, sources before -# targets), with each message an exact contraction of `qf[v]` against the already-computed -# incoming messages from `v`'s other neighbors. No fixed-point iteration is needed. - -function incoming_subtree_messages(messages, graph, v, w) - return [ - NDA.state(messages[edgetype(graph)(u, v)]) for - u in neighbors(graph, v) if u != w - ] +select_dmrg_stopping_criterion(c::AI.StoppingCriterion) = c +function select_dmrg_stopping_criterion(::Nothing) + return throw( + ArgumentError( + "`stopping_criterion` must be specified, e.g.\n" * + " `stopping_criterion = (; maxsweeps = 10)`,\n" * + " `stopping_criterion = (; maxsweeps = 10, tol = 1.0e-10)`, or\n" * + " `stopping_criterion = AI.StopAfterIteration(10) | StopWhenEnergyConverged(1.0e-10)`." + ) + ) end - -function environment_operator(message, link_index_map) - ketnames = [n for n in dimnames(message) if haskey(link_index_map, n)] - branames = [link_index_map[n] for n in ketnames] - return NDA.operator(message, branames, ketnames) +function select_dmrg_stopping_criterion(kwargs::NamedTuple) + return select_dmrg_stopping_criterion(; kwargs...) end - -""" - quadratic_form_environments(qf::QuadraticFormNetwork; root) -> MessageCache - -Exact projected-Hamiltonian environments of `⟨ψ|H|ψ⟩` on a tree, as a `MessageCache` of -`NamedDimsArrays` operators keyed by directed edges. The message on `v → w` is the -contraction of the `⟨ψ|H|ψ⟩` subtree on `v`'s side of `(v, w)`, wrapped as an operator -recording the bra ↔ ket link correspondence (see [`environment_operator`](@ref)). -""" -function quadratic_form_environments(qf::QuadraticFormNetwork; kwargs...) - sequence = forest_cover_edge_sequence(qf; kwargs...) - messages = Dict{edgetype(qf), Any}() - for e in sequence - v, w = src(e), dst(e) - incoming = incoming_subtree_messages(messages, qf, v, w) - message = contract_network([[qf[v]]; incoming]) - messages[e] = environment_operator(message, qf.link_index_map) +function select_dmrg_stopping_criterion(; maxsweeps = nothing, tol = nothing, kwargs...) + if !isempty(kwargs) + throw( + ArgumentError( + "Unrecognized `stopping_criterion` kwargs: $(keys(kwargs)). " * + "Supported: `maxsweeps`, `tol`." + ) + ) end - return messagecache(messages) -end - -# ============================ Effective Hamiltonian ===================================== -# -# The effective Hamiltonian for a `region` (a vector of vertices) is the projected operator -# obtained by contracting the operator tensors on the region with the incoming environment -# messages on the region's boundary. Contracting those gives a single tensor whose codomain -# (output) names are the region's bra names and whose domain (input) names are its ket names -# — i.e. exactly a `NamedDimsArrays` operator. So `effective_hamiltonian` returns that -# operator, and its action on a region ket tensor `T` is `NamedDimsArrays.apply(H, T)` (which -# contracts the ket names and renames the resulting bra names back to ket names). Returning -# an operator object (rather than a bare closure) follows the effective-Hamiltonian designs -# in ITensorMPS (`ProjMPO`), MPSKit (`AC_hamiltonian`), and TeNPy (`EffectiveH`). - -""" - effective_hamiltonian(qf::QuadraticFormNetwork, env, region) -> NamedDimsOperator - -Effective (projected) Hamiltonian for `region` (a vector of vertices) as a `NamedDimsArrays` -operator: its domain (input) names are the region's ket names and its codomain (output) -names are the matching bra names. Apply it to a region ket tensor `T` with -`NamedDimsArrays.apply`. The environment `env` is a `MessageCache` of -[`quadratic_form_environments`](@ref). -""" -function effective_hamiltonian(qf::QuadraticFormNetwork, env, region) - operators = [operator_tensor(qf, v) for v in region] - boundary = [NDA.state(m) for m in incoming_edge_data(env, region)] - h = contract_network([operators; boundary]) - sitemap = site_index_map(qf.operator) - ketnames = [ - n for n in dimnames(h) if haskey(sitemap, n) || haskey(qf.link_index_map, n) - ] - branames = [bra_name_map(qf)[n] for n in ketnames] - return NDA.operator(h, branames, ketnames) -end - -# ============================ Orthogonalization ========================================= -# -# Isometric (QR) gauge on a tree. The orthogonality center is *not* stored on the network -# — it is tracked by the caller. `orthogonalize(state, center)` canonicalizes the whole -# tree so every tensor but `center` is an isometry pointing toward `center`; -# `orthogonalize(state, source, dest)` moves the center along the tree path, returning the -# walked edges so environments can be refreshed incrementally. -# -# Each gauge step keeps the link name on its edge stable (the fresh QR bond is renamed back -# to the original link name), so name maps built against the ket — e.g. a -# `QuadraticFormNetwork`'s `link_index_map` — stay valid across gauging. - -# Make `state[v]` an isometry whose only non-isometric leg points toward `w`, pushing the -# `R` factor into `state[w]`. Mutates `state`. -function gauge_move!(state, v, w) - ln = only(linknames(state, v => w)) - tv = state[v] - rows = collect(setdiff(dimnames(tv), [ln])) - Q, R = TA.qr(tv, rows) - r = only(setdiff(dimnames(Q), rows)) - new_w = R * state[w] - setindex_preserve_graph!(state, replacedimnames(Q, r => ln), v) - setindex_preserve_graph!(state, replacedimnames(new_w, r => ln), w) - return state -end - -""" - orthogonalize(state, center) -> state - -Canonicalize the tree tensor network `state` so that every tensor except `center` (a -vertex) is an isometry pointing toward `center`; the orthogonality center is then `center`. -""" -function orthogonalize(state, center) - state = copy(state) - for e in post_order_dfs_edges(state, center) - gauge_move!(state, src(e), dst(e)) + if isnothing(maxsweeps) && isnothing(tol) + throw(ArgumentError("At least one of `maxsweeps` or `tol` must be specified.")) end - return state -end - -""" - orthogonalize(state, source, dest) -> (state, walked_edges) - -Move the orthogonality center of the tree tensor network `state` from vertex `source` to -vertex `dest` by QR gauge steps along the tree path, assuming `state` is already canonical -with center `source`. Returns the re-gauged `state` and the directed edges walked (for -incremental environment refresh). -""" -function orthogonalize(state, source, dest) - state = copy(state) - walked = collect(edge_path(state, source, dest)) - for e in walked - gauge_move!(state, src(e), dst(e)) + criterion = nothing + if !isnothing(maxsweeps) + criterion = AI.StopAfterIteration(maxsweeps) end - return state, walked -end - -# ============================ Local update (TwoSiteEigsolve) ============================ -# -# Per-region update strategy, resolved through `select_algorithm`/`default_algorithm` like -# the belief-propagation and apply-operator strategies. The default `TwoSiteEigsolve` -# extracts the region's two-site ket tensor, finds the lowest eigenpair of the effective -# Hamiltonian with `KrylovKit.eigsolve`, and SVD-splits the result back onto the two -# vertices (truncating with `trunc`). Assumes the orthogonality center sits on `region`, so -# the effective Hamiltonian is the exact projected Hamiltonian and the eigenvalue is the -# energy. Mutates the ket of `qf`. - -abstract type DMRGUpdateAlgorithm <: AbstractAlgorithm end - -function region_update! end - -@kwdef struct TwoSiteEigsolve{Trunc} <: DMRGUpdateAlgorithm - trunc::Trunc = nothing - which::Symbol = :SR -end - -function default_algorithm(::typeof(region_update!), ::Type{<:Tuple}; kwargs...) - return TwoSiteEigsolve(; kwargs...) -end - -""" - region_update!(qf::QuadraticFormNetwork, env, region; alg=nothing, kwargs...) -> (qf, energy) - -Optimize the ket tensors on `region` (a vector of vertices) against the effective -Hamiltonian of `qf` with environments `env`, returning the updated `qf` and the energy. -The orthogonality center must sit on `region`. -""" -function region_update!(qf, env, region; alg = nothing, kwargs...) - algorithm = select_algorithm(region_update!, alg, (qf, env, region); kwargs...) - return region_update!(algorithm, qf, env, region) -end - -function region_update!(alg::TwoSiteEigsolve, qf, env, region) - return region_update_nsite!(Val(length(region)), alg, qf, env, region) -end - -function region_update_nsite!(::Val{N}, alg, qf, env, region) where {N} - return throw(ArgumentError("$N-site region update not implemented")) -end - -function svd_split(T, rows; trunc) - return isnothing(trunc) ? TA.svd(T, rows) : TA.svd(T, rows; trunc) -end - -# --- Temporary `VectorInterface` overloads for named arrays --------------------------- -# -# `KrylovKit.eigsolve` drives its Krylov vectors through `VectorInterface`. The generic -# `AbstractArray` fallbacks broadcast in a way that fails on a named array (e.g. -# `zerovector(x, S)`), so we provide name-aware methods here. This is type piracy on -# `AbstractNamedDimsArray` and is intended to move into `NamedDimsArrays` before this work -# merges. -VI.scalartype(::Type{<:AbstractNamedDimsArray{T}}) where {T} = T -function VI.zerovector(x::AbstractNamedDimsArray, ::Type{S}) where {S <: Number} - return fill!(similar(x, S), zero(S)) -end -VI.scale(x::AbstractNamedDimsArray, α::Number) = x * α -VI.scale!!(x::AbstractNamedDimsArray, α::Number) = x * α -VI.scale!!(::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number) = x * α -function VI.add!!( - y::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number, β::Number - ) - return x * α + y * β -end -VI.inner(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) = (conj(x) * y)[] -# -------------------------------------------------------------------------------------- - -# Lowest eigenpair of the effective Hamiltonian operator `H_eff` acting on the named tensor -# `T` (via `NamedDimsArrays.apply`), using `KrylovKit.eigsolve`. -function eigsolve_named(H_eff, T, which) - vals, vecs = eigsolve(x -> NDA.apply(H_eff, x), T, 1, which; ishermitian = true) - return real(vals[1]), vecs[1] + if !isnothing(tol) + converged = StopWhenEnergyConverged(; tol) + criterion = isnothing(criterion) ? converged : criterion | converged + end + return criterion end -function region_update_nsite!(::Val{2}, alg::TwoSiteEigsolve, qf, env, region) - v1, v2 = region - ln = only(linknames(qf.ket, v1 => v2)) - rows = collect(setdiff(dimnames(qf.ket[v1]), [ln])) - - T = qf.ket[v1] * qf.ket[v2] - H_eff = effective_hamiltonian(qf, env, region) - energy, T_opt = eigsolve_named(H_eff, T, alg.which) - - U, S, Vt = svd_split(T_opt, rows; alg.trunc) - bond = only(setdiff(dimnames(U), rows)) - # Center on `v2` (the direction the sweep moves): `v1` is left isometric. - setindex_preserve_graph!(qf.ket, replacedimnames(U, bond => ln), v1) - setindex_preserve_graph!(qf.ket, replacedimnames(S * Vt, bond => ln), v2) - return qf, energy +# Default sweep plan: every tree edge as a 2-site region, both directions (a back-and-forth +# sweep), via `forest_cover_edge_sequence`. +function default_dmrg_regions(ket; kwargs...) + return [[src(e), dst(e)] for e in forest_cover_edge_sequence(ket; kwargs...)] end # ============================ Layered DMRG algorithm ==================================== @@ -346,34 +122,78 @@ end # forward message on each walked edge) and refreshes the message leaving the region after the # local solve. So a sweep costs O(N), not O(N^2). -# Default sweep plan: every tree edge as a 2-site region, both directions (a back-and-forth -# sweep), via `forest_cover_edge_sequence`. -function default_dmrg_regions(ket; kwargs...) - return [[src(e), dst(e)] for e in forest_cover_edge_sequence(ket; kwargs...)] +# === Layer 1: outer sweep loop === + +struct DMRGProblem{Operator, LinkMap} <: AI.Problem + operator::Operator + link_index_map::LinkMap end -# Recompute the environment message on `v → w` from the (current) subtree on `v`'s side: -# `qf[v]` contracted with the incoming messages from `v`'s other neighbors. Used to refresh -# a message after the tensor at `v` has been re-gauged or updated. -function set_environment_message!(env, qf, v, w) - incoming = incoming_subtree_messages(env, qf, v, w) - message = contract_network([[qf[v]]; incoming]) - env[edgetype(qf)(v, w)] = environment_operator(message, qf.link_index_map) - return env +# `sweep_update_algorithm(n)` returns the per-region update algorithm for sweep `n`, so the +# truncation (or any update-algorithm parameter) can vary by sweep. +@kwdef struct DMRGAlgorithm{ + Regions, SweepUpdateAlgorithm, StoppingCriterion <: AI.StoppingCriterion, + } <: AIE.NestedAlgorithm + regions::Regions + sweep_update_algorithm::SweepUpdateAlgorithm + stopping_criterion::StoppingCriterion end -# Walk the orthogonality center from its current vertex to `target` along the tree path, -# QR-gauging each edge and refreshing the forward environment message so the env stays -# exact. No-op when the center is already at `target`. -function move_center!(state, qf, target) - for e in edge_path(state.iterate, state.center, target) - gauge_move!(state.iterate, src(e), dst(e)) - set_environment_message!(state.env, qf, src(e), dst(e)) - end - state.center = target +function sweep_algorithm(algorithm::DMRGAlgorithm, n::Int) + return DMRGSweepAlgorithm(; + update_algorithm = algorithm.sweep_update_algorithm(n), + stopping_criterion = AI.StopAfterIteration(length(algorithm.regions)) + ) +end + +@kwdef mutable struct DMRGState{ + Substate <: AI.State, StoppingCriterionState <: AI.StoppingCriterionState, + } <: AIE.NestedState + substate::Substate + iteration::Int = 0 + stopping_criterion_state::StoppingCriterionState +end + +# The current energy lives on the substate; expose it on the outer state too. +energy(state::DMRGState) = state.substate.energy + +function AI.initialize_state( + problem::DMRGProblem, algorithm::DMRGAlgorithm; iterate, iteration::Int = 0 + ) + subproblem = DMRGSweepProblem( + problem.operator, algorithm.regions, problem.link_index_map + ) + substate = AI.initialize_state(subproblem, sweep_algorithm(algorithm, 1); iterate) + stopping_criterion_state = AI.initialize_state( + problem, algorithm, algorithm.stopping_criterion; iterate + ) + return DMRGState(; iteration, stopping_criterion_state, substate) +end + +function AI.initialize_state!( + problem::DMRGProblem, algorithm::DMRGAlgorithm, state::DMRGState; + iteration::Int = 0 + ) + state.iteration = iteration + AI.initialize_state!( + problem, algorithm, algorithm.stopping_criterion, state.stopping_criterion_state + ) return state end +function AIE.initialize_subsolve( + problem::DMRGProblem, algorithm::DMRGAlgorithm, state::DMRGState + ) + subproblem = DMRGSweepProblem( + problem.operator, algorithm.regions, problem.link_index_map + ) + return subproblem, sweep_algorithm(algorithm, state.iteration), state.substate +end + +function AI.finalize_state!(::DMRGProblem, ::DMRGAlgorithm, state::DMRGState) + return state.iterate, energy(state) +end + # === Layer 2: one sweep over regions === struct DMRGSweepProblem{Operator, Regions, LinkMap} <: AI.Problem @@ -459,79 +279,129 @@ function AI.step!( return state end -# === Layer 1: outer sweep loop === +# Walk the orthogonality center from its current vertex to `target` along the tree path, +# QR-gauging each edge and refreshing the forward environment message so the env stays +# exact. No-op when the center is already at `target`. +function move_center!(state, qf, target) + for e in edge_path(state.iterate, state.center, target) + gauge_move!(state.iterate, src(e), dst(e)) + set_environment_message!(state.env, qf, src(e), dst(e)) + end + state.center = target + return state +end -struct DMRGProblem{Operator, LinkMap} <: AI.Problem - operator::Operator - link_index_map::LinkMap +# Recompute the environment message on `v → w` from the (current) subtree on `v`'s side: +# `qf[v]` contracted with the incoming messages from `v`'s other neighbors. Used to refresh +# a message after the tensor at `v` has been re-gauged or updated. +function set_environment_message!(env, qf, v, w) + incoming = incoming_subtree_messages(env, qf, v, w) + message = contract_network([[qf[v]]; incoming]) + env[edgetype(qf)(v, w)] = environment_operator(message, qf.link_index_map) + return env end -# `sweep_update_algorithm(n)` returns the per-region update algorithm for sweep `n`, so the -# truncation (or any update-algorithm parameter) can vary by sweep. -@kwdef struct DMRGAlgorithm{ - Regions, SweepUpdateAlgorithm, StoppingCriterion <: AI.StoppingCriterion, - } <: AIE.NestedAlgorithm - regions::Regions - sweep_update_algorithm::SweepUpdateAlgorithm - stopping_criterion::StoppingCriterion +# ============================ Local update (TwoSiteEigsolve) ============================ +# +# Per-region update strategy, resolved through `select_algorithm`/`default_algorithm` like +# the belief-propagation and apply-operator strategies. The default `TwoSiteEigsolve` +# extracts the region's two-site ket tensor, finds the lowest eigenpair of the effective +# Hamiltonian with `KrylovKit.eigsolve`, and SVD-splits the result back onto the two +# vertices (truncating with `trunc`). Assumes the orthogonality center sits on `region`, so +# the effective Hamiltonian is the exact projected Hamiltonian and the eigenvalue is the +# energy. Mutates the ket of `qf`. + +abstract type DMRGUpdateAlgorithm <: AbstractAlgorithm end + +function region_update! end + +@kwdef struct TwoSiteEigsolve{Trunc} <: DMRGUpdateAlgorithm + trunc::Trunc = nothing + which::Symbol = :SR end -function sweep_algorithm(algorithm::DMRGAlgorithm, n::Int) - return DMRGSweepAlgorithm(; - update_algorithm = algorithm.sweep_update_algorithm(n), - stopping_criterion = AI.StopAfterIteration(length(algorithm.regions)) - ) +function default_algorithm(::typeof(region_update!), ::Type{<:Tuple}; kwargs...) + return TwoSiteEigsolve(; kwargs...) end -@kwdef mutable struct DMRGState{ - Substate <: AI.State, StoppingCriterionState <: AI.StoppingCriterionState, - } <: AIE.NestedState - substate::Substate - iteration::Int = 0 - stopping_criterion_state::StoppingCriterionState +""" + region_update!(qf::QuadraticFormNetwork, env, region; alg=nothing, kwargs...) -> (qf, energy) + +Optimize the ket tensors on `region` (a vector of vertices) against the effective +Hamiltonian of `qf` with environments `env`, returning the updated `qf` and the energy. +The orthogonality center must sit on `region`. +""" +function region_update!(qf, env, region; alg = nothing, kwargs...) + algorithm = select_algorithm(region_update!, alg, (qf, env, region); kwargs...) + return region_update!(algorithm, qf, env, region) end -# The current energy lives on the substate; expose it on the outer state too. -energy(state::DMRGState) = state.substate.energy +function region_update!(alg::TwoSiteEigsolve, qf, env, region) + return region_update_nsite!(Val(length(region)), alg, qf, env, region) +end -function AI.initialize_state( - problem::DMRGProblem, algorithm::DMRGAlgorithm; iterate, iteration::Int = 0 - ) - subproblem = DMRGSweepProblem( - problem.operator, algorithm.regions, problem.link_index_map - ) - substate = AI.initialize_state(subproblem, sweep_algorithm(algorithm, 1); iterate) - stopping_criterion_state = AI.initialize_state( - problem, algorithm, algorithm.stopping_criterion; iterate - ) - return DMRGState(; iteration, stopping_criterion_state, substate) +function region_update_nsite!(::Val{N}, alg, qf, env, region) where {N} + return throw(ArgumentError("$N-site region update not implemented")) end -function AI.initialize_state!( - problem::DMRGProblem, algorithm::DMRGAlgorithm, state::DMRGState; - iteration::Int = 0 - ) - state.iteration = iteration - AI.initialize_state!( - problem, algorithm, algorithm.stopping_criterion, state.stopping_criterion_state - ) - return state +function region_update_nsite!(::Val{2}, alg::TwoSiteEigsolve, qf, env, region) + v1, v2 = region + ln = only(linknames(qf.ket, v1 => v2)) + rows = collect(setdiff(dimnames(qf.ket[v1]), [ln])) + + T = qf.ket[v1] * qf.ket[v2] + H_eff = effective_hamiltonian(qf, env, region) + energy, T_opt = eigsolve_named(H_eff, T, alg.which) + + U, S, Vt = svd_split(T_opt, rows; alg.trunc) + bond = only(setdiff(dimnames(U), rows)) + # Center on `v2` (the direction the sweep moves): `v1` is left isometric. + setindex_preserve_graph!(qf.ket, replacedimnames(U, bond => ln), v1) + setindex_preserve_graph!(qf.ket, replacedimnames(S * Vt, bond => ln), v2) + return qf, energy end -function AIE.initialize_subsolve( - problem::DMRGProblem, algorithm::DMRGAlgorithm, state::DMRGState - ) - subproblem = DMRGSweepProblem( - problem.operator, algorithm.regions, problem.link_index_map - ) - return subproblem, sweep_algorithm(algorithm, state.iteration), state.substate +function svd_split(T, rows; trunc) + return isnothing(trunc) ? TA.svd(T, rows) : TA.svd(T, rows; trunc) end -function AI.finalize_state!(::DMRGProblem, ::DMRGAlgorithm, state::DMRGState) - return state.iterate, energy(state) +# Lowest eigenpair of the effective Hamiltonian operator `H_eff` acting on the named tensor +# `T` (via `NamedDimsArrays.apply`), using `KrylovKit.eigsolve`. +function eigsolve_named(H_eff, T, which) + vals, vecs = eigsolve(x -> NDA.apply(H_eff, x), T, 1, which; ishermitian = true) + return real(vals[1]), vecs[1] +end + +# The effective Hamiltonian for a `region` (a vector of vertices) is the projected operator +# obtained by contracting the operator tensors on the region with the incoming environment +# messages on the region's boundary. Contracting those gives a single tensor whose codomain +# (output) names are the region's bra names and whose domain (input) names are its ket names +# — i.e. exactly a `NamedDimsArrays` operator. So `effective_hamiltonian` returns that +# operator, and its action on a region ket tensor `T` is `NamedDimsArrays.apply(H, T)` (which +# contracts the ket names and renames the resulting bra names back to ket names). + +""" + effective_hamiltonian(qf::QuadraticFormNetwork, env, region) -> NamedDimsOperator + +Effective (projected) Hamiltonian for `region` (a vector of vertices) as a `NamedDimsArrays` +operator: its domain (input) names are the region's ket names and its codomain (output) +names are the matching bra names. Apply it to a region ket tensor `T` with +`NamedDimsArrays.apply`. The environment `env` is a `MessageCache` of +[`quadratic_form_environments`](@ref). +""" +function effective_hamiltonian(qf::QuadraticFormNetwork, env, region) + operators = [operator_tensor(qf, v) for v in region] + boundary = [NDA.state(m) for m in incoming_edge_data(env, region)] + h = contract_network([operators; boundary]) + sitemap = site_index_map(qf.operator) + ketnames = [ + n for n in dimnames(h) if haskey(sitemap, n) || haskey(qf.link_index_map, n) + ] + branames = [bra_name_map(qf)[n] for n in ketnames] + return NDA.operator(h, branames, ketnames) end -# === Energy-based convergence === +# ============================ Energy-based convergence ================================== @kwdef struct StopWhenEnergyConverged <: AI.StoppingCriterion tol::Float64 @@ -581,98 +451,3 @@ function AI.is_finished( ) return st.delta < c.tol end - -# === Top-level entry point === - -select_dmrg_stopping_criterion(c::AI.StoppingCriterion) = c -function select_dmrg_stopping_criterion(::Nothing) - return throw( - ArgumentError( - "`stopping_criterion` must be specified, e.g.\n" * - " `stopping_criterion = (; maxsweeps = 10)`,\n" * - " `stopping_criterion = (; maxsweeps = 10, tol = 1.0e-10)`, or\n" * - " `stopping_criterion = AI.StopAfterIteration(10) | StopWhenEnergyConverged(1.0e-10)`." - ) - ) -end -function select_dmrg_stopping_criterion(kwargs::NamedTuple) - return select_dmrg_stopping_criterion(; kwargs...) -end -function select_dmrg_stopping_criterion(; maxsweeps = nothing, tol = nothing, kwargs...) - if !isempty(kwargs) - throw( - ArgumentError( - "Unrecognized `stopping_criterion` kwargs: $(keys(kwargs)). " * - "Supported: `maxsweeps`, `tol`." - ) - ) - end - if isnothing(maxsweeps) && isnothing(tol) - throw(ArgumentError("At least one of `maxsweeps` or `tol` must be specified.")) - end - criterion = nothing - if !isnothing(maxsweeps) - criterion = AI.StopAfterIteration(maxsweeps) - end - if !isnothing(tol) - converged = StopWhenEnergyConverged(; tol) - criterion = isnothing(criterion) ? converged : criterion | converged - end - return criterion -end - -""" - dmrg(operator, ket; stopping_criterion, kwargs...) -> (ket, energy) - -Two-site DMRG for the ground state of `operator` (a [`TensorNetworkOperator`](@ref) on the -same tree as `ket`) starting from the ket tensor network `ket`. Returns the optimized ket -and its energy. - -Keyword arguments: - - - `stopping_criterion` (required): an `AI.StoppingCriterion`, or a `NamedTuple` accepting - `maxsweeps` (cap on back-and-forth passes) and/or `tol` (stop once the energy change - between consecutive sweeps drops below `tol`), e.g. `(; maxsweeps = 10, tol = 1.0e-10)`. - - `trunc`: a function `n -> trunc` giving the SVD-split truncation for sweep `n`. Cannot - be combined with an explicit `alg`. - - `regions`: sweep plan, a vector of vertex-vector regions (default: every tree edge in - both directions). - - `link_index_map`: ket→bra link name map (default: a fresh name per ket link). - - `alg`: per-region update algorithm (fixed across sweeps). Other keyword arguments - (e.g. `which`) are forwarded to the default `TwoSiteEigsolve`. -""" -function dmrg( - operator::TensorNetworkOperator, ket; - stopping_criterion = nothing, - trunc = nothing, - regions = default_dmrg_regions(ket), - link_index_map = Dict( - ln => randname(ln) for e in edges(ket) for ln in linknames(ket, e) - ), - alg = nothing, - kwargs... - ) - sweep_update_algorithm = if isnothing(alg) - sweep_trunc = isnothing(trunc) ? Returns(nothing) : trunc - n -> select_algorithm( - region_update!, - nothing, - (ket,); - trunc = sweep_trunc(n), - kwargs... - ) - else - isnothing(trunc) || throw( - ArgumentError( - "`trunc` cannot be combined with an explicit `alg`; set it on `alg`." - ) - ) - Returns(select_algorithm(region_update!, alg, (ket,))) - end - algorithm = DMRGAlgorithm(; - regions, sweep_update_algorithm, - stopping_criterion = select_dmrg_stopping_criterion(stopping_criterion) - ) - problem = DMRGProblem(operator, link_index_map) - return AI.solve(problem, algorithm; iterate = copy(ket)) -end diff --git a/src/nameddimsarrays_extensions.jl b/src/nameddimsarrays_extensions.jl new file mode 100644 index 00000000..b53b9e5c --- /dev/null +++ b/src/nameddimsarrays_extensions.jl @@ -0,0 +1,22 @@ +using NamedDimsArrays: AbstractNamedDimsArray +using VectorInterface: VectorInterface as VI + +# Temporary `VectorInterface` methods for named arrays. +# +# `KrylovKit.eigsolve` drives its Krylov vectors through `VectorInterface`. The generic +# `AbstractArray` fallbacks broadcast in a way that fails on a named array (e.g. +# `zerovector(x, S)`), so we provide name-aware methods here. This is type piracy on +# `AbstractNamedDimsArray` and is intended to move into `NamedDimsArrays`. +VI.scalartype(::Type{<:AbstractNamedDimsArray{T}}) where {T} = T +function VI.zerovector(x::AbstractNamedDimsArray, ::Type{S}) where {S <: Number} + return fill!(similar(x, S), zero(S)) +end +VI.scale(x::AbstractNamedDimsArray, α::Number) = x * α +VI.scale!!(x::AbstractNamedDimsArray, α::Number) = x * α +VI.scale!!(::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number) = x * α +function VI.add!!( + y::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number, β::Number + ) + return x * α + y * β +end +VI.inner(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) = (conj(x) * y)[] diff --git a/src/orthogonalize.jl b/src/orthogonalize.jl new file mode 100644 index 00000000..999aa233 --- /dev/null +++ b/src/orthogonalize.jl @@ -0,0 +1,59 @@ +using Graphs: dst, src +using NamedDimsArrays: dimnames, replacedimnames +using NamedGraphs.GraphsExtensions: edge_path, post_order_dfs_edges +using TensorAlgebra: TensorAlgebra as TA + +# Isometric (QR) gauge on a tree. The orthogonality center is *not* stored on the network +# — it is tracked by the caller. `orthogonalize(state, center)` canonicalizes the whole +# tree so every tensor but `center` is an isometry pointing toward `center`; +# `orthogonalize(state, source, dest)` moves the center along the tree path, returning the +# walked edges so environments can be refreshed incrementally. +# +# Each gauge step keeps the link name on its edge stable (the fresh QR bond is renamed back +# to the original link name), so name maps built against the ket — e.g. a +# `QuadraticFormNetwork`'s `link_index_map` — stay valid across gauging. + +# Make `state[v]` an isometry whose only non-isometric leg points toward `w`, pushing the +# `R` factor into `state[w]`. Mutates `state`. +function gauge_move!(state, v, w) + ln = only(linknames(state, v => w)) + tv = state[v] + rows = collect(setdiff(dimnames(tv), [ln])) + Q, R = TA.qr(tv, rows) + r = only(setdiff(dimnames(Q), rows)) + new_w = R * state[w] + setindex_preserve_graph!(state, replacedimnames(Q, r => ln), v) + setindex_preserve_graph!(state, replacedimnames(new_w, r => ln), w) + return state +end + +""" + orthogonalize(state, center) -> state + +Canonicalize the tree tensor network `state` so that every tensor except `center` (a +vertex) is an isometry pointing toward `center`; the orthogonality center is then `center`. +""" +function orthogonalize(state, center) + state = copy(state) + for e in post_order_dfs_edges(state, center) + gauge_move!(state, src(e), dst(e)) + end + return state +end + +""" + orthogonalize(state, source, dest) -> (state, walked_edges) + +Move the orthogonality center of the tree tensor network `state` from vertex `source` to +vertex `dest` by QR gauge steps along the tree path, assuming `state` is already canonical +with center `source`. Returns the re-gauged `state` and the directed edges walked (for +incremental environment refresh). +""" +function orthogonalize(state, source, dest) + state = copy(state) + walked = collect(edge_path(state, source, dest)) + for e in walked + gauge_move!(state, src(e), dst(e)) + end + return state, walked +end diff --git a/src/quadraticformnetwork.jl b/src/quadraticformnetwork.jl new file mode 100644 index 00000000..84eddb77 --- /dev/null +++ b/src/quadraticformnetwork.jl @@ -0,0 +1,111 @@ +using .LazyNamedDimsArrays: lazy +using DataGraphs: DataGraphs, underlying_graph +using Graphs: dst, edgetype, neighbors, src, vertices +using NamedDimsArrays: NamedDimsArrays as NDA, dimnames, replacedimnames +using NamedGraphs.GraphsExtensions: vertextype + +# A lazy `⟨ψ|H|ψ⟩` network: a ket `TensorNetwork`, a `TensorNetworkOperator` (which carries +# the ket → bra *site* name map), and a forward ket → bra map for the *link* names. The bra +# layer is derived from the ket (`conj` + index renaming), never stored, so updating a ket +# tensor is reflected in the bra. As an `AbstractTensorNetwork`, the data on vertex `v` is +# the lazy product `lazy(ket) * lazy(operator) * lazy(bra)`, so the existing +# `contract_network` / `MessageCache` machinery treats it like any other tensor network. +struct QuadraticFormNetwork{V, VD, Ket, Operator, LinkMap} <: + AbstractTensorNetwork{V, VD} + ket::Ket + operator::Operator + link_index_map::LinkMap +end + +function bra_name_map(qf::QuadraticFormNetwork) + return merge(site_index_map(qf.operator), qf.link_index_map) +end + +ket_tensor(qf::QuadraticFormNetwork, v) = qf.ket[v] +operator_tensor(qf::QuadraticFormNetwork, v) = qf.operator[v] +function bra_tensor(qf::QuadraticFormNetwork, v) + m = bra_name_map(qf) + return replacedimnames(n -> get(m, n, n), conj(qf.ket[v])) +end + +# === AbstractTensorNetwork / DataGraphs interface === + +DataGraphs.underlying_graph(qf::QuadraticFormNetwork) = underlying_graph(qf.ket) + +function DataGraphs.is_vertex_assigned(qf::QuadraticFormNetwork, v) + return DataGraphs.is_vertex_assigned(qf.ket, v) +end +DataGraphs.is_edge_assigned(::QuadraticFormNetwork, _e) = false + +function DataGraphs.get_vertex_data(qf::QuadraticFormNetwork, v) + return lazy(ket_tensor(qf, v)) * lazy(operator_tensor(qf, v)) * lazy(bra_tensor(qf, v)) +end + +function Base.copy(qf::QuadraticFormNetwork) + return QuadraticFormNetwork(copy(qf.ket), copy(qf.operator), qf.link_index_map) +end + +# === constructor === + +function QuadraticFormNetwork(ket, operator::TensorNetworkOperator, link_index_map) + V = vertextype(ket) + tmp = QuadraticFormNetwork{ + V, Any, typeof(ket), typeof(operator), typeof(link_index_map), + }( + ket, + operator, + link_index_map + ) + VD = typeof(DataGraphs.get_vertex_data(tmp, first(vertices(ket)))) + return QuadraticFormNetwork{ + V, VD, typeof(ket), typeof(operator), typeof(link_index_map), + }( + ket, + operator, + link_index_map + ) +end + +# === Environments === +# +# On a tree, the projected-Hamiltonian environment of `⟨ψ|H|ψ⟩` is exact: the message on +# directed edge `v → w` is the contraction of the entire `⟨ψ|H|ψ⟩` subtree on `v`'s side +# of the cut `(v, w)`. It carries the three bonds crossing the cut (ket, operator, bra). +# +# The messages are computed by a single dependency-ordered pass over every directed edge +# (`forest_cover_edge_sequence` visits each tree edge in both directions, sources before +# targets), with each message an exact contraction of `qf[v]` against the already-computed +# incoming messages from `v`'s other neighbors. No fixed-point iteration is needed. + +function incoming_subtree_messages(messages, graph, v, w) + return [ + NDA.state(messages[edgetype(graph)(u, v)]) for + u in neighbors(graph, v) if u != w + ] +end + +function environment_operator(message, link_index_map) + ketnames = [n for n in dimnames(message) if haskey(link_index_map, n)] + branames = [link_index_map[n] for n in ketnames] + return NDA.operator(message, branames, ketnames) +end + +""" + quadratic_form_environments(qf::QuadraticFormNetwork; root) -> MessageCache + +Exact projected-Hamiltonian environments of `⟨ψ|H|ψ⟩` on a tree, as a `MessageCache` of +`NamedDimsArrays` operators keyed by directed edges. The message on `v → w` is the +contraction of the `⟨ψ|H|ψ⟩` subtree on `v`'s side of `(v, w)`, wrapped as an operator +recording the bra ↔ ket link correspondence (see [`environment_operator`](@ref)). +""" +function quadratic_form_environments(qf::QuadraticFormNetwork; kwargs...) + sequence = forest_cover_edge_sequence(qf; kwargs...) + messages = Dict{edgetype(qf), Any}() + for e in sequence + v, w = src(e), dst(e) + incoming = incoming_subtree_messages(messages, qf, v, w) + message = contract_network([[qf[v]]; incoming]) + messages[e] = environment_operator(message, qf.link_index_map) + end + return messagecache(messages) +end diff --git a/src/tensornetworkoperator.jl b/src/tensornetworkoperator.jl new file mode 100644 index 00000000..4fe8cdf4 --- /dev/null +++ b/src/tensornetworkoperator.jl @@ -0,0 +1,38 @@ +using DataGraphs: DataGraphs, underlying_graph +using NamedDimsArrays: NamedDimsArrays as NDA +using NamedGraphs.GraphsExtensions: vertextype + +# A tensor-network operator: an operator `TensorNetwork` together with the map between its +# bra-side and ket-side physical names. The tensor-network analogue of a `NamedDimsArrays` +# operator, with `state` the underlying operator network (mirroring `NamedDimsArrays.state` +# on a `NamedDimsOperator`) and the tensor-network interface forwarded to it. +struct TensorNetworkOperator{V, VD, State, SiteMap} <: AbstractTensorNetwork{V, VD} + state::State + site_index_map::SiteMap +end + +NDA.state(o::TensorNetworkOperator) = o.state +site_index_map(o::TensorNetworkOperator) = o.site_index_map + +DataGraphs.underlying_graph(o::TensorNetworkOperator) = underlying_graph(NDA.state(o)) +function DataGraphs.is_vertex_assigned(o::TensorNetworkOperator, v) + return DataGraphs.is_vertex_assigned(NDA.state(o), v) +end +function DataGraphs.is_edge_assigned(o::TensorNetworkOperator, e) + return DataGraphs.is_edge_assigned(NDA.state(o), e) +end +function DataGraphs.get_vertex_data(o::TensorNetworkOperator, v) + return DataGraphs.get_vertex_data(NDA.state(o), v) +end + +function Base.copy(o::TensorNetworkOperator) + return TensorNetworkOperator(copy(NDA.state(o)), site_index_map(o)) +end + +function TensorNetworkOperator(state, site_index_map) + return TensorNetworkOperator{ + vertextype(state), eltype(state), typeof(state), typeof(site_index_map), + }( + state, site_index_map + ) +end From 0fbe80cb42c6ca08c1ca2c4127ea31cbfb752907 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 25 Jun 2026 09:26:51 -0400 Subject: [PATCH 3/4] Upgrade tree-DMRG to the ITensorBase 0.8 and TensorAlgebra 0.11 stack The DMRG code predated the INN v0.7 upgrade. Move it onto the released stack: named arrays and operators now come from ITensorBase rather than NamedDimsArrays (state, operator, apply, lazy, dimnames, replacedimnames, uniquename), the QR and SVD factorizations go through MatrixAlgebraKit (qr_compact, svd_trunc), and the tensor network types are ITensorNetwork and AbstractITensorNetwork. The VectorInterface methods that KrylovKit needs were dropped here and moved into ITensorBase (https://github.com/ITensor/ITensorBase.jl/pull/186), so the Aqua piracy check is back to the full default. --- Project.toml | 4 ++++ src/ITensorNetworksNext.jl | 1 - src/dmrg.jl | 31 +++++++++++++++---------------- src/nameddimsarrays_extensions.jl | 22 ---------------------- src/orthogonalize.jl | 6 +++--- src/quadraticformnetwork.jl | 15 +++++++-------- src/tensornetworkoperator.jl | 24 ++++++++++++------------ test/test_aqua.jl | 6 +----- test/test_dmrg.jl | 27 +++++++++++++-------------- 9 files changed, 55 insertions(+), 81 deletions(-) delete mode 100644 src/nameddimsarrays_extensions.jl diff --git a/Project.toml b/Project.toml index cb7aaf65..8b17183e 100644 --- a/Project.toml +++ b/Project.toml @@ -24,6 +24,10 @@ SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" +[sources.ITensorBase] +rev = "mf/vectorinterface-itensor" +url = "https://github.com/ITensor/ITensorBase.jl" + [compat] Adapt = "4.3" AlgorithmsInterface = "0.1" diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index 48192b44..f46db3d2 100644 --- a/src/ITensorNetworksNext.jl +++ b/src/ITensorNetworksNext.jl @@ -21,7 +21,6 @@ include("beliefpropagation/normnetwork.jl") include("apply/apply_operators.jl") -include("nameddimsarrays_extensions.jl") include("tensornetworkoperator.jl") include("quadraticformnetwork.jl") include("orthogonalize.jl") diff --git a/src/dmrg.jl b/src/dmrg.jl index b779baf1..ce5b9c65 100644 --- a/src/dmrg.jl +++ b/src/dmrg.jl @@ -2,10 +2,11 @@ import .AlgorithmsInterfaceExtensions as AIE import AlgorithmsInterface as AI using Base: @kwdef using Graphs: dst, edges, edgetype, src +using ITensorBase: + ITensorBase as ITB, dimnames, operator, replacedimnames, state, uniquename using KrylovKit: eigsolve -using NamedDimsArrays: NamedDimsArrays as NDA, dimnames, replacedimnames +using MatrixAlgebraKit: svd_trunc using NamedGraphs.GraphsExtensions: edge_path -using TensorAlgebra: TensorAlgebra as TA # ============================ Top-level entry point ===================================== @@ -35,7 +36,7 @@ function dmrg( trunc = nothing, regions = default_dmrg_regions(ket), link_index_map = Dict( - ln => randname(ln) for e in edges(ket) for ln in linknames(ket, e) + ln => uniquename(ln) for e in edges(ket) for ln in linknames(ket, e) ), alg = nothing, kwargs... @@ -361,14 +362,12 @@ function region_update_nsite!(::Val{2}, alg::TwoSiteEigsolve, qf, env, region) return qf, energy end -function svd_split(T, rows; trunc) - return isnothing(trunc) ? TA.svd(T, rows) : TA.svd(T, rows; trunc) -end +svd_split(T, rows; trunc) = svd_trunc(T, rows; trunc) # Lowest eigenpair of the effective Hamiltonian operator `H_eff` acting on the named tensor -# `T` (via `NamedDimsArrays.apply`), using `KrylovKit.eigsolve`. +# `T` (via `ITensorBase.apply`), using `KrylovKit.eigsolve`. function eigsolve_named(H_eff, T, which) - vals, vecs = eigsolve(x -> NDA.apply(H_eff, x), T, 1, which; ishermitian = true) + vals, vecs = eigsolve(x -> ITB.apply(H_eff, x), T, 1, which; ishermitian = true) return real(vals[1]), vecs[1] end @@ -376,29 +375,29 @@ end # obtained by contracting the operator tensors on the region with the incoming environment # messages on the region's boundary. Contracting those gives a single tensor whose codomain # (output) names are the region's bra names and whose domain (input) names are its ket names -# — i.e. exactly a `NamedDimsArrays` operator. So `effective_hamiltonian` returns that -# operator, and its action on a region ket tensor `T` is `NamedDimsArrays.apply(H, T)` (which -# contracts the ket names and renames the resulting bra names back to ket names). +# — i.e. exactly an ITensor operator. So `effective_hamiltonian` returns that operator, and +# its action on a region ket tensor `T` is `ITensorBase.apply(H, T)` (which contracts the ket +# names and renames the resulting bra names back to ket names). """ - effective_hamiltonian(qf::QuadraticFormNetwork, env, region) -> NamedDimsOperator + effective_hamiltonian(qf::QuadraticFormNetwork, env, region) -> ITensorOperator -Effective (projected) Hamiltonian for `region` (a vector of vertices) as a `NamedDimsArrays` +Effective (projected) Hamiltonian for `region` (a vector of vertices) as an ITensor operator: its domain (input) names are the region's ket names and its codomain (output) names are the matching bra names. Apply it to a region ket tensor `T` with -`NamedDimsArrays.apply`. The environment `env` is a `MessageCache` of +`ITensorBase.apply`. The environment `env` is a `MessageCache` of [`quadratic_form_environments`](@ref). """ function effective_hamiltonian(qf::QuadraticFormNetwork, env, region) operators = [operator_tensor(qf, v) for v in region] - boundary = [NDA.state(m) for m in incoming_edge_data(env, region)] + boundary = [state(m) for m in incoming_edge_data(env, region)] h = contract_network([operators; boundary]) sitemap = site_index_map(qf.operator) ketnames = [ n for n in dimnames(h) if haskey(sitemap, n) || haskey(qf.link_index_map, n) ] branames = [bra_name_map(qf)[n] for n in ketnames] - return NDA.operator(h, branames, ketnames) + return operator(h, branames, ketnames) end # ============================ Energy-based convergence ================================== diff --git a/src/nameddimsarrays_extensions.jl b/src/nameddimsarrays_extensions.jl deleted file mode 100644 index b53b9e5c..00000000 --- a/src/nameddimsarrays_extensions.jl +++ /dev/null @@ -1,22 +0,0 @@ -using NamedDimsArrays: AbstractNamedDimsArray -using VectorInterface: VectorInterface as VI - -# Temporary `VectorInterface` methods for named arrays. -# -# `KrylovKit.eigsolve` drives its Krylov vectors through `VectorInterface`. The generic -# `AbstractArray` fallbacks broadcast in a way that fails on a named array (e.g. -# `zerovector(x, S)`), so we provide name-aware methods here. This is type piracy on -# `AbstractNamedDimsArray` and is intended to move into `NamedDimsArrays`. -VI.scalartype(::Type{<:AbstractNamedDimsArray{T}}) where {T} = T -function VI.zerovector(x::AbstractNamedDimsArray, ::Type{S}) where {S <: Number} - return fill!(similar(x, S), zero(S)) -end -VI.scale(x::AbstractNamedDimsArray, α::Number) = x * α -VI.scale!!(x::AbstractNamedDimsArray, α::Number) = x * α -VI.scale!!(::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number) = x * α -function VI.add!!( - y::AbstractNamedDimsArray, x::AbstractNamedDimsArray, α::Number, β::Number - ) - return x * α + y * β -end -VI.inner(x::AbstractNamedDimsArray, y::AbstractNamedDimsArray) = (conj(x) * y)[] diff --git a/src/orthogonalize.jl b/src/orthogonalize.jl index 999aa233..03baaf4e 100644 --- a/src/orthogonalize.jl +++ b/src/orthogonalize.jl @@ -1,7 +1,7 @@ using Graphs: dst, src -using NamedDimsArrays: dimnames, replacedimnames +using ITensorBase: dimnames, replacedimnames +using MatrixAlgebraKit: qr_compact using NamedGraphs.GraphsExtensions: edge_path, post_order_dfs_edges -using TensorAlgebra: TensorAlgebra as TA # Isometric (QR) gauge on a tree. The orthogonality center is *not* stored on the network # — it is tracked by the caller. `orthogonalize(state, center)` canonicalizes the whole @@ -19,7 +19,7 @@ function gauge_move!(state, v, w) ln = only(linknames(state, v => w)) tv = state[v] rows = collect(setdiff(dimnames(tv), [ln])) - Q, R = TA.qr(tv, rows) + Q, R = qr_compact(tv, rows) r = only(setdiff(dimnames(Q), rows)) new_w = R * state[w] setindex_preserve_graph!(state, replacedimnames(Q, r => ln), v) diff --git a/src/quadraticformnetwork.jl b/src/quadraticformnetwork.jl index 84eddb77..7e90e227 100644 --- a/src/quadraticformnetwork.jl +++ b/src/quadraticformnetwork.jl @@ -1,17 +1,16 @@ -using .LazyNamedDimsArrays: lazy using DataGraphs: DataGraphs, underlying_graph using Graphs: dst, edgetype, neighbors, src, vertices -using NamedDimsArrays: NamedDimsArrays as NDA, dimnames, replacedimnames +using ITensorBase: dimnames, lazy, operator, replacedimnames, state using NamedGraphs.GraphsExtensions: vertextype -# A lazy `⟨ψ|H|ψ⟩` network: a ket `TensorNetwork`, a `TensorNetworkOperator` (which carries +# A lazy `⟨ψ|H|ψ⟩` network: a ket `ITensorNetwork`, a `TensorNetworkOperator` (which carries # the ket → bra *site* name map), and a forward ket → bra map for the *link* names. The bra # layer is derived from the ket (`conj` + index renaming), never stored, so updating a ket -# tensor is reflected in the bra. As an `AbstractTensorNetwork`, the data on vertex `v` is +# tensor is reflected in the bra. As an `AbstractITensorNetwork`, the data on vertex `v` is # the lazy product `lazy(ket) * lazy(operator) * lazy(bra)`, so the existing # `contract_network` / `MessageCache` machinery treats it like any other tensor network. struct QuadraticFormNetwork{V, VD, Ket, Operator, LinkMap} <: - AbstractTensorNetwork{V, VD} + AbstractITensorNetwork{V, VD} ket::Ket operator::Operator link_index_map::LinkMap @@ -79,7 +78,7 @@ end function incoming_subtree_messages(messages, graph, v, w) return [ - NDA.state(messages[edgetype(graph)(u, v)]) for + state(messages[edgetype(graph)(u, v)]) for u in neighbors(graph, v) if u != w ] end @@ -87,14 +86,14 @@ end function environment_operator(message, link_index_map) ketnames = [n for n in dimnames(message) if haskey(link_index_map, n)] branames = [link_index_map[n] for n in ketnames] - return NDA.operator(message, branames, ketnames) + return operator(message, branames, ketnames) end """ quadratic_form_environments(qf::QuadraticFormNetwork; root) -> MessageCache Exact projected-Hamiltonian environments of `⟨ψ|H|ψ⟩` on a tree, as a `MessageCache` of -`NamedDimsArrays` operators keyed by directed edges. The message on `v → w` is the +ITensor operators keyed by directed edges. The message on `v → w` is the contraction of the `⟨ψ|H|ψ⟩` subtree on `v`'s side of `(v, w)`, wrapped as an operator recording the bra ↔ ket link correspondence (see [`environment_operator`](@ref)). """ diff --git a/src/tensornetworkoperator.jl b/src/tensornetworkoperator.jl index 4fe8cdf4..90cada30 100644 --- a/src/tensornetworkoperator.jl +++ b/src/tensornetworkoperator.jl @@ -1,32 +1,32 @@ using DataGraphs: DataGraphs, underlying_graph -using NamedDimsArrays: NamedDimsArrays as NDA +using ITensorBase: ITensorBase as ITB using NamedGraphs.GraphsExtensions: vertextype -# A tensor-network operator: an operator `TensorNetwork` together with the map between its -# bra-side and ket-side physical names. The tensor-network analogue of a `NamedDimsArrays` -# operator, with `state` the underlying operator network (mirroring `NamedDimsArrays.state` -# on a `NamedDimsOperator`) and the tensor-network interface forwarded to it. -struct TensorNetworkOperator{V, VD, State, SiteMap} <: AbstractTensorNetwork{V, VD} +# A tensor-network operator: an operator `ITensorNetwork` together with the map between its +# bra-side and ket-side physical names. The tensor-network analogue of an ITensor operator, +# with `state` the underlying operator network (mirroring `ITensorBase.state` on an +# `ITensorOperator`) and the tensor-network interface forwarded to it. +struct TensorNetworkOperator{V, VD, State, SiteMap} <: AbstractITensorNetwork{V, VD} state::State site_index_map::SiteMap end -NDA.state(o::TensorNetworkOperator) = o.state +ITB.state(o::TensorNetworkOperator) = o.state site_index_map(o::TensorNetworkOperator) = o.site_index_map -DataGraphs.underlying_graph(o::TensorNetworkOperator) = underlying_graph(NDA.state(o)) +DataGraphs.underlying_graph(o::TensorNetworkOperator) = underlying_graph(ITB.state(o)) function DataGraphs.is_vertex_assigned(o::TensorNetworkOperator, v) - return DataGraphs.is_vertex_assigned(NDA.state(o), v) + return DataGraphs.is_vertex_assigned(ITB.state(o), v) end function DataGraphs.is_edge_assigned(o::TensorNetworkOperator, e) - return DataGraphs.is_edge_assigned(NDA.state(o), e) + return DataGraphs.is_edge_assigned(ITB.state(o), e) end function DataGraphs.get_vertex_data(o::TensorNetworkOperator, v) - return DataGraphs.get_vertex_data(NDA.state(o), v) + return DataGraphs.get_vertex_data(ITB.state(o), v) end function Base.copy(o::TensorNetworkOperator) - return TensorNetworkOperator(copy(NDA.state(o)), site_index_map(o)) + return TensorNetworkOperator(copy(ITB.state(o)), site_index_map(o)) end function TensorNetworkOperator(state, site_index_map) diff --git a/test/test_aqua.jl b/test/test_aqua.jl index 6d3276da..8eb46122 100644 --- a/test/test_aqua.jl +++ b/test/test_aqua.jl @@ -3,9 +3,5 @@ using ITensorNetworksNext: ITensorNetworksNext using Test: @testset @testset "Code quality (Aqua.jl)" begin - # Piracy is checked separately as `broken`: `dmrg.jl` temporarily pirates a few - # `VectorInterface` methods on `AbstractNamedDimsArray` (needed by `KrylovKit.eigsolve`). - # These are slated to move into `NamedDimsArrays`; drop the `broken` marker once they do. - Aqua.test_all(ITensorNetworksNext; persistent_tasks = false, piracies = false) - Aqua.test_piracies(ITensorNetworksNext; broken = true) + Aqua.test_all(ITensorNetworksNext; persistent_tasks = false) end diff --git a/test/test_dmrg.jl b/test/test_dmrg.jl index 24e99a17..57a74746 100644 --- a/test/test_dmrg.jl +++ b/test/test_dmrg.jl @@ -1,9 +1,8 @@ using Graphs: add_edge!, edges, vertices -using ITensorBase: Index -using ITensorNetworksNext: TensorNetwork, TensorNetworkOperator, dmrg, insertlink! +using ITensorBase: Index, name, nameddims, setname, uniquename +using ITensorNetworksNext: ITensorNetwork, TensorNetworkOperator, dmrg, insertlink! using LinearAlgebra: eigen using MatrixAlgebraKit: truncrank -using NamedDimsArrays: name, nameddims, randname, setname using NamedGraphs.NamedGraphGenerators: named_path_graph using NamedGraphs: NamedGraph using StableRNGs: StableRNG @@ -34,7 +33,7 @@ function tfim_path_operator(N, sites, sitemap; J, h) verts = collect(1:N) bond_edges = [(verts[i], verts[i + 1]) for i in 1:(N - 1)] bonds = Dict( - e => setname(Index(Base.OneTo(3)), randname(name(sites[verts[1]]))) + e => setname(Index(Base.OneTo(3)), uniquename(name(sites[verts[1]]))) for e in bond_edges ) W = ising_mpo_tensor(; J, h) @@ -58,7 +57,7 @@ function tfim_path_operator(N, sites, sitemap; J, h) W, (name(bonds[bond_edges[left]]), name(bonds[bond_edges[right]]), k, b) ) end - tn = TensorNetwork(g) do v + tn = ITensorNetwork(g) do v return tensor(v) end return TensorNetworkOperator(tn, sitemap) @@ -83,7 +82,7 @@ function tfim_star_operator(sites, sitemap; J, h) add_edge!(g, 1, v) end bonds = Dict( - v => setname(Index(Base.OneTo(3)), randname(name(sites[1]))) for v in leaves + v => setname(Index(Base.OneTo(3)), uniquename(name(sites[1]))) for v in leaves ) # Center tensor: one bond per leaf, with `-hX` once and `-JZ`/`I` on each leaf channel. center = zeros(3, 3, 3, 2, 2) @@ -105,7 +104,7 @@ function tfim_star_operator(sites, sitemap; J, h) end tensors = Dict(1 => nameddims(center, center_names), (v => leaf_tensor(v) for v in leaves)...) - tn = TensorNetwork(g) do v + tn = ITensorNetwork(g) do v return tensors[v] end return TensorNetworkOperator(tn, sitemap) @@ -124,7 +123,7 @@ end function random_ket(rng, g) sites = Dict(v => Index(Base.OneTo(2)) for v in vertices(g)) - ket = TensorNetwork(NamedGraph(collect(vertices(g)))) do v + ket = ITensorNetwork(NamedGraph(collect(vertices(g)))) do v return randn(rng, Float64, (sites[v],)) end for e in edges(g) @@ -138,7 +137,7 @@ end rng = StableRNG(8) g = named_path_graph(N) ket0, sites = random_ket(rng, g) - sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + sitemap = Dict(name(sites[v]) => uniquename(name(sites[v])) for v in vertices(g)) operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) @@ -151,7 +150,7 @@ end rng = StableRNG(8) g = named_path_graph(N) ket0, sites = random_ket(rng, g) - sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + sitemap = Dict(name(sites[v]) => uniquename(name(sites[v])) for v in vertices(g)) operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) @@ -165,7 +164,7 @@ end rng = StableRNG(8) g = named_path_graph(N) ket0, sites = random_ket(rng, g) - sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + sitemap = Dict(name(sites[v]) => uniquename(name(sites[v])) for v in vertices(g)) operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) @@ -180,7 +179,7 @@ end rng = StableRNG(8) g = named_path_graph(N) ket0, sites = random_ket(rng, g) - sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + sitemap = Dict(name(sites[v]) => uniquename(name(sites[v])) for v in vertices(g)) operator = tfim_path_operator(N, sites, sitemap; J = 1.0, h = 0.7) exact = minimum(eigen(tfim_path_dense(N; J = 1.0, h = 0.7)).values) @@ -197,7 +196,7 @@ end add_edge!(g, 1, v) end ket0, sites = random_ket(rng, g) - sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + sitemap = Dict(name(sites[v]) => uniquename(name(sites[v])) for v in vertices(g)) operator = tfim_star_operator(sites, sitemap; J = 1.0, h = 0.6) exact = minimum(eigen(tfim_star_dense(; J = 1.0, h = 0.6)).values) @@ -209,7 +208,7 @@ end rng = StableRNG(8) g = named_path_graph(4) ket0, sites = random_ket(rng, g) - sitemap = Dict(name(sites[v]) => randname(name(sites[v])) for v in vertices(g)) + sitemap = Dict(name(sites[v]) => uniquename(name(sites[v])) for v in vertices(g)) operator = tfim_path_operator(4, sites, sitemap; J = 1.0, h = 0.7) @test_throws ArgumentError dmrg(operator, ket0) end From f17e10cae2dc15c9684e75ae74575ba9d0722841 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 25 Jun 2026 14:34:51 -0400 Subject: [PATCH 4/4] Drop the ITensorBase sources pin now that 0.8.2 is registered --- Project.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/Project.toml b/Project.toml index 8b17183e..cb7aaf65 100644 --- a/Project.toml +++ b/Project.toml @@ -24,10 +24,6 @@ SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -[sources.ITensorBase] -rev = "mf/vectorinterface-itensor" -url = "https://github.com/ITensor/ITensorBase.jl" - [compat] Adapt = "4.3" AlgorithmsInterface = "0.1"