Use MPSKit's logging macros and adapt_solver

Project.toml

@@ -14,7 +14,6 @@ DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 KrylovKit = "0b1a1467-8014-51b9-945f-bf0ae24f4b77"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 MPSKit = "bb1c41ca-d63c-52ed-829e-0820dda26502"
 MPSKitModels = "ca635005-6f8c-4cd1-b51d-8491250ef2ab"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
@@ -38,7 +37,6 @@ DocStringExtensions = "0.9.3"
 FiniteDifferences = "0.12"
 KrylovKit = "0.9.5, 0.10"
 LinearAlgebra = "1"
-LoggingExtras = "1"
 MPSKit = "0.14"
 MPSKitModels = "0.4"
 MatrixAlgebraKit = "0.6.5"

src/PEPSKit.jl

@@ -24,13 +24,15 @@ using KrylovKit: Lanczos, BlockLanczos
 
 using TensorOperations, OptimKit
 using ChainRulesCore, Zygote
-using LoggingExtras
 import TupleTools
 
 using MPSKit
 using MPSKit: MPSTensor, MPOTensor, GenericMPSTensor, MPSBondTensor, ProductTransferMatrix
 using MPSKit: InfiniteEnvironments
-using MPSKit: DynamicTol, updatetol
+using MPSKit: DynamicTol
+using MPSKit.DynamicTols: adapt_solver
+using MPSKit: @log_initialization, @log_iteration, @log_convergence, @log_nonconvergence
+using MPSKit.IterativeLoggers: with_verbosity
 import MPSKit.DynamicTols: _updatetol
 import MPSKit: tensorexpr, leading_boundary, loginit!, logiter!, logfinish!, logcancel!, physicalspace
 import MPSKit: infinite_temperature_density_matrix

src/algorithms/bp/beliefpropagation.jl

