simplify dolinsolve

Author: oscarddssmith

Project.toml

@@ -121,8 +121,8 @@ OrdinaryDiffEqQPRK = "1"
 OrdinaryDiffEqRKN = "1"
 OrdinaryDiffEqRosenbrock = "1"
 OrdinaryDiffEqSDIRK = "1"
-OrdinaryDiffEqStabilizedIRK = "1"
 OrdinaryDiffEqSSPRK = "1"
+OrdinaryDiffEqStabilizedIRK = "1"
 OrdinaryDiffEqStabilizedRK = "1"
 OrdinaryDiffEqSymplecticRK = "1"
 OrdinaryDiffEqTsit5 = "1"

lib/OrdinaryDiffEqBDF/src/OrdinaryDiffEqBDF.jl

@@ -7,8 +7,7 @@ import OrdinaryDiffEqCore: alg_order, calculate_residuals!,
                            OrdinaryDiffEqMutableCache, OrdinaryDiffEqConstantCache,
                            OrdinaryDiffEqNewtonAdaptiveAlgorithm,
                            OrdinaryDiffEqNewtonAlgorithm,
-                           AbstractController, DEFAULT_PRECS,
-                           CompiledFloats, uses_uprev,
+                           AbstractController, CompiledFloats, uses_uprev,
                            alg_cache, _vec, _reshape, @cache,
                            isfsal, full_cache,
                            constvalue, isadaptive, error_constant,

lib/OrdinaryDiffEqBDF/src/algorithms.jl

@@ -105,10 +105,10 @@ end
     Divergence = -2
 end
 const TryAgain = SlowConvergence
-
-DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, solverdata) = nothing, nothing
 isdiscretecache(cache) = false
 
+# unused. Delete this once StocasticDiffEq doesn't use it
+DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, solverdata) = nothing, nothing
 include("doc_utils.jl")
 include("misc_utils.jl")
 

lib/OrdinaryDiffEqCore/src/OrdinaryDiffEqCore.jl

