From 079969ee5de1e073693591be788578c0ee91546a Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Mon, 13 Jul 2026 18:14:30 +0200
Subject: [PATCH 1/9] refactor: unify left/right DMRG sweeps into a single loop

Replace the two near-duplicate sweep loops in `find_groundstate!` with a
single loop over `direction in (:right, :left)`, dispatching on `Val`
via `_sweep_range` and `dmrg_gauge!`.

No behavioral change: `_sweep_range` reproduces the original site
ranges, and `dmrg_gauge!` preserves the original left_gauge!/right_gauge!
pairing per direction.
---
 src/algorithms/groundstate/dmrg.jl | 81 ++++++++++++++++--------------
 1 file changed, 44 insertions(+), 37 deletions(-)

diff --git a/src/algorithms/groundstate/dmrg.jl b/src/algorithms/groundstate/dmrg.jl
index 61f3097a0..5c3829f66 100644
--- a/src/algorithms/groundstate/dmrg.jl
+++ b/src/algorithms/groundstate/dmrg.jl
@@ -64,6 +64,32 @@ function DMRG(;
     return DMRG(tol, maxiter, verbosity, alg_eigsolve′, finalize, alg_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)
+
+"""
+    dmrg_gauge!(ψ, pos, AC, alg, direction::Val; normalize::Bool = false)
+
+Gauge `ψ` at `pos` using updated tensor `AC`, choosing `left_gauge!`/`right_gauge!`
+based on sweep `direction` (`:right`/`:left`).
+"""
+dmrg_gauge!(::AbstractFiniteMPS, ::Int, AC, alg, ::Val; normalize::Bool)
+
+dmrg_gauge!(ψ::AbstractFiniteMPS, pos::Int, AC, alg, ::Val{:right}; normalize::Bool = false) =
+       left_gauge!(ψ, pos, AC, alg; normalize)
+
+dmrg_gauge!(ψ::AbstractFiniteMPS, pos::Int, AC, alg, ::Val{:left}; normalize::Bool = false) =
+    right_gauge!(ψ, pos, AC, alg; normalize)
+
 function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environments(ψ, H, ψ))
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
     ϵ = maximum(ϵs)
@@ -78,44 +104,25 @@ function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environme
 
             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
+                        
+                        # 3. gauge (QR center-move or truncated SVD, selected by `alg_gauge`)
+                        @timeit timeroutput "gauge" dmrg_gauge!(ψ, pos, AC′, alg.alg_gauge, Val(direction); normalize = true)
                     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)

From f0f29426243c2e05e3a322407ab0991f37b44da8 Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 16:33:43 +0200
Subject: [PATCH 2/9] feat: add post expansion algorithms and integrate into
 DMRG

---
 src/MPSKit.jl                             |  5 ++
 src/algorithms/groundstate/dmrg.jl        | 33 ++++-------
 src/algorithms/post_expand/dmrg3s.jl      | 72 +++++++++++++++++++++++
 src/algorithms/post_expand/post_expand.jl | 11 ++++
 4 files changed, 101 insertions(+), 20 deletions(-)
 create mode 100644 src/algorithms/post_expand/dmrg3s.jl
 create mode 100644 src/algorithms/post_expand/post_expand.jl

diff --git a/src/MPSKit.jl b/src/MPSKit.jl
index b7387f3da..7e88dd690 100644
--- a/src/MPSKit.jl
+++ b/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, 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")
diff --git a/src/algorithms/groundstate/dmrg.jl b/src/algorithms/groundstate/dmrg.jl
index 5c3829f66..6c14ecae4 100644
--- a/src/algorithms/groundstate/dmrg.jl
+++ b/src/algorithms/groundstate/dmrg.jl
@@ -25,7 +25,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,13 +44,16 @@ 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 subspace after the optimization step"
+    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...) :
@@ -58,10 +61,10 @@ function DMRG(;
     # 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
 
 
@@ -76,19 +79,6 @@ _sweep_range(::DMRG, chain_length::Int, ::Val{:right}) = 1:(chain_length-1)
 
 _sweep_range(::DMRG, chain_length::Int, ::Val{:left}) = reverse(2:chain_length)
 
-"""
-    dmrg_gauge!(ψ, pos, AC, alg, direction::Val; normalize::Bool = false)
-
-Gauge `ψ` at `pos` using updated tensor `AC`, choosing `left_gauge!`/`right_gauge!`
-based on sweep `direction` (`:right`/`:left`).
-"""
-dmrg_gauge!(::AbstractFiniteMPS, ::Int, AC, alg, ::Val; normalize::Bool)
-
-dmrg_gauge!(ψ::AbstractFiniteMPS, pos::Int, AC, alg, ::Val{:right}; normalize::Bool = false) =
-       left_gauge!(ψ, pos, AC, alg; normalize)
-
-dmrg_gauge!(ψ::AbstractFiniteMPS, pos::Int, AC, alg, ::Val{:left}; normalize::Bool = false) =
-    right_gauge!(ψ, pos, AC, alg; normalize)
 
 function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environments(ψ, H, ψ))
     ϵs = map(pos -> calc_galerkin(pos, ψ, H, ψ, envs), 1:length(ψ))
