diff --git a/Project.toml b/Project.toml
index 8b9d043d..1dc59d9a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -9,12 +9,14 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
+LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
 MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
+Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
+Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
 OhMyThreads = "67456a42-1dca-4109-a031-0a68de7e3ad5"
-OptimKit = "77e91f04-9b3b-57a6-a776-40b61faaebe0"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -29,12 +31,14 @@ ChainRulesCore = "1.0"
 Compat = "3.46, 4.2"
 FiniteDifferences = "0.12"
 KrylovKit = "0.8"
+LineSearches = "7.3.0"
 LinearAlgebra = "1"
 LoggingExtras = "1"
 MPSKit = "0.11"
 MPSKitModels = "0.3.5"
+Manifolds = "0.10.8"
+Manopt = "0.5.4"
 OhMyThreads = "0.7"
-OptimKit = "0.3"
 Printf = "1"
 QuadGK = "2.11.1"
 Random = "1"
diff --git a/README.md b/README.md
index 25f9a924..8dd2d139 100644
--- a/README.md
+++ b/README.md
@@ -54,7 +54,7 @@ opt_alg = PEPSOptimize(;
 
 # ground state search
 state = InfinitePEPS(2, D)
-ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
+ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
 result = fixedpoint(state, H, opt_alg, ctm)
 
 @show result.E # -0.6625...
diff --git a/docs/src/index.md b/docs/src/index.md
index 709f4222..61775310 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -36,7 +36,7 @@ opt_alg = PEPSOptimize(;
 
 # ground state search
 state = InfinitePEPS(2, D)
-ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
+ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
 result = fixedpoint(state, H, opt_alg, ctm)
 
 @show result.E # -0.6625...
diff --git a/examples/boundary_mps.jl b/examples/boundary_mps.jl
index 51af6dea..6e82091d 100644
--- a/examples/boundary_mps.jl
+++ b/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))
diff --git a/examples/heisenberg.jl b/examples/heisenberg.jl
index 34baed18..651f46aa 100644
--- a/examples/heisenberg.jl
+++ b/examples/heisenberg.jl
@@ -1,6 +1,6 @@
 using LinearAlgebra
-using TensorKit, OptimKit
-using PEPSKit, KrylovKit
+using TensorKit, KrylovKit, PEPSKit
+using Manopt
 
 # Square lattice Heisenberg Hamiltonian
 # We use the parameters (J₁, J₂, J₃) = (-1, 1, -1) by default to capture
@@ -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)
+env₀, = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg)
 result = fixedpoint(ψ₀, H, opt_alg, env₀)
 @show result.E
diff --git a/examples/hubbard_su.jl b/examples/hubbard_su.jl
index 6478a28e..d65512f5 100644
--- a/examples/hubbard_su.jl
+++ b/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
diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl
index 68ade133..9c894581 100644
--- a/src/PEPSKit.jl
+++ b/src/PEPSKit.jl
@@ -5,7 +5,8 @@ using Base: @kwdef
 using Compat
 using Accessors: @set
 using VectorInterface
-using TensorKit, KrylovKit, MPSKit, OptimKit, TensorOperations
+using TensorKit, KrylovKit, MPSKit, TensorOperations
+using TensorKit: ℂ, ℝ  # To avoid conflict with Manifolds
 using ChainRulesCore, Zygote
 using LoggingExtras
 using MPSKit: loginit!, logiter!, logfinish!, logcancel!
