From cdde6e1e15d0da4b868ec61e8a10a8ab32ee23a6 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Thu, 31 Oct 2024 00:13:19 -0400 Subject: [PATCH] fix: forwarddiff support --- ext/NonlinearSolveLeastSquaresOptimExt.jl | 6 +- ext/NonlinearSolveSundialsExt.jl | 18 +- lib/NonlinearSolveBase/src/linear_solve.jl | 4 +- lib/NonlinearSolveBase/src/utils.jl | 2 +- .../src/NonlinearSolveFirstOrder.jl | 2 +- .../test/core_tests.jl | 12 +- src/NonlinearSolve.jl | 28 +-- src/forward_diff.jl | 68 +++--- test/23_test_problems_tests.jl | 9 +- test/cuda_tests.jl | 3 +- test/forward_ad_tests.jl | 212 +++++++++--------- 11 files changed, 203 insertions(+), 161 deletions(-) diff --git a/ext/NonlinearSolveLeastSquaresOptimExt.jl b/ext/NonlinearSolveLeastSquaresOptimExt.jl index 0df265963..6e27f87a9 100644 --- a/ext/NonlinearSolveLeastSquaresOptimExt.jl +++ b/ext/NonlinearSolveLeastSquaresOptimExt.jl @@ -35,9 +35,9 @@ function SciMLBase.__solve( ) end - linsolve = alg.ls === :qr ? LSO.QR() : - (alg.ls === :cholesky ? LSO.Cholesky() : - (alg.ls === :lsmr ? LSO.LSMR() : nothing)) + linsolve = alg.linsolve === :qr ? LSO.QR() : + (alg.linsolve === :cholesky ? LSO.Cholesky() : + (alg.linsolve === :lsmr ? LSO.LSMR() : nothing)) lso_solver = if alg.alg === :lm LSO.LevenbergMarquardt(linsolve) diff --git a/ext/NonlinearSolveSundialsExt.jl b/ext/NonlinearSolveSundialsExt.jl index edea3a49b..303da2167 100644 --- a/ext/NonlinearSolveSundialsExt.jl +++ b/ext/NonlinearSolveSundialsExt.jl @@ -1,16 +1,28 @@ module NonlinearSolveSundialsExt +using Sundials: KINSOL + +using CommonSolve: CommonSolve using NonlinearSolveBase: NonlinearSolveBase, nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution -using NonlinearSolve: DualNonlinearProblem +using NonlinearSolve: NonlinearSolve, DualNonlinearProblem using SciMLBase: SciMLBase -using Sundials: KINSOL function SciMLBase.__solve(prob::DualNonlinearProblem, alg::KINSOL, args...; kwargs...) sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) +end + +function SciMLBase.__init(prob::DualNonlinearProblem, alg::KINSOL, args...; kwargs...) + p = NonlinearSolveBase.nodual_value(prob.p) + newprob = SciMLBase.remake(prob; u0 = NonlinearSolveBase.nodual_value(prob.u0), p) + cache = CommonSolve.init(newprob, alg, args...; kwargs...) + return NonlinearSolveForwardDiffCache( + cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p) + ) end end diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl index ab70e2d6c..f3a2338c9 100644 --- a/lib/NonlinearSolveBase/src/linear_solve.jl +++ b/lib/NonlinearSolveBase/src/linear_solve.jl @@ -69,9 +69,7 @@ function construct_linear_solver(alg, linsolve, A, b, u; stats, kwargs...) error("Default Julia Backsolve Operator `\\` doesn't support Preconditioners") return NativeJLLinearSolveCache(A, b, stats) elseif no_preconditioner && linsolve === nothing - # Non-allocating linear solve exists in StaticArrays.jl - if (A isa SMatrix || A isa WrappedArray{<:Any, <:SMatrix}) && - Core.Compiler.return_type(\, Tuple{typeof(A), typeof(b)}) <: SArray + if (A isa SMatrix || A isa WrappedArray{<:Any, <:SMatrix}) return NativeJLLinearSolveCache(A, b, stats) end end diff --git a/lib/NonlinearSolveBase/src/utils.jl b/lib/NonlinearSolveBase/src/utils.jl index 44c7e957a..99e8e50c5 100644 --- a/lib/NonlinearSolveBase/src/utils.jl +++ b/lib/NonlinearSolveBase/src/utils.jl @@ -222,7 +222,7 @@ function maybe_pinv!!(workspace, A::StridedMatrix) !issingular && return LinearAlgebra.tril!(parent(inv(A_))) else F = LinearAlgebra.lu(A; check = false) - if issuccess(F) + if LinearAlgebra.issuccess(F) Ai = LinearAlgebra.inv!(F) return convert(typeof(parent(Ai)), Ai) end diff --git a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl index 789c13ab7..b611b9051 100644 --- a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl +++ b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl @@ -7,7 +7,7 @@ using Setfield: @set! using ADTypes: ADTypes using ArrayInterface: ArrayInterface -using LinearAlgebra: LinearAlgebra, Diagonal, dot +using LinearAlgebra: LinearAlgebra, Diagonal, dot, diagind using StaticArraysCore: SArray using CommonSolve: CommonSolve diff --git a/lib/NonlinearSolveQuasiNewton/test/core_tests.jl b/lib/NonlinearSolveQuasiNewton/test/core_tests.jl index 6225a07d1..5e70661d8 100644 --- a/lib/NonlinearSolveQuasiNewton/test/core_tests.jl +++ b/lib/NonlinearSolveQuasiNewton/test/core_tests.jl @@ -170,11 +170,14 @@ end LiFukushimaLineSearch() ) @testset "[OOP] u0: $(typeof(u0))" for u0 in (ones(32), @SVector(ones(2)), 1.0) + broken = Sys.iswindows() && u0 isa Vector{Float64} && + linesearch isa BackTracking && ad isa AutoFiniteDiff + solver = LimitedMemoryBroyden(; linesearch) sol = solve_oop(quadratic_f, u0; solver) - @test SciMLBase.successful_retcode(sol) + @test SciMLBase.successful_retcode(sol) broken=broken err = maximum(abs, quadratic_f(sol.u, 2.0)) - @test err < 1e-9 + @test err<1e-9 broken=broken cache = init( NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol = 1e-9 @@ -185,11 +188,14 @@ end @testset "[IIP] u0: $(typeof(u0))" for u0 in (ones(32),) ad isa AutoZygote && continue + broken = Sys.iswindows() && u0 isa Vector{Float64} && + linesearch isa BackTracking && ad isa AutoFiniteDiff + solver = LimitedMemoryBroyden(; linesearch) sol = solve_iip(quadratic_f!, u0; solver) @test SciMLBase.successful_retcode(sol) err = maximum(abs, quadratic_f(sol.u, 2.0)) - @test err < 1e-9 + @test err<1e-9 broken=broken cache = init( NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol = 1e-9 diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index cc71ef03d..2ee0ac575 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -28,8 +28,8 @@ using NonlinearSolveQuasiNewton: Broyden, Klement using SimpleNonlinearSolve: SimpleBroyden, SimpleKlement # Default AD Support -using FiniteDiff: FiniteDiff # Default Finite Difference Method -using ForwardDiff: ForwardDiff # Default Forward Mode AD +using FiniteDiff: FiniteDiff # Default Finite Difference Method +using ForwardDiff: ForwardDiff, Dual # Default Forward Mode AD # Sparse AD Support: Implemented via extensions using SparseArrays: SparseArrays @@ -39,9 +39,9 @@ using SparseMatrixColorings: SparseMatrixColorings using BracketingNonlinearSolve: BracketingNonlinearSolve using LineSearch: LineSearch using LinearSolve: LinearSolve -using NonlinearSolveFirstOrder: NonlinearSolveFirstOrder -using NonlinearSolveQuasiNewton: NonlinearSolveQuasiNewton -using NonlinearSolveSpectralMethods: NonlinearSolveSpectralMethods +using NonlinearSolveFirstOrder: NonlinearSolveFirstOrder, GeneralizedFirstOrderAlgorithm +using NonlinearSolveQuasiNewton: NonlinearSolveQuasiNewton, QuasiNewtonAlgorithm +using NonlinearSolveSpectralMethods: NonlinearSolveSpectralMethods, GeneralizedDFSane using SimpleNonlinearSolve: SimpleNonlinearSolve const SII = SymbolicIndexingInterface @@ -53,16 +53,16 @@ include("extension_algs.jl") include("default.jl") -# const ALL_SOLVER_TYPES = [ -# Nothing, AbstractNonlinearSolveAlgorithm, GeneralizedDFSane, -# GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, -# LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, -# SpeedMappingJL, FixedPointAccelerationJL, SIAMFANLEquationsJL, -# CMINPACK, PETScSNES, -# NonlinearSolvePolyAlgorithm{:NLLS, <:Any}, NonlinearSolvePolyAlgorithm{:NLS, <:Any} -# ] +const ALL_SOLVER_TYPES = [ + Nothing, AbstractNonlinearSolveAlgorithm, + GeneralizedDFSane, GeneralizedFirstOrderAlgorithm, QuasiNewtonAlgorithm, + LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, + SpeedMappingJL, FixedPointAccelerationJL, SIAMFANLEquationsJL, + CMINPACK, PETScSNES, + NonlinearSolvePolyAlgorithm +] -# include("internal/forward_diff.jl") # we need to define after the algorithms +include("forward_diff.jl") @setup_workload begin include("../common/nonlinear_problem_workloads.jl") diff --git a/src/forward_diff.jl b/src/forward_diff.jl index 269ebc99a..28534c59e 100644 --- a/src/forward_diff.jl +++ b/src/forward_diff.jl @@ -1,22 +1,30 @@ -const DualNonlinearProblem = NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, - <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}} where {iip, T, V, P} +const DualNonlinearProblem = NonlinearProblem{ + <:Union{Number, <:AbstractArray}, iip, + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {iip, T, V, P} const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{ <:Union{Number, <:AbstractArray}, iip, - <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}} where {iip, T, V, P} + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {iip, T, V, P} const DualAbstractNonlinearProblem = Union{ - DualNonlinearProblem, DualNonlinearLeastSquaresProblem} + DualNonlinearProblem, DualNonlinearLeastSquaresProblem +} for algType in ALL_SOLVER_TYPES @eval function SciMLBase.__solve( - prob::DualNonlinearProblem, alg::$(algType), args...; kwargs...) - sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) - dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) + prob::DualAbstractNonlinearProblem, alg::$(algType), args...; kwargs... + ) + sol, partials = NonlinearSolveBase.nonlinearsolve_forwarddiff_solve( + prob, alg, args...; kwargs... + ) + dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, prob.p) return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) end end -@concrete mutable struct NonlinearSolveForwardDiffCache +@concrete mutable struct NonlinearSolveForwardDiffCache <: AbstractNonlinearSolveCache cache prob alg @@ -25,36 +33,41 @@ end partials_p end -@internal_caches NonlinearSolveForwardDiffCache :cache - -function reinit_cache!(cache::NonlinearSolveForwardDiffCache; - p = cache.p, u0 = get_u(cache.cache), kwargs...) - inner_cache = reinit_cache!(cache.cache; p = __value(p), u0 = __value(u0), kwargs...) +function InternalAPI.reinit!( + cache::NonlinearSolveForwardDiffCache, args...; + p = cache.p, u0 = NonlinearSolveBase.get_u(cache.cache), kwargs... +) + inner_cache = InternalAPI.reinit!( + cache.cache; p = nodual_value(p), u0 = nodual_value(u0), kwargs... + ) cache.cache = inner_cache cache.p = p - cache.values_p = __value(p) + cache.values_p = nodual_value(p) cache.partials_p = ForwardDiff.partials(p) return cache end for algType in ALL_SOLVER_TYPES + # XXX: Extend to DualNonlinearLeastSquaresProblem @eval function SciMLBase.__init( - prob::DualNonlinearProblem, alg::$(algType), args...; kwargs...) - p = __value(prob.p) - newprob = NonlinearProblem(prob.f, __value(prob.u0), p; prob.kwargs...) + prob::DualNonlinearProblem, alg::$(algType), args...; kwargs... + ) + p = nodual_value(prob.p) + newprob = SciMLBase.remake(prob; u0 = nodual_value(prob.u0), p) cache = init(newprob, alg, args...; kwargs...) return NonlinearSolveForwardDiffCache( - cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p)) + cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p) + ) end end -function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache) +function CommonSolve.solve!(cache::NonlinearSolveForwardDiffCache) sol = solve!(cache.cache) prob = cache.prob uu = sol.u - Jₚ = nonlinearsolve_∂f_∂p(prob, prob.f, uu, cache.values_p) - Jᵤ = nonlinearsolve_∂f_∂u(prob, prob.f, uu, cache.values_p) + Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(prob, prob.f, uu, cache.values_p) + Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, prob.f, uu, cache.values_p) z_arr = -Jᵤ \ Jₚ @@ -65,11 +78,12 @@ function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache) partials = sum(sumfun, zip(eachcol(z_arr), cache.p)) end - dual_soln = nonlinearsolve_dual_solution(sol.u, partials, cache.p) + dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, cache.p) return SciMLBase.build_solution( - prob, cache.alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + prob, cache.alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) end -@inline __value(x) = x -@inline __value(x::Dual) = ForwardDiff.value(x) -@inline __value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) +nodual_value(x) = x +nodual_value(x::Dual) = ForwardDiff.value(x) +nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) diff --git a/test/23_test_problems_tests.jl b/test/23_test_problems_tests.jl index 23b9a3e2a..20c2c3b5c 100644 --- a/test/23_test_problems_tests.jl +++ b/test/23_test_problems_tests.jl @@ -106,10 +106,13 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "Broyden" setup=[RobustnessTesting] tags=[:core] begin - alg_ops = (Broyden(), Broyden(; init_jacobian = Val(:true_jacobian)), +@testitem "Broyden" setup=[RobustnessTesting] tags=[:core] retries=3 begin + alg_ops = ( + Broyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), Broyden(; update_rule = Val(:bad_broyden)), - Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden))) + Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden)) + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) if Sys.isapple() diff --git a/test/cuda_tests.jl b/test/cuda_tests.jl index 91d6178a4..bd992e3fc 100644 --- a/test/cuda_tests.jl +++ b/test/cuda_tests.jl @@ -15,8 +15,7 @@ SOLVERS = ( NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), - # XXX: Fails currently - # LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), + LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), PseudoTransient(), Klement(), Broyden(; linesearch = LiFukushimaLineSearch()), diff --git a/test/forward_ad_tests.jl b/test/forward_ad_tests.jl index e40b0a91d..dcd1c2e3a 100644 --- a/test/forward_ad_tests.jl +++ b/test/forward_ad_tests.jl @@ -1,116 +1,126 @@ -# @testsetup module ForwardADTesting -# using Reexport, NonlinearSolve -# @reexport using ForwardDiff, MINPACK, NLsolve, StaticArrays, Sundials, LinearAlgebra +@testsetup module ForwardADTesting +using Reexport, NonlinearSolve +@reexport using ForwardDiff, MINPACK, NLsolve, StaticArrays, Sundials, LinearAlgebra -# test_f!(du, u, p) = (@. du = u^2 - p) -# test_f(u, p) = (@. u^2 - p) +test_f!(du, u, p) = (@. du = u^2 - p) +test_f(u, p) = (@. u^2 - p) -# jacobian_f(::Number, p) = 1 / (2 * √p) -# jacobian_f(::Number, p::Number) = 1 / (2 * √p) -# jacobian_f(u, p::Number) = one.(u) .* (1 / (2 * √p)) -# jacobian_f(u, p::AbstractArray) = diagm(vec(@. 1 / (2 * √p))) +jacobian_f(::Number, p) = 1 / (2 * √p) +jacobian_f(::Number, p::Number) = 1 / (2 * √p) +jacobian_f(u, p::Number) = one.(u) .* (1 / (2 * √p)) +jacobian_f(u, p::AbstractArray) = diagm(vec(@. 1 / (2 * √p))) -# function solve_with(::Val{mode}, u, alg) where {mode} -# f = if mode === :iip -# solve_iip(p) = solve(NonlinearProblem(test_f!, u, p), alg).u -# elseif mode === :iip_cache -# function solve_iip_init(p) -# cache = SciMLBase.init(NonlinearProblem(test_f!, u, p), alg) -# return SciMLBase.solve!(cache).u -# end -# elseif mode === :oop -# solve_oop(p) = solve(NonlinearProblem(test_f, u, p), alg).u -# elseif mode === :oop_cache -# function solve_oop_init(p) -# cache = SciMLBase.init(NonlinearProblem(test_f, u, p), alg) -# return SciMLBase.solve!(cache).u -# end -# end -# return f -# end +function solve_with(::Val{mode}, u, alg) where {mode} + f = if mode === :iip + solve_iip(p) = solve(NonlinearProblem(test_f!, u, p), alg).u + elseif mode === :iip_cache + function solve_iip_init(p) + cache = SciMLBase.init(NonlinearProblem(test_f!, u, p), alg) + return SciMLBase.solve!(cache).u + end + elseif mode === :oop + solve_oop(p) = solve(NonlinearProblem(test_f, u, p), alg).u + elseif mode === :oop_cache + function solve_oop_init(p) + cache = SciMLBase.init(NonlinearProblem(test_f, u, p), alg) + return SciMLBase.solve!(cache).u + end + end + return f +end -# __compatible(::Any, ::Val{:oop}) = true -# __compatible(::Any, ::Val{:oop_cache}) = true -# __compatible(::Number, ::Val{:iip}) = false -# __compatible(::AbstractArray, ::Val{:iip}) = true -# __compatible(::StaticArray, ::Val{:iip}) = false -# __compatible(::Number, ::Val{:iip_cache}) = false -# __compatible(::AbstractArray, ::Val{:iip_cache}) = true -# __compatible(::StaticArray, ::Val{:iip_cache}) = false +compatible(::Any, ::Val{:oop}) = true +compatible(::Any, ::Val{:oop_cache}) = true +compatible(::Number, ::Val{:iip}) = false +compatible(::AbstractArray, ::Val{:iip}) = true +compatible(::StaticArray, ::Val{:iip}) = false +compatible(::Number, ::Val{:iip_cache}) = false +compatible(::AbstractArray, ::Val{:iip_cache}) = true +compatible(::StaticArray, ::Val{:iip_cache}) = false -# __compatible(::Any, ::Number) = true -# __compatible(::Number, ::AbstractArray) = false -# __compatible(u::AbstractArray, p::AbstractArray) = size(u) == size(p) +compatible(::Any, ::Number) = true +compatible(::Number, ::AbstractArray) = false +compatible(u::AbstractArray, p::AbstractArray) = size(u) == size(p) -# __compatible(u::Number, ::SciMLBase.AbstractNonlinearAlgorithm) = true -# __compatible(u::Number, ::Union{CMINPACK, NLsolveJL, KINSOL}) = true -# __compatible(u::AbstractArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true -# __compatible(u::AbstractArray{T, N}, ::KINSOL) where {T, N} = N == 1 # Need to be fixed upstream -# __compatible(u::StaticArray{S, T, N}, ::KINSOL) where {S <: Tuple, T, N} = false -# __compatible(u::StaticArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true -# __compatible(u::StaticArray, ::Union{CMINPACK, NLsolveJL, KINSOL}) = false -# __compatible(u, ::Nothing) = true +compatible(u::Number, ::SciMLBase.AbstractNonlinearAlgorithm) = true +compatible(u::Number, ::Union{CMINPACK, NLsolveJL, KINSOL}) = true +compatible(u::AbstractArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true +compatible(u::AbstractArray{T, N}, ::KINSOL) where {T, N} = N == 1 # Need to be fixed upstream +compatible(u::StaticArray{S, T, N}, ::KINSOL) where {S <: Tuple, T, N} = false +compatible(u::StaticArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true +compatible(u::StaticArray, ::Union{CMINPACK, NLsolveJL, KINSOL}) = false +compatible(u, ::Nothing) = true -# __compatible(::Any, ::Any) = true -# __compatible(::CMINPACK, ::Val{:iip_cache}) = false -# __compatible(::CMINPACK, ::Val{:oop_cache}) = false -# __compatible(::NLsolveJL, ::Val{:iip_cache}) = false -# __compatible(::NLsolveJL, ::Val{:oop_cache}) = false -# __compatible(::KINSOL, ::Val{:iip_cache}) = false -# __compatible(::KINSOL, ::Val{:oop_cache}) = false +compatible(::Any, ::Any) = true +compatible(::CMINPACK, ::Val{:iip_cache}) = false +compatible(::CMINPACK, ::Val{:oop_cache}) = false +compatible(::NLsolveJL, ::Val{:iip_cache}) = false +compatible(::NLsolveJL, ::Val{:oop_cache}) = false +compatible(::KINSOL, ::Val{:iip_cache}) = false +compatible(::KINSOL, ::Val{:oop_cache}) = false -# export test_f!, test_f, jacobian_f, solve_with, __compatible -# end +export test_f!, test_f, jacobian_f, solve_with, compatible +end -# @testitem "ForwardDiff.jl Integration" setup=[ForwardADTesting] tags=[:core] begin -# for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), -# PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(), DFSane(), -# nothing, NLsolveJL(), CMINPACK(), KINSOL(; globalization_strategy = :LineSearch)) -# us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2)) +@testitem "ForwardDiff.jl Integration" setup=[ForwardADTesting] tags=[:core] begin + @testset for alg in ( + NewtonRaphson(), + TrustRegion(), + LevenbergMarquardt(), + PseudoTransient(; alpha_initial = 10.0), + Broyden(), + Klement(), + DFSane(), + nothing, + NLsolveJL(), + CMINPACK(), + KINSOL(; globalization_strategy = :LineSearch) + ) + us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2)) -# alg isa CMINPACK && Sys.isapple() && continue + alg isa CMINPACK && Sys.isapple() && continue -# @testset "Scalar AD" begin -# for p in 1.0:0.1:100.0, u0 in us, mode in (:iip, :oop, :iip_cache, :oop_cache) -# __compatible(u0, alg) || continue -# __compatible(u0, Val(mode)) || continue -# __compatible(alg, Val(mode)) || continue + @testset "Scalar AD" begin + for p in 1.0:0.1:100.0, u0 in us, mode in (:iip, :oop, :iip_cache, :oop_cache) + compatible(u0, alg) || continue + compatible(u0, Val(mode)) || continue + compatible(alg, Val(mode)) || continue -# sol = solve(NonlinearProblem(test_f, u0, p), alg) -# if SciMLBase.successful_retcode(sol) -# gs = abs.(ForwardDiff.derivative(solve_with(Val{mode}(), u0, alg), p)) -# gs_true = abs.(jacobian_f(u0, p)) -# if !(isapprox(gs, gs_true, atol = 1e-5)) -# @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true -# else -# @test abs.(gs)≈abs.(gs_true) atol=1e-5 -# end -# end -# end -# end + sol = solve(NonlinearProblem(test_f, u0, p), alg) + if SciMLBase.successful_retcode(sol) + gs = abs.(ForwardDiff.derivative(solve_with(Val{mode}(), u0, alg), p)) + gs_true = abs.(jacobian_f(u0, p)) + if !(isapprox(gs, gs_true, atol = 1e-5)) + @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_gradient=gs true_gradient=gs_true + else + @test abs.(gs)≈abs.(gs_true) atol=1e-5 + end + end + end + end -# @testset "Jacobian" begin -# for u0 in us, -# p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0]), -# mode in (:iip, :oop, :iip_cache, :oop_cache) + @testset "Jacobian" begin + for u0 in us, + p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0]), + mode in (:iip, :oop, :iip_cache, :oop_cache) -# __compatible(u0, p) || continue -# __compatible(u0, alg) || continue -# __compatible(u0, Val(mode)) || continue -# __compatible(alg, Val(mode)) || continue + compatible(u0, p) || continue + compatible(u0, alg) || continue + compatible(u0, Val(mode)) || continue + compatible(alg, Val(mode)) || continue -# sol = solve(NonlinearProblem(test_f, u0, p), alg) -# if SciMLBase.successful_retcode(sol) -# gs = abs.(ForwardDiff.jacobian(solve_with(Val{mode}(), u0, alg), p)) -# gs_true = abs.(jacobian_f(u0, p)) -# if !(isapprox(gs, gs_true, atol = 1e-5)) -# @show sol.retcode, sol.u -# @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_jacobian=gs true_jacobian=gs_true -# else -# @test abs.(gs)≈abs.(gs_true) atol=1e-5 -# end -# end -# end -# end -# end -# end + sol = solve(NonlinearProblem(test_f, u0, p), alg) + if SciMLBase.successful_retcode(sol) + gs = abs.(ForwardDiff.jacobian(solve_with(Val{mode}(), u0, alg), p)) + gs_true = abs.(jacobian_f(u0, p)) + if !(isapprox(gs, gs_true, atol = 1e-5)) + @show sol.retcode, sol.u + @error "ForwardDiff Failed for u0=$(u0) and p=$(p) with $(alg)" forwardiff_jacobian=gs true_jacobian=gs_true + else + @test abs.(gs)≈abs.(gs_true) atol=1e-5 + end + end + end + end + end +end