Split off norm-preserving vector retractions (#160)

* Split off vector retractions, remove diffset

* Add dedicated retraction test

* Apply suggestions from code review

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

* Address remaining review comments

* Fix typos

* Fix typo

* Distinguishing fancy ticks is hard

---------

Co-authored-by: Lukas Devos <ldevos98@gmail.com>

Author: leburgel

src/PEPSKit.jl

@@ -19,9 +19,9 @@ include("utility/diffable_threads.jl")
 include("utility/svd.jl")
 include("utility/rotations.jl")
 include("utility/mirror.jl")
-include("utility/diffset.jl")
 include("utility/hook_pullback.jl")
 include("utility/autoopt.jl")
+include("utility/retractions.jl")
 
 include("networks/tensors.jl")
 include("networks/local_sandwich.jl")

src/algorithms/optimization/peps_optimization.jl

@@ -197,9 +197,9 @@ function fixedpoint(
     if isnothing(alg.symmetrization)
         retract = peps_retract
     else
-        retract, symm_finalize! = symmetrize_retract_and_finalize!(alg.symmetrization)
-        fin! = finalize!  # Previous finalize!
-        finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter)
+        retract, finalize! = symmetrize_retract_and_finalize!(
+            alg.symmetrization, peps_retract, finalize!
+        )
     end
 
     # :fixed mode compatibility
@@ -284,30 +284,21 @@ end
 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{η}{||η||},
+A' ← \\cos ( α ‖η‖ / ‖A‖ ) A + \\sin ( α ‖η‖ / ‖A‖ ) ‖A‖ η / ‖η‖,
 ```
 and corresponding directional derivative `ξ`,
 ```math
-ξ = \\cos \\left( α \\frac{||η||}{||A||} \\right) η - \\sin \\left( α \\frac{||η||}{||A||} \\right) ||η|| \\frac{A}{||A||},
+ξ = \\cos ( α ‖η‖ / ‖A‖ ) η - \\sin ( α ‖η‖ / ‖A‖ ) ‖η‖ A / ‖A‖,
 ```
-such that ``\\langle A', ξ \\rangle = 0`` and ``||A'|| = ||A||``.
+such that ``⟨ A', ξ ⟩ = 0`` and ``‖A'‖ = ‖A‖``.
 """
 function peps_retract(x, η, α)
     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])
 
-    ξ = similar(η)
-    ξ.A .=
-        cos.(α .* norms_η ./ norms_peps) .* η.A .-
-        sin.(α .* norms_η ./ norms_peps) .* norms_η .* peps.A ./ norms_peps
+    retractions = norm_preserving_retract.(unitcell(peps), unitcell(η), α)
+    peps´ = InfinitePEPS(map(first, retractions))
+    ξ = InfinitePEPS(map(last, retractions))
 
     return (peps´, env), ξ
 end
@@ -319,31 +310,21 @@ Transports a direction at `A = x[1]` to a valid direction at `A´ = x´[1]` corr
 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
+where ``⟨ Δξ, A ⟩ = ⟨ Δξ, η ⟩ = 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) + Δξ
+ξ(α) = ⟨ η / ‖η‖, ξ ⟩ ( \\cos ( α ‖η‖ / ‖A‖ ) η / ‖η‖ - \\sin( α ‖η‖ / ‖A‖ ) A / ‖A‖ ) + Δξ
 ```
-such that ``||ξ(α)|| = ||ξ||, \\langle A', ξ(α) \\rangle = 0``.
+such that ``‖ξ(α)‖ = ‖ξ‖, ⟨ A', ξ(α) ⟩ = 0``.
 """
 function peps_transport!(ξ, x, η, α, x´)
     peps = x[1]
-    norms_peps = norm.(peps.A)
-
-    norms_η = norm.(η.A)
-    normalized_η = η.A ./ norms_η
-    overlaps_η_ξ = inner.(normalized_η, ξ.A)
+    peps´ = x´[1]
 
-    # 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
-        ) .+ Δξ
+    norm_preserving_transport!.(
+        unitcell(ξ), unitcell(peps), unitcell(η), α, unitcell(peps´)
+    )
 
     return ξ
 end
@@ -356,17 +337,22 @@ real_inner(_, η₁, η₂) = real(dot(η₁, η₂))
 
 Return the `retract` and `finalize!` function for symmetrizing the `peps` and `grad` tensors.
 """
-function symmetrize_retract_and_finalize!(symm::SymmetrizationStyle)
-    finf = function symmetrize_finalize!((peps, env), E, grad, _)
+function symmetrize_retract_and_finalize!(
+    symm::SymmetrizationStyle, retract=peps_retract, (finalize!)=OptimKit._finalize!
+)
+    function symmetrize_then_finalize!((peps, env), E, grad, numiter)
+        # symmetrize the gradient
         grad_symm = symmetrize!(grad, symm)
-        return (peps, env), E, grad_symm
+        # then finalize
+        return finalize!((peps, env), E, grad_symm, numiter)
     end
-    retf = function symmetrize_retract((peps, env), η, α)
-        peps_symm = deepcopy(peps)
-        peps_symm.A .+= η.A .* α
-        env′ = deepcopy(env)
-        symmetrize!(peps_symm, symm)
-        return (peps_symm, env′), η
+    function retract_then_symmetrize((peps, env), η, α)
+        # retract
+        (peps´, env´), ξ = retract((peps, env), η, α)
+        # symmetrize retracted point and directional derivative
+        peps´_symm = symmetrize!(peps´, symm)
+        ξ_symm = symmetrize!(ξ, symm)
+        return (peps´_symm, env´), ξ_symm
     end
-    return retf, finf
+    return retract_then_symmetrize, symmetrize_then_finalize!
 end

src/utility/diffset.jl

@@ -0,0 +1,65 @@
+#=
+Utilities for preserving the norm of (VectorInterface-compliant) vectors during optimization.
+=#
+
+"""
+    norm_preserving_retract(A, η, α)
+
+Performs a norm-preserving retraction of vector `A` along the direction `η` with step size
+`α`, giving a new vector `A´`,
+```math
+A' ← \\cos ( α ‖η‖ / ‖A‖ ) A + \\sin ( α ‖η‖ / ‖A‖ ) ‖A‖ η / ‖η‖,
+```
+and corresponding directional derivative `ξ`,
+```math
+ξ = \\cos ( α ‖η‖ / ‖A‖ ) η - \\sin ( α ‖η‖ / ‖A‖ ) ‖η‖ A / ‖A‖,
+```
+such that ``⟨ A', ξ ⟩ = 0`` and ``‖A'‖ = ‖A‖``.
+
+!!! note
+    The vectors `A` and `η` should satisfy the interface specified by
+    [VectorInterface.jl](https://github.com/Jutho/VectorInterface.jl)
+
+"""
+function norm_preserving_retract(A, η, α)
+    n_A = norm(A)
+    n_η = norm(η)
+    sn, cs = sincos(α * n_η / n_A)
+
+    A´ = add(A, η, sn * n_A / n_η, cs)
+    ξ = add(A, η, cs, -sn * n_η / n_A)
+
+    return A´, ξ
+end
+
+"""
+    norm_preserving_transport!(ξ, A, η, α, A′)
+
+Transports a direction `ξ` at `A` to a valid direction at `A´` corresponding to
+the norm-preserving retraction of `A` along `η` with step size `α`. In particular, starting
+from a direction `η` of the form
+```math
+ξ = ⟨ η / ‖η‖, ξ ⟩ η / ‖η‖ + Δξ
+```
+where ``⟨ Δξ, A ⟩ = ⟨ Δξ, η ⟩ = 0``, it returns
+```math
+ξ(α) = ⟨ η / ‖η‖, ξ ⟩ ( \\cos ( α ‖η‖ / ‖A‖ ) η / ‖η‖ - \\sin ( α ‖η‖ / ‖A‖ ) A / ‖A‖ ) + Δξ
+```
+such that ``‖ξ(α)‖ = ‖ξ‖, ⟨ A', ξ(α) ⟩ = 0``.
+
+!!! note
+    The vectors `A` and `η` should satisfy the interface specified by
+    [VectorInterface.jl](https://github.com/Jutho/VectorInterface.jl)
+
+"""
+function norm_preserving_transport!(ξ, A, η, α, A´)
+    n_A = norm(A)
+    n_η = norm(η)
+    sn, cs = sincos(α * n_η / n_A)
+
+    overlaps_η_ξ = inner(η, ξ) / n_η
+    add!(ξ, η, (cs - 1) * overlaps_η_ξ / n_η)
+    add!(ξ, A, -sn * overlaps_η_ξ / n_A)
+
+    return ξ
+end

src/utility/retractions.jl

@@ -58,6 +58,9 @@ end
         @time @safetestset "Differentiable tmap" begin
             include("utility/diff_maps.jl")
         end
+        @time @safetestset "Norm-preserving retractions" begin
+            include("utility/retractions.jl")
+        end
     end
     if GROUP == "ALL" || GROUP == "EXAMPLES"
         @time @safetestset "Transverse Field Ising model" begin

test/runtests.jl

@@ -0,0 +1,33 @@
+using Test
+using Random
+using LinearAlgebra
+using TensorKit
+using VectorInterface
+using PEPSKit
+
+dtype = ComplexF64
+Vphyss = [ℂ^2, U1Space(0 => 1, -1 => 1, 1 => 1)]
+Vpepss = [ℂ^4, U1Space(0 => 2, -1 => 1, 1 => 1)]
+
+@testset "Norm-preserving tensor retractions for sectortype $(sectortype(Vphyss[i]))" for i in
+                                                                                          eachindex(
+    Vphyss
+)
+    Vphys = Vphyss[i]
+    Vpeps = Vpepss[i]
+    peps_space = Vphys ← Vpeps ⊗ Vpeps ⊗ Vpeps' ⊗ Vpeps'
+
+    α = 1e-1 * randn(Float64)
+    A = randn(dtype, peps_space)
+    normalized_A = scale(A, inv(norm(A)))
+    η = randn(dtype, peps_space)
+    ζ = randn(dtype, peps_space)
+    add!(η, normalized_A, -inner(normalized_A, η))
+    add!(ζ, normalized_A, -inner(normalized_A, ζ))
+
+    A´, ξ = PEPSKit.norm_preserving_retract(A, η, α)
+    @test norm(A´) ≈ norm(A) rtol = 1e-12
+
+    PEPSKit.norm_preserving_transport!(ζ, A, η, α, A´)
+    @test inner(ζ, A´) ≈ 0 atol = 1e-12
+end