Clean up and simplify optimization interface

README.md

@@ -55,8 +55,8 @@ opt_alg = PEPSOptimize(;
 # ground state search
 state = InfinitePEPS(2, D)
 ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
-result = fixedpoint(H, state, ctm, opt_alg)
+peps, env, E, = fixedpoint(H, state, ctm, opt_alg)
 
-@show result.E # -0.6625...
+@show E # -0.6625...
 ```
 

docs/src/index.md

@@ -37,7 +37,7 @@ opt_alg = PEPSOptimize(;
 # ground state search
 state = InfinitePEPS(2, D)
 ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
-result = fixedpoint(H, state, ctm, opt_alg)
+peps, env, E, = fixedpoint(H, state, ctm, opt_alg)
 
-@show result.E # -0.6625...
+@show E # -0.6625...
 ```

examples/heisenberg.jl

@@ -27,5 +27,5 @@ opt_alg = PEPSOptimize(;
 # Of course there is a noticable bias for small χbond and χenv.
 ψ₀ = InfinitePEPS(2, χbond)
 env₀, = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg)
-result = fixedpoint(H, ψ, env₀, opt_alg₀)
-@show result.E
+peps, env, E, = fixedpoint(H, ψ, env₀, opt_alg₀)
+@show E

src/algorithms/ctmrg/sequential.jl

@@ -55,10 +55,10 @@ Compute CTMRG projectors in the `:sequential` scheme either for an entire column
 for a specific `coordinate` (where `dir=WEST` is already implied in the `:sequential` scheme).
 """
 function sequential_projectors(
-    col::Int, state::InfiniteSquareNetwork, envs::CTMRGEnv, alg::ProjectorAlgorithm
-)
+    col::Int, state::InfiniteSquareNetwork{T}, envs::CTMRGEnv, alg::ProjectorAlgorithm
+) where {T}
     # SVD half-infinite environment column-wise
-    ϵ = Zygote.Buffer(zeros(real(scalartype(envs)), size(envs, 2)))
+    ϵ = Zygote.Buffer(zeros(real(scalartype(T)), size(envs, 2)))
     S = Zygote.Buffer(
         zeros(size(envs, 2)), tensormaptype(spacetype(T), 1, 1, real(scalartype(T)))
     )

src/algorithms/optimization/peps_optimization.jl

@@ -43,44 +43,46 @@ function PEPSOptimize(;
 end
 
 """
-
     fixedpoint(operator, peps₀::InfinitePEPS{F}, [env₀::CTMRGEnv]; kwargs...)
     fixedpoint(operator, peps₀::InfinitePEPS{T}, alg::PEPSOptimize, [env₀::CTMRGEnv];
-               finalize!=OptimKit._finalize!, symmetrization=nothing) where {T}
+               finalize!=OptimKit._finalize!) where {T}
 
-Optimize `peps₀` with respect to the `operator` according to the parameters supplied
-in `alg`. The initial environment `env₀` serves as an initial guess for the first CTMRG run.
+Optimize `operator` starting from `peps₀` according to the parameters supplied in `alg`.
+The initial environment `env₀` serves as an initial guess for the first CTMRG run.
 By default, a random initial environment is used.
 
 The `finalize!` kwarg can be used to insert a function call after each optimization step
 by utilizing the `finalize!` kwarg of `OptimKit.optimize`.
 The function maps `(peps, envs), f, g = finalize!((peps, envs), f, g, numiter)`.
 The `symmetrization` kwarg accepts `nothing` or a `SymmetrizationStyle`, in which case the
 PEPS and PEPS gradient are symmetrized after each optimization iteration. Note that this
-requires a symmmetric `ψ₀` and `env₀` to converge properly.
+requires a symmmetric `peps₀` and `env₀` to converge properly.
 
