Invalidate jac and friends' caches when necessary

Author: YingboMa

src/structural_transformation/pantelides.jl

@@ -149,5 +149,5 @@ function dae_index_lowering(sys::ODESystem; kwargs...)
     s = get_structure(sys)
     (s isa SystemStructure) || (sys = initialize_system_structure(sys))
     sys, var_eq_matching, eq_to_diff = pantelides!(sys; kwargs...)
-    return pantelides_reassemble(sys, eq_to_diff, var_eq_matching)
+    return invalidate_cache!(pantelides_reassemble(sys, eq_to_diff, var_eq_matching))
 end

src/structural_transformation/symbolics_tearing.jl

@@ -150,7 +150,7 @@ function tearing(sys; simplify=false)
     sys = init_for_tearing(sys)
     var_eq_matching = tear_graph(sys)
 
-    tearing_reassemble(sys, var_eq_matching; simplify=simplify)
+    invalidate_cache!(tearing_reassemble(sys, var_eq_matching; simplify=simplify))
 end
 
 function init_for_tearing(sys)

src/systems/abstractsystem.jl

@@ -232,6 +232,23 @@ for prop in [
     end
 end
 
+const EMPTY_TGRAD = Vector{Num}(undef, 0)
+const EMPTY_JAC = Matrix{Num}(undef, 0, 0)
+function invalidate_cache!(sys::AbstractSystem)
+    if isdefined(sys, :tgrad)
+        sys.tgrad[] = EMPTY_TGRAD
+    elseif isdefined(sys, :jac)
+        sys.jac[] = EMPTY_JAC
+    elseif isdefined(sys, :ctrl_jac)
+        sys.jac[] = EMPTY_JAC
+    elseif isdefined(sys, :Wfact)
+        sys.jac[] = EMPTY_JAC
+    elseif isdefined(sys, :Wfact_t)
+        sys.jac[] = EMPTY_JAC
+    end
+    return sys
+end
+
 Setfield.get(obj::AbstractSystem, ::Setfield.PropertyLens{field}) where {field} = getfield(obj, field)
 @generated function ConstructionBase.setproperties(obj::AbstractSystem, patch::NamedTuple)
     if issubset(fieldnames(patch), fieldnames(obj))

src/systems/alias_elimination.jl

@@ -68,7 +68,7 @@ function alias_elimination(sys)
     @set! sys.states = newstates
     @set! sys.observed = [observed(sys); [lhs ~ rhs for (lhs, rhs) in pairs(subs)]]
     @set! sys.structure = nothing
-    return sys
+    return invalidate_cache!(sys)
 end
 
 """

src/systems/diffeqs/odesystem.jl

@@ -147,11 +147,11 @@ function ODESystem(
     process_variables!(var_to_name, defaults, ps′)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
-    tgrad = RefValue(Vector{Num}(undef, 0))
-    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
-    ctrl_jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
-    Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
-    Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
+    tgrad = RefValue(EMPTY_TGRAD)
+    jac = RefValue{Any}(EMPTY_JAC)
+    ctrl_jac = RefValue{Any}(EMPTY_JAC)
+    Wfact   = RefValue(EMPTY_JAC)
+    Wfact_t = RefValue(EMPTY_JAC)
     sysnames = nameof.(systems)
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))

src/systems/diffeqs/sdesystem.jl

@@ -130,11 +130,11 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     process_variables!(var_to_name, defaults, ps′)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
-    tgrad = RefValue(Vector{Num}(undef, 0))
-    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
-    ctrl_jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
-    Wfact   = RefValue(Matrix{Num}(undef, 0, 0))
-    Wfact_t = RefValue(Matrix{Num}(undef, 0, 0))
+    tgrad = RefValue(EMPTY_TGRAD)
+    jac = RefValue{Any}(EMPTY_JAC)
+    ctrl_jac = RefValue{Any}(EMPTY_JAC)
+    Wfact   = RefValue(EMPTY_JAC)
+    Wfact_t = RefValue(EMPTY_JAC)
     SDESystem(deqs, neqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type, checks = checks)
 end
 

src/systems/nonlinear/nonlinearsystem.jl

@@ -94,7 +94,7 @@ function NonlinearSystem(eqs, states, ps;
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))
+    jac = RefValue{Any}(EMPTY_JAC)
     defaults = todict(defaults)
     defaults = Dict{Any,Any}(value(k) => value(v) for (k, v) in pairs(defaults))