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

Project.toml

@@ -36,7 +36,7 @@ MPSKitAdaptExt = "Adapt"
 [compat]
 Accessors = "0.1"
 Adapt = "4"
-BlockTensorKit = "0.3.14"
+BlockTensorKit = "0.3.15"
 Compat = "3.47, 4.10"
 DocStringExtensions = "0.9.3"
 HalfIntegers = "1.6.0"

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 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/approximate/fvomps.jl

@@ -8,7 +8,7 @@ function approximate!(ψ::AbstractFiniteMPS, Oϕ, alg::DMRG2, envs = environment
             ϵ = 0.0
             for pos in [1:(length(ψ) - 1); (length(ψ) - 2):-1:1]
                 AC2′ = AC2_projection(pos, ψ, Oϕ, envs)
-                al, c, ar, = svd_trunc!(AC2′, alg.alg_gauge)
+                al, c, ar, = svd_trunc!(AC2′, inner_alg_gauge(alg))
 
                 AC2 = ψ.AC[pos] * _transpose_tail(ψ.AR[pos + 1])
                 ϵ = max(ϵ, norm(al * c * ar - AC2) / norm(AC2))

src/algorithms/groundstate/dmrg.jl

@@ -2,24 +2,54 @@
 # center-move as in textbook single-site DMRG)
 _truncates(::MatrixAlgebraKit.AbstractAlgorithm) = false
 _truncates(::MatrixAlgebraKit.TruncatedAlgorithm) = true
+_truncates(alg::NoExpand) = _truncates(alg.alg_gauge)
+_truncates(alg::DMRG3S) = _truncates(alg.alg_gauge)
+
+# a no-truncation `trscheme` selects a (bond-preserving) QR gauge, anything else a truncated SVD
+_build_inner_gauge(trscheme, alg_svd, alg_orth) =
+    trscheme isa MatrixAlgebraKit.NoTruncation ? alg_orth :
+    MatrixAlgebraKit.TruncatedAlgorithm(alg_svd, trscheme)
+
+_expands(::NoExpand) = false
+_expands(::DMRG3S) = true
 
 """
 $(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.
-
-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
-MatrixAlgebraKit.Householder) by default), otherwise it is a truncated SVD (`alg_svd` with the
-given `trscheme`).
+eigensolve, and (3) a gauge step (`alg_gauge`). With the defaults (`alg_expand = nothing` and
+`alg_gauge = NoExpand()`, a non-truncating QR 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_gauge` to a bond-expanding gauge algorithm (e.g. [`DMRG3S`](@ref))
+instead enriches the bond as part of the gauge step, after the eigensolve. Either way, a
+truncating gauge (see below) is then desirable to cut the enlarged bond back down.
+
+# Choosing the gauge
+
+By default, `alg_gauge` is built for you from `trscheme`/`alg_svd`/`alg_orth`: `trscheme =
+notrunc()` (the default) gives a QR decomposition (`alg_orth`, [`Householder`](@extref
+MatrixAlgebraKit.Householder) by default), any other `trscheme` gives a truncated SVD (`alg_svd`
+with that `trscheme`).
+
+```julia
+DMRG()                            # QR gauge, no truncation
+DMRG(; trscheme = truncdim(50))   # truncated SVD gauge
+```
+
+To use a bond-expanding gauge such as [`DMRG3S`](@ref), pass it directly as `alg_gauge`; `trscheme`
+etc. are still routed through to build the *inner* gauge it wraps, exactly as above:
+
+```julia
+DMRG(; alg_gauge = DMRG3S(0.1, ExponentialDecay(0.7)), trscheme = truncdim(50))
+```
+
+If `alg_gauge` is instead given with its inner gauge already set (e.g. `DMRG3S(0.1, sched,
+some_gauge)`), `trscheme`/`alg_svd`/`alg_orth` must be left at their defaults — passing both is an
+error, since it leaves two conflicting sources for the same setting.
 
 ## Fields
 
@@ -44,24 +74,39 @@ 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
 
-    "factorization used for the post-update gauge: a QR algorithm (no truncation) or a truncated SVD"
+    "gauge algorithm applied after each local update: `NoExpand` for a plain gauge step (a QR
+    algorithm with no truncation, or a truncated SVD), or an algorithm that additionally expands
+    the bond beforehand (e.g. [`DMRG3S`](@ref))"
     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_gauge = NoExpand(), trscheme = nothing,
         alg_svd = Defaults.alg_svd(), alg_orth = Defaults.alg_orth()
     )
     # single-site DMRG defaults to the per-bond adaptive controller (`AdaptiveKrylov`); pass
     # `alg_eigsolve = (; adaptive = false, ...)` to opt out (the splat overrides the default).
     alg_eigsolve′ = alg_eigsolve isa NamedTuple ?
         Defaults.alg_eigsolve(; adaptive = true, 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."
+
+    inner_gauge = alg_gauge.alg_gauge
+    if isnothing(inner_gauge)
+        trscheme = something(trscheme, notrunc()) # enforce trscheme default here
+        inner_gauge = _build_inner_gauge(trscheme, alg_svd, alg_orth)
+        alg_gauge = set_alg_gauge(alg_gauge, inner_gauge)
+    else
+        isnothing(trscheme) || throw(
+            ArgumentError(
+                "`trscheme` was given together with an `alg_gauge` that already carries its own " *
+                    "gauge algorithm (e.g. `DMRG3S(noise, schedule, some_gauge)`); set the truncation " *
+                    "via one or the other, not both."
+            )
+        )
+    end
+
+    if (!isnothing(alg_expand) || _expands(alg_gauge)) && !_truncates(alg_gauge)
+        @warn "DMRG with a bond-expanding `alg_expand` and/or `alg_gauge` 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)
 end
@@ -71,7 +116,7 @@ function local_update!(
         site, direction,
         ψ, O, alg::DMRG, envs,
         ϵ_global, ϵ_trunc, decay_rate,
-        timeroutput
+        iter, timeroutput
     )
     ϵ_local = calc_galerkin(site, ψ, O, ψ, envs)
 
@@ -87,8 +132,10 @@ function local_update!(
         fixedpoint(H_effective, ac_old, :SR, alg_eigsolve)
     end
 
+    alg_gauge = _update_alg_gauge(alg.alg_gauge, iter, ϵ_global)
+
     # 3. gauge
-    ψ, ϵ_trunc = @timeit timeroutput "gauge" gauge!(ψ, site, direction, AC′, alg.alg_gauge; normalize = true)
+    ψ, ϵ_trunc = @timeit timeroutput "gauge" gauge!(site, direction, ψ, O, envs, AC′, alg_gauge; normalize = true)
 
     # 4. bookkeeping: measured contraction factor per matvec, kept a strict contraction in (0, 1)
     decay_rate = clamp((first(info.normres) / ϵ_local)^(1 / max(1, info.numops)), 1.0e-3, 0.999)
@@ -137,15 +184,15 @@ function DMRG2(;
     alg_eigsolve′ = alg_eigsolve isa NamedTuple ?
         Defaults.alg_eigsolve(; adaptive = true, alg_eigsolve...) : alg_eigsolve
     # two-site DMRG always truncates the enlarged bond back down, so the gauge is a truncated SVD
-    alg_gauge = MatrixAlgebraKit.TruncatedAlgorithm(alg_svd, trscheme)
+    alg_gauge = NoExpand(MatrixAlgebraKit.TruncatedAlgorithm(alg_svd, trscheme))
     return DMRG2(tol, maxiter, verbosity, alg_eigsolve′, alg_gauge, finalize)
 end
 
 function local_update!(
         pos, direction,
         ψ, O, alg::DMRG2, envs,
         ϵ_global, ϵ_trunc, decay_rate,
-        timeroutput
+        iter, timeroutput
     )
     Heff = @timeit timeroutput "AC2_hamiltonian" AC2_hamiltonian(pos, ψ, O, ψ, envs)
 
@@ -162,9 +209,11 @@ function local_update!(
         (newA2center, info)
     end
 
+    alg_gauge = _update_alg_gauge(alg.alg_gauge, iter, ϵ_global)
+
     # 2. gauge: truncated SVD split back into single-site tensors and install;
     #           the discarded weight is the truncation error
-    ψ, ϵ_trunc = @timeit timeroutput "gauge" gauge2!(ψ, pos, direction, newA2center, alg.alg_gauge; normalize = true)
+    ψ, ϵ_trunc = @timeit timeroutput "gauge" gauge2!(pos, direction, ψ, O, envs, newA2center, alg_gauge; normalize = true)
 
     # 3. bookkeeping: measured contraction factor per matvec, kept a strict contraction in (0, 1)
     decay_rate = clamp((first(info.normres) / ϵ_local)^(1 / max(1, info.numops)), 1.0e-3, 0.999)
@@ -184,6 +233,8 @@ _num_updates(::DMRG2, ψ) = length(ψ) - 1
 _sweep_ranges(::DMRG, ψ) = (1:(length(ψ) - 1), length(ψ):-1:2)
 _sweep_ranges(::DMRG2, ψ) = (1:(length(ψ) - 1), (length(ψ) - 2):-1:1)
 
+inner_alg_gauge(alg::Union{DMRG, DMRG2}) = alg.alg_gauge.alg_gauge
+
 function find_groundstate!(
         ψ::AbstractFiniteMPS, H, alg::Union{DMRG, DMRG2}, envs = environments(ψ, H, ψ)
     )
@@ -211,7 +262,7 @@ function find_groundstate!(
                         pos, Val(:right),
                         ψ, H, alg, envs,
                         ϵ_global, ϵ_truncs[pos], decay_rates[pos],
-                        timeroutput
+                        iter, timeroutput
                     )
                     ϵ_global = maximum(ϵ_locals)
                 end
@@ -223,7 +274,7 @@ function find_groundstate!(
                         pos, Val(:left),
                         ψ, H, alg, envs,
                         ϵ_global, ϵ_truncs[pos], decay_rates[pos],
-                        timeroutput
+                        iter, timeroutput
                     )
                     ϵ_global = maximum(ϵ_locals)
                 end

src/algorithms/post_expand/dmrg3s.jl

@@ -0,0 +1,165 @@
+"""
+    NoiseSchedule
+
+Supertype for controlling how a bond-expansion algorithm's noise/perturbation amplitude
+evolves across DMRG sweeps. A schedule `s` is called as `s(noise, iter, ϵ) -> noise`, given
+the current noise amplitude, the outer iteration count, and the current global convergence
+error, and returns the amplitude to use for the next iteration. Schedules can be composed
+with `∘`; see [`ExponentialDecay`](@ref), [`Warmup`](@ref), [`FunctionalSchedule`](@ref).
+"""
+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::Int)
+
+Noise schedule that keeps the noise amplitude constant for the first `iters` outer
+iterations, then drops it to exactly zero. Useful composed with a decaying schedule (e.g.
+`ExponentialDecay(0.7) ∘ Warmup(5)`) to hold the perturbation at full strength for a few
+sweeps before starting to taper it off.
+"""
+struct Warmup <: NoiseSchedule
+    iters::Int
+end
+
+(s::Warmup)(noise, iter, ϵ) = iter ≤ s.iters ? noise : zero(noise)
+
+"""
+    DMRG3S(noise, schedule::NoiseSchedule)
+
+Gauge algorithm wrapper that, at every site update, injects a Hamiltonian-derived
+perturbation of the just-optimized tensor before gauging — the "strictly single-site DMRG with
+subspace expansion" scheme of Hubig, McCulloch, Schollwöck & Wolf,
+[Phys. Rev. B 91, 155115 (2015)](https://doi.org/10.1103/PhysRevB.91.155115). This lets
+single-site DMRG introduce basis states/quantum-number sectors absent from the initial
+state, helping it escape local minima that plain single-site DMRG can get stuck in.
+
+`noise` is the initial perturbation amplitude; `schedule` (see [`ExponentialDecay`](@ref),
+[`Warmup`](@ref)) controls how it evolves across outer iterations, and once it decays to
+exactly zero the gauge step reverts to a plain [`NoExpand`](@ref) for the remainder of the
+run. As with [`NoExpand`](@ref), the actual factorization used to gauge the enriched tensor
+is filled in by `DMRG`'s constructor, not supplied here directly — see `DMRG`'s docstring
+for the calling convention:
+
+```julia
+DMRG(; alg_gauge = DMRG3S(0.1, ExponentialDecay(0.7)), trscheme = truncdim(50))
+```
+
+A truncating `trscheme` is strongly recommended alongside `DMRG3S`, to cut the perturbed
+bond back down each sweep — `DMRG`'s constructor warns if none is given.
+"""
+struct DMRG3S{N, S <: NoiseSchedule, A} <: Algorithm
+    noise::N
+    schedule::S
+    alg_gauge::A
+end
+
+DMRG3S(noise, schedule) = DMRG3S(noise, schedule, nothing)
+
+set_alg_gauge(alg::DMRG3S, inner_gauge) = DMRG3S(alg.noise, alg.schedule, inner_gauge)
+
+function _update_alg_gauge(alg::DMRG3S, iter, ϵ)
+    noise = alg.schedule(alg.noise, iter, ϵ)
+    return iszero(noise) ? NoExpand(alg.alg_gauge) : DMRG3S(noise, alg.schedule, alg.alg_gauge)
+end
+
+# similar to `fuser` but the contracted leg is flattened
+function _get_combiner(::Type{TorA}, V1, V2) where {TorA}
+    Vprod = V1 ⊗ V2
+    return isomorphism(TorA, oplus(fuse(Vprod)), Vprod)
+end
+
+function gauge!(pos::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, envs, AC, alg::DMRG3S; normalize = true)
+    El = leftenv(envs, pos, ψ)
+    Hi = H[pos]
+    α = alg.noise
+    T = promote_type(scalartype(ψ), scalartype(Hi))
+    V = right_virtualspace(AC)
+
+    combiner = _get_combiner(T, V, right_virtualspace(Hi))'
+    Vpert = only(domain(combiner))
+
+    mpo_ac = MPO_AC_Hamiltonian(El, Hi, combiner)
+    pert = α * mpo_ac(AC)
+
+    AC_expanded = catdomain(AC, pert)
+
+    AL, C, ϵ = left_gauge(AC_expanded, alg.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 gauge!(pos::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, envs, AC, alg::DMRG3S; normalize = true)
+    Er = rightenv(envs, pos, ψ)
+    Hi = H[pos]
+    α = alg.noise
+    T = promote_type(scalartype(ψ), scalartype(Hi))
+    V = left_virtualspace(AC)
+
+    combiner = _get_combiner(T, V, left_virtualspace(Hi))
+    Vpert = only(codomain(combiner))
+    combiner = _transpose_front(combiner)
+
+    mpo_ac = MPO_AC_Hamiltonian(combiner, Hi, Er)
+    pert = α * mpo_ac(AC)
+    AC = _transpose_tail(AC)
+    pert = _transpose_tail(pert)
+
+    AC_expanded = catcodomain(AC, pert)
+
+    C, AR, ϵ = right_gauge(AC_expanded, alg.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