-The function returns a `NamedTuple` which contains the following entries:
-- `peps`: final `InfinitePEPS`
-- `env`: `CTMRGEnv` corresponding to the final PEPS
-- `E`: final energy
-- `E_history`: convergence history of the energy function
-- `grad`: final energy gradient
-- `gradnorm_history`: convergence history of the energy gradient norms
-- `numfg`: total number of calls to the energy function
+The function returns the final PEPS, CTMRG environment and cost value, as well as an
+information `NamedTuple` which contains the following entries:
+- `last_gradient`: last gradient of the cost function
+- `fg_evaluations`: number of evaluations of the cost and gradient function
+- `costs`: history of cost values
+- `gradnorms`: history of gradient norms
+- `truncation_errors`: history of truncation errors of the boundary algorithm
+- `condition_numbers`: history of condition numbers of the CTMRG environments
+- `gradnorms_unitcell`: history of gradient norms for each respective unit cell entry
+- `times`: history of times each optimization step took
 """
 function fixedpoint(
-    operator, peps₀::InfinitePEPS{F}, env₀::CTMRGEnv=CTMRGEnv(peps₀, field(F)^20); kwargs...
-) where {F}
+    operator, peps₀::InfinitePEPS{T}, env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20); kwargs...
+) where {T}
     alg = fixedpoint_selector(; kwargs...) # TODO: implement fixedpoint_selector
     return fixedpoint(operator, peps₀, env₀, alg)
 end
 function fixedpoint(
     operator,
-    peps₀::InfinitePEPS,
+    peps₀::InfinitePEPS{T},
     env₀::CTMRGEnv,
     alg::PEPSOptimize;
     (finalize!)=OptimKit._finalize!,
-)
+) where {T}
+    # setup retract and finalize! for symmetrization
     if isnothing(alg.symmetrization)
         retract = peps_retract
     else
@@ -89,42 +91,57 @@ function fixedpoint(
         finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter)
     end
 
+    # check realness compatibility
     if scalartype(env₀) <: Real && iterscheme(alg.gradient_alg) == :fixed
         env₀ = complex(env₀)
         @warn "the provided real environment was converted to a complex environment since \
         :fixed mode generally produces complex gauges; use :diffgauge mode instead to work \
         with purely real environments"
     end
 
-    (peps, env), E, ∂E, numfg, convhistory = optimize(
+    # initialize info collection vectors
+    truncation_errors = Vector{real(scalartype(T))}()
+    condition_numbers = Vector{real(scalartype(T))}()
+    gradnorms_unitcell = Vector{Matrix{real(scalartype(T))}}()
+    times = Float64[]
+
+    # optimize operator cost function
+    (peps_final, env_final), cost, ∂cost, numfg, convergence_history = optimize(
         (peps₀, env₀), alg.optimizer; retract, inner=real_inner, finalize!
-    ) do (peps, envs)
+    ) do (peps, env)
+        start_time = time_ns()
         E, gs = withgradient(peps) do ψ
-            envs´, = hook_pullback(
+            env′, info = hook_pullback(
                 leading_boundary,
-                envs,
+                env,
                 ψ,
                 alg.boundary_alg;
                 alg_rrule=alg.gradient_alg,
             )
             ignore_derivatives() do
-                alg.reuse_env && update!(envs, envs´)
+                alg.reuse_env && update!(env, env′)
+                push!(truncation_errors, info.truncation_error)
+                push!(condition_numbers, info.condition_number)
             end
-            return cost_function(ψ, envs´, operator)
+            return cost_function(ψ, env′, operator)
         end
         g = only(gs)  # `withgradient` returns tuple of gradients `gs`
+        push!(gradnorms_unitcell, norm.(g.A))
+        push!(times, (time_ns() - start_time) * 1e-9)
         return E, g
     end
 
-    return (;
-        peps,
-        env,
-        E,
-        E_history=convhistory[:, 1],
-        grad=∂E,
-        gradnorm_history=convhistory[:, 2],
-        numfg,
+    info = (
+        last_gradient=∂cost,
+        fg_evaluations=numfg,
+        costs=convergence_history[:, 1],
+        gradnorms=convergence_history[:, 2],
+        truncation_errors,
+        condition_numbers,
+        gradnorms_unitcell,
+        times,
     )
+    return peps_final, env_final, cost, info
 end
 
 # Update PEPS unit cell in non-mutating way
@@ -138,3 +155,23 @@ end
 
 # Take real valued part of dot product
 real_inner(_, η₁, η₂) = real(dot(η₁, η₂))
+
+"""
+    symmetrize_retract_and_finalize!(symm::SymmetrizationStyle)
+
+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, envs), E, grad, _)
+        grad_symm = symmetrize!(grad, symm)
+        return (peps, envs), E, grad_symm
+    end
+    retf = function symmetrize_retract((peps, envs), η, α)
+        peps_symm = deepcopy(peps)
+        peps_symm.A .+= η.A .* α
+        envs′ = deepcopy(envs)
+        symmetrize!(peps_symm, symm)
+        return (peps_symm, envs′), η
+    end
+    return retf, finf
+end

