Differentiate general eigen and eigvals

src/dual.jl

@@ -823,6 +823,29 @@ function _eigvals_hermitian(A::DualMatrixRealComplex{T,<:Real,N}) where {T,N}
     return _to_duals(Val(T), λ, parts)
 end
 
+@noinline function _throw_repeated_eigvals(i, j, λ)
+    throw(ArgumentError(lazy"eigenvector derivatives are not defined for repeated eigenvalues, but λ[$i] == λ[$j] == $λ"))
+end
+
+# `_lyap_div_zero_diag!!` zeroes the diagonal, where `λ[j] - λ[i]` vanishes by construction.
+# Two equal eigenvalues make an off-diagonal denominator vanish as well, and the eigenvector
+# partials come out as `Inf` or `NaN`: the eigenvectors of a repeated eigenvalue are not
+# unique and hence not differentiable. The `eigen` methods below call this once, right after
+# the decomposition and before anything can fail for another reason, rather than from inside
+# `_lyap_div_zero_diag!!`, which runs once per partial direction.
+#
+# The eigenvalues are compared as they are, without extracting primal values: for nested
+# `Dual`s two of them can be `!=` here and still divide to `Inf`, but then their primals
+# coincide, and the decomposition of the values -- which every method computes first, one
+# `Dual` level down -- has already seen them as exact duplicates and thrown.
+function _check_distinct_eigvals(λ::AbstractVector)
+    for j in eachindex(λ), i in eachindex(λ)
+        i < j || continue
+        λ[i] == λ[j] && _throw_repeated_eigvals(i, j, λ[i])
+    end
+    return nothing
+end
+
 # A ./ (λ' .- λ) but with diagonal elements zeroed out
 # Default out-of-place method
 function _lyap_div_zero_diag!!(A::AbstractMatrix, λ::AbstractVector)
@@ -856,6 +879,7 @@ function _eigen_hermitian(A::DualMatrixRealComplex{T,<:Real,N}) where {T,N}
     F = eigen(_structured_value(A))
     λ = F.values
     Q = F.vectors