@@ -56,87 +57,153 @@ 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/manopt.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=2)
-        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
+- `optim_alg=quasi_Newton`: Manopt optimizer function
+- `optim_maxiter=100`: Maximal number of optimization iterations
+- `optim_tol=1e-4`: Gradient norm convergence tolerance
+- `stopping_criterion`: Manopt stopping criterion
+
+    ```
+    stopping_criterion = StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol)
+    ```
+
+- `record_group`: Values (`RecordAction`s) that are recorded during Manopt optimization
+
+    ```
+    record_group = [
+        RecordCost() => :Cost,
+        RecordGradientNorm() => :GradientNorm,
+        RecordTruncationError() => :TruncationError,
+        RecordConditionNumber() => :ConditionNumber,
+        RecordUnitCellGradientNorm() => :UnitcellGradientNorm,
+        RecordTime(; mode=:iterative) => :Time,
+    ]
+    ```
+
+- `debug_group = [DebugPEPSOptimize(), :Stop]`: Optimization iteration info printing
+- `stepsize = AdaptiveWNGradient()`:
+
+- `optim_kwargs`: All keyword arguments that are passed onto a Manopt optimization call
+
+    ```
+    optim_kwargs = (;
+        stopping_criterion, record=record_group, debug=debug_group, stepsize
+    )
+    ```
+
+- `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
+
+# OhMyThreads scheduler
+- `scheduler=Ref{Scheduler}(...)`: Multi-threading scheduler which can be accessed via `set_scheduler!`
 """
 module Defaults
-    using TensorKit, KrylovKit, OptimKit, OhMyThreads
+    using TensorKit, KrylovKit, OhMyThreads
+    using Manopt
     using PEPSKit:
         LinSolver,
         FixedSpaceTruncation,
         SVDAdjoint,
         HalfInfiniteProjector,
-        SimultaneousCTMRG
+        SimultaneousCTMRG,
+        RecordTruncationError,
+        RecordConditionNumber,
+        RecordUnitCellGradientNorm,
+        DebugPEPSOptimize
+
+    # 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=2)
-    const gradient_linsolver = KrylovKit.BiCGStab(;
-        maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
+
+    # Optimization
+    const optim_alg = quasi_Newton
+    const optim_maxiter = 100
+    const optim_tol = 1e-4
+    const stopping_criterion =
+        StopAfterIteration(optim_maxiter) | StopWhenGradientNormLess(optim_tol)
+    const record_group = [
+        RecordCost() => :Cost,
+        RecordGradientNorm() => :GradientNorm,
+        RecordTruncationError() => :TruncationError,
+        RecordConditionNumber() => :ConditionNumber,
+        RecordUnitCellGradientNorm() => :UnitcellGradientNorm,
+        RecordTime(; mode=:iterative) => :Time,
+    ]
+    const debug_group = [DebugPEPSOptimize(), :Stop]
+    const stepsize = AdaptiveWNGradient()
+    const optim_kwargs = (;
+        stopping_criterion, record=record_group, debug=debug_group, stepsize
     )
+    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
 
     # OhMyThreads scheduler defaults
     const scheduler = Ref{Scheduler}()
@@ -192,6 +259,8 @@ export LocalOperator
 export expectation_value, costfun, product_peps, correlation_length
 export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver
+export RecordTruncationError, RecordConditionNumber, RecordUnitCellGradientNorm
+export DebugPEPSOptimize
 export fixedpoint
 
 export absorb_weight
diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl
index c5225f9d..ec1cd8e5 100644
--- a/src/algorithms/ctmrg/ctmrg.jl
+++ b/src/algorithms/ctmrg/ctmrg.jl
@@ -38,10 +38,14 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm)
     env = deepcopy(envinit)
     log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG"))
 
+    truncation_error = 0.0
+    condition_number = 1.0
     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
+            truncation_error = info.truncation_error
+            condition_number = info.condition_number
             η, CS, TS = calc_convergence(env, CS, TS)
 
             if η ≤ alg.tol && iter ≥ alg.miniter
@@ -54,7 +58,7 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRGAlgorithm)
                 ctmrg_logiter!(log, iter, η, state, env)
             end
         end
-        return env
+        return env, (; truncation_error, condition_number)
     end
 end
 
diff --git a/src/algorithms/ctmrg/sequential.jl b/src/algorithms/ctmrg/sequential.jl
index 4404a7f5..8402aad0 100644
--- a/src/algorithms/ctmrg/sequential.jl
+++ b/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,10 +55,13 @@ 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])
@@ -64,10 +69,14 @@ function sequential_projectors(
             (WEST, r, c), state, envs, @set(alg.trscheme = trscheme)
         )
         ϵ[r] = info.err / norm(info.S)
+        S[r] = info.S
         return proj
     end
 
-    return (map(first, projectors), map(last, projectors)), (; err=maximum(copy(ϵ)))
+    truncation_error = maximum(copy(ϵ))
+    condition_number = maximum(_condition_number, copy(S))
+    info = (; truncation_error, condition_number)
+    return (map(first, projectors), map(last, projectors)), info
 end
 function sequential_projectors(
     coordinate::NTuple{3,Int},
diff --git a/src/algorithms/ctmrg/simultaneous.jl b/src/algorithms/ctmrg/simultaneous.jl
index 12d72960..31f30567 100644
--- a/src/algorithms/ctmrg/simultaneous.jl
+++ b/src/algorithms/ctmrg/simultaneous.jl
@@ -44,10 +44,20 @@ function _prealloc_svd(edges::Array{E,N}) where {E,N}
     Sc = scalartype(E)
     U = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, space(e)), edges))
     V = Zygote.Buffer(map(e -> TensorMap(zeros, Sc, domain(e), codomain(e)), edges))
-    S = Zygote.Buffer(U.data, tensormaptype(spacetype(E), 1, 1, real(Sc)))  # Corner type but with real numbers
+    S = Zygote.Buffer(edges, tensormaptype(spacetype(E), 1, 1, real(Sc)))  # Corner type but with real numbers
     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)
@@ -76,7 +86,11 @@ function simultaneous_projectors(
 
     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))
+    S = copy(S)
+    truncation_error = maximum(copy(ϵ))  # TODO: This makes Zygote error on first execution?
+    condition_number = maximum(_condition_number, S)
+    info = (; truncation_error, condition_number, U=copy(U), S, V=copy(V))
+    return (P_left, P_right), info
 end
 function simultaneous_projectors(
     coordinate, enlarged_corners::Array{E,3}, alg::HalfInfiniteProjector
diff --git a/src/algorithms/peps_opt.jl b/src/algorithms/optimization/fixed_point_differentiation.jl
similarity index 59%
rename from src/algorithms/peps_opt.jl
rename to src/algorithms/optimization/fixed_point_differentiation.jl
index 7be93c92..8d13981e 100644
--- a/src/algorithms/peps_opt.jl
+++ b/src/algorithms/optimization/fixed_point_differentiation.jl
@@ -76,137 +76,6 @@ function LinSolver(;
     return LinSolver{iterscheme}(solver)
 end
 
-"""
-    PEPSOptimize{G}(; boundary_alg=Defaults.ctmrg_alg, optimizer::OptimKit.OptimizationAlgorithm=Defaults.optimizer
-                    reuse_env::Bool=true, gradient_alg::G=Defaults.gradient_alg)
-
-Algorithm struct that represent PEPS ground-state optimization using AD.
-Set the algorithm to contract the infinite PEPS in `boundary_alg`;
-currently only `CTMRGAlgorithm`s are supported. The `optimizer` computes the gradient directions
-based on the CTMRG gradient and updates the PEPS parameters. In this optimization,
-the CTMRG runs can be started on the converged environments of the previous optimizer
-step by setting `reuse_env` to true. Otherwise a random environment is used at each
-step. The CTMRG gradient itself is computed using the `gradient_alg` algorithm.
-"""
-struct PEPSOptimize{G}
-    boundary_alg::CTMRGAlgorithm
-    optimizer::OptimKit.OptimizationAlgorithm
-    reuse_env::Bool
-    gradient_alg::G
-
-    function PEPSOptimize(  # Inner constructor to prohibit illegal setting combinations
-        boundary_alg::CTMRGAlgorithm,
-        optimizer,
-        reuse_env,
-        gradient_alg::G,
-    ) where {G}
-        if gradient_alg isa GradMode
-            if boundary_alg isa SequentialCTMRG && iterscheme(gradient_alg) === :fixed
-                throw(ArgumentError(":sequential and :fixed are not compatible"))
-            end
-        end
-        return new{G}(boundary_alg, optimizer, reuse_env, gradient_alg)
-    end
-end
-function PEPSOptimize(;
-    boundary_alg=Defaults.ctmrg_alg,
-    optimizer=Defaults.optimizer,
-    reuse_env=Defaults.reuse_env,
-    gradient_alg=Defaults.gradient_alg,
-)
-    return PEPSOptimize(boundary_alg, optimizer, reuse_env, gradient_alg)
-end
-
-"""
-    fixedpoint(ψ₀::InfinitePEPS{T}, H, alg::PEPSOptimize, [env₀::CTMRGEnv];
-               finalize!=OptimKit._finalize!, symmetrization=nothing) where {T}
-    
-Optimize `ψ₀` with respect to the Hamiltonian `H` 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.
-
-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
-"""
-function fixedpoint(
-    ψ₀::InfinitePEPS{F},
-    H,
-    alg::PEPSOptimize,
-    env₀::CTMRGEnv=CTMRGEnv(ψ₀, field(F)^20);
-    (finalize!)=OptimKit._finalize!,
-    symmetrization=nothing,
-) where {F}
-    if isnothing(symmetrization)
-        retract = peps_retract
-    else
-        retract, symm_finalize! = symmetrize_retract_and_finalize!(symmetrization)
-        fin! = finalize!  # Previous finalize!
-        finalize! = (x, f, g, numiter) -> fin!(symm_finalize!(x, f, g, numiter)..., numiter)
-    end
-
-    if scalartype(env₀) <: Real
-        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(
-        (ψ₀, env₀), alg.optimizer; retract, inner=real_inner, finalize!
-    ) do (peps, envs)
-        E, gs = withgradient(peps) do ψ
-            envs´ = hook_pullback(
-                leading_boundary,
-                envs,
-                ψ,
-                alg.boundary_alg;
-                alg_rrule=alg.gradient_alg,
-            )
-            ignore_derivatives() do
-                alg.reuse_env && update!(envs, envs´)
-            end
-            return costfun(ψ, envs´, H)
-        end
-        g = only(gs)  # `withgradient` returns tuple of gradients `gs`
-        return E, g
-    end
-
-    return (;
-        peps,
-        env,
-        E,
-        E_history=convhistory[:, 1],
-        grad=∂E,
-        gradnorm_history=convhistory[:, 2],
-        numfg,
-    )
-end
-
-# Update PEPS unit cell in non-mutating way
-# Note: Both x and η are InfinitePEPS during optimization
-function peps_retract(x, η, α)
-    peps = deepcopy(x[1])
-    peps.A .+= η.A .* α
-    env = deepcopy(x[2])
-    return (peps, env), η
-end
-
-# Take real valued part of dot product
-real_inner(_, η₁, η₂) = real(dot(η₁, η₂))
-
 #=
 Evaluating the gradient of the cost function for CTMRG:
 - The gradient of the cost function for CTMRG can be computed using automatic differentiation (AD) or explicit evaluation of the geometric sum.
@@ -222,9 +91,9 @@ function _rrule(
     state,
     alg::CTMRGAlgorithm,
 )
-    envs = leading_boundary(envinit, state, alg)
+    envs, info = leading_boundary(envinit, state, alg)
 
-    function leading_boundary_diffgauge_pullback(Δenvs′)
+    function leading_boundary_diffgauge_pullback((Δenvs′, Δinfo))
         Δenvs = unthunk(Δenvs′)
 
         # find partial gradients of gauge_fixed single CTMRG iteration
@@ -239,7 +108,7 @@ function _rrule(
         return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent()
     end
 
-    return envs, leading_boundary_diffgauge_pullback
+    return (envs, info), leading_boundary_diffgauge_pullback
 end
 
 # Here f is differentiated from an pre-computed SVD with fixed U, S and V
@@ -252,9 +121,9 @@ function _rrule(
     alg::SimultaneousCTMRG,
 )
     @assert !isnothing(alg.projector_alg.svd_alg.rrule_alg)
-    envs = leading_boundary(envinit, state, alg)
-    envsconv, info = ctmrg_iteration(state, envs, alg)
-    envs_fixed, signs = gauge_fix(envs, envsconv)
+    envs, = leading_boundary(envinit, state, alg)
+    envs_conv, info = ctmrg_iteration(state, envs, alg)
+    envs_fixed, signs = gauge_fix(envs, envs_conv)
 
     # Fix SVD
     Ufix, Vfix = fix_relative_phases(info.U, info.V, signs)
@@ -264,7 +133,7 @@ function _rrule(
     alg_fixed = @set alg.projector_alg.svd_alg = svd_alg_fixed
     alg_fixed = @set alg_fixed.projector_alg.trscheme = notrunc()
 
-    function leading_boundary_fixed_pullback(Δenvs′)
+    function leading_boundary_fixed_pullback((Δenvs′, Δinfo))
         Δenvs = unthunk(Δenvs′)
 
         f(A, x) = fix_global_phases(x, ctmrg_iteration(A, x, alg_fixed)[1])
@@ -278,7 +147,7 @@ function _rrule(
         return NoTangent(), ZeroTangent(), ∂F∂envs, NoTangent()
     end
 
-    return envs_fixed, leading_boundary_fixed_pullback
+    return (envs_fixed, info), leading_boundary_fixed_pullback
 end
 
 @doc """
@@ -352,4 +221,4 @@ function fpgrad(∂F∂x, ∂f∂x, ∂f∂A, y₀, alg::LinSolver)
     end
 
     return ∂f∂A(y)
-end
+end
\ No newline at end of file
diff --git a/src/algorithms/optimization/manopt.jl b/src/algorithms/optimization/manopt.jl
new file mode 100644
index 00000000..142f4640
--- /dev/null
+++ b/src/algorithms/optimization/manopt.jl
@@ -0,0 +1,116 @@
+using Printf
+using Manifolds: Manifolds
+using Manifolds:
+    AbstractManifold,
+    AbstractRetractionMethod,
+    Euclidean,
+    default_retraction_method,
+    retract
+using Manopt
+
+"""
+    mutable struct RecordTruncationError <: RecordAction
+
+Record the maximal truncation error of all `boundary_alg` runs of the corresponding
+optimization iteration.
+"""
+mutable struct RecordTruncationError <: RecordAction
+    recorded_values::Vector{Float64}
+    RecordTruncationError() = new(Vector{Float64}())
+end
+function (r::RecordTruncationError)(
+    p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int
+)
+    cache = Manopt.get_cost_function(get_objective(p))
+    return Manopt.record_or_reset!(r, cache.env_info.truncation_error, i)
+end
+
+"""
+    mutable struct RecordConditionNumber <: RecordAction
+
+Record the maximal condition number of all `boundary_alg` runs of the corresponding
+optimization iteration.
+"""
+mutable struct RecordConditionNumber <: RecordAction
+    recorded_values::Vector{Float64}
+    RecordConditionNumber() = new(Vector{Float64}())
+end
+function (r::RecordConditionNumber)(
+    p::AbstractManoptProblem, ::AbstractManoptSolverState, i::Int
+)
+    cache = Manopt.get_cost_function(get_objective(p))
+    return Manopt.record_or_reset!(r, cache.env_info.condition_number, i)
+end
+
+"""
+    mutable struct RecordUnitCellGradientNorm <: RecordAction
+        
+Record the PEPS gradient norms unit cell entry-wise, i.e. an array 
+of norms `norm.(peps.A)`.
+"""
+mutable struct RecordUnitCellGradientNorm <: RecordAction
+    recorded_values::Vector{Matrix{Float64}}
+    RecordUnitCellGradientNorm() = new(Vector{Matrix{Float64}}())
+end
+function (r::RecordUnitCellGradientNorm)(
+    p::AbstractManoptProblem, s::AbstractManoptSolverState, i::Int
+)
+    cache = Manopt.get_cost_function(get_objective(p))
+    grad = cache.from_vec(get_gradient(s))
+    return Manopt.record_or_reset!(r, norm.(grad.A), i)
+end
+
+"""
+    mutable struct DebugPEPSOptimize <: DebugAction
+    
+Custom `DebugAction` printing for PEPS optimization runs.
+
+The debug info is output using `@info` and by default prints the optimization iteration,
+the cost function value, the gradient norm, the last step size and the time elapsed during
+the current iteration.
+"""
+mutable struct DebugPEPSOptimize <: DebugAction
+    last_time::UInt64
+    DebugPEPSOptimize() = new(time_ns())
+end
+function (d::DebugPEPSOptimize)(
+    p::AbstractManoptProblem, s::AbstractManoptSolverState, k::Int
+)
+    time_new = time_ns()
+    cost = get_cost(p, get_iterate(s))
+    if k == 0
+        @info @sprintf("Initial f(x) = %.8f", cost)
+    elseif k > 0
+        gradient_norm = norm(get_manifold(p), get_iterate(s), get_gradient(s))
+        stepsize = get_last_stepsize(p, s, k)
+        time = (time_new - d.last_time) * 1e-9
+        @info @sprintf(
+            "Optimization %d: f(x) = %.8f   ‖∂f‖ = %.8f   step = %.4f   time = %.2f sec",
+            k,
+            cost,
+            gradient_norm,
+            stepsize,
+            time
+        )
+    end
+    d.last_time = time_new  # update time stamp
+    return nothing
+end
+
+"""
+    SymmetrizeExponentialRetraction <: AbstractRetractionMethod
+    
+Exponential retraction followed by a symmetrization step.
+"""
+struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod
+    symmetrization::SymmetrizationStyle
+    from_vec::Function
+end
+
+function Manifolds.retract!(
+    M::Euclidean, p, q, X, t::Number, sr::SymmetrizeExponentialRetraction
+)
+    v = Manifolds.retract!(M, p, q, X, t)
+    v_symm_peps = symmetrize!(sr.from_vec(v), sr.symmetrization)
+    return to_vec(v_symm_peps)
+end
diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl
new file mode 100644
index 00000000..93bdd3b8
--- /dev/null
+++ b/src/algorithms/optimization/peps_optimization.jl
@@ -0,0 +1,245 @@
+"""
+    struct PEPSOptimize{G}
+
+Algorithm struct for PEPS optimization using automatic differentiation.
+
+# Fields
+- `boundary_alg::CTMRGAlgorithm`: algorithm for determining the PEPS environment
+- `optim_alg::Function`: Manopt optimization algorithm
+- `optim_kwargs::NamedTuple`: Keyword arguments provided to the Manopt `optim_alg` call
+- `gradient_alg::G`: Algorithm computing the cost function gradient in reverse-mode
+- `reuse_env::Bool`: If `true` the previous environment is used to initialize the next
+    `leading_boundary` call
+- `symmetrization::Union{Nothing,SymmetrizationStyle}`: Symmetrize the PEPS and PEPS
+    gradient after each optimization iteration (does nothing if `nothing` is provided)
+"""
+struct PEPSOptimize{G}
+    boundary_alg::CTMRGAlgorithm
+    optim_alg::Function
+    optim_kwargs::NamedTuple
+    gradient_alg::G
+    reuse_env::Bool
+    # reuse_env_tol::Float64  # TODO: add option for reuse tolerance
+    symmetrization::Union{Nothing,SymmetrizationStyle}
+
+    function PEPSOptimize(  # Inner constructor to prohibit illegal setting combinations
+        boundary_alg::CTMRGAlgorithm,
+        optim_alg,
+        optim_kwargs,
+        gradient_alg::G,
+        reuse_env,
+        symmetrization,
+    ) where {G}
+        if gradient_alg isa GradMode
+            if boundary_alg isa SequentialCTMRG && iterscheme(gradient_alg) === :fixed
+                throw(ArgumentError(":sequential and :fixed are not compatible"))
+            end
+        end
+        return new{G}(
+            boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization
+        )
+    end
+end
+
+"""
+    PEPSOptimize(;
+        boundary_alg=Defaults.ctmrg_alg,
+        optim_alg=Defaults.optim_alg,
+        maxiter=Defaults.optim_maxiter,
+        tol=Defaults.optim_tol,
+        gradient_alg=Defaults.gradient_alg,
+        reuse_env=Defaults.reuse_env,
+        symmetrization=nothing,
+        kwargs...,
+    )
+
+Convenience keyword argument constructor for `PEPSOptimize` algorithms.
+Here, `maxiter` and `tol` are passed onto `StopAfterIteration` and `StopWhenGradientNormLess`
+stopping criteria, respectively. Additionally, any keyword arguments can be provided which 
+are then stored inside `optim_kwargs` and passed to the Manopt optimization call, such that
+that all arguments of the respective `optim_alg` can be used.
+"""
+function PEPSOptimize(;
+    boundary_alg=Defaults.ctmrg_alg,
+    optim_alg=Defaults.optim_alg,
+    maxiter=Defaults.optim_maxiter,
+    tol=Defaults.optim_tol,
+    gradient_alg=Defaults.gradient_alg,
+    reuse_env=Defaults.reuse_env,
+    symmetrization=nothing,
+    kwargs...,
+)
+    stopping_criterion = StopAfterIteration(maxiter) | StopWhenGradientNormLess(tol)
+    optim_kwargs = merge(Defaults.optim_kwargs, (; stopping_criterion, kwargs...))
+    return PEPSOptimize(
+        boundary_alg, optim_alg, optim_kwargs, gradient_alg, reuse_env, symmetrization
+    )
+end
+
+"""
+    mutable struct PEPSCostFunctionCache{T}
+
+Stores objects used for computing PEPS cost functions during optimization that are
+needed apart from the PEPS that is being optimized.
+
+# Fields
+- `operator::LocalOperator`: cost function operator
+- `alg::PEPSOptimize`: optimization parameters
+- `env::CTMRGEnv`: environment of the current PEPS
+- `from_vec`: map which returns vectorized PEPS as an `InfinitePEPS`
+- `peps_vec::Vector{T}`: current vectorized PEPS
+- `grad_vec::Vector{T}`: current vectorized gradient
+- `cost::Float64`: current cost function value
+- `env_info::NamedTuple`: return info of `leading_boundary` used by `RecordAction`s
+"""
+mutable struct PEPSCostFunctionCache{T}
+    operator::LocalOperator
+    alg::PEPSOptimize
+    env::CTMRGEnv
+    from_vec
+    peps_vec::Vector{T}
+    grad_vec::Vector{T}
+    cost::Float64
+    env_info::NamedTuple
+end
+
+"""
+    PEPSCostFunctionCache(
+        operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv
+    )
+
+Initialize a `PEPSCostFunctionCache` using `peps_vec` from which the vector dimension
+and scalartype are derived.
+"""
+function PEPSCostFunctionCache(
+    operator::LocalOperator, alg::PEPSOptimize, peps_vec::Vector, from_vec, env::CTMRGEnv
+)
+    return PEPSCostFunctionCache(
+        operator,
+        alg,
+        env,
+        from_vec,
+        similar(peps_vec),
+        similar(peps_vec),
+        0.0,
+        (; truncation_error=0.0, condition_number=1.0),
+    )
+end
+
+"""
+    cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T}
+
+Update the cost and gradient of the `PEPSCostFunctionCache` with respect to the new point
+`peps_vec`.
+"""
+function cost_and_grad!(cache::PEPSCostFunctionCache{T}, peps_vec::Vector{T}) where {T}
+    cache.peps_vec .= peps_vec  # update point in manifold
+    peps = cache.from_vec(peps_vec)  # convert back to InfinitePEPS
+    env₀ =
+        cache.alg.reuse_env ? cache.env : CTMRGEnv(randn, scalartype(cache.env), cache.env)
+
+    # compute cost and gradient
+    cost, grads = withgradient(peps) do ψ
+        env, info = hook_pullback(
+            leading_boundary,
+            env₀,
+            ψ,
+            cache.alg.boundary_alg;
+            alg_rrule=cache.alg.gradient_alg,
+        )
+        cost = expectation_value(ψ, cache.operator, env)
+        ignore_derivatives() do
+            update!(cache.env, env)  # update environment in-place
+            cache.env_info = info  # update environment information (truncation error, ...)
+            isapprox(imag(cost), 0; atol=sqrt(eps(real(cost)))) ||
+                @warn "Expectation value is not real: $cost."
+        end
+        return real(cost)
+    end
+    grad = only(grads)  # `withgradient` returns tuple of gradients `grads`
+
+    # symmetrize gradient
+    if !isnothing(cache.alg.symmetrization)
+        grad = symmetrize!(grad, cache.alg.symmetrization)
+    end
+
+    cache.cost = cost  # update cost function value
+    cache.grad_vec .= to_vec(grad)[1]  # update vectorized gradient
+    return cache.cost, cache.grad_vec
+end
+
+"""
+    (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T}
+    
+Return the cost of `cache` and recompute if `peps_vec` is a new point.
+"""
+function (cache::PEPSCostFunctionCache{T})(::Euclidean, peps_vec::Vector{T}) where {T}
+    # Note that it is necessary to implement the cost function as a functor so that the
+    # `PEPSCostFunctionCache` is available through `get_objective(::AbstractManoptProblem)`
+    # to the `RecordAction`s during optimization
+    if !(peps_vec == cache.peps_vec) # update cache if at new point
+        cost_and_grad!(cache, peps_vec)
+    end
+    return cache.cost
+end
+
+"""
+    gradient_function(cache::PEPSCostFunctionCache{T}) where {T}
+
+Get the gradient function of `cache` which returns the gradient vector
+and recomputes it if the provided point is new.
+"""
+function gradient_function(cache::PEPSCostFunctionCache{T}) where {T}
+    return function gradient_function(::Euclidean, peps_vec::Vector{T})
+        if !(peps_vec == cache.peps_vec) # update cache if at new point
+            cost_and_grad!(cache, peps_vec)
+        end
+        return cache.grad_vec
+    end
+end
+
+"""
+    fixedpoint(
+        peps₀::InfinitePEPS{T},
+        operator::LocalOperator,
+        alg::PEPSOptimize,
+        env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20);
+    ) where {T}
+
+Optimize a PEPS starting from `peps₀` and `env₀` by minimizing the cost function given
+by expectation value of `operator`. All optimization parameters are provided through `alg`.
+"""
+function fixedpoint(
+    peps₀::InfinitePEPS{T},
+    operator::LocalOperator,
+    alg::PEPSOptimize,
+    env₀::CTMRGEnv=CTMRGEnv(peps₀, field(T)^20);
+) where {T}
+    if scalartype(env₀) <: Real
+        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
+
+    # construct cost and grad functions
+    peps₀_vec, from_vec = to_vec(peps₀)
+    cache = PEPSCostFunctionCache(operator, alg, peps₀_vec, from_vec, deepcopy(env₀))
+    cost = cache
+    grad = gradient_function(cache)
+
+    # optimize
+    M = Euclidean(length(peps₀_vec))
+    retraction_method = if isnothing(alg.symmetrization)
+        default_retraction_method(M)
+    else
+        SymmetrizeExponentialRetraction(alg.symmetrization, from_vec)
+    end
+    result = alg.optim_alg(
+        M, cost, grad, peps₀_vec; alg.optim_kwargs..., retraction_method, return_state=true
+    )
+
+    # extract final result
+    peps_final = from_vec(get_solver_result(result))
+    return peps_final, cost.env, cache.cost, result
+end
diff --git a/src/algorithms/toolbox.jl b/src/algorithms/toolbox.jl
index d740ab1d..7dffed19 100644
--- a/src/algorithms/toolbox.jl
+++ b/src/algorithms/toolbox.jl
@@ -37,15 +37,6 @@ function MPSKit.expectation_value(
     return expectation_value(pf, CartesianIndex(op[1]) => op[2], envs)
 end
 
-function costfun(peps::InfinitePEPS, envs::CTMRGEnv, O::LocalOperator)
-    E = MPSKit.expectation_value(peps, O, envs)
-    ignore_derivatives() do
-        isapprox(imag(E), 0; atol=sqrt(eps(real(E)))) ||
-            @warn "Expectation value is not real: $E."
-    end
-    return real(E)
-end
-
 function LinearAlgebra.norm(peps::InfinitePEPS, env::CTMRGEnv)
     total = one(scalartype(peps))
 
diff --git a/src/environments/ctmrg_environments.jl b/src/environments/ctmrg_environments.jl
index 4b73e98c..26768675 100644
--- a/src/environments/ctmrg_environments.jl
+++ b/src/environments/ctmrg_environments.jl
@@ -449,6 +449,25 @@ end
 @non_differentiable CTMRGEnv(peps::InfinitePEPS, args...)
 @non_differentiable CTMRGEnv(peps::InfinitePartitionFunction, args...)
 
+"""
+    CTMRGEnv([f=randn, T=scalartype(env)], env::CTMRGEnv)
+
+Copy constructor for easily initializing a similar environment, optionally providing a
+function `f` for element generation with type `T`.
+"""
+function CTMRGEnv(env::CTMRGEnv)
+    return CTMRTEnv(randn, scalartype(env), env)
+end
+function CTMRGEnv(f, T, env::CTMRGEnv)
+    corners = map(env.corners) do corner
+        return TensorMap(f, T, space(corner))
+    end
+    edges = map(env.edges) do edge
+        return TensorMap(f, T, space(edge))
+    end
+    return CTMRGEnv(corners, edges)
+end
+
 # Custom adjoint for CTMRGEnv constructor, needed for fixed-point differentiation
 function ChainRulesCore.rrule(::Type{CTMRGEnv}, corners, edges)
     ctmrgenv_pullback(ē) = NoTangent(), ē.corners, ē.edges
diff --git a/src/states/infinitepeps.jl b/src/states/infinitepeps.jl
index 29409014..9243fa1e 100644
--- a/src/states/infinitepeps.jl
+++ b/src/states/infinitepeps.jl
@@ -110,7 +110,9 @@ function InfinitePEPS(
 end
 
 unitcell(t::InfinitePEPS) = t.A
+
 TensorKit.space(t::InfinitePEPS, i, j) = space(t[i, j], 1)
+TensorKit.dim(t::InfinitePEPS) = sum(dim.(t.A))
 
 # Chainrules
 function ChainRulesCore.rrule(
diff --git a/src/utility/symmetrization.jl b/src/utility/symmetrization.jl
index 5bec6a96..abfdb474 100644
--- a/src/utility/symmetrization.jl
+++ b/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
diff --git a/src/utility/util.jl b/src/utility/util.jl
index 075f39ea..402880ad 100644
--- a/src/utility/util.jl
+++ b/src/utility/util.jl
@@ -56,10 +56,7 @@ function sdiag_pow(S::AbstractTensorMap, pow::Real; tol::Real=eps(scalartype(S))
     tol *= norm(S, Inf)  # Relative tol w.r.t. largest singular value (use norm(∘, Inf) to make differentiable)
     Spow = similar(S)
     for (k, b) in blocks(S)
-        copyto!(
-            blocks(Spow)[k],
-            LinearAlgebra.diagm(_safe_pow.(LinearAlgebra.diag(b), pow, tol)),
-        )
+        copyto!(blocks(Spow)[k], diagm(_safe_pow.(diag(b), pow, tol)))
     end
     return Spow
 end
diff --git a/test/boundarymps/vumps.jl b/test/boundarymps/vumps.jl
index 2ec6265c..efe544f2 100644
--- a/test/boundarymps/vumps.jl
+++ b/test/boundarymps/vumps.jl
@@ -17,7 +17,7 @@ const vumps_alg = VUMPS(; alg_eigsolve=MPSKit.Defaults.alg_eigsolve(; ishermitia
     mps, envs, ϵ = leading_boundary(mps, T, vumps_alg)
     N = abs(sum(expectation_value(mps, T)))
 
-    ctm = leading_boundary(
+    ctm, = leading_boundary(
         CTMRGEnv(psi, ComplexSpace(20)), psi, SimultaneousCTMRG(; verbosity=1)
     )
     N´ = abs(norm(psi, ctm))
@@ -33,7 +33,7 @@ end
     mps, envs, ϵ = leading_boundary(mps, T, vumps_alg)
     N = abs(prod(expectation_value(mps, T)))
 
-    ctm = leading_boundary(
+    ctm, = leading_boundary(
         CTMRGEnv(psi, ComplexSpace(20)), psi, SimultaneousCTMRG(; verbosity=1)
     )
     N´ = abs(norm(psi, ctm))
diff --git a/test/ctmrg/fixed_iterscheme.jl b/test/ctmrg/fixed_iterscheme.jl
index 208fdece..39eccf5c 100644
--- a/test/ctmrg/fixed_iterscheme.jl
+++ b/test/ctmrg/fixed_iterscheme.jl
@@ -30,7 +30,7 @@ atol = 1e-5
     # initialize states
     Random.seed!(2394823842)
     psi = InfinitePEPS(2, χbond; unitcell)
-    env_conv1 = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg)
+    env_conv1, = leading_boundary(CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg)
 
     # do extra iteration to get SVD
     env_conv2, info = ctmrg_iteration(psi, env_conv1, ctm_alg)
@@ -61,7 +61,7 @@ end
     Random.seed!(91283219347)
     psi = InfinitePEPS(2, χbond)
     env_init = CTMRGEnv(psi, ComplexSpace(χenv))
-    env_conv1 = leading_boundary(env_init, psi, ctm_alg_iter)
+    env_conv1, = leading_boundary(env_init, psi, ctm_alg_iter)
 
     # do extra iteration to get SVD
     env_conv2_iter, info_iter = ctmrg_iteration(psi, env_conv1, ctm_alg_iter)
diff --git a/test/ctmrg/flavors.jl b/test/ctmrg/flavors.jl
index 5d8b440c..1c84f7e9 100644
--- a/test/ctmrg/flavors.jl
+++ b/test/ctmrg/flavors.jl
@@ -19,10 +19,10 @@ projector_algs = [HalfInfiniteProjector, FullInfiniteProjector]
     # compute environments
     Random.seed!(32350283290358)
     psi = InfinitePEPS(2, χbond; unitcell)
-    env_sequential = leading_boundary(
+    env_sequential, = leading_boundary(
         CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg_sequential
     )
-    env_simultaneous = leading_boundary(
+    env_simultaneous, = leading_boundary(
         CTMRGEnv(psi, ComplexSpace(χenv)), psi, ctm_alg_simultaneous
     )
 
@@ -68,7 +68,7 @@ end
     χs = [16 17 18; 15 20 21; 14 19 22]
     psi = InfinitePEPS(Ds, Ds, Ds)
     env = CTMRGEnv(psi, rand(10:20, 3, 3), rand(10:20, 3, 3))
-    env2 = leading_boundary(env, psi, ctm_alg)
+    env2, = leading_boundary(env, psi, ctm_alg)
 
     # check that the space is fixed
     @test all(space.(env.corners) .== space.(env2.corners))
diff --git a/test/ctmrg/gaugefix.jl b/test/ctmrg/gaugefix.jl
index 13d9a15b..52231498 100644
--- a/test/ctmrg/gaugefix.jl
+++ b/test/ctmrg/gaugefix.jl
@@ -20,7 +20,7 @@ function _pre_converge_env(
     Random.seed!(seed)  # Seed RNG to make random environment consistent
     psi = InfinitePEPS(rand, T, physical_space, peps_space; unitcell)
     env₀ = CTMRGEnv(psi, ctm_space)
-    env_conv = leading_boundary(env₀, psi, SequentialCTMRG(; tol))
+    env_conv, = leading_boundary(env₀, psi, SequentialCTMRG(; tol))
     return env_conv, psi
 end
 
@@ -45,7 +45,7 @@ end
     alg = ctmrg_alg(; tol, projector_alg)
     env_pre, psi = preconv[(S, T, unitcell)]
     env_pre
-    env = leading_boundary(env_pre, psi, alg)
+    env, = leading_boundary(env_pre, psi, alg)
     env′, = ctmrg_iteration(psi, env, alg)
     env_fixed, = gauge_fix(env, env′)
     @test calc_elementwise_convergence(env, env_fixed) ≈ 0 atol = atol
diff --git a/test/ctmrg/gradients.jl b/test/ctmrg/gradients.jl
index 648d48bd..0fa7c419 100644
--- a/test/ctmrg/gradients.jl
+++ b/test/ctmrg/gradients.jl
@@ -1,10 +1,10 @@
 using Test
 using Random
-using PEPSKit
 using TensorKit
-using Zygote
-using OptimKit
 using KrylovKit
+using PEPSKit
+using PEPSKit: to_vec, PEPSCostFunctionCache, gradient_function
+using Manopt, Manifolds
 
 ## Test models, gradmodes and CTMRG algorithm
 # -------------------------------------------
@@ -41,7 +41,6 @@ gradmodes = [
         LinSolver(; solver=KrylovKit.BiCGStab(; tol=gradtol), iterscheme=:diffgauge),
     ],
 ]
-steps = -0.01:0.005:0.01
 
 ## Tests
 # ------
@@ -55,27 +54,19 @@ steps = -0.01:0.005:0.01
     gms = gradmodes[i]
     calgs = ctmrg_algs[i]
     psi_init = InfinitePEPS(Pspace, Vspace, Vspace)
-    @testset "$ctmrg_alg and $alg_rrule" for (ctmrg_alg, alg_rrule) in
-                                             Iterators.product(calgs, gms)
-        @info "optimtest of $ctmrg_alg and $alg_rrule on $(names[i])"
+    @testset "$ctmrg_alg and $gradient_alg" for (ctmrg_alg, gradient_alg) in
+                                                Iterators.product(calgs, gms)
+        @info "gradient check of $ctmrg_alg and $alg_rrule on $(names[i])"
         Random.seed!(42039482030)
-        dir = InfinitePEPS(Pspace, Vspace, Vspace)
         psi = InfinitePEPS(Pspace, Vspace, Vspace)
-        env = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg)
-        alphas, fs, dfs1, dfs2 = OptimKit.optimtest(
-            (psi, env),
-            dir;
-            alpha=steps,
-            retract=PEPSKit.peps_retract,
-            inner=PEPSKit.real_inner,
-        ) do (peps, envs)
-            E, g = Zygote.withgradient(peps) do psi
-                envs2 = PEPSKit.hook_pullback(leading_boundary, envs, psi, ctmrg_alg; alg_rrule)
-                return costfun(psi, envs2, models[i])
-            end
+        env, = leading_boundary(CTMRGEnv(psi, Espace), psi, ctmrg_alg)
 
-            return E, only(g)
-        end
-        @test dfs1 ≈ dfs2 atol = 1e-2
+        psi_vec, from_vec = to_vec(psi)
+        opt_alg = PEPSOptimize(; boundary_alg=ctmrg_alg, gradient_alg)
+        cache = PEPSCostFunctionCache(models[i], opt_alg, psi_vec, from_vec, deepcopy(env))
+        cost = cache
+        grad = gradient_function(cache)
+        M = Euclidean(length(psi_vec))
+        @test check_gradient(M, cost, grad; N=5, exactness_tol=1e-4, io=stdout)
     end
 end
diff --git a/test/heisenberg.jl b/test/heisenberg.jl
index 5865ba46..a7df0db3 100644
--- a/test/heisenberg.jl
+++ b/test/heisenberg.jl
@@ -1,18 +1,15 @@
 using Test
 using Random
 using Accessors
-using PEPSKit
 using TensorKit
 using KrylovKit
-using OptimKit
+using PEPSKit
 
 # initialize parameters
 Dbond = 2
 χenv = 16
 ctm_alg = SimultaneousCTMRG()
-opt_alg = PEPSOptimize(;
-    boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2)
-)
+opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, tol=1e-3)
 # compare against Juraj Hasik's data:
 # https://github.com/jurajHasik/j1j2_ipeps_states/blob/main/single-site_pg-C4v-A1/j20.0/state_1s_A1_j20.0_D2_chi_opt48.dat
 E_ref = -0.6602310934799577
@@ -22,13 +19,13 @@ E_ref = -0.6602310934799577
     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)
+    env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
 
     # optimize energy and compute correlation lengths
-    result = fixedpoint(psi_init, H, opt_alg, env_init)
-    ξ_h, ξ_v, = correlation_length(result.peps, result.env)
+    peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init)
+    ξ_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
 
@@ -36,18 +33,16 @@ end
     # initialize states
     Random.seed!(456)
     unitcell = (1, 2)
-    H_1x2 = heisenberg_XYZ(InfiniteSquare(unitcell...))
-    psi_init_1x2 = InfinitePEPS(2, Dbond; unitcell)
-    env_init_1x2 = leading_boundary(
-        CTMRGEnv(psi_init_1x2, ComplexSpace(χenv)), psi_init_1x2, ctm_alg
-    )
+    H = heisenberg_XYZ(InfiniteSquare(unitcell...))
+    psi_init = InfinitePEPS(2, Dbond; unitcell)
+    env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
 
     # optimize energy and compute correlation lengths
-    result_1x2 = fixedpoint(psi_init_1x2, H_1x2, opt_alg, env_init_1x2)
-    ξ_h_1x2, ξ_v_1x2, = correlation_length(result_1x2.peps, result_1x2.env)
+    peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init)
+    ξ_h, ξ_v, = correlation_length(peps, env)
 
-    @test result_1x2.E ≈ 2 * E_ref atol = 1e-2
-    @test all(@. ξ_h_1x2 > 0 && ξ_v_1x2 > 0)
+    @test E ≈ 2 * E_ref atol = 1e-2
+    @test all(@. ξ_h > 0 && ξ_v > 0)
 end
 
 @testset "Simple update into AD optimization" begin
@@ -83,7 +78,7 @@ end
     # absorb weight into site tensors and CTMRG
     peps = InfinitePEPS(peps)
     envs₀ = CTMRGEnv(rand, Float64, peps, Espace)
-    envs = leading_boundary(envs₀, peps, SimultaneousCTMRG())
+    envs, = leading_boundary(envs₀, peps, SimultaneousCTMRG())
 
     # measure physical quantities
     e_site = costfun(peps, envs, ham) / (N1 * N2)
diff --git a/test/j1j2_model.jl b/test/j1j2_model.jl
index 8dd88397..f1c405ba 100644
--- a/test/j1j2_model.jl
+++ b/test/j1j2_model.jl
@@ -1,9 +1,8 @@
 using Test
 using Random
-using PEPSKit
 using TensorKit
 using KrylovKit
-using OptimKit
+using PEPSKit
 
 # initialize parameters
 χbond = 2
@@ -11,8 +10,9 @@ using OptimKit
 ctm_alg = SimultaneousCTMRG()
 opt_alg = PEPSOptimize(;
     boundary_alg=ctm_alg,
-    optimizer=LBFGS(4; gradtol=1e-3, verbosity=2),
+    tol=1e-3,
     gradient_alg=LinSolver(; iterscheme=:diffgauge),
+    symmetrization=RotateReflect(),
 )
 
 # initialize states
@@ -20,10 +20,10 @@ 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);
+env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg);
 
 # find fixedpoint
-result = fixedpoint(psi_init, H, opt_alg, env_init; symmetrization=RotateReflect())
+result = fixedpoint(psi_init, H, opt_alg, env_init)
 ξ_h, ξ_v, = correlation_length(result.peps, result.env)
 
 # compare against Juraj Hasik's data:
diff --git a/test/pwave.jl b/test/pwave.jl
index dcdecf7e..fd33a944 100644
--- a/test/pwave.jl
+++ b/test/pwave.jl
@@ -1,30 +1,35 @@
 using Test
 using Random
-using PEPSKit
 using TensorKit
 using KrylovKit
-using OptimKit
+using PEPSKit
+using Manopt
+using LineSearches
 
 # Initialize parameters
 unitcell = (2, 2)
 H = pwave_superconductor(InfiniteSquare(unitcell...))
-χbond = 2
+Dbond = 2
 χenv = 16
-ctm_alg = SimultaneousCTMRG(; maxiter=150)
+ctm_alg = SimultaneousCTMRG()
 opt_alg = PEPSOptimize(;
-    boundary_alg=ctm_alg, optimizer=LBFGS(4; maxiter=10, gradtol=1e-3, verbosity=2)
+    boundary_alg=ctm_alg,
+    maxiter=10,
+    gradient_alg=LinSolver(; iterscheme=:diffgauge),
+    stepsize=ConstantLength(1.5),
+    memory_size=4,
 )
 
 # initialize states
 Pspace = Vect[FermionParity](0 => 1, 1 => 1)
-Vspace = Vect[FermionParity](0 => χbond ÷ 2, 1 => χbond ÷ 2)
+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);
+env_init, = leading_boundary(CTMRGEnv(psi_init, Envspace), psi_init, ctm_alg);
 
 # find fixedpoint
-result = fixedpoint(psi_init, H, opt_alg, env_init)
+peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init)
 
 # 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(psi_init)) ≈ -2.6241 atol = 5e-2
diff --git a/test/tf_ising.jl b/test/tf_ising.jl
index 8b84696f..aa4181e8 100644
--- a/test/tf_ising.jl
+++ b/test/tf_ising.jl
@@ -1,9 +1,8 @@
 using Test
 using Random
-using PEPSKit
 using TensorKit
 using KrylovKit
-using OptimKit
+using PEPSKit
 
 # References
 # ----------
@@ -20,32 +19,30 @@ mˣ = 0.91
 χbond = 2
 χenv = 16
 ctm_alg = SimultaneousCTMRG()
-opt_alg = PEPSOptimize(;
-    boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2)
-)
+opt_alg = PEPSOptimize(; boundary_alg=ctm_alg, tol=1e-3)
 
 # initialize states
 H = transverse_field_ising(InfiniteSquare(); g)
 Random.seed!(91283219347)
 psi_init = InfinitePEPS(2, χbond)
-env_init = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
+env_init, = leading_boundary(CTMRGEnv(psi_init, ComplexSpace(χenv)), psi_init, ctm_alg)
 
 # find fixedpoint
-result = fixedpoint(psi_init, H, opt_alg, env_init)
-ξ_h, ξ_v, = correlation_length(result.peps, result.env)
+peps, env, E, = fixedpoint(psi_init, H, opt_alg, env_init)
+ξ_h, ξ_v, = correlation_length(peps, env)
 
 # compute magnetization
 σx = TensorMap(scalartype(psi_init)[0 1; 1 0], ℂ^2, ℂ^2)
 M = LocalOperator(H.lattice, (CartesianIndex(1, 1),) => σx)
 magn = expectation_value(result.peps, M, result.env)
 
-@test result.E ≈ e atol = 1e-2
+@test 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(psi_init, H_polar, opt_alg, env_init)
+peps_polar, env_polar, E_polar, = fixedpoint(psi_init, H_polar, opt_alg, env_init)
 ξ_h_polar, ξ_v_polar, = correlation_length(result_polar.peps, result_polar.env)
 @test ξ_h_polar < ξ_h
 @test ξ_v_polar < ξ_v