Use MatrixAlgebraKit's SVD pullbacks (#335)

* Add `FullSVDPullback` and `TruncSVDPullback` wrapping MAK's pullbacks

* Add broadening to `eigh` pullbacks

* Update defaults and adjust rest of code to name changes

* Add links in docstrings

* Fix eigh tests and add eigh broadening tests

* Make `eigh` default consistent with MAK

* Fix SVD derivative unthunks

* Fix gauge fixing index ordering

* Format

* Fix unit cell indices in SVDAdjoint gauge fixing

* Fix IterSVD `truncation_indices` return

* Increase test coverage and improve broadening tests

* Add regression test for CTMRG gradient accuracy

* Add `:trunc` description and revert `:sdd` and `:svd` name changes

* Apply suggestions

---------

Co-authored-by: leburgel <lander.burgelman@gmail.com>

Author: pbrehmer

src/Defaults.jl

@@ -9,8 +9,8 @@ Module containing default algorithm parameter values and arguments.
 * `ctmrg_maxiter=$(Defaults.ctmrg_maxiter)` : Maximal number of CTMRG iterations per run.
 * `ctmrg_miniter=$(Defaults.ctmrg_miniter)` : Minimal number of CTMRG carried out.
 * `ctmrg_alg=:$(Defaults.ctmrg_alg)` : Default CTMRG algorithm variant.
-    - `:simultaneous`: Simultaneous expansion and renormalization of all sides.
-    - `:sequential`: Sequential application of left moves and rotations.
+    - `:simultaneous` : Simultaneous expansion and renormalization of all sides.
+    - `:sequential` : Sequential application of left moves and rotations.
 * `ctmrg_verbosity=$(Defaults.ctmrg_verbosity)` : CTMRG output information verbosity
 
 ## SVD forward & reverse
@@ -22,25 +22,40 @@ Module containing default algorithm parameter values and arguments.
     - `:truncrank` : Additionally supply truncation dimension `η`; truncate such that the 2-norm of the truncated values is smaller than `η`
     - `:truncspace` : Additionally supply truncation space `η`; truncate according to the supplied vector space 
     - `:trunctol` : Additionally supply singular value cutoff `η`; truncate such that every retained singular value is larger than `η`
+* `rrule_degeneracy_atol=$(Defaults.rrule_degeneracy_atol)` : Broadening amplitude which smoothens the divergent term in the retained contributions of an SVD or eigh pullback, in case of (pseudo) degenerate singular values
 * `svd_fwd_alg=:$(Defaults.svd_fwd_alg)` : SVD algorithm that is used in the forward pass.
-    - `:sdd`: MatrixAlgebraKit's `LAPACK_DivideAndConquer`
-    - `:svd`: MatrixAlgebraKit's `LAPACK_QRIteration`
-    - `:iterative`: Iterative SVD only computing the specifed number of singular values and vectors, see [`IterSVD`](@ref PEPSKit.IterSVD)
+    - `:sdd` : MatrixAlgebraKit's `LAPACK_DivideAndConquer`
+    - `:svd` : MatrixAlgebraKit's `LAPACK_QRIteration`
+    - `:iterative` : Iterative SVD only computing the specifed number of singular values and vectors, see [`IterSVD`](@ref PEPSKit.IterSVD)
 * `svd_rrule_tol=$(Defaults.svd_rrule_tol)` : Accuracy of SVD reverse-rule.
 * `svd_rrule_min_krylovdim=$(Defaults.svd_rrule_min_krylovdim)` : Minimal Krylov dimension of the reverse-rule algorithm (if it is a Krylov algorithm).
 * `svd_rrule_verbosity=$(Defaults.svd_rrule_verbosity)` : SVD gradient output verbosity.
 * `svd_rrule_alg=:$(Defaults.svd_rrule_alg)` : Reverse-rule algorithm for the SVD gradient.
