Refactor CTMRG routines and add PEPSOptimize struct, clean up some bits

Project.toml

@@ -15,6 +15,7 @@ Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 VectorInterface = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 MPSKit = "0.10"

README.md

@@ -3,7 +3,7 @@
     <img alt="PEPSKit.jl logo" src="https://github.com/quantumghent/PEPSKit.jl/blob/master/docs/src/assets/logo.svg" width="150">
 </picture>
 
-# PepsKit.jl
+# PEPSKit.jl
 
 [![docs][docs-dev-img]][docs-dev-url] ![CI][ci-url] [![codecov][codecov-img]][codecov-url]
 

examples/heisenberg.jl

@@ -1,140 +1,47 @@
-using Revise, PEPSKit, TensorKit, Zygote, MPSKit
-using MPSKitModels, LinearAlgebra, OptimKit
-using PEPSKit:
-    NORTH, SOUTH, WEST, EAST, NORTHWEST, NORTHEAST, SOUTHEAST, SOUTHWEST, @diffset
-using JLD2, ChainRulesCore
+using LinearAlgebra
+using TensorKit, MPSKitModels, OptimKit
+using PEPSKit
 
-function two_site_rho(r::Int, c::Int, ψ::InfinitePEPS, env::PEPSKit.CTMRGEnv)
-    cp = mod1(c + 1, size(ψ, 2))
-    @tensor ρ[-11, -20; -12, -18] :=
-        env.corners[NORTHWEST, r, c][1, 3] *
-        env.edges[WEST, r, c][2, 7, 9, 1] *
-        env.corners[SOUTHWEST, r, c][4, 2] *
-        env.edges[NORTH, r, c][3, 5, 8, 13] *
-        env.edges[SOUTH, r, c][14, 6, 10, 4] *
-        ψ[r, c][-12, 5, 15, 6, 7] *
-        conj(ψ[r, c][-11, 8, 19, 10, 9]) *
-        env.edges[NORTH, r, cp][13, 16, 22, 23] *
-        env.edges[SOUTH, r, cp][28, 17, 21, 14] *
-        ψ[r, cp][-18, 16, 25, 17, 15] *
-        conj(ψ[r, cp][-20, 22, 26, 21, 19]) *
-        env.corners[NORTHEAST, r, cp][23, 24] *
-        env.edges[EAST, r, cp][24, 25, 26, 27] *
-        env.corners[SOUTHEAST, r, cp][27, 28]
-    return ρ
-end
-
-function iCtmGsEh(
-    ψ::InfinitePEPS, env::PEPSKit.CTMRGEnv, H::AbstractTensorMap{S,2,2}
-) where {S}
-    #Es = Matrix{eltype(H)}(undef,size(ψ,1),size(ψ,2))
-    E = 0.0
-    for r in 1:size(ψ, 1), c in 1:size(ψ, 2)
-        ρ = two_site_rho(r, c, ψ, env)
-        @tensor nn = ρ[1 2; 1 2]
-        @tensor Eh = H[1 2; 3 4] * ρ[1 2; 3 4]
-        Eh = Eh / nn
-        E = E + Eh
-        #@diffset Es[r,c] = Eh;
-    end
-    return real(E)
-end
-
-function H_expectation_value(
-    ψ::InfinitePEPS, env::PEPSKit.CTMRGEnv, H::AbstractTensorMap{S,2,2}
-) where {S}
-    Eh = iCtmGsEh(ψ, env, H)
-    ψ1 = rotl90(ψ)
-    env1 = PEPSKit.rotate_north(env, EAST)
-    Ev = iCtmGsEh(ψ1, env1, H)
-    E = real(Eh + Ev)
-    return E
-end
-
-function SqLatHeisenberg()
+# Square lattice Heisenberg Hamiltonian
+# Sublattice-rotate to get (1, 1, 1) → (-1, 1, -1), transformed to GS with single-site unit cell
+function square_lattice_heisenberg(; Jx=-1.0, Jy=1.0, Jz=-1.0)
     Sx, Sy, Sz, _ = spinmatrices(1//2)
 
-    Dphys = ComplexSpace(2)
-    σx = TensorMap(Sx, Dphys, Dphys)
-    σy = TensorMap(Sy, Dphys, Dphys)
-    σz = TensorMap(Sz, Dphys, Dphys)
+    Vphys = ℂ^2
+    σx = TensorMap(Sx, Vphys, Vphys)
+    σy = TensorMap(Sy, Vphys, Vphys)
+    σz = TensorMap(Sz, Vphys, Vphys)
 
     @tensor H[-1 -3; -2 -4] :=
-        -σx[-1, -2] * σx[-3, -4] + σy[-1, -2] * σy[-3, -4] + -σz[-1, -2] * σz[-3, -4]
+        Jx * σx[-1, -2] * σx[-3, -4] +
+        Jy * σy[-1, -2] * σy[-3, -4] +
+        Jz * σz[-1, -2] * σz[-3, -4]
 
     return H
 end
 
-H = SqLatHeisenberg()
-
-function cfun(x)
-    (ψ, env) = x
-
-    function fun(peps)
-        env = leading_boundary(peps, alg_ctm, env)
-        x = H_expectation_value(peps, env, H)
-        return x
-    end
-    env = leading_boundary(ψ, alg_ctm, env)
-    E = H_expectation_value(ψ, env, H)
-    ∂E = fun'(ψ)
-
-    @assert !isnan(norm(∂E))
-    return E, ∂E
-end
-
-# my_retract is not an in place function which should not change x
-function my_retract(x, dx, α::Number)
-    (ϕ, env0) = x
-    ψ = deepcopy(ϕ)
-    env = deepcopy(env0)
-    ψ.A .+= dx.A .* α
-    #env = leading_boundary(ψ, alg_ctm,env)
-    return (ψ, env), dx
-end
-
-my_inner(x, dx1, dx2) = real(dot(dx1, dx2))
-
-function my_add!(Y, X, a)
-    Y.A .+= X.A .* a
-    return Y
-end
-
-function my_scale!(η, β)
-    rmul!(η.A, β)
-    return η
-end
-
-function init_psi(d::Int, D::Int, Lx::Int, Ly::Int)
+# Initialize InfinitePEPS with random & complex entries by default
+function init_peps(d, D, Lx, Ly; finit=randn, dtype=ComplexF64)
     Pspaces = fill(ℂ^d, Lx, Ly)
     Nspaces = fill(ℂ^D, Lx, Ly)
     Espaces = fill(ℂ^D, Lx, Ly)
-
-    Sspaces = adjoint.(circshift(Nspaces, (1, 0)))
-    Wspaces = adjoint.(circshift(Espaces, (0, -1)))
-
-    A = map(Pspaces, Nspaces, Espaces, Sspaces, Wspaces) do P, N, E, S, W
-        return TensorMap(rand, ComplexF64, P ← N * E * S * W)
-    end
-
-    return InfinitePEPS(A)
-end
-
-alg_ctm = CTMRG(; verbose=1, tol=1e-4, trscheme=truncdim(10), miniter=4, maxiter=200)
-
-function main(; d=2, D=2, Lx=1, Ly=1)
-    ψ = init_psi(d, D, Lx, Ly)
-    env = leading_boundary(ψ, alg_ctm)
-    optimize(
-        cfun,
-        (ψ, env),
-        ConjugateGradient(; verbosity=2);
-        inner=my_inner,
-        retract=my_retract,
-        (scale!)=my_scale!,
-        (add!)=my_add!,
-    )
-    return ψ
-end
-
-main()
+    return InfinitePEPS(Pspaces, Nspaces, Espaces; finit, dtype)
+end
+
+# Parameters
+H = square_lattice_heisenberg()
+χbond = 2
+χenv = 20
+ctmalg = CTMRG(; trscheme=truncdim(χenv), tol=1e-10, miniter=4, maxiter=100, verbose=2)
+optalg = PEPSOptimize{NaiveAD}(;
+    optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
+    fpgrad_tol=1e-6,
+    fpgrad_maxiter=100,
+    verbose=2,
+)
+
+# Ground state search
+ψinit = init_peps(2, χbond, 1, 1)
+envinit = leading_boundary(ψinit, ctmalg, CTMRGEnv(ψinit; χenv))
+result = groundsearch(H, ctmalg, optalg, ψinit, envinit)
+@show result.E₀

src/PEPSKit.jl

@@ -1,17 +1,13 @@
 module PEPSKit
 
+using LinearAlgebra, Base.Threads, Base.Iterators, Printf
 using Accessors
 using VectorInterface
-using TensorKit,
-    KrylovKit, MPSKit, OptimKit, Base.Threads, Base.Iterators, Parameters, Printf
-using ChainRulesCore
-
-using LinearAlgebra: LinearAlgebra
-
-export CTMRG, CTMRG2
-export leading_boundary
+using TensorKit, KrylovKit, MPSKit, OptimKit
+using ChainRulesCore, Zygote
 
 include("utility/util.jl")
+include("utility/rotations.jl")
 
 include("states/abstractpeps.jl")
 include("states/infinitepeps.jl")
@@ -28,19 +24,24 @@ include("environments/ctmrgenv.jl")
 include("environments/boundarympsenv.jl")
 
 include("algorithms/ctmrg.jl")
-include("algorithms/expval.jl")
+include("algorithms/peps_opt.jl")
 
 include("utility/symmetrization.jl")
 include("algorithms/pepo_opt.jl")
 
-include("utility/rotations.jl")
-
-#default settings
+# Default settings
 module Defaults
-    const maxiter = 100
-    const tol = 1e-12
+    const ctmrg_maxiter = 100
+    const ctmrg_miniter = 4
+    const ctmrg_tol = 1e-12
+    const grad_maxiter = 100
+    const grad_tol = 4
 end
 
+export CTMRG, CTMRGEnv
+export leading_boundary, expectation_value
+export PEPSOptimize, NaiveAD, ManualIter, LinSolve
+export groundsearch
 export InfinitePEPS, InfiniteTransferPEPS
 export InfinitePEPO, InfiniteTransferPEPO
 export initializeMPS, initializePEPS