Add PEPS initialization based on fidelity maximization

Project.toml

@@ -41,7 +41,7 @@ Printf = "1"
 QuadGK = "2.11.1"
 Random = "1"
 Statistics = "1"
-TensorKit = "0.14.9"
+TensorKit = "0.14.11"
 TensorOperations = "5"
 TestExtras = "0.3"
 VectorInterface = "0.4, 0.5"

src/Defaults.jl

@@ -73,6 +73,14 @@ Module containing default algorithm parameter values and arguments.
 * `ls_maxfg=$(Defaults.ls_maxfg)` : Maximum number of function evaluations for the line search in each step of the optimization.
 * `lbfgs_memory=$(Defaults.lbfgs_memory)` : Size of limited memory representation of BFGS Hessian matrix.
 
+## Approximate
+
+* `approximate_alg=:$(Defaults.approximate_alg)` : Default `ApproximateAlgorithm` for PEPS approximation.
+* `approximate_tol=$(Defaults.approximate_tol)` : Infidelity tolerance of the approximation iteration.
+* `approximate_maxiter=$(Defaults.approximate_maxiter)` : Maximal number of approximation steps.
+* `approximate_miniter=$(Defaults.approximate_miniter)` : Minimal number of approximation steps.
+* `approximate_verbosity=$(Defaults.approximate_verbosity)` : Approximator output information verbosity.
+
 ## OhMyThreads scheduler
 
 - `scheduler=Ref{Scheduler}(...)` : Multithreading scheduler which can be accessed via `set_scheduler!`.
@@ -125,6 +133,13 @@ const ls_maxiter = 10
 const ls_maxfg = 20
 const lbfgs_memory = 20
 
+# Approximate
+const approximate_alg = :fidelitymaxcrude
+const approximate_tol = 1.0e-2
+const approximate_maxiter = 10
+const approximate_miniter = 3
+const approximate_verbosity = 3
+
 # OhMyThreads scheduler defaults
 const scheduler = Ref{Scheduler}()
 

src/PEPSKit.jl

@@ -80,6 +80,7 @@ include("utility/symmetrization.jl")
 
 include("algorithms/optimization/fixed_point_differentiation.jl")
 include("algorithms/optimization/peps_optimization.jl")
+include("algorithms/optimization/fidelity_initialize.jl")
 
 include("algorithms/select_algorithm.jl")
 
@@ -96,6 +97,7 @@ export correlator, correlation_length
 export leading_boundary
 export PEPSOptimize, GeomSum, ManualIter, LinSolver, EigSolver
 export fixedpoint
+export FidelityMaxCrude, approximate
 
 export absorb_weight
 export ALSTruncation, FullEnvTruncation

src/algorithms/optimization/fidelity_initialize.jl

