Clean up and simplify optimization interface (#127)

- Return truncation error and condition number of projectors in `ctmrg_iter` and `leading_boundary`
- Track truncation errors, condition numbers, unit-cell-wise gradient norms and optimization step times during `fixedpoint` optimization
- Rename `costfun` to `cost_function`
- Add `symmetrization::SymmetrizationStyle` field to `PEPSOptimize`
- Fix naming convention for `CTMRGEnv` to `env`
- Break up `peps_opt.jl` file

---------

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

Author: pbrehmer

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

@@ -26,13 +26,13 @@ T = InfiniteTransferPEPS(peps, 1, 1)
 mps = PEPSKit.initializeMPS(T, [ComplexSpace(20)])
 
 # We then find the leading boundary MPS fixed point using the VUMPS algorithm
-mps, envs, ϵ = leading_boundary(mps, T, VUMPS())
+mps, env, ϵ = leading_boundary(mps, T, VUMPS())
 
 # The norm of the state per unit cell is then given by the expectation value <mps|T|mps>
 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))
@@ -52,10 +52,10 @@ peps2 = InfinitePEPS(ComplexSpace(2), ComplexSpace(2); unitcell=(2, 2))
 T2 = PEPSKit.MultilineTransferPEPS(peps2, 1)
 
 mps2 = PEPSKit.initializeMPS(T2, fill(ComplexSpace(20), 2, 2))
-mps2, envs2, ϵ = leading_boundary(mps2, T2, VUMPS())
+mps2, env2, ϵ = 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))
@@ -77,7 +77,7 @@ pepo = ising_pepo(1)
 T3 = InfiniteTransferPEPO(peps, pepo, 1, 1)
 
 mps3 = PEPSKit.initializeMPS(T3, [ComplexSpace(20)])
-mps3, envs3, ϵ = leading_boundary(mps3, T3, VUMPS())
+mps3, env3, ϵ = leading_boundary(mps3, T3, VUMPS())
 @show N3 = abs(prod(expectation_value(mps3, T3)))
 
 # These objects and routines can be used to optimize PEPS fixed points of 3D partition

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/heisenberg_evol/heis_tools.jl

@@ -16,7 +16,7 @@ using PEPSKit
 """
 Measure magnetization on each site
 """
-function cal_mags(peps::InfinitePEPS, envs::CTMRGEnv)
+function cal_mags(peps::InfinitePEPS, env::CTMRGEnv)
     N1, N2 = size(peps)
     lattice = collect(space(t, 1) for t in peps.A)
     # detect symmetry on physical axis
@@ -32,7 +32,7 @@ function cal_mags(peps::InfinitePEPS, envs::CTMRGEnv)
     return [
         collect(
             expectation_value(
-                peps, LocalOperator(lattice, (CartesianIndex(r, c),) => Sa), envs
+                peps, LocalOperator(lattice, (CartesianIndex(r, c),) => Sa), env
             ) for (r, c) in Iterators.product(1:N1, 1:N2)
         ) for Sa in Sas
     ]
@@ -41,11 +41,11 @@ end
 """
 Measure physical quantities for Heisenberg model
 """
-function measure_heis(peps::InfinitePEPS, H::LocalOperator, envs::CTMRGEnv)
+function measure_heis(peps::InfinitePEPS, H::LocalOperator, env::CTMRGEnv)
     results = Dict{String,Any}()
     N1, N2 = size(peps)
-    results["e_site"] = costfun(peps, envs, H) / (N1 * N2)
-    results["mag"] = cal_mags(peps, envs)
+    results["e_site"] = costfun(peps, env, H) / (N1 * N2)
+    results["mag"] = cal_mags(peps, env)
     results["mag_norm"] = collect(
         norm([mags[r, c] for mags in results["mag"]]) for
         (r, c) in Iterators.product(1:N1, 1:N2)

examples/heisenberg_evol/simpleupdate.jl

@@ -3,7 +3,7 @@ include("heis_tools.jl")
 # benchmark data is from Phys. Rev. B 94, 035133 (2016)
 Dbond, χenv, symm = 4, 16, Trivial
 trscheme_peps = truncerr(1e-10) & truncdim(Dbond)
-trscheme_envs = truncerr(1e-10) & truncdim(χenv)
+trscheme_env = truncerr(1e-10) & truncdim(χenv)
 N1, N2 = 2, 2
 # Heisenberg model Hamiltonian
 # (already only includes nearest neighbor terms)
@@ -41,10 +41,10 @@ for (dt, tol) in zip(dts, tols)
 end
 # measure physical quantities with CTMRG
 peps_ = InfinitePEPS(peps)
-envs = CTMRGEnv(rand, Float64, peps_, Espace)
-ctm_alg = SequentialCTMRG(; tol=1e-10, verbosity=2, trscheme=trscheme_envs)
-envs = leading_boundary(envs, peps_, ctm_alg)
-meas = measure_heis(peps_, ham, envs)
+env₀ = CTMRGEnv(rand, Float64, peps_, Espace)
+ctm_alg = SequentialCTMRG(; tol=1e-10, verbosity=2, trscheme=trscheme_env)
+env = leading_boundary(env₀, peps_, ctm_alg)
+meas = measure_heis(peps_, ham, env)
 display(meas)
 @info @sprintf("Energy = %.8f\n", meas["e_site"])
 @info @sprintf("Staggered magnetization = %.8f\n", mean(meas["mag_norm"]))

examples/hubbard_su.jl

@@ -40,10 +40,10 @@ peps = InfinitePEPS(peps)
 # CTMRG
 χenv0, χenv = 6, 20
 Espace = Vect[fℤ₂](0 => χenv0 / 2, 1 => χenv0 / 2)
-envs = CTMRGEnv(randn, Float64, peps, Espace)
+env = CTMRGEnv(randn, Float64, peps, Espace)
 for χ in [χenv0, χenv]
     ctm_alg = SequentialCTMRG(; maxiter=300, tol=1e-7)
-    envs = leading_boundary(envs, peps, ctm_alg)
+    env, = leading_boundary(env, 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, env, 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