Merge pull request #4547 from SciML/as/cache-linear-expander

 feat: allow caching `LinearExpander`s in `System`

Author: AayushSabharwal

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -998,7 +998,12 @@ Invalidate cached jacobians, etc.
 """
 function invalidate_cache!(sys::AbstractSystem)
     has_metadata(sys) || return sys
-    empty!(getmetadata(sys, MutableCacheKey, nothing))
+    cache = getmetadata(sys, MutableCacheKey, nothing)
+    if cache isa MutableCacheT
+        # Avoid clearing the linear expansion cache. It doesn't depend on anything in the
+        # system, just useful to have around.
+        filter!(Base.Fix2(===, LinearExpansionCache) ∘ first, cache)
+    end
     return sys
 end
 

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -1450,7 +1450,7 @@ function calculate_A_b(sys::System; sparse = false, throw = true)
     # `linear_expansion` caches values based on `var`. This loop ordering helps
     # avoid invalidating the cache frequently.
     for (j, var) in enumerate(dvs)
-        lex = Symbolics.LinearExpander(var; strict = true)
+        lex = get_linear_expander_for!(sys, var, true)
         for (i, resid) in enumerate(rhss)
             p, q, islinear = lex(resid)
             # An equation such as `0 ~ 1 + x * y` will return `(x, 1, true)` for `y`.

lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl

@@ -199,7 +199,12 @@ function generate_initializesystem_timevarying(
         kwargs...
     )
     diffcache_params = SU.getmetadata(sys, DiffCacheParams, Dict{SymbolicT, Int}())::Dict{SymbolicT, Int}
-    isys = SU.setmetadata(isys, DiffCacheParams, diffcache_params)
+    isys = SU.setmetadata(isys, DiffCacheParams, copy(diffcache_params))
+    # Reuse `LinearExpander` cache for the initialization system
+    linear_expansion_cache = check_mutable_cache(sys, LinearExpansionCache, LinearExpansionCacheT, nothing)
+    if linear_expansion_cache !== nothing
+        store_to_mutable_cache!(isys, LinearExpansionCache, linear_expansion_cache)
+    end
     return isys
 end
 

lib/ModelingToolkitBase/src/systems/system.jl

@@ -1627,3 +1627,29 @@ function should_invalidate_mutable_cache_entry(
     )
     return true
 end
+
+abstract type LinearExpansionCache end
+
+const LinearExpansionCacheT = Dict{Tuple{SymbolicT, Bool}, Symbolics.LinearExpander}
+
+function should_invalidate_mutable_cache_entry(
+        ::Type{LinearExpansionCache}, @nospecialize(patch::NamedTuple)
+    )
+    return false
+end
+
+"""
+    $TYPEDSIGNATURES
+
+Get a cached `Symbolics.LinearExpander` for variable `var` and strictness `strict`. This
+allows reusing significant cached information between `LinearExpander` calls. This cache
+is never invalidated.
+"""
+function get_linear_expander_for!(sys::System, var::SymbolicT, strict::Bool)
+    cache = check_mutable_cache(sys, LinearExpansionCache, LinearExpansionCacheT, nothing)
+    if cache === nothing
+        cache = LinearExpansionCacheT()
+        store_to_mutable_cache!(sys, LinearExpansionCache, cache)
+    end
+    return get!(() -> Symbolics.LinearExpander(var; strict), cache, (var, strict))
+end