@@ -0,0 +1,248 @@
+"""
+$(TYPEDEF)
+
+Abstract super type for tensor approximation algorithms.
+"""
+abstract type ApproximateAlgorithm end
+
+const APPROXIMATE_SYMBOLS = IdDict{Symbol, Type{<:ApproximateAlgorithm}}()
+
+"""
+$(TYPEDEF)
+
+PEPS approximation algorithm maximizing the fidelity by successively applying the fidelity
+derivative with respect to the approximator PEPS.
+
+## Constructors
+
+    FidelityMaxCrude(; kwargs...)
+
+Construct the approximation algorithm from the following keyword arguments:
+
+* `tol::Float64=$(Defaults.approximate_tol)` : Infidelity tolerance of the approximation iteration.
+* `maxiter::Int=$(Defaults.approximate_maxiter)` : Maximal number of approximation steps.
+* `miniter::Int=$(Defaults.approximate_miniter)` : Minimal number of approximation steps.
+* `verbosity::Int=$(Defaults.approximate_verbosity)` : Approximator output information verbosity.
+* `boundary_alg::Union{<:CTMRGAlgorithm,NamedTuple}` : CTMRG algorithm used for contracting the norm and fidelity networks.
+"""
+struct FidelityMaxCrude <: ApproximateAlgorithm
+    tol::Float64
+    maxiter::Int
+    miniter::Int
+    verbosity::Int
+    boundary_alg::CTMRGAlgorithm
+end
+function FidelityMaxCrude(; kwargs...)
+    return ApproximateAlgorithm(; alg = :fidelitymaxcrude, kwargs...)
+end
+
+APPROXIMATE_SYMBOLS[:fidelitymaxcrude] = FidelityMaxCrude
+
+"""
+    ApproximateAlgorithm(; kwargs...)
+
+Keyword argument parser returning the appropriate `ApproximateAlgorithm` algorithm struct.
+"""
+function ApproximateAlgorithm(;
+        alg = Defaults.approximate_alg,
+        tol = Defaults.approximate_tol,
+        maxiter = Defaults.approximate_maxiter, miniter = Defaults.approximate_miniter,
+        verbosity = Defaults.approximate_verbosity,
+        boundary_alg = (; verbosity = max(1, verbosity - 2)),  # shouldn't be smaller than one by default
+    )
+    # replace symbol with projector alg type
+    haskey(APPROXIMATE_SYMBOLS, alg) || throw(ArgumentError("unknown approximate algorithm: $alg"))
+    alg_type = APPROXIMATE_SYMBOLS[alg]
+
+    boundary_algorithm = _alg_or_nt(CTMRGAlgorithm, boundary_alg)
+
+    return alg_type(tol, maxiter, miniter, verbosity, boundary_algorithm)
+end
+
+"""
+    single_site_fidelity_initialize(
+        peps::InfinitePEPS, [bondspace = _maxspace(peps₀)]; kwargs...
+    )
+
+Generate a single-site unit cell PEPS from a (possibly) multi-site `peps` by approximating
+the respective entries using [`approximate!`](@ref). By default, the maximal bond space of
+`peps₀` is used for all virtual legs of the single-site PEPS.
+
+## Keyword arguments
+
+* `noise_amp=1.0-1` : Gaussian noise amplitude of initial single-site PEPS
+
+All additional keyword arguments will be passed to the [`approximate!`](@ref) call.
+"""
+function single_site_fidelity_initialize(
+        peps::InfinitePEPS, bondspace = _spacemax(peps);
+        noise_amp = 1.0e-1, kwargs...
+    )
+    @assert allequal(map(p -> space(p, 1), unitcell(peps))) "PEPS must have uniform physical spaces"
+
+    physspace = space(unitcell(peps)[1], 1)
+    peps_single = noise_amp * InfinitePEPS(randn, scalartype(peps), physspace, bondspace) # single-site unit cell with random noise
+
+    # absorb peps₀ tensors into single-site tensors in-place
+    peps_uc = InfinitePEPS(fill(only(unitcell(peps_single)), size(peps))) # fill unit cell with peps_single tensors
+    absorb!(peps_uc[1], peps[1]) # absorb (1, 1) tensor of peps₀ (applies to all peps_uc entries since absorb! is mutating)
+    peps_single, = approximate(peps_uc, peps; kwargs...)
+
+    return InfinitePEPS([peps_single[1];;])
+end
+
+# maximal virtual space over unit cell
+function _spacemax(peps::InfinitePEPS)
+    return reduce(supremum, map(p -> supremum(domain(p)[1], domain(p)[2]), unitcell(peps)))
+end
+
+@doc """
+    approximate(ψ₀::InfinitePEPS, ψ::InfinitePEPS; kwargs...)
+    # expert version
+    approximate(ψ₀::InfinitePEPS, ψ::InfinitePEPS, alg::ApproximateAlgorithm)
+
+Approximate `ψ` from the initial guess `ψ₀`. The approximation algorithm is specified via
+the keyword arguments or directly by passing an [`ApproximateAlgorithm`](@ref) struct.
+
+## Keyword arguments
+
+* `alg::Symbol=:$(Defaults.approximate_alg)` : Approximation algorithm, which can be one of the following:
+    - `:fidelitymaxcrude` : Maximize the fidelity from the fidelity gradient in a power-method fashion.
+* `tol::Float64=$(Defaults.approximate_tol)` : Infidelity tolerance of the approximation iteration.
+* `maxiter::Int=$(Defaults.approximate_maxiter)` : Maximal number of approximation steps.
+* `miniter::Int=$(Defaults.approximate_miniter)` : Minimal number of approximation steps.
+* `verbosity::Int=$(Defaults.approximate_verbosity)` : Approximator output information verbosity.
+* `boundary_alg::Union{<:CTMRGAlgorithm,NamedTuple}` : CTMRG algorithm used for contracting the norm and fidelity networks.
+
+## Return values
+
+The final approximator and its environment are returned.
+"""
+approximate
+
+
+function MPSKit.approximate(ψ₀::InfinitePEPS, ψ::InfinitePEPS; kwargs...)
+    alg = ApproximateAlgorithm(; kwargs...)
+    return approximate(ψ₀, ψ, alg)
+end
+function MPSKit.approximate(ψ₀::InfinitePEPS, ψ::InfinitePEPS, alg::FidelityMaxCrude)
+    @assert size(ψ₀) == size(ψ) "incompatible unit cell sizes"
+    @assert all(map((p₀, p) -> space(p₀, 1) == space(p, 1), unitcell(ψ₀), unitcell(ψ))) "incompatible physical spaces"
+
+    log = MPSKit.IterLog("Approximate")
+    return LoggingExtras.withlevel(; alg.verbosity) do
+        # normalize reference PEPS
+        peps_init = ψ # smaller bond spaces
+        envspace = domain(ψ₀[1])[1]
+        env₀, = leading_boundary(CTMRGEnv(peps_init, envspace), peps_init, alg.boundary_alg)
+        peps_init /= sqrt(abs(_local_norm(peps_init, peps_init, env₀))) # normalize to ensure that fidelity is bounded by 1
+
+        # normalize maximizer PEPS
+        peps = ψ₀
+        env, = leading_boundary(CTMRGEnv(peps, envspace), peps, alg.boundary_alg)
+        peps /= sqrt(abs(_local_norm(peps, peps, env)))
+
+        approximate_loginit!(log, one(real(scalartype(peps))), zero(real(scalartype(peps))))
+        nw₀ = InfiniteSquareNetwork(peps_init, peps) # peps₀ has different virtual spaces than peps
+        envnw, = leading_boundary(CTMRGEnv(nw₀, envspace), nw₀, alg.boundary_alg)
+        peps′ = _∂local_norm(peps_init, envnw)
+        for iter in 1:(alg.maxiter)
+            # compute fidelity from ∂norm
+            fid = abs2(_local_norm(peps, peps′))
+            infid = 1 - fid
+            if abs(infid) ≤ alg.tol && iter ≥ alg.miniter
+                approximate_logfinish!(log, iter, infid, fid)
+                break
+            end
+            if iter == alg.maxiter
+                approximate_logcancel!(log, iter, infid, fid)
+                break
+            else
+                approximate_logiter!(log, iter, infid, fid)
+            end
+
+            # contract boundary of fidelity network
+            # initialize CTMRG on environment of peps′ (must have matching virtual spaces!)
+            envnw, = leading_boundary(env, InfiniteSquareNetwork(peps, peps′), alg.boundary_alg)
+            ∂norm = _∂local_norm(peps, envnw)
+
+            # renormalize current PEPS
+            peps = peps′
+            env, = leading_boundary(env, peps, alg.boundary_alg)
+            peps /= sqrt(abs(_local_norm(peps, peps, env)))
+
+            peps′ = ∂norm
+        end
+
+        return peps, env
+    end
+end
+
+# custom fidelity maximization logging
+function approximate_loginit!(log, infid, fid)
+    return @infov 2 loginit!(log, infid, fid)
+end
+function approximate_logiter!(log, iter, infid, fid)
+    return @infov 3 logiter!(log, iter, infid, fid)
+end
+function approximate_logfinish!(log, iter, infid, fid)
+    return @infov 2 logfinish!(log, iter, infid, fid)
+end
+function approximate_logcancel!(log, iter, infid, fid)
+    return @warnv 1 logcancel!(log, iter, infid, fid)
+end
+
+
+"""
+$(SIGNATURES)
+
+Sum over `contract_local_norm` values of all unit cell entries.
+"""
+function _local_norm(ket::InfinitePEPS, bra::InfinitePEPS, env::CTMRGEnv)
+    return sum(ind -> contract_local_norm((ind,), ket, bra, env), eachcoordinate(ket))
+end
+function _local_norm(peps::InfinitePEPS, ∂norm::InfinitePEPS)
+    return sum(eachcoordinate(peps)) do (r, c)
+        @tensor conj(peps[r, c][d; D_N D_E D_S D_W]) * ∂norm[r, c][d; D_N D_E D_S D_W]
+    end
+end
+
+
+"""
+$(SIGNATURES)
+
+Compute the `InfinitePEPS` resulting from removing the bra PEPS tensors in `_local_norm`.
+"""
+function _∂local_norm(peps::InfinitePEPS, env::CTMRGEnv)
+    return InfinitePEPS(map(ind -> _∂contract_site(ind, peps, env), eachcoordinate(peps)))
+end
+
+# contract CTMRG environment leaving open the bra-PEPS virtual and physical bonds
+function _∂contract_site(ind::Tuple{Int, Int}, peps::InfinitePEPS, env::CTMRGEnv)
+    r, c = ind
+    return _∂contract_site(
+        env.corners[NORTHWEST, _prev(r, end), _prev(c, end)],
+        env.corners[NORTHEAST, _prev(r, end), _next(c, end)],
+        env.corners[SOUTHEAST, _next(r, end), _next(c, end)],
+        env.corners[SOUTHWEST, _next(r, end), _prev(c, end)],
+        env.edges[NORTH, _prev(r, end), c], env.edges[EAST, r, _next(c, end)],
+        env.edges[SOUTH, _next(r, end), c], env.edges[WEST, r, _prev(c, end)],
+        peps[r, c],
+    )
+end
+function _∂contract_site(
+        C_northwest, C_northeast, C_southeast, C_southwest,
+        E_north::CTMRG_PEPS_EdgeTensor, E_east::CTMRG_PEPS_EdgeTensor,
+        E_south::CTMRG_PEPS_EdgeTensor, E_west::CTMRG_PEPS_EdgeTensor, ψ,
+    )
+    return @autoopt @tensor peps′[d; D_N_below D_E_below D_S_below D_W_below] :=
+        E_west[χ_WSW D_W_above D_W_below; χ_WNW] *
+        C_northwest[χ_WNW; χ_NNW] *
+        E_north[χ_NNW D_N_above D_N_below; χ_NNE] *
+        C_northeast[χ_NNE; χ_ENE] *
+        E_east[χ_ENE D_E_above D_E_below; χ_ESE] *
+        C_southeast[χ_ESE; χ_SSE] *
+        E_south[χ_SSE D_S_above D_S_below; χ_SSW] *
+        C_southwest[χ_SSW; χ_WSW] *
+        ψ[d; D_N_above D_E_above D_S_above D_W_above]
+end

