diff --git a/src/Defaults.jl b/src/Defaults.jl index 13ba67ce2..7c7453c74 100644 --- a/src/Defaults.jl +++ b/src/Defaults.jl @@ -17,6 +17,18 @@ Module containing default algorithm parameter values and arguments. * `ctmrg_tol_max=$(Defaults.ctmrg_tol_max)` : Maximal CTMRG tolerance used by `ctmrg_dynamic_tols`. * `ctmrg_tol_factor=$(Defaults.ctmrg_tol_factor)` : Tolerance scaling factor used by `ctmrg_dynamic_tols`. +## Boundary contraction + +* `boundary_alg=:$(Defaults.boundary_alg)` : Default algorithm family used to contract a network, see [`PEPSKit.BoundaryAlgorithm`](@ref). + - `:SimultaneousCTMRG`, `:SequentialCTMRG`, `:C4vCTMRG` : CTMRG variants. + - `:SymmetricBoundaryMPS` : Boundary MPS contraction of a fully symmetric single-site network. + +## Boundary MPS + +* `boundarymps_mps_alg=:$(Defaults.boundarymps_mps_alg)` : Default MPS optimization algorithm driving a boundary MPS contraction. + - `:VUMPS` : Variational uniform MPS. + - `:VOMPS` : Variational optimization of the MPS through MPO-MPS overlap maximization. + ## SVD forward & reverse * `trunc=:$(Defaults.trunc)` : Truncation scheme for SVDs and other decompositions. @@ -150,6 +162,12 @@ const ctmrg_tol_min = 1.0e-12 const ctmrg_tol_max = 1.0e-4 const ctmrg_tol_factor = 1.0e-3 +# Boundary contraction +const boundary_alg = :SimultaneousCTMRG # ∈ {:SimultaneousCTMRG, :SequentialCTMRG, :C4vCTMRG, :SymmetricBoundaryMPS} + +# Boundary MPS +const boundarymps_mps_alg = :VUMPS # ∈ {:VUMPS, :VOMPS} + # SVD forward & reverse const trunc = :FixedSpaceTruncation # ∈ {:FixedSpaceTruncation, :notrunc, :truncerror, :truncspace, :trunctol} const rrule_degeneracy_atol = 1.0e-13 diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 827a16689..e58d66b82 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -84,6 +84,7 @@ include("operators/models.jl") include("environments/ctmrg_environments.jl") include("environments/vumps_environments.jl") +include("environments/boundarymps_environments.jl") include("environments/suweight.jl") include("environments/bp_environments.jl") include("environments/product_state_environments.jl") @@ -105,6 +106,7 @@ include("algorithms/contractions/absorb.jl") include("algorithms/contractions/absorb_weight.jl") include("algorithms/contractions/transfer.jl") include("algorithms/contractions/vumps_contractions.jl") +include("algorithms/contractions/boundarymps_contractions.jl") include("algorithms/contractions/bp_messages.jl") include("algorithms/contractions/local_patch/expr_utils.jl") include("algorithms/contractions/local_patch/network_expr.jl") @@ -120,6 +122,8 @@ include("algorithms/contractions/correlator/peps.jl") include("algorithms/contractions/correlator/pepo_purified.jl") include("algorithms/contractions/correlator/pepo_1layer.jl") +include("algorithms/boundary_algorithm.jl") + include("algorithms/ctmrg/sparse_environments.jl") include("algorithms/ctmrg/ctmrg.jl") include("algorithms/ctmrg/projectors/projectors.jl") @@ -161,6 +165,10 @@ include("algorithms/expectation_value/correlator_adapters.jl") include("algorithms/expectation_value/correlators.jl") include("algorithms/toolbox.jl") +include("algorithms/boundarymps/symmetric_boundarymps.jl") +include("algorithms/boundarymps/characteristic_equations.jl") +include("algorithms/boundarymps/observables.jl") + include("algorithms/optimization/implicit_differentiation.jl") include("algorithms/optimization/preconditioning.jl") include("algorithms/optimization/peps_optimization.jl") @@ -177,6 +185,7 @@ export corner, edge, setcorner!, setedge! export FixedSpaceTruncation, SiteDependentTruncation export HalfInfiniteProjector, FullInfiniteProjector export C4vCTMRG, C4vEighProjector, C4vQRProjector +export SymmetricBoundaryMPS, SymmetricBoundaryMPSEnv export initialize_random_c4v_env, initialize_singlet_c4v_env export LocalOperator, physicalspace export product_peps diff --git a/src/algorithms/boundary_algorithm.jl b/src/algorithms/boundary_algorithm.jl new file mode 100644 index 000000000..aca445027 --- /dev/null +++ b/src/algorithms/boundary_algorithm.jl @@ -0,0 +1,45 @@ +""" +$(TYPEDEF) + +Abstract super type for all algorithms that contract an infinite square network by computing +a boundary fixed point, such as CTMRG and boundary MPS algorithms. + +This is the type of the `boundary_alg` field of [`PEPSOptimize`](@ref), and hence the +supertype every algorithm which can be used to contract a network during a variational +optimization must belong to. +""" +abstract type BoundaryAlgorithm end + +const BOUNDARY_ALGORITHM_SYMBOLS = IdDict{Symbol, Type{<:BoundaryAlgorithm}}() + +""" + _tol(alg::BoundaryAlgorithm) + +Effective convergence tolerance of a boundary contraction algorithm. Defaults to the `tol` +field; algorithms which keep their tolerance elsewhere (such as boundary MPS algorithms, +where it lives on the wrapped MPS optimization algorithm) must overload this. +""" +_tol(alg::BoundaryAlgorithm) = alg.tol + +""" + BoundaryAlgorithm(; alg=:$(Defaults.boundary_alg), kwargs...) + +Keyword argument parser returning the appropriate [`BoundaryAlgorithm`](@ref) struct, where +`alg` selects the boundary contraction *family*: + +* `:SimultaneousCTMRG`, `:SequentialCTMRG`, `:C4vCTMRG` : dispatch to [`CTMRGAlgorithm`](@ref) +* `:SymmetricBoundaryMPS` : dispatch to [`SymmetricBoundaryMPS`](@ref) + +All remaining keyword arguments are forwarded to the corresponding parser. Note that the +underlying MPS optimization algorithm of a boundary MPS contraction is *not* selected here, +but through that parser's own `mps_alg` keyword, since several boundary MPS families can be +driven by the same MPS algorithm. +""" +function BoundaryAlgorithm(; alg = Defaults.boundary_alg, kwargs...) + # CTMRG variants go through their own parser, which handles the projector and + # decomposition keyword arguments + haskey(CTMRG_SYMBOLS, alg) && return CTMRGAlgorithm(; alg, kwargs...) + haskey(BOUNDARY_ALGORITHM_SYMBOLS, alg) || + throw(ArgumentError("unknown boundary algorithm: $alg")) + return BOUNDARY_ALGORITHM_SYMBOLS[alg](; kwargs...) +end diff --git a/src/algorithms/boundarymps/characteristic_equations.jl b/src/algorithms/boundarymps/characteristic_equations.jl new file mode 100644 index 000000000..260fd9e17 --- /dev/null +++ b/src/algorithms/boundarymps/characteristic_equations.jl @@ -0,0 +1,68 @@ +# +# Characteristic equation used in implicit differentiation of symmetric boundary MPS +# contractions +# + +""" + generate_boundary_mps_characteristic_equation( + env::SymmetricBoundaryMPSEnv, (VLfp, VRfp)::Tuple{<:EdgeTensor, <:RightProjector} + ) + +Takes the fixed-point values of a converged symmetric boundary MPS contraction, given by the +environment `env`, along with the left null space `VLfp` of its left-gauged MPS tensor and +the right null space `VRfp` of its right-gauged MPS tensor, and generates a function +``F(s, l, r, c, gl, gr)`` which characterizes the convergence of the boundary MPS contraction +in terms of the characteristic equation ``F(s, l, r, c, gl, gr) = 0``. + +Here, ``s`` corresponds to a state variable (e.g. an `InfinitePEPS` that is being optimized), +``c``, ``gl`` and ``gr`` directly represent the MPS bond tensor and the left and right +environments, while ``l`` and ``r`` parametrize differentiable left- and +right-gauged MPS tensors as ``AL = AL_{fp} + V_{L,fp} * l`` and +``AR = AR_{fp} + (r * V_{R,fp})``. + +``F`` returns a tuple of five tensors, corresponding to an equation for each of ``l``, ``r``, +``c``, ``gl`` and ``gr``. The first three are obtained by projecting the effective site +operator applied to the center-gauged MPS tensor onto the respective tangent directions, and +the last two express that the environments are fixed points of the MPS-network-MPS transfer +matrix. All equations are normalized by the eigenvalue ``λ`` of the effective site operator. + +See also [`generate_symmetric_characteristic_equation`](@ref). +""" +function generate_boundary_mps_characteristic_equation( + env::SymmetricBoundaryMPSEnv, (VLfp, VRfp)::Tuple{TE, TP} + ) where {TE <: EdgeTensor, TP <: RightProjector} + ALfp, ARfp, Cfp, = _unpack(env) + + # constant preconditioner + # NOTE: this relies on the bond tensor having been diagonalized, which + # `leading_boundary` takes care of + iCfp = sdiag_pow(real(DiagonalTensorMap(Cfp)), -1) + + function boundary_mps_characteristic_equation(state, l, r, c, gl, gr) + network = InfiniteSquareNetwork(state) + O = network[1, 1] + + # prepare appropriately parametrized MPS tensors + AL = ALfp + VLfp * l + AR = ARfp + repartition_left(r * VRfp) + ARR = repartition_right(AR) # 'right isometry' form + + # construct the center tensor in a symmetric way + AC = (absorb_right(AL, c) + absorb_left(AR, c)) / 2 + + # main partial contractions to reuse + AC´L = ∂AC(AC, gl, O, gr) # 'left isometry' form + AC´R = repartition_right(AC´L) # 'right isometry' form + λ = dot(AC, AC´L) + + F1 = VLfp' * AC´L * iCfp / λ - l + F2 = iCfp * AC´R * VRfp' / λ - r + F3 = (AL' * AC´L + AC´R * ARR') / (2 * λ) - c + F4 = MPSKit.transfer_left(gl, O, AL, AL) / λ - gl + F5 = MPSKit.transfer_right(gr, O, AR, AR) / λ - gr + + return F1, F2, F3, F4, F5 + end + + return boundary_mps_characteristic_equation +end diff --git a/src/algorithms/boundarymps/observables.jl b/src/algorithms/boundarymps/observables.jl new file mode 100644 index 000000000..4b3e594d2 --- /dev/null +++ b/src/algorithms/boundarymps/observables.jl @@ -0,0 +1,34 @@ +# +# Observables evaluated using a symmetric boundary MPS environment +# + +""" + network_value(network::InfiniteSquareNetwork, env::SymmetricBoundaryMPSEnv) + +Return the value (per unit cell) of a contractible network contracted using a symmetric +boundary MPS environment. +""" +function network_value(network::InfiniteSquareNetwork, env::SymmetricBoundaryMPSEnv) + size(network) == (1, 1) || + throw(ArgumentError("symmetric boundary MPS environments require a single-site unit cell")) + AC = get_AC(env) + AC´ = PEPS_AC_Hamiltonian(env.GL, network[1, 1], env.GR) * AC + return dot(AC, AC´) +end +function network_value(state, env::SymmetricBoundaryMPSEnv) + return network_value(InfiniteSquareNetwork(state), env) +end + +function LinearAlgebra.norm(peps::InfinitePEPS, env::SymmetricBoundaryMPSEnv) + return network_value(InfiniteSquareNetwork(peps), env) +end + +## Partition function tensor insertions + +function contract_local_tensor( + ::Union{CartesianIndex{2}, Tuple{Int, Int}}, O::PartitionFunctionTensor, + env::SymmetricBoundaryMPSEnv, + ) + # the index is irrelevant here: symmetric boundary MPS environments are single-site + return _contract_site(get_AC(env), env.GL, env.GR, O) +end diff --git a/src/algorithms/boundarymps/symmetric_boundarymps.jl b/src/algorithms/boundarymps/symmetric_boundarymps.jl new file mode 100644 index 000000000..44876cd0a --- /dev/null +++ b/src/algorithms/boundarymps/symmetric_boundarymps.jl @@ -0,0 +1,204 @@ +# +# Boundary MPS contraction of networks with a single-site unit cell and full spatial symmetry +# + +# MPS optimization algorithms which can drive a boundary MPS contraction; these are the +# workhorse of every boundary MPS family, and are selected separately from the family itself +# +# NOTE: `MPSKit.GradientGrassmann` is deliberately absent. Its statmech objective function is +# `-log(real(⟨ψ|O|ψ⟩))`, which is only meaningful for a Hermitian transfer matrix; for a +# generic network the expectation value is complex along the line search and its real part +# can turn negative, throwing a `DomainError` out of `log`. It can still be used by passing +# an instance as `mps_alg` for networks where the transfer matrix is known to be Hermitian. +const MPS_ALGORITHM_SYMBOLS = IdDict{Symbol, Type{<:MPSKit.Algorithm}}( + :VUMPS => VUMPS, :VOMPS => VOMPS, +) + +# add algorithm-specific keyword arguments to the MPS algorithm kwargs if needed +_pad_mps_kwargs(::Type, mps_kwargs) = mps_kwargs +function _pad_mps_kwargs(::Type{<:VUMPS}, mps_kwargs) + # the network transfer matrix is not Hermitian, so neither is the effective eigenvalue + # problem solved at every VUMPS iteration + return (; alg_eigsolve = MPSKit.Defaults.alg_eigsolve(; ishermitian = false), mps_kwargs...) +end + +""" +$(TYPEDEF) + +Algorithm for contracting an infinite square network with a single-site unit cell which is +invariant under rotations and Hermitian reflections, using a uniform boundary MPS. + +The actual contraction is carried out by an MPSKit MPS optimization algorithm which is +wrapped by this struct, and which acts on the row-to-row transfer matrix of the network. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + SymmetricBoundaryMPS(; kwargs...) + SymmetricBoundaryMPS(mps_alg::MPSKit.Algorithm) + +Construct a symmetric boundary MPS algorithm either from an MPSKit MPS optimization +algorithm directly, or based on the following keyword arguments: + +* `tol::Real=$(Defaults.ctmrg_tol)` : Convergence tolerance of the boundary MPS contraction. +* `maxiter::Int=$(Defaults.ctmrg_maxiter)` : Maximal number of boundary MPS iterations. +* `verbosity::Int=$(Defaults.ctmrg_verbosity)` : Output information verbosity. +* `mps_alg::Union{Symbol,NamedTuple,MPSKit.Algorithm}=(; alg::Symbol=:$(Defaults.boundarymps_mps_alg))` : MPS optimization algorithm driving the contraction, where `alg` can be one of the following: + - `:VUMPS` : Variational uniform MPS, see [`MPSKit.VUMPS`](@extref) for details. + - `:VOMPS` : Variational optimization of the MPS through MPO-MPS overlap maximization, see [`MPSKit.VOMPS`](@extref) for details. + + A bare `Symbol` is shorthand for `(; alg = symbol)`; any further entries of the `NamedTuple` + are passed on to the MPS algorithm constructor and override `tol`, `maxiter` and + `verbosity`. Supplying an `MPSKit.Algorithm` instance uses it as is, in which case + `tol`, `maxiter` and `verbosity` are ignored. + +!!! note + The row-to-row transfer matrix of a network is generally not Hermitian, which restricts + which MPS optimization algorithms are applicable. The eigensolver used by + [`MPSKit.VUMPS`](@extref) must be able to handle non-Hermitian effective operators, so + the default constructed here sets `ishermitian = false` accordingly. For the same reason + [`MPSKit.GradientGrassmann`](@extref) is not offered as a `mps_alg` symbol: its + objective function assumes a Hermitian transfer matrix and errors otherwise. +""" +struct SymmetricBoundaryMPS{A} <: BoundaryAlgorithm + "wrapped MPSKit MPS optimization algorithm" + alg::A +end +BOUNDARY_ALGORITHM_SYMBOLS[:SymmetricBoundaryMPS] = SymmetricBoundaryMPS + +function SymmetricBoundaryMPS(; + tol = Defaults.ctmrg_tol, + maxiter = Defaults.ctmrg_maxiter, + verbosity = Defaults.ctmrg_verbosity, + mps_alg = (;), + ) + return SymmetricBoundaryMPS(_mps_algorithm(mps_alg; tol, maxiter, verbosity)) +end + +""" + _mps_algorithm(mps_alg; tol, maxiter, verbosity) + +Parse the `mps_alg` keyword argument of a boundary MPS algorithm into an MPSKit MPS +optimization algorithm. Accepts a `Symbol`, a `NamedTuple` carrying an `alg` symbol along +with further algorithm keyword arguments, or an `MPSKit.Algorithm` instance which is +returned unchanged. +""" +_mps_algorithm(mps_alg::MPSKit.Algorithm; kwargs...) = mps_alg +_mps_algorithm(mps_alg::Symbol; kwargs...) = _mps_algorithm((; alg = mps_alg); kwargs...) +function _mps_algorithm(mps_alg::NamedTuple; tol, maxiter, verbosity) + mps_kwargs = (; alg = Defaults.boundarymps_mps_alg, tol, maxiter, verbosity, mps_alg...) + + haskey(MPS_ALGORITHM_SYMBOLS, mps_kwargs.alg) || + throw(ArgumentError("unknown MPS optimization algorithm: $(mps_kwargs.alg)")) + alg_type = MPS_ALGORITHM_SYMBOLS[mps_kwargs.alg] + + # pad kwargs based on algorithm requirements and remove the `alg` keyword argument + mps_kwargs = _pad_mps_kwargs(alg_type, Base.structdiff(mps_kwargs, (; alg = nothing))) + + return alg_type(; mps_kwargs...) +end + +_tol(alg::SymmetricBoundaryMPS) = alg.alg.tol +_tol(alg::SymmetricBoundaryMPS{<:MPSKit.GradientGrassmann}) = alg.alg.method.gradtol + +# `SymmetricBoundaryMPS` has no top-level `tol` field (it lives on the wrapped algorithm), so +# the default `MPSKit.DynamicTols._updatetol` (which sets `alg.tol`) doesn't apply +_updatetol(alg::SymmetricBoundaryMPS, tol::Real) = @set alg.alg.tol = tol +function _updatetol(alg::SymmetricBoundaryMPS{<:MPSKit.GradientGrassmann}, tol::Real) + return @set alg.alg.method.gradtol = tol +end + +# +## contraction +# + +""" + diagonalize_center(AL, AR, C, GL, GR) + +Gauge a symmetric boundary MPS environment such that its bond tensor is diagonal, by +absorbing the singular vectors of `C` into the neighboring tensors. +""" +function diagonalize_center( + AL::TE, AR::TE, C::TC, GL::TE, GR::TE + ) where {TE <: EdgeTensor, TC <: CornerTensor} + U, C´, V = svd_compact(C) + AL´ = absorb_left_right(AL, U', U) + AR´ = absorb_left_right(AR, V, V') + GL´ = absorb_left_right(GL, U', U) + GR´ = absorb_left_right(GR, V, V') + return AL´, AR´, C´, GL´, GR´ +end + +""" + leading_boundary(env₀::SymmetricBoundaryMPSEnv, network; kwargs...) -> env, info + # expert version: + leading_boundary(env₀::SymmetricBoundaryMPSEnv, network, alg::SymmetricBoundaryMPS) + +Contract a single-site `network` which is invariant under rotations and Hermitian +reflections using a uniform boundary MPS, and return the resulting environment. + +## Return values + +* `env` : The final environment. +* `info` : A `NamedTuple` containing information about the contraction, with fields + `converged`, `convergence_error`, `contraction_metrics` and the network value `N`. +""" +function MPSKit.leading_boundary( + env₀::SymmetricBoundaryMPSEnv, network::InfiniteSquareNetwork, + alg::SymmetricBoundaryMPS, + ) + size(network) == (1, 1) || + throw(ArgumentError("symmetric boundary MPS contraction requires a single-site unit cell")) + + # convert to MPSKit language + mps₀ = InfiniteMPS([env₀.AL], env₀.C) + envs₀ = MPSKit.InfiniteEnvironments( + PeriodicVector([env₀.GL]), PeriodicVector([env₀.GR]) + ) + O = InfiniteMPO(unitcell(network)[1, :]) + + # run the actual boundary MPS algorithm + mps, envs, ϵ = MPSKit.leading_boundary(mps₀, O, alg.alg, envs₀) + + # unpack + AL, AR, C = only(mps.AL), only(mps.AR), only(mps.C) + GL, GR = only(envs.GLs), only(envs.GRs) + + # diagonalize the bond tensor (optional, but assumed in the current implementation of the characteristic equations) + AL, AR, C, GL, GR = diagonalize_center(AL, AR, C, GL, GR) + + # HACK: keep pretending the bond tensor is a dense and (possibly) complex tensor, such + # that backpropagation through subsequent observable evaluations yields dense and + # complex bond tensor cotangents, which implicit differentiation requires here + C = TensorMap(C) + if real(scalartype(network)) != scalartype(network) + C = complex(C) + end + + env = SymmetricBoundaryMPSEnv(AL, AR, C, GL, GR) + N = network_value(network, env) + + info = (; + converged = ϵ < _tol(alg), + convergence_error = ϵ, + contraction_metrics = (;), + N, + ) + + return env, info +end +function MPSKit.leading_boundary( + env₀::SymmetricBoundaryMPSEnv, network::InfiniteSquareNetwork; kwargs... + ) + return MPSKit.leading_boundary( + env₀, network, select_algorithm(leading_boundary, env₀; kwargs...) + ) +end +function MPSKit.leading_boundary(env₀::SymmetricBoundaryMPSEnv, state, args...; kwargs...) + return MPSKit.leading_boundary( + env₀, InfiniteSquareNetwork(state), args...; kwargs... + ) +end diff --git a/src/algorithms/contractions/boundarymps_contractions.jl b/src/algorithms/contractions/boundarymps_contractions.jl new file mode 100644 index 000000000..f00f4cae1 --- /dev/null +++ b/src/algorithms/contractions/boundarymps_contractions.jl @@ -0,0 +1,352 @@ +# +# Contractions used in boundary MPS contractions of infinite square networks +# + +## Repartitions + +# NOTE: duplicates of MPSKit._transpose_front and MPSKit._transpose_tail internals + +""" + repartition_right(A::EdgeTensor; copy = true) + +Move the physical legs of an (N, 1) tensor map representing a left isometry from the +codomain into the domain, mapping the `(χ, D...) ← χ` partition used for left-gauged tensors +onto the `χ ← (χ, D...)` partition used for right-gauged tensors. +""" +repartition_right(A::EdgeTensor; copy = true) = repartition(A, 1; copy) + +""" + repartition_left(A::RightProjector; copy = true) + +Move the physical legs of a (1, N) tensor map representing a right isometry from the +codomain into the domain, mapping the `χ ← (χ, D...)` partition used for right-gauged +tensors onto the `(χ, D...) ← χ` partition used for left-gauged tensors. +""" +repartition_left(A::RightProjector; copy = true) = repartition(A, numind(A) - 1; copy) + +""" + _repartition(t::AbstractTensorMap, N₁::Int, N₂::Int=numind(t) - N₁; copy=false) + +Differentiable stand-in for [`repartition`](@extref `TensorKit.repartition-Tuple{AbstractTensorMap, Int64, Int64}`). Identical to it, except that +it avoids a bug with the `backend` and `allocator` keyword arguments in the TensorKit +`rrule` implementation. + +NOTE: to be removed once https://github.com/QuantumKitHub/TensorKit.jl/pull/513 is merged +and released. +""" +function _repartition( + t::AbstractTensorMap, N₁::Int, N₂::Int = numind(t) - N₁; copy::Bool = false + ) + N₁ + N₂ == numind(t) || + throw(ArgumentError("Invalid repartition: $(numind(t)) to ($N₁, $N₂)")) + p₁, p₂ = let all_inds = (codomainind(t)..., reverse(domainind(t))...) + ntuple(i -> all_inds[i], N₁), reverse(ntuple(i -> all_inds[i + N₁], N₂)) + end + return transpose(t, (p₁, p₂); copy) +end +function ChainRulesCore.rrule( + config::RuleConfig, ::typeof(repartition_right), A::EdgeTensor + ) + A´, repartition_pullback = rrule_via_ad(config, _repartition, A, 1) + repartition_right_pullback(ΔA´) = NoTangent(), repartition_pullback(ΔA´)[2] + return A´, repartition_right_pullback +end +function ChainRulesCore.rrule( + config::RuleConfig, ::typeof(repartition_left), A::RightProjector + ) + A´, repartition_pullback = rrule_via_ad(config, _repartition, A, numind(A) - 1) + repartition_left_pullback(ΔA´) = NoTangent(), repartition_pullback(ΔA´)[2] + return A´, repartition_left_pullback +end + +## Effective operators + +""" + ∂C(C::CornerTensor, GL::EdgeTensor, GR::EdgeTensor) + +Apply the effective bond operator defined by the left and right environments `GL` and `GR` +to a bond tensor `C`. +This is exactly equivalent to the action of a [`MPSKit.C_hamiltonian`](@extref), but avoids +issues with AD through an [`MPSKit.MPODerivativeOperator`](@extref) and planar contractions. +""" +function ∂C(C::CornerTensor{S}, GL::EdgeTensor{S}, GR::EdgeTensor{S}) where {S} + GR = twistdual(GR, numind(GR)) + return _∂C(C, GL, GR) +end +@generated function _∂C( + C::CornerTensor{S}, GL::EdgeTensor{S, N}, GR::EdgeTensor{S, N} + ) where {S, N} + C´_e = tensorexpr(:C´, -1, -2) + C_e = tensorexpr(:C, 1, 2) + GL_e = tensorexpr(:GL, (-1, (3:(N + 1))...), 1) + GR_e = tensorexpr(:GR, (2:(N + 1)...,), -2) + return macroexpand(@__MODULE__, :(return @tensor $C´_e := $GL_e * $C_e * $GR_e)) +end + +""" + ∂AC(AC::EdgeTensor, GL::EdgeTensor, O, GR::EdgeTensor) + +Apply the effective site operator defined by the left and right environments `GL` and `GR` +and the local network tensor `O` to a center-gauged MPS tensor `AC`. +This is exactly equivalent to the action of an [`MPSKit.AC_hamiltonian`](@extref), but +avoids issues with AD through an [`MPSKit.MPODerivativeOperator`](@extref) and planar +contractions. +""" +function ∂AC(AC::E, GL::E, O, GR::E) where {E <: EdgeTensor} + GR = twistdual(GR, numind(GR)) + return _∂AC(AC, GL, O, GR) +end +function _∂AC( + AC::EdgeTensor{S, 3}, GL::EdgeTensor{S, 3}, O::PEPSSandwich, GR::EdgeTensor{S, 3} + ) where {S} + return @autoopt @tensor AC′[χ_SW D_S_above D_S_below; χ_SE] := + GL[χ_SW D_W_above D_W_below; χ_NW] * + AC[χ_NW D_N_above D_N_below; χ_NE] * + GR[χ_NE D_E_above D_E_below; χ_SE] * + ket(O)[d; D_N_above D_E_above D_S_above D_W_above] * + conj(bra(O)[d; D_N_below D_E_below D_S_below D_W_below]) +end +function _∂AC( + AC::EdgeTensor{S, 2}, GL::EdgeTensor{S, 2}, O::PartitionFunctionTensor, + GR::EdgeTensor{S, 2}, + ) where {S} + return @autoopt @tensor AC′[χ_SW D_S; χ_SE] := + GL[χ_SW D_W; χ_NW] * + AC[χ_NW D_N; χ_NE] * + GR[χ_NE D_E; χ_SE] * + O[D_W D_S; D_N D_E] +end + +## Local contractions + +""" + get_AC(env::SymmetricBoundaryMPSEnv) + +Construct the center-gauged MPS tensor of a symmetric boundary MPS environment by +symmetrically combining its left- and right-gauged forms. +""" +function get_AC(env::SymmetricBoundaryMPSEnv) + return (absorb_right(env.AL, env.C) + absorb_left(env.AR, env.C)) / 2 +end + +# Accessors emitted by the generated contractions below. The north and south rows carry the +# gauge center in their westmost slot and right-gauged tensors elsewhere, while a symmetric +# boundary MPS has a single left and right environment for every row. +_north_edge(env::SymmetricBoundaryMPSEnv, i::Int) = isone(i) ? get_AC(env) : env.AR +_south_edge(env::SymmetricBoundaryMPSEnv, i::Int) = isone(i) ? get_AC(env) : env.AR +_east_edge(env::SymmetricBoundaryMPSEnv, ::Int) = env.GR +_west_edge(env::SymmetricBoundaryMPSEnv, ::Int) = env.GL + +""" + boundary_contraction_expr(::Type{<:SymmetricBoundaryMPSEnv}, rowrange, colrange) + +Build the contraction expressions for the boundary MPS tensors surrounding the patch spanned +by `rowrange` and `colrange`. + +Two things distinguish it from the CTMRG version. There are no corner tensors, since the +gauge center of the boundary MPS already absorbs them, so edges of neighboring sides share a +single environment label. And the boundary MPS acts as its own bra, so the south row is the +conjugate of the north row traversed in the opposite direction, which flips its codomain and +domain relative to a CTMRG south edge. +""" +function boundary_contraction_expr( + ::Type{<:SymmetricBoundaryMPSEnv{TC, TE}}, rowrange, colrange + ) where {TC, TE} + # the MPS tensors carry one virtual leg per layer of the sandwich, on top of the two + # environment indices threading the ring, so the height follows from the tensor type + height = numout(TE) - 1 + rmin, rmax = extrema(rowrange) + cmin, cmax = extrema(colrange) + gridsize = (rmax - rmin + 1, cmax - cmin + 1) + + # corners are absorbed into the gauge center, so neighboring sides share a label + north_labels = [ + envlabel(:NW), (envlabel(NORTH, i) for i in 1:(gridsize[2] - 1))..., envlabel(:NE), + ] + east_labels = [ + envlabel(:NE), (envlabel(EAST, i) for i in 1:(gridsize[1] - 1))..., envlabel(:SE), + ] + south_labels = [ + envlabel(:SW), (envlabel(SOUTH, i) for i in 1:(gridsize[2] - 1))..., envlabel(:SE), + ] + west_labels = [ + envlabel(:NW), (envlabel(WEST, i) for i in 1:(gridsize[1] - 1))..., envlabel(:SW), + ] + + edges_N = map(1:gridsize[2]) do i + return tensorexpr( + :(_north_edge(env, $i)), + (north_labels[i], virtuallabel.(NORTH, ntuple(identity, height), i)...), + north_labels[i + 1], + ) + end + + edges_E = map(1:gridsize[1]) do i + return tensorexpr( + :(_east_edge(env, $i)), + (east_labels[i], virtuallabel.(EAST, ntuple(identity, height), i)...), + east_labels[i + 1], + ) + end + + edges_S = map(1:gridsize[2]) do i + edge = tensorexpr( + :(_south_edge(env, $i)), + (south_labels[i], virtuallabel.(SOUTH, ntuple(identity, height), i)...), + south_labels[i + 1], + ) + return Expr(:call, :conj, edge) + end + + edges_W = map(1:gridsize[1]) do i + return tensorexpr( + :(_west_edge(env, $i)), + (west_labels[i + 1], virtuallabel.(WEST, ntuple(identity, height), i)...), + west_labels[i], + ) + end + + return [edges_N..., edges_E..., edges_S..., edges_W...] +end + +## Density matrices and local contractions + +# `reduced_densitymatrix` itself is generic in the environment; only the patch-size special +# cases it dispatches to need a symmetric boundary MPS implementation. + +# Special cases mirroring the corresponding CTMRG density matrices, which keep the same +# contraction order but avoid unnecessary intermediate permutations. Since the gauge center +# of the boundary MPS already absorbs the corners, no corner absorption step is needed here; +# the only structural difference from the CTMRG versions is that the south edges are the +# conjugates of the north ones, and hence carry swapped codomain and domain. + +function reduced_densitymatrix1x1( + inds::CartesianIndex{2}, ket::InfinitePEPS, bra::InfinitePEPS, + env::SymmetricBoundaryMPSEnv, + ) + row, col = Tuple(inds) + + A = ket[row, col] + Ā = bra[row, col] + + E_north = _north_edge(env, 1) + E_east = _east_edge(env, 1) + E_south = _south_edge(env, 1) + E_west = _west_edge(env, 1) + + @tensor EE_SW[χSE χNW DSb DWb; DSt DWt] := + conj(E_south[χSW DSt DSb; χSE]) * E_west[χSW DWt DWb; χNW] + + @tensor EE_SWA[χSE χNW DNt DEt; dt DSb DWb] := + EE_SW[χSE χNW DSb DWb; DSt DWt] * A[dt; DNt DEt DSt DWt] + + @tensor EE_NE[DNb DEb; χSE χNW DNt DEt] := + E_north[χNW DNt DNb; χNE] * E_east[χNE DEt DEb; χSE] + + @tensor EEAEE[dt; DNb DEb DSb DWb] := + EE_NE[DNb DEb; χSE χNW DNt DEt] * EE_SWA[χSE χNW DNt DEt; dt DSb DWb] + + @tensor ρ[dt; db] := EEAEE[dt; DNb DEb DSb DWb] * conj(Ā[db; DNb DEb DSb DWb]) + + return ρ / str(ρ) +end + +function reduced_densitymatrix2x1( + ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, + env::SymmetricBoundaryMPSEnv, + ) + row, col = Tuple(ind) + + A_north = ket[row, col] + Ā_north = bra[row, col] + A_south = ket[row + 1, col] + Ā_south = bra[row + 1, col] + + E_north = _north_edge(env, 1) + E_northeast = _east_edge(env, 1) + E_southeast = _east_edge(env, 2) + E_south = _south_edge(env, 1) + E_southwest = _west_edge(env, 2) + E_northwest = _west_edge(env, 1) + + @tensor EE_NW[χW χNE DNWt DNt; DNWb DNb] := + E_northwest[χW DNWt DNWb; χNW] * E_north[χNW DNt DNb; χNE] + @tensor EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] := + EE_NW[χW χNE DNWt DNt; DNWb DNb] * conj(Ā_north[dNb; DNb DNEb DMb DNWb]) + @tensor EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] := + EEA_NW[χW DMb dNb χNE DNEb; DNWt DNt] * A_north[dNt; DNt DNEt DMt DNWt] + @tensor EEEAA_N[dNt dNb; χW DMt DMb χE] := + EEAA_NW[χW DMb dNb dNt DMt; χNE DNEt DNEb] * E_northeast[χNE DNEt DNEb; χE] + + @tensor EE_SE[χE χSW DSEt DSt; DSEb DSb] := + E_southeast[χE DSEt DSEb; χSE] * conj(E_south[χSW DSt DSb; χSE]) + @tensor EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] := + EE_SE[χE χSW DSEt DSt; DSEb DSb] * conj(Ā_south[dSb; DMb DSEb DSb DSWb]) + @tensor EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] := + EEA_SE[χE DMb dSb χSW DSWb; DSEt DSt] * A_south[dSt; DMt DSEt DSt DSWt] + @tensor EEEAA_S[χW DMt DMb χE; dSt dSb] := + EEAA_SE[χE DMb dSb dSt DMt; χSW DSWt DSWb] * E_southwest[χSW DSWt DSWb; χW] + + @tensor ρ[dNt dSt; dNb dSb] := + EEEAA_N[dNt dNb; χW DMt DMb χE] * EEEAA_S[χW DMt DMb χE; dSt dSb] + + return ρ / str(ρ) +end + +function reduced_densitymatrix1x2( + ind::CartesianIndex, ket::InfinitePEPS, bra::InfinitePEPS, + env::SymmetricBoundaryMPSEnv, + ) + row, col = Tuple(ind) + + A_west = ket[row, col] + Ā_west = bra[row, col] + A_east = ket[row, col + 1] + Ā_east = bra[row, col + 1] + + E_northwest = _north_edge(env, 1) + E_northeast = _north_edge(env, 2) + E_east = _east_edge(env, 1) + E_southeast = _south_edge(env, 2) + E_southwest = _south_edge(env, 1) + E_west = _west_edge(env, 1) + + @tensor EE_SW[χS χNW DSWt DWt; DSWb DWb] := + conj(E_southwest[χSW DSWt DSWb; χS]) * E_west[χSW DWt DWb; χNW] + @tensor EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] := + EE_SW[χS χNW DSWt DWt; DSWb DWb] * conj(Ā_west[dWb; DNWb DMb DSWb DWb]) + @tensor EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] := + EEA_SW[χS DMb dWb χNW DNWb; DSWt DWt] * A_west[dWt; DNWt DMt DSWt DWt] + @tensor EEEAA_W[dWt dWb; χS DMt DMb χN] := + EEAA_SW[χS DMb dWb dWt DMt; χNW DNWt DNWb] * E_northwest[χNW DNWt DNWb; χN] + + @tensor EE_NE[χN χSE DNEt DEt; DNEb DEb] := + E_northeast[χN DNEt DNEb; χNE] * E_east[χNE DEt DEb; χSE] + @tensor EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] := + EE_NE[χN χSE DNEt DEt; DNEb DEb] * conj(Ā_east[dEb; DNEb DEb DSEb DMb]) + @tensor EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] := + EEA_NE[χN DMb dEb χSE DSEb; DNEt DEt] * A_east[dEt; DNEt DEt DSEt DMt] + @tensor EEEAA_E[χS DMt DMb χN; dEt dEb] := + EEAA_NE[χN DMb dEb dEt DMt; χSE DSEt DSEb] * conj(E_southeast[χS DSEt DSEb; χSE]) + + @tensor ρ[dWt dEt; dWb dEb] := + EEEAA_W[dWt dWb; χS DMt DMb χN] * EEEAA_E[χS DMt DMb χN; dEt dEb] + + return ρ / str(ρ) +end + +## Partition function tensor insertions + +function _contract_site( + AC::EdgeTensor{S, 2}, GL::EdgeTensor{S, 2}, GR::EdgeTensor{S, 2}, + O::PartitionFunctionTensor, + ) where {S} + @autoopt @tensor o = + AC[χ_NW D_N; χ_NE] * + GR[χ_NE D_E; χ_SE] * + conj(AC[χ_SW D_S; χ_SE]) * + GL[χ_SW D_W; χ_NW] * + O[D_W D_S; D_N D_E] + + return o +end diff --git a/src/algorithms/ctmrg/ctmrg.jl b/src/algorithms/ctmrg/ctmrg.jl index f8cc6e53f..96880eed3 100644 --- a/src/algorithms/ctmrg/ctmrg.jl +++ b/src/algorithms/ctmrg/ctmrg.jl @@ -4,7 +4,7 @@ $(TYPEDEF) Abstract super type for the corner transfer matrix renormalization group (CTMRG) algorithm for contracting infinite PEPS. """ -abstract type CTMRGAlgorithm end +abstract type CTMRGAlgorithm <: BoundaryAlgorithm end const CTMRG_SYMBOLS = IdDict{Symbol, Type{<:CTMRGAlgorithm}}() diff --git a/src/algorithms/optimization/implicit_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl index f00632504..d42e21c10 100644 --- a/src/algorithms/optimization/implicit_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -194,6 +194,21 @@ function _check_algorithm_combination(::C4vCTMRG, symm::Union{Nothing, <:Symmetr end return nothing end +function _check_algorithm_combination( + ::SymmetricBoundaryMPS, symm::Union{Nothing, <:SymmetrizationStyle} + ) + if !(symm isa RotateReflect) + msg = "SymmetricBoundaryMPS optimization is compatible only with combined Hermitian reflection and rotation symmetrization. \ + Make sure to set `symmetrization = RotateReflect()` or implement an equivalent custom symmetrization scheme." + @warn msg + end + return nothing +end +function _check_algorithm_combination(::SymmetricBoundaryMPS, ::FixedPointGradient) + msg = "The `:FixedPointGradient` algorithm is not implemented for `SymmetricBoundaryMPS`; \ + select `:ImplicitGradient` instead" + throw(ArgumentError(msg)) +end #= Evaluating the gradient of the cost function for CTMRG: @@ -684,3 +699,87 @@ function PEPSKit._rrule( return (env, info), leading_boundary_characteristic_pullback end + + +## Symmetric boundary MPS gradient through implicit differentiation + +# extract a cotangent from an environment `Tangent`, filling in explicit zeros where needed +_env_cotangent(Δ, t) = Δ isa AbstractZero ? zerovector(t) : unthunk(Δ) + +""" + _rrule( + gradmode::ImplicitGradient, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + env₀::SymmetricBoundaryMPSEnv, + state, + alg::SymmetricBoundaryMPS, + ) + +Reverse rule for a boundary MPS contraction of a network with a single-site unit cell which +is invariant under rotations and Hermitian reflections. Uses an implicit differentiation +approach based on the algebraic characteristic equations which characterize convergence of +the contraction algorithm. + +See [`generate_boundary_mps_characteristic_equation`](@ref). +""" +function _rrule( + gradmode::ImplicitGradient, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + env₀::SymmetricBoundaryMPSEnv, + state, + alg::SymmetricBoundaryMPS, + ) + # the forward computation can be used as is + env, info = MPSKit.leading_boundary(env₀, state, alg) + AL, AR, C, GL, GR = _unpack(env) + + # null spaces of the left- and right-gauged MPS tensors + VL = left_null(AL) + VR = right_null(repartition_right(AR)) + + # instantiate the differentiable variables parametrizing the MPS tensors + l = zeros(scalartype(AL), space(VL, numind(VL))' ← space(AL, numind(AL))') + r = zeros(scalartype(AR), space(AR, 1) ← space(VR, 1)) + + # initialize the characteristic equations and check that they are actually satisfied + F = generate_boundary_mps_characteristic_equation(env, (VL, VR)) + F_norms = norm.(F(state, l, r, C, GL, GR)) + # the equations for `l` and `r` are the only ones containing an inverse bond tensor, so + # their residuals are amplified by its condition number; scale their tolerances to match + amplification = norm(sdiag_pow(real(DiagonalTensorMap(C)), -1), Inf) + F_tols = (1.0e2 * _tol(alg)) .* (amplification, amplification, 1, 1, 1) + any(F_norms .> F_tols) && @warn( + "Characteristic equations not satisfied, still using the gradient: $F_norms" + ) + + # get the partial pullbacks of the characteristic equations + _, F_vjp = rrule_via_ad(config, F, state, l, r, C, GL, GR) + vjp_env(x) = F_vjp(x)[3:end] # MPS tensor, bond tensor and environment pullback + vjp_state(x) = F_vjp(x)[2] # state pullback + + function leading_boundary_implicit_pullback((_Δenv, _Δinfo)) + Δenv = unthunk(_Δenv) + Δenv isa AbstractZero && return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent() + + # extract the environment cotangents and map them onto the parametrization + ΔAL = _env_cotangent(Δenv.AL, AL) + ΔAR = _env_cotangent(Δenv.AR, AR) + Δc = _env_cotangent(Δenv.C, C) + Δgl = _env_cotangent(Δenv.GL, GL) + Δgr = _env_cotangent(Δenv.GR, GR) + Δl = VL' * ΔAL + Δr = repartition_right(ΔAR) * VR' + + # collect cotangents of the characteristic equations + Δy = (Δl, Δr, Δc, Δgl, Δgr) + + # evaluate the implicit gradient + Δstate = implicit_gradient(Δy, vjp_env, vjp_state, Δy, gradmode.solver_alg) + + return NoTangent(), ZeroTangent(), Δstate, NoTangent() + end + + return (env, info), leading_boundary_implicit_pullback +end diff --git a/src/algorithms/optimization/peps_optimization.jl b/src/algorithms/optimization/peps_optimization.jl index 634292d3e..7d661756f 100644 --- a/src/algorithms/optimization/peps_optimization.jl +++ b/src/algorithms/optimization/peps_optimization.jl @@ -29,7 +29,7 @@ $(TYPEDFIELDS) Construct a PEPS optimization algorithm struct based on keyword arguments. For a full description, see [`fixedpoint`](@ref). The supported keywords are: -* `boundary_alg::Union{NamedTuple,<:CTMRGAlgorithm,...}` +* `boundary_alg::Union{NamedTuple,<:BoundaryAlgorithm}` * `gradient_alg::Union{NamedTuple,Nothing,<:GradientAlgorithm}` * `optimizer_alg::Union{NamedTuple,<:OptimKit.OptimizationAlgorithm}` * `precondition_alg::Union{NamedTuple,Nothing,<:PreconditionAlgorithm}` @@ -62,7 +62,7 @@ function PEPSOptimize(; boundary_alg = (;), gradient_alg = (;), optimizer_alg = (;), precondition_alg = (;), reuse_env = Defaults.reuse_env, symmetrization = nothing, ) - boundary_algorithm = _alg_or_nt(CTMRGAlgorithm, boundary_alg) + boundary_algorithm = _alg_or_nt(BoundaryAlgorithm, boundary_alg) gradient_algorithm = _alg_or_nt(GradientAlgorithm, gradient_alg) optimizer_algorithm = _alg_or_nt(OptimKit.OptimizationAlgorithm, optimizer_alg) precondition_algorithm = _alg_or_nt(PreconditionAlgorithm, precondition_alg) @@ -288,7 +288,7 @@ function fixedpoint( # 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 + alg.gradient_alg, tracked_finalizer.tol_state.iter, _tol(boundary_alg) ) E, gs = withgradient(peps) do ψ env′, info = hook_pullback( diff --git a/src/algorithms/select_algorithm.jl b/src/algorithms/select_algorithm.jl index c121fb3fd..db94d43e0 100644 --- a/src/algorithms/select_algorithm.jl +++ b/src/algorithms/select_algorithm.jl @@ -81,8 +81,8 @@ function select_algorithm( boundary_alg = _dynamic_tol_or_alg(boundary_alg; dynamic_tol_kwargs...) end - # C4vCTMRG-specific defaults - if parent_alg(boundary_alg) isa C4vCTMRG + # defaults specific to fully symmetric contraction algorithms + if parent_alg(boundary_alg) isa Union{C4vCTMRG, SymmetricBoundaryMPS} # symmetrize state and gradient if isnothing(symmetrization) symmetrization = RotateReflect() @@ -132,6 +132,16 @@ function select_algorithm( ) end +function select_algorithm( + ::typeof(leading_boundary), ::SymmetricBoundaryMPSEnv; + tol = Defaults.ctmrg_tol, + maxiter = Defaults.ctmrg_maxiter, + verbosity = Defaults.ctmrg_verbosity, + mps_alg = (;), + ) + return SymmetricBoundaryMPS(; tol, maxiter, verbosity, mps_alg) +end + function select_algorithm( ::typeof(leading_boundary), env₀::CTMRGEnv; diff --git a/src/environments/boundarymps_environments.jl b/src/environments/boundarymps_environments.jl new file mode 100644 index 000000000..dc6948d2e --- /dev/null +++ b/src/environments/boundarymps_environments.jl @@ -0,0 +1,199 @@ +""" +$(TYPEDEF) + +Environment of a boundary MPS contraction of a network with a single-site unit cell which is +invariant under rotations and Hermitian reflections. + +Such an environment consists of a single uniform MPS approximating the leading eigenvector +of the row-to-row transfer matrix of the network, given in its left- and right-gauged form +`AL` and `AR` along with the corresponding bond tensor `C`, as well as the left and right +environments `GL` and `GR`. The latter are obtained as the left and right fixed points of +the MPS-network-MPS transfer matrix. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + SymmetricBoundaryMPSEnv([f=randn, T=scalartype(network)], network, Venv::ElementarySpace) + +Construct a symmetric boundary MPS environment for a given `network` with all virtual +environment spaces given by `Venv`, where the tensor entries are generated by `f` with +scalar type `T`. Here, `network` can be an [`InfiniteSquareNetwork`](@ref) or anything that +can be converted into one, such as an [`InfinitePEPS`](@ref) or an +[`InfinitePartitionFunction`](@ref). +""" +struct SymmetricBoundaryMPSEnv{TC, TE} + "left-gauged MPS tensor" + AL::TE + "right-gauged MPS tensor" + AR::TE + "MPS bond tensor" + C::TC + "left environment" + GL::TE + "right environment" + GR::TE +end + +# unpack an environment into its constituent tensors +_unpack(env::SymmetricBoundaryMPSEnv) = (env.AL, env.AR, env.C, env.GL, env.GR) + +# +## initialization +# + +""" + _initialize_GL([f=randn, T=scalartype(A)], A::GenericMPSTensor, O) + _initialize_GR([f=randn, T=scalartype(A)], A::GenericMPSTensor, O) + +Initialize a left or right environment for an MPS tensor `A` and a local network tensor +`O`, with entries generated by `f` and scalar type `T`. +""" +function _initialize_GL(f, ::Type{T}, A::GenericMPSTensor, O) where {T} + V = left_virtualspace(A) ⊗ _elementwise_dual(west_virtualspace(O)) ← left_virtualspace(A) + return f(T, V) +end +function _initialize_GR(f, ::Type{T}, A::GenericMPSTensor, O) where {T} + V = right_virtualspace(A) ⊗ _elementwise_dual(east_virtualspace(O)) ← right_virtualspace(A) + return f(T, V) +end + +function SymmetricBoundaryMPSEnv(network::InfiniteSquareNetwork, Venv::ElementarySpace) + return SymmetricBoundaryMPSEnv(randn, scalartype(network), network, Venv) +end +function SymmetricBoundaryMPSEnv( + state::Union{InfinitePEPS, InfinitePartitionFunction}, args... + ) + return SymmetricBoundaryMPSEnv(InfiniteSquareNetwork(state), args...) +end +function SymmetricBoundaryMPSEnv( + f, T, state::Union{InfinitePEPS, InfinitePartitionFunction}, args... + ) + return SymmetricBoundaryMPSEnv(f, T, InfiniteSquareNetwork(state), args...) +end +function SymmetricBoundaryMPSEnv( + f, ::Type{T}, network::InfiniteSquareNetwork, Venv::ElementarySpace + ) where {T} + size(network) == (1, 1) || + throw(ArgumentError("symmetric boundary MPS environments require a single-site unit cell")) + O = network[1, 1] + + # initialize the MPS tensors by constructing an actual MPS, this improves stability + A = f(T, Venv ⊗ _elementwise_dual(north_virtualspace(O)) ← Venv) + mps = InfiniteMPS([A]) + AL, AR, C = only(mps.AL), only(mps.AR), only(mps.C) + + GL = _initialize_GL(f, T, AL, O) + GR = _initialize_GR(f, T, AR, O) + + return SymmetricBoundaryMPSEnv(AL, AR, C, GL, GR) +end + +# +## interface +# + +function Base.:(==)(env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv) + return all(splat(==), zip(_unpack(env₁), _unpack(env₂))) +end +function Base.isapprox( + env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv; kwargs... + ) + return all(zip(_unpack(env₁), _unpack(env₂))) do (t₁, t₂) + return isapprox(t₁, t₂; kwargs...) + end +end + +Base.real(env::SymmetricBoundaryMPSEnv) = SymmetricBoundaryMPSEnv(map(real, _unpack(env))...) +Base.complex(env::SymmetricBoundaryMPSEnv) = SymmetricBoundaryMPSEnv(map(complex, _unpack(env))...) +Base.similar(env::SymmetricBoundaryMPSEnv) = SymmetricBoundaryMPSEnv(map(similar, _unpack(env))...) +Base.copy(env::SymmetricBoundaryMPSEnv) = SymmetricBoundaryMPSEnv(map(copy, _unpack(env))...) + +function Base.:+(env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv) + return SymmetricBoundaryMPSEnv(map(+, _unpack(env₁), _unpack(env₂))...) +end +function Base.:-(env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv) + return SymmetricBoundaryMPSEnv(map(-, _unpack(env₁), _unpack(env₂))...) +end +function Base.:*(α::Number, env::SymmetricBoundaryMPSEnv) + return SymmetricBoundaryMPSEnv(map(Base.Fix1(*, α), _unpack(env))...) +end + +""" + update!(env::SymmetricBoundaryMPSEnv, env´::SymmetricBoundaryMPSEnv) + +Update the environment `env` in-place with the contents of `env´`. +""" +function update!( + env::SymmetricBoundaryMPSEnv{TC, TE}, env´::SymmetricBoundaryMPSEnv{TC, TE} + ) where {TC, TE} + foreach(splat(copy!), zip(_unpack(env), _unpack(env´))) + return env +end + +# +## VectorInterface +# + +function VI.scalartype(::Type{SymmetricBoundaryMPSEnv{TC, TE}}) where {TC, TE} + return promote_type(scalartype(TC), scalartype(TE)) +end + +function VI.zerovector(env::SymmetricBoundaryMPSEnv, ::Type{S}) where {S <: Number} + return SymmetricBoundaryMPSEnv(map(t -> zerovector(t, S), _unpack(env))...) +end +function VI.zerovector!(env::SymmetricBoundaryMPSEnv) + foreach(zerovector!, _unpack(env)) + return env +end +VI.zerovector!!(env::SymmetricBoundaryMPSEnv) = zerovector!(env) + +function VI.scale(env::SymmetricBoundaryMPSEnv, α::Number) + return SymmetricBoundaryMPSEnv(map(t -> scale(t, α), _unpack(env))...) +end +function VI.scale!(env::SymmetricBoundaryMPSEnv, α::Number) + foreach(t -> scale!(t, α), _unpack(env)) + return env +end +function VI.scale!( + env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv, α::Number + ) + foreach(zip(_unpack(env₁), _unpack(env₂))) do (t₁, t₂) + return scale!(t₁, t₂, α) + end + return env₁ +end +VI.scale!!(env::SymmetricBoundaryMPSEnv, α::Number) = scale!(env, α) +function VI.scale!!( + env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv, α::Number + ) + return scale!(env₁, env₂, α) +end + +function VI.add( + env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv, α::Number, β::Number + ) + return SymmetricBoundaryMPSEnv( + map((t₁, t₂) -> add(t₁, t₂, α, β), _unpack(env₁), _unpack(env₂))... + ) +end +function VI.add!( + env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv, α::Number, β::Number + ) + foreach(zip(_unpack(env₁), _unpack(env₂))) do (t₁, t₂) + return add!(t₁, t₂, α, β) + end + return env₁ +end +function VI.add!!( + env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv, α::Number, β::Number + ) + return add!(env₁, env₂, α, β) +end + +function VI.inner(env₁::SymmetricBoundaryMPSEnv, env₂::SymmetricBoundaryMPSEnv) + return sum(splat(inner), zip(_unpack(env₁), _unpack(env₂))) +end +VI.norm(env::SymmetricBoundaryMPSEnv) = norm(_unpack(env)) diff --git a/test/boundarymps/symmetric_gradients.jl b/test/boundarymps/symmetric_gradients.jl new file mode 100644 index 000000000..6a8a77897 --- /dev/null +++ b/test/boundarymps/symmetric_gradients.jl @@ -0,0 +1,10 @@ +using PEPSKit + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +if !is_buildkite + TestSuite.boundary_mps_symmetric_gradients(Vector) +end diff --git a/test/boundarymps/symmetric_observables.jl b/test/boundarymps/symmetric_observables.jl new file mode 100644 index 000000000..925a74869 --- /dev/null +++ b/test/boundarymps/symmetric_observables.jl @@ -0,0 +1,10 @@ +using PEPSKit + +@isdefined(TestSuite) || include("../testsuite/TestSuite.jl") +using .TestSuite + +is_buildkite = get(ENV, "BUILDKITE", "false") == "true" + +if !is_buildkite + TestSuite.boundary_mps_symmetric_observables(Vector) +end diff --git a/test/testsuite/TestSuite.jl b/test/testsuite/TestSuite.jl index 018a68a82..744a5a784 100644 --- a/test/testsuite/TestSuite.jl +++ b/test/testsuite/TestSuite.jl @@ -46,6 +46,18 @@ module BoundaryMPSVUMPS end using .BoundaryMPSVUMPS +module BoundaryMPSSymmetricGradients + include("boundarymps/symmetric_gradients.jl") + export boundary_mps_symmetric_gradients +end +using .BoundaryMPSSymmetricGradients + +module BoundaryMPSSymmetricObservables + include("boundarymps/symmetric_observables.jl") + export boundary_mps_symmetric_observables +end +using .BoundaryMPSSymmetricObservables + # BP # -------------- module BPExpVals diff --git a/test/testsuite/boundarymps/symmetric_gradients.jl b/test/testsuite/boundarymps/symmetric_gradients.jl new file mode 100644 index 000000000..0f4d5d315 --- /dev/null +++ b/test/testsuite/boundarymps/symmetric_gradients.jl @@ -0,0 +1,115 @@ +using Test +using Random +using LinearAlgebra +using PEPSKit +using Adapt +using MPSKit +using TensorKit +using Zygote +using OptimKit +using KrylovKit + +sd = 42039482049 + +## Test symmetric boundary MPS gradients through implicit differentiation +# ----------------------------------------------------------------------- +χbond = 2 +χenv = 8 +symmetry = RotateReflect() +Pspaces = [ComplexSpace(2), ComplexSpace(2)] +Vspaces = [ComplexSpace(χbond), ComplexSpace(χbond)] +Espaces = [ComplexSpace(χenv), ComplexSpace(χenv)] +models = [heisenberg_XYZ(InfiniteSquare()), transverse_field_ising(InfiniteSquare())] +names = ["Heisenberg", "Ising"] + +boundary_tol = 1.0e-10 +boundary_maxiter = 400 +boundary_verbosity = 1 +gradtol = 1.0e-8 +comp_tol = 1.0e-5 +steps = -0.01:0.005:0.01 + +mps_algs = [:VUMPS, :VOMPS] + +## Tests +# ------ +function boundary_mps_symmetric_gradients(AT) + return @testset "Symmetric boundary MPS energy gradients for $(names[i]) model ($AT)" verbose = true for i in + eachindex( + models + ) + H = models[i] + Pspace = Pspaces[i] + Vspace = Vspaces[i] + Espace = Espaces[i] + + Random.seed!(sd) + psi = adapt(AT, symmetrize!(InfinitePEPS(Pspace, Vspace), symmetry)) + dir = adapt(AT, symmetrize!(InfinitePEPS(Pspace, Vspace), symmetry)) + + gradient_alg = ImplicitGradient(; solver_alg = (; alg = :GMRES, tol = gradtol)) + + # reference gradient from a CTMRG contraction differentiated through the fixed point + ctmrg_alg = SimultaneousCTMRG(; + tol = boundary_tol, maxiter = boundary_maxiter, verbosity = boundary_verbosity + ) + ctmrg_gradient_alg = FixedPointGradient(; solver_alg = (; alg = :GMRES, tol = gradtol)) + ctmrg_env₀ = CTMRGEnv(psi, Espace) + ctmrg_env, = leading_boundary(ctmrg_env₀, psi, ctmrg_alg) + N_ref = norm(psi, ctmrg_env) + E_ref, g_ref = Zygote.withgradient(psi) do ψ + env, = PEPSKit.hook_pullback( + leading_boundary, ctmrg_env₀, ψ, ctmrg_alg; alg_rrule = ctmrg_gradient_alg + ) + return cost_function(ψ, env, H) + end + g_ref = only(g_ref) + symmetrize!(g_ref, symmetry) + + @testset "mps_alg=:$mps_alg" for mps_alg in mps_algs + boundary_alg = SymmetricBoundaryMPS(; + mps_alg, tol = boundary_tol, + maxiter = boundary_maxiter, verbosity = boundary_verbosity, + ) + env₀ = SymmetricBoundaryMPSEnv(psi, Espace) + env, info = leading_boundary(env₀, psi, boundary_alg) + + # the forward contraction should agree with CTMRG + @test info.converged + @test network_value(psi, env) ≈ N_ref rtol = 1.0e-6 + + @info "optimtest of mps_alg=:$mps_alg on $(names[i])" + alphas, fs, dfs1, dfs2 = OptimKit.optimtest( + (psi, env), + dir; + alpha = steps, + retract = PEPSKit.peps_retract, + inner = PEPSKit.real_inner, + ) do (peps, e) + E, g = Zygote.withgradient(peps) do ψ + env´, = PEPSKit.hook_pullback( + leading_boundary, e, ψ, boundary_alg; alg_rrule = gradient_alg + ) + return cost_function(ψ, env´, H) + end + g = only(g) + symmetrize!(g, symmetry) + return E, g + end + @test dfs1 ≈ dfs2 atol = 1.0e-2 + + @info "direct gradient comparison of mps_alg=:$mps_alg on $(names[i])" + E_trial, g_trial = Zygote.withgradient(psi) do ψ + env´, = PEPSKit.hook_pullback( + leading_boundary, env₀, ψ, boundary_alg; alg_rrule = gradient_alg + ) + return cost_function(ψ, env´, H) + end + g_trial = only(g_trial) + symmetrize!(g_trial, symmetry) + + @test E_trial ≈ E_ref rtol = comp_tol + @test norm(g_ref - g_trial) / norm(g_ref) < comp_tol + end + end +end diff --git a/test/testsuite/boundarymps/symmetric_observables.jl b/test/testsuite/boundarymps/symmetric_observables.jl new file mode 100644 index 000000000..586c513e5 --- /dev/null +++ b/test/testsuite/boundarymps/symmetric_observables.jl @@ -0,0 +1,67 @@ +using Test +using Random +using LinearAlgebra +using PEPSKit +using Adapt +using MPSKit +using TensorKit + +sd = 42039482049 + +## Test observables evaluated with a symmetric boundary MPS environment against CTMRG +# ----------------------------------------------------------------------------------- +# The boundary MPS environment absorbs the corners into its gauge center and acts as its own +# bra, so it reaches the generic `expectation_value` machinery through a different set of +# contractions than CTMRG does. Contracting the same symmetric state both ways must give the +# same energy. + +χbond = 2 +χenv = 16 +symmetry = RotateReflect() +Pspace = ComplexSpace(2) +Vspace = ComplexSpace(χbond) +Espace = ComplexSpace(χenv) + +H = heisenberg_XYZ(InfiniteSquare()) + +boundary_tol = 1.0e-10 +boundary_maxiter = 400 +boundary_verbosity = 0 +comp_tol = 1.0e-8 + +mps_algs = [:VUMPS, :VOMPS] + +## Tests +# ------ +function boundary_mps_symmetric_observables(AT) + Random.seed!(sd) + # a symmetric boundary MPS contraction is only meaningful for a fully symmetric network + psi = adapt(AT, symmetrize!(InfinitePEPS(Pspace, Vspace), symmetry)) + + # CTMRG reference + ctmrg_env, = leading_boundary( + CTMRGEnv(psi, Espace), psi; + alg = :SimultaneousCTMRG, tol = boundary_tol, + maxiter = boundary_maxiter, verbosity = boundary_verbosity, + ) + norm_ref = norm(psi, ctmrg_env) + E_ref = expectation_value(psi, H, ctmrg_env) + + return @testset "Symmetric boundary MPS Heisenberg energy, mps_alg=:$mps_alg ($AT)" for mps_alg in mps_algs + boundary_alg = SymmetricBoundaryMPS(; + mps_alg, tol = boundary_tol, + maxiter = boundary_maxiter, verbosity = boundary_verbosity, + ) + env, info = leading_boundary(SymmetricBoundaryMPSEnv(psi, Espace), psi, boundary_alg) + @test info.converged + + # the network value normalizing the expectation value must agree first + @test norm(psi, env) ≈ norm_ref rtol = comp_tol + + E = expectation_value(psi, H, env) + @test E ≈ E_ref rtol = comp_tol + + # `cost_function` is the real part of the same quantity, and is what the optimizer sees + @test cost_function(psi, env, H) ≈ real(E_ref) rtol = comp_tol + end +end