From e652e81deca8ef012dac9b4e8b8db5ef3058c557 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Mon, 7 Oct 2024 19:49:53 -0400
Subject: [PATCH] test: NonlinearProblem forward diff testing

---
 .../test/core/forward_diff_tests.jl           | 83 +++++++++++++++++++
 1 file changed, 83 insertions(+)

diff --git a/lib/SimpleNonlinearSolve/test/core/forward_diff_tests.jl b/lib/SimpleNonlinearSolve/test/core/forward_diff_tests.jl
index 0005796f9..2b392e39c 100644
--- a/lib/SimpleNonlinearSolve/test/core/forward_diff_tests.jl
+++ b/lib/SimpleNonlinearSolve/test/core/forward_diff_tests.jl
@@ -1,3 +1,86 @@
+@testitem "ForwardDiff.jl Integration: NonlinearProblem" tags=[:core] begin
+    using ArrayInterface
+    using ForwardDiff, FiniteDiff, SimpleNonlinearSolve, StaticArrays, LinearAlgebra,
+          Zygote, ReverseDiff, SciMLBase
+    using DifferentiationInterface
+
+    const DI = DifferentiationInterface
+
+    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)))
+
+    @testset for alg in (
+        SimpleNewtonRaphson(),
+        SimpleTrustRegion(),
+        SimpleTrustRegion(; nlsolve_update_rule = Val(true)),
+        SimpleHalley(),
+        SimpleBroyden(),
+        SimpleKlement(),
+        SimpleDFSane()
+    )
+        us = (
+            2.0,
+            @SVector([1.0, 1.0]),
+            [1.0, 1.0],
+            ones(2, 2),
+            @SArray(ones(2, 2))
+        )
+
+        @testset "Scalar AD" begin
+            for p in 1.0:0.1:100.0, u0 in us
+                sol = solve(NonlinearProblem{false}(test_f, u0, p), alg)
+                if SciMLBase.successful_retcode(sol)
+                    gs = abs.(ForwardDiff.derivative(p) do pᵢ
+                        solve(NonlinearProblem{false}(test_f, u0, pᵢ), alg).u
+                    end)
+                    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_gradient=gs true_gradient=gs_true
+                    else
+                        @test abs.(gs)≈abs.(gs_true) atol=1e-5
+                    end
+                end
+            end
+        end
+
+        @testset "Jacobian" begin
+            @testset "$(typeof(u0))" for u0 in us[2:end],
+                p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0])
+
+                if u0 isa AbstractArray && p isa AbstractArray
+                    size(u0) != size(p) && continue
+                end
+
+                @testset for (iip, fn) in ((false, test_f), (true, test_f!))
+                    iip && (u0 isa Number || !ArrayInterface.can_setindex(u0)) && continue
+
+                    sol = solve(NonlinearProblem{iip}(fn, u0, p), alg)
+                    if SciMLBase.successful_retcode(sol)
+                        gs = abs.(ForwardDiff.jacobian(p) do pᵢ
+                            solve(NonlinearProblem{iip}(fn, u0, pᵢ), alg).u
+                        end)
+                        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
+end
+
 @testitem "ForwardDiff.jl Integration NonlinearLeastSquaresProblem" tags=[:core] begin
     using ForwardDiff, FiniteDiff, SimpleNonlinearSolve, StaticArrays, LinearAlgebra,
           Zygote, ReverseDiff