Clean up and simplify optimization interface

README.md

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

docs/src/index.md

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

examples/boundary_mps.jl

@@ -32,7 +32,7 @@ mps, envs, ϵ = leading_boundary(mps, T, VUMPS())
 N = abs(prod(expectation_value(mps, T)))
 
 # This can be compared to the result obtained using the CTMRG algorithm
-ctm = leading_boundary(
+ctm, = leading_boundary(
     peps, SimultaneousCTMRG(; verbosity=1), CTMRGEnv(peps, ComplexSpace(20))
 )
 N´ = abs(norm(peps, ctm))
@@ -55,7 +55,7 @@ mps2 = PEPSKit.initializeMPS(T2, fill(ComplexSpace(20), 2, 2))
 mps2, envs2, ϵ = leading_boundary(mps2, T2, VUMPS())
 N2 = abs(prod(expectation_value(mps2, T2)))
 
-ctm2 = leading_boundary(
+ctm2, = leading_boundary(
     peps2, SimultaneousCTMRG(; verbosity=1), CTMRGEnv(peps2, ComplexSpace(20))
 )
 N2´ = abs(norm(peps2, ctm2))

examples/heisenberg.jl

@@ -26,6 +26,6 @@ opt_alg = PEPSOptimize(;
 # E/N = −0.6694421, which is a QMC estimate from https://arxiv.org/abs/1101.3281.
 # 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, opt_alg, env₀)
-@show result.E
+env₀, = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg)
+peps, env, E, = fixedpoint(H, ψ, env₀, opt_alg₀)
+@show E

examples/hubbard_su.jl

@@ -43,7 +43,7 @@ Espace = Vect[fℤ₂](0 => χenv0 / 2, 1 => χenv0 / 2)
 envs = CTMRGEnv(randn, Float64, peps, Espace)
 for χ in [χenv0, χenv]
     ctm_alg = SequentialCTMRG(; maxiter=300, tol=1e-7)
-    envs = leading_boundary(envs, peps, ctm_alg)
+    envs, = leading_boundary(envs, peps, ctm_alg)
 end
 
 # Benchmark values of the ground state energy from
@@ -53,7 +53,7 @@ Es_exact = Dict(0 => -1.62, 2 => -0.176, 4 => 0.8603, 6 => -0.6567, 8 => -0.5243
 E_exact = Es_exact[U] - U / 2
 
 # measure energy
-E = costfun(peps, envs, ham) / (N1 * N2)
+E = cost_function(peps, envs, ham) / (N1 * N2)
 @info "Energy           = $E"
 @info "Benchmark energy = $E_exact"
 @test isapprox(E, E_exact; atol=5e-2)

src/PEPSKit.jl

@@ -54,56 +54,65 @@ include("algorithms/time_evolution/simpleupdate.jl")
 
 include("algorithms/toolbox.jl")
 
-include("algorithms/peps_opt.jl")
-
 include("utility/symmetrization.jl")
 
+include("algorithms/optimization/fixed_point_differentiation.jl")
+include("algorithms/optimization/peps_optimization.jl")
+
 """
     module Defaults
-        const ctmrg_maxiter = 100
-        const ctmrg_miniter = 4
-        const ctmrg_tol = 1e-8
-        const fpgrad_maxiter = 30
-        const fpgrad_tol = 1e-6
-        const reuse_env = true
-        const trscheme = FixedSpaceTruncation()
-        const fwd_alg = TensorKit.SDD()
-        const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
-        const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
-        const projector_alg_type = HalfInfiniteProjector
-        const projector_alg = projector_alg_type(svd_alg, trscheme, 2)
-        const ctmrg_alg = SimultaneousCTMRG(
-            ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg
-        )
-        const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)
-        const gradient_linsolver = KrylovKit.BiCGStab(;
-            maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
-        )
-        const iterscheme = :fixed
-        const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
-        const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler())
-    end
 
-Module containing default values that represent typical algorithm parameters.
+Module containing default algorithm parameter values and arguments.
 
-- `ctmrg_maxiter`: Maximal number of CTMRG iterations per run
-- `ctmrg_miniter`: Minimal number of CTMRG carried out
-- `ctmrg_tol`: Tolerance checking singular value and norm convergence
-- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient
-- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration
-- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment
-- `trscheme`: Truncation scheme for SVDs and other decompositions
-- `fwd_alg`: SVD algorithm that is used in the forward pass
+# CTMRG
+- `ctmrg_tol=1e-8`: Tolerance checking singular value and norm convergence
+- `ctmrg_maxiter=100`: Maximal number of CTMRG iterations per run
+- `ctmrg_miniter=4`: Minimal number of CTMRG carried out
+- `trscheme=FixedSpaceTruncation()`: Truncation scheme for SVDs and other decompositions
+- `fwd_alg=TensorKit.SDD()`: SVD algorithm that is used in the forward pass
 - `rrule_alg`: Reverse-rule for differentiating that SVD
-- `svd_alg`: Combination of `fwd_alg` and `rrule_alg`
-- `projector_alg_type`: Default type of projector algorithm
+
+    ```
+    rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1)
+    ```
+
+- `svd_alg=SVDAdjoint(; fwd_alg, rrule_alg)`: Combination of `fwd_alg` and `rrule_alg`
+- `projector_alg_type=HalfInfiniteProjector`: Default type of projector algorithm
 - `projector_alg`: Algorithm to compute CTMRG projectors
+
+    ```
+    projector_alg = projector_alg_type(; svd_alg, trscheme, verbosity=0)
+    ```
+
 - `ctmrg_alg`: Algorithm for performing CTMRG runs
-- `optimizer`: Optimization algorithm for PEPS ground-state optimization
+
+    ```
+    ctmrg_alg = SimultaneousCTMRG(
+        ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg
+    )
+    ```
+
+# Optimization
+- `fpgrad_maxiter=30`: Maximal number of iterations for computing the CTMRG fixed-point gradient
+- `fpgrad_tol=1e-6`: Convergence tolerance for the fixed-point gradient iteration
+- `iterscheme=:fixed`: Scheme for differentiating one CTMRG iteration
 - `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm
-- `iterscheme`: Scheme for differentiating one CTMRG iteration
+
+    ```
+    gradient_linsolver=KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol)
+    ```
+
 - `gradient_alg`: Algorithm to compute the gradient fixed-point
-- `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!`
+
+    ```
+    gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+    ```
+
+- `reuse_env=true`: If `true`, the current optimization step is initialized on the previous environment
+- `optimizer=LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)`: Default `OptimKit.OptimizerAlgorithm` for PEPS optimization
+
+# OhMyThreads scheduler
+- `scheduler=Ref{Scheduler}(...)`: Multi-threading scheduler which can be accessed via `set_scheduler!`
 """
 module Defaults
     using TensorKit, KrylovKit, OptimKit, OhMyThreads
@@ -113,28 +122,30 @@ module Defaults
         SVDAdjoint,
         HalfInfiniteProjector,
         SimultaneousCTMRG
+
+    # CTMRG
     const ctmrg_tol = 1e-8
     const ctmrg_maxiter = 100
     const ctmrg_miniter = 4
-    const fpgrad_maxiter = 30
-    const fpgrad_tol = 1e-6
     const sparse = false
-    const reuse_env = true
     const trscheme = FixedSpaceTruncation()
     const fwd_alg = TensorKit.SDD()
-    const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
+    const rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1)
     const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
     const projector_alg_type = HalfInfiniteProjector
-    const projector_alg = projector_alg_type(svd_alg, trscheme, 2)
+    const projector_alg = projector_alg_type(; svd_alg, trscheme, verbosity=0)
     const ctmrg_alg = SimultaneousCTMRG(
         ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg
     )
-    const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)
-    const gradient_linsolver = KrylovKit.BiCGStab(;
-        maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
-    )
+
+    # Optimization
+    const fpgrad_maxiter = 30
+    const fpgrad_tol = 1e-6
+    const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol)
     const iterscheme = :fixed
     const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