@@ -33,24 +33,24 @@ end
 Contract `network` in the BP approximation and return the corresponding messages.
 """
 function leading_boundary(env₀::BPEnv, network::InfiniteSquareNetwork, alg::BeliefPropagation)
-    return LoggingExtras.withlevel(; alg.verbosity) do
+    return with_verbosity(; alg.verbosity) do
         env = deepcopy(env₀)
         log = MPSKit.IterLog("BP")
         ϵ = Inf
-        @infov 1 loginit!(log, ϵ)
+        @log_initialization loginit!(log, ϵ)
         for iter in 1:(alg.maxiter)
             env′ = bp_iteration(network, env, alg)
             ϵ = oftype(ϵ, tr_distance(env, env′))
             env = env′
 
             if ϵ <= alg.tol && iter >= alg.miniter
-                @infov 2 logfinish!(log, iter, ϵ)
+                @log_convergence logfinish!(log, iter, ϵ)
                 break
             end
             if iter ≥ alg.maxiter
-                @warnv 1 logcancel!(log, iter, ϵ)
+                @log_nonconvergence logcancel!(log, iter, ϵ)
             else
-                @infov 3 logiter!(log, iter, ϵ)
+                @log_iteration logiter!(log, iter, ϵ)
             end
         end
 

src/algorithms/ctmrg/ctmrg.jl

@@ -112,7 +112,7 @@ function leading_boundary(
     )
     check_input(leading_boundary, network, env₀, alg)
     log = ignore_derivatives(() -> MPSKit.IterLog("CTMRG"))
-    return LoggingExtras.withlevel(; alg.verbosity) do
+    return with_verbosity(; alg.verbosity) do
         env = deepcopy(env₀)
         CS, TS = ignore_derivatives() do
             return convergence_spectra(env₀, alg)
@@ -158,16 +158,16 @@ function check_input(::typeof(leading_boundary), network, env, alg::CTMRGAlgorit
 
 # custom CTMRG logging
 function ctmrg_loginit!(log, η, network, env)
-    return @infov 2 loginit!(log, η, network_value(network, env))
+    return @log_initialization loginit!(log, η, network_value(network, env))
 end
 function ctmrg_logiter!(log, iter, η, network, env)
-    return @infov 3 logiter!(log, iter, η, network_value(network, env))
+    return @log_iteration logiter!(log, iter, η, network_value(network, env))
 end
 function ctmrg_logfinish!(log, iter, η, network, env)
-    return @infov 2 logfinish!(log, iter, η, network_value(network, env))
+    return @log_convergence logfinish!(log, iter, η, network_value(network, env))
 end
 function ctmrg_logcancel!(log, iter, η, network, env)
-    return @warnv 1 logcancel!(log, iter, η, network_value(network, env))
+    return @log_nonconvergence logcancel!(log, iter, η, network_value(network, env))
 end
 
 @non_differentiable ctmrg_loginit!(args...)

src/algorithms/optimization/peps_optimization.jl

@@ -253,8 +253,8 @@ function fixedpoint(
     times = Vector{Float64}(undef, 1)
 
     # `tol_state` tracks (iter, gradnorm) of the last accepted optimization step, used to
-    # adjust `alg.boundary_alg`/`alg.gradient_alg` via `MPSKit.updatetol` if they are
-    # wrapped in a `MPSKit.DynamicTol`.
+    # adjust `alg.boundary_alg`/`alg.gradient_alg` via `MPSKit.DynamicTols.adapt_solver`
+    # if they are wrapped in a `MPSKit.DynamicTol`.
     tol_state = (iter = 0, gradnorm = one(T))
 
     # set up finalizer to update `tol_state` and record info after each accepted optimization step
@@ -267,8 +267,9 @@ function fixedpoint(
 
     # initialize the preconditioner
     function precondition(x, g)
-        precondition_alg = updatetol(
-            alg.precondition_alg, tracked_finalizer.tol_state.iter, tracked_finalizer.tol_state.gradnorm
+        precondition_alg = adapt_solver(
+            alg.precondition_alg;
+            iter = tracked_finalizer.tol_state.iter, g_global = tracked_finalizer.tol_state.gradnorm
         )
         return peps_precondition(x, g, tracked_finalizer.tol_state, precondition_alg)
     end
@@ -280,15 +281,15 @@ function fixedpoint(
         hasconverged, shouldstop, (finalize!) = tracked_finalizer,
     ) do (peps, env)
         start_time = time()
-        boundary_alg = updatetol(
-            alg.boundary_alg,
-            tracked_finalizer.tol_state.iter,
-            tracked_finalizer.tol_state.gradnorm
+        boundary_alg = adapt_solver(
+            alg.boundary_alg;
+            iter = tracked_finalizer.tol_state.iter,
+            g_global = tracked_finalizer.tol_state.gradnorm
         )
         # gradient tolerance is scaled relative to the boundary algorithm's own
         # (just-updated) effective tolerance, not directly to the gradient norm
-        gradient_alg = updatetol(
-            alg.gradient_alg, tracked_finalizer.tol_state.iter, boundary_alg.tol
+        gradient_alg = adapt_solver(
+            alg.gradient_alg; iter = tracked_finalizer.tol_state.iter, g_global = boundary_alg.tol
         )
         E, gs = withgradient(peps) do ψ
             env′, info = hook_pullback(

src/algorithms/time_evolution/simpleupdate.jl

@@ -209,9 +209,9 @@ function MPSKit.time_evolve(
         verbosity::Int = 2, check_interval::Int = 500
     )
     @assert check_interval >= 0
-    return LoggingExtras.withlevel(; verbosity) do
+    return with_verbosity(; verbosity) do
         time_start = time()
-        @infov 1 "--- Time evolution (simple update), dt = $(it.dt) ---"
+        @log_initialization "--- Time evolution (simple update), dt = $(it.dt) ---"
         env0, time0 = it.state.env, time()
         for (psi, env, info) in it
             iter = it.state.iter
@@ -220,12 +220,12 @@ function MPSKit.time_evolve(
                 ((iter % check_interval == 0) || (iter == 1) || stop)
             time1 = time()
             if showinfo
-                @infov 2 "Space of x-weight at [1, 1] = $(space(env[1, 1, 1], 1))"
-                @infov 2 @sprintf("SU iter %-7d: |Δλ| = %.3e. Time = %.3f s/it", iter, compare_weights(env0, env), time1 - time0)
+                @log_convergence "Space of x-weight at [1, 1] = $(space(env[1, 1, 1], 1))"
+                @log_convergence @sprintf("SU iter %-7d: |Δλ| = %.3e. Time = %.3f s/it", iter, compare_weights(env0, env), time1 - time0)
             end
             if stop
                 time_end = time()
-                @infov 1 @sprintf("Time evolution finished in %.2f s", time_end - time_start)
+                @log_convergence @sprintf("Time evolution finished in %.2f s", time_end - time_start)
                 return psi, env, info
             else
                 env0 = env
@@ -240,9 +240,9 @@ function MPSKit.time_evolve(
         tol::Float64 = 1.0e-8, verbosity::Int = 2, check_interval::Int = 500
     ) where {G, S <: SUState{<:InfinitePEPS}}
     @assert check_interval >= 0
-    return LoggingExtras.withlevel(; verbosity) do
+    return with_verbosity(; verbosity) do
         time_start = time()
-        @infov 1 "--- Time evolution (simple update), dt = $(it.dt) ---"
+        @log_initialization "--- Time evolution (simple update), dt = $(it.dt) ---"
         @assert it.alg.imaginary_time "Only imaginary time evolution of InfinitePEPS allows convergence checking."
         env0, time0 = it.state.env, time()
         for (psi, env, info) in it
@@ -256,8 +256,8 @@ function MPSKit.time_evolve(
                 # TODO: convert to BPEnv instead
                 ctmenv = CTMRGEnv(env)
                 energy = real(expectation_value(psi, H, ctmenv)) / prod(size(psi))
-                @infov 2 "Space of x-weight at [1, 1] = $(space(env[1, 1, 1], 1))"
-                @infov 2 @sprintf(
+                @log_convergence "Space of x-weight at [1, 1] = $(space(env[1, 1, 1], 1))"
+                @log_convergence @sprintf(
                     "SU iter %-7d: E ≈ %.5f, |Δλ| = %.3e. Time = %.3f s/it",
                     iter, energy, diff, time1 - time0
                 )
@@ -266,11 +266,11 @@ function MPSKit.time_evolve(
                 @warn "SU: bond weights have not converged."
             end
             if diff < tol
-                @infov 2 "SU: bond weights have converged."
+                @log_convergence "SU: bond weights have converged."
             end
             if stop
                 time_end = time()
-                @infov 1 @sprintf("Time evolution finished in %.2f s", time_end - time_start)
+                @log_convergence @sprintf("Time evolution finished in %.2f s", time_end - time_start)
                 return psi, env, info
             else
                 env0 = env