From 56063bcebae3ae720886a639a26d165c9a6bae00 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Tue, 28 May 2024 23:57:35 +0200 Subject: [PATCH 1/2] Add support for StaticArrays >= 1.7 --- Project.toml | 2 +- ext/ForwardDiffStaticArraysExt.jl | 16 ++++++------- src/dual.jl | 31 ++++++++++++++++++++----- test/JacobianTest.jl | 10 ++++++-- test/MiscTest.jl | 38 +++++++++++++++++++++++++++++++ 5 files changed, 79 insertions(+), 18 deletions(-) diff --git a/Project.toml b/Project.toml index cb88b6cb..45782afe 100644 --- a/Project.toml +++ b/Project.toml @@ -32,7 +32,7 @@ LogExpFunctions = "0.3" NaNMath = "1" Preferences = "1" SpecialFunctions = "1, 2" -StaticArrays = "1.5 - 1.6" +StaticArrays = "1.5" julia = "1.6" [extras] diff --git a/ext/ForwardDiffStaticArraysExt.jl b/ext/ForwardDiffStaticArraysExt.jl index 52914cd6..f2b1540b 100644 --- a/ext/ForwardDiffStaticArraysExt.jl +++ b/ext/ForwardDiffStaticArraysExt.jl @@ -7,7 +7,7 @@ using ForwardDiff: Dual, partials, GradientConfig, JacobianConfig, HessianConfig gradient, hessian, jacobian, gradient!, hessian!, jacobian!, extract_gradient!, extract_jacobian!, extract_value!, vector_mode_gradient, vector_mode_gradient!, - vector_mode_jacobian, vector_mode_jacobian!, valtype, value, _lyap_div! + vector_mode_jacobian, vector_mode_jacobian!, valtype, value using DiffResults: DiffResult, ImmutableDiffResult, MutableDiffResult @generated function dualize(::Type{T}, x::StaticArray) where T @@ -23,19 +23,17 @@ end @inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x)) +# To fix method ambiguity issues: function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) 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...)) + return ForwardDiff._eigvals(A) end - function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}, <:StaticArrays.StaticMatrix}) where {Tg,T<:Real,N} - λ = eigvals(A) - _,Q = eigen(Symmetric(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...))) + return ForwardDiff._eigen(A) end +# For `MMatrix` we can use the in-place method +ForwardDiff._lyap_div!!(A::StaticArrays.MMatrix, λ::AbstractVector) = ForwardDiff._lyap_div!(A, λ) + # Gradient @inline ForwardDiff.gradient(f, x::StaticArray) = vector_mode_gradient(f, x) @inline ForwardDiff.gradient(f, x::StaticArray, cfg::GradientConfig) = gradient(f, x) diff --git a/src/dual.jl b/src/dual.jl index 2ff5366f..419c27b4 100644 --- a/src/dual.jl +++ b/src/dual.jl @@ -719,7 +719,11 @@ end # Symmetric eigvals # #-------------------# -function LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} +# To be able to reuse this default definition in the StaticArrays extension +# (has to be re-defined to avoid method ambiguity issues) +# we forward the call to an internal method that can be shared and reused +LinearAlgebra.eigvals(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} = _eigvals(A) +function _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...)) @@ -737,8 +741,19 @@ function LinearAlgebra.eigvals(A::SymTridiagonal{<:Dual{Tg,T,N}}) where {Tg,T<:R Dual{Tg}.(λ, tuple.(parts...)) end -# A ./ (λ - λ') but with diag special cased -function _lyap_div!(A, λ) +# A ./ (λ' .- λ) but with diag special cased +# Default out-of-place method +function _lyap_div!!(A::AbstractMatrix, λ::AbstractVector) + return map( + (a, b, idx) -> a / (idx[1] == idx[2] ? oneunit(b) : b), + A, + λ' .- λ, + CartesianIndices(A), + ) +end +# For `Matrix` (and e.g. `StaticArrays.MMatrix`) we can use an in-place method +_lyap_div!!(A::Matrix, λ::AbstractVector) = _lyap_div!(A, λ) +function _lyap_div!(A::AbstractMatrix, λ::AbstractVector) for (j,μ) in enumerate(λ), (k,λ) in enumerate(λ) if k ≠ j A[k,j] /= μ - λ @@ -747,17 +762,21 @@ function _lyap_div!(A, λ) A end -function LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} +# To be able to reuse this default definition in the StaticArrays extension +# (has to be re-defined to avoid method ambiguity issues) +# we forward the call to an internal method that can be shared and reused +LinearAlgebra.eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} = _eigen(A) +function _eigen(A::Symmetric{<:Dual{Tg,T,N}}) where {Tg,T<:Real,N} λ = eigvals(A) _,Q = eigen(Symmetric(value.(parent(A)))) - parts = ntuple(j -> Q*_lyap_div!(Q' * getindex.(partials.(A), j) * Q - Diagonal(getindex.(partials.(λ), j)), value.(λ)), N) + 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) + 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 diff --git a/test/JacobianTest.jl b/test/JacobianTest.jl index ef51efc1..8a767a60 100644 --- a/test/JacobianTest.jl +++ b/test/JacobianTest.jl @@ -243,8 +243,14 @@ end @test ForwardDiff.jacobian(ev1, x0) ≈ Calculus.finite_difference_jacobian(ev1, x0) ev2(x) = eigen(SymTridiagonal(x, x[1:end-1])).vectors[:,1] @test ForwardDiff.jacobian(ev2, x0) ≈ Calculus.finite_difference_jacobian(ev2, x0) - x0_static = SVector{2}(x0) - @test ForwardDiff.jacobian(ev1, x0_static) ≈ Calculus.finite_difference_jacobian(ev1, x0) + + x0_svector = SVector{2}(x0) + @test ForwardDiff.jacobian(ev1, x0_svector) isa SMatrix{2, 2} + @test ForwardDiff.jacobian(ev1, x0_svector) ≈ Calculus.finite_difference_jacobian(ev1, x0) + + x0_mvector = MVector{2}(x0) + @test ForwardDiff.jacobian(ev1, x0_mvector) isa MMatrix{2, 2} + @test ForwardDiff.jacobian(ev1, x0_mvector) ≈ Calculus.finite_difference_jacobian(ev1, x0) end @testset "type stability" begin diff --git a/test/MiscTest.jl b/test/MiscTest.jl index 66b71da5..a2888deb 100644 --- a/test/MiscTest.jl +++ b/test/MiscTest.jl @@ -6,6 +6,7 @@ using Test using ForwardDiff using DiffTests using SparseArrays: sparse +using StaticArrays using IrrationalConstants include(joinpath(dirname(@__FILE__), "utils.jl")) @@ -180,4 +181,41 @@ end # example from https://github.com/JuliaDiff/DiffRules.jl/pull/98#issuecomment-1574420052 @test only(ForwardDiff.hessian(t -> abs(t[1])^2, [0.0])) == 2 +@testset "_lyap_div!!" begin + # In-place version for `Matrix` + A = rand(3, 3) + Acopy = copy(A) + λ = rand(3) + B = @inferred(ForwardDiff._lyap_div!!(A, λ)) + @test B === A + @test B[diagind(B)] == Acopy[diagind(Acopy)] + no_diag(X) = [X[i] for i in eachindex(X) if !(i in diagind(X))] + @test no_diag(B) == no_diag(Acopy ./ (λ' .- λ)) + + # Immutable static arrays + A_smatrix = SMatrix{3,3}(Acopy) + λ_svector = SVector{3}(λ) + B_smatrix = @inferred(ForwardDiff._lyap_div!!(A_smatrix, λ_svector)) + @test B_smatrix !== A_smatrix + @test B_smatrix isa SMatrix{3,3} + @test B_smatrix == B + λ_mvector = MVector{3}(λ) + B_smatrix = @inferred(ForwardDiff._lyap_div!!(A_smatrix, λ_mvector)) + @test B_smatrix !== A_smatrix + @test B_smatrix isa SMatrix{3,3} + @test B_smatrix == B + + # Mutable static arrays + A_mmatrix = MMatrix{3,3}(Acopy) + λ_svector = SVector{3}(λ) + B_mmatrix = @inferred(ForwardDiff._lyap_div!!(A_mmatrix, λ_svector)) + @test B_mmatrix === A_mmatrix + @test B_mmatrix == B + A_mmatrix = MMatrix{3,3}(Acopy) + λ_mvector = MVector{3}(λ) + B_mmatrix = @inferred(ForwardDiff._lyap_div!!(A_mmatrix, λ_mvector)) + @test B_mmatrix === A_mmatrix + @test B_mmatrix == B +end + end # module From b84c2d9494f4901e172bd655a443e272a903bf32 Mon Sep 17 00:00:00 2001 From: David Widmann Date: Wed, 29 May 2024 00:12:37 +0200 Subject: [PATCH 2/2] Add missing import --- test/MiscTest.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/test/MiscTest.jl b/test/MiscTest.jl index a2888deb..e7e0d91a 100644 --- a/test/MiscTest.jl +++ b/test/MiscTest.jl @@ -8,6 +8,7 @@ using DiffTests using SparseArrays: sparse using StaticArrays using IrrationalConstants +using LinearAlgebra include(joinpath(dirname(@__FILE__), "utils.jl"))