Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Clean up and simplify optimization interface #127

Merged
merged 29 commits into from
Feb 11, 2025
Merged
Show file tree
Hide file tree
Changes from 18 commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
6211f4d
Update PEPSKit.Defaults
pbrehmer Feb 4, 2025
1916e2b
Break up peps_opt.jl into two files
pbrehmer Feb 4, 2025
c4e0436
Add symmetrization to PEPSOptimize
pbrehmer Feb 4, 2025
ab794ba
Fix realness warning in fixedpoint
pbrehmer Feb 4, 2025
42766e3
Rename `costfun` to `cost_function` and add docstring
pbrehmer Feb 5, 2025
9afaf5f
Clean up tests
pbrehmer Feb 5, 2025
dba99e4
Make PEPSKit.fixedpoint argument order consistent with MPSKit.fixedpo…
pbrehmer Feb 5, 2025
80f003d
Add truncation_error and condition_number to CTMRG info return tuples
pbrehmer Feb 5, 2025
6a8e739
Adapt leading_boundary calls to additional info return value
pbrehmer Feb 5, 2025
31e32bd
Adapt leading_boundary rrule to info tuple
pbrehmer Feb 5, 2025
ba17a81
Make changes runnable
pbrehmer Feb 6, 2025
3adbc67
Merge branch 'master' into pb-clean-opt-interface
pbrehmer Feb 6, 2025
3dc602f
Add collection of truncation errors, condition numbers, etc. during o…
pbrehmer Feb 6, 2025
e773be5
Adapt tests and examples to new fixedpoint return values
pbrehmer Feb 6, 2025
a027a3f
Fix tests (1st try)
pbrehmer Feb 6, 2025
799d9e8
Fix tests (2nd try)
pbrehmer Feb 6, 2025
a5c5e2a
Apply suggestions
pbrehmer Feb 7, 2025
3bad3da
Rename `envs` to `env`
pbrehmer Feb 7, 2025
1267077
Backpropagate _condition_number with ZeroTangent
pbrehmer Feb 7, 2025
d4c620b
Update src/algorithms/ctmrg/simultaneous.jl
pbrehmer Feb 7, 2025
8c93c6e
Relocate computation of condition to projectors
pbrehmer Feb 7, 2025
38d42ac
Merge branch 'pb-clean-opt-interface' of github.com:quantumghent/PEPS…
pbrehmer Feb 7, 2025
58b4570
Use `@ignore_derivatives` for _condition_number
pbrehmer Feb 7, 2025
bf09931
Apply local info suggestion
pbrehmer Feb 7, 2025
6afe263
Replace Zygote.Buffers in CTMRG
pbrehmer Feb 10, 2025
534e6f9
Merge branch 'pb-clean-opt-interface' of github.com:quantumghent/PEPS…
pbrehmer Feb 10, 2025
ea46eb2
Fix condition number computation
pbrehmer Feb 10, 2025
7d4f51f
Fix _condition_number
pbrehmer Feb 10, 2025
e7077a3
Compose first and last directly
pbrehmer Feb 10, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,9 @@ opt_alg = PEPSOptimize(;

# ground state search
state = InfinitePEPS(2, D)
ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
result = fixedpoint(state, H, opt_alg, ctm)
ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
peps, env, E, = fixedpoint(H, state, ctm, opt_alg)

@show result.E # -0.6625...
@show E # -0.6625...
```

6 changes: 3 additions & 3 deletions docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@ opt_alg = PEPSOptimize(;

# ground state search
state = InfinitePEPS(2, D)
ctm = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
result = fixedpoint(state, H, opt_alg, ctm)
ctm, = leading_boundary(CTMRGEnv(state, ComplexSpace(chi)), state, ctm_alg)
peps, env, E, = fixedpoint(H, state, ctm, opt_alg)

@show result.E # -0.6625...
@show E # -0.6625...
```
10 changes: 5 additions & 5 deletions examples/boundary_mps.jl
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,13 @@ T = InfiniteTransferPEPS(peps, 1, 1)
mps = PEPSKit.initializeMPS(T, [ComplexSpace(20)])

# We then find the leading boundary MPS fixed point using the VUMPS algorithm
mps, envs, ϵ = leading_boundary(mps, T, VUMPS())
mps, env, ϵ = leading_boundary(mps, T, VUMPS())

# The norm of the state per unit cell is then given by the expectation value <mps|T|mps>
N = abs(prod(expectation_value(mps, T)))

# This can be compared to the result obtained using the CTMRG algorithm
ctm = leading_boundary(
ctm, = leading_boundary(
peps, SimultaneousCTMRG(; verbosity=1), CTMRGEnv(peps, ComplexSpace(20))
)
N´ = abs(norm(peps, ctm))
Expand All @@ -52,10 +52,10 @@ peps2 = InfinitePEPS(ComplexSpace(2), ComplexSpace(2); unitcell=(2, 2))
T2 = PEPSKit.MultilineTransferPEPS(peps2, 1)

mps2 = PEPSKit.initializeMPS(T2, fill(ComplexSpace(20), 2, 2))
mps2, envs2, ϵ = leading_boundary(mps2, T2, VUMPS())
mps2, env2, ϵ = leading_boundary(mps2, T2, VUMPS())
N2 = abs(prod(expectation_value(mps2, T2)))

ctm2 = leading_boundary(
ctm2, = leading_boundary(
peps2, SimultaneousCTMRG(; verbosity=1), CTMRGEnv(peps2, ComplexSpace(20))
)
N2´ = abs(norm(peps2, ctm2))
Expand All @@ -77,7 +77,7 @@ pepo = ising_pepo(1)
T3 = InfiniteTransferPEPO(peps, pepo, 1, 1)

mps3 = PEPSKit.initializeMPS(T3, [ComplexSpace(20)])
mps3, envs3, ϵ = leading_boundary(mps3, T3, VUMPS())
mps3, env3, ϵ = leading_boundary(mps3, T3, VUMPS())
@show N3 = abs(prod(expectation_value(mps3, T3)))

# These objects and routines can be used to optimize PEPS fixed points of 3D partition
Expand Down
6 changes: 3 additions & 3 deletions examples/heisenberg.jl
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,6 @@ opt_alg = PEPSOptimize(;
# E/N = −0.6694421, which is a QMC estimate from https://arxiv.org/abs/1101.3281.
# Of course there is a noticable bias for small χbond and χenv.
ψ₀ = InfinitePEPS(2, χbond)
env₀ = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg)
result = fixedpoint(ψ₀, H, opt_alg, env₀)
@show result.E
env₀, = leading_boundary(CTMRGEnv(ψ₀, ℂ^χenv), ψ₀, ctm_alg)
peps, env, E, = fixedpoint(H, ψ, env₀, opt_alg₀)
@show E
10 changes: 5 additions & 5 deletions examples/heisenberg_evol/heis_tools.jl
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ using PEPSKit
"""
Measure magnetization on each site
"""
function cal_mags(peps::InfinitePEPS, envs::CTMRGEnv)
function cal_mags(peps::InfinitePEPS, env::CTMRGEnv)
N1, N2 = size(peps)
lattice = collect(space(t, 1) for t in peps.A)
# detect symmetry on physical axis
Expand All @@ -32,7 +32,7 @@ function cal_mags(peps::InfinitePEPS, envs::CTMRGEnv)
return [
collect(
expectation_value(
peps, LocalOperator(lattice, (CartesianIndex(r, c),) => Sa), envs
peps, LocalOperator(lattice, (CartesianIndex(r, c),) => Sa), env
) for (r, c) in Iterators.product(1:N1, 1:N2)
) for Sa in Sas
]
Expand All @@ -41,11 +41,11 @@ end
"""
Measure physical quantities for Heisenberg model
"""
function measure_heis(peps::InfinitePEPS, H::LocalOperator, envs::CTMRGEnv)
function measure_heis(peps::InfinitePEPS, H::LocalOperator, env::CTMRGEnv)
results = Dict{String,Any}()
N1, N2 = size(peps)
results["e_site"] = costfun(peps, envs, H) / (N1 * N2)
results["mag"] = cal_mags(peps, envs)
results["e_site"] = costfun(peps, env, H) / (N1 * N2)
results["mag"] = cal_mags(peps, env)
results["mag_norm"] = collect(
norm([mags[r, c] for mags in results["mag"]]) for
(r, c) in Iterators.product(1:N1, 1:N2)
Expand Down
10 changes: 5 additions & 5 deletions examples/heisenberg_evol/simpleupdate.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ include("heis_tools.jl")
# benchmark data is from Phys. Rev. B 94, 035133 (2016)
Dbond, χenv, symm = 4, 16, Trivial
trscheme_peps = truncerr(1e-10) & truncdim(Dbond)
trscheme_envs = truncerr(1e-10) & truncdim(χenv)
trscheme_env = truncerr(1e-10) & truncdim(χenv)
N1, N2 = 2, 2
# Heisenberg model Hamiltonian
# (already only includes nearest neighbor terms)
Expand Down Expand Up @@ -41,10 +41,10 @@ for (dt, tol) in zip(dts, tols)
end
# measure physical quantities with CTMRG
peps_ = InfinitePEPS(peps)
envs = CTMRGEnv(rand, Float64, peps_, Espace)
ctm_alg = SequentialCTMRG(; tol=1e-10, verbosity=2, trscheme=trscheme_envs)
envs = leading_boundary(envs, peps_, ctm_alg)
meas = measure_heis(peps_, ham, envs)
env₀ = CTMRGEnv(rand, Float64, peps_, Espace)
ctm_alg = SequentialCTMRG(; tol=1e-10, verbosity=2, trscheme=trscheme_env)
env = leading_boundary(env₀, peps_, ctm_alg)
meas = measure_heis(peps_, ham, env)
display(meas)
@info @sprintf("Energy = %.8f\n", meas["e_site"])
@info @sprintf("Staggered magnetization = %.8f\n", mean(meas["mag_norm"]))
Expand Down
6 changes: 3 additions & 3 deletions examples/hubbard_su.jl
Original file line number Diff line number Diff line change
Expand Up @@ -40,10 +40,10 @@ peps = InfinitePEPS(peps)
# CTMRG
χenv0, χenv = 6, 20
Espace = Vect[fℤ₂](0 => χenv0 / 2, 1 => χenv0 / 2)
envs = CTMRGEnv(randn, Float64, peps, Espace)
env = CTMRGEnv(randn, Float64, peps, Espace)
for χ in [χenv0, χenv]
ctm_alg = SequentialCTMRG(; maxiter=300, tol=1e-7)
envs = leading_boundary(envs, peps, ctm_alg)
env, = leading_boundary(env, peps, ctm_alg)
end

# Benchmark values of the ground state energy from
Expand All @@ -53,7 +53,7 @@ Es_exact = Dict(0 => -1.62, 2 => -0.176, 4 => 0.8603, 6 => -0.6567, 8 => -0.5243
E_exact = Es_exact[U] - U / 2

# measure energy
E = costfun(peps, envs, ham) / (N1 * N2)
E = cost_function(peps, env, ham) / (N1 * N2)
@info "Energy = $E"
@info "Benchmark energy = $E_exact"
@test isapprox(E, E_exact; atol=5e-2)
109 changes: 60 additions & 49 deletions src/PEPSKit.jl
Original file line number Diff line number Diff line change
Expand Up @@ -54,56 +54,65 @@ include("algorithms/time_evolution/simpleupdate.jl")

include("algorithms/toolbox.jl")

include("algorithms/peps_opt.jl")

include("utility/symmetrization.jl")

include("algorithms/optimization/fixed_point_differentiation.jl")
include("algorithms/optimization/peps_optimization.jl")

"""
module Defaults
const ctmrg_maxiter = 100
const ctmrg_miniter = 4
const ctmrg_tol = 1e-8
const fpgrad_maxiter = 30
const fpgrad_tol = 1e-6
const reuse_env = true
const trscheme = FixedSpaceTruncation()
const fwd_alg = TensorKit.SDD()
const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
const projector_alg_type = HalfInfiniteProjector
const projector_alg = projector_alg_type(svd_alg, trscheme, 2)
const ctmrg_alg = SimultaneousCTMRG(
ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg
)
const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)
const gradient_linsolver = KrylovKit.BiCGStab(;
maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
)
const iterscheme = :fixed
const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
const scheduler = Ref{Scheduler}(Threads.nthreads() == 1 ? SerialScheduler() : DynamicScheduler())
end

Module containing default values that represent typical algorithm parameters.
Module containing default algorithm parameter values and arguments.

- `ctmrg_maxiter`: Maximal number of CTMRG iterations per run
- `ctmrg_miniter`: Minimal number of CTMRG carried out
- `ctmrg_tol`: Tolerance checking singular value and norm convergence
- `fpgrad_maxiter`: Maximal number of iterations for computing the CTMRG fixed-point gradient
- `fpgrad_tol`: Convergence tolerance for the fixed-point gradient iteration
- `reuse_env`: If `true`, the current optimization step is initialized on the previous environment
- `trscheme`: Truncation scheme for SVDs and other decompositions
- `fwd_alg`: SVD algorithm that is used in the forward pass
# CTMRG
- `ctmrg_tol=1e-8`: Tolerance checking singular value and norm convergence
- `ctmrg_maxiter=100`: Maximal number of CTMRG iterations per run
- `ctmrg_miniter=4`: Minimal number of CTMRG carried out
- `trscheme=FixedSpaceTruncation()`: Truncation scheme for SVDs and other decompositions
- `fwd_alg=TensorKit.SDD()`: SVD algorithm that is used in the forward pass
- `rrule_alg`: Reverse-rule for differentiating that SVD
- `svd_alg`: Combination of `fwd_alg` and `rrule_alg`
- `projector_alg_type`: Default type of projector algorithm

```
rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1)
```

- `svd_alg=SVDAdjoint(; fwd_alg, rrule_alg)`: Combination of `fwd_alg` and `rrule_alg`
- `projector_alg_type=HalfInfiniteProjector`: Default type of projector algorithm
- `projector_alg`: Algorithm to compute CTMRG projectors

```
projector_alg = projector_alg_type(; svd_alg, trscheme, verbosity=0)
```

- `ctmrg_alg`: Algorithm for performing CTMRG runs
- `optimizer`: Optimization algorithm for PEPS ground-state optimization

```
ctmrg_alg = SimultaneousCTMRG(
ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg
)
```

# Optimization
- `fpgrad_maxiter=30`: Maximal number of iterations for computing the CTMRG fixed-point gradient
- `fpgrad_tol=1e-6`: Convergence tolerance for the fixed-point gradient iteration
- `iterscheme=:fixed`: Scheme for differentiating one CTMRG iteration
- `gradient_linsolver`: Default linear solver for the `LinSolver` gradient algorithm
- `iterscheme`: Scheme for differentiating one CTMRG iteration

```
gradient_linsolver=KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol)
```

- `gradient_alg`: Algorithm to compute the gradient fixed-point
- `scheduler`: Multi-threading scheduler which can be accessed via `set_scheduler!`

```
gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
```

- `reuse_env=true`: If `true`, the current optimization step is initialized on the previous environment
- `optimizer=LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)`: Default `OptimKit.OptimizerAlgorithm` for PEPS optimization

# OhMyThreads scheduler
- `scheduler=Ref{Scheduler}(...)`: Multi-threading scheduler which can be accessed via `set_scheduler!`
"""
module Defaults
using TensorKit, KrylovKit, OptimKit, OhMyThreads
Expand All @@ -113,28 +122,30 @@ module Defaults
SVDAdjoint,
HalfInfiniteProjector,
SimultaneousCTMRG

# CTMRG
const ctmrg_tol = 1e-8
const ctmrg_maxiter = 100
const ctmrg_miniter = 4
const fpgrad_maxiter = 30
const fpgrad_tol = 1e-6
const sparse = false
const reuse_env = true
const trscheme = FixedSpaceTruncation()
const fwd_alg = TensorKit.SDD()
const rrule_alg = Arnoldi(; tol=1e-2fpgrad_tol, krylovdim=48, verbosity=-1)
const rrule_alg = Arnoldi(; tol=ctmrg_tol, krylovdim=48, verbosity=-1)
const svd_alg = SVDAdjoint(; fwd_alg, rrule_alg)
const projector_alg_type = HalfInfiniteProjector
const projector_alg = projector_alg_type(svd_alg, trscheme, 2)
const projector_alg = projector_alg_type(; svd_alg, trscheme, verbosity=0)
const ctmrg_alg = SimultaneousCTMRG(
ctmrg_tol, ctmrg_maxiter, ctmrg_miniter, 2, projector_alg
)
const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)
const gradient_linsolver = KrylovKit.BiCGStab(;
maxiter=Defaults.fpgrad_maxiter, tol=Defaults.fpgrad_tol
)

# Optimization
const fpgrad_maxiter = 30
const fpgrad_tol = 1e-6
const gradient_linsolver = KrylovKit.BiCGStab(; maxiter=fpgrad_maxiter, tol=fpgrad_tol)
const iterscheme = :fixed
const gradient_alg = LinSolver(; solver=gradient_linsolver, iterscheme)
const reuse_env = true
const optimizer = LBFGS(32; maxiter=100, gradtol=1e-4, verbosity=3)

# OhMyThreads scheduler defaults
const scheduler = Ref{Scheduler}()
Expand Down Expand Up @@ -187,7 +198,7 @@ export SVDAdjoint, IterSVD
export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
export FixedSpaceTruncation, HalfInfiniteProjector, FullInfiniteProjector
export LocalOperator
export expectation_value, costfun, product_peps, correlation_length
export expectation_value, cost_function, product_peps, correlation_length
export leading_boundary
export PEPSOptimize, GeomSum, ManualIter, LinSolver
export fixedpoint
Expand Down
Loading