feat: add reversible transformation API

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"

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

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[]

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)

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

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