Implement norm-preserving retractions

src/algorithms/optimization/peps_optimization.jl

@@ -218,9 +218,17 @@
     gradnorms_unitcell = Vector{Matrix{T}}()
     times = Vector{Float64}()
 
+    # normalize the initial guess
+    peps₀ = peps_normalize(peps₀)
+
     # optimize operator cost function
     (peps_final, env_final), cost, ∂cost, numfg, convergence_history = optimize(
-        (peps₀, env₀), alg.optimizer_alg; retract, inner=real_inner, finalize!
+        (peps₀, env₀),
+        alg.optimizer_alg;
+        retract,
+        inner=real_inner,
+        finalize!,
+        (transport!)=(peps_transport!),
     ) do (peps, env)
         start_time = time_ns()
         E, gs = withgradient(peps) do ψ
@@ -257,13 +265,85 @@
     return peps_final, env_final, cost, info
 end
 
-# Update PEPS unit cell in non-mutating way
-# Note: Both x and η are InfinitePEPS during optimization
+"""
+    peps_normalize(A::InfinitePEPS)
+
+Normalize the individual tensors in the unit cell of an `InfinitePEPS` such that they each
+have unit Euclidean norm.
+"""
+function peps_normalize(A::InfinitePEPS)
+    normalized_tensors = normalize.(unitcell(A))
+    return InfinitePEPS(normalized_tensors)
+end
+
+"""
+    peps_retract(x, η, α)
+
+Performs a norm-preserving retraction of an infinite PEPS `A = x[1]` along `η` with step
+size `α`, giving a new PEPS `A´`,
+```math
+A' \\leftarrow \\cos \\left( α \\frac{||η||}{||A||} \\right) A + \\sin \\left( α \\frac{||η||}{||A||} \\right) ||A|| \\frac{η}{||η||},
+```
+and corresponding directional derivative `ξ`,
+```math
+ξ = \\cos \\left( α \\frac{||η||}{||A||} \\right) η - \\sin \\left( α \\frac{||η||}{||A||} \\right) ||η|| \\frac{A}{||A||},
+```
+such that ``\\langle A', ξ \\rangle = 0`` and ``||A'|| = ||A||``.
+"""
 function peps_retract(x, η, α)
-    peps = deepcopy(x[1])
-    peps.A .+= η.A .* α
+    peps = x[1]
+    norms_peps = norm.(peps.A)
+    norms_η = norm.(η.A)
+
+    peps´ = similar(x[1])
+    peps´.A .=
+        cos.(α .* norms_η ./ norms_peps) .* peps.A .+
+        sin.(α .* norms_η ./ norms_peps) .* norms_peps .* η.A ./ norms_η
+
     env = deepcopy(x[2])
-    return (peps, env), η
+
+    ξ = similar(η)
+    ξ.A .=
+        cos.(α .* norms_η ./ norms_peps) .* η.A .-
+        sin.(α .* norms_η ./ norms_peps) .* norms_η .* peps.A ./ norms_peps
+
+    return (peps´, env), ξ
+end
+
+"""
+    peps_transport!(ξ, x, η, α, x′)
+
+Transports a direction at `A = x[1]` to a valid direction at `A´ = x´[1]` corresponding to
+the norm-preserving retraction of `A` along `η` with step size `α`. In particular, starting
+from a direction `η` of the form
+```math
+ξ = \\left\\langle \\frac{η}{||η||}, ξ \\right\\rangle \\frac{η}{||η||} + Δξ
+```
+where ``\\langle Δξ, A \\rangle = \\langle Δξ, η \\rangle = 0``, it returns
+```math
+ξ(α) = \\left\\langle \\frac{η}{||η||}, ξ \\right \\rangle \\left( \\cos \\left( α \\frac{||η||}{||A||} \\right) \\frac{η}{||η||} - \\sin( \\left( α \\frac{||η||}{||A||} \\right) \\frac{A}{||A||} \\right) + Δξ
+```
+such that ``||ξ(α)|| = ||ξ||, \\langle A', ξ(α) \\rangle = 0``.
+"""
+function peps_transport!(ξ, x, η, α, x´)
+    peps = x[1]
+    norms_peps = norm.(peps.A)
+
+    norms_η = norm.(η.A)
+    normalized_η = η.A ./ norms_η
+    overlaps_η_ξ = inner.(normalized_η, ξ.A)
+
+    # isolate the orthogonal component
+    Δξ = ξ.A .- overlaps_η_ξ .* normalized_η
+
+    # keep orthogonal component fixed, modify the rest by the proper directional derivative
+    ξ.A .=
+        overlaps_η_ξ .* (
+            cos.(α .* norms_η ./ norms_peps) .* normalized_η .-
+            sin.(α .* norms_η ./ norms_peps) .* peps.A ./ norms_peps
+        ) .+ Δξ
+
+    return ξ
 end
 
 # Take real valued part of dot product

src/utility/svd.jl

@@ -290,8 +290,10 @@ function ChainRulesCore.rrule(
     Ũ, S̃, Ṽ⁺, info = tsvd(t, alg; trunc, p)
     U, S, V⁺ = info.U_full, info.S_full, info.V_full # untruncated SVD decomposition
 
-    smallest_sval = minimum(minimum(abs.(diag(b))) for (_, b) in blocks(S̃))
-    pullback_tol = max(1e-14, 1e-2 * smallest_sval)
+    smallest_sval = minimum(((_, b),) -> minimum(diag(b)), blocks(S̃))
+    pullback_tol = clamp(
+        smallest_sval, eps(scalartype(S̃))^(3 / 4), eps(scalartype(S̃))^(1 / 2)
+    )
 
     function tsvd!_nothing_pullback(ΔUSVi)
         ΔU, ΔS, ΔV⁺, = unthunk.(ΔUSVi)