Merge pull request #2309 from SciML/nonlinearsolve_diff

Split off Nonlinear Solvers and Differentiation

Author: ChrisRackauckas

lib/OrdinaryDiffEqBDF/src/OrdinaryDiffEqBDF.jl

@@ -9,25 +9,26 @@ import OrdinaryDiffEq: alg_order, calculate_residuals!,
                        OrdinaryDiffEqNewtonAlgorithm,
                        AbstractController, DEFAULT_PRECS,
                        CompiledFloats, uses_uprev,
-                       NLNewton, alg_cache, _vec, _reshape, @cache,
-                       isfsal, full_cache, build_nlsolver,
-                       nlsolve!, nlsolvefail, isnewton,
+                       alg_cache, _vec, _reshape, @cache,
+                       isfsal, full_cache,
                        constvalue, isadaptive, error_constant,
-                       DIRK, set_new_W!, has_special_newton_error,
-                       du_alias_or_new, trivial_limiter!,
+                       has_special_newton_error,
+                       trivial_limiter!,
                        ImplicitEulerConstantCache,
-                       compute_step!,
-                       ImplicitEulerCache, COEFFICIENT_MULTISTEP,
-                       markfirststage!, UJacobianWrapper, mul!,
+
+                       ImplicitEulerCache,
                        issplit, qsteady_min_default, qsteady_max_default,
                        get_current_alg_order, get_current_adaptive_order,
                        default_controller, stepsize_controller!, step_accept_controller!,
                        step_reject_controller!, post_newton_controller!,
                        u_modified!, DAEAlgorithm, _unwrap_val, DummyController
 using TruncatedStacktraces, MuladdMacro, MacroTools, FastBroadcast, RecursiveArrayTools
 import StaticArrays: SArray, MVector, SVector, @SVector, StaticArray, MMatrix, SA
-using LinearAlgebra: I
+using LinearAlgebra: mul!, I
 using ArrayInterface
+using OrdinaryDiffEq.OrdinaryDiffEqDifferentiation: UJacobianWrapper
+using OrdinaryDiffEq.OrdinaryDiffEqNonlinearSolve: NLNewton, du_alias_or_new, build_nlsolver,
+nlsolve!, nlsolvefail, isnewton, markfirststage!, set_new_W!, DIRK,  compute_step!, COEFFICIENT_MULTISTEP
 
 include("algorithms.jl")
 include("alg_utils.jl")

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,55 @@
+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
+import LinearSolve: OperatorAssumptions
+import FunctionWrappersWrappers
+using DiffEqBase
+
+import LinearAlgebra
+import LinearAlgebra: Diagonal, I, UniformScaling, diagind, mul!,lmul!, axpby!, opnorm
+import SparseArrays: SparseMatrixCSC, AbstractSparseMatrix, nonzeros
+
+import InteractiveUtils
+import ArrayInterface
+
+import StaticArrayInterface
+import StaticArrays: SArray, MVector, SVector, @SVector, StaticArray, MMatrix, SA,
+                     StaticMatrix
+
+using DiffEqBase: TimeGradientWrapper,
+                  UJacobianWrapper, TimeDerivativeWrapper,
+                  UDerivativeWrapper
+using SciMLBase: AbstractSciMLOperator
+using OrdinaryDiffEq: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveImplicitAlgorithm, DAEAlgorithm,
+OrdinaryDiffEqImplicitAlgorithm, CompositeAlgorithm, OrdinaryDiffEqExponentialAlgorithm,
+OrdinaryDiffEqAdaptiveExponentialAlgorithm, @unpack, AbstractNLSolver, nlsolve_f, issplit,
+concrete_jac, unwrap_alg, OrdinaryDiffEqCache, _vec, standardtag, isnewton, _unwrap_val,
+set_new_W!, set_W_γdt!, alg_difftype, unwrap_cache, diffdir, get_W, isfirstcall, isfirststage, isJcurrent, get_new_W_γdt_cutoff,
+TryAgain, DIRK, COEFFICIENT_MULTISTEP, NORDSIECK_MULTISTEP, GLM, FastConvergence, Convergence, SlowConvergence, VerySlowConvergence, Divergence, NLStatus, MethodType, constvalue
+
+import OrdinaryDiffEq: get_chunksize, resize_J_W!, resize_nlsolver!
+
+import OrdinaryDiffEq: alg_autodiff
+
+using FastBroadcast: @..
+
+@static if isdefined(DiffEqBase, :OrdinaryDiffEqTag)
+    import DiffEqBase: OrdinaryDiffEqTag
+else
+    struct OrdinaryDiffEqTag end
+end
+
+include("alg_utils.jl")
+include("linsolve_utils.jl")
+include("derivative_utils.jl")
+include("derivative_wrappers.jl")
+
+end

