diff --git a/Project.toml b/Project.toml index b4af29cf8c..8ce7150091 100644 --- a/Project.toml +++ b/Project.toml @@ -50,10 +50,6 @@ OrdinaryDiffEqBDF = "6ad6398a-0878-4a85-9266-38940aa047c8" OrdinaryDiffEqDefault = "50262376-6c5a-4cf5-baba-aaf4f84d72d7" OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" -[sources] -ModelingToolkitBase = {path = "lib/ModelingToolkitBase"} -SciCompDSL = {path = "lib/SciCompDSL"} - [extensions] MTKFMIExt = "FMIImport" MTKOrdinaryDiffEqBDFExt = "OrdinaryDiffEqBDF" diff --git a/lib/ModelingToolkitBase/src/systems/abstractsystem.jl b/lib/ModelingToolkitBase/src/systems/abstractsystem.jl index 67cc284c60..1e170c03d0 100644 --- a/lib/ModelingToolkitBase/src/systems/abstractsystem.jl +++ b/lib/ModelingToolkitBase/src/systems/abstractsystem.jl @@ -567,7 +567,7 @@ function add_initialization_parameters(sys::AbstractSystem; split = true, _unhac all_initialvars = Set{SymbolicT}() # time-independent systems don't initialize unknowns # but may initialize parameters using guesses for unknowns - _sys = _unhack_sys === nothing ? unhack_system(sys) : _unhack_sys + _sys = _unhack_sys === nothing ? reverse_all_default_reversible_transformations(sys) : _unhack_sys obs = observed(_sys) eqs = equations(_sys) for x in unknowns(_sys) @@ -698,7 +698,7 @@ function complete( end sys = newsys check_no_parameter_equations(sys) - _unhack_sys = unhack_system(sys) + _unhack_sys = reverse_all_default_reversible_transformations(sys) if add_initial_parameters sys = add_initialization_parameters(sys; split, _unhack_sys)::T end @@ -1022,6 +1022,11 @@ function refreshed_metadata(meta::Base.ImmutableDict, patch = nothing) newmeta = MetadataT() hascache = false for (k, v) in meta + # `Base.ImmutableDict` is like a stack. New key-value pairs (even with the same key) + # are inserted at the top of the stack. `getindex` returns the first matching key. + # We ignore repetitions of a key we've already encountered to avoid overwriting + # newer entries with lingering older ones. + haskey(newmeta, k) && continue if k === MutableCacheKey hascache = true new_cache = maybe_invalidate_cache((v::MutableCacheTLVT)[]::MutableCacheT, patch)::MutableCacheT diff --git a/lib/ModelingToolkitBase/src/systems/callbacks.jl b/lib/ModelingToolkitBase/src/systems/callbacks.jl index 62cc7635d5..5e78a74c00 100644 --- a/lib/ModelingToolkitBase/src/systems/callbacks.jl +++ b/lib/ModelingToolkitBase/src/systems/callbacks.jl @@ -193,7 +193,7 @@ function AffectSystem( union!(accessed_params, sys_params) # add scalarized unknowns to the map. - _obs = observed(unhack_system(affectsys)) + _obs = observed(reverse_all_default_reversible_transformations(affectsys)) _dvs = vcat(unknowns(affectsys), map(eq -> eq.lhs, _obs)) _dvs = __safe_scalarize_vars(_dvs) _discs = __safe_scalarize_vars(discretes) @@ -1415,7 +1415,7 @@ Base.@nospecializeinfer function compile_explicit_affect( ps_to_update = discretes(aff) dvs_to_update = setdiff(unknowns(aff), getfield.(observed(sys), :lhs)) - _affsys = unhack_system(affsys) + _affsys = reverse_all_default_reversible_transformations(affsys) aff_ir_info = get_ir_info(affsys) aff_ir = get_irstructure(affsys) obseqs = Equation[] diff --git a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl index 917aba67a4..46f5e8c316 100644 --- a/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl +++ b/lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl @@ -55,7 +55,7 @@ function generate_initializesystem_timevarying( check_units = true, check_defguess = false, name = nameof(sys), kwargs... ) - _sys = unhack_system(sys) + _sys = reverse_transformations_for_initialization(sys) trueobs = observed(_sys) eqs = equations(_sys) # remove any observed equations that directly or indirectly contain @@ -226,7 +226,7 @@ function generate_initializesystem_timeindependent( name = nameof(sys), kwargs... ) eqs = equations(sys) - _sys = unhack_system(sys) + _sys = reverse_transformations_for_initialization(sys) trueobs = observed(_sys) eqs = equations(_sys) # remove any observed equations that directly or indirectly contain @@ -885,18 +885,56 @@ function unhack_observed(obseqs, eqs) return obseqs, eqs end +""" + $TYPEDEF + +Default reversible transformation applied to all systems in `mtkcompile` for backward +compatibility. +""" +abstract type UnhackSystemTransformation <: ReversibleTransformations end + +function reverse_transformation(sys::AbstractSystem, ::Type{UnhackSystemTransformation}) + return unhack_system(sys) +end + +# NOTE: +# 1. Only this MTKBase is loaded. `__mtkcompile` here will remove `UnhackSystemTransformation` +# from the transformations to avoid it becoming an infinite loop. `unhack_system` also +# checks to make sure it doesn't rerun itself. +# +# 2. This MTKBase + old MTKTearing + old MTK: MTKTearing will anyway do its own array +# observed equations and remove them in `unhack_system`. The `UnhackSystemTransformation` +# added in `__mtkcompile` will run `unhack_system` in `reverse_all_default_reversible_transformations`. +# +# 3. This MTKBase + new MTKTearing + old MTK: MTKTearing doesn't implement `unhack_system` +# anymore and removes `UnhackSystemTransformation`. MTKTearing's own reversible transformations +# will run as intended. Calls to `unhack_system` from MTK will not early-exit. +# +# 4. This MTKBase + new MTKTearing + new MTK: `unhack_system` is never called and +# `UnhackSystemTransformation` is not present post-`mtkcompile`. + +function remove_unhack_system_transformation(sys::AbstractSystem) + return filter_reversible_transformations( + tf -> tf !== UnhackSystemTransformation, sys + )::System +end + """ unhack_system(sys::AbstractSystem) Given a system, remove any codegen oddities applied to it. This is typically used as a precursor to generating the initialization system. It is the successor of the now deprecated `unhack_observed`. + +DEPRECATED: See `with_reversible_transformation` and the reversible transformation API. """ function unhack_system(sys) - obs, eqs = unhack_observed(observed(sys), equations(sys)) - @set! sys.observed = obs - @set! sys.eqs = eqs - return sys + # See note above + tfs = getmetadata(sys, ReversibleTransformations, [])::Vector{Any} + for tf in tfs + tf === UnhackSystemTransformation && return sys + end + return reverse_all_default_reversible_transformations(sys) end function UnknownsInTimeIndependentInitializationError(eq, non_params) diff --git a/lib/ModelingToolkitBase/src/systems/problem_utils.jl b/lib/ModelingToolkitBase/src/systems/problem_utils.jl index 5c2dafe7cd..fcf2283279 100644 --- a/lib/ModelingToolkitBase/src/systems/problem_utils.jl +++ b/lib/ModelingToolkitBase/src/systems/problem_utils.jl @@ -1873,7 +1873,7 @@ function __process_SciMLProblem( add_initials!(sys, op) - _sys = unhack_system(sys) + _sys = reverse_all_default_reversible_transformations(sys) obs = observed(_sys) guesses = operating_point_preprocess(sys, guesses; name = "guesses") @@ -2367,7 +2367,7 @@ function get_u0(sys::AbstractSystem, varmap; kwargs...) op = build_operating_point(sys, varmap) binds = bindings(sys) no_override_merge_except_missing!(op, binds) - obs = observed(unhack_system(sys)) + obs = observed(reverse_all_default_reversible_transformations(sys)) add_observed_equations!(op, obs) return varmap_to_vars(op, unknowns(sys); kwargs...) @@ -2385,7 +2385,7 @@ function get_p(sys::AbstractSystem, varmap; split = is_split(sys), kwargs...) binds = bindings(sys) no_override_merge_except_missing!(op, binds) add_initials!(sys, op) - obs = observed(unhack_system(sys)) + obs = observed(reverse_all_default_reversible_transformations(sys)) add_observed_equations!(op, obs) return if split diff --git a/lib/ModelingToolkitBase/src/systems/system.jl b/lib/ModelingToolkitBase/src/systems/system.jl index b2bccb0cf7..13b15fd6ea 100644 --- a/lib/ModelingToolkitBase/src/systems/system.jl +++ b/lib/ModelingToolkitBase/src/systems/system.jl @@ -1688,3 +1688,103 @@ Get the `IRStructure` associated with the system. function get_irstructure(sys::System) return get_irstructure_tlv(sys)[]::IRStructure{SymReal} end + +# Reversible transformation API + +abstract type ReversibleTransformations end + +""" + $TYPEDSIGNATURES + +Add `tf` to the list of reversible transformations applied to `sys`. Returns a modified +copy of `sys`. Reversible transformations must define [`reverse_transformation`](@ref). +They can also define several functions to determine when they are reversed. Transformations +are reversed in the reverse order of application. +""" +function with_reversible_transformation(sys::System, @nospecialize(tf)) + prev_tfs = getmetadata(sys, ReversibleTransformations, nothing)::Union{Nothing, Vector{Any}} + new_tfs::Vector{Any} = if prev_tfs === nothing + [] + else + copy(prev_tfs) + end + push!(new_tfs, tf) + return setmetadata(sys, ReversibleTransformations, new_tfs) +end + +""" + $TYPEDSIGNATURES + +Filter the reversible transformations applied to `sys`. Return a new system with the +updated list of transformations. +""" +function filter_reversible_transformations(fn::F, sys::System) where {F} + prev_tfs = getmetadata(sys, ReversibleTransformations, [])::Union{Nothing, Vector{Any}} + new_tfs::Vector{Any} = if prev_tfs === nothing + [] + else + copy(prev_tfs) + end + filter!(fn, new_tfs) + return setmetadata(sys, ReversibleTransformations, new_tfs) +end + +""" + function reverse_transformation(sys::System, tf) + +Reverse a reversible transformation ([`with_reversible_transformation`](@ref)) applied to `sys`. +""" +function reverse_transformation end + +""" + $TYPEDSIGNATURES + +Return a boolean indicating whether the transformation `tf` is reversed by default +when performing operations that require reversing such transformations. Defaults to +`true` for all transformations. +""" +reverse_transformation_by_default(@nospecialize(tf)) = true + +""" + $TYPEDSIGNATURES + +Return a boolean indicating whether the transformation `tf` is reversed by default +when building the initialization system. +""" +function reverse_transformation_during_initialization(@nospecialize(tf)) + return reverse_transformation_by_default(tf) +end + +function __reverse_transformations_helper(fn::F, sys::System) where {F} + tfs = getmetadata(sys, ReversibleTransformations, [])::Vector{Any} + new_tfs = [] + for tf in Iterators.reverse(tfs) + if !fn(tf)::Bool + push!(new_tfs, tf) + continue + end + sys = reverse_transformation(sys, tf)::System + end + reverse!(new_tfs) + return setmetadata(sys, ReversibleTransformations, new_tfs)::System +end + +""" + $TYPEDSIGNATURES + +Reverse all reversible transformations ([`with_reversible_transformation`](@ref)) applied +to `sys` for which [`reverse_transformation_by_default`](@ref) is `true`. +""" +function reverse_all_default_reversible_transformations(sys::System) + return __reverse_transformations_helper(reverse_transformation_by_default, sys) +end + +""" + $TYPEDSIGNATURES + +Reverse all reversible transformations ([`with_reversible_transformation`](@ref)) applied +to `sys` for which [`reverse_transformation_during_initialization`](@ref) is `true`. +""" +function reverse_transformations_for_initialization(sys::System) + return __reverse_transformations_helper(reverse_transformation_during_initialization, sys) +end diff --git a/lib/ModelingToolkitBase/src/systems/systems.jl b/lib/ModelingToolkitBase/src/systems/systems.jl index 26068c2519..3f7568acce 100644 --- a/lib/ModelingToolkitBase/src/systems/systems.jl +++ b/lib/ModelingToolkitBase/src/systems/systems.jl @@ -88,6 +88,11 @@ function mtkcompile( split = true, kwargs... ) isscheduled(sys) && throw(RepeatedStructuralSimplificationError()) + + # For backward compatibility with old ModelingToolkit which does not + # integrate with the reversible transformation API. + sys = with_reversible_transformation(sys, UnhackSystemTransformation) + # Canonicalize types of arguments to prevent repeated compilation of inner methods inputs = canonicalize_io(unwrap_vars(inputs), "input") outputs = canonicalize_io(unwrap_vars(outputs), "output") @@ -243,14 +248,16 @@ function __mtkcompile( if !has_derivatives && !has_shifts obseqs = copy(original_obs) get_trivial_observed_equations!(Equation[], eqs, obseqs, all_dvs, nothing) - add_array_observed!(obseqs) - obseqs = topsort_equations(sys, obseqs, [eq.lhs for eq in obseqs]) map!(eq -> Symbolics.COMMON_ZERO ~ (eq.rhs - eq.lhs), eqs, eqs) observables = Set{SymbolicT}() for eq in obseqs push!(observables, eq.lhs) end setdiff!(flat_dvs, observables) + sys = remove_unhack_system_transformation(sys) + tf = add_array_observed!(obseqs, flat_dvs) + sys = with_reversible_transformation(sys, tf) + obseqs = topsort_equations(sys, obseqs, [eq.lhs for eq in obseqs]) new_ps = [get_ps(sys); collect(inputs)] @set! sys.eqs = eqs @set! sys.unknowns = flat_dvs @@ -370,8 +377,6 @@ function __mtkcompile( ) end get_trivial_observed_equations!(diffeqs, alg_eqs, obseqs, all_dvs, iv) - add_array_observed!(obseqs) - obseqs = topsort_equations(sys, obseqs, [eq.lhs for eq in obseqs]) for i in eachindex(alg_eqs) eq = alg_eqs[i] alg_eqs[i] = 0 ~ subst(eq.rhs - eq.lhs) @@ -386,6 +391,11 @@ function __mtkcompile( new_dvs = [diffvars; alg_vars] new_ps = [get_ps(sys); collect(inputs)] + sys = remove_unhack_system_transformation(sys) + tf = add_array_observed!(obseqs, new_dvs) + sys = with_reversible_transformation(sys, tf) + obseqs = topsort_equations(sys, obseqs, [eq.lhs for eq in obseqs]) + for eq in new_eqs if SU.query(eq.rhs) do v Moshi.Match.@match v begin @@ -531,26 +541,70 @@ end ndims = ndims(arr) end -function add_array_observed!(obseqs::Vector{Equation}) - array_obsvars = Set{SymbolicT}() - for eq in obseqs - arr, isarr = split_indexed_var(eq.lhs) - isarr && push!(array_obsvars, arr) +struct ScalarizedArrayObserved <: ReversibleTransformations + arrvars::Set{SymbolicT} +end + +reverse_transformation_during_initialization(::ScalarizedArrayObserved) = true + +function add_array_observed!(obseqs::Vector{Equation}, unknowns::Vector{SymbolicT}) + # map of array observed variable (unscalarized) to number of its + # scalarized terms that appear as observed variables + arr_obs_occurrences = Dict{SymbolicT, Int}() + for (i, eq) in enumerate(obseqs) + lhs = eq.lhs + rhs = eq.rhs + unscal, isarr = split_indexed_var(lhs) + isarr || continue + cnt = get(arr_obs_occurrences, unscal, 0) + arr_obs_occurrences[unscal] = cnt + 1 end - for var in array_obsvars - firstind = first(SU.stable_eachindex(var))::SU.StableIndex{Int} - firstind = Tuple(firstind.idxs) - scal = SymbolicT[] - for i in SU.stable_eachindex(var) - push!(scal, var[i]) + + # count variables in unknowns if they are scalarized forms of variables + # also present as observed. e.g. if `x[1]` is an unknown and `x[2] ~ (..)` + # is an observed equation. + for sym in unknowns + unscal, isarr = split_indexed_var(sym) + isarr || continue + cnt = get(arr_obs_occurrences, unscal, 0) + iszero(cnt) && continue + arr_obs_occurrences[unscal] = cnt + 1 + end + + obs_arr_eqs = Equation[] + arrvars = Set{SymbolicT}() + for (arrvar, cnt) in arr_obs_occurrences + cnt == length(arrvar) || continue + firstind = (first(SU.stable_eachindex(arrvar))::SU.StableIndex{Int}).idxs + scal_args = Symbolics.SArgsT() + sizehint!(scal_args, length(arrvar)::Int) + push!(scal_args, Symbolics.SConst(size(arrvar))) + for i in SU.stable_eachindex(arrvar) + push!(scal_args, arrvar[i]) end - if all(isone, firstind) - push!(obseqs, var ~ reshape(scal, size(var))) - continue + rhs = Symbolics.STerm( + SU.array_literal, scal_args; + type = SU.symtype(arrvar), shape = SU.shape(arrvar) + ) + if !all(isone, firstind) + rhs = Symbolics.STerm( + offset_array, + Symbolics.SArgsT((Symbolics.SConst(Tuple(firstind)), rhs)); + type = SU.symtype(arrvar), shape = SU.shape(arrvar) + ) end - push!(obseqs, var ~ offset_array(firstind, reshape(scal, size(var)))) + push!(obs_arr_eqs, arrvar ~ rhs) + push!(arrvars, arrvar) end - return + append!(obseqs, obs_arr_eqs) + + return ScalarizedArrayObserved(arrvars) +end + +function reverse_transformation(sys::AbstractSystem, tf::ScalarizedArrayObserved) + obs = copy(observed(sys)) + filter!(eq -> !(eq.lhs in tf.arrvars), obs) + return @set! sys.observed = obs end """ @@ -842,7 +896,7 @@ function simplify_optimization_system(sys::System; split = true, kwargs...) snlsys = mtkcompile(nlsys; kwargs..., fully_determined = false)::System obs = observed(snlsys) seqs = equations(snlsys) - trueobs = observed(unhack_system(snlsys)) + trueobs = observed(reverse_all_default_reversible_transformations(snlsys)) subs = Dict{SymbolicT, SymbolicT}() for eq in trueobs subs[eq.lhs] = eq.rhs