Skip to content

Commit

Permalink
Fix left over merge issues, update default verbosities
Browse files Browse the repository at this point in the history
  • Loading branch information
@pbrehmer
pbrehmer committed Jul 15, 2024
1 parent 741f4de commit fe17ed7
Show file tree
Hide file tree
Showing 12 changed files with 60 additions and 237 deletions.
1 change: 0 additions & 1 deletion README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,6 @@ opt_alg = PEPSOptimize(;
optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
gradient_alg=GMRES(; tol=1e-6, maxiter=100),
reuse_env=true,
verbosity=2,
)

# ground state search
Expand Down
1 change: 0 additions & 1 deletion docs/src/index.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@ opt_alg = PEPSOptimize(;
optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
gradient_alg=GMRES(; tol=1e-6, maxiter=100),
reuse_env=true,
verbosity=2,
)

# ground state search
Expand Down
3 changes: 1 addition & 2 deletions examples/heisenberg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,12 @@ H = square_lattice_heisenberg(; Jx=-1, Jy=1, Jz=-1)
# Parameters
χbond = 2
χenv = 20
ctm_alg = CTMRG(; tol=1e-10, miniter=4, maxiter=100, verbosity=1, trscheme=truncdim(χenv))
ctm_alg = CTMRG(; tol=1e-10, miniter=4, maxiter=100, verbosity=2)
opt_alg = PEPSOptimize(;
boundary_alg=ctm_alg,
optimizer=LBFGS(4; maxiter=100, gradtol=1e-4, verbosity=2),
gradient_alg=GMRES(; tol=1e-6, maxiter=100),
reuse_env=true,
verbosity=2,
)

# Ground state search
Expand Down
215 changes: 9 additions & 206 deletions src/algorithms/ctmrg.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ function CTMRG(;
tol=Defaults.ctmrg_tol,
maxiter=Defaults.ctmrg_maxiter,
miniter=Defaults.ctmrg_miniter,
verbosity=1,
verbosity=2,
svd_alg=SVDAdjoint(),
trscheme=FixedSpaceTruncation(),
ctmrgscheme=Defaults.ctmrgscheme,
Expand All @@ -77,7 +77,7 @@ Per default, a random initial environment is used.
function MPSKit.leading_boundary(state, alg::CTMRG)
return MPSKit.leading_boundary(CTMRGEnv(state, oneunit(spacetype(state))), state, alg)
end
function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
function MPSKit.leading_boundary(envinit, state, alg::CTMRG{S}) where {S}
CS = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.corners)
TS = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envinit.edges)

Expand All @@ -99,9 +99,14 @@ function MPSKit.leading_boundary(envinit, state, alg::CTMRG)
end

# Do one final iteration that does not change the spaces
alg_fixed = @set alg.projector_alg.trscheme = FixedSpaceTruncation()
alg_fixed = CTMRG(;
verbosity=alg.verbosity,
svd_alg=alg.projector_alg.svd_alg,
trscheme=FixedSpaceTruncation(),
ctmrgscheme=S,
)
env′, = ctmrg_iter(state, env, alg_fixed)
envfix = gauge_fix(env, env′)
envfix, = gauge_fix(env, env′)

η = calc_elementwise_convergence(envfix, env; atol=alg.tol^(1 / 2))
N = norm(state, envfix)
Expand All @@ -125,208 +130,6 @@ ctmrg_logcancel!(log, iter, η, N) = @warnv 1 logcancel!(log, iter, η, N)
@non_differentiable ctmrg_logfinish!(args...)
@non_differentiable ctmrg_logcancel!(args...)

"""
gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
Fix the gauge of `envfinal` based on the previous environment `envprev`.
This assumes that the `envfinal` is the result of one CTMRG iteration on `envprev`.
Given that the CTMRG run is converged, the returned environment will be
element-wise converged to `envprev`.
"""
function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
# Check if spaces in envprev and envfinal are the same
same_spaces = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) &&
space(envfinal.corners[dir, r, c]) == space(envprev.corners[dir, r, c])
end
@assert all(same_spaces) "Spaces of envprev and envfinal are not the same"

