diff --git a/src/TNRKit.jl b/src/TNRKit.jl index 78f47824..1bccfd87 100644 --- a/src/TNRKit.jl +++ b/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 diff --git a/src/schemes/symmetric_looptnr.jl b/src/schemes/symmetric_looptnr.jl index 2481cbae..499fba0a 100644 --- a/src/schemes/symmetric_looptnr.jl +++ b/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,7 +48,7 @@ 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 @@ -56,7 +56,7 @@ end 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 @@ -64,36 +64,77 @@ 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 diff --git a/test/schemes/schemes.jl b/test/schemes/schemes.jl index b428a0a6..130b7ae4 100644 --- a/test/schemes/schemes.jl +++ b/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