@@ -101,6 +91,7 @@ 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
@@ -120,8 +111,10 @@ function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environme
                             _, AC′ = fixedpoint(Hac, ψ.AC[pos], :SR, alg_eigsolve)
                         end
                         
-                        # 3. gauge (QR center-move or truncated SVD, selected by `alg_gauge`)
-                        @timeit timeroutput "gauge" dmrg_gauge!(ψ, pos, AC′, alg.alg_gauge, Val(direction); normalize = true)
+                        # 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
                 end
             end
diff --git a/src/algorithms/post_expand/dmrg3s.jl b/src/algorithms/post_expand/dmrg3s.jl
new file mode 100644
index 000000000..9e613c814
--- /dev/null
+++ b/src/algorithms/post_expand/dmrg3s.jl
@@ -0,0 +1,72 @@
+abstract type NoiseSchedule end
+
+struct ExponentialDecay <: NoiseSchedule
+    decay_rate::Float64
+end
+
+struct Warmup <: NoiseSchedule
+    iters::Int
+end
+
+(s::ExponentialDecay)(noise, iter, ϵ) = noise * s.decay_rate^iter
+(s::Warmup)(noise, iter, ϵ) = iter ≤ s.iters ? noise : zero(noise)
+
+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
\ No newline at end of file
diff --git a/src/algorithms/post_expand/post_expand.jl b/src/algorithms/post_expand/post_expand.jl
new file mode 100644
index 000000000..32fc6a34c
--- /dev/null
+++ b/src/algorithms/post_expand/post_expand.jl
@@ -0,0 +1,11 @@
+struct NoExpand <: Algorithm end
+
+_update_post_expand(::NoExpand, iter, ϵ) = NoExpand()
+
+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

From a3e455d36ed7cfabde0db5fa8e9314f6504e086d Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 16:34:34 +0200
Subject: [PATCH 3/9] feat: add DMRG3S test cases and implement
 bad_initial_state function

---
 test/algorithms/groundstate.jl | 52 ++++++++++++++++++++++++++++++++++
 test/setup/testsetup.jl        | 31 ++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/test/algorithms/groundstate.jl b/test/algorithms/groundstate.jl