lib/OrdinaryDiffEqDifferentiation/src/alg_utils.jl

@@ -0,0 +1,79 @@
+# 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
+
+function DiffEqBase.prepare_alg(
+        alg::Union{
+            OrdinaryDiffEqAdaptiveImplicitAlgorithm{0, AD,
+                FDT},
+            OrdinaryDiffEqImplicitAlgorithm{0, AD, FDT},
+            DAEAlgorithm{0, AD, FDT},
+            OrdinaryDiffEqExponentialAlgorithm{0, AD, FDT}},
+        u0::AbstractArray{T},
+        p, prob) where {AD, FDT, T}
+
+    # If not using autodiff or norecompile mode or very large bitsize (like a dual number u0 already)
+    # don't use a large chunksize as it will either error or not be beneficial
+    if !(alg_autodiff(alg) isa AutoForwardDiff) ||
+       (isbitstype(T) && sizeof(T) > 24) ||
+       (prob.f isa ODEFunction &&
+        prob.f.f isa FunctionWrappersWrappers.FunctionWrappersWrapper)
+        return remake(alg, chunk_size = Val{1}())
+    end
+
+    L = StaticArrayInterface.known_length(typeof(u0))
+    if L === nothing # dynamic sized
+        # If chunksize is zero, pick chunksize right at the start of solve and
+        # then do function barrier to infer the full solve
+        x = if prob.f.colorvec === nothing
+            length(u0)
+        else
+            maximum(prob.f.colorvec)
+        end
+
+        cs = ForwardDiff.pickchunksize(x)
+        return remake(alg, chunk_size = Val{cs}())
+    else # statically sized
+        cs = pick_static_chunksize(Val{L}())
+        return remake(alg, chunk_size = cs)
+    end
+end
+
+@generated function pick_static_chunksize(::Val{chunksize}) where {chunksize}
+    x = ForwardDiff.pickchunksize(chunksize)
+    :(Val{$x}())
+end

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

