WIP: Split off Nonlinear Solvers and Differentiation

These subpackages will be dependencies of some of the other subpackages. The reason to excise this from the core is to down dependencies required for explicit RK methods.

Todo:

- [ ] DAE initialization needs a good error message for if nonlinear solvers haven't been loaded but you need to run the initialization (because of a initializeprob)
- [ ] Dependencies need to be setup in the Project.tomls of the sublibraries

Author: ChrisRackauckas

lib/OrdinaryDiffEqDifferentiation/Project.toml

@@ -0,0 +1,19 @@
+name = "OrdinaryDiffEqDifferentiation"
+uuid = "4302a76b-040a-498a-8c04-15b101fed76b"
+authors = ["Chris Rackauckas <[email protected]>", "Yingbo Ma <[email protected]>"]
+version = "1.0.0"
+
+[deps]
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+
+[compat]
+julia = "1.10"
+
+[extras]
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["DiffEqDevTools", "Random", "SafeTestsets", "Test"]

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -0,0 +1,28 @@
+module OrdinaryDiffEqDifferentiation
+
+import ADTypes: AutoFiniteDiff, AutoForwardDiff
+
+import SparseDiffTools: SparseDiffTools, matrix_colors, forwarddiff_color_jacobian!,
+                forwarddiff_color_jacobian, ForwardColorJacCache,
+                default_chunk_size, getsize, JacVec
+
+import ForwardDiff, FiniteDiff
+import ForwardDiff.Dual
+import LinearSolve
+
+using DiffEqBase: TimeGradientWrapper,
+                  UJacobianWrapper, TimeDerivativeWrapper,
+                  UDerivativeWrapper
+
+@static if isdefined(DiffEqBase, :OrdinaryDiffEqTag)
+    import DiffEqBase: OrdinaryDiffEqTag
+else
+    struct OrdinaryDiffEqTag end
+end
+
+include("alg_utils.jl")
+include("ilnsolve_utils.jl")
+include("derivative_utils.jl")
+include("derivative_wrappers.jl")
+
+end

lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

