DMRG3S: strictly single-site DMRG subspace expansion (Hubig et al. 2015)

src/MPSKit.jl

@@ -38,6 +38,8 @@ export time_evolve, timestep, timestep!, make_time_mpo
 export TDVP, TDVP2, WI, WII, TaylorCluster
 export changebonds, changebonds!
 export VUMPSSvdCut, OptimalExpand, SvdCut, RandExpand, SketchedExpand
+export post_expand!, NoExpand
+export NoiseSchedule, FunctionalSchedule, ExponentialDecay, Warmup, DMRG3S
 export propagator
 export DynamicalDMRG, NaiveInvert, Jeckelmann
 export exact_diagonalization, fidelity_susceptibility
@@ -160,6 +162,9 @@ include("algorithms/changebonds/svdcut.jl")
 include("algorithms/changebonds/randexpand.jl")
 include("algorithms/changebonds/sketchedexpand.jl")
 
+include("algorithms/post_expand/post_expand.jl")
+include("algorithms/post_expand/dmrg3s.jl")
+
 include("algorithms/timestep/tdvp.jl")
 include("algorithms/timestep/taylorcluster.jl")
 include("algorithms/timestep/wii.jl")

src/algorithms/groundstate/dmrg.jl

@@ -9,12 +9,17 @@ $(TYPEDEF)
 Density Matrix Renormalization Group algorithm for finding the dominant eigenvector.
 
 Each site update is, in order: (1) an optional bond expansion (`alg_expand`), (2) a single-site
