Add OrthAdjoint for left_orth, right_orth

src/Defaults.jl

@@ -162,11 +162,11 @@ const eigh_fwd_alg = :DefaultAlgorithm # ∈ {:<MatrixAlgebraKit.EighAlgorithms>
 const eigh_rrule_alg = :FullPullback # ∈ {:FullPullback, :TruncPullback}
 const eigh_rrule_verbosity = 0
 
-# QR forward & reverse
-const qr_fwd_alg = :DefaultAlgorithm
-const qr_fwd_positive = true
-const qr_rrule_alg = :FullPullback
-const qr_rrule_verbosity = 0
+# Orthogonalization forward & reverse
+const orth_fwd_alg = :DefaultAlgorithm
+const orth_fwd_positive = true
+const orth_rrule_alg = :FullPullback
+const orth_rrule_verbosity = 0
 
 # Projectors
 const projector_alg = :HalfInfiniteProjector # ∈ {:HalfInfiniteProjector, :FullInfiniteProjector}

src/PEPSKit.jl

@@ -62,7 +62,7 @@ include("utility/diffable_threads.jl")
 include("utility/twistdual.jl")
 include("utility/eigh.jl")
 include("utility/svd.jl")
-include("utility/qr.jl")
+include("utility/orth.jl")
 include("utility/rotations.jl")
 include("utility/hook_pullback.jl")
 include("utility/autoopt.jl")
@@ -172,7 +172,7 @@ include("algorithms/select_algorithm.jl")
 
 using .Defaults: set_scheduler!
 export set_scheduler!
-export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, QRAdjoint
+export EighAdjoint, IterEigh, SVDAdjoint, IterSVD, OrthAdjoint, QRAdjoint
 export CTMRGEnv, SequentialCTMRG, SimultaneousCTMRG
 export initialize_ctmrg_environment,
     RandomInitialization, ProductStateInitialization, ApplicationInitialization, IdentityInitialization

src/algorithms/ctmrg/c4v.jl

@@ -21,7 +21,7 @@ For a full description, see [`leading_boundary`](@ref). The supported keywords a
 * `miniter::Int=$(Defaults.ctmrg_miniter)`
 * `verbosity::Int=$(Defaults.ctmrg_verbosity)`
 * `trunc::Union{TruncationStrategy,NamedTuple}=(; alg::Symbol=:$(Defaults.trunc))`
-* `decomposition_alg::Union{NamedTuple,<:EighAdjoint,<:QRAdjoint}=(;)`
+* `decomposition_alg::Union{NamedTuple,<:EighAdjoint,<:OrthAdjoint}=(;)`
 * `projector_alg::Symbol=:$(Defaults.projector_alg_c4v)`
 """
 struct C4vCTMRG{P <: ProjectorAlgorithm} <: CTMRGAlgorithm

src/algorithms/ctmrg/ctmrg.jl

@@ -89,11 +89,9 @@ supplied via the keyword arguments or directly as an [`CTMRGAlgorithm`](@ref) st
     - `:FullInfiniteProjector` : Projection via SVDs of full-infinite (all four enlarged corners) CTMRG environments.
     - `:C4vEighProjector` : Projection via `eigh` of the Hermitian enlarged corner, works only for [`C4vCTMRG`](@ref).
     - `:C4vQRProjector` : Projection via QR decomposition of the lower-rank column-enlarged corner, works only for [`C4vCTMRG`](@ref).
-* `decomposition_alg::Union{NamedTuple,<:SVDAdjoint,<:EighAdjoint,<:QRAdjoint}` : Tensor
-  decomposition algorithm used for computing projectors. When specified as a `NamedTuple`,
-  the settings are passed a the forward algorithm to the appropriate decomposition
-  for the given projector algorithm. For information on which forward algorithms are
-  available, and how to specify them, see [`SVDAdjoint`](@ref), [`EighAdjoint`](@ref) and [`QRAdjoint`](@ref).
+* `decomposition_alg::Union{NamedTuple,<:SVDAdjoint,<:EighAdjoint,<:OrthAdjoint}` : Tensor decomposition algorithm used for computing projectors.
+  When specified as a `NamedTuple`, the settings are passed as the forward algorithm to the appropriate decomposition for the given projector algorithm.
+  For information on which forward algorithms are available, and how to specify them, see [`SVDAdjoint`](@ref), [`EighAdjoint`](@ref) and [`OrthAdjoint`](@ref).
 
 ## Return values
 

src/algorithms/ctmrg/projectors/c4v_qr.jl

@@ -13,7 +13,7 @@ $(TYPEDFIELDS)
 
 Construct the C₄ᵥ `qr`-based projector algorithm based on the following keyword arguments:
 
-* `decomposition_alg=QRAdjoint()` : `left_orth` algorithm including the reverse rule. See [`QRAdjoint`](@ref).
+* `decomposition_alg=OrthAdjoint()` : `left_orth` algorithm including the reverse rule. See [`OrthAdjoint`](@ref).
 """
 struct C4vQRProjector{S} <: ProjectorAlgorithm
     # TODO: support all `left_orth` algorithms