# Try the "general" algorithm from https://arxiv.org/abs/2311.11894
signs = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
# Gather edge tensors and pretend they're InfiniteMPSs
if dir == NORTH
Tsprev = circshift(envprev.edges[dir, r, :], 1 - c)
Tsfinal = circshift(envfinal.edges[dir, r, :], 1 - c)
elseif dir == EAST
Tsprev = circshift(envprev.edges[dir, :, c], 1 - r)
Tsfinal = circshift(envfinal.edges[dir, :, c], 1 - r)
elseif dir == SOUTH
Tsprev = circshift(reverse(envprev.edges[dir, r, :]), c)
Tsfinal = circshift(reverse(envfinal.edges[dir, r, :]), c)
elseif dir == WEST
Tsprev = circshift(reverse(envprev.edges[dir, :, c]), r)
Tsfinal = circshift(reverse(envfinal.edges[dir, :, c]), r)
end

# Random MPS of same bond dimension
M = map(Tsfinal) do t
TensorMap(randn, scalartype(t), codomain(t) ← domain(t))
end

# Find right fixed points of mixed transfer matrices
ρinit = TensorMap(
randn,
scalartype(T),
MPSKit._lastspace(Tsfinal[end])' ← MPSKit._lastspace(M[end])',
)
ρprev = transfermatrix_fixedpoint(Tsprev, M, ρinit)
ρfinal = transfermatrix_fixedpoint(Tsfinal, M, ρinit)

# Decompose and multiply
Qprev, = leftorth(ρprev)
Qfinal, = leftorth(ρfinal)

return Qprev * Qfinal'
end

cornersfix, edgesfix = fix_relative_phases(envfinal, signs)

# Fix global phase
cornersgfix = map(envprev.corners, cornersfix) do Cprev, Cfix
return dot(Cfix, Cprev) * Cfix
end
edgesgfix = map(envprev.edges, edgesfix) do Tprev, Tfix
return dot(Tfix, Tprev) * Tfix
end
return CTMRGEnv(cornersgfix, edgesgfix)
end

# this is a bit of a hack to get the fixed point of the mixed transfer matrix
# because MPSKit is not compatible with AD
function transfermatrix_fixedpoint(tops, bottoms, ρinit)
_, vecs, info = eigsolve(ρinit, 1, :LM, Arnoldi()) do ρ
return foldr(zip(tops, bottoms); init=ρ) do (top, bottom), ρ
return @tensor ρ′[-1; -2] := top[-1 4 3; 1] * conj(bottom[-2 4 3; 2]) * ρ[1; 2]
end
end
info.converged > 0 || @warn "eigsolve did not converge"
return first(vecs)
end

# Explicit fixing of relative phases (doing this compactly in a loop is annoying)
function _contract_gauge_corner(corner, σ_in, σ_out)
@autoopt @tensor corner_fix[χ_in; χ_out] :=
σ_in[χ_in; χ1] * corner[χ1; χ2] * conj(σ_out[χ_out; χ2])
end
function _contract_gauge_edge(edge, σ_in, σ_out)
@autoopt @tensor edge_fix[χ_in D_above D_below; χ_out] :=
σ_in[χ_in; χ1] * edge[χ1 D_above D_below; χ2] * conj(σ_out[χ_out; χ2])
end
function fix_relative_phases(envfinal::CTMRGEnv, signs)
C1 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
_contract_gauge_corner(
envfinal.corners[NORTHWEST, r, c],
signs[WEST, r, c],
signs[NORTH, r, _next(c, end)],
)
end
T1 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
_contract_gauge_edge(
envfinal.edges[NORTH, r, c],
signs[NORTH, r, c],
signs[NORTH, r, _next(c, end)],
)
end
C2 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
_contract_gauge_corner(
envfinal.corners[NORTHEAST, r, c],
signs[NORTH, r, c],
signs[EAST, _next(r, end), c],
)
end
T2 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
_contract_gauge_edge(
envfinal.edges[EAST, r, c], signs[EAST, r, c], signs[EAST, _next(r, end), c]
)
end
C3 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
_contract_gauge_corner(
envfinal.corners[SOUTHEAST, r, c],
signs[EAST, r, c],
signs[SOUTH, r, _prev(c, end)],
)
end
T3 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
_contract_gauge_edge(
envfinal.edges[SOUTH, r, c],
signs[SOUTH, r, c],
signs[SOUTH, r, _prev(c, end)],
)
end
C4 = map(Iterators.product(axes(envfinal.corners)[2:3]...)) do (r, c)
_contract_gauge_corner(
envfinal.corners[SOUTHWEST, r, c],
signs[SOUTH, r, c],
signs[WEST, _prev(r, end), c],
)
end
T4 = map(Iterators.product(axes(envfinal.edges)[2:3]...)) do (r, c)
_contract_gauge_edge(
envfinal.edges[WEST, r, c], signs[WEST, r, c], signs[WEST, _prev(r, end), c]
)
end

