Merge pull request #185 from ReactionMechanismGenerator/precondition

Use preconditioner if sparsity is larger than an empirical threshold

Author: mjohnson541

Project.toml

@@ -13,6 +13,7 @@ DiffEqSensitivity = "41bf760c-e81c-5289-8e54-58b1f1f8abe2"
 FastGaussQuadrature = "442a2c76-b920-505d-bb47-c5924d526838"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
+IncompleteLU = "40713840-3770-5561-ab4c-a76e7d0d7895"
 IterTools = "c8e1da08-722c-5040-9ed9-7db0dc04731e"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
@@ -46,6 +47,7 @@ DiffEqBase = "6"
 DiffEqSensitivity = "6"
 ForwardDiff = "0.10"
 Images = "0.23"
+IncompleteLU = "0.2.0"
 IterTools = "1.3.0"
 LsqFit = "0.12"
 ModelingToolkit = "3"

src/Reactor.jl

@@ -3,38 +3,68 @@ using DiffEqBase
 using ForwardDiff
 using Sundials
 using ModelingToolkit
+using IncompleteLU
+using LinearAlgebra
+using SparseArrays
 abstract type AbstractReactor end
 export AbstractReactor
 
-struct Reactor{D,Q} <: AbstractReactor
+struct Reactor{D,Q,F1,F2,F3} <: AbstractReactor
     domain::D
     ode::ODEProblem
     recommendedsolver::Q
     forwardsensitivities::Bool
+    precsundials::F1 #function to calculate preconditioner for Sundials solvers
+    psetupsundials::F2 #function to compute preconditioner \ residue for Sundials solvers
+    precsjulia::F3 #function to calculate preconditioner for Julia solvers
 end
 
-function Reactor(domain::T,y0::Array{W,1},tspan::Tuple,interfaces::Z=[];p::X=DiffEqBase.NullParameters(),forwardsensitivities=false,forwarddiff=false,modelingtoolkit=false) where {T<:AbstractDomain,W<:Real,Z<:AbstractArray,X}
+function Reactor(domain::T,y0::Array{T1,1},tspan::Tuple,interfaces::Z=[];p::X=DiffEqBase.NullParameters(),forwardsensitivities=false,forwarddiff=false,modelingtoolkit=false,tau=1e-3) where {T<:AbstractDomain,T1<:Real,Z<:AbstractArray,X}
     dydt(dy::X,y::T,p::V,t::Q) where {X,T,Q,V} = dydtreactor!(dy,y,t,domain,interfaces,p=p)
     jacy!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobiany!(J,y,p,t,domain,interfaces,nothing)
     jacyforwarddiff!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianyforwarddiff!(J,y,p,t,domain,interfaces,nothing)
     jacp!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianp!(J,y,p,t,domain,interfaces,nothing)
     jacpforwarddiff!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianpforwarddiff!(J,y,p,t,domain,interfaces,nothing)
 
+    psetupsundials(p::T1, t::T2, u::T3, du::T4, jok::Bool, jcurPtr::T5, gamma::T6) where {T1,T2,T3,T4,T5,T6} = _psetup(p, t, u, du, jok, jcurPtr, gamma, jacy!, W::SparseMatrixCSC{Float64, Int64}, preccache::Base.RefValue{IncompleteLU.ILUFactorization{Float64, Int64}}, tau::Float64)
+    precsundials(z::T1, r::T2, p::T3, t::T4, y::T5, fy::T6, gamma::T7, delta::T8, lr::T9) where {T1,T2,T3,T4,T5,T6,T7,T8,T9} = _prec(z, r, p, t, y, fy, gamma, delta, lr, preccache)
+    precsjulia(W::T1,du::T2,u::T3,p::T4,t::T5,newW::T6,Plprev::T7,Prprev::T8,solverdata::T9) where {T1,T2,T3,T4,T5,T6,T7,T8,T9} =  _precsjulia(W,du,u,p,t,newW,Plprev,Prprev,solverdata,tau)
+    
+    # determine worst sparsity
+    y0length = length(y0)
+    J = spzeros(y0length,y0length)
+    jacyforwarddiff!(J,NaN*ones(y0length),p,0.0)
+    @. J.nzval = 1.0
+    sparsity = 1.0 - length(J.nzval)/(y0length*y0length)
+
+    # preconditioner caches for Sundials solver
+    W = spzeros(y0length,y0length)
+    jacyforwarddiff!(W,y0,p,0.0)
+    @. W.nzval = -1.0*W.nzval
+    idxs = diagind(W)
+    @inbounds @views @. W[idxs] = W[idxs] + 1
+    prectmp = ilu(W, τ = tau)
+    preccache = Ref(prectmp)
+    
     if (forwardsensitivities || !forwarddiff) && domain isa Union{ConstantTPDomain,ConstantVDomain,ConstantPDomain,ParametrizedTPDomain,ParametrizedVDomain,ParametrizedPDomain,ConstantTVDomain,ParametrizedTConstantVDomain,ConstantTAPhiDomain}
         if !forwardsensitivities
             odefcn = ODEFunction(dydt;jac=jacy!,paramjac=jacp!)
         else
             odefcn = ODEFunction(dydt;paramjac=jacp!)
         end
     else