@@ -0,0 +1,37 @@
+# Extract AD type parameter from algorithm, returning as Val to ensure type stability for boolean options.
+function _alg_autodiff(alg::OrdinaryDiffEqAlgorithm)
+    error("This algorithm does not have an autodifferentiation option defined.")
+end
+_alg_autodiff(::OrdinaryDiffEqAdaptiveImplicitAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+_alg_autodiff(::DAEAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+_alg_autodiff(::OrdinaryDiffEqImplicitAlgorithm{CS, AD}) where {CS, AD} = Val{AD}()
+_alg_autodiff(alg::CompositeAlgorithm) = _alg_autodiff(alg.algs[end])
+function _alg_autodiff(::Union{OrdinaryDiffEqExponentialAlgorithm{CS, AD},
+        OrdinaryDiffEqAdaptiveExponentialAlgorithm{CS, AD}
+}) where {
+        CS, AD
+}
+    Val{AD}()
+end
+
+function alg_autodiff(alg)
+    autodiff = _alg_autodiff(alg)
+    if autodiff == Val(false)
+        return AutoFiniteDiff()
+    elseif autodiff == Val(true)
+        return AutoForwardDiff()
+    else
+        return _unwrap_val(autodiff)
+    end
+end
+
+Base.@pure function determine_chunksize(u, alg::DiffEqBase.DEAlgorithm)
+    determine_chunksize(u, get_chunksize(alg))
+end
+Base.@pure function determine_chunksize(u, CS)
+    if CS != 0
+        return CS
+    else
+        return ForwardDiff.pickchunksize(length(u))
+    end
+end

src/derivative_utils.jl renamed to lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

@@ -0,0 +1,57 @@
+issuccess_W(W::LinearAlgebra.Factorization) = LinearAlgebra.issuccess(W)
+issuccess_W(W::Number) = !iszero(W)
+issuccess_W(::Any) = true
+
+function dolinsolve(integrator, linsolve; A = nothing, linu = nothing, b = nothing,
+        du = nothing, u = nothing, p = nothing, t = nothing,
+        weight = nothing, solverdata = nothing,
+        reltol = integrator === nothing ? nothing : integrator.opts.reltol)
+    A !== nothing && (linsolve.A = A)
+    b !== nothing && (linsolve.b = b)
+    linu !== nothing && (linsolve.u = linu)
+
+    Plprev = linsolve.Pl isa LinearSolve.ComposePreconditioner ? linsolve.Pl.outer :
+             linsolve.Pl
+    Prprev = linsolve.Pr isa LinearSolve.ComposePreconditioner ? linsolve.Pr.outer :
+             linsolve.Pr
+
+    _alg = unwrap_alg(integrator, true)
+
+    _Pl, _Pr = _alg.precs(linsolve.A, du, u, p, t, A !== nothing, Plprev, Prprev,
+        solverdata)
+    if (_Pl !== nothing || _Pr !== nothing)
+        __Pl = _Pl === nothing ? SciMLOperators.IdentityOperator(length(integrator.u)) : _Pl
+        __Pr = _Pr === nothing ? SciMLOperators.IdentityOperator(length(integrator.u)) : _Pr
+        linsolve.Pl = __Pl
+        linsolve.Pr = __Pr
+    end
+
+    linres = solve!(linsolve; reltol)
+
+    # TODO: this ignores the add of the `f` count for add_steps!
+    if integrator isa SciMLBase.DEIntegrator && _alg.linsolve !== nothing &&
+       !LinearSolve.needs_concrete_A(_alg.linsolve) &&
+       linsolve.A isa WOperator && linsolve.A.J isa AbstractSciMLOperator
+        if alg_autodiff(_alg) isa AutoForwardDiff
+            integrator.stats.nf += linres.iters
+        elseif alg_autodiff(_alg) isa AutoFiniteDiff
+            integrator.stats.nf += 2 * linres.iters
+        else
+            error("$alg_autodiff not yet supported in dolinsolve function")
+        end
+    end
+
+    return linres
+end
+
+function wrapprecs(_Pl::Nothing, _Pr::Nothing, weight, u)
+    Pl = LinearSolve.InvPreconditioner(Diagonal(_vec(weight)))
+    Pr = Diagonal(_vec(weight))
+    Pl, Pr
+end
+
+function wrapprecs(_Pl, _Pr, weight, u)
+    Pl = _Pl === nothing ? SciMLOperators.IdentityOperator(length(u)) : _Pl
+    Pr = _Pr === nothing ? SciMLOperators.IdentityOperator(length(u)) : _Pr
+    Pl, Pr
+end

src/derivative_wrappers.jl renamed to lib/OrdinaryDiffEqDifferentiation/src/derivative_wrappers.jl

@@ -0,0 +1,20 @@
+name = "OrdinaryDiffEqNonlinearSolve"
+uuid = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
+authors = ["Chris Rackauckas <[email protected]>", "Yingbo Ma <[email protected]>"]
+version = "1.0.0"
+
+[deps]
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqDifferentiation = "4302a76b-040a-498a-8c04-15b101fed76b"
+
+[compat]
+julia = "1.10"
+
+[extras]
+DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["DiffEqDevTools", "Random", "SafeTestsets", "Test"]

lib/OrdinaryDiffEqDifferentiation/src/linsolve_utils.jl

@@ -0,0 +1,21 @@
+module OrdinaryDiffEqNonlinearSolve
+
+import ADTypes: AutoFiniteDiff, AutoForwardDiff
+
+import SciMLBase
+import DiffEqBase
+import PreallocationTools
+using SimpleNonlinearSolve: SimpleTrustRegion, SimpleGaussNewton
+using NonlinearSolve: FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg
+
+import OrdinaryDiffEq: resize_nlsolver!, _initialize_dae!
+import OrdinaryDiffEqDifferentiation: update_W!, isnewton
+
+include("type.jl")
+include("utils.jl")
+include("nlsolve.jl")
+include("functional.jl")
+include("newton.jl")
+include("initialize_dae.jl")
+
+end