+    _check_distinct_eigvals(λ)
     # `Q' * (∂A * Q)`, not `(Q' * ∂A) * Q`: the latter hits `Adjoint * Symmetric`, which has no BLAS
     # specialization and so allocates an extra temporary and skips `symm`/`hemm`
     Qt_∂A_Q = ntuple(j -> Q' * (_structured_partials(A, j) * Q), N)
@@ -864,6 +888,136 @@ function _eigen_hermitian(A::DualMatrixRealComplex{T,<:Real,N}) where {T,N}
     return Eigen(_to_duals(Val(T), λ, λ_partials), _to_duals(Val(T), Q, Q_partials))
 end
 
+# General eigvals and eigen #
+#---------------------------#
+
+# The derivatives are computed entirely from the values of `A`, i.e. one `Dual` level
+# below the entries of `A`, and are only assembled into `Dual`s at the very end. Mixing
+# the two levels in a single expression (e.g. multiplying the eigenvectors of `value.(A)`
+# by `A` itself) would make the same-tag product rule apply at the outer level and hence
+# give wrong results for nested `Dual`s, as in `ForwardDiff.hessian`.
+#
+# The formulas are the ones of https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf:
+# with `A = U * Diagonal(λ) * inv(U)` and `M = inv(U) * Ȧ * U`, we have `λ̇ = diag(M)`
+# and `U̇ = U * (F .* M)`, where `F[i,j] = inv(λ[j] - λ[i])` off the diagonal and zero on it.
+
+@noinline function _throw_no_real_eigvec_entry()
+    throw(ArgumentError("no entry of an eigenvector is real, so the entry that carries the phase convention cannot be identified; the derivatives assume the normalization of `eigen`, whose eigenvectors have a real entry of largest magnitude"))
+end
+
+# Index of the entry of largest magnitude among the real entries of `v`. LAPACK normalizes
+# the eigenvectors it returns to unit 2-norm with their largest entry real, so this is the
+# entry that carries the phase convention. Modelled after `_findrealmaxabs2` in ChainRules,
+# except that reaching the end without a real entry is an error rather than a silent
+# fallback: it means the assumed normalization does not hold.
+function _findrealmaxabs2(v)
+    imax = firstindex(v)
+    amax = abs2(zero(eltype(v)))
+    found = false
+    for i in eachindex(v)
+        vi = v[i]
+        isreal(vi) || continue
+        a = abs2(vi)
+        a < amax && continue
+        amax, imax = a, i
+        found = true
+    end
+    found || _throw_no_real_eigvec_entry()
+    return imax
+end
+
+# The diagonal of `F` is a free gauge parameter: `U̇ + U * Diagonal(ċ)` solves the same
+# differentiated eigenvalue equation for any `ċ`. Setting it to zero, as `U̇ = U * (F .* M)`
+# does, normalizes the eigenvectors by `diag(inv(U) * U̇) == 0`, which for a non-normal `A`
+# is not the convention `eigen` itself returns. Choose `ċ` instead such that the derivatives
+# belong to LAPACK's convention, i.e. differentiate its two constraints at `t = 0`:
+#
+#   `u' * u == 1`             ⇒  `real(ċ) = -real(u' * u̇)`
+#   `imag(u[k]) == 0`         ⇒  `imag(ċ) = -imag(u̇[k]) / real(u[k])`,  `k = _findrealmaxabs2(u)`
+#
+# This mirrors `_eigen_norm_phase_fwd!` in ChainRules, so that both agree on `eigen`.
+function _eigen_norm_phase!(U̇, U)
+    for i in axes(U, 2)
+        u, u̇ = view(U, :, i), view(U̇, :, i)
+        ċ_norm = -real(dot(u, u̇))
+        if eltype(U) <: Real
+            u̇ .+= u .* ċ_norm
+        else
+            k = _findrealmaxabs2(u)
+            u̇ .+= u .* complex(ċ_norm, -imag(u̇[k]) / real(u[k]))
+            # `imag(u[k]) == 0` holds identically along the curve, so `imag(u̇[k])` is exactly
+            # zero. Leaving it as the cancellation `q + fl(r * fl(-q / r))` keeps a rounding
+            # residue instead, and one differentiation level up that residue sits in the
+            # partials, where `isreal` sees it: `_findrealmaxabs2` then finds no real entry
+            # in the column and fixes the phase at the wrong one. Being exact rather than
+            # approximate, this also keeps third and higher derivatives right.
+            u̇[k] = complex(real(u̇[k]), zero(real(u̇[k])))
+        end
+    end
+    return U̇
+end
+
+# A matrix that is Hermitian in both its values and its partials is handled by the
+# `Symmetric` methods, which is what `LinearAlgebra.eigen!` and `eigvals!` do for the
+# values as well. Their eigenvalues are real, so the eltype no longer depends on whether
+# `geevx!` happens to produce a nonzero imaginary part.
+#
+# Those methods take no keyword arguments. `permute` and `scale` only control the balancing
+# that the symmetric algorithm does not use, so they can be ignored the way
+# `LinearAlgebra.eigen!` ignores them, but any `sortby` other than the ascending order that
+# is returned anyway has to go through the general path.
+function _use_symmetric(A; permute::Bool=true, scale::Bool=true,
+                        sortby::Union{Function,Nothing}=LinearAlgebra.eigsortby)
+    return (sortby === nothing || sortby === LinearAlgebra.eigsortby) && ishermitian(A)
+end
+
+# `value.(A)` is a temporary of our own, so the decomposition may consume it. `eigen!` only
+# exists for BLAS element types, which is the innermost level of the nesting; above it the
+# recursion goes through the copying `eigen` below. `!!` as in `_lyap_div_zero_diag!!`.
+_eigen!!(B::StridedMatrix{<:LinearAlgebra.BlasFloat}; kwargs...) = eigen!(B; kwargs...)
+_eigen!!(B::AbstractMatrix; kwargs...) = eigen(B; kwargs...)
+
+# `permute`, `scale` and `sortby` are forwarded to the underlying decomposition of
+# `value.(A)`; the derivatives are assembled in whatever order it returns, and for nested
+# `Dual`s every level is decomposed with the same keyword arguments
+function LinearAlgebra.eigvals(A::StridedMatrix{Dual{Tg,T,N}}; kwargs...) where {Tg,T<:Real,N}
+    _use_symmetric(A; kwargs...) && return _eigvals_hermitian(Symmetric(A))
+    return _eigvals_general(A; kwargs...)
+end
+function _eigvals_general(A::StridedMatrix{Dual{Tg,T,N}}; kwargs...) where {Tg,T<:Real,N}
+    λ, U = _eigen!!(value.(A); kwargs...)
+    luU = lu(U)
+    # `diag(inv(U) * Ȧ * U)` without forming the second product: entry `i` is row `i` of
+    # `inv(U) * Ȧ` against column `i` of `U`, which is `n` multiplications rather than a
+    # matmul. Only the diagonal is wanted here, unlike in `_eigen_general`. One `n^2` buffer
+    # serves every direction, so the loop allocates only the result vectors.
+    B = similar(U)
+    parts = ntuple(N) do j
+        B .= partials.(A, j)
+        ldiv!(luU, B)
+        map((b, u) -> sum(prod, zip(b, u)), eachrow(B), eachcol(U))
+    end
+    return _to_duals(Val(Tg), λ, parts)
+end
+
+function LinearAlgebra.eigen(A::StridedMatrix{Dual{Tg,T,N}}; kwargs...) where {Tg,T<:Real,N}
+    _use_symmetric(A; kwargs...) && return _eigen_hermitian(Symmetric(A))
+    return _eigen_general(A; kwargs...)
+end
+function _eigen_general(A::StridedMatrix{Dual{Tg,T,N}}; kwargs...) where {Tg,T<:Real,N}
+    λ, U = _eigen!!(value.(A); kwargs...)
+    # before `lu`, so that a repeated eigenvalue is reported as such rather than surfacing
+    # as a `SingularException` from the factorization of a defective `U`
+    _check_distinct_eigvals(λ)
+    luU = lu(U)
+    M = ntuple(j -> ldiv!(luU, partials.(A, j) * U), N)
+    λ_parts = map(diag, M)
+    # `_lyap_div_zero_diag!!` zeroes the diagonal itself, so `M[j]` no longer has to have
+    # `Diagonal(λ_parts[j])` subtracted from it first
+    U_parts = ntuple(j -> _eigen_norm_phase!(U * _lyap_div_zero_diag!!(M[j], λ), U), N)
+    return Eigen(_to_duals(Val(Tg), λ, λ_parts), _to_duals(Val(Tg), U, U_parts))
+end
+
 # Functions in SpecialFunctions which return tuples #
 # Their derivatives are not defined in DiffRules    #
 #---------------------------------------------------#

test/HessianTest.jl

@@ -163,6 +163,156 @@ end
     @test ForwardDiff.hessian(x->dot(x,H,x), zeros(3)) ≈ [2 6 10; 6 10 14; 10 14 18]
 end
 
+@testset "nested duals in general eigen" begin
+    # The eigenvalue derivatives have to be computed from `value.(A)` alone; mixing that
+    # level with `A` itself applies the product rule at the wrong level for nested `Dual`s
+    B(w) = [3.0+w[1] 1.0+w[2]; 0.4+w[2] 2.0-2*w[1]]
+    w = [0.11, -0.07]
+    # sum(eigvals(B(w))) == tr(B(w)) == 5 - w[1] is linear in `w`
+    @test ForwardDiff.hessian(w -> sum(eigvals(B(w))), w) ≈ zeros(2, 2) atol=1e-12
+    # sum(eigvals(B(w)) .^ 2) == tr(B(w)^2) is quadratic in `w`
+    @test ForwardDiff.hessian(w -> sum(eigvals(B(w)) .^ 2), w) ≈ [10 0; 0 4]
+
+    # keyword arguments reach every level of the nesting
+    @test ForwardDiff.hessian(w -> sum(eigvals(B(w); sortby = nothing)), w) ≈ zeros(2, 2) atol=1e-12
+    @test ForwardDiff.hessian(w -> sum(eigvals(B(w); permute = false, scale = false) .^ 2), w) ≈ [10 0; 0 4]
+
+    # complex eigenvalues: λ = w[1] ± im*(1 + w[2]), i.e. sum(abs2, λ) == 2*(w[1]^2 + (1 + w[2])^2)
+    C(w) = [w[1] -1.0-w[2]; 1.0+w[2] w[1]]
+    @test ForwardDiff.hessian(w -> sum(abs2, eigvals(C(w))), [0.3, 0.2]) ≈ [4 0; 0 4]
+
+    # https://github.com/JuliaDiff/ForwardDiff.jl/issues/111. `S(w)` is Hermitian in its
+    # values *and* its partials, so both sides take the `Symmetric` shortcut: this pins the
+    # dispatch, not the general path, and the two Hessians are bitwise identical.
+    S(w) = [w[1]^2 w[1]*w[2]*w[3]; w[1]*w[2]*w[3] w[2]^2]
+    g(w) = sum(log, eigvals(S(w)))
+    gsym(w) = sum(log, eigvals(Symmetric(S(w))))
+    w111 = [0.9, 1.4, 0.3]
+    @test ishermitian(S(w111))
+    @test ForwardDiff.hessian(g, w111) == ForwardDiff.hessian(gsym, w111)
+
+    # The same log-determinant Hessian on a matrix that is *not* Hermitian, so that the
+    # general path is what is under test. `det(Bgen(w)) == w[1]^2 * (1 + w[2]^2 / 2)`, which
+    # gives a closed-form reference that never goes through `eigen`.
+    Bgen(w) = [w[1]^2 w[1]*w[2]; 0.5*w[1]*w[2] w[2]^2+1]
+    hgen(w) = sum(log, eigvals(Bgen(w)))
+    wgen = [0.9, 1.4]
+    @test !ishermitian(Bgen(wgen))
+    @test isreal(eigvals(Bgen(wgen)))
+    @test ForwardDiff.hessian(hgen, wgen) ≈
+          [-2/wgen[1]^2 0; 0 (1 - wgen[2]^2/2)/(1 + wgen[2]^2/2)^2]
+    @test ForwardDiff.hessian(hgen, wgen) ≈ ForwardDiff.hessian(w -> log(det(Bgen(w))), wgen)
+
+    # eigenvectors
+    A0 = [2.0 1.0 0.5; 0.5 3.0 1.5; 0.25 0.75 4.0]
+    function v1(x)
+        v = eigen(reshape(x, 3, 3)).vectors[:, 1]
+        return sum(abs2, v .- [1.0, 0.5, -0.2])
+    end
+    x0 = vec(A0)
+    @test ForwardDiff.hessian(v1, x0) ≈ ForwardDiff.jacobian(x -> ForwardDiff.gradient(v1, x), x0)
+    @test ForwardDiff.hessian(v1, x0) ≈ Calculus.finite_difference_jacobian(x -> ForwardDiff.gradient(v1, x), x0) atol=1e-5
+end
+
+# Helpers for the eigenvector phase gauge tests below
+struct GaugeTag end
+
+# A deterministic `n x n` matrix with a complex spectrum and well separated eigenvalues
+function complex_spectrum_matrix(n)
+    A = zeros(n, n)
+    for b in 1:(n ÷ 2)
+        i = 2b - 1
+        A[i, i] = A[i+1, i+1] = 0.5 + b / 4
+        A[i, i+1] = -1.0 - b / 8
+        A[i+1, i] = 1.0 + b / 8
+    end
+    isodd(n) && (A[n, n] = 2.0)
+    for i in 1:n, j in 1:n
+        A[i, j] += 0.15 * sinpi((i + 2j) / (n + 1))
+    end
+    return A
+end
+
+seed_matrix(n, i, j) = (S = zeros(n, n); S[i, j] = 1.0; S)
+
+# `A` lifted to a `Dual` of nesting depth `k` with vanishing partials
+lift_dual(A, k) = k == 0 ? A : Dual{GaugeTag}.(lift_dual(A, k - 1), lift_dual(zero(A), k - 1))
+# `A` seeded with one partial per level, innermost seed first, the way `hessian` nests them
+function nest_dual(A, seeds...)
+    M = A
+    for (k, seed) in enumerate(seeds)
+        M = Dual{GaugeTag}.(M, lift_dual(seed, k - 1))
+    end
+    return M
+end
+
+# The entry `eigen` made real, i.e. the one of largest magnitude
+phase_index(v) = argmax(j -> abs2(v[j]), eachindex(v))
+
+# All components of a (possibly nested) `Dual`: its value and every partial, recursively
+function dual_components!(out, x)
+    if x isa Dual
+        dual_components!(out, ForwardDiff.value(x))
+        for i in 1:ForwardDiff.npartials(x)
+            dual_components!(out, ForwardDiff.partials(x, i))
+        end
+    else
+        push!(out, x)
+    end
+    return out
+end
+dual_components(x) = dual_components!(Float64[], x)
+
+# The eigenvectors and their first derivative in direction `seed`, with the columns flipped
+# to match the signs of `Vref`
+function eigvecs_and_derivative(X, seed, Vref)
+    V = eigen(Dual{GaugeTag}.(X, seed)).vectors
+    V0 = map(z -> complex(ForwardDiff.value(real(z)), ForwardDiff.value(imag(z))), V)
+    V1 = map(z -> complex(ForwardDiff.partials(real(z), 1), ForwardDiff.partials(imag(z), 1)), V)
+    for i in axes(V0, 2)
+        k = phase_index(view(Vref, :, i))
+        if real(V0[k, i]) * real(Vref[k, i]) < 0
+            V0[:, i] .*= -1
+            V1[:, i] .*= -1
+        end
+    end
+    return V0, V1
+end
+
+@testset "eigenvector phase gauge under nesting, n = $n" for n in 2:6
+    A = complex_spectrum_matrix(n)
+    # otherwise the phase convention never comes up
+    @test !isreal(eigvals(A))
+    seeds = (seed_matrix(n, 1, 1), seed_matrix(n, 2, min(3, n)), seed_matrix(n, min(3, n), 1))
+
+    # `eigen` returns eigenvectors of unit 2-norm whose largest-magnitude entry is real, so
+    # `imag(u[k]) == 0` and `u' * u == 1` hold identically along the curve and every
+    # derivative of them has to vanish. A rounding residue in the first is invisible at
+    # first order, but one level up it sits in the partials, where `isreal` sees it, and
+    # `_findrealmaxabs2` then fixes the phase at the wrong entry.
+    @testset "nesting depth $depth" for depth in 1:3
+        V = eigen(nest_dual(A, seeds[1:depth]...)).vectors
+        for i in axes(V, 2)
+            v = view(V, :, i)
+            @test iszero(imag(V[phase_index(v), i]))
+            @test all(x -> abs(x) < 1e-12, dual_components(real(dot(v, v)) - 1))
+        end
+    end
+
+    # second derivatives of the eigenvectors, against central differences of the first
+    # derivatives. `eigen` does not pin the sign of the real entry, so the columns have to
+    # be realigned before differencing; plain central differences are off by `O(1)`.
+    Ea, Eb = seeds[1], seeds[2]
+    Vref = eigen(A).vectors
+    V = eigen(nest_dual(A, Ea, Eb)).vectors
+    ad = map(z -> complex(ForwardDiff.partials(ForwardDiff.partials(real(z), 1), 1),
+                          ForwardDiff.partials(ForwardDiff.partials(imag(z), 1), 1)), V)
+    h = 1e-5
+    _, Dp = eigvecs_and_derivative(A .+ h .* Eb, Ea, Vref)
+    _, Dm = eigvecs_and_derivative(A .- h .* Eb, Ea, Vref)
+    @test maximum(abs, ad .- (Dp .- Dm) ./ (2h)) < 1e-6
+end
+
 #https://github.com/JuliaDiff/ForwardDiff.jl/issues/720
 @testset "allocation-free hessian with StaticArrays" begin
     function hessian_allocs()