Jacobian sparsity for ODEProblem

Author: YingboMa

src/structural_transformation/StructuralTransformations.jl

@@ -16,11 +16,12 @@ using ModelingToolkit: ODESystem, AbstractSystem,var_from_nested_derivative, Dif
                        isdiffeq, isdifferential, isinput,
                        empty_substitutions, get_substitutions,
                        get_structure, get_iv, independent_variables,
-                       get_structure, defaults, InvalidSystemException,
+                       has_structure, defaults, InvalidSystemException,
                        ExtraEquationsSystemException,
                        ExtraVariablesSystemException,
                        get_postprocess_fbody, vars!,
-                       IncrementalCycleTracker, add_edge_checked!, topological_sort
+                       IncrementalCycleTracker, add_edge_checked!, topological_sort,
+                       invalidate_cache!
 
 using ModelingToolkit.BipartiteGraphs
 import .BipartiteGraphs: invview
@@ -41,6 +42,7 @@ export tearing, dae_index_lowering, check_consistency
 export tearing_assignments, tearing_substitution
 export build_torn_function, build_observed_function, ODAEProblem
 export sorted_incidence_matrix
+export torn_system_jacobian_sparsity
 
 include("utils.jl")
 include("pantelides.jl")

src/structural_transformation/codegen.jl

@@ -4,7 +4,7 @@ using ModelingToolkit: isdifferenceeq, has_continuous_events, generate_rootfindi
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 
-function torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs)
+function torn_system_with_nlsolve_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs)
     s = structure(sys)
     @unpack fullvars, graph = s
 
@@ -329,7 +329,7 @@ function build_torn_function(
 
         ODEFunction{true}(
                           @RuntimeGeneratedFunction(expr),
-                          sparsity = jacobian_sparsity ? torn_system_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs) : nothing,
+                          sparsity = jacobian_sparsity ? torn_system_with_nlsolve_jacobian_sparsity(sys, var_eq_matching, var_sccs, nlsolve_scc_idxs, eqs_idxs, states_idxs) : nothing,
                           syms = syms,
                           observed = observedfun,
                           mass_matrix = mass_matrix,

src/structural_transformation/utils.jl

@@ -234,6 +234,25 @@ function uneven_invmap(n::Int, list)
     return rename
 end
 
+function torn_system_jacobian_sparsity(sys)
+    has_structure(sys) || return nothing
+    s = structure(sys)
+    @unpack fullvars, graph = s
+
+    states_idxs = findall(!isdifferential, fullvars)
+    var2idx = Dict{Int,Int}(v => i for (i, v) in enumerate(states_idxs))
+    I = Int[]; J = Int[]
+    for ieq in 𝑠vertices(graph)
+        for ivar in 𝑠neighbors(graph, ieq)
+            nivar = get(var2idx, ivar, 0)
+            nivar == 0 && continue
+            push!(I, ieq)
+            push!(J, nivar)
+        end
+    end
+    return sparse(I, J, true)
+end
+
 ###
 ### Nonlinear equation(s) solving
 ###

src/systems/diffeqs/abstractodesystem.jl

@@ -281,9 +281,13 @@ function calculate_massmatrix(sys::AbstractODESystem; simplify=false)
     M == I ? I : M
 end
 
-jacobian_sparsity(sys::AbstractODESystem) =
+function jacobian_sparsity(sys::AbstractODESystem)
+    sparsity = torn_system_jacobian_sparsity(sys)
+    sparsity === nothing || return sparsity
+
     jacobian_sparsity([eq.rhs for eq ∈ equations(sys)],
                       [dv for dv in states(sys)])
+end
 
 function isautonomous(sys::AbstractODESystem)
     tgrad = calculate_tgrad(sys;simplify=true)
@@ -319,6 +323,7 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      eval_module = @__MODULE__,
                                      steady_state = false,
                                      checkbounds=false,
+                                     sparsity=false,
                                      kwargs...) where {iip}
 
     f_gen = generate_function(sys, dvs, ps; expression=Val{eval_expression}, expression_module=eval_module, checkbounds=checkbounds, kwargs...)
@@ -394,6 +399,7 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                      syms = Symbol.(states(sys)),
                      indepsym = Symbol(get_iv(sys)),
                      observed = observedfun,
+                     sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
                     )
 end