-    - `:full`: Uses a modified version of MatrixAlgebraKit's reverse-rule for `svd_compact` which doesn't solve any linear problem and instead requires access to the full SVD, see [`PEPSKit.FullSVDReverseRule`](@ref).
-    - `:gmres`: GMRES iterative linear solver, see the [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.GMRES) for details
-    - `:bicgstab`: BiCGStab iterative linear solver, see the [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.BiCGStab) for details
-    - `:arnoldi`: Arnoldi Krylov algorithm, see the [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.Arnoldi) for details
-* `svd_rrule_broadening=$(Defaults.svd_rrule_broadening)` : Lorentzian broadening amplitude which smoothens the divergent term in the SVD adjoint in case of (pseudo) degenerate singular values
+    - `:full` : Uses a modified version of MatrixAlgebraKit's reverse-rule for `svd_compact` which doesn't solve any linear problem and instead requires access to the full SVD, see [`PEPSKit.FullSVDPullback`](@ref).
+    - `:trunc` : MatrixAlgebraKit's `svd_trunc_pullback!` solving a Sylvester equation on the truncated subspace and therefore only requires access to the truncated SVD.
+    - `:gmres` : GMRES iterative linear solver, see the [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.GMRES) for details
+    - `:bicgstab` : BiCGStab iterative linear solver, see the [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.BiCGStab) for details
+    - `:arnoldi` : Arnoldi Krylov algorithm, see the [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.Arnoldi) for details
+
+## `eigh` forward & reverse
+
+* `eigh_fwd_alg=:$(Defaults.eigh_fwd_alg)` : `eigh` algorithm that is used in the forward pass.
+    - `:qriteration` : MatrixAlgebraKit's `LAPACK_QRIteration`.
+    - `:bisection` : MatrixAlgebraKit's `LAPACK_Bisection`.
+    - `:divideandconquer` : MatrixAlgebraKit's `LAPACK_DivideAndConquer`.
+    - `:multiple` : MatrixAlgebraKit's `LAPACK_MultipleRelativelyRobustRepresentations`. 
+    - `:lanczos` : Lanczos algorithm, see [`KrylovKit.Lanczos`](@extref) for details.
+    - `:blocklanczos` : Block Lanczos algorithm, see [`KrylovKit.BlockLanczos`](@extref) for details.
+* `eigh_rrule_alg=:$(Defaults.eigh_rrule_alg)` : Reverse-rule algorithm for the `eigh` gradient.
+    - `:full` : Full pullback algorithm for eigendecompositions, see [`PEPSKit.FullEighPullback`](@ref).
+    - `:trunc` : Truncated reverse-mode algorithm for eigendecompositions, see [`PEPSKit.TruncEighPullback`](@ref).
+* `eigh_rrule_verbosity=$(Defaults.eigh_rrule_verbosity)` : eigh gradient output verbosity.
 
 ## Projectors
 
 * `projector_alg=:$(Defaults.projector_alg)` : Default variant of the CTMRG projector algorithm.
-    - `:halfinfinite`: Projection via SVDs of half-infinite (two enlarged corners) CTMRG environments.
-    - `:fullinfinite`: Projection via SVDs of full-infinite (all four enlarged corners) CTMRG environments.
+    - `:halfinfinite` : Projection via SVDs of half-infinite (two enlarged corners) CTMRG environments.
+    - `:fullinfinite` : Projection via SVDs of full-infinite (all four enlarged corners) CTMRG environments.
 * `projector_verbosity=$(Defaults.projector_verbosity)` : Projector output information verbosity.
 
 ## Fixed-point gradient
@@ -93,17 +108,17 @@ const sparse = false # TODO: implement sparse CTMRG
 
 # SVD forward & reverse
 const trunc = :fixedspace # ∈ {:fixedspace, :notrunc, :truncerror, :truncspace, :trunctol}
-const svd_fwd_alg = :sdd # ∈ {:sdd, :svd, :iterative}
+const rrule_degeneracy_atol = 1.0e-13
+const svd_fwd_alg = :sdd # ∈ {:sdd, :svd, :bisection, :jacobi, :iterative}
 const svd_rrule_tol = ctmrg_tol
 const svd_rrule_min_krylovdim = 48
 const svd_rrule_verbosity = -1
-const svd_rrule_alg = :full # ∈ {:full, :gmres, :bicgstab, :arnoldi}
-const svd_rrule_broadening = 1.0e-13
+const svd_rrule_alg = :full # ∈ {:full, :trunc, :gmres, :bicgstab, :arnoldi}
 const krylovdim_factor = 1.4
 
 # eigh forward & reverse
 const eigh_fwd_alg = :qriteration # ∈ {:qriteration, :bisection, :divideandconquer, :multiple, :lanczos, :blocklanczos}
-const eigh_rrule_alg = :trunc # ∈ {:trunc, :full}
+const eigh_rrule_alg = :full # ∈ {:full, :trunc}
 const eigh_rrule_verbosity = 0
 
 # QR forward & reverse

src/PEPSKit.jl

@@ -12,6 +12,7 @@ using MatrixAlgebraKit: LAPACK_DivideAndConquer, LAPACK_QRIteration
 using MatrixAlgebraKit:
     TruncationStrategy, NoTruncation, truncate, findtruncated, truncation_error, diagview
 using MatrixAlgebraKit: LAPACK_EighAlgorithm, eigh_pullback!, eigh_trunc_pullback!
