Patch broken solvers + better testing

SciML · Sep 8, 2023 · e87a82d · e87a82d
1 parent 4fe75e8
commit e87a82d
Show file tree

Hide file tree

Showing 10 changed files with 1,476 additions and 2,070 deletions.
diff --git a/Project.toml b/Project.toml
@@ -27,6 +27,7 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 ArrayInterface = "6.0.24, 7"
 DiffEqBase = "6"
 EnumX = "1"
+Enzyme = "0.11"
 FiniteDiff = "2"
 ForwardDiff = "0.10.3"
 LinearSolve = "2"
@@ -38,19 +39,23 @@ SimpleNonlinearSolve = "0.1"
 SparseDiffTools = "1, 2"
 StaticArraysCore = "1.4"
 UnPack = "1.0"
+Zygote = "0.6"
 julia = "1.6"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra"]
+test = ["Enzyme", "BenchmarkTools", "SafeTestsets", "Pkg", "Test", "ForwardDiff", "StaticArrays", "Symbolics", "LinearSolve", "Random", "LinearAlgebra", "Zygote", "SparseDiffTools"]
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -8,15 +8,16 @@ using DiffEqBase, LinearAlgebra, LinearSolve, SparseDiffTools
 import ForwardDiff
 
 import ADTypes: AbstractFiniteDifferencesMode
-import ArrayInterface: undefmatrix, matrix_colors
+import ArrayInterface: undefmatrix, matrix_colors, parameterless_type, ismutable
 import ConcreteStructs: @concrete
 import EnumX: @enumx
 import ForwardDiff: Dual
 import LinearSolve: ComposePreconditioner, InvPreconditioner, needs_concrete_A
-import RecursiveArrayTools: AbstractVectorOfArray, recursivecopy!, recursivefill!
+import RecursiveArrayTools: ArrayPartition,
+    AbstractVectorOfArray, recursivecopy!, recursivefill!
 import Reexport: @reexport
 import SciMLBase: AbstractNonlinearAlgorithm, NLStats, _unwrap_val, has_jac, isinplace
-import StaticArraysCore: StaticArray, SVector
+import StaticArraysCore: StaticArray, SVector, SArray, MArray
 import UnPack: @unpack
 
 @reexport using ADTypes, SciMLBase, SimpleNonlinearSolve
@@ -33,8 +34,6 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::AbstractNonlinearSolveAl
     return solve!(cache)
 end
 
-# FIXME: Scalar Case is Completely Broken
-
 include("utils.jl")
 include("raphson.jl")
 include("trustRegion.jl")

diff --git a/src/jacobian.jl b/src/jacobian.jl
@@ -36,12 +36,16 @@ function jacobian!!(J::Union{AbstractMatrix{<:Number}, Nothing}, cache)
     @unpack f, uf, u, p, jac_cache, alg, fu2 = cache
     iip = isinplace(cache)
     if iip
-        has_jac(f) ? f.jac(J, u, p) : sparse_jacobian!(J, alg.ad, jac_cache, uf, fu2, u)
+        has_jac(f) ? f.jac(J, u, p) :
+        sparse_jacobian!(J, alg.ad, jac_cache, uf, fu2, _maybe_mutable(u, alg.ad))
     else
-        return has_jac(f) ? f.jac(u, p) : sparse_jacobian!(J, alg.ad, jac_cache, uf, u)
+        return has_jac(f) ? f.jac(u, p) :
+               sparse_jacobian!(J, alg.ad, jac_cache, uf, _maybe_mutable(u, alg.ad))
     end
-    return nothing
+    return J
 end
+# Scalar case
+jacobian!!(::Number, cache) = last(value_derivative(cache.uf, cache.u))
 
 # Build Jacobian Caches
 function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p,
@@ -54,15 +58,16 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p,
     linsolve_needs_jac = (concrete_jac(alg) === nothing &&
                           (!haslinsolve || (haslinsolve && (alg.linsolve === nothing ||
                              needs_concrete_A(alg.linsolve)))))
-    alg_wants_jac = (concrete_jac(alg) === nothing && concrete_jac(alg))
+    alg_wants_jac = (concrete_jac(alg) !== nothing && concrete_jac(alg))
 
     # NOTE: The deepcopy is needed here since we are using the resid_prototype elsewhere
-    fu = f.resid_prototype === nothing ? (iip ? zero(u) : f(u, p)) :
-         deepcopy(f.resid_prototype)
+    fu = f.resid_prototype === nothing ? (iip ? _mutable_zero(u) : _mutable(f(u, p))) :
+         (iip ? deepcopy(f.resid_prototype) : f.resid_prototype)
     if !has_analytic_jac && (linsolve_needs_jac || alg_wants_jac)
         sd = sparsity_detection_alg(f, alg.ad)
-        jac_cache = iip ? sparse_jacobian_cache(alg.ad, sd, uf, fu, u) :
-                    sparse_jacobian_cache(alg.ad, sd, uf, u; fx = fu)
+        ad = alg.ad
+        jac_cache = iip ? sparse_jacobian_cache(ad, sd, uf, fu, _maybe_mutable(u, ad)) :
+                    sparse_jacobian_cache(ad, sd, uf, _maybe_mutable(u, ad); fx = fu)
     else
         jac_cache = nothing
     end
@@ -78,7 +83,7 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p,
         end
     end
 
-    du = zero(u)
+    du = _mutable_zero(u)
     linprob = LinearProblem(J, _vec(fu); u0 = _vec(du))
 
     weight = similar(u)
@@ -90,3 +95,11 @@ function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u, p,
 
     return uf, linsolve, J, fu, jac_cache, du
 end
+
+## Special Handling for Scalars
+function jacobian_caches(alg::AbstractNonlinearSolveAlgorithm, f, u::Number, p,
+    ::Val{false})
+    # NOTE: Scalar `u` assumes scalar output from `f`
+    uf = JacobianWrapper(f, p)
+    return uf, nothing, u, nothing, nothing, u
+end