-        odefcn = ODEFunction(dydt;jac=jacyforwarddiff!,paramjac=jacpforwarddiff!) 
+        odefcn = ODEFunction(dydt;jac=jacyforwarddiff!,paramjac=jacpforwarddiff!,jac_prototype=float.(J)) #jac_prototype is not needed/used for Sundials solvers but maybe needed for Julia solvers
     end
     if forwardsensitivities
         ode = ODEForwardSensitivityProblem(odefcn,y0,tspan,p)
         recsolver = Sundials.CVODE_BDF(linear_solver=:GMRES)
     else
         ode = ODEProblem(odefcn,y0,tspan,p)
-        recsolver  = Sundials.CVODE_BDF()
+        if sparsity > 0.8 #empirical threshold to use preconditioner
+            recsolver  = Sundials.CVODE_BDF(linear_solver=:GMRES,prec=precsundials,psetup=psetupsundials,prec_side=1)
+        else
+            recsolver  = Sundials.CVODE_BDF()
+        end
     end
     if modelingtoolkit
         sys = modelingtoolkitize(ode)
@@ -50,9 +80,9 @@ function Reactor(domain::T,y0::Array{W,1},tspan::Tuple,interfaces::Z=[];p::X=Dif
             ode = ODEProblem(odefcn,y0,tspan,p)
         end
     end
-    return Reactor(domain,ode,recsolver,forwardsensitivities)
+    return Reactor(domain,ode,recsolver,forwardsensitivities,precsundials,psetupsundials,precsjulia)
 end
-function Reactor(domains::T,y0s::W,tspan::W2,interfaces::Z=Tuple(),ps::X=DiffEqBase.NullParameters();forwardsensitivities=false,modelingtoolkit=false) where {T<:Tuple,W<:Tuple,Z,X,W2}
+function Reactor(domains::T,y0s::W1,tspan::W2,interfaces::Z=Tuple(),ps::X=DiffEqBase.NullParameters();forwardsensitivities=false,modelingtoolkit=false,tau=1e-3) where {T<:Tuple,W1<:Tuple,Z,X,W2}
     #adjust indexing
     y0 = zeros(sum(length(y) for y in y0s))
     Nvars = 0
@@ -137,6 +167,25 @@ function Reactor(domains::T,y0s::W,tspan::W2,interfaces::Z=Tuple(),ps::X=DiffEqB
     jacy!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianyforwarddiff!(J,y,p,t,domains,interfaces,nothing)
     jacp!(J::Q2,y::T,p::V,t::Q) where {Q2,T,Q,V} = jacobianpforwarddiff!(J,y,p,t,domains,interfaces,nothing)
 
+    psetupsundials(p::T1, t::T2, u::T3, du::T4, jok::Bool, jcurPtr::T5, gamma::T6) where {T1,T2,T3,T4,T5,T6} = _psetup(p, t, u, du, jok, jcurPtr, gamma, jacy!, W::SparseMatrixCSC{Float64, Int64}, preccache::Base.RefValue{IncompleteLU.ILUFactorization{Float64, Int64}}, tau::Float64)
+    precsundials(z::T1, r::T2, p::T3, t::T4, y::T5, fy::T6, gamma::T7, delta::T8, lr::T9) where {T1,T2,T3,T4,T5,T6,T7,T8,T9} = _prec(z, r, p, t, y, fy, gamma, delta, lr, preccache)
+    precsjulia(W::T1,du::T2,u::T3,p::T4,t::T5,newW::T6,Plprev::T7,Prprev::T8,solverdata::T9) where {T1,T2,T3,T4,T5,T6,T7,T8,T9} =  _precsjulia(W,du,u,p,t,newW,Plprev,Prprev,solverdata,tau)
+    
+    # determine worst sparsity
+    y0length = length(y0)
+    J = spzeros(y0length,y0length)
+    jacy!(J,NaN*ones(y0length),p,0.0)
+    @. J.nzval = 1.0
+    sparsity = 1.0 - length(J.nzval)/(y0length*y0length)
+
+    # preconditioner caches for Sundials solver
+    W = spzeros(y0length,y0length)
+    jacy!(W,y0,p,0.0)
+    @. W.nzval = -1.0*W.nzval
+    idxs = diagind(W)
+    @inbounds @views @. W[idxs] = W[idxs] + 1
+    prectmp = ilu(W, τ = tau)
+    preccache = Ref(prectmp)
 
     if forwardsensitivities
         odefcn = ODEFunction(dydt;paramjac=jacp!)
@@ -149,20 +198,91 @@ function Reactor(domains::T,y0s::W,tspan::W2,interfaces::Z=Tuple(),ps::X=DiffEqB
             ode = ODEForwardSensitivityProblem(odefcn,y0,tspan,p)
         end
     else
-        odefcn = ODEFunction(dydt;jac=jacy!,paramjac=jacp!)
+        odefcn = ODEFunction(dydt;jac=jacy!,paramjac=jacp!,jac_prototype=float.(J))
         ode = ODEProblem(odefcn,y0,tspan,p)
-        recsolver  = Sundials.CVODE_BDF()
+        if sparsity > 0.8 #empirical threshold to use preconditioner
+            recsolver  = Sundials.CVODE_BDF(linear_solver=:GMRES,prec=precsundials,psetup=psetupsundials,prec_side=1)
+        else
+            recsolver  = Sundials.CVODE_BDF()
+        end
         if modelingtoolkit
             sys = modelingtoolkitize(ode)
             jac = eval(ModelingToolkit.generate_jacobian(sys)[2])
             odefcn = ODEFunction(dydt;jac=jac,paramjac=jacp!)
             ode = ODEProblem(odefcn,y0,tspan,p)
         end
     end
-    return Reactor(domains,ode,recsolver,forwardsensitivities),y0,p
+    return Reactor(domains,ode,recsolver,forwardsensitivities,precsundials,psetupsundials,precsjulia),y0,p
 end
 export Reactor
 
+#preconditioner related functions
+@inline function _psetupsundials(p::T1, t::T2, u::T3, du::T4, jok::Bool, jcurPtr::T5, gamma::T6, jac!::T7, W::T8, preccache::T9, tau::T10) where {T1,T2,T3,T4,T5,T6,T7,T8,T9,T10}
+    """
+    Update preconditioner when Jacobian needs to be updated for Sundials solvers. Credit to tutorial of DifferentialEquations.jl.
+    p: the parameters
+    t: the current independent variable
+    u: the current state
+    du: the current f(u,p,t)
+    jok: a bool indicating whether the Jacobian needs to be updated
+    jcurPtr: a reference to an Int for whether the Jacobian was updated. jcurPtr[]=true should be set if the Jacobian was updated, and jcurPtr[]=false should be set if the Jacobian was not updated.
+    gamma: the gamma of W = M - gamma*J
+    """
+    if jok
+        @. W = 0.0
+        jac!(W,u,p,t)
+        jcurPtr[] = true
+
+        # W = I - gamma*J
+        @. W.nzval = -gamma*W.nzval
+        idxs = diagind(W)
+        @inbounds @views @. W[idxs] = W[idxs] + 1
+
+        # Build preconditioner on W
+        preccache[] = ilu(W, τ = tau)
+    end
+    nothing
+end
+@inline function _precsundials(z::T1, r::T2, p::T3, t::T4, y::T5, fy::T6, gamma::T7, delta::T8, lr::T9, preccache::T10) where {T1,T2,T3,T4,T5,T6,T7,T8,T9,T10}
+    """
+    Compute preccache \\ r in-place and store the result in z for Sundials solver. Credit to tutorial of DifferentialEquations.jl.
+    z: the computed output vector
+    r: the right-hand side vector of the linear system
+    p: the parameters
+    t: the current independent variable
+    du: the current value of f(u,p,t)
+    gamma: the gamma of W = M - gamma*J
+    delta: the iterative method tolerance
+    lr: a flag for whether lr=1 (left) or lr=2 (right) preconditioning
+    preccache: preconditioner cache
+    """
+    ldiv!(z,preccache[],r)
+end
+@inline function _precsjulia(W::T1,du::T2,u::T3,p::T4,t::T5,newW::T6,Plprev::T7,Prprev::T8,solverdata::T9,tau::T10) where {T1,T2,T3,T4,T5,T6,T7,T8,T9,T10}
+    """
+    Update preconditioner when Jacobian needs to be updated for Julia solvers. Credit to tutorial of DifferentialEquations.jl.
+    W: I - gamma*J or I/gamma - J depending on the algorithm.
+       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.
+    """
+    if newW === nothing || newW
+        Pl = ilu(convert(AbstractMatrix,W), τ = tau)
+    else
+        Pl = Plprev
+    end
+    Pl,nothing
+end
+
+
 @inline function getrate(rxn::T,cs::Array{W,1},kfs::Array{Q,1},krevs::Array{Q,1}) where {T<:AbstractReaction,Q,W<:Real}
     Nreact = length(rxn.reactantinds)
     Nprod = length(rxn.productinds)