+using MatrixAlgebraKit: svd_pullback!, svd_trunc_pullback!
 
 using TensorKit
 using TensorKit: AdjointTensorMap, SectorDict
@@ -123,11 +124,11 @@ include("algorithms/select_algorithm.jl")
 
 using .Defaults: set_scheduler!
 export set_scheduler!
-export SVDAdjoint, FullSVDReverseRule, IterSVD
+export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
 export FixedSpaceTruncation, SiteDependentTruncation
 export HalfInfiniteProjector, FullInfiniteProjector
-export EighAdjoint, IterEigh, QRAdjoint, C4vCTMRG, C4vEighProjector, C4vQRProjector
+export C4vCTMRG, C4vEighProjector, C4vQRProjector
 export initialize_random_c4v_env, initialize_singlet_c4v_env
 export LocalOperator, physicalspace
 export product_peps

src/algorithms/ctmrg/projectors.jl

@@ -56,7 +56,7 @@ end
 
 Return the tensor decomposition algorithm of the `alg` projector algorithm.
 Additionally, the multi-index `(dir, r, c)` can be supplied which will return the
-decomposition performed at that index, e.g. when using `FixedEigh` or `FixedSVD`.
+decomposition performed at that index, e.g. when using [`FixedEig`](@ref) or [`FixedSVD`](@ref).
 """
 decomposition_algorithm(alg::ProjectorAlgorithm) = alg.decomposition_alg
 function decomposition_algorithm(alg::ProjectorAlgorithm, (dir, r, c))
@@ -67,11 +67,12 @@ function decomposition_algorithm(alg::ProjectorAlgorithm, (dir, r, c))
             FixedSVD(
                 fwd_alg.U[dir, r, c], fwd_alg.S[dir, r, c], fwd_alg.V[dir, r, c],
                 fwd_alg.U_full[dir, r, c], fwd_alg.S_full[dir, r, c], fwd_alg.V_full[dir, r, c],
+                fwd_alg.truncation_indices[dir, r, c],
             )
         else
             FixedSVD(
                 fwd_alg.U[dir, r, c], fwd_alg.S[dir, r, c], fwd_alg.V[dir, r, c],
-                nothing, nothing, nothing,
+                nothing, nothing, nothing, nothing,
             )
         end
         return SVDAdjoint(; fwd_alg = fix_svd, rrule_alg = decomposition_alg.rrule_alg)

src/algorithms/ctmrg/simultaneous.jl

@@ -64,7 +64,8 @@ function _split_proj_and_info(proj_and_info)
     U_full = map(x -> x[2].U_full, proj_and_info)
     S_full = map(x -> x[2].S_full, proj_and_info)
     V_full = map(x -> x[2].V_full, proj_and_info)
-    info = (; truncation_error, condition_number, U, S, V, U_full, S_full, V_full)
+    truncation_indices = map(x -> x[2].truncation_indices, proj_and_info)
+    info = (; truncation_error, condition_number, U, S, V, U_full, S_full, V_full, truncation_indices)
     return (P_left, P_right), info
 end
 

src/algorithms/optimization/fixed_point_differentiation.jl

@@ -320,24 +320,25 @@ end
 function gauge_fix(alg::SVDAdjoint, signs, info)
     # embed gauge signs in larger space to fix gauge of full U and V on truncated subspace
     rowsize, colsize = size(signs, 2), size(signs, 3)
-    signs_full = map(Iterators.product(1:4, 1:rowsize, 1:colsize)) do (dir, r, c)
-        σ = signs[dir, r, c]
-        r_sign, c_sign = if dir == NORTH # take unit cell interdependency of signs into account
-            r, _prev(c, colsize)
+    inds = info.truncation_indices
+    signs_full = map(Iterators.product(1:4, 1:rowsize, 1:colsize)) do (dir, row, col)
+        σ = signs[dir, row, col]
+        row_sign, col_sign = if dir == NORTH # take unit cell interdependency of signs into account
+            row, _prev(col, colsize)
         elseif dir == EAST
-            _prev(r, rowsize), c
+            _prev(row, rowsize), col
         elseif dir == SOUTH
-            r, _next(c, colsize)
+            row, _next(col, colsize)
         elseif dir == WEST
-            _next(r, rowsize), c
+            _next(row, rowsize), col
         end
-        extended_space = domain(info.U_full[dir, r_sign, c_sign]) ← codomain(info.V_full[dir, r_sign, c_sign])
-        extended_σ = zeros(scalartype(σ), extended_space)
-        for (c, b) in blocks(extended_σ)
-            σc = block(σ, c)
-            kept_dim = size(σc, 1)
-            b[diagind(b)] .= one(scalartype(σ)) # put ones on the diagonal
-            b[1:kept_dim, 1:kept_dim] .= σc # set to σ on kept subspace
+
+        ind = inds[dir, row_sign, col_sign]
+        extended_σ = id(scalartype(σ), domain(info.S_full[dir, row_sign, col_sign]))
+        for (c, b) in blocks(σ)
+            I = get(ind, c, nothing)
+            @assert !isnothing(I)
+            block(extended_σ, c)[I, I] = b
         end
         return extended_σ
     end
@@ -346,15 +347,15 @@ function gauge_fix(alg::SVDAdjoint, signs, info)
     U_fixed, V_fixed = fix_relative_phases(info.U, info.V, signs)
     U_full_fixed, V_full_fixed = fix_relative_phases(info.U_full, info.V_full, signs_full)
     return SVDAdjoint(;
-        fwd_alg = FixedSVD(U_fixed, info.S, V_fixed, U_full_fixed, info.S_full, V_full_fixed),
+        fwd_alg = FixedSVD(U_fixed, info.S, V_fixed, U_full_fixed, info.S_full, V_full_fixed, inds),
         rrule_alg = alg.rrule_alg,
     )
 end
 function gauge_fix(alg::SVDAdjoint{F}, signs, info) where {F <: IterSVD}
     # fix kept U and V only since iterative SVD doesn't have access to full spectrum
     U_fixed, V_fixed = fix_relative_phases(info.U, info.V, signs)
     return SVDAdjoint(;
-        fwd_alg = FixedSVD(U_fixed, info.S, V_fixed, nothing, nothing, nothing),
+        fwd_alg = FixedSVD(U_fixed, info.S, V_fixed, nothing, nothing, nothing, nothing),
         rrule_alg = alg.rrule_alg,
     )
 end
@@ -374,7 +375,7 @@ function gauge_fix(alg::EighAdjoint, signs, info)
     V_fixed = info.V * σ'
     V_full_fixed = info.V_full * extended_σ'
     return EighAdjoint(;
-        fwd_alg = FixedEig(info.D, V_fixed, info.D_full, V_full_fixed, info.truncation_indices),
+        fwd_alg = FixedEig(info.D, V_fixed, info.D_full, V_full_fixed, inds),
         rrule_alg = alg.rrule_alg,
     )
 end

src/utility/eigh.jl

@@ -16,6 +16,7 @@ Construct a `FullEighPullback` algorithm struct from the following keyword argum
 * `verbosity::Int=0` : Suppresses all output if `≤0`, prints gauge dependency warnings if `1`, and always prints gauge dependency if `≥2`.
 """
 @kwdef struct FullEighPullback
