Fix tests (mostly), fix _condition_number

src/algorithms/ctmrg/simultaneous.jl

@@ -50,8 +50,11 @@ end
 
 # Compute condition number σ_max / σ_min for diagonal singular value TensorMap
 function _condition_number(S::AbstractTensorMap)
-    S_diag = diag(S.data)
-    return maximum(S_diag) / minimum(S_diag)
+    # 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)
 

src/algorithms/peps_opt.jl

@@ -175,6 +175,24 @@ function (d::DebugPEPSOptimize)(
     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
+
 """
     struct PEPSOptimize{G}
 
@@ -375,24 +393,6 @@ function gradient_function(cache::PEPSCostFunctionCache{T}) where {T}
     end
 end
 
-"""
-    SymmetrizeExponentialRetraction <: AbstractRetractionMethod
-    
-Exponential retraction followed by a symmetrization step.
-"""
-struct SymmetrizeExponentialRetraction <: AbstractRetractionMethod
-    symmetrization::SymmetrizationStyle
-    from_vec::Function
-end
-
-function Manifolds.retract(
-    M::AbstractManifold, p, X, t::Number, sr::SymmetrizeExponentialRetraction
-)
-    q = retract(M, p, X, t, ExponentialRetraction())
-    q_symm_peps = symmetrize!(sr.from_vec(q), sr.symmetrization)
-    return to_vec(q_symm_peps)
-end
-
 """
     fixedpoint(
         peps₀::InfinitePEPS{T},

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

test/ctmrg/gradients.jl

@@ -3,7 +3,8 @@ using Random
 using TensorKit
 using KrylovKit
 using PEPSKit
-using Zygote
+using PEPSKit: to_vec, PEPSCostFunctionCache, gradient_function
+using Manopt, Manifolds
 
 ## Test models, gradmodes and CTMRG algorithm
 # -------------------------------------------
@@ -40,7 +41,6 @@ gradmodes = [
         LinSolver(; solver=KrylovKit.BiCGStab(; tol=gradtol), iterscheme=:diffgauge),
     ],
 ]
-steps = -0.01:0.005:0.01
 
 ## Tests
 # ------
@@ -54,28 +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)
-        # TODO: redo this test using Manopt
-        # 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
 
-        #     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

test/heisenberg.jl

@@ -4,7 +4,6 @@ using Accessors
 using TensorKit
 using KrylovKit
 using PEPSKit
-using Manopt
 
 # initialize parameters
 Dbond = 2
@@ -23,29 +22,29 @@ E_ref = -0.6602310934799577
     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
 
 @testset "(1, 2) unit cell AD optimization" begin
     # 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

test/j1j2_model.jl

@@ -10,8 +10,9 @@ using PEPSKit
 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
@@ -22,7 +23,7 @@ psi_init = symmetrize!(psi_init, RotateReflect())
 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:

test/pwave.jl

@@ -3,27 +3,33 @@ using Random
 using TensorKit
 using KrylovKit
 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=20,
+    gradient_alg=LinSolver(; iterscheme=:diffgauge),
+    stepsize=WolfePowellLinesearch(),
+    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);
 
 # 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

test/tf_ising.jl

@@ -20,7 +20,7 @@ mˣ = 0.91
 χenv = 16
 ctm_alg = SimultaneousCTMRG()
 opt_alg = PEPSOptimize(;
-    boundary_alg=ctm_alg, optimizer=LBFGS(4; gradtol=1e-3, verbosity=2)
+    boundary_alg=ctm_alg, tol=1e-3, stepsize=WolfePowellBinaryLinesearch()
 )
 
 # initialize states
@@ -30,21 +30,21 @@ psi_init = InfinitePEPS(2, χbond)
 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