index 605d45814..c42b7bb62 100644
--- a/test/algorithms/groundstate.jl
+++ b/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 = 20
+        H = heisenberg_XXX(ComplexF64, U1Irrep; spin = 1 // 2, L)
+
+        Random.seed!(1234)
+        ψ_bad = bad_initial_state(H, L)
+
+        ψ_stuck, envs_stuck, δ_stuck = find_groundstate(
+            ψ_bad, H, DMRG(; verbosity = verbosity_conv, maxiter = 30)
+        )
+        E_stuck = real(expectation_value(ψ_stuck, H, envs_stuck))
+
+        alg_post_expand = DMRG3S(0.1, ExponentialDecay(0.8))
+        ψ_escape, envs_escape, δ_escape = find_groundstate(
+            ψ_bad, H, DMRG(;
+                verbosity = verbosity_conv, maxiter = 30,
+                alg_post_expand, trscheme = truncrank(20),
+            )
+        )
+        E_escape = real(expectation_value(ψ_escape, H, envs_escape))
+
+        # paper reports E(α=0) = -6.35479, E(α≠0) = -8.6824724 for this L = 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)
diff --git a/test/setup/testsetup.jl b/test/setup/testsetup.jl
index 053d272f3..7e921395d 100644
--- a/test/setup/testsetup.jl
+++ b/test/setup/testsetup.jl
@@ -12,6 +12,7 @@ using LinearAlgebra: Diagonal
 using Combinatorics: permutations
 using TensorKitTensors.SpinOperators: S_x, S_y, S_z, S_x_S_x, S_y_S_y, S_z_S_z, S_exchange, S_plus_S_min, S_min_S_plus
 using TensorKitTensors.FermionOperators: f_plus_f_min, f_min_f_plus, f_plus_f_plus, f_min_f_min, f_num
+using Random: shuffle
 
 # exports
 export S_x, S_y, S_z
@@ -22,6 +23,7 @@ export symm_mul_mpo
 export transverse_field_ising, heisenberg_XXX, bilinear_biquadratic_model, XY_model,
     kitaev_model
 export classical_ising_tensors, classical_ising, sixvertex
+export bad_initial_state
 
 # using TensorOperations
 
@@ -242,4 +244,33 @@ function sixvertex(; a = 1.0, b = 1.0, c = 1.0)
     return InfiniteMPO([permute(TensorMap(d, ℂ^2 ⊗ ℂ^2, ℂ^2 ⊗ ℂ^2), ((1, 2), (4, 3)))])
 end
 
+# Functions for testing DMRG3S
+function distinct_bitstrings(n::Int, n_up::Int, n_states::Int)
+    trivial_state = vcat(ones(Int, n_up), zeros(Int, n - n_up))
+    results = Set{Vector{Int}}()
+    while length(results) < n_states
+        push!(results, shuffle(trivial_state))
+    end
+    return collect(results)
+end
+
+function bad_initial_state(H, L; T = ComplexF64, n_states = 20, n_fixed = 3)
+    physd = U1Space(-1 // 2 => 1, 1 // 2 => 1)
+    n_free = L - n_fixed
+    n_up = (n_free - n_fixed) ÷ 2  # 7 of 17 up, so the fixed +3/2 tail nets to Sz_total = 0
+
+    configs = distinct_bitstrings(n_free, n_up, n_states)
+
+    return sum(configs) do bits
+        charges = [b == 1 ? 1 // 2 : -1 // 2 for b in bits]
+        append!(charges, fill(1 // 2, n_fixed))  # the 3 fixed, all-up sites
+
+        q = cumsum(charges)
+        @assert q[end] == 0 "configuration doesn't land in the Sz_total = 0 sector"
+        Vs = [U1Space(qi => 1) for qi in q[1:(end - 1)]]  # L-1 internal bonds only;
+
+        return normalize!(FiniteMPS(T, fill(physd, L), Vs))
+    end
+end
+
 end

From c84307686a01825acf384a90b4f57143d2ef819f Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 17:08:12 +0200
Subject: [PATCH 4/9] Add docstrings

---
 src/algorithms/post_expand/dmrg3s.jl      | 37 +++++++++++++++++++++++
 src/algorithms/post_expand/post_expand.jl | 18 +++++++++++
 2 files changed, 55 insertions(+)

diff --git a/src/algorithms/post_expand/dmrg3s.jl b/src/algorithms/post_expand/dmrg3s.jl
index 9e613c814..a4272eb8a 100644
--- a/src/algorithms/post_expand/dmrg3s.jl
+++ b/src/algorithms/post_expand/dmrg3s.jl
@@ -1,9 +1,31 @@
+"""
+    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
 
+"""
+    ExponentialDecay(decay_rate)
+
+Noise schedule that shrinks geometrically: `noise -> noise * decay_rate^iter`.
+Use `decay_rate < 1` for a standalone [`DMRG3S`](@ref) run that gradually turns
+enrichment off as the state converges.
+"""
 struct ExponentialDecay <: NoiseSchedule
     decay_rate::Float64
 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
@@ -11,6 +33,21 @@ end
 (s::ExponentialDecay)(noise, iter, ϵ) = noise * s.decay_rate^iter
 (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
diff --git a/src/algorithms/post_expand/post_expand.jl b/src/algorithms/post_expand/post_expand.jl
index 32fc6a34c..ed8076914 100644
--- a/src/algorithms/post_expand/post_expand.jl
+++ b/src/algorithms/post_expand/post_expand.jl
@@ -1,7 +1,25 @@
+"""
+    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

From c5de44515f13cd4acfe351ac2ad2ca353dc10b76 Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 17:08:34 +0200
Subject: [PATCH 5/9] modify dmrg docstring to include alg_post_expand

---
 src/algorithms/groundstate/dmrg.jl | 20 +++++++++++++-------
 1 file changed, 13 insertions(+), 7 deletions(-)

diff --git a/src/algorithms/groundstate/dmrg.jl b/src/algorithms/groundstate/dmrg.jl
index 6c14ecae4..59cb49b14 100644
--- a/src/algorithms/groundstate/dmrg.jl
+++ b/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
@@ -44,7 +49,8 @@ struct DMRG{A, F, E, PE, 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 subspace after the optimization step"
+    "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"

From 7264145bb7aae4ce7f51a6ae666bcc74e62f03ff Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 17:31:20 +0200
Subject: [PATCH 6/9] feat: add NoiseSchedule composition and ExponentialDecay
 threshold
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Add `FunctionalSchedule`, wrapping an arbitrary `(noise, iter, ϵ) -> noise`
callable as a `NoiseSchedule`. Overload `∘` on `NoiseSchedule` to compose
two schedules into a `FunctionalSchedule`, following the usual function
composition convention (`s2` applied first, `s1` applied to its result).

Add a `threshold` keyword to `ExponentialDecay`, snapping the decayed noise
to exactly zero once it falls below it -- avoids running the (cheap but
non-free) expansion step indefinitely on a vanishingly small amplitude.
This also makes `ExponentialDecay ∘ Warmup` a natural way to combine a
hard iteration cutoff with smooth decay.
---
 src/MPSKit.jl                        |  2 +-
 src/algorithms/post_expand/dmrg3s.jl | 45 +++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/src/MPSKit.jl b/src/MPSKit.jl
index 7e88dd690..52727b455 100644
--- a/src/MPSKit.jl
+++ b/src/MPSKit.jl
@@ -39,7 +39,7 @@ export TDVP, TDVP2, WI, WII, TaylorCluster
 export changebonds, changebonds!
 export VUMPSSvdCut, OptimalExpand, SvdCut, RandExpand, SketchedExpand
 export post_expand!, NoExpand
-export NoiseSchedule, ExponentialDecay, Warmup, DMRG3S
+export NoiseSchedule, FunctionalSchedule, ExponentialDecay, Warmup, DMRG3S
 export propagator
 export DynamicalDMRG, NaiveInvert, Jeckelmann
 export exact_diagonalization, fidelity_susceptibility
diff --git a/src/algorithms/post_expand/dmrg3s.jl b/src/algorithms/post_expand/dmrg3s.jl
index a4272eb8a..b2cedea21 100644
--- a/src/algorithms/post_expand/dmrg3s.jl
+++ b/src/algorithms/post_expand/dmrg3s.jl
@@ -8,14 +8,46 @@ returning the noise amplitude to use for the next iteration.
 abstract type NoiseSchedule end
 
 """
-    ExponentialDecay(decay_rate)
+    FunctionalSchedule(f) <: NoiseSchedule
 
-Noise schedule that shrinks geometrically: `noise -> noise * decay_rate^iter`.
-Use `decay_rate < 1` for a standalone [`DMRG3S`](@ref) run that gradually turns
-enrichment off as the state converges.
+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 ExponentialDecay <: NoiseSchedule
-    decay_rate::Float64
+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
 
 """
@@ -30,7 +62,6 @@ struct Warmup <: NoiseSchedule
     iters::Int
 end
 
-(s::ExponentialDecay)(noise, iter, ϵ) = noise * s.decay_rate^iter
 (s::Warmup)(noise, iter, ϵ) = iter ≤ s.iters ? noise : zero(noise)
 
 """

From 12f25b3628ff0a38d135e767e9c6af5457ae46ba Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 18:29:57 +0200
Subject: [PATCH 7/9] style: apply Runic formatting

---
 src/algorithms/groundstate/dmrg.jl   |  8 ++++----
 src/algorithms/post_expand/dmrg3s.jl | 20 ++++++++++----------
 2 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/src/algorithms/groundstate/dmrg.jl b/src/algorithms/groundstate/dmrg.jl
index 59cb49b14..197c6de09 100644
--- a/src/algorithms/groundstate/dmrg.jl
+++ b/src/algorithms/groundstate/dmrg.jl
@@ -81,7 +81,7 @@ 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{:right}) = 1:(chain_length - 1)
 
 _sweep_range(::DMRG, chain_length::Int, ::Val{:left}) = reverse(2:chain_length)
 
@@ -116,10 +116,10 @@ function find_groundstate!(ψ::AbstractFiniteMPS, H, alg::DMRG, envs = environme
                             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)
+                        @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
                 end
diff --git a/src/algorithms/post_expand/dmrg3s.jl b/src/algorithms/post_expand/dmrg3s.jl
index b2cedea21..a8320e97b 100644
--- a/src/algorithms/post_expand/dmrg3s.jl
+++ b/src/algorithms/post_expand/dmrg3s.jl
@@ -38,7 +38,7 @@ to exactly zero once it falls below `threshold`. Use `decay_rate < 1` for a stan
 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
+struct ExponentialDecay{T <: Real} <: NoiseSchedule
     decay_rate::T
     threshold::T
 end
@@ -94,20 +94,20 @@ function _get_combiner(::Type{T}, V1, V2) where {T}
     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)
+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]    
+
+    @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])
+    B = _transpose_tail(ψ.AR[pos + 1])
     AR = _transpose_front(catcodomain(B, zeros(T, Vpert ← domain(B))))
 
     normalize && normalize!(C)
@@ -116,25 +116,25 @@ function post_expand!(pos::Int, ::Val{:right}, ψ::AbstractFiniteMPS, H, alg::DM
     return ψ
 end
 
-function post_expand!(pos::Int, ::Val{:left}, ψ::AbstractFiniteMPS, H, alg::DMRG3S, envs, AC, alg_gauge; normalize=true)
+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]
+    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
\ No newline at end of file
+end

From ac34bf288a255827a31b04ffe12fa68988077fd8 Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 23:10:49 +0200
Subject: [PATCH 8/9] fix test: rename variables to don't overwrite initial
 ones

---
 test/algorithms/groundstate.jl | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/test/algorithms/groundstate.jl b/test/algorithms/groundstate.jl
index c42b7bb62..a496e57ba 100644
--- a/test/algorithms/groundstate.jl
+++ b/test/algorithms/groundstate.jl
@@ -142,27 +142,27 @@ verbosity_conv = 1
     end
 
     @testset "DMRG3S escapes local minimum (Hubig et al. 2015, Sec. VII A)" begin
-        L = 20
-        H = heisenberg_XXX(ComplexF64, U1Irrep; spin = 1 // 2, L)
+        L_heis = 20
+        H_heis = heisenberg_XXX(ComplexF64, U1Irrep; spin = 1 // 2, L=L_heis)
 
         Random.seed!(1234)
-        ψ_bad = bad_initial_state(H, L)
+        ψ_bad = bad_initial_state(H_heis, L_heis)
 
         ψ_stuck, envs_stuck, δ_stuck = find_groundstate(
-            ψ_bad, H, DMRG(; verbosity = verbosity_conv, maxiter = 30)
+            ψ_bad, H_heis, DMRG(; verbosity = verbosity_conv, maxiter = 30)
         )
-        E_stuck = real(expectation_value(ψ_stuck, H, envs_stuck))
+        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, DMRG(;
+            ψ_bad, H_heis, DMRG(;
                 verbosity = verbosity_conv, maxiter = 30,
                 alg_post_expand, trscheme = truncrank(20),
             )
         )
-        E_escape = real(expectation_value(ψ_escape, H, envs_escape))
+        E_escape = real(expectation_value(ψ_escape, H_heis, envs_escape))
 
-        # paper reports E(α=0) = -6.35479, E(α≠0) = -8.6824724 for this L = 20, S = 1/2 AFM setup
+        # 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

From 7cb333723c147f11af60db25502a575a3f7a6b2f Mon Sep 17 00:00:00 2001
From: VinceNeede <v.bisogno3@studenti.unipi.it>
Date: Tue, 14 Jul 2026 23:21:11 +0200
Subject: [PATCH 9/9] style: runic format

---
 test/algorithms/groundstate.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/algorithms/groundstate.jl b/test/algorithms/groundstate.jl
index a496e57ba..6279551b0 100644
--- a/test/algorithms/groundstate.jl
+++ b/test/algorithms/groundstate.jl
@@ -143,7 +143,7 @@ verbosity_conv = 1
 
     @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)
+        H_heis = heisenberg_XXX(ComplexF64, U1Irrep; spin = 1 // 2, L = L_heis)
 
         Random.seed!(1234)
         ψ_bad = bad_initial_state(H_heis, L_heis)