+    degeneracy_atol::Real = Defaults.rrule_degeneracy_atol
     verbosity::Int = 0
 end
 
@@ -37,6 +38,7 @@ Construct a `TruncEighPullback` algorithm struct from the following keyword argu
 * `verbosity::Int=0` : Suppresses all output if `≤0`, prints gauge dependency warnings if `1`, and always prints gauge dependency if `≥2`.
 """
 @kwdef struct TruncEighPullback
+    degeneracy_atol::Real = Defaults.rrule_degeneracy_atol
     verbosity::Int = 0
 end
 
@@ -64,8 +66,8 @@ Construct a `EighAdjoint` algorithm struct based on the following keyword argume
     - `:lanczos` : Lanczos algorithm for symmetric/Hermitian matrices, see [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.Lanczos)
     - `:blocklanczos` : Block version of `:lanczos` for repeated extremal eigenvalues, see [KrylovKit docs](https://jutho.github.io/KrylovKit.jl/stable/man/algorithms/#KrylovKit.BlockLanczos)
 * `rrule_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.eigh_rrule_alg))`: Reverse-rule algorithm for differentiating the eigenvalue decomposition. Can be supplied by an `Algorithm` instance directly or as a `NamedTuple` where `alg` is one of the following:
-    - `:trunc` : MatrixAlgebraKit's `eigh_trunc_pullback` solving a Sylvester equation on the truncated subspace
-    - `:full` : MatrixAlgebraKit's `eigh_pullback` that requires access to the full spectrum
+    - `:full` : MatrixAlgebraKit's `eigh_pullback!` that requires access to the full spectrum
+    - `:trunc` : MatrixAlgebraKit's `eigh_trunc_pullback!` solving a Sylvester equation on the truncated subspace
 """
 struct EighAdjoint{F, R}
     fwd_alg::F