return stack([C1, C2, C3, C4]; dims=1), stack([T1, T2, T3, T4]; dims=1)
end

function calc_convergence(envs, CSold, TSold)
CSnew = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envs.corners)
ΔCS = maximum(zip(CSold, CSnew)) do (c_old, c_new)
# only compute the difference on the smallest part of the spaces
smallest = infimum(MPSKit._firstspace(c_old), MPSKit._firstspace(c_new))
e_old = isometry(MPSKit._firstspace(c_old), smallest)
e_new = isometry(MPSKit._firstspace(c_new), smallest)
return norm(e_new' * c_new * e_new - e_old' * c_old * e_old)
end

TSnew = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envs.edges)
ΔTS = maximum(zip(TSold, TSnew)) do (t_old, t_new)
MPSKit._firstspace(t_old) == MPSKit._firstspace(t_new) ||
return scalartype(t_old)(Inf)
return norm(t_new - t_old)
end

@debug "maxᵢ|Cⁿ⁺¹ - Cⁿ|ᵢ = $ΔCS maxᵢ|Tⁿ⁺¹ - Tⁿ|ᵢ = $ΔTS"

return max(ΔCS, ΔTS), CSnew, TSnew
end

@non_differentiable calc_convergence(args...)

"""
calc_elementwise_convergence(envfinal, envfix; atol=1e-6)
Check if the element-wise difference of the corner and edge tensors of the final and fixed
CTMRG environments are below some tolerance.
"""
function calc_elementwise_convergence(envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6)
ΔC = envfinal.corners .- envfix.corners
ΔCmax = norm(ΔC, Inf)
ΔCmean = norm(ΔC)
@debug "maxᵢⱼ|Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmax mean |Cⁿ⁺¹ - Cⁿ|ᵢⱼ = $ΔCmean"

ΔT = envfinal.edges .- envfix.edges
ΔTmax = norm(ΔT, Inf)
ΔTmean = norm(ΔT)
@debug "maxᵢⱼ|Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmax mean |Tⁿ⁺¹ - Tⁿ|ᵢⱼ = $ΔTmean"

# Check differences for all tensors in unit cell to debug properly
for (dir, r, c) in Iterators.product(axes(envfinal.edges)...)
@debug(
"$((dir, r, c)): all |Cⁿ⁺¹ - Cⁿ|ᵢⱼ < ϵ: ",
all(x -> abs(x) < atol, convert(Array, ΔC[dir, r, c])),
)
@debug(
"$((dir, r, c)): all |Tⁿ⁺¹ - Tⁿ|ᵢⱼ < ϵ: ",
all(x -> abs(x) < atol, convert(Array, ΔT[dir, r, c])),
)
end

return max(ΔCmax, ΔTmax)
end

@non_differentiable calc_elementwise_convergence(args...)

"""
ctmrg_iter(state, env::CTMRGEnv{C,T}, alg::CTMRG) where {C,T}
Expand Down
43 changes: 33 additions & 10 deletions src/algorithms/ctmrg_gauge_fix.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,15 @@ This assumes that the `envfinal` is the result of one CTMRG iteration on `envpre
Given that the CTMRG run is converged, the returned environment will be
element-wise converged to `envprev`.
"""
function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C′,T′}) where {C,C′,T,T′}
function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C,T}) where {C,T}
# Check if spaces in envprev and envfinal are the same
same_spaces = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
space(envfinal.edges[dir, r, c]) == space(envprev.edges[dir, r, c]) &&
space(envfinal.corners[dir, r, c]) == space(envprev.corners[dir, r, c])
end
@assert all(same_spaces) "Spaces of envprev and envfinal are not the same"

# "general" algorithm from https://arxiv.org/abs/2311.11894
# Try the "general" algorithm from https://arxiv.org/abs/2311.11894
signs = map(Iterators.product(axes(envfinal.edges)...)) do (dir, r, c)
# Gather edge tensors and pretend they're InfiniteMPSs
if dir == NORTH
Expand Down Expand Up @@ -46,8 +46,8 @@ function gauge_fix(envprev::CTMRGEnv{C,T}, envfinal::CTMRGEnv{C′,T′}) where
ρfinal = transfermatrix_fixedpoint(Tsfinal, M, ρinit)

