From 0c5ade6904608c1c6a593b7fba6de1ecd797ee14 Mon Sep 17 00:00:00 2001 From: Kristoffer Carlsson Date: Tue, 30 Jun 2026 09:45:19 +0200 Subject: [PATCH] cache `parameters`, `namespacemap` and `reorder_parameters` --- .../src/systems/abstractsystem.jl | 44 +++++++++++----- .../src/systems/codegen.jl | 51 ++++++++++++++----- .../src/systems/index_cache.jl | 33 ++++++++++-- 3 files changed, 100 insertions(+), 28 deletions(-) diff --git a/lib/ModelingToolkitBase/src/systems/abstractsystem.jl b/lib/ModelingToolkitBase/src/systems/abstractsystem.jl index 1e170c03d0..9a66492f88 100644 --- a/lib/ModelingToolkitBase/src/systems/abstractsystem.jl +++ b/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 diff --git a/lib/ModelingToolkitBase/src/systems/codegen.jl b/lib/ModelingToolkitBase/src/systems/codegen.jl index f3aee7f6f6..8707edbd62 100644 --- a/lib/ModelingToolkitBase/src/systems/codegen.jl +++ b/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)) diff --git a/lib/ModelingToolkitBase/src/systems/index_cache.jl b/lib/ModelingToolkitBase/src/systems/index_cache.jl index ed0dcc48a6..d2eb9c8c66 100644 --- a/lib/ModelingToolkitBase/src/systems/index_cache.jl +++ b/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