Merge pull request #4685 from KristofferC/kc/cache_maps

cache `parameters`, `namespacemap` and `reorder_parameters`

Author: AayushSabharwal

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

@@ -1613,27 +1613,47 @@ function __no_initial_params_pred(x::SymbolicT)
     end
 end
 
-"""
-$(TYPEDSIGNATURES)
+struct CachedSystemParameters
+    value::Vector{SymbolicT}
+end
 
-Get the parameters of the system `sys` and its subsystems.
+function should_invalidate_mutable_cache_entry(::Type{CachedSystemParameters}, patch::NamedTuple)
+    return haskey(patch, :ps) || haskey(patch, :systems)
+end
 
-See also [`@parameters`](@ref) and [`ModelingToolkitBase.get_ps`](@ref).
-"""
-function parameters(sys::AbstractSystem; initial_parameters = false)
+function _unfiltered_parameters(sys::AbstractSystem)
     ps = get_ps(sys)
-    if ps === SciMLBase.NullParameters()
-        return SymbolicT[]
-    end
+    ps === SciMLBase.NullParameters() && return SymbolicT[]
     if eltype(ps) <: Pair
         ps = Vector{SymbolicT}(unwrap.(first.(ps)))
     end
-    systems = get_systems(sys)
     result = OrderedSet{SymbolicT}(ps)
-    for subsys in systems
+    for subsys in get_systems(sys)
         union!(result, namespace_parameters(subsys))
     end
-    result = collect(result)
+    return collect(result)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Get the parameters of the system `sys` and its subsystems.
+
+See also [`@parameters`](@ref) and [`ModelingToolkitBase.get_ps`](@ref).
+"""
+function parameters(sys::AbstractSystem; initial_parameters = false)
+    if sys isa System && iscomplete(sys)
+        cached = check_mutable_cache(sys, CachedSystemParameters, CachedSystemParameters, nothing)
+        if cached isa CachedSystemParameters
+            result = copy(cached.value)
+        else
+            base = _unfiltered_parameters(sys)
+            store_to_mutable_cache!(sys, CachedSystemParameters, CachedSystemParameters(base))
+            result = copy(base)
+        end
+    else
+        result = _unfiltered_parameters(sys)
+    end
     if !initial_parameters && !is_initializesystem(sys)
         filter!(__no_initial_params_pred, result)
     end

lib/ModelingToolkitBase/src/systems/codegen.jl

@@ -1178,6 +1178,43 @@ function (pred::CheckInvalidAndTrackNamespaced)(x::SymbolicT)
     return false
 end
 
+struct NamespaceMap
+    value::Dict{SymbolicT, SymbolicT}
+end
+
+function should_invalidate_mutable_cache_entry(::Type{NamespaceMap}, patch::NamedTuple)
+    return haskey(patch, :name) || haskey(patch, :observed) || haskey(patch, :unknowns) ||
+           haskey(patch, :ps) || haskey(patch, :systems)
+end
+
+function _build_namespace_map(sys)
+    ns_map = Dict{SymbolicT, SymbolicT}(renamespace(sys, eq.lhs) => eq.lhs for eq in observed(sys))
+    for sym in unknowns(sys)
+        ns_map[renamespace(sys, sym)] = sym
+        if iscall(sym) && operation(sym) === getindex
+            ns_map[renamespace(sys, arguments(sym)[1])] = arguments(sym)[1]
+        end
+    end
+    for sym in parameters(sys; initial_parameters = true)
+        ns_map[renamespace(sys, sym)] = sym
+        if iscall(sym) && operation(sym) === getindex
+            ns_map[renamespace(sys, arguments(sym)[1])] = arguments(sym)[1]
+        end
+    end
+    return ns_map
+end
+
+function get_namespace_map(sys)
+    if sys isa System && iscomplete(sys)
+        cached = check_mutable_cache(sys, NamespaceMap, NamespaceMap, nothing)
+        cached isa NamespaceMap && return cached.value
+        m = _build_namespace_map(sys)
+        store_to_mutable_cache!(sys, NamespaceMap, NamespaceMap(m))
+        return m
+    end
+    return _build_namespace_map(sys)
+end
+
 """
     build_explicit_observed_function(sys, ts; kwargs...) -> Function(s)
 
@@ -1283,19 +1320,7 @@ Base.@nospecializeinfer function build_explicit_observed_function(
     end
 
     namespace_subs = Dict{SymbolicT, SymbolicT}()
-    ns_map = Dict{SymbolicT, SymbolicT}(renamespace(sys, eq.lhs) => eq.lhs for eq in observed(sys))
-    for sym in unknowns(sys)
-        ns_map[renamespace(sys, sym)] = sym
-        if iscall(sym) && operation(sym) === getindex
-            ns_map[renamespace(sys, arguments(sym)[1])] = arguments(sym)[1]
-        end
-    end
-    for sym in parameters(sys; initial_parameters = true)
-        ns_map[renamespace(sys, sym)] = sym
-        if iscall(sym) && operation(sym) === getindex
-            ns_map[renamespace(sys, arguments(sym)[1])] = arguments(sym)[1]
-        end
-    end
+    ns_map = get_namespace_map(sys)
     allsyms = as_atomic_array_set(unknowns(sys))
     foreach(Base.Fix1(push_as_atomic_array!, allsyms), observables(sys))
     foreach(Base.Fix1(push_as_atomic_array!, allsyms), parameters(sys; initial_parameters = true))

lib/ModelingToolkitBase/src/systems/index_cache.jl

@@ -588,9 +588,36 @@ end
 
 const ReorderedParametersT = Vector{Union{Vector{SymbolicT}, Vector{Vector{SymbolicT}}}}
 
-function reorder_parameters(
-        sys::AbstractSystem, ps = parameters(sys; initial_parameters = true); kwargs...
-    )
+struct ReorderedDefaultParameters
+    value::ReorderedParametersT
+end
+
+function should_invalidate_mutable_cache_entry(::Type{ReorderedDefaultParameters}, patch::NamedTuple)
+    return haskey(patch, :ps) || haskey(patch, :index_cache)
+end
+
+function _copy_reordered(v::ReorderedParametersT)
+    return ReorderedParametersT(map(v) do inner
+        inner isa Vector{SymbolicT} ? copy(inner) : map(copy, inner)
+    end)
+end
+
+function reorder_parameters(sys::AbstractSystem; kwargs...)
+    if isempty(kwargs) && sys isa System && has_index_cache(sys) && get_index_cache(sys) !== nothing
+        cached = check_mutable_cache(sys, ReorderedDefaultParameters, ReorderedDefaultParameters, nothing)
+        if cached isa ReorderedDefaultParameters
+            return _copy_reordered(cached.value)
+        end
+        val = reorder_parameters(
+            get_index_cache(sys)::IndexCache, parameters(sys; initial_parameters = true)
+        )
+        store_to_mutable_cache!(sys, ReorderedDefaultParameters, ReorderedDefaultParameters(val))
+        return _copy_reordered(val)
+    end
+    return reorder_parameters(sys, parameters(sys; initial_parameters = true); kwargs...)
+end
+
+function reorder_parameters(sys::AbstractSystem, ps; kwargs...)
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
         return reorder_parameters(get_index_cache(sys)::IndexCache, ps; kwargs...)
     elseif ps isa Tuple