From 8c8118b66ce0d2234d916b758e614ee43505c7b2 Mon Sep 17 00:00:00 2001
From: Sheehan Olver <solver@mac.com>
Date: Tue, 23 Mar 2021 15:53:49 +0000
Subject: [PATCH] eigen and eigvals for Symmetric, Hermitian, and
 SymTridiagonal

---
 Project.toml         |  4 ++--
 src/ForwardDiff.jl   |  1 +
 src/dual.jl          | 46 ++++++++++++++++++++++++++++++++++++++++++++
 test/JacobianTest.jl |  6 ++++++
 4 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index 7c3bbce8..37bbc96d 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.10.17"
 CommonSubexpressions = "bbf7d656-a473-5ed7-a52c-81e309532950"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 DiffRules = "b552c78f-8df3-52c6-915a-8e097449b14b"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
@@ -26,9 +27,8 @@ julia = "1"
 Calculus = "49dc2e85-a5d0-5ad3-a950-438e2897f1b9"
 DiffTests = "de460e47-3fe3-5279-bb4a-814414816d5d"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Calculus", "DiffTests", "LinearAlgebra", "SparseArrays", "Test", "InteractiveUtils"]
+test = ["Calculus", "DiffTests", "SparseArrays", "Test", "InteractiveUtils"]
diff --git a/src/ForwardDiff.jl b/src/ForwardDiff.jl
index cfafe6a5..5389beb6 100644
--- a/src/ForwardDiff.jl
+++ b/src/ForwardDiff.jl
@@ -4,6 +4,7 @@ using DiffRules, DiffResults
 using DiffResults: DiffResult, MutableDiffResult, ImmutableDiffResult
 using StaticArrays
 using Random
+using LinearAlgebra
 
 import NaNMath
 import SpecialFunctions
diff --git a/src/dual.jl b/src/dual.jl
index 5c9e3628..580b5197 100644
--- a/src/dual.jl
+++ b/src/dual.jl
@@ -625,6 +625,52 @@ if VERSION >= v"1.6.0-DEV.292"
     end
 end
 
+# Symmetric eigvals #
+#-------------------#
+
+function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
+    λ,Q = eigen(Symmetric(value.(parent(A))))
+    parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
+    Dual{Tg}.(λ, tuple.(parts...))
+end
+
+function LinearAlgebra.eigvals(A::Hermitian{<:Complex{<:Dual{Tg,T,N}}}) where {Tg,T<:Real,N}
+    λ,Q = eigen(Hermitian(value.(real.(parent(A))) .+ im .* value.(imag.(parent(A)))))
+    parts = ntuple(j -> diag(real.(Q' * (getindex.(partials.(real.(A)) .+ im .* partials.(imag.(A)), j)) * Q)), N)
+    Dual{Tg}.(λ, tuple.(parts...))
+end
+
+function LinearAlgebra.eigvals(A::SymTridiagonal{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
+    λ,Q = eigen(SymTridiagonal(value.(parent(A).dv),value.(parent(A).ev)))
+    parts = ntuple(j -> diag(Q' * getindex.(partials.(A), j) * Q), N)
+    Dual{Tg}.(λ, tuple.(parts...))
+end
+
+# A ./ (λ - λ') but with diag special cased
+function _lyap_div!(A, λ)
+    for (j,μ) in enumerate(λ), (k,λ) in enumerate(λ)
+        if k ≠ j
+            A[k,j] /= μ - λ
+        end
+    end
+    A
+end
+
+function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
+    λ = eigvals(A)
+    _,Q = eigen(SymTridiagonal(value.(parent(A).dv),value.(parent(A).ev)))
+    parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
+    Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
+end
+
+function LinearAlgebra.eigen(A::SymTridiagonal{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N}
+    λ = eigvals(A)
+    _,Q = eigen(SymTridiagonal(value.(parent(A))))
+    parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N)
+    Eigen(λ,Dual{Tg}.(Q, tuple.(parts...)))
+end
+
+
 ###################
 # Pretty Printing #
 ###################
diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl
index 33a32f2d..befa69ac 100644
--- a/test/JacobianTest.jl
+++ b/test/JacobianTest.jl
@@ -7,6 +7,7 @@ using ForwardDiff
 using ForwardDiff: Dual, Tag, JacobianConfig
 using StaticArrays
 using DiffTests
+using LinearAlgebra
 
 include(joinpath(dirname(@__FILE__), "utils.jl"))
 
@@ -231,4 +232,9 @@ end
     @test_throws DimensionMismatch ForwardDiff.jacobian(sum, fill(2pi, 10^6)) # chunk_mode_jacobian
 end
 
+@testset "eigen" begin
+    @test ForwardDiff.jacobian(x -> eigvals(SymTridiagonal(x, x[1:end-1])), [1.,2.]) ≈ [(1 - 3/sqrt(5))/2 (1 - 1/sqrt(5))/2 ; (1 + 3/sqrt(5))/2 (1 + 1/sqrt(5))/2]
+    @test ForwardDiff.jacobian(x -> eigvals(Symmetric(x*x')), [1.,2.]) ≈ [0 0; 2 4]
+end
+
 end # module