# Decompose and multiply
Qprev, = leftorth!(ρprev)
Qfinal, = leftorth!(ρfinal)
Qprev, = leftorth(ρprev)
Qfinal, = leftorth(ρfinal)

return Qprev * Qfinal'
end
Expand Down Expand Up @@ -177,15 +177,38 @@ function fix_global_phases(envprev::CTMRGEnv, envfix::CTMRGEnv)
return CTMRGEnv(cornersgfix, edgesgfix)
end


function calc_convergence(envs, CSold, TSold)
CSnew = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envs.corners)
ΔCS = maximum(zip(CSold, CSnew)) do (c_old, c_new)
# only compute the difference on the smallest part of the spaces
smallest = infimum(MPSKit._firstspace(c_old), MPSKit._firstspace(c_new))
e_old = isometry(MPSKit._firstspace(c_old), smallest)
e_new = isometry(MPSKit._firstspace(c_new), smallest)
return norm(e_new' * c_new * e_new - e_old' * c_old * e_old)
end

TSnew = map(x -> tsvd(x; alg=TensorKit.SVD())[2], envs.edges)
ΔTS = maximum(zip(TSold, TSnew)) do (t_old, t_new)
MPSKit._firstspace(t_old) == MPSKit._firstspace(t_new) ||
return scalartype(t_old)(Inf)
return norm(t_new - t_old)
end

@debug "maxᵢ|Cⁿ⁺¹ - Cⁿ|ᵢ = $ΔCS maxᵢ|Tⁿ⁺¹ - Tⁿ|ᵢ = $ΔTS"

return max(ΔCS, ΔTS), CSnew, TSnew
end

@non_differentiable calc_convergence(args...)

"""
check_elementwise_convergence(envfinal, envfix; atol=1e-6)
calc_elementwise_convergence(envfinal, envfix; atol=1e-6)
Check if the element-wise difference of the corner and edge tensors of the final and fixed
CTMRG environments are below some tolerance.
"""
function check_elementwise_convergence(
envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6
)
function calc_elementwise_convergence(envfinal::CTMRGEnv, envfix::CTMRGEnv; atol::Real=1e-6)
ΔC = envfinal.corners .- envfix.corners
ΔCmax = norm(ΔC, Inf)
ΔCmean = norm(ΔC)
Expand All @@ -208,7 +231,7 @@ function check_elementwise_convergence(
)
end

return isapprox(ΔCmax, 0; atol) && isapprox(ΔTmax, 0; atol)
return max(ΔCmax, ΔTmax)
end

@non_differentiable check_elementwise_convergence(args...)
@non_differentiable calc_elementwise_convergence(args...)
2 changes: 1 addition & 1 deletion src/algorithms/peps_opt.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ end

"""
PEPSOptimize{G}(; boundary_alg=CTMRG(), optimizer::OptimKit.OptimizationAlgorithm=Defaults.optimizer
reuse_env::Bool=true, gradient_alg::G=LinSolver(), verbosity::Int=0)
reuse_env::Bool=true, gradient_alg::G=LinSolver())
Algorithm struct that represent PEPS ground-state optimization using AD.
Set the algorithm to contract the infinite PEPS in `boundary_alg`;
Expand Down
4 changes: 2 additions & 2 deletions test/ctmrg/ctmrgschemes.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,8 @@ using PEPSKit
# initialize parameters
χbond = 2
χenv = 16
ctm_alg_sequential = CTMRG(; tol=1e-10, verbosity=1, ctmrgscheme=:sequential)
ctm_alg_simultaneous = CTMRG(; tol=1e-10, verbosity=1, ctmrgscheme=:simultaneous)
ctm_alg_sequential = CTMRG(; tol=1e-10, verbosity=2, ctmrgscheme=:sequential)
ctm_alg_simultaneous = CTMRG(; tol=1e-10, verbosity=2, ctmrgscheme=:simultaneous)
unitcells = [(1, 1), (3, 4)]

@testset "$(unitcell) unit cell" for unitcell in unitcells
Expand Down
Loading

0 comments on commit fe17ed7

Please sign in to comment.