Expose sparse matrix generation in higher level prob constructors

Author: ChrisRackauckas

src/systems/diffeqs/abstractodesystem.jl

@@ -39,8 +39,11 @@ function generate_tgrad(sys::AbstractODESystem, dvs = states(sys), ps = paramete
     return build_function(tgrad, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
 end
 
-function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; sparse = false, kwargs...)
     jac = calculate_jacobian(sys)
+    if sparse
+        jac = sparse(jac)
+    end
     return build_function(jac, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
 end
 
@@ -121,7 +124,8 @@ respectively.
 function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      ps = parameters(sys);
                                      version = nothing, tgrad=false,
-                                     jac = false, Wfact = false, kwargs...) where {iip}
+                                     jac = false, Wfact = false,
+                                     kwargs...) where {iip}
     f_oop,f_iip = generate_function(sys, dvs, ps, Val{false}; kwargs...)
 
     f(u,p,t) = f_oop(u,p,t)
@@ -173,11 +177,12 @@ function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     parammap=DiffEqBase.NullParameters();
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
-                                    checkbounds = false,
+                                    checkbounds = false, sparse = false,
                                     linenumbers = true, multithread=false,
                                     kwargs...) where iip
     f = ODEFunction(sys;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
-                        linenumbers=linenumbers,multithread=multithread)
+                        linenumbers=linenumbers,multithread=multithread,
+                        sparse=sparse)
     u0 = varmap_to_vars(u0map,states(sys))
     p = varmap_to_vars(parammap,parameters(sys))
     ODEProblem(f,u0,tspan,p;kwargs...)

src/systems/diffeqs/sdesystem.jl

@@ -67,33 +67,33 @@ respectively.
 """
 function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.states, ps = sys.ps;
                                      version = nothing, tgrad=false,
-                                     jac = false, Wfact = false) where {iip}
-    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
-    g_oop,g_iip = generate_diffusion_function(sys, dvs, ps, Val{false})
+                                     jac = false, Wfact = false, kwargs...) where {iip}
+    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false}; kwargs...)
+    g_oop,g_iip = generate_diffusion_function(sys, dvs, ps, Val{false}; kwargs...)
 
     f(u,p,t) = f_oop(u,p,t)
     f(du,u,p,t) = f_iip(du,u,p,t)
     g(u,p,t) = g_oop(u,p,t)
     g(du,u,p,t) = g_iip(du,u,p,t)
 
     if tgrad
-        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false})
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false}; kwargs...)
         _tgrad(u,p,t) = tgrad_oop(u,p,t)
         _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
     else
         _tgrad = nothing
     end
 
     if jac
-        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false}; kwargs...)
         _jac(u,p,t) = jac_oop(u,p,t)
         _jac(J,u,p,t) = jac_iip(J,u,p,t)
     else
         _jac = nothing
     end
 
     if Wfact
-        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false})
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false}; kwargs...)
         Wfact_oop, Wfact_iip = tmp_Wfact
         Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
         _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)
@@ -125,12 +125,13 @@ end
 function DiffEqBase.SDEProblem{iip}(sys::SDESystem,u0map,tspan,p=parammap;
                                     version = nothing, tgrad=false,
                                     jac = false, Wfact = false,
-                                    checkbounds = false,
+                                    checkbounds = false, sparse = false,
                                     linenumbers = true, multithread=false,
                                     kwargs...) where iip
 
     f = SDEFunction(sys;tgrad=tgrad,jac=jac,Wfact=Wfact,checkbounds=checkbounds,
-                        linenumbers=linenumbers,multithread=multithread)
+                        linenumbers=linenumbers,multithread=multithread,
+                        sparse=false)
     u0 = varmap_to_vars(u0map,states(sys))
     p = varmap_to_vars(parammap,parameters(sys))
     SDEProblem(f,f.g,u0,tspan,p;kwargs...)

src/systems/nonlinear/nonlinearsystem.jl

@@ -47,8 +47,12 @@ function calculate_jacobian(sys::NonlinearSystem)
     return jac
 end
 
-function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys), expression = Val{true};
+                           sparse = false, kwargs...)
     jac = calculate_jacobian(sys)
+    if sparse
+        jac = sparse(jac)
+    end
     return build_function(jac, convert.(Variable,vs), convert.(Variable,ps), (), NLSysToExpr(sys))
 end
 
@@ -74,15 +78,15 @@ end
 
 function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,tspan,
                                           parammap=DiffEqBase.NullParameters();
-                                          jac = false,
+                                          jac = false, sparse=false,
                                           checkbounds = false,
                                           linenumbers = true, multithread=false,
                                           kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
 
     f = generate_function(sys;checkbounds=checkbounds,linenumbers=linenumbers,
-                              multithread=multithread)
+                              multithread=multithread,sparse=sparse)
     u0 = varmap_to_vars(u0map,dvs)
     p = varmap_to_vars(parammap,ps)
     NonlinearProblem(f,u0,tspan,p;kwargs...)

src/systems/optimization/optimizationsystem.jl

@@ -52,8 +52,12 @@ function calculate_hessian(sys::OptimizationSystem)
     expand_derivatives.(hessian(equations(sys), [dv() for dv in states(sys)]))
 end
 
-function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+function generate_hessian(sys::OptimizationSystem, vs = states(sys), ps = parameters(sys), expression = Val{true};
+                          sparse = false, kwargs...)
     hes = calculate_hessian(sys)
+    if sparse
+        hes = sparse(hes)
+    end
     return build_function(hes, convert.(Variable,vs), convert.(Variable,ps), (), x->convert(Expr, x); kwargs...)
 end
 
@@ -69,15 +73,15 @@ namespace_operation(sys::OptimizationSystem) = namespace_operation(sys.op,sys.na
 function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem,
                                           parammap=DiffEqBase.NullParameters();
                                           u0=nothing, lb=nothing, ub=nothing,
-                                          hes = false,
+                                          hes = false, sparse = false,
                                           checkbounds = false,
                                           linenumbers = true, multithread=false,
                                           kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
 
     f = generate_function(sys,checkbounds=checkbounds,linenumbers=linenumbers,
-                              multithread=multithread)
+                              multithread=multithread,sparse=sparse)
     u0 = varmap_to_vars(u0,dvs)
     p = varmap_to_vars(parammap,ps)
     lb = varmap_to_vars(lb,dvs)