refactor: use cached LinearExpanders

Author: AayushSabharwal

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -1446,7 +1446,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