diff --git a/src/Defaults.jl b/src/Defaults.jl index 9bcbb6244..d52971faf 100644 --- a/src/Defaults.jl +++ b/src/Defaults.jl @@ -85,6 +85,9 @@ Module containing default algorithm parameter values and arguments. - `:GeomSum` : Geometric sum approximation of the Neumann series of the inverse Jacobian, see [`PEPSKit.GeomSum`](@ref) for details - `:ManualIter` : Manual fixed-point iteration, see [`PEPSKit.ManualIter`](@ref) for details * `gradient_fixedpoint_solver_eager=$(Defaults.gradient_fixedpoint_solver_eager)` : Enables `:Arnoldi` solver algorithm to finish before the full Krylov dimension is reached. +* `gradient_implicit_solver_alg=:$(Defaults.gradient_implicit_solver_alg)` : Default solver algorithm for the `ImplicitGradient` gradient algorithm. + - `:GMRES` : GMRES iterative linear solver, see [`KrylovKit.GMRES`](@extref) for details + - `:BiCGStab` : BiCGStab iterative linear solver, see [`KrylovKit.BiCGStab`](@extref) for details ## Optimization @@ -151,6 +154,7 @@ const gradient_verbosity = -1 const gradient_alg = :FixedPointGradient const gradient_fixedpoint_solver_alg = :Arnoldi # ∈ {:GMRES, :BiCGStab, :Arnoldi, :GeomSum, :ManualIter} const gradient_fixedpoint_solver_eager = true +const gradient_implicit_solver_alg = :GMRES # ∈ {:GMRES, :BiCGStab} # Optimization const reuse_env = true diff --git a/src/PEPSKit.jl b/src/PEPSKit.jl index 0139d60d0..c2e53dfa7 100644 --- a/src/PEPSKit.jl +++ b/src/PEPSKit.jl @@ -93,6 +93,7 @@ include("algorithms/contractions/ctmrg/renormalize_corner.jl") include("algorithms/contractions/ctmrg/renormalize_edge.jl") include("algorithms/contractions/ctmrg/contract_site.jl") include("algorithms/contractions/ctmrg/gaugefix.jl") +include("algorithms/contractions/ctmrg/characteristic_equations.jl") include("algorithms/contractions/absorb_weight.jl") include("algorithms/contractions/transfer.jl") @@ -140,7 +141,7 @@ include("algorithms/toolbox.jl") include("algorithms/correlator_adapters.jl") include("algorithms/correlators.jl") -include("algorithms/optimization/fixed_point_differentiation.jl") +include("algorithms/optimization/implicit_differentiation.jl") include("algorithms/optimization/peps_optimization.jl") include("algorithms/select_algorithm.jl") @@ -161,7 +162,7 @@ export product_peps export reduced_densitymatrix, expectation_value, network_value, cost_function export correlator, correlation_length export leading_boundary -export PEPSOptimize, FixedPointGradient, GeomSum, ManualIter +export PEPSOptimize, FixedPointGradient, GeomSum, ManualIter, ImplicitGradient export fixedpoint export LocalTruncation diff --git a/src/algorithms/contractions/ctmrg/characteristic_equations.jl b/src/algorithms/contractions/ctmrg/characteristic_equations.jl new file mode 100644 index 000000000..91488fd28 --- /dev/null +++ b/src/algorithms/contractions/ctmrg/characteristic_equations.jl @@ -0,0 +1,684 @@ +# +# Characteristic equations used in implicit differentiation of CMTRG contractions +# + +## Convenience aliases + +const CornerTensor{S, N} = AbstractTensorMap{T, S, 1, 1} where {T} +const EdgeTensor{S, N} = AbstractTensorMap{T, S, N, 1} where {T} +const LeftProjector{S, N} = AbstractTensorMap{T, S, N, 1} where {T} +const RightProjector{S, N} = AbstractTensorMap{T, S, 1, N} where {T} + +const CornerTensors{TC} = Array{TC, 3} where {TC <: CornerTensor} +const EdgeTensors{TE} = Array{TE, 3} where {TE <: EdgeTensor} +const LeftProjectors{TP} = Array{TP, 3} where {TP <: LeftProjector} +const RightProjectors{TP} = Array{TP, 3} where {TP <: RightProjector} + + +## C4v symmetric case + +# partial contractions for different networks + +# north edge and its left projector +function contract_EnVd( + En::EdgeTensor{S, 3}, Vd::LeftProjector{S, 3}, O::PEPSSandwich + ) where {S} + @autoopt @tensor EnVd[χ_NNW D_W_a D_W_b D_S_a D_S_b; χ_E] := + Vd[χ_NNE D_E_a D_E_b; χ_E] * + En[χ_NNW D_N_a D_N_b; χ_NNE] * + ket(O)[d; D_N_a D_E_a D_S_a D_W_a] * + conj(bra(O)[d; D_N_b D_E_b D_S_b D_W_b]) + return EnVd +end +function contract_EnVd( + En::EdgeTensor{S, 2}, Vd::LeftProjector{S, 2}, O::PartitionFunctionTensor + ) where {S} + @autoopt @tensor EnVd[χ_NNW D_W D_S; χ_E] := + Vd[χ_NNE D_E; χ_E] * + En[χ_NNW D_N; χ_NNE] * + O[D_W D_S; D_N D_E] + return EnVd +end + +# northwest enlarged corner with its left projector +function contract_EwCEnVd(Ew::EdgeTensor{S, 3}, C::CornerTensor{S}, EnVd::LeftProjector{S, 5}) where {S} + @autoopt @tensor EwCEnVd[χ_WSW D_S_a D_S_b; χ_E] := + Ew[χ_WSW D_W_a D_W_b; χ_WNW] * + C[χ_WNW; χ_NNW] * + EnVd[χ_NNW D_W_a D_W_b D_S_a D_S_b; χ_E] + return EwCEnVd +end +function contract_EwCEnVd(Ew::EdgeTensor{S, 2}, C::CornerTensor{S}, EnVd::LeftProjector{S, 3}) where {S} + @autoopt @tensor EwCEnVd[χ_WSW D_S; χ_E] := + Ew[χ_WSW D_W; χ_WNW] * + C[χ_WNW; χ_NNW] * + EnVd[χ_NNW D_W D_S; χ_E] + return EwCEnVd +end + +# north edge renormalization +function contract_E´(Ud::RightProjector{S, 3}, EnVd::LeftProjector{S, 5}) where {S} + @tensor E´[-1 -2 -3; -4] := Ud[-1; 1 2 3] * EnVd[1 2 3 -2 -3; -4] + return E´ +end +function contract_E´(Ud::RightProjector{S, 2}, EnVd::LeftProjector{S, 3}) where {S} + @tensor E´[-1 -2; -3] := Ud[-1; 1 2] * EnVd[1 2 -2; -3] + return E´ +end + +""" + generate_symmetric_characteristic_equation( + Cfp::CornerTensor, + Efp::EdgeTensor, # unused + Ufp::LeftProjector, + ULfp::LeftProjector, + ) + +Takes the fixed-point values of the corner tensor `Cfp`, edge tensor `Efp`, left isometry +`Ufp` and its left null space `ULfp` corresponding to a converged C4v CTMRG contraction, and +generates a function ``F(s, C, E, u)`` which characterizes the convergence of the C4v CTMRG +algorithm in terms of the characteristic equation ``F(s, C, E, u) = 0``. +Here, ``s`` corresponds to a state variable (e.g. an `InfinitePEPS` that is being optimized), +and ``(C, E, u)`` represents a C4v symmetric contraction environment. +``C`` and ``E`` directly represent the corner and edge tensors, while ``u`` parametrizes +a differentiable projector ``U`` as ``U = U_{fp} + U_{L,fp} * u``. + +``F`` returns a tuple of three tensors, corresponding to an equation for ``C``, ``E`` and +``u`` respectively: +``` + C---E---|~~~| + | | | U |--- - λC * --C-- + E---O---|~~~| + | | + [ U†] + | +``` +``` + |~~~|---E---|~~~| + ---| U†| | | U |--- - λE * --E-- + |~~~|---O---|~~~| | + | +``` +``` + C---E---|~~~| + | | | U |---Cfp^{-1} - λC * --u-- + E---O---|~~~| + | | + [ ULfp†] + | +``` +where ``λ_C`` and ``λ_E`` which are defined as the inner product of the first term in the +first two contractions given here with ``C`` and ``E`` respectively. +""" +function generate_symmetric_characteristic_equation( + Cfp::CornerTensor, + Efp::EdgeTensor, # unused + Ufp::LeftProjector, + ULfp::LeftProjector, + ) + + iC = sdiag_pow(real(DiagonalTensorMap(Cfp)), -1) + ULd = ULfp' + + function symmetric_characteristic_equation(state, C, E, u) + network = InfiniteSquareNetwork(state) + O = network[1, 1] + + # project input + C = _project_hermitian(C) + E = _project_hermitian(E) + + # assuming 'eigenvalue decomposition' of enlarged corner as ECE = U * S * V + # where now sectretly V = U† + U = Ufp + ULfp * u # left isometry + Ep = physical_flip(E) # edge (west or south) + + # then the 'right projector' (which goes on the left side) is U† + Ud = U' + + # evaluate partial contractions to reuse + # north edge and its left projector + EnVd = contract_EnVd(E, U, O) + EwCEnVd = contract_EwCEnVd(Ep, C, EnVd) + + # F1: corner + C´ = Ud * EwCEnVd + C´ = _project_hermitian(C´) # project output + λ_C = dot(C, C´) + F1 = C´ / λ_C - C + + # F2: edge + E´ = contract_E´(Ud, EnVd) + E´ = _project_hermitian(E´) # project output + λ_E = dot(E, E´) + F2 = E´ / λ_E - E + + # F3: u + ULdEwCEnVd = ULd * EwCEnVd + F3 = (ULdEwCEnVd * iC) / λ_C - u + + return F1, F2, F3 + end + + return symmetric_characteristic_equation +end + + +## Generic asymmetric case + +# Index helpers +# ------------- + +# unit cell indices for absorbing S⁻¹ into left and right projectors +function _proj_sinv_indices(coordinate, nrows, ncols) + dir, r, c = coordinate + r′, c′ = if dir == NORTH + _prev(r, nrows), c + elseif dir == EAST + r, _next(c, ncols) + elseif dir == SOUTH + _next(r, nrows), c + elseif dir == WEST + r, _prev(c, ncols) + end + return dir, r′, c′ +end +# unit cell indices for absorbing fourthroot of inverse S^2 into U isometry +function _leftvec_invfroot_indices(coordinate, nrows, ncols) + dir, r, c = coordinate + r′, c′ = if dir == NORTH + _next(r, nrows), _prev(c, ncols) + elseif dir == EAST + _prev(r, nrows), _prev(c, ncols) + elseif dir == SOUTH + _prev(r, nrows), _next(c, ncols) + elseif dir == WEST + _next(r, nrows), _next(c, ncols) + end + return _prev(dir, 4), r′, c′ +end +# unit cell indices for absorbing forthroot of inverse S^2 into V isometry +function _rightvec_invfroot_indices(coordinate, nrows, ncols) + dir, r, c = coordinate + r′, c′ = if dir == NORTH + r, _next(_next(c, ncols), ncols) + elseif dir == EAST + _next(_next(r, nrows), nrows), c + elseif dir == SOUTH + r, _prev(_prev(c, ncols), ncols) + elseif dir == WEST + _prev(_prev(r, nrows), nrows), c + end + return _next(dir, 4), r′, c′ +end + +# corner coordinate relative to enlarged corner position +function _above_left(co, nrows, ncols) + dir, r, c = co + if dir == 1 + return dir, _prev(r, nrows), _prev(c, ncols) + elseif dir == 2 + return dir, _prev(r, nrows), _next(c, ncols) + elseif dir == 3 + return dir, _next(r, nrows), _next(c, ncols) + elseif dir == 4 + return dir, _next(r, nrows), _prev(c, ncols) + end +end + +# edge coordinate relative to enlarged corner position +function _left(co, nrows, ncols) + dir, r, c = co + if dir == 1 + return _prev(dir, 4), r, _prev(c, ncols) + elseif dir == 2 + return _prev(dir, 4), _prev(r, nrows), c + elseif dir == 3 + return _prev(dir, 4), r, _next(c, ncols) + elseif dir == 4 + return _prev(dir, 4), _next(r, nrows), c + end +end +function _above(co, nrows, ncols) + dir, r, c = co + if dir == 1 + return dir, _prev(r, nrows), c + elseif dir == 2 + return dir, r, _next(c, ncols) + elseif dir == 3 + return dir, _next(r, nrows), c + elseif dir == 4 + return dir, r, _prev(c, ncols) + end +end + +# projector coordinate relative to enlarged corner position +function _left_projector(co, nrows, ncols) + dir, r, c = co + if dir == 1 + return dir, r, _prev(c, ncols) + elseif dir == 2 + return dir, _prev(r, nrows), c + elseif dir == 3 + return dir, r, _next(c, ncols) + elseif dir == 4 + return dir, _next(r, nrows), c + end +end + +# Custom roots +# ------------ + +# square root of a diagonal TensorMap, but with complex non-diagonal adjoint +squareroot(t::AbstractTensorMap) = sdiag_pow(t, 0.5) + +function ChainRulesCore.rrule( + ::typeof(squareroot), t::AbstractTensorMap{T, S, 1, 1} + ) where {T, S} + domain(t) == codomain(t) || + error("Square root of a tensor only exists when domain == codomain.") + P_t = ProjectTo(t) # does this projection project down to a diagonal variation, if t is diagonal? + C = squareroot(t) + + function squareroot_pullback(ΔC_) + ΔC = unthunk(ΔC_) + F = similar(t) + for (c, b) in blocks(C) + copyto!(block(F, c), _squareroot_pullback(b)) + end + return NoTangent(), P_t(_elementwise_mult(ΔC, F)) + end + return C, squareroot_pullback +end +function _squareroot_pullback(C::AbstractMatrix) + Fdata = similar(C) + for i in axes(Fdata, 1), j in axes(Fdata, 2) + # Taking the diagonal only is okay, when dA is diagonal anyway: Fdata[i, i] = 1 / (2 * conj(C[i, i])) + Fdata[i, j] = 1 / conj(C[i, i] + C[j, j]) + end + return Fdata +end + +# take fourth root of diagonal TensorMap, but with complex non-diagonal adjoint +fourthroot(t::AbstractTensorMap) = sdiag_pow(t, 0.25) + +function ChainRulesCore.rrule( + ::typeof(fourthroot), t::AbstractTensorMap{T, S, 1, 1} + ) where {T, S} + domain(t) == codomain(t) || + error("Square root of a tensor only exist when domain == codomain.") + P_t = ProjectTo(t) # does this projection project down to a diagonal variation, if t is diagonal? + C = fourthroot(t) + + function fourthroot_pullback(ΔC_) + ΔC = unthunk(ΔC_) + F = similar(t) + for (c, b) in blocks(C) + copyto!(block(F, c), _fourthroot_pullback(b)) + end + return NoTangent(), P_t(_elementwise_mult(ΔC, F)) + end + return C, fourthroot_pullback +end +function _fourthroot_pullback(C::AbstractMatrix) + Fdata = similar(C) + for i in axes(Fdata, 1), j in axes(Fdata, 2) + # Taking the diagonal only is okay, when dA is diagonal anyway: Fdata[i, i] = 1 / (4 * conj(C[i, i]^3)) + # However, for Q-deformed CTMRG, we need the full version: + Fdata[i, j] = 1 / conj(C[i, i]^3 + C[i, i] * C[j, j]^2 + C[i, i]^2 * C[j, j] + C[j, j]^3) + end + return Fdata +end + +# Util +# ---- + +function eachcoordinate(tensor_unitcell::Array{<:AbstractTensorMap, 3}) + return collect(Iterators.product(axes(tensor_unitcell)...)) +end + +# add rrule for twistdual through dummy out-of-place function +function _twistdual(t::AbstractTensorMap, i::Int) + isdual(space(t, i)) || return t + return twist(t, i) +end +function _twistdual(t::AbstractTensorMap, is) + is′ = filter(i -> isdual(space(t, i)), is) + return twist(t, is′) +end +function ChainRulesCore.rrule( + config::RuleConfig, ::typeof(twistdual), t::AbstractTensorMap, i + ) + tout, twistdual_pullback = rrule_via_ad(config, _twistdual, t, i) + return tout, twistdual_pullback +end + +# add rrule for twistnondual through dummy out-of-place function +function _twistnondual(t::AbstractTensorMap, i::Int) + !isdual(space(t, i)) || return t + return twist(t, i) +end +function _twistnondual(t::AbstractTensorMap, is) + is′ = filter(i -> !isdual(space(t, i)), is) + return twist(t, is′) +end +function ChainRulesCore.rrule( + config::RuleConfig, ::typeof(twistnondual), t::AbstractTensorMap, i + ) + tout, twistnondual_pullback = rrule_via_ad(config, _twistnondual, t, i) + return tout, twistnondual_pullback +end + +function absorb_left( + E::AbstractTensorMap{T, S}, C::CornerTensor{S} + ) where {T, S} + pC = (codomainind(C), domainind(C)) + pE = ((codomainind(E)[1],), (codomainind(E)[2:end]..., domainind(E)...)) + pCE = (codomainind(E), domainind(E)) + return tensorcontract(C, pC, false, E, pE, false, pCE) +end +function absorb_right( + P::AbstractTensorMap{T, S}, C::CornerTensor{S} + ) where {T, S} + pP = ((codomainind(P)..., domainind(P)[2:end]...), (domainind(P)[1],)) + pC = (codomainind(C), domainind(C)) + pPC = (codomainind(P), (domainind(P)[end], domainind(P)[1:(end - 1)]...)) + return tensorcontract(P, pP, false, C, pC, false, pPC) +end +# specialized versions; TODO: probably remove, this is a terrible idea for fermionic tensors... +function absorb_left( + P::RightProjector{S}, C::CornerTensor{S} + ) where {S} + return C * P +end +function absorb_right( + E::EdgeTensor{S}, C::CornerTensor{S} + ) where {S} + return E * C +end +function absorb_left_right(T::AbstractTensorMap, CL::CornerTensor, CR::CornerTensor) + return absorb_right(absorb_left(T, CL), CR) +end + +# Partial contractions +# -------------------- + +# corner eigenvalue equation +function _corners_fixed_point(coordinate, env, EC, PL, PR) + dir, r, c = coordinate + C′ = if dir == NORTHWEST + renormalize_northwest_corner((r, c), EC, PL, PR) + elseif dir == NORTHEAST + renormalize_northeast_corner((r, c), EC, PL, PR) + elseif dir == SOUTHEAST + renormalize_southeast_corner((r, c), EC, PL, PR) + elseif dir == SOUTHWEST + renormalize_southwest_corner((r, c), EC, PL, PR) + end + C′ /= dot(env.corners[coordinate...], C′) + return C′ - env.corners[coordinate...] +end + +# edge eigenvalue equation +function _edges_fixed_point(coordinate, env, O, PL, PR) + dir, r, c = coordinate + E′ = if dir == NORTH + renormalize_north_edge((r, c), env, PL, PR, O) + elseif dir == EAST + renormalize_east_edge((r, c), env, PL, PR, O) + elseif dir == SOUTH + renormalize_south_edge((r, c), env, PL, PR, O) + elseif dir == WEST + renormalize_west_edge((r, c), env, PL, PR, O) + end + E′ /= dot(env.edges[coordinate...], E′) + return E′ - env.edges[coordinate...] +end + +function contract_EPL( + E::EdgeTensor{S, 3}, PL::LeftProjector{S, 3}, O::PEPSSandwich, + ) where {S} + @autoopt @tensor eipl[χ_W D_W_above D_W_below D_S_above D_S_below; χ_S] := + PL[χ_N D_E_above D_E_below; χ_S] * + E[χ_W D_N_above D_N_below; χ_N] * + 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]) + return eipl +end +function contract_EPL( + E::EdgeTensor{S, 2}, PL::LeftProjector{S, 2}, O::PartitionFunctionTensor, + ) where {S} + @autoopt @tensor eipl[χ_W D_W D_S; χ_S] := + PL[χ_N D_E; χ_S] * + E[χ_W D_N; χ_N] * + O[D_W D_S; D_N D_E] + return eipl +end + +function contract_EiCiEPL( + EiCi::EdgeTensor{S, 3}, EPL::LeftProjector{S, 5}, + ) where {S} + @tensor ecepl[-1 -2 -3; -4] := + EPL[1 2 3 -2 -3; -4] * + EiCi[-1 2 3; 1] + return ecepl +end +function contract_EiCiEPL( + EiCi::EdgeTensor{S, 2}, EPL::LeftProjector{S, 3}, + ) where {S} + @tensor ecepl[-1 -2; -4] := + EPL[1 2 -2; -4] * + EiCi[-1 2; 1] + return ecepl +end + +function contract_PREPL( + PR::RightProjector{S, 3}, EPL::LeftProjector{S, 5}, + ) where {S} + @tensor E´[-1 -2 -3; -4] := + PR[-1; 1 2 3] * EPL[1 2 3 -2 -3; -4] + return E´ +end +function contract_PREPL( + PR::RightProjector{S, 2}, EPL::LeftProjector{S, 3}, + ) where {S} + @tensor E´[-1 -2; -4] := + PR[-1; 1 2] * EPL[1 2 -2; -4] + return E´ +end + +# contract two projectors into bond matrix; proxy for multiplication, but with fermion consistency +function _contract_PR_PL(PR::RightProjector{S, N}, PL::LeftProjector{S, N}) where {S, N} + pPR = (codomainind(PR), domainind(PR)) + pPL = (codomainind(PL), domainind(PL)) + pPLR = ((1,), (2,)) + return tensorcontract(PR, pPR, false, PL, pPL, false, pPLR) +end + +function _contract_PR_M(PR::RightProjector{S, N}, M::AbstractTensorMap{T, S, N, N}) where {T, S, N} + pPR = (codomainind(PR), domainind(PR)) + pM = (codomainind(M), domainind(M)) + return tensorcontract(PR, pPR, false, M, pM, false, pPR) +end + +# computation of characteristic equations by constructing full enlarged corners +function contract_halfinfinite_characteristic_equation( + C::CornerTensors, E::EdgeTensors, + is::CornerTensors, s::CornerTensors, + u::CornerTensors, v::CornerTensors, + n::InfiniteSquareNetwork, + Ud::RightProjectors, Vd::LeftProjectors, + iCi::CornerTensors, + ULd::RightProjectors, VRd::LeftProjectors, + iSfp::CornerTensors, + ) + coordinates = eachcoordinate(n, 1:4) + nrows, ncols = size(n) + + # precompute rotated local sandwiches, enlarged corners, and projectors + Or = map(coordinates) do co + dir, r, c = co + return _rotate_north_localsandwich(n[r, c], dir) + end + EC = map(coordinates) do co + return TensorMap(EnlargedCorner(n, CTMRGEnv(iCi, E), co)) + end + PR = map(coordinates) do co + co′ = _proj_sinv_indices(co, nrows, ncols) + return absorb_right(Ud[co...] * EC[co...], is[co′...]) + end + PLpart = map(coordinates) do co + co′ = _next_coordinate(co, nrows, ncols) + return EC[co′...] * Vd[co...] + end + PL = map(coordinates) do co + co′ = _proj_sinv_indices(co, nrows, ncols) + return absorb_left(PLpart[co...], is[co′...]) + end + + # prepare partial contractions with projectors for F1 and F2 + + # absorb corner into right side of edge + EiCi = map(coordinates) do co + return E[_left(co, nrows, ncols)...] * iCi[_above_left(co, nrows, ncols)...] + end + # pre-contract the top edge and its sandwich into the left projector + EPL = map(coordinates) do co + return contract_EPL(E[_above(co, nrows, ncols)...], PL[co...], Or[co...]) + end + # pre-contract the top-left corner and left edge into the previous to complete the + # top-left enlarged corner + EiCiEPL = map(coordinates) do co + return contract_EiCiEPL(EiCi[co...], EPL[co...]) + end + + # corners + F1 = map(coordinates) do co + co´ = _prev_coordinate(co, nrows, ncols) + C´ = _contract_PR_PL(PR[co´...], EiCiEPL[co...]) + λC = dot(C[co...], C´) + return C´ / λC - C[co...] + end + + # edges + F2 = map(coordinates) do co + E´ = contract_PREPL(PR[_left_projector(co, nrows, ncols)...], EPL[co...]) + λ_E = dot(E[co...], E´) + return E´ / λ_E - E[co...] + end + + # halfinfinite environment + F345 = map(coordinates) do co + s´ = _contract_PR_PL(PR[co...], PLpart[co...]) + λs = dot(s[co...], s´) + fp4 = s´ / λs - s[co...] + + co´ = _next_coordinate(co, nrows, ncols) + fp3 = (ULd[co...] * EiCiEPL[co...]) * iSfp[co...] / λs - u[co...] + fp5 = iSfp[co...] * (_contract_PR_M(PR[co...], EC[co´...]) * VRd[co...]) / λs - v[co...] + + return fp3, fp4, fp5 + end + F3 = map(x -> x[1], F345) + F4 = map(x -> x[2], F345) + F5 = map(x -> x[3], F345) + + return F1, F2, F3, F4, F5 +end + + +# Combine into characteristic equations +# ------------------------------------- + +""" + generate_halfinfinite_characteristic_equation( + ::CTMRGAlgorithm{<:HalfInfiniteProjector}, + iSfp::CornerTensors, + Ufp::LeftProjectors, + Vfp::RightProjectors, + ULfp::LeftProjectors, + VRfp::RightProjectors, + ) + +TODO: write this up. +""" +function generate_halfinfinite_characteristic_equation( + iSfp::CornerTensors, + Ufp::LeftProjectors, + Vfp::RightProjectors, + ULfp::LeftProjectors, + VRfp::RightProjectors, + ) + + iSfp = real.(DiagonalTensorMap.(iSfp)) # use as constant preconditioner? + coordinates = eachcoordinate(iSfp) + nrows, ncols = size(iSfp)[2:3] + + # the main routine which uses both the singular values and their inverses + function asymmetric_characteristic_equation(state, C, E, u, s, v) + ## Prepare all the objects we need in the right parametrization + is = map(inv, s) + + # outspace variation parametrization of isometries + U = map(coordinates) do co + return Ufp[co...] + ULfp[co...] * u[co...] + end + V = map(coordinates) do co + return Vfp[co...] + v[co...] * VRfp[co...] + end + + isqsR = map(fourthroot, adjoint.(is) .* is) # root that goes into the left projector + isqsL = map(fourthroot, is .* adjoint.(is)) # root that goes into the right projector + + # pre-dagger the isometries, absorb the square roots + Ud = map(coordinates) do co + co′ = _leftvec_invfroot_indices(co, nrows, ncols) + absorb_right(U[co...]', isqsR[co′...]) + end + Vd = map(coordinates) do co + co′ = _rightvec_invfroot_indices(co, nrows, ncols) + absorb_left(V[co...]', isqsL[co′...]) + end + ULd = map(coordinates) do co + co′ = _leftvec_invfroot_indices(co, nrows, ncols) + absorb_right(ULfp[co...]', isqsR[co′...]) + end + VRd = map(coordinates) do co + co′ = _rightvec_invfroot_indices(co, nrows, ncols) + absorb_left(VRfp[co...]', isqsL[co′...]) + end + + # pre-contract full inverses into corners from both sides + iCi = map(coordinates) do co + co′ = _prev_coordinate(co, nrows, ncols) + return is[co′...] * C[co...] * is[co...] + end + + ## Perform the actual contractions + F1, F2, F3, F4, F5 = contract_halfinfinite_characteristic_equation( + C, E, + is, s, + u, v, + InfiniteSquareNetwork(state), + Ud, Vd, + iCi, + ULd, VRd, + iSfp, + ) + + return F1, F2, F3, F4, F5 + end + + return asymmetric_characteristic_equation +end + +function generate_fullinfinite_characteristic_equation( + iSfp::CornerTensors, + Ufp::LeftProjectors, + Vfp::RightProjectors, + ULfp::LeftProjectors, + VRfp::RightProjectors, + ) + + throw(ArgumentError("Characteristic equations for CTMRGAlgorithm{<:FullInfiniteProjector} are not yet implemented.")) +end diff --git a/src/algorithms/optimization/fixed_point_differentiation.jl b/src/algorithms/optimization/implicit_differentiation.jl similarity index 51% rename from src/algorithms/optimization/fixed_point_differentiation.jl rename to src/algorithms/optimization/implicit_differentiation.jl index 2c2a54991..6cfaf768a 100644 --- a/src/algorithms/optimization/fixed_point_differentiation.jl +++ b/src/algorithms/optimization/implicit_differentiation.jl @@ -180,12 +180,6 @@ function _check_algorithm_combination(boundary_alg, gradient_alg, symmetrization _check_algorithm_combination(boundary_alg, symmetrization) return nothing end -function _check_algorithm_combination(::SequentialCTMRG, ::FixedPointGradient) - msg = "The `:FixedPointGradient` algorithm is not compatible with `SequentialCTMRG` since the sequential \ - application of SVDs does not allow to differentiate through a fixed set of \ - gauges; select SimultaneousCTMRG instead to use :fixed mode" - throw(ArgumentError(msg)) -end function _check_algorithm_combination(::C4vCTMRG, symm::Union{Nothing, <:SymmetrizationStyle}) if !(symm isa RotateReflect) msg = "C4vCTMRG optimization is compatible only with combined Hermitian reflection and rotation symmetrization. \ @@ -215,7 +209,7 @@ function _rrule( gradmode::FixedPointGradient, config::RuleConfig, ::typeof(MPSKit.leading_boundary), - envinit, + envinit::CTMRGEnv, state, alg::CTMRGAlgorithm, ) @@ -360,3 +354,331 @@ function fixedpoint_gradient(∂E∂x, ∂f∂x, ∂f∂A, x₀, alg::KrylovKit. return ∂f∂A(y) end + +# +# Implicit gradient computation +# + + +""" +$(TYPEDEF) + +Gradient algorithm for computing the gradient of an optimization procedure by +differentiating an implicit algebraic characterization of the solution. + +## Fields + +$(TYPEDFIELDS) + +## Constructors + + ImplicitGradient(; kwargs...) + +Construct an implicit gradient algorithm struct based on keyword arguments. +The supported keywords are: + +* `tol::Real=$(Defaults.gradient_tol)` +* `maxiter::Int=$(Defaults.gradient_maxiter)` +* `verbosity::Int=$(Defaults.gradient_verbosity)` +* `solver_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=:$(Defaults.gradient_fixedpoint_solver_alg))`: solver algorithm for the `ImplicitGradient` gradient algorithm. + - `:GMRES` : GMRES iterative linear solver, see [`KrylovKit.GMRES`](@extref) for details + - `:BiCGStab` : BiCGStab iterative linear solver, see [`KrylovKit.BiCGStab`](@extref) for details +""" +struct ImplicitGradient{A} <: GradientAlgorithm{A} + solver_alg::A +end +ImplicitGradient(; kwargs...) = GradientAlgorithm(; alg = :ImplicitGradient, kwargs...) +GRADIENT_ALGORITHM_SYMBOLS[:ImplicitGradient] = ImplicitGradient + +const IMPLICIT_SOLVER_SYMBOLS = IdDict{Symbol, Type{<:Any}}( + :GMRES => GMRES, :BiCGStab => BiCGStab, +) + +_default_solver_alg(::Type{<:ImplicitGradient}) = Defaults.gradient_implicit_solver_alg +_select_solver_alg_symbol(::Type{<:ImplicitGradient}, solver_alg) = + IMPLICIT_SOLVER_SYMBOLS[solver_alg] + +function _check_algorithm_combination(::CTMRGAlgorithm, ::ImplicitGradient) + msg = "The `:ImplicitGradient` algorithm is currently only implemented for the `C4vCTMRG` algorithm." + throw(ArgumentError(msg)) +end +_check_algorithm_combination(::C4vCTMRG, ::ImplicitGradient) = nothing + +function _gauge_fix_c4v_projector(::C4vCTMRG{<:C4vEighProjector}, signs, info) + return info.V * signs[1]' +end +function _gauge_fix_c4v_projector(::C4vCTMRG{<:C4vQRProjector}, signs, info) + return info.Q * signs[1]' +end + +function _check_algorithm_combination(::SequentialCTMRG, ::ImplicitGradient) + msg = "The `:ImplicitGradient` algorithm is not yet compatible with the `SequentialCTMRG` algorithm; select \ + SimultaneousCTMRG instead to use :ImplicitGradient" + throw(ArgumentError(msg)) +end + +function _check_algorithm_combination(::SimultaneousCTMRG{<:FullInfiniteProjector}, ::ImplicitGradient) + msg = "The `:ImplicitGradient` algorithm is not yet compatible with the `FullInfiniteProjector` scheme, \ + since the corresponding charageristic equations have not yet been implemented; select \ + SimultaneousCTMRG with HalfInfiniteProjector instead to use :ImplicitGradient" + throw(ArgumentError(msg)) +end + + +@doc raw""" + implicit_gradient(x̆, ∂ₓF, ∂ₚF, y₀, alg) + +Evaluates the VJP action ``x̆ ∂ₚx`` for an intermediate variable ``x \equiv x(p)`` +which satisfies the characteristic equation ``F(x, p) = 0``, given the VJP actions ``∂ₓF`` +and ``∂ₚF`` of the function ``F`` encoding the characteristic equation. + +More specifically, given a cost function ``E(x(p), p)`` defined in terms of a set of +variational parameters ``p`` and a set of intermediate variables ``x`` that depend on ``p``, +``x \equiv x(p)``, the gradient of the cost function is given by + +```math +dE/dp = ∂ₓE ∂ₚx + ∂ₚE. +``` + +Given the characteristic equation ``F(x, p) = 0``, the VJP action of the Jacobian ``∂ₚx``` on +the adjoint ``x̆ = ∂ₓE`` in the first term of this expression can be evaluated through +implicit differentiation of the characteristic equation as +```math +x̆ ∂ₚx = x̆ (∂ₓF)⁻¹ (-∂ₚF). +``` + +This can be used to differentiate contraction routines, where ``p`` are the variational +parameters of a tensor network, ``x̆ = ∂ₓE`` is the partial +derivative of the cost function with respect to the contraction environment ``x``, and ``F`` +is an algebraic function characterizing the convergence of the algorithm. +""" +implicit_gradient + +function implicit_gradient(∂E∂x, ∂F∂x, ∂F∂A, y₀, alg::KrylovKit.LinearSolver) + y, info = reallinsolve(∂F∂x, ∂E∂x, y₀, alg) + if alg.verbosity > 0 && info.converged != 1 + @warn("Implicit gradient linear solve reached maximal number of iterations:", info) + end + + return (-1) * ∂F∂A(y) +end + +## C4vCTMRG gradient through implicit differentiation + +function _rrule( + gradmode::ImplicitGradient, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit::CTMRGEnv, + state, + alg::C4vCTMRG, + ) + _check_algorithm_combination(alg, gradmode) + + env, = leading_boundary(envinit, state, alg) + + # prepare iterating function corresponding to a single gauge-fixed CTMRG iteration + alg_fixed = _set_fixed_truncation(alg) # fix spaces during differentiation + alg_gauge = _scrambling_env_gauge(alg) # select appropriate gauge-fixing algorithm + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg_fixed) + signs, corner_phases, edge_phases = compute_gauge_fix_gauge(env_conv, env, alg_gauge) + env_fixed = fix_phases(env_conv, signs, corner_phases, edge_phases) + + # NOTE: explicitly keeping corner non-diagonal and (possibly) complex, for use in the backwards pass + C, E = only(env_fixed.corners[NORTHWEST, :, :]), only(env_fixed.edges[NORTH, :, :]) + + # gauge-fix projector accordingly + U = _gauge_fix_c4v_projector(alg, signs, info) + + # get the projector nullspace + UL = left_null(U) + + # instantiate the differentiable variables corresponding to the intermediate projector of the contraction algorithm + u = zeros(scalartype(U), space(UL, numind(UL))' ← space(U, numind(U))') + + # prepare pullback of C4v CTMRG environment constructor (artefact of reusing asymmetric environment type for C4v symmetric contraction) + _, c4v_env_vjp = rrule_via_ad(config, CTMRGEnv, C, E) + + # initialize the partial pullbacks of the characteristic equations + F = generate_symmetric_characteristic_equation(C, E, U, UL) + + # get the partial pullback of the characteristic equations + _, F_vjp = rrule_via_ad(config, F, state, C, E, u) + vjp_env(x) = F_vjp(x)[3:end] # environment and isometry pullback + vjp_state(x) = F_vjp(x)[2] # state pullback + + function leading_boundary_implicit_pullback((_Δenv, _Δinfo)) + Δenv, Δinfo = unthunk(_Δenv), unthunk(_Δinfo) + + # accumulate adjoints from all corners & edges through _c4v_env constructor + _, ΔC, ΔE = c4v_env_vjp(Δenv) + + # apply Hermitian projectors to edge and corner cotangents + ΔE = _project_hermitian(ΔE) + ΔC = _project_hermitian(ΔC) + + # extract cotengent of the isometry variable and verify that it is trivial + ΔU = isa(Δinfo.U, AbstractZero) ? zerovector(U) : unthunk(Δinfo.U) + Δu = UL' * ΔU + norm(Δu) < 1.0e4 * alg.tol || @warn "Nonzero Δu cotangent: norm(Δu)=$(norm(Δu))" + + # collect cotangents of characteristic equation + Δy = (ΔC, ΔE, Δu) + + # evaluate the implicit gradient + Δstate = implicit_gradient(Δy, vjp_env, vjp_state, Δy, gradmode.solver_alg) + + return NoTangent(), ZeroTangent(), Δstate, NoTangent() + end + + # HACK: return gauge-fixed environment with potentially non-diagonal and complex corners + # this ensures that backpropagation through any subsequent observable evaluations will + # give non-diagonal and complex corner adjoints, which is required for root + # differentiation to work here + # NOTE: this very bad practice, since the return type here is manifestly different from + # what you would get from a regular forward run without differentiation + return (env_fixed, info), leading_boundary_implicit_pullback +end + + +## CTMRG gradient through implicit differentiation + +# removing inverse roots of singular values from corners and edges +# i.e. what we actually do to get the modified corners and edges +function remove_inverse_roots( + C::CornerTensors, E::EdgeTensors, S::CornerTensors + ) + sqS = sdiag_pow.(S, 0.5) + coordinates = eachcoordinate(C) + nrows, ncols = size(C)[2:3] + C′ = map(coordinates) do co + co′ = _prev_coordinate(co, nrows, ncols) + return sqS[co′...] * C[co...] * sqS[co...] + end + E′ = map(coordinates) do co + co′ = _edge_sinv_indices(co, nrows, ncols) + return absorb_left_right( + E[co...], sqS[co′...], sqS[co...] + ) + end + return C′ ./ norm.(C′), E′ ./ norm.(E′) +end + +# absorbing inverse roots of singular values into modified corners and edges to obtain +# regular corners and edges +# i.e. the thing we pretend we did right before the cost function evaluation +function absorb_inverse_roots( + C::CornerTensors, E::EdgeTensors, S::CornerTensors + ) + # absorb inverse roots + is = map(inv, S) + isqS = map(squareroot, is) # do this instead of sdiag_pow(S, -0.5) since dS may be non-diagonal + coordinates = eachcoordinate(C) + nrows, ncols = size(C)[2:3] + C′ = map(coordinates) do co + co′ = _prev_coordinate(co, nrows, ncols) + return isqS[co′...] * C[co...] * isqS[co...] + end + E′ = map(coordinates) do co + co′ = _edge_sinv_indices(co, nrows, ncols) + return absorb_left_right(E[co...], isqS[co′...], isqS[co...]) + end + return C′ ./ norm.(C′), E′ ./ norm.(E′) +end + +# unit cell indices for absorbing square roots of S or S⁻¹ into corners and edges +function _edge_sinv_indices(coordinate, nrows, ncols) + dir, r, c = coordinate + r, c = if dir == NORTH + r, _prev(c, ncols) + elseif dir == EAST + _prev(r, nrows), c + elseif dir == SOUTH + r, _next(c, ncols) + elseif dir == WEST + _next(r, nrows), c + end + return (dir, r, c) +end + +# characteristic equation for asymmetric simultaneous CTMRG with generic unit cells and sparse SVD, +# i.e. where computation of the pullback requires solving a Sylvester equation for dU and dV +function PEPSKit._rrule( + gradmode::ImplicitGradient, + config::RuleConfig, + ::typeof(MPSKit.leading_boundary), + envinit::CTMRGEnv, + state, + alg::SimultaneousCTMRG{<:HalfInfiniteProjector}, # SequentialCTMRG doesn't return U, S, V (yet) + ) + env, = leading_boundary(envinit, state, alg) + + # gauge-fix SVD isometries + env_conv, info = ctmrg_iteration(InfiniteSquareNetwork(state), env, alg) + signs, = compute_gauge_fix_gauge(env_conv, env, ScramblingEnvGauge()) + S = normalize.(info.S) + U, V = fix_relative_phases(info.U, info.V, signs) + + # pretend the singular values matrices are just arbitrary complex tensors + s = TensorMap.(S) + if real(scalartype(state)) != scalartype(state) + # complex network, complex things, required to make the derivatives work... + s = complex.(s) + end + + # remove the inverse square roots here to obtain modified corners and edges + C̃, Ẽ = remove_inverse_roots(env.corners, env.edges, S) + # which is fine as long as we pretend we did the absorption as the last step in the + # forward pass, and then explicitly backpropagate through that in the pullback here + _, absorb_inverse_roots_vjp = rrule_via_ad(config, absorb_inverse_roots, C̃, Ẽ, s) + + # get the projector nullspaces + UL = left_null.(U) + VR = right_null.(V) + + # instantiate the variables used in the characteristic equations + u = map(zip(U, UL)) do (Uc, ULc) + return zeros(scalartype(Uc), space(ULc, numind(ULc))' ← space(Uc, numind(Uc))') + end + v = map(zip(V, VR)) do (Vc, VRc) + return zeros(scalartype(Vc), space(Vc, 1) ← space(VRc, 1)) + end + is = sdiag_pow.(s, -1) # also treat them as general complex tensors + + # generate the characteristic equations + F = generate_halfinfinite_characteristic_equation(is, U, V, UL, VR) + + # check if characteristic equations are actually satisfied + FS = F(state, C̃, Ẽ, u, s, v) + F_nrms = norm.(FS) + any(F_nrms .> 1.0e2 * alg.tol) && + @warn "Characteristic equations not satisfied, still using the gradient: $F_nrms" + + # get the partial gradients of the characteristic equations + _, F_vjp = rrule_via_ad(config, F, state, C̃, Ẽ, u, s, v) # full automatic pullback + vjp_env(x) = F_vjp(x)[3:end] # environment and SVD pullback + vjp_state(x) = F_vjp(x)[2] # state pullback + + function leading_boundary_characteristic_pullback((_Δenv, _Δinfo)) + # unpack incoming cotangents + Δenv = unthunk(_Δenv) + + # apply pullback of the dummy inverse root absorption to get adjoints of modified + # corners and edges and singular values + _, ΔC̃, ΔẼ, Δs = absorb_inverse_roots_vjp((Δenv.corners, Δenv.edges)) + + # U and V should have zero cotangents + Δu = zerovector.(u) + Δv = zerovector.(v) + + # collect cotangents of characteristic equation + Δy = (ΔC̃, ΔẼ, Δu, Δs, Δv) + + # 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_characteristic_pullback +end diff --git a/test/gradients/c4v_ctmrg_gradients.jl b/test/gradients/c4v_ctmrg_gradients.jl index 2a0c4dab7..5e139c46f 100644 --- a/test/gradients/c4v_ctmrg_gradients.jl +++ b/test/gradients/c4v_ctmrg_gradients.jl @@ -24,7 +24,7 @@ ctmrg_verbosity = 1 ctmrg_algs = [[:C4vCTMRG]] projector_algs = [[:C4vEighProjector, :C4vQRProjector]] decomposition_rrule_algs = [[:FullPullback, :TruncPullback]] -gradient_algs = [[nothing, :FixedPointGradient]] +gradient_algs = [[nothing, :FixedPointGradient, :ImplicitGradient]] gradient_solver_algs = [[:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi]] steps = -0.01:0.005:0.01 @@ -67,6 +67,12 @@ naive_gradient_done = Set() gradient_solver_alg = nothing # unused in naive gradient, so set to nothing to avoid confusion end + # only run GMRES for the implicit gradient, and skip distinction between decomposition rrule algs + if gradient_alg == :ImplicitGradient + gradient_solver_alg == :GMRES || continue + decomposition_rrule_alg == first(dalgs) || continue + end + # check for allowed combinations of projector alg and decomposition rrule alg decomposition_rrule_alg in allowed_rrule_algs[projector_alg] || continue diff --git a/test/gradients/ctmrg_gradients.jl b/test/gradients/ctmrg_gradients.jl index 09e3f47c4..aad274455 100644 --- a/test/gradients/ctmrg_gradients.jl +++ b/test/gradients/ctmrg_gradients.jl @@ -21,7 +21,7 @@ ctmrg_verbosity = 0 ctmrg_algs = [[:SequentialCTMRG, :SimultaneousCTMRG], [:SequentialCTMRG, :SimultaneousCTMRG]] projector_algs = [[:HalfInfiniteProjector, :FullInfiniteProjector], [:HalfInfiniteProjector, :FullInfiniteProjector]] svd_rrule_algs = [[:FullPullback, :TruncPullback, :Arnoldi], [:FullPullback, :Arnoldi]] -gradient_algs = [[nothing, :FixedPointGradient], [:FixedPointGradient]] +gradient_algs = [[nothing, :FixedPointGradient, :ImplicitGradient], [:FixedPointGradient, :ImplicitGradient]] gradient_solver_algs = [ [:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi], [:GeomSum, :ManualIter, :GMRES, :BiCGStab, :Arnoldi], @@ -36,11 +36,13 @@ naive_gradient_combinations = [ ] naive_gradient_done = Set() -# fixed-point differentiation is incompatible with sequential CTMRG function _check_disallowed_combination( ctmrg_alg, projector_alg, decomposition_rrule_alg, gradient_alg ) - ctmrg_alg == :SequentialCTMRG && !isnothing(gradient_alg) && return true + # characteristic equations for full infinite projector are not implemented + projector_alg == :FullInfiniteProjector && gradient_alg == :ImplicitGradient && return true + # sequential CTMRG doesn't give access to the SVD decompositions + ctmrg_alg == :SequentialCTMRG && gradient_alg == :ImplicitGradient && return true return false end @@ -65,6 +67,21 @@ end calgs, palgs, salgs, galgs, gsalgs ) + # only run GMRES for the implicit gradient, and skip distinction between decomposition rrule algs + if gradient_alg == :ImplicitGradient + gradient_solver_alg == :GMRES || continue + svd_rrule_alg == first(salgs) || continue + end + + # check for allowed algorithm combinations when testing naive gradient + if isnothing(gradient_alg) + combo = (ctmrg_alg, projector_alg, svd_rrule_alg) + combo in naive_gradient_combinations || continue + combo in naive_gradient_done && continue + push!(naive_gradient_done, combo) + gradient_solver_alg = nothing # unused in naive gradient, so set to nothing to avoid confusion + end + # filter disallowed algorithm combinations if _check_disallowed_combination( ctmrg_alg, projector_alg, svd_rrule_alg, gradient_alg @@ -77,15 +94,6 @@ end continue end - # check for allowed algorithm combinations when testing naive gradient - if isnothing(gradient_alg) - combo = (ctmrg_alg, projector_alg, svd_rrule_alg) - combo in naive_gradient_combinations || continue - combo in naive_gradient_done && continue - push!(naive_gradient_done, combo) - gradient_solver_alg = nothing # unused in naive gradient, so set to nothing to avoid confusion - end - @info "optimtest of ctmrg_alg=:$ctmrg_alg, projector_alg=:$projector_alg, svd_rrule_alg=:$svd_rrule_alg and gradient_alg=(; alg = :$gradient_alg, solver_alg = (; alg = :$gradient_solver_alg)) on $(names[i])" Random.seed!(42039482030) dir = InfinitePEPS(Pspace, Vspace)