+    const reuse_env = true
+    const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)
 
     # OhMyThreads scheduler defaults
     const scheduler = Ref{Scheduler}()
@@ -187,7 +198,7 @@ export SVDAdjoint, IterSVD
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
 export FixedSpaceTruncation, HalfInfiniteProjector, FullInfiniteProjector
 export LocalOperator
-export expectation_value, costfun, product_peps, correlation_length
+export expectation_value, cost_function, product_peps, correlation_length
 export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver
 export fixedpoint

src/algorithms/ctmrg/ctmrg.jl

@@ -38,10 +38,13 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm)
     env = deepcopy(envinit)
     log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG"))
 
+    truncation_error = zero(real(scalartype(state)))
+    condition_number = one(real(scalartype(state)))
+    info = (; truncation_error, condition_number)
     return LoggingExtras.withlevel(; alg.verbosity) do
         ctmrg_loginit!(log, η, state, envinit)
         for iter in 1:(alg.maxiter)
-            env, = ctmrg_iteration(state, env, alg)  # Grow and renormalize in all 4 directions
+            env, info = ctmrg_iteration(state, env, alg)  # Grow and renormalize in all 4 directions
             η, CS, TS = calc_convergence(env, CS, TS)
 
             if η ≤ alg.tol && iter ≥ alg.miniter