src/utility/symmetrization.jl

@@ -216,23 +216,3 @@ function symmetrize!(peps::InfinitePEPS, symm::RotateReflect)
     end
     return peps
 end
-
-"""
-    symmetrize_retract_and_finalize!(symm::SymmetrizationStyle)
-
-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, envs), E, grad, _)
-        grad_symm = symmetrize!(grad, symm)
-        return (peps, envs), E, grad_symm
-    end
-    retf = function symmetrize_retract((peps, envs), η, α)
-        peps_symm = deepcopy(peps)
-        peps_symm.A .+= η.A .* α
-        e = deepcopy(envs)
-        symmetrize!(peps_symm, symm)
-        return (peps_symm, e), η
-    end
-    return retf, finf
-end

test/heisenberg.jl

@@ -21,14 +21,14 @@ E_ref = -0.6602310934799577
     # initialize states
     Random.seed!(123)
     H = heisenberg_XYZ(InfiniteSquare())
-    psi_init = InfinitePEPS(2, Dbond)
-    env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
+    peps₀ = InfinitePEPS(2, Dbond)
+    env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χenv)), peps₀, ctm_alg)
 
     # optimize energy and compute correlation lengths
-    result = fixedpoint(H, psi_init, env_init, opt_alg)
-    ξ_h, ξ_v, = correlation_length(result.peps, result.env)
+    peps, env, E, = fixedpoint(H, peps₀, env₀, opt_alg)
+    ξ_h, ξ_v, = correlation_length(peps, env)
 
-    @test result.E ≈ E_ref atol = 1e-2
+    @test E ≈ E_ref atol = 1e-2
     @test all(@. ξ_h > 0 && ξ_v > 0)
 end
 
@@ -37,14 +37,14 @@ end
     Random.seed!(456)
     unitcell = (1, 2)
     H = heisenberg_XYZ(InfiniteSquare(unitcell...))
-    psi_init = InfinitePEPS(2, Dbond; unitcell)
-    env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
+    peps₀ = InfinitePEPS(2, Dbond; unitcell)
+    env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χenv)), peps₀, ctm_alg)
 
     # optimize energy and compute correlation lengths
-    result = fixedpoint(H, psi_init, env_init, opt_alg)
-    ξ_h, ξ_v, = correlation_length(result.peps, result.env)
+    peps, env, E, = fixedpoint(H, peps₀, env₀, opt_alg)
+    ξ_h, ξ_v, = correlation_length(peps, env)
 
-    @test result.E ≈ 2 * E_ref atol = 1e-2
+    @test E ≈ 2 * E_ref atol = 1e-2
     @test all(@. ξ_h > 0 && ξ_v > 0)
 end
 
@@ -92,9 +92,9 @@ end
     # continue with auto differentiation
     svd_alg_gmres = SVDAdjoint(; rrule_alg=GMRES(; tol=1e-5))
     opt_alg_gmres = @set opt_alg.boundary_alg.projector_alg.svd_alg = svd_alg_gmres
-    result_final = fixedpoint(ham, peps, envs, opt_alg_gmres)  # sensitivity warnings and degeneracies due to SU(2)?
-    ξ_h, ξ_v, = correlation_length(result_final.peps, result_final.env)
-    e_site2 = result_final.E / (N1 * N2)
+    peps_final, env_final, E_final, = fixedpoint(ham, peps, envs, opt_alg_gmres)  # sensitivity warnings and degeneracies due to SU(2)?
+    ξ_h, ξ_v, = correlation_length(peps_final, env_final)
+    e_site2 = E_final / (N1 * N2)
     @info "Auto diff energy = $e_site2"
     @test e_site2 ≈ E_ref atol = 1e-2
     @test all(@. ξ_h > 0 && ξ_v > 0)

test/j1j2_model.jl

@@ -19,17 +19,17 @@ opt_alg = PEPSOptimize(;
 # initialize states
 Random.seed!(91283219347)
 H = j1_j2(InfiniteSquare(); J2=0.25)
-psi_init = product_peps(2, χbond; noise_amp=1e-1)
-psi_init = symmetrize!(psi_init, RotateReflect())
-env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg);
+peps₀ = product_peps(2, χbond; noise_amp=1e-1)
+peps₀ = symmetrize!(peps₀, RotateReflect())
+env_init, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χenv)), peps₀, ctm_alg);
 
 # find fixedpoint
-result = fixedpoint(H, psi_init, env_init, opt_alg)
-ξ_h, ξ_v, = correlation_length(result.peps, result.env)
+peps, env, E, = fixedpoint(H, peps₀, env₀, opt_alg)
+ξ_h, ξ_v, = correlation_length(peps, env)
 
 # compare against Juraj Hasik's data:
 # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.25/state_1s_A1_j20.25_D2_chi_opt48.dat
 ξ_ref = -1 / log(0.2723596743547324)
-@test result.E ≈ -0.5618837021945925 atol = 1e-3
+@test E ≈ -0.5618837021945925 atol = 1e-3
 @test all(@. isapprox(ξ_h, ξ_ref; atol=1e-1) && isapprox(ξ_v, ξ_ref; atol=1e-1))
 @test ξ_h ≈ ξ_v atol = 1e-6  # Test symmetrization of optimized PEPS and environment

test/pwave.jl

@@ -20,11 +20,11 @@ Pspace = Vect[FermionParity](0 => 1, 1 => 1)
 Vspace = Vect[FermionParity](0 => Dbond ÷ 2, 1 => Dbond ÷ 2)
 Envspace = Vect[FermionParity](0 => χenv ÷ 2, 1 => χenv ÷ 2)
 Random.seed!(91283219347)
-psi_init = InfinitePEPS(Pspace, Vspace, Vspace; unitcell)
-env_init, = leading_boundary(CTMRGEnv(psi_init, Envspace), psi_init, ctm_alg);
+peps₀ = InfinitePEPS(Pspace, Vspace, Vspace; unitcell)
+env₀, = leading_boundary(CTMRGEnv(peps₀, Envspace), peps₀, ctm_alg);
 
 # find fixedpoint
-result = fixedpoint(H, psi_init, env_init, opt_alg)
+_, _, E, = fixedpoint(H, peps₀, env₀, opt_alg)
 
 # comparison with Gaussian PEPS minimum at D=2 on 1000x1000 square lattice with aPBC
-@test result.E / prod(size(psi_init)) ≈ -2.6241 atol = 5e-2
+@test E / prod(size(peps₀)) ≈ -2.6241 atol = 5e-2

test/tf_ising.jl

@@ -27,25 +27,25 @@ opt_alg = PEPSOptimize(;
 # initialize states
 H = transverse_field_ising(InfiniteSquare(); g)
 Random.seed!(2928528935)
-psi_init = InfinitePEPS(2, χbond)
-env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
+peps₀ = InfinitePEPS(2, χbond)
+env₀, = leading_boundary(CTMRGEnv(peps₀, ComplexSpace(χenv)), peps₀, ctm_alg)
 
 # find fixedpoint
-result = fixedpoint(H, psi_init, env_init, opt_alg)
-ξ_h, ξ_v, = correlation_length(result.peps, result.env)
+peps, env, E, = fixedpoint(H, peps₀, env₀, opt_alg)
+ξ_h, ξ_v, = correlation_length(peps, env)
 
 # compute magnetization
-σx = TensorMap(scalartype(psi_init)[0 1; 1 0], ℂ^2, ℂ^2)
+σx = TensorMap(scalartype(peps₀)[0 1; 1 0], ℂ^2, ℂ^2)
 M = LocalOperator(H.lattice, (CartesianIndex(1, 1),) => σx)
-magn = expectation_value(result.peps, M, result.env)
+magn = expectation_value(peps, M, env)
 
 @test result.E ≈ e atol = 1e-2
 @test imag(magn) ≈ 0 atol = 1e-6
 @test abs(magn) ≈ mˣ atol = 5e-2
 
 # find fixedpoint in polarized phase and compute correlations lengths
 H_polar = transverse_field_ising(InfiniteSquare(); g=4.5)
-result_polar = fixedpoint(H_polar, psi_init, env_init, opt_alg)
-ξ_h_polar, ξ_v_polar, = correlation_length(result_polar.peps, result_polar.env)
+peps_polar, env_polar, = fixedpoint(H_polar, peps₀, env₀, opt_alg)
+ξ_h_polar, ξ_v_polar, = correlation_length(peps_polar, env_polar)
 @test ξ_h_polar < ξ_h
 @test ξ_v_polar < ξ_v