@@ -626,6 +626,8 @@ function jacobian2W(mass_matrix, dtgamma::Number, J::AbstractMatrix,
     return W
 end
 
+is_always_new(alg) = isdefined(alg, :always_new) ? alg.always_new : false
+
 function calc_W!(W, integrator, nlsolver::Union{Nothing, AbstractNLSolver}, cache, dtgamma,
         repeat_step, W_transform = false, newJW = nothing)
     @unpack t, dt, uprev, u, f, p = integrator
@@ -920,7 +922,7 @@ function LinearSolve.init_cacheval(
         assumptions::OperatorAssumptions)
 end
 
-for alg in InteractiveUtils.subtypes(OrdinaryDiffEq.LinearSolve.AbstractFactorization)
+for alg in InteractiveUtils.subtypes(LinearSolve.AbstractFactorization)
     @eval function LinearSolve.init_cacheval(alg::$alg, A::WOperator, b, u, Pl, Pr,
             maxiters::Int, abstol, reltol, verbose::Bool,
             assumptions::OperatorAssumptions)
@@ -929,3 +931,39 @@ for alg in InteractiveUtils.subtypes(OrdinaryDiffEq.LinearSolve.AbstractFactoriz
             assumptions::OperatorAssumptions)
     end
 end
+
+function resize_J_W!(cache, integrator, i)
+    (isdefined(cache, :J) && isdefined(cache, :W)) || return
+
+    @unpack f = integrator
+
+    if cache.W isa WOperator
+        nf = nlsolve_f(f, integrator.alg)
+        islin = f isa Union{ODEFunction, SplitFunction} && islinear(nf.f)
+        if !islin
+            if cache.J isa AbstractSciMLOperator
+                resize!(cache.J, i)
+            elseif f.jac_prototype !== nothing
+                J = similar(f.jac_prototype, i, i)
+                J = MatrixOperator(J; update_func! = f.jac)
+            end
+            if cache.W.jacvec isa AbstractSciMLOperator
+                resize!(cache.W.jacvec, i)
+            end
+            cache.W = WOperator{DiffEqBase.isinplace(integrator.sol.prob)}(f.mass_matrix,
+                integrator.dt,
+                cache.J,
+                integrator.u,
+                cache.W.jacvec;
+                transform = cache.W.transform)
+            cache.J = cache.W.J
+        end
+    else
+        if cache.J !== nothing
+            cache.J = similar(cache.J, i, i)
+        end
+        cache.W = similar(cache.W, i, i)
+    end
+
+    nothing
+end

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

@@ -0,0 +1,59 @@
+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
+
+Base.resize!(p::LinearSolve.LinearCache, i) = p

lib/OrdinaryDiffEqDifferentiation/src/linsolve_utils.jl

@@ -2,8 +2,8 @@ module OrdinaryDiffEqExponentialRK
 
 import OrdinaryDiffEq: alg_order, alg_adaptive_order, ismultistep, OrdinaryDiffEqExponentialAlgorithm,
                        _unwrap_val, OrdinaryDiffEqMutableCache, OrdinaryDiffEqConstantCache,
-                       build_jac_config, UJacobianWrapper, @cache, alg_cache, UDerivativeWrapper,
-                       initialize!, perform_step!, @unpack, unwrap_alg, calc_J, calc_J!,
+                        @cache, alg_cache,
+                       initialize!, perform_step!, @unpack, unwrap_alg,
                        OrdinaryDiffEqAdaptiveExponentialAlgorithm, CompositeAlgorithm,
                        ExponentialAlgorithm, fsal_typeof, isdtchangeable, calculate_residuals, calculate_residuals!
 using RecursiveArrayTools
@@ -12,6 +12,7 @@ using LinearAlgebra: axpy!, mul!
 using DiffEqBase, SciMLBase
 using ExponentialUtilities
 import RecursiveArrayTools: recursivecopy!
+using OrdinaryDiffEq.OrdinaryDiffEqDifferentiation: build_jac_config, UJacobianWrapper, UDerivativeWrapper, calc_J, calc_J!
 
 include("algorithms.jl")
 include("alg_utils.jl")

lib/OrdinaryDiffEqDifferentiation/test/runtests.jl

@@ -14,10 +14,10 @@ import OrdinaryDiffEq: alg_order, alg_maximum_order, get_current_adaptive_order,
                        DEFAULT_PRECS,
                        constvalue, PolyesterThreads, Sequential, BaseThreads,
                        _digest_beta1_beta2, timedepentdtmin, _unwrap_val,
-                       TimeDerivativeWrapper, UDerivativeWrapper, calc_J, _reshape, _vec,
-                       WOperator, TimeGradientWrapper, UJacobianWrapper, build_grad_config,
-                       build_jac_config, calc_J!, jacobian2W!, dolinsolve
+                       _reshape, _vec
 using DiffEqBase, FastBroadcast, Polyester, MuladdMacro, RecursiveArrayTools, LinearSolve
+import OrdinaryDiffEq.OrdinaryDiffEqDifferentiation: TimeDerivativeWrapper, UDerivativeWrapper, calc_J, WOperator, TimeGradientWrapper, UJacobianWrapper,                        build_grad_config,
+build_jac_config, calc_J!, jacobian2W!, dolinsolve
 
 include("algorithms.jl")
 include("alg_utils.jl")