test/runtests.jl

@@ -69,6 +69,9 @@ end
         @time @safetestset "Density matrix from double-layer PEPO" begin
             include("toolbox/densitymatrices.jl")
         end
+        @time @safetestset "Fidelity initialization" begin
+            include("toolbox/fidelity_initialize.jl")
+        end
     end
     if GROUP == "ALL" || GROUP == "UTILITY"
         @time @safetestset "LocalOperator" begin

test/toolbox/fidelity_initialize.jl

@@ -0,0 +1,63 @@
+using TensorKit
+using PEPSKit
+using PEPSKit: single_site_fidelity_initialize, _spacemax
+using Random
+using Test
+
+@testset "Fidelity initialization with single_site_fidelity_initialize" begin
+    Random.seed!(18092025)
+
+    peps_1x1 = InfinitePEPS(randn, ComplexF64, 2, 2)
+    peps_single1 = single_site_fidelity_initialize(
+        peps_1x1, ℂ^3; noise_amp = 0.3, tol = 1.0e-2, miniter = 5,
+        boundary_alg = (; tol = 1.0e-6, trscheme = truncdim(10), verbosity = 1)
+    )
+    @test peps_single1 isa InfinitePEPS
+    @test size(peps_single1) == (1, 1)
+
+    unitcell = (3, 3)
+    Pspaces = fill(2, unitcell...)
+    Nspaces = rand(2:4, unitcell...)
+    Espaces = rand(2:4, unitcell...)
+    peps_3x3 = InfinitePEPS(randn, ComplexF64, Pspaces, Nspaces, Espaces)
+
+    peps_single2 = single_site_fidelity_initialize(
+        peps_3x3; maxiter = 10, noise_amp = 0.1,
+        boundary_alg = (; tol = 1.0e-6, maxiter = 20, verbosity = 1)
+    )
+    @test peps_single2 isa InfinitePEPS
+    @test size(peps_single2) == (1, 1)
+end
+
+@testset "approximate methods with FidelityMaxCrude" begin
+    peps1 = InfinitePEPS(2, 2)
+    peps2 = InfinitePEPS(2, 2)
+    Random.seed!(1234)
+    peps_approx1, = approximate(
+        peps1, peps2; maxiter = 3,
+        boundary_alg = (; trscheme = truncdim(6), verbosity = 1)
+    )
+    Random.seed!(1234)
+    peps_approx2, = approximate(
+        peps1, peps2, FidelityMaxCrude(;
+            maxiter = 3,
+            boundary_alg = (; trscheme = truncdim(6), verbosity = 1)
+        )
+    )
+
+    @test peps_approx1 ≈ peps_approx2
+end
+
+@testset "_spacemax for non-uniform unit cell and symmetry sectors" begin
+    Pspaces = fill(Z2Space(0 => 1, 1 => 1), (2, 2))
+    Nspaces = [
+        Z2Space(0 => 2, 1 => 3) Z2Space(0 => 2, 1 => 5)
+        Z2Space(0 => 4, 1 => 3) Z2Space(0 => 2, 1 => 1)
+    ]
+    Espaces = [
+        Z2Space(0 => 6, 1 => 3) Z2Space(0 => 2, 1 => 1)
+        Z2Space(0 => 2, 1 => 3) Z2Space(0 => 2, 1 => 4)
+    ]
+    peps = InfinitePEPS(Pspaces, Nspaces, Espaces)
+    @test _spacemax(peps) == Z2Space(0 => 6, 1 => 5)
+end