src/algorithms/ctmrg/projectors/projectors.jl

@@ -30,7 +30,7 @@ function ProjectorAlgorithm(;
     elseif alg in [:C4vEighProjector]
         _alg_or_nt(EighAdjoint, decomposition_alg)
     elseif alg in [:C4vQRProjector]
-        _alg_or_nt(QRAdjoint, decomposition_alg)
+        _alg_or_nt(OrthAdjoint, decomposition_alg)
     end # TODO: how do we solve this in a proper way?
 
     # qr-ctmrg does not need truncation or degeneracy checks

src/utility/orth.jl

@@ -0,0 +1,144 @@
+"""
+$(TYPEDEF)
+
+Wrapper for an orthogonalization algorithm `fwd_alg` with a defined reverse rule `rrule_alg`.
+Use the same instance with `left_orth(!)` for QR decomposition and `right_orth(!)` for LQ decomposition.
+
+## Fields
+
+$(TYPEDFIELDS)
+
+## Constructors
+
+    OrthAdjoint(; kwargs...)
+
+Construct an `OrthAdjoint` algorithm struct based on the following keyword arguments:
+
+* `fwd_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.orth_fwd_alg))`: Orthogonalization algorithm of the forward pass which can either be passed as an `Algorithm` instance or a `NamedTuple` where the algorithm is specified by the `alg` keyword.
+  The available algorithms are provided through MatrixAlgebraKit and include:
+    - `:DefaultAlgorithm` : MatrixAlgebraKit's [default orthogonalization algorithm](@extref MatrixAlgebraKit.DefaultAlgorithm) for a given matrix type.
+    - `:Householder` : MatrixAlgebraKit's [`Householder`](@extref MatrixAlgebraKit.Householder)
+* `rrule_alg::Union{Algorithm,NamedTuple}=(; alg::Symbol=$(Defaults.orth_rrule_alg))`: Reverse-rule algorithm for differentiating the QR or LQ decomposition.
+  Can be supplied by an `Algorithm` instance directly or as a `NamedTuple` where `alg` is one of the following:
+    - `:FullPullback` : MatrixAlgebraKit's [`qr_pullback!`](@extref MatrixAlgebraKit.qr_pullback!) or [`lq_pullback!`](@extref MatrixAlgebraKit.lq_pullback!), selected by the orthogonalization direction.
+
+!!! note
+    Manually specifying a `rrule_alg` is considered expert-mode usage, and should only be done when full control over the implementation is desired.
+    For all regular use cases, the default reverse rule algorithms, automatically chosen based on the forward algorithm, should be sufficient.
+"""
+struct OrthAdjoint{F, R}
+    fwd_alg::F
+    rrule_alg::R
+end
+
+"""Compatibility alias for [`OrthAdjoint`](@ref), supporting both QR and LQ decompositions."""
+const QRAdjoint = OrthAdjoint
+
+const ORTH_FWD_SYMBOLS = IdDict{Symbol, Any}(
+    :DefaultAlgorithm => DefaultAlgorithm,
+    :Householder => Householder,
+)
+const ORTH_RRULE_SYMBOLS = IdDict{Symbol, Type{<:Any}}(
+    :FullPullback => FullPullback
+)
+
+function OrthAdjoint(; fwd_alg = (;), rrule_alg = (;))
+    # parse forward algorithm
+    fwd_algorithm = if fwd_alg isa NamedTuple
+        fwd_kwargs = (;
+            alg = Defaults.orth_fwd_alg,
+            positive = Defaults.orth_fwd_positive,
+            fwd_alg...,
+        ) # overwrite with specified kwargs
+        haskey(ORTH_FWD_SYMBOLS, fwd_kwargs.alg) ||
+            throw(ArgumentError("unknown forward algorithm: $(fwd_kwargs.alg)"))
+        fwd_type = ORTH_FWD_SYMBOLS[fwd_kwargs.alg]
+        fwd_kwargs = Base.structdiff(fwd_kwargs, (; alg = nothing)) # remove `alg` keyword argument
+        fwd_type(; fwd_kwargs...)
+    else
+        fwd_alg
+    end
+
+    # parse reverse-rule algorithm
+    rrule_algorithm = if rrule_alg isa NamedTuple
+        rrule_kwargs = (;
+            alg = Defaults.orth_rrule_alg,
+            verbosity = Defaults.orth_rrule_verbosity,
+            rrule_alg...,
+        ) # overwrite with specified kwargs
+
+        haskey(ORTH_RRULE_SYMBOLS, rrule_kwargs.alg) ||
+            throw(ArgumentError("unknown rrule algorithm: $(rrule_kwargs.alg)"))
+        rrule_type = ORTH_RRULE_SYMBOLS[rrule_kwargs.alg]
+        rrule_type()
+    else
+        rrule_alg
+    end
+
+    return OrthAdjoint(fwd_algorithm, rrule_algorithm)
+end
+
+"""
+    left_orth(t, alg::OrthAdjoint)
+    left_orth!(t, alg::OrthAdjoint)
+
+Wrapper around `left_orth(!)` which dispatches on the `OrthAdjoint` algorithm.
+This is needed since a custom adjoint may be defined, depending on the `alg`.
+"""
+MatrixAlgebraKit.left_orth(t, alg::OrthAdjoint) = left_orth!(copy(t), alg)
+MatrixAlgebraKit.left_orth!(t, alg::OrthAdjoint) = left_orth!(t, alg.fwd_alg)
+
+# left_orth! rrule wrapping MatrixAlgebraKit's qr_pullback!
+# https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/blob/b76c7bb60014ecfead6925d0df6cb4b8d7c2668a/src/pullbacks/qr.jl#L49
+function ChainRulesCore.rrule(
+        ::typeof(left_orth!),
+        t::AbstractTensorMap,
+        alg::OrthAdjoint{F, R},
+    ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback}
+    QR = left_orth(t, alg)
+    gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity)
+
+    """Propagate QR cotangents with the configured gauge-sensitivity tolerance."""
+    function left_orth!_pullback(ΔQR)
+        Δt = zeros(storagetype(t), space(t))
+        MatrixAlgebraKit.qr_pullback!(Δt, t, QR, unthunk.(ΔQR); gauge_atol = gtol(ΔQR))
+        return NoTangent(), Δt, NoTangent()
+    end
+    function left_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent})
+        return NoTangent(), ZeroTangent(), NoTangent()
+    end
+
+    return QR, left_orth!_pullback
+end
+
+"""
+    right_orth(t, alg::OrthAdjoint)
+    right_orth!(t, alg::OrthAdjoint)
+
+Wrapper around `right_orth(!)` which dispatches on the `OrthAdjoint` algorithm.
+This is needed since a custom adjoint may be defined, depending on the `alg`.
+"""
+MatrixAlgebraKit.right_orth(t, alg::OrthAdjoint) = right_orth!(copy(t), alg)
+MatrixAlgebraKit.right_orth!(t, alg::OrthAdjoint) = right_orth!(t, alg.fwd_alg)
+
+# right_orth! rrule wrapping MatrixAlgebraKit's lq_pullback!
+function ChainRulesCore.rrule(
+        ::typeof(right_orth!),
+        t::AbstractTensorMap,
+        alg::OrthAdjoint{F, R},
+    ) where {F <: MatrixAlgebraKit.Algorithm, R <: FullPullback}
+    LQ = right_orth(t, alg)
+    gtol = _get_pullback_gauge_tol(alg.rrule_alg.verbosity)
+
+    """Propagate LQ cotangents with the configured gauge-sensitivity tolerance."""
+    function right_orth!_pullback(ΔLQ)
+        Δt = zeros(storagetype(t), space(t))
+        MatrixAlgebraKit.lq_pullback!(Δt, t, LQ, unthunk.(ΔLQ); gauge_atol = gtol(ΔLQ))
+        return NoTangent(), Δt, NoTangent()
+    end
+    function right_orth!_pullback(::Tuple{ZeroTangent, ZeroTangent})
+        return NoTangent(), ZeroTangent(), NoTangent()
+    end
+
+    return LQ, right_orth!_pullback
+end

test/testsuite/TestSuite.jl

@@ -266,6 +266,12 @@ module UtilityEighWrapper
 end
 using .UtilityEighWrapper
 
+module UtilityOrthWrapper
+    include("utility/orth_wrapper.jl")
+    export utility_orth_wrapper
+end
+using .UtilityOrthWrapper
+
 module UtilitySVDWrapper
     include("utility/svd_wrapper.jl")
     export utility_svd_wrapper

test/testsuite/gradients/c4v_ctmrg_gradients.jl

@@ -29,7 +29,7 @@ projector_algs = (:C4vEighProjector, :C4vQRProjector)
 decomposition_alg(projector_alg, rrule_alg) = if projector_alg == :C4vEighProjector
     EighAdjoint(; rrule_alg = (; alg = rrule_alg))
 elseif projector_alg == :C4vQRProjector
-    QRAdjoint(; rrule_alg = (; alg = rrule_alg))
+    OrthAdjoint(; rrule_alg = (; alg = rrule_alg))
 else
     error("unknown projector alg: $projector_alg")
 end