@@ -77,12 +79,8 @@ const EIGH_FWD_SYMBOLS = IdDict{Symbol, Any}(
     :bisection => LAPACK_Bisection,
     :divideandconquer => LAPACK_DivideAndConquer,
     :multiple => LAPACK_MultipleRelativelyRobustRepresentations,
-    :lanczos =>
-        (; tol = 1.0e-14, krylovdim = 30, kwargs...) ->
-    IterEigh(; alg = Lanczos(; tol, krylovdim), kwargs...),
-    :blocklanczos =>
-        (; tol = 1.0e-14, krylovdim = 30, kwargs...) ->
-    IterEigh(; alg = BlockLanczos(; tol, krylovdim), kwargs...),
+    :lanczos => (; tol = 1.0e-14, krylovdim = 30, kwargs...) -> IterEigh(; alg = Lanczos(; tol, krylovdim), kwargs...),
+    :blocklanczos => (; tol = 1.0e-14, krylovdim = 30, kwargs...) -> IterEigh(; alg = BlockLanczos(; tol, krylovdim), kwargs...),
 )
 const EIGH_RRULE_SYMBOLS = IdDict{Symbol, Type{<:Any}}(
     :full => FullEighPullback, :trunc => TruncEighPullback,
@@ -105,6 +103,7 @@ function EighAdjoint(; fwd_alg = (;), rrule_alg = (;))
     rrule_algorithm = if rrule_alg isa NamedTuple
         rrule_kwargs = (;
             alg = Defaults.eigh_rrule_alg,
+            degeneracy_atol = Defaults.rrule_degeneracy_atol,
             verbosity = Defaults.eigh_rrule_verbosity,
             rrule_alg...,
         ) # overwrite with specified kwargs
@@ -113,7 +112,7 @@ function EighAdjoint(; fwd_alg = (;), rrule_alg = (;))
             throw(ArgumentError("unknown rrule algorithm: $(rrule_kwargs.alg)"))
         rrule_type = EIGH_RRULE_SYMBOLS[rrule_kwargs.alg]
         if rrule_type <: Union{FullEighPullback, TruncEighPullback}
-            rrule_kwargs = (; rrule_kwargs.verbosity)
+            rrule_kwargs = (; rrule_kwargs.degeneracy_atol, rrule_kwargs.verbosity)
         end
 
         rrule_type(; rrule_kwargs...)
@@ -249,7 +248,10 @@ function _eigh_trunc!(f, alg::IterEigh, trunc::TruncationStrategy)
     truncation_error =
         trunc isa NoTruncation ? abs(zero(scalartype(f))) : norm(V * D * V' - f)
     condition_number = cond(D)
-    info = (; truncation_error, condition_number, D_full = nothing, V_full = nothing)
+    info = (;
+        truncation_error, condition_number, D_full = nothing, V_full = nothing,
+        truncation_indices = nothing,
+    )
 
     return D, V, info
 end
@@ -329,7 +331,7 @@ function _get_pullback_gauge_tol(verbosity::Int)
     if verbosity ≤ 0 # never print gauge sensitivity
         return (_) -> Inf
     elseif verbosity == 1 # print gauge sensitivity above default atol
-        MatrixAlgebraKit.default_pullback_gaugetol
+        MatrixAlgebraKit.default_pullback_gauge_atol
     else # always print gauge sensitivity
         return (_) -> 0.0
     end
@@ -349,7 +351,7 @@ function ChainRulesCore.rrule(
     function eigh_trunc!_full_pullback(ΔDV)
         Δt = eigh_pullback!(
             zeros(scalartype(t), space(t)), t, (D, V), ΔDV, inds;
-            gauge_atol = gtol(ΔDV)
+            gauge_atol = gtol(ΔDV), degeneracy_atol = alg.rrule_alg.degeneracy_atol,
         )
         return NoTangent(), Δt, NoTangent()
     end
@@ -373,7 +375,7 @@ function ChainRulesCore.rrule(
     function eigh_trunc!_trunc_pullback(ΔDV)
         Δf = eigh_trunc_pullback!(
             zeros(scalartype(t), space(t)), t, (D, V), ΔDV;
-            gauge_atol = gtol(ΔDV)
+            gauge_atol = gtol(ΔDV), degeneracy_atol = alg.rrule_alg.degeneracy_atol,
         )
         return NoTangent(), Δf, NoTangent()
     end