-eigensolve, and (3) a gauge step (`alg_gauge`). With the defaults (`alg_expand = nothing` and a
-non-truncating QR `alg_gauge`) this is textbook single-site DMRG, which cannot change the bond
-dimension. Setting `alg_expand` to a bond-expansion algorithm (e.g. [`OptimalExpand`](@ref),
-[`RandExpand`](@ref), [`SketchedExpand`](@ref)) enriches the bond with directions orthogonal to
-the current state ahead of each eigensolve, recovering Controlled Bond Expansion (CBE) DMRG; a
-truncating `alg_gauge` is then desirable to cut the enlarged bond back down.
+eigensolve, and (3) a post-expansion and gauge step (`alg_post_expand`, `alg_gauge`). With the
+defaults (`alg_expand = nothing`, `alg_post_expand = NoExpand()`, and a non-truncating QR
+`alg_gauge`) this is textbook single-site DMRG, which cannot change the bond dimension. Setting
+`alg_expand` to a bond-expansion algorithm (e.g. [`OptimalExpand`](@ref), [`RandExpand`](@ref),
+[`SketchedExpand`](@ref)) enriches the bond with directions orthogonal to the current state ahead
+of each eigensolve, recovering Controlled Bond Expansion (CBE) DMRG. Setting `alg_post_expand` to
+[`DMRG3S`](@ref) instead enriches the bond *after* the eigensolve, using a Hamiltonian-derived
+perturbation of the just-optimized tensor, following Hubig, McCulloch, Schollwöck & Wolf,
+[Phys. Rev. B 91, 155115 (2015)](https://doi.org/10.1103/PhysRevB.91.155115); this can be used on
+its own or alongside `alg_expand`. A truncating `alg_gauge` is desirable whenever either expansion
+mechanism is active, to cut the enlarged bond back down.
 
 The gauge algorithm is selected in the keyword constructor from the `trscheme` argument: when it
 is `notrunc()` the gauge is a QR decomposition (`alg_orth`, [`Householder`](@extref
@@ -25,7 +30,7 @@ given `trscheme`).
 
 $(TYPEDFIELDS)
 """
-struct DMRG{A, F, E, G} <: Algorithm
+struct DMRG{A, F, E, PE, G} <: Algorithm
     "tolerance for convergence criterium"
     tol::Float64
 
@@ -44,26 +49,43 @@ struct DMRG{A, F, E, G} <: Algorithm
     "algorithm used to expand the bond ahead of each local update, or `nothing` for none"
     alg_expand::E
 
+    "algorithm used to expand the bond after the local update, e.g. `NoExpand()` (default) or
+    [`DMRG3S`](@ref)"
+    alg_post_expand::PE
+
     "factorization used for the post-update gauge: a QR algorithm (no truncation) or a truncated SVD"
     alg_gauge::G
 end
 function DMRG(;
         tol = Defaults.tol, maxiter = Defaults.maxiter, alg_eigsolve = (;),
         verbosity = Defaults.verbosity, finalize = Defaults._finalize,
-        alg_expand = nothing, trscheme = notrunc(),
+        alg_expand = nothing, alg_post_expand = NoExpand(), trscheme = notrunc(),
         alg_svd = Defaults.alg_svd(), alg_orth = Defaults.alg_orth()
     )
     alg_eigsolve′ = alg_eigsolve isa NamedTuple ? Defaults.alg_eigsolve(; alg_eigsolve...) :
         alg_eigsolve
     # a no-truncation `trscheme` selects a (bond-preserving) QR gauge, anything else a truncated SVD
     alg_gauge = trscheme isa MatrixAlgebraKit.NoTruncation ? alg_orth :
         MatrixAlgebraKit.TruncatedAlgorithm(alg_svd, trscheme)
-    if !isnothing(alg_expand) && !_truncates(alg_gauge)
-        @warn "DMRG with `alg_expand` but no truncation (`trscheme = notrunc()`): the bond dimension will grow unboundedly each sweep."
+    if (!isnothing(alg_expand) || alg_post_expand ≠ NoExpand()) && !_truncates(alg_gauge)
+        @warn "DMRG with `alg_expand` or `alg_post_expand` but no truncation (`trscheme = notrunc()`): the bond dimension will grow unboundedly each sweep."
     end
-    return DMRG(tol, maxiter, verbosity, alg_eigsolve′, finalize, alg_expand, alg_gauge)
+    return DMRG(tol, maxiter, verbosity, alg_eigsolve′, finalize, alg_expand, alg_post_expand, alg_gauge)
 end
 
+
+"""
+    _sweep_range(alg::DMRG, chain_length::Int, direction::Val)
+
+Site positions to visit for one sweep in `direction` (`:right` or `:left`).
+"""
+_sweep_range(::DMRG, chain_length::Int, ::Val)
+
+_sweep_range(::DMRG, chain_length::Int, ::Val{:right}) = 1:(chain_length - 1)
+
+_sweep_range(::DMRG, chain_length::Int, ::Val{:left}) = reverse(2:chain_length)
+
+
 function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environments(ψ, H, ψ))
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
     ϵ = maximum(ϵs)
@@ -75,47 +97,31 @@ function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environme
         @infov 2 loginit!(log, ϵ, expectation_value(ψ, H, envs))
         for iter in 1:(alg.maxiter)
             alg_eigsolve = updatetol(alg.alg_eigsolve, iter, ϵ)
+            alg_post_expand = _update_post_expand(alg.alg_post_expand, iter, ϵ)
 
             zerovector!(ϵs)
             @timeit timeroutput "sweep" begin
-                # left-to-right
-                for pos in 1:(length(ψ) - 1)
-                    local AC′
-                    # convergence: pre-expansion single-site Galerkin error
-                    ϵs[pos] = max(ϵs[pos], calc_galerkin(pos, ψ, H, ψ, envs))
-
-                    # 1. expand
-                    isnothing(alg.alg_expand) ||
-                        @timeit timeroutput "expand" changebond!(pos, Val(:right), ψ, H, alg.alg_expand, envs)
-
-                    # 2. local update
-                    @timeit timeroutput "AC_eigsolve" begin
-                        Hac = AC_hamiltonian(pos, ψ, H, ψ, envs)
-                        _, AC′ = fixedpoint(Hac, ψ.AC[pos], :SR, alg_eigsolve)
-                    end
-
-                    # 3. gauge (QR center-move or truncated SVD, selected by `alg_gauge`)
-                    @timeit timeroutput "gauge" left_gauge!(ψ, pos, AC′, alg.alg_gauge; normalize = true)
-                end
-
-                # right-to-left
-                for pos in length(ψ):-1:2
-                    local AC′
-                    # convergence: pre-expansion single-site Galerkin error
-                    ϵs[pos] = max(ϵs[pos], calc_galerkin(pos, ψ, H, ψ, envs))
-
-                    # 1. expand
-                    isnothing(alg.alg_expand) ||
-                        @timeit timeroutput "expand" changebond!(pos, Val(:left), ψ, H, alg.alg_expand, envs)
-
-                    # 2. local update
-                    @timeit timeroutput "AC_eigsolve" begin
-                        Hac = AC_hamiltonian(pos, ψ, H, ψ, envs)
-                        _, AC′ = fixedpoint(Hac, ψ.AC[pos], :SR, alg_eigsolve)
+                for direction in (:right, :left)
+                    for pos in _sweep_range(alg, length(ψ), Val(direction))
+                        local AC′
+                        # convergence: pre-expansion single-site Galerkin error
+                        ϵs[pos] = max(ϵs[pos], calc_galerkin(pos, ψ, H, ψ, envs))
+
+                        # 1. expand
+                        isnothing(alg.alg_expand) ||
+                            @timeit timeroutput "expand" changebond!(pos, Val(direction), ψ, H, alg.alg_expand, envs)
+
+                        # 2. local update
+                        @timeit timeroutput "AC_eigsolve" begin
+                            Hac = AC_hamiltonian(pos, ψ, H, ψ, envs)
+                            _, AC′ = fixedpoint(Hac, ψ.AC[pos], :SR, alg_eigsolve)
+                        end
+
+                        # expand post optimization
+                        @timeit timeroutput "post expand and gauge" begin
+                            post_expand!(pos, Val(direction), ψ, H, alg_post_expand, envs, AC′, alg.alg_gauge; normalize = true)
+                        end
                     end
-
-                    # 3. gauge (QR center-move or truncated SVD, selected by `alg_gauge`)
-                    @timeit timeroutput "gauge" right_gauge!(ψ, pos, AC′, alg.alg_gauge; normalize = true)
                 end
             end
             ϵ = maximum(ϵs)

src/algorithms/post_expand/dmrg3s.jl

@@ -0,0 +1,140 @@
+"""
+    NoiseSchedule
+
+Supertype for rules that evolve [`DMRG3S`](@ref)'s noise amplitude `α` between outer
+DMRG iterations. A schedule `s::NoiseSchedule` is called as `s(noise, iter, ϵ)`,
+returning the noise amplitude to use for the next iteration.
+"""
+abstract type NoiseSchedule end
+
+"""
+    FunctionalSchedule(f) <: NoiseSchedule
+
+Wrap an arbitrary callable `f(noise, iter, ϵ) -> noise` as a [`NoiseSchedule`](@ref).
+Used internally by `∘` to compose schedules, but can also be constructed directly for
+ad hoc schedules that don't warrant their own named type.
+"""
+struct FunctionalSchedule{F} <: NoiseSchedule
+    f::F
+end
+(s::FunctionalSchedule)(noise, iter, ϵ) = s.f(noise, iter, ϵ)
+
+"""
+    s1::NoiseSchedule ∘ s2::NoiseSchedule
+
+Compose two noise schedules: `s2` is applied first, and its result is fed through `s1`,
+i.e. `(s1 ∘ s2)(noise, iter, ϵ) == s1(s2(noise, iter, ϵ), iter, ϵ)` — the same convention
+as function composition in `Base`.
+"""
+Base.:∘(s1::NoiseSchedule, s2::NoiseSchedule) =
+    FunctionalSchedule((noise, iter, ϵ) -> s1(s2(noise, iter, ϵ), iter, ϵ))
+
+"""
+    ExponentialDecay(decay_rate; threshold = 0.0)
+
+Noise schedule that shrinks geometrically: `noise -> noise * decay_rate^iter`, snapped
+to exactly zero once it falls below `threshold`. Use `decay_rate < 1` for a standalone
+[`DMRG3S`](@ref) run that gradually turns enrichment off as the state converges; a
+nonzero `threshold` avoids running the (cheap, but non-free) expansion step
+indefinitely on a noise amplitude too small to matter.
+"""
+struct ExponentialDecay{T <: Real} <: NoiseSchedule
+    decay_rate::T
+    threshold::T
+end
+ExponentialDecay(decay_rate, threshold) = ExponentialDecay(promote(decay_rate, threshold)...)
+ExponentialDecay(decay_rate; threshold = zero(decay_rate)) = ExponentialDecay(decay_rate, threshold)
+
+function (s::ExponentialDecay)(noise, iter, ϵ)
+    decayed = noise * s.decay_rate^iter
+    return decayed < s.threshold ? zero(decayed) : decayed
+end
+
+"""
+    Warmup(iters)
+
+Noise schedule that holds the noise amplitude constant for `iters` outer iterations,
+then drops it to zero. Useful when [`DMRG3S`](@ref) is combined with another bond
+expansion algorithm (e.g. [`SketchedExpand`](@ref)) and is only needed to help escape local minima during
+the first few sweeps.
+"""
+struct Warmup <: NoiseSchedule
+    iters::Int
+end
+
+(s::Warmup)(noise, iter, ϵ) = iter ≤ s.iters ? noise : zero(noise)
+
+"""
+    DMRG3S(noise, schedule::NoiseSchedule) <: Algorithm
+
+Strictly single-site DMRG subspace expansion (Hubig, McCulloch, Schollwöck & Wolf,
+[Phys. Rev. B 91, 155115 (2015)](https://doi.org/10.1103/PhysRevB.91.155115)).
+Enriches the bond at each site update with a term built from the local Hamiltonian
+and environment, scaled by mixing factor `noise`, before truncating via the sweep's
+`alg_gauge`. This lets single-site DMRG introduce basis states/quantum-number sectors
+that weren't present in the initial state, avoiding the local minima plain single-site
+DMRG can get stuck in (see the paper's Section VII A for a worked example).
+
+`schedule` determines how `noise` evolves between outer iterations; see
+[`ExponentialDecay`](@ref) and [`Warmup`](@ref). Once the schedule returns zero,
+`DMRG3S` is replaced by [`NoExpand`](@ref) for the remainder of the run.
+"""
+struct DMRG3S{N, S <: NoiseSchedule} <: Algorithm
+    noise::N
+    schedule::S
+end
+
+function _update_post_expand(alg::DMRG3S, iter, ϵ)
+    noise = alg.schedule(alg.noise, iter, ϵ)
+    return iszero(noise) ? NoExpand() : DMRG3S(noise, alg.schedule)
+end
+
+function _get_combiner(::Type{T}, V1, V2) where {T}
+    Vf = fuse(V1 ⊗ V2)
+    return isomorphism(T, (V1 ⊗ V2) ← Vf), Vf
+end
+
+function post_expand!(pos::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::DMRG3S, envs, AC, alg_gauge; normalize = true)
+    El = leftenv(envs, pos, ψ)
+    Hi = H[pos]
+    α = alg.noise
+    T = promote_type(scalartype(ψ), scalartype(Hi))
+    V = right_virtualspace(AC)
+    combiner, Vpert = _get_combiner(T, V, right_virtualspace(Hi))
+
+    @plansor pert[-1 -2; -3] := α * El[-1 1; 2] * AC[2 3; 4] * Hi[1 -2; 3 5] * combiner[4, 5; -3]
+
+    AC_expanded = catdomain(AC, pert)
+
+    AL, C = left_gauge(AC_expanded, alg_gauge)
+    B = _transpose_tail(ψ.AR[pos + 1])
+    AR = _transpose_front(catcodomain(B, zeros(T, Vpert ← domain(B))))
+
+    normalize && normalize!(C)
+    ψ.AC[pos] = (AL, C)
+    ψ.AC[pos + 1] = (C, AR)
+    return ψ
+end
+
+function post_expand!(pos::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::DMRG3S, envs, AC, alg_gauge; normalize = true)
+    Er = rightenv(envs, pos, ψ)
+    Hi = H[pos]
+    α = alg.noise
+    T = promote_type(scalartype(ψ), scalartype(Hi))
+    V = left_virtualspace(AC)
+    combiner, Vpert = _get_combiner(T, V, left_virtualspace(Hi))
+
+    @plansor pert[l; r s] := α * (combiner')[l; li lh] * AC[li, si; ri] * Hi[lh, s; si, rh] * Er[ri rh; r]
+
+    AC = _transpose_tail(AC)
+    AC_expanded = catcodomain(AC, pert)
+
+    C, AR = right_gauge(AC_expanded, alg_gauge)
+    B = ψ.AL[pos - 1]
+    AL = catdomain(B, zeros(T, codomain(B) ← Vpert))
+
+    normalize && normalize!(C)
+    ψ.AC[pos] = (C, AR)
+    ψ.AC[pos - 1] = (AL, C)
+    return ψ
+end

src/algorithms/post_expand/post_expand.jl

@@ -0,0 +1,29 @@
+"""
+    NoExpand <: Algorithm
+
+Sentinel algorithm: perform the standard center-move gauge with no subspace expansion.
+"""
+struct NoExpand <: Algorithm end
+
+_update_post_expand(::NoExpand, iter, ϵ) = NoExpand()
+
+"""
+    post_expand!(pos, direction::Val, ψ, H, alg, envs, AC, alg_gauge; normalize=false)
+
+Advance the DMRG sweep past `pos` in `direction` (`Val(:right)` or `Val(:left)`),
+using `alg` to optionally enrich the bond before gauging `AC` into `ψ` with `alg_gauge`.
+Mutates `ψ.AC[pos]` (and, if `alg` expands the bond, the neighboring tensor as well)
+and returns `ψ`.
+
+`alg::NoExpand` performs a plain gauge move with no enrichment; see [`DMRG3S`](@ref)
+for an algorithm that also injects a Hamiltonian-derived perturbation.
+"""
+post_expand!(::Int, ::Val, ::AbstractFiniteMPS, H, ::NoExpand, envs, AC, alg_gauge)
+
+function post_expand!(pos::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, ::NoExpand, envs, AC, alg_gauge; normalize::Bool = false)
+    return left_gauge!(ψ, pos, AC, alg_gauge; normalize)
+end
+
+function post_expand!(pos::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, ::NoExpand, envs, AC, alg_gauge; normalize::Bool = false)
+    return right_gauge!(ψ, pos, AC, alg_gauge; normalize)
+end

test/algorithms/groundstate.jl

@@ -115,6 +115,58 @@ verbosity_conv = 1
         @test dim(left_virtualspace(ψ, L ÷ 2)) == D
     end
 
+    @testset "DMRG3S" begin
+        # start from a small bond so the post-expansion is exercised, mirroring CBEDMRG above
+        Random.seed!(1234)
+        ψ₀ = FiniteMPS(randn, ComplexF64, L, ℙ^2, ℙ^(D ÷ 2))
+        v₀ = variance(ψ₀, H)
+        alg_post_expand = DMRG3S(0.1, ExponentialDecay(0.7))  # TODO: match final constructor API
+        trscheme = truncrank(D)
+
+        # test logging
+        ψ, envs, δ = find_groundstate(
+            ψ₀, H, DMRG(; verbosity = verbosity_full, maxiter = 2, alg_post_expand, trscheme)
+        )
+
+        ψ, envs, δ = find_groundstate(
+            ψ, H, DMRG(; verbosity = verbosity_conv, maxiter = 10, alg_post_expand, trscheme), envs
+        )
+        v = variance(ψ, H)
+
+        # test using low variance
+        @test sum(δ) ≈ 0 atol = 1.0e-3
+        @test v < v₀
+        @test v < 1.0e-2
+        # the bond should have grown to the truncation target
+        @test dim(left_virtualspace(ψ, L ÷ 2)) == D
+    end
+
+    @testset "DMRG3S escapes local minimum (Hubig et al. 2015, Sec. VII A)" begin
+        L_heis = 20
+        H_heis = heisenberg_XXX(ComplexF64, U1Irrep; spin = 1 // 2, L = L_heis)
+
+        Random.seed!(1234)
+        ψ_bad = bad_initial_state(H_heis, L_heis)
+
+        ψ_stuck, envs_stuck, δ_stuck = find_groundstate(
+            ψ_bad, H_heis, DMRG(; verbosity = verbosity_conv, maxiter = 30)
+        )
+        E_stuck = real(expectation_value(ψ_stuck, H_heis, envs_stuck))
+
+        alg_post_expand = DMRG3S(0.1, ExponentialDecay(0.8))
+        ψ_escape, envs_escape, δ_escape = find_groundstate(
+            ψ_bad, H_heis, DMRG(;
+                verbosity = verbosity_conv, maxiter = 30,
+                alg_post_expand, trscheme = truncrank(20),
+            )
+        )
+        E_escape = real(expectation_value(ψ_escape, H_heis, envs_escape))
+
+        # paper reports E(α=0) = -6.35479, E(α≠0) = -8.6824724 for this L_heis = 20, S = 1/2 AFM setup
+        @test E_escape < E_stuck - 1.0
+        @test isapprox(E_escape, -8.6824724; atol = 1.0e-4)
+    end
+
     @testset "GradientGrassmann" begin
         ψ₀ = FiniteMPS(randn, ComplexF64, 10, ℙ^2, ℙ^D)
         v₀ = variance(ψ₀, H)