Sloop speedup

src/TNRKit.jl

@@ -5,7 +5,7 @@ using MatrixAlgebraKit
 using MatrixAlgebraKit: TruncationStrategy
 using LoggingExtras, Printf
 using KrylovKit
-using OptimKit, Zygote
+using OptimKit
 using DocStringExtensions
 using SpecialFunctions
 using FastGaussQuadrature

src/schemes/symmetric_looptnr.jl

@@ -27,7 +27,7 @@ mutable struct SLoopTNR{E, S, TT <: AbstractTensorMap{E, S, 4, 0}} <: TNRScheme{
 
     "Gradient optimization algorithm"
     gradalg::OptimKit.LBFGS
-    function SLoopTNR(T::TT; gradalg = LBFGS(10; verbosity = 0, gradtol = 6.0e-7, maxiter = 40000)) where {E, S, TT <: AbstractTensorMap{E, S, 4, 0}}
+    function SLoopTNR(T::TT; gradalg = LBFGS(10; verbosity = 0, gradtol = 1.0e-6, maxiter = 200)) where {E, S, TT <: AbstractTensorMap{E, S, 4, 0}}
         return new{E, S, TT}(T, gradalg)
     end
 end
@@ -48,52 +48,93 @@ classical_ising_inv() = classical_ising_inv(ising_βc)
 
 ########## utility functions ##########
 function trnorm_2x2(T)
-    @tensor TT[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(T[1 -2 2 -4])
+    @tensoropt TT[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(T[1 -2 2 -4])
     return sqrt(TTtoNorm(TT))
 end
 
 ########## Cost function ##########
 function StoSS(S)
     V = domain(S)[1]
     b = isomorphism(V, V')
-    @tensor SS[-1 -2 -3 -4] := S[-1 -2; 1] * S[-3 -4; 2] * b[1 2]
+    @tensoropt SS[-1 -2 -3 -4] := S[-1 -2; 1] * S[-3 -4; 2] * b[1 2]
     return SS
 end
 
 function TTtoNorm(TT)
     V = domain(TT)
     b = isomorphism(V[1] ⊗ V[2], V[1]' ⊗ V[2]')
     TTb = TT * b
-    @tensor T4[-1 -2; -3 -4] := TT[-1 -2; 1 2] * TTb[-3 -4; 1 2]
+    @tensoropt T4[-1 -2; -3 -4] := TT[-1 -2; 1 2] * TTb[-3 -4; 1 2]
     V = domain(T4)
     b = isomorphism(V[1] ⊗ V[2], V[1]' ⊗ V[2]')
     T4b = T4 * b
-    @tensor T8[-1 -2; -3 -4] := T4[-1 -2; 1 2] * T4b[-3 -4; 1 2]
+    @tensoropt T8[-1 -2; -3 -4] := T4[-1 -2; 1 2] * T4b[-3 -4; 1 2]
     V = domain(T8)
     b = isomorphism(V[1] ⊗ V[2], V[1]' ⊗ V[2]')
     return tr(T8 * b)
 end
 
-function cost_looptnr(S, T)
+function TtoNorm(T)
+    @tensoropt TT[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(T[1 -2 2 -4])
+    return TTtoNorm(TT)
+end
+
+function cost_looptnr(S, T, n_TT)
     @assert eltype(S) == Float64 "Modification is needed for complex numbers!"
     SS = StoSS(S)
-    @tensor TT[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(T[1 -2 2 -4])
-    @tensor TSS[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(SS[1 -2 2 -4])
-    @tensor S4[-1 -2; -3 -4] := SS[1 -1 2 -3] * conj(SS[1 -2 2 -4])
-    # T
-    return TTtoNorm(TT) + TTtoNorm(S4) - 2 * TTtoNorm(TSS)
+
+    @tensoropt TSS[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(SS[1 -2 2 -4])
+    @tensoropt S4[-1 -2; -3 -4] := SS[1 -1 2 -3] * conj(SS[1 -2 2 -4])
+    return n_TT + TTtoNorm(S4) - 2 * TTtoNorm(TSS)
 end
 
 ########## Gradient Optimization ##########
-function fg(f, A)
-    f_out, g = Zygote.withgradient(f, A)
+function loop_environment(TT)
+    V = domain(TT)
+    b = isomorphism(V[1] ⊗ V[2], V[1]' ⊗ V[2]')
+    TTb = TT * b
+    @tensoropt T4[-1 -2; -3 -4] := TT[-1 -2; 1 2] * TTb[-3 -4; 1 2]
+    return T4 * b * TT
+end
+
+function StoSS_pullback(S, gSS)
+    V = domain(S)[1]
+    b = isomorphism(V, V')
+    Sb = S * b
+    @tensoropt gS[-1 -2; -3] := gSS[-1 -2 1 2] * conj(Sb[1 2; -3])
+    return 2 * gS
+end
+
+function gradient_looptnr(S, T, SS, TSS, S4)
+    # One of the eight equal contributions from the S-S overlap.
+    env_SS = loop_environment(S4)
+    @tensoropt gSS[-1 -2 -3 -4] := SS[-1 1 -3 2] * env_SS[-2 1; -4 2]
+    gS_SS = StoSS_pullback(S, gSS)
+
+    # One of the four equal contributions from the T-S overlap.
+    env_TS = loop_environment(TSS)
+    env_TS_transposed = permute(env_TS, ((2, 1), (4, 3)))
+    @tensoropt gTS_dual[-1 -2 -3 -4] := conj(T[-1 1 -3 2]) * env_TS_transposed[-2 1; -4 2]
+    gTS = flip(gTS_dual, (1, 2, 3, 4))
+    gS_TS = StoSS_pullback(S, gTS)
+
+    return 8 * gS_SS - 2 * 4 * gS_TS
+end
+
+function cost_looptnr_fg(S, T, n_TT)
+    @assert eltype(S) == Float64 "Modification is needed for complex numbers!"
+    SS = StoSS(S)
 
-    return f_out, g[1]
+    @tensoropt TSS[-1 -2; -3 -4] := T[1 -1 2 -3] * conj(SS[1 -2 2 -4])
+    @tensoropt S4[-1 -2; -3 -4] := SS[1 -1 2 -3] * conj(SS[1 -2 2 -4])
+    cost = n_TT + TTtoNorm(S4) - 2 * TTtoNorm(TSS)
+    grad = gradient_looptnr(S, T, SS, TSS, S4)
+    return cost, grad
 end
 
 function optimize_S(scheme, S)
-    opt_fun(x) = cost_looptnr(x, scheme.T)
-    opt_fg(x) = fg(opt_fun, x)
+    n_TT = TtoNorm(scheme.T)
+    opt_fg(x) = cost_looptnr_fg(x, scheme.T, n_TT)
     Sopt, fx, gx, numfg, normgradhistory = optimize(
         opt_fg, S,
         scheme.gradalg
@@ -126,7 +167,7 @@ end
 
 function entanglement_filtering(T; trunc = trunctol(atol = 1.0e-12))
     entanglement_function(steps, data) = abs(data[end])
-    entanglement_criterion = maxiter(100) & convcrit(1.0e-12, entanglement_function)
+    entanglement_criterion = maxiter(200) & convcrit(1.0e-12, entanglement_function)
 
     psi_center = Ψ_center(T)
     psi_corner = Ψ_corner(T)
@@ -146,7 +187,7 @@ function entanglement_filtering(T; trunc = trunctol(atol = 1.0e-12))
     P_top = PL_list[3]
     P_left = PL_list[3]
 
-    @tensor T_new[-1 -2 -3 -4] := T[1 2 3 4] * P_left[-1; 1] * P_bottom[-2; 2] *
+    @tensoropt T_new[-1 -2 -3 -4] := T[1 2 3 4] * P_left[-1; 1] * P_bottom[-2; 2] *
         P_top[-3; 3] * P_right[-4; 4]
     return T_new
 end
@@ -168,14 +209,14 @@ function ef_oneloop(T, trunc::TruncationStrategy)
         criterion, trunc
     )
     i = 1
-    @tensor S[-2 -1; -3] := ΨB[i][-1; -2 2] * PRs[mod(i, 8) + 1][2; -3]
+    @tensoropt S[-2 -1; -3] := ΨB[i][-1; -2 2] * PRs[mod(i, 8) + 1][2; -3]
     return S
 end
 
 ########## Updating the tensor ##########
 function combine_4S(S)
     Sflip = flip(S, (1, 2))
-    @tensor Tnew[-1 -2 -3 -4] := S[1 2; -4] * Sflip[1 4; -3] * S[3 4; -1] * Sflip[3 2; -2]
+    @tensoropt Tnew[-1 -2 -3 -4] := S[1 2; -4] * Sflip[1 4; -3] * S[3 4; -1] * Sflip[3 2; -2]
     return Tnew
 end
 

test/schemes/schemes.jl

@@ -1,6 +1,8 @@
 using Test
 using TNRKit
 using TensorKit
+using LinearAlgebra
+using OptimKit
 
 # This tests every scheme in the library on the Z2 symmetric Ising model.
 println("---------------------")
@@ -368,7 +370,24 @@ end
     @test free_energy(data, ising_βc; initial_size = 2) ≈ f_onsager rtol = 1.0e-6
 end
 
-# SLoopTNR
+@testset "SLoopTNR - Manual gradient" begin
+    V = ℝ^2
+    T_inv = ones(Float64, V ⊗ V ⊗ V ⊗ V ← one(V))
+    S = ones(Float64, V ⊗ V ← V)
+    dS = ones(Float64, space(S))
+    n_TT = TNRKit.TtoNorm(T_inv)
+
+    cost, grad = TNRKit.cost_looptnr_fg(S, T_inv, n_TT)
+    ϵ = 1.0e-6
+    directional_derivative = (
+        TNRKit.cost_looptnr(S + ϵ * dS, T_inv, n_TT) -
+            TNRKit.cost_looptnr(S - ϵ * dS, T_inv, n_TT)
+    ) / (2ϵ)
+
+    @test cost ≈ TNRKit.cost_looptnr(S, T_inv, n_TT)
+    @test real(dot(grad, dS)) ≈ directional_derivative rtol = 1.0e-9
+end
+
 @testset "SLoopTNR - Ising Model" begin
     @info "SLoopTNR ising free energy"
     T_inv = classical_ising_inv()
@@ -377,6 +396,18 @@ end
     data = run!(scheme, truncrank(4), maxiter(25))
 
     @test free_energy(data, ising_βc) ≈ f_onsager rtol = 1.0e-5
+
+    @info "SLoopTNR Ising fixed-point tensor"
+    gradalg = LBFGS(10; verbosity = 0, gradtol = 6.0e-7, maxiter = 2000)
+    scheme = SLoopTNR(classical_ising_inv(); gradalg)
+    run!(scheme, truncrank(16), maxiter(16))
+    fp = fixed_point_tensor(scheme)
+
+    # The finite-χ value approaches the exact 0.645 from below (the paper
+    # reports 0.610 at D = 96 and finite L).
+    σ4 = real(fp[2, 2, 2, 2])
+    @test σ4 ≈ 0.5967 atol = 5.0e-3
+    @test σ4 ≈ 0.645 atol = 6.0e-2
 end
 
 # ctm