@@ -54,7 +57,7 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm)
                 ctmrg_logiter!(log, iter, η, state, env)
             end
         end
-        return env
+        return env, info
     end
 end
 

src/algorithms/ctmrg/sequential.jl

@@ -31,18 +31,20 @@ function SequentialCTMRG(;
 end
 
 function ctmrg_iteration(state, envs::CTMRGEnv, alg::SequentialCTMRG)
-    ϵ = zero(real(scalartype(state)))
+    truncation_error = zero(real(scalartype(state)))
+    condition_number = zero(real(scalartype(state)))
     for _ in 1:4 # rotate
         for col in 1:size(state, 2) # left move column-wise
             projectors, info = sequential_projectors(col, state, envs, alg.projector_alg)
             envs = renormalize_sequentially(col, projectors, state, envs)
-            ϵ = max(ϵ, info.err)
+            truncation_error = max(truncation_error, info.truncation_error)
+            condition_number = max(condition_number, info.condition_number)
         end
         state = rotate_north(state, EAST)
         envs = rotate_north(envs, EAST)
     end
 
-    return envs, (; err=ϵ)
+    return envs, (; truncation_error, condition_number)
 end
 
 """
@@ -53,21 +55,28 @@ 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)))
+    )
     coordinates = eachcoordinate(envs)[:, col]
     projectors = dtmap(coordinates) do (r, c)
         trscheme = truncation_scheme(alg, envs.edges[WEST, _prev(r, size(envs, 2)), c])
         proj, info = sequential_projectors(
             (WEST, r, c), state, envs, @set(alg.trscheme = trscheme)
         )
+        S[r] = info.S
         ϵ[r] = info.err / norm(info.S)
         return proj
     end
 
-    return (map(first, projectors), map(last, projectors)), (; err=maximum(copy(ϵ)))
+    truncation_error = maximum(copy(ϵ))
+    condition_number = maximum(_condition_number, S)
+    info = (; truncation_error, condition_number)
+    return (map(first, projectors), map(last, projectors)), info
 end
 function sequential_projectors(
     coordinate::NTuple{3,Int},

src/algorithms/ctmrg/simultaneous.jl

@@ -48,6 +48,16 @@ function _prealloc_svd(edges::Array{E,N}) where {E,N}
     return U, S, V
 end
 
+# Compute condition number σ_max / σ_min for diagonal singular value TensorMap
+function _condition_number(S::AbstractTensorMap)
+    # Take maximal condition number over all blocks
+    return maximum(blocks(S)) do (_, b)
+        b_diag = diag(b)
+        return maximum(b_diag) / minimum(b_diag)
+    end
+end
+@non_differentiable _condition_number(S::AbstractTensorMap)
+
 """
     simultaneous_projectors(enlarged_corners::Array{E,3}, envs::CTMRGEnv, alg::ProjectorAlgorithm)
     simultaneous_projectors(coordinate, enlarged_corners::Array{E,3}, alg::ProjectorAlgorithm)
@@ -74,9 +84,10 @@ function simultaneous_projectors(
         return proj
     end
 
-    P_left = map(first, projectors)
-    P_right = map(last, projectors)
-    return (P_left, P_right), (; err=maximum(copy(ϵ)), U=copy(U), S=copy(S), V=copy(V))
+    truncation_error = maximum(copy(ϵ))
+    condition_number = maximum(_condition_number, S)
+    info = (; truncation_error, condition_number, U=copy(U), S=copy(S), V=copy(V))
+    return (map(first, projectors), map(last, projectors)), info
 end
 function simultaneous_projectors(
     coordinate, enlarged_corners::Array{E,3}, alg::HalfInfiniteProjector