@@ -87,7 +87,6 @@ function differentiation_rk_docstring(description::String,
     concrete_jac = nothing,
     diff_type = Val{:forward},
     linsolve = nothing,
-    precs = DEFAULT_PRECS,
     """ * extra_keyword_default
 
     keyword_default_description = """
@@ -111,40 +110,7 @@ function differentiation_rk_docstring(description::String,
       For example, to use [KLU.jl](https://github.com/JuliaSparse/KLU.jl), specify
       `$name(linsolve = KLUFactorization()`).
        When `nothing` is passed, uses `DefaultLinearSolver`.
-    - `precs`: Any [LinearSolve.jl-compatible preconditioner](https://docs.sciml.ai/LinearSolve/stable/basics/Preconditioners/)
-      can be used as a left or right preconditioner.
-      Preconditioners are specified by the `Pl,Pr = precs(W,du,u,p,t,newW,Plprev,Prprev,solverdata)`
-      function where the arguments are defined as:
-        - `W`: the current Jacobian of the nonlinear system. Specified as either
-            ``I - \\gamma J`` or ``I/\\gamma - J`` depending on the algorithm. This will
-            commonly be a `WOperator` type defined by OrdinaryDiffEq.jl. It is a lazy
-            representation of the operator. Users can construct the W-matrix on demand
-            by calling `convert(AbstractMatrix,W)` to receive an `AbstractMatrix` matching
-            the `jac_prototype`.
-        - `du`: the current ODE derivative
-        - `u`: the current ODE state
-        - `p`: the ODE parameters
-        - `t`: the current ODE time
-        - `newW`: a `Bool` which specifies whether the `W` matrix has been updated since
-            the last call to `precs`. It is recommended that this is checked to only
-            update the preconditioner when `newW == true`.
-        - `Plprev`: the previous `Pl`.
-        - `Prprev`: the previous `Pr`.
-        - `solverdata`: Optional extra data the solvers can give to the `precs` function.
-            Solver-dependent and subject to change.
-      The return is a tuple `(Pl,Pr)` of the LinearSolve.jl-compatible preconditioners.
-      To specify one-sided preconditioning, simply return `nothing` for the preconditioner
-      which is not used. Additionally, `precs` must supply the dispatch:
-      ```julia
-      Pl, Pr = precs(W, du, u, p, t, ::Nothing, ::Nothing, ::Nothing, solverdata)
-      ```
-      which is used in the solver setup phase to construct the integrator
-      type with the preconditioners `(Pl,Pr)`.
-      The default is `precs=DEFAULT_PRECS` where the default preconditioner function
-      is defined as:
-      ```julia
-      DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, solverdata) = nothing, nothing
-      ```
+
     """ * extra_keyword_description
 
     generic_solver_docstring(

lib/OrdinaryDiffEqCore/src/doc_utils.jl

@@ -49,7 +49,7 @@ rosensol = solve(prob_rober, AutoTsit5(Rosenbrock23(autodiff = false)))
 sol = solve(prob_rober, reltol = 1e-7, abstol = 1e-7)
 rosensol = solve(
     prob_rober, AutoVern7(Rodas5P(autodiff = false)), reltol = 1e-7, abstol = 1e-7)
-# test that default has the same performance as AutoTsit5(Rosenbrock23()) (which we expect it to use for this).
+# test that default has the same performance as AutoTsit5(Rodas5P()) (which we expect it to use for this).
 @test sol.stats.naccept == rosensol.stats.naccept
 @test sol.stats.nf == rosensol.stats.nf
 @test unique(sol.alg_choice) == [2, 4]
@@ -75,7 +75,7 @@ for n in (100, 600)
         vcat([1.0, 0.0, 0.0], ones(n)), (0.0, 100.0), (0.04, 3e7, 1e4))
     global sol = solve(prob_ex_rober)
     fsol = solve(prob_ex_rober, AutoTsit5(FBDF(; autodiff = false, linsolve)))
-    # test that default has the same performance as AutoTsit5(Rosenbrock23()) (which we expect it to use for this).
+    # test that default has the same performance as AutoTsit5(FBDF()) (which we expect it to use for this).
     @test sol.stats.naccept == fsol.stats.naccept
     @test sol.stats.nf == fsol.stats.nf
     @test unique(sol.alg_choice) == [1, stiffalg]

lib/OrdinaryDiffEqDefault/test/default_solver_tests.jl

@@ -27,7 +27,7 @@ import StaticArrays: SArray, MVector, SVector, @SVector, StaticArray, MMatrix, S
 using DiffEqBase: TimeGradientWrapper,
                   UJacobianWrapper, TimeDerivativeWrapper,
                   UDerivativeWrapper
-using SciMLBase: AbstractSciMLOperator
+using SciMLBase: AbstractSciMLOperator, DEIntegrator
 import OrdinaryDiffEqCore
 using OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, OrdinaryDiffEqAdaptiveImplicitAlgorithm,
                           DAEAlgorithm,

lib/OrdinaryDiffEqDifferentiation/src/OrdinaryDiffEqDifferentiation.jl

@@ -883,7 +883,7 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
         # Thus setup JacVec and a concrete J, using sparsity when possible
         _f = islin ? (isode ? f.f : f.f1.f) : f
         J = if f.jac_prototype === nothing
-            ArrayInterface.undefmatrix(u)
+            ArrayInterface.zeromatrix(u)
         else
             deepcopy(f.jac_prototype)
         end
@@ -907,7 +907,7 @@ function build_J_W(alg, u, uprev, p, t, dt, f::F, ::Type{uEltypeNoUnits},
                 f.jac(uprev, p, t)
             end
         elseif f.jac_prototype === nothing
-            ArrayInterface.undefmatrix(u)
+            ArrayInterface.zeromatrix(u)
         else
             deepcopy(f.jac_prototype)
         end
@@ -1003,4 +1003,4 @@ function resize_J_W!(cache, integrator, i)
     end
 
     nothing
-end
+end

lib/OrdinaryDiffEqDifferentiation/src/derivative_utils.jl

@@ -3,27 +3,20 @@ 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
+    if !isnothing(A)
+        if integrator isa DEIntegrator
+            (;u, p, t) = integrator
+            du = hasproperty(integrator, :du) ? integrator.du : nothing
+            p = (du, u, p, t)
+            reinit!(linsolve; A, p)
+        else
+            reinit!(linsolve; A)
+        end
     end
 
     linres = solve!(linsolve; reltol)
@@ -44,16 +37,15 @@ function dolinsolve(integrator, linsolve; A = nothing, linu = nothing, b = nothi
     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
+function wrapprecs(linsolver, W, weight)
+    if hasproperty(linsolver, :precs) && isnothing(linsolver.precs)
+        Pl = LinearSolve.InvPreconditioner(Diagonal(_vec(weight)))
+        Pr = Diagonal(_vec(weight))
+        precs = Returns((Pl, Pr))
+        return remake(linsolver; precs)
+    else
+        return linsolver
+    end
 end
 
 Base.resize!(p::LinearSolve.LinearCache, i) = p

lib/OrdinaryDiffEqDifferentiation/src/linsolve_utils.jl

@@ -13,7 +13,7 @@ import OrdinaryDiffEqCore: alg_order, alg_maximum_order, get_current_adaptive_or
                            OrdinaryDiffEqAdaptiveAlgorithm,
                            OrdinaryDiffEqAdaptiveImplicitAlgorithm,
                            alg_cache, CompiledFloats, @threaded, stepsize_controller!,
-                           DEFAULT_PRECS, full_cache,
+                           full_cache,
                            constvalue, PolyesterThreads, Sequential, BaseThreads,
                            _digest_beta1_beta2, timedepentdtmin, _unwrap_val,
                            _reshape, _vec, get_fsalfirstlast, generic_solver_docstring,