Add some basic support for svd_trunc_no_error (#390)

* Add some basic support for svd_trunc_no_error

* Import svd_trunc_no_error into test

Author: kshyatt

src/utility/svd.jl

@@ -133,6 +133,24 @@ function MatrixAlgebraKit.svd_trunc!(t::AdjointTensorMap, alg::SVDAdjoint)
     return adjoint(vt), adjoint(s), adjoint(u), ϵ
 end
 
+"""
+    svd_trunc_no_error(t, alg::SVDAdjoint)
+    svd_trunc_no_error!(t, alg::SVDAdjoint)
+
+Wrapper around `svd_trunc_no_error(!)` which dispatches on the `SVDAdjoint` algorithm.
+This is needed since a custom adjoint may be defined, depending on the `alg`.
+E.g., for `IterSVD` the adjoint for a truncated SVD from `KrylovKit.svdsolve` is used.
+The `_no_error(!)` versions of `svd_trunc(!)` do not compute the truncation error.
+"""
+MatrixAlgebraKit.svd_trunc_no_error(t, alg::SVDAdjoint) = svd_trunc_no_error!(copy(t), alg)
+function MatrixAlgebraKit.svd_trunc_no_error!(t, alg::SVDAdjoint)
+    return svd_trunc_no_error!(t, alg.fwd_alg)
+end
+function MatrixAlgebraKit.svd_trunc_no_error!(t::AdjointTensorMap, alg::SVDAdjoint)
+    u, s, vt = svd_trunc_no_error!(adjoint(t), alg)
+    return adjoint(vt), adjoint(s), adjoint(u)
+end
+
 #
 ## Forward algorithms
 #
@@ -172,22 +190,24 @@ deterministic_start_vector(t::AbstractMatrix) = ones(scalartype(t), size(t, 1))
 
 # Compute SVD data block-wise using KrylovKit algorithm
 # TODO: redefine _empty_svdtensors, _create_svdtensors
-function MatrixAlgebraKit.svd_trunc!(f, alg::TruncatedAlgorithm{<:IterSVD})
+function MatrixAlgebraKit.svd_trunc_no_error!(f, alg::TruncatedAlgorithm{<:IterSVD})
     fwd_alg = alg.alg
     trunc = alg.trunc
     U, S, V = if isempty(blocksectors(f))
         # early return
-        truncation_error = zero(real(scalartype(f)))
         MatrixAlgebraKit.initialize_output(svd_compact!, f, DefaultAlgorithm()) # specified algorithm doesn't matter here
     else
         SVDdata, dims = _compute_svddata!(f, fwd_alg, trunc)
         _create_svdtensors(f, SVDdata, dims)
     end
+    return U, S, V
+end
 
+function MatrixAlgebraKit.svd_trunc!(f, alg::TruncatedAlgorithm{<:IterSVD})
+    U, S, Vᴴ = svd_trunc_no_error!(f, alg)
     truncation_error =
-        trunc isa NoTruncation ? abs(zero(scalartype(f))) : norm(U * S * V - f)
-
-    return U, S, V, truncation_error
+        (trunc isa NoTruncation || isempty(blocksectors(f))) ? abs(zero(scalartype(f))) : norm(U * S * Vᴴ - f)
+    return U, S, Vᴴ, truncation_error
 end
 
 # Copy from TensorKit v0.14 internal functions
@@ -295,6 +315,36 @@ function ChainRulesCore.rrule(
     return (Ũ, S̃, Ṽ⁺, truncerror), svd_trunc!_full_pullback
 end
 
+# svd_trunc_no_error! rrule wrapping MatrixAlgebraKit's svd_pullback!
+# https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/blob/b76c7bb60014ecfead6925d0df6cb4b8d7c2668a/src/pullbacks/svd.jl#L33
+function ChainRulesCore.rrule(
+        ::typeof(svd_trunc_no_error!),
+        t::AbstractTensorMap,
+        alg::SVDAdjoint{F, R}
+    ) where {F <: TruncatedAlgorithm{<:MatrixAlgebraKit.Algorithm}, R <: FullPullback}
+    # TODO: filter out any decomposition algorithm that doesn't give access to the full spectrum
+
+    # requires access to the full decomposition
+    U, S, V⁺ = svd_compact!(t, alg.fwd_alg.alg)
+    (Ũ, S̃, Ṽ⁺), inds = truncate(svd_trunc!, (U, S, V⁺), alg.fwd_alg.trunc)
+
+    gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity)
+
+    function svd_trunc!_full_pullback(ΔUSV′)
+        ΔUSV = unthunk.(ΔUSV′)
+        Δt = svd_pullback!(
+            zeros(scalartype(t), space(t)), t, (U, S, V⁺), ΔUSV, inds;
+            gauge_atol = gtol(ΔUSV), degeneracy_atol = alg.rrule_alg.degeneracy_atol,
+        )
+        return NoTangent(), Δt, NoTangent()
+    end
+    function svd_trunc!_full_pullback(::Tuple{ZeroTangent, ZeroTangent, ZeroTangent})
+        return NoTangent(), ZeroTangent(), NoTangent()
+    end
+
+    return (Ũ, S̃, Ṽ⁺), svd_trunc!_full_pullback
+end
+
 # svd_trunc! rrule wrapping MatrixAlgebraKit's svd_trunc_pullback! (also works for IterSVD)
 # https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/blob/b76c7bb60014ecfead6925d0df6cb4b8d7c2668a/src/pullbacks/svd.jl#L143
 function ChainRulesCore.rrule(
@@ -320,6 +370,31 @@ function ChainRulesCore.rrule(
     return (U, S, V⁺, ϵ), svd_trunc!_trunc_pullback
 end
 
+# svd_trunc_no_error! rrule wrapping MatrixAlgebraKit's svd_trunc_pullback! (also works for IterSVD)
+# https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/blob/b76c7bb60014ecfead6925d0df6cb4b8d7c2668a/src/pullbacks/svd.jl#L143
+function ChainRulesCore.rrule(
+        ::typeof(svd_trunc_no_error!),
+        t,
+        alg::SVDAdjoint{F, R},
+    ) where {F, R <: TruncPullback}
+    U, S, V⁺ = svd_trunc_no_error(t, alg)
+    gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity)
+
+    function svd_trunc!_trunc_pullback(ΔUSV′)
+        ΔUSV = unthunk.(ΔUSV′)
+        Δf = svd_trunc_pullback!(
+            zeros(scalartype(t), space(t)), t, (U, S, V⁺), ΔUSV;
+            gauge_atol = gtol(ΔUSV), degeneracy_atol = alg.rrule_alg.degeneracy_atol,
+        )
+        return NoTangent(), Δf, NoTangent()
+    end
+    function svd_trunc!_trunc_pullback(::Tuple{ZeroTangent, ZeroTangent, ZeroTangent})
+        return NoTangent(), ZeroTangent(), NoTangent()
+    end
+
+    return (U, S, V⁺), svd_trunc!_trunc_pullback
+end
+
 # KrylovKit rrule compatible with TensorMaps & function handles
 function ChainRulesCore.rrule(
         ::typeof(svd_trunc!),
@@ -388,3 +463,72 @@ function ChainRulesCore.rrule(
 
     return (U, S, V, ϵ), svd_trunc!_itersvd_pullback
 end
+
+# KrylovKit rrule compatible with TensorMaps & function handles
+function ChainRulesCore.rrule(
+        ::typeof(svd_trunc_no_error!),
+        f,
+        alg::SVDAdjoint{F, R}
+    ) where {F, R <: Union{GMRES, BiCGStab, Arnoldi}}
+    U, S, V = svd_trunc_no_error(f, alg)
+
+    # update rrule_alg tolerance to be compatible with smallest singular value
+    rrule_alg = alg.rrule_alg
+    smallest_sval = minimum(((_, b),) -> minimum(diag(b)), blocks(S))
+    proper_tol = clamp(rrule_alg.tol, eps(scalartype(S))^(3 / 4), 1.0e-2 * smallest_sval)
+    rrule_alg = @set rrule_alg.tol = proper_tol
+
+    function svd_trunc!_itersvd_pullback(ΔUSVi)
+        Δf = similar(f)
+        ΔU, ΔS, ΔV, = unthunk.(ΔUSVi)
+
+        for (c, b) in blocks(Δf)
+            Uc, Sc, Vc = block(U, c), block(S, c), block(V, c)
+            ΔUc, ΔSc, ΔVc = block(ΔU, c), block(ΔS, c), block(ΔV, c)
+            Sdc = view(Sc, diagind(Sc))
+            ΔSdc = ΔSc isa AbstractZero ? ΔSc : view(ΔSc, diagind(ΔSc))
+
+            n_vals = length(Sdc)
+            lvecs = Vector{Vector{scalartype(f)}}(eachcol(Uc))
+            rvecs = Vector{Vector{scalartype(f)}}(eachcol(Vc'))
+
+            # Dummy objects only used for warnings
+            minimal_info = KrylovKit.ConvergenceInfo(n_vals, nothing, nothing, -1, -1)  # Only num. converged is used
+            minimal_alg = GKL(; tol = rrule_alg.tol, verbosity = 1)  # Tolerance is used for gauge sensitivity, verbosity is used for warnings
+
+            if ΔUc isa AbstractZero && ΔVc isa AbstractZero  # Handle ZeroTangent singular vectors
+                Δlvecs = fill(ZeroTangent(), n_vals)
+                Δrvecs = fill(ZeroTangent(), n_vals)
+            else
+                Δlvecs = Vector{Vector{scalartype(f)}}(eachcol(ΔUc))
+                Δrvecs = Vector{Vector{scalartype(f)}}(eachcol(ΔVc'))
+            end
+
+            xs, ys = KrylovKitCRCExt.compute_svdsolve_pullback_data(
+                ΔSc isa AbstractZero ? fill(zero(Sc[1]), n_vals) : ΔSdc,
+                Δlvecs,
+                Δrvecs,
+                Sdc,
+                lvecs,
+                rvecs,
+                minimal_info,
+                block(f, c),
+                :LR,
+                minimal_alg,
+                rrule_alg,
+            )
+            copyto!(
+                b,
+                KrylovKitCRCExt.construct∂f_svd(
+                    HasReverseMode(), block(f, c), lvecs, rvecs, xs, ys
+                ),
+            )
+        end
+        return NoTangent(), Δf, NoTangent()
+    end
+    function svd_trunc!_itersvd_pullback(::Tuple{ZeroTangent, ZeroTangent, ZeroTangent})
+        return NoTangent(), ZeroTangent(), NoTangent()
+    end
+
+    return (U, S, V), svd_trunc!_itersvd_pullback
+end

test/utility/svd_wrapper.jl

@@ -6,12 +6,13 @@ using ChainRulesCore, Zygote
 using Accessors
 using PEPSKit
 
-using MatrixAlgebraKit: TruncatedAlgorithm, diagview
+using MatrixAlgebraKit: TruncatedAlgorithm, diagview, svd_trunc_no_error
 
 # Gauge-invariant loss function
-function lossfun(A, alg, R = randn(space(A)), trunc = notrunc())
+function lossfun(svd_trunc_f, A, alg, R = randn(space(A)), trunc = notrunc())
     alg = @set alg.fwd_alg = TruncatedAlgorithm(alg.fwd_alg, trunc)
-    U, S, V, = svd_trunc(A, alg)
+    USV = svd_trunc_f(A, alg)
+    U, S, V = USV[1:3] # avoid looking at ϵ if present
     return real(dot(R, U * V)) + dot(S, S)  # Overlap with random tensor R is gauge-invariant and differentiable, also for m≠n
 end
 
@@ -28,29 +29,29 @@ full_alg = SVDAdjoint(; rrule_alg = (; alg = :FullPullback, degeneracy_atol = 1.
 trunc_alg = SVDAdjoint(; rrule_alg = (; alg = :TruncPullback, degeneracy_atol = 1.0e-13))
 iter_alg = SVDAdjoint(; fwd_alg = (; alg = :GKL))
 
-@testset "Non-truncated SVD" begin
-    l_full, g_full = withgradient(A -> lossfun(A, full_alg, R), r)
-    l_trunc, g_trunc = withgradient(A -> lossfun(A, trunc_alg, R), r)
-    l_iter, g_iter = withgradient(A -> lossfun(A, iter_alg, R), r)
+@testset "Non-truncated SVD $f" for f in (svd_trunc, svd_trunc_no_error)
+    l_full, g_full = withgradient(A -> lossfun(f, A, full_alg, R), r)
+    l_trunc, g_trunc = withgradient(A -> lossfun(f, A, trunc_alg, R), r)
+    l_iter, g_iter = withgradient(A -> lossfun(f, A, iter_alg, R), r)
 
     @test l_full ≈ l_trunc ≈ l_iter
     @test g_full[1] ≈ g_trunc[1] rtol = rtol
     @test g_full[1] ≈ g_iter[1] rtol = rtol
     @test g_trunc[1] ≈ g_iter[1] rtol = rtol
 end
 
-@testset "Truncated SVD with χ=$χ" begin
-    l_full, g_full = withgradient(A -> lossfun(A, full_alg, R, trunc), r)
-    l_trunc, g_trunc = withgradient(A -> lossfun(A, trunc_alg, R, trunc), r)
-    l_iter, g_iter = withgradient(A -> lossfun(A, iter_alg, R, trunc), r)
+@testset "Truncated SVD $f with χ=$χ" for f in (svd_trunc, svd_trunc_no_error)
+    l_full, g_full = withgradient(A -> lossfun(f, A, full_alg, R, trunc), r)
+    l_trunc, g_trunc = withgradient(A -> lossfun(f, A, trunc_alg, R, trunc), r)
+    l_iter, g_iter = withgradient(A -> lossfun(f, A, iter_alg, R, trunc), r)
 
     @test l_full ≈ l_trunc ≈ l_iter
     @test g_full[1] ≈ g_trunc[1] rtol = rtol
     @test g_full[1] ≈ g_iter[1] rtol = rtol
     @test g_trunc[1] ≈ g_iter[1] rtol = rtol
 end
 
-@testset "Truncated SVD broadening for $(alg.rrule_alg)" for alg in [full_alg, trunc_alg]
+@testset "Truncated SVD broadening for $f, $(alg.rrule_alg)" for f in (svd_trunc, svd_trunc_no_error), alg in [full_alg, trunc_alg]
     u, s, v, = svd_compact(r)
     s.data[1:2:m] .= s.data[2:2:m] # make every singular value two-fold degenerate
     r_degen = u * s * v
@@ -59,13 +60,13 @@ end
     small_broadening_alg = @set full_alg.rrule_alg.degeneracy_atol = 1.0e-13
 
     l_only_cutoff, g_only_cutoff = withgradient(
-        A -> lossfun(A, full_alg, R, trunc), r_degen
+        A -> lossfun(f, A, full_alg, R, trunc), r_degen
     ) # cutoff sets degenerate difference to zero
     l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = withgradient( # degenerate singular value differences lead to divergent contributions
-        A -> lossfun(A, no_broadening_no_cutoff_alg, R, trunc), r_degen,
+        A -> lossfun(f, A, no_broadening_no_cutoff_alg, R, trunc), r_degen,
     )
     l_small_broadening, g_small_broadening = withgradient( # broadening smoothens divergent contributions
-        A -> lossfun(A, small_broadening_alg, R, trunc), r_degen,
+        A -> lossfun(f, A, small_broadening_alg, R, trunc), r_degen,
     )
 
     @test l_only_cutoff ≈ l_no_broadening_no_cutoff ≈ l_small_broadening
@@ -79,23 +80,23 @@ symm_trspace = truncspace(Z2Space(0 => symm_m ÷ 2, 1 => symm_n ÷ 3))
 symm_r = randn(dtype, symm_space, symm_space)
 symm_R = randn(dtype, space(symm_r))
 
-@testset "IterSVD of symmetric tensors" begin
-    l_full, g_full = withgradient(A -> lossfun(A, full_alg, symm_R), symm_r)
-    l_trunc, g_trunc = withgradient(A -> lossfun(A, trunc_alg, symm_R), symm_r)
-    l_iter, g_iter = withgradient(A -> lossfun(A, iter_alg, symm_R), symm_r)
+@testset "IterSVD of symmetric tensors $f" for f in (svd_trunc, svd_trunc_no_error)
+    l_full, g_full = withgradient(A -> lossfun(f, A, full_alg, symm_R), symm_r)
+    l_trunc, g_trunc = withgradient(A -> lossfun(f, A, trunc_alg, symm_R), symm_r)
+    l_iter, g_iter = withgradient(A -> lossfun(f, A, iter_alg, symm_R), symm_r)
     @test l_full ≈ l_trunc ≈ l_iter
     @test g_full[1] ≈ g_trunc[1] rtol = rtol
     @test g_full[1] ≈ g_iter[1] rtol = rtol
     @test g_trunc[1] ≈ g_iter[1] rtol = rtol
 
     l_full_tr, g_full_tr = withgradient(
-        A -> lossfun(A, full_alg, symm_R, symm_trspace), symm_r
+        A -> lossfun(f, A, full_alg, symm_R, symm_trspace), symm_r
     )
     l_trunc_tr, g_trunc_tr = withgradient(
-        A -> lossfun(A, trunc_alg, symm_R, symm_trspace), symm_r
+        A -> lossfun(f, A, trunc_alg, symm_R, symm_trspace), symm_r
     )
     l_iter_tr, g_iter_tr = withgradient(
-        A -> lossfun(A, iter_alg, symm_R, symm_trspace), symm_r
+        A -> lossfun(f, A, iter_alg, symm_R, symm_trspace), symm_r
     )
     @test l_full_tr ≈ l_trunc_tr ≈ l_iter_tr
     @test g_full_tr[1] ≈ g_trunc_tr[1] rtol = rtol
@@ -104,14 +105,14 @@ symm_R = randn(dtype, space(symm_r))
 
     iter_alg_fallback = @set iter_alg.fwd_alg.fallback_threshold = 0.4  # Do dense decomposition in one block, sparse one in the other
     l_iter_fb, g_iter_fb = withgradient(
-        A -> lossfun(A, iter_alg_fallback, symm_R, symm_trspace), symm_r
+        A -> lossfun(f, A, iter_alg_fallback, symm_R, symm_trspace), symm_r
     )
     @test l_iter_fb ≈ l_trunc_tr ≈ l_full_tr
     @test g_full_tr[1] ≈ g_iter_fb[1] rtol = rtol
     @test g_trunc_tr[1] ≈ g_iter_fb[1] rtol = rtol
 end
 
-@testset "Truncated symmetric SVD broadening for $(alg.rrule_alg)" for alg in [full_alg, trunc_alg]
+@testset "Truncated symmetric SVD broadening for $f, $(alg.rrule_alg)" for f in (svd_trunc, svd_trunc_no_error), alg in [full_alg, trunc_alg]
     u, s, v, = svd_compact(symm_r)
     # make every singular value in the 0-sector three-fold degenerate
     b0 = diagview(block(s, Z2Irrep(0)))
@@ -126,14 +127,14 @@ end
     small_broadening_alg = @set alg.rrule_alg.degeneracy_atol = 1.0e-13
 
     l_only_cutoff, g_only_cutoff = withgradient(
-        A -> lossfun(A, alg, symm_R, symm_trspace), symm_r_degen
+        A -> lossfun(f, A, alg, symm_R, symm_trspace), symm_r_degen
     ) # cutoff sets degenerate difference to zero
     l_no_broadening_no_cutoff, g_no_broadening_no_cutoff = withgradient( # degenerate singular value differences lead to divergent contributions
-        A -> lossfun(A, no_broadening_no_cutoff_alg, symm_R, symm_trspace),
+        A -> lossfun(f, A, no_broadening_no_cutoff_alg, symm_R, symm_trspace),
         symm_r_degen,
     )
     l_small_broadening, g_small_broadening = withgradient( # broadening smoothens divergent contributions
-        A -> lossfun(A, small_broadening_alg, symm_R, symm_trspace),
+        A -> lossfun(f, A, small_broadening_alg, symm_R, symm_trspace),
         symm_r_degen,
     )