Skip to content

Commit 6fe09b7

Browse files
Merge pull request #4655 from SciML/as/reversible-tfs
feat: add reversible transformation API
2 parents cc75d01 + b489560 commit 6fe09b7

7 files changed

Lines changed: 231 additions & 38 deletions

File tree

Project.toml

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,6 @@ OrdinaryDiffEqBDF = "6ad6398a-0878-4a85-9266-38940aa047c8"
5050
OrdinaryDiffEqDefault = "50262376-6c5a-4cf5-baba-aaf4f84d72d7"
5151
OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
5252

53-
[sources]
54-
ModelingToolkitBase = {path = "lib/ModelingToolkitBase"}
55-
SciCompDSL = {path = "lib/SciCompDSL"}
56-
5753
[extensions]
5854
MTKFMIExt = "FMIImport"
5955
MTKOrdinaryDiffEqBDFExt = "OrdinaryDiffEqBDF"

lib/ModelingToolkitBase/src/systems/abstractsystem.jl

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -567,7 +567,7 @@ function add_initialization_parameters(sys::AbstractSystem; split = true, _unhac
567567
all_initialvars = Set{SymbolicT}()
568568
# time-independent systems don't initialize unknowns
569569
# but may initialize parameters using guesses for unknowns
570-
_sys = _unhack_sys === nothing ? unhack_system(sys) : _unhack_sys
570+
_sys = _unhack_sys === nothing ? reverse_all_default_reversible_transformations(sys) : _unhack_sys
571571
obs = observed(_sys)
572572
eqs = equations(_sys)
573573
for x in unknowns(_sys)
@@ -698,7 +698,7 @@ function complete(
698698
end
699699
sys = newsys
700700
check_no_parameter_equations(sys)
701-
_unhack_sys = unhack_system(sys)
701+
_unhack_sys = reverse_all_default_reversible_transformations(sys)
702702
if add_initial_parameters
703703
sys = add_initialization_parameters(sys; split, _unhack_sys)::T
704704
end
@@ -1022,6 +1022,11 @@ function refreshed_metadata(meta::Base.ImmutableDict, patch = nothing)
10221022
newmeta = MetadataT()
10231023
hascache = false
10241024
for (k, v) in meta
1025+
# `Base.ImmutableDict` is like a stack. New key-value pairs (even with the same key)
1026+
# are inserted at the top of the stack. `getindex` returns the first matching key.
1027+
# We ignore repetitions of a key we've already encountered to avoid overwriting
1028+
# newer entries with lingering older ones.
1029+
haskey(newmeta, k) && continue
10251030
if k === MutableCacheKey
10261031
hascache = true
10271032
new_cache = maybe_invalidate_cache((v::MutableCacheTLVT)[]::MutableCacheT, patch)::MutableCacheT

lib/ModelingToolkitBase/src/systems/callbacks.jl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -193,7 +193,7 @@ function AffectSystem(
193193
union!(accessed_params, sys_params)
194194

195195
# add scalarized unknowns to the map.
196-
_obs = observed(unhack_system(affectsys))
196+
_obs = observed(reverse_all_default_reversible_transformations(affectsys))
197197
_dvs = vcat(unknowns(affectsys), map(eq -> eq.lhs, _obs))
198198
_dvs = __safe_scalarize_vars(_dvs)
199199
_discs = __safe_scalarize_vars(discretes)
@@ -1415,7 +1415,7 @@ Base.@nospecializeinfer function compile_explicit_affect(
14151415
ps_to_update = discretes(aff)
14161416
dvs_to_update = setdiff(unknowns(aff), getfield.(observed(sys), :lhs))
14171417

1418-
_affsys = unhack_system(affsys)
1418+
_affsys = reverse_all_default_reversible_transformations(affsys)
14191419
aff_ir_info = get_ir_info(affsys)
14201420
aff_ir = get_irstructure(affsys)
14211421
obseqs = Equation[]

lib/ModelingToolkitBase/src/systems/nonlinear/initializesystem.jl

Lines changed: 44 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ function generate_initializesystem_timevarying(
5555
check_units = true, check_defguess = false,
5656
name = nameof(sys), kwargs...
5757
)
58-
_sys = unhack_system(sys)
58+
_sys = reverse_transformations_for_initialization(sys)
5959
trueobs = observed(_sys)
6060
eqs = equations(_sys)
6161
# remove any observed equations that directly or indirectly contain
@@ -226,7 +226,7 @@ function generate_initializesystem_timeindependent(
226226
name = nameof(sys), kwargs...
227227
)
228228
eqs = equations(sys)
229-
_sys = unhack_system(sys)
229+
_sys = reverse_transformations_for_initialization(sys)
230230
trueobs = observed(_sys)
231231
eqs = equations(_sys)
232232
# remove any observed equations that directly or indirectly contain
@@ -885,18 +885,56 @@ function unhack_observed(obseqs, eqs)
885885
return obseqs, eqs
886886
end
887887

888+
"""
889+
$TYPEDEF
890+
891+
Default reversible transformation applied to all systems in `mtkcompile` for backward
892+
compatibility.
893+
"""
894+
abstract type UnhackSystemTransformation <: ReversibleTransformations end
895+
896+
function reverse_transformation(sys::AbstractSystem, ::Type{UnhackSystemTransformation})
897+
return unhack_system(sys)
898+
end
899+
900+
# NOTE:
901+
# 1. Only this MTKBase is loaded. `__mtkcompile` here will remove `UnhackSystemTransformation`
902+
# from the transformations to avoid it becoming an infinite loop. `unhack_system` also
903+
# checks to make sure it doesn't rerun itself.
904+
#
905+
# 2. This MTKBase + old MTKTearing + old MTK: MTKTearing will anyway do its own array
906+
# observed equations and remove them in `unhack_system`. The `UnhackSystemTransformation`
907+
# added in `__mtkcompile` will run `unhack_system` in `reverse_all_default_reversible_transformations`.
908+
#
909+
# 3. This MTKBase + new MTKTearing + old MTK: MTKTearing doesn't implement `unhack_system`
910+
# anymore and removes `UnhackSystemTransformation`. MTKTearing's own reversible transformations
911+
# will run as intended. Calls to `unhack_system` from MTK will not early-exit.
912+
#
913+
# 4. This MTKBase + new MTKTearing + new MTK: `unhack_system` is never called and
914+
# `UnhackSystemTransformation` is not present post-`mtkcompile`.
915+
916+
function remove_unhack_system_transformation(sys::AbstractSystem)
917+
return filter_reversible_transformations(
918+
tf -> tf !== UnhackSystemTransformation, sys
919+
)::System
920+
end
921+
888922
"""
889923
unhack_system(sys::AbstractSystem)
890924
891925
Given a system, remove any codegen oddities applied to it. This is typically used as a
892926
precursor to generating the initialization system. It is the successor of the now
893927
deprecated `unhack_observed`.
928+
929+
DEPRECATED: See `with_reversible_transformation` and the reversible transformation API.
894930
"""
895931
function unhack_system(sys)
896-
obs, eqs = unhack_observed(observed(sys), equations(sys))
897-
@set! sys.observed = obs
898-
@set! sys.eqs = eqs
899-
return sys
932+
# See note above
933+
tfs = getmetadata(sys, ReversibleTransformations, [])::Vector{Any}
934+
for tf in tfs
935+
tf === UnhackSystemTransformation && return sys
936+
end
937+
return reverse_all_default_reversible_transformations(sys)
900938
end
901939

902940
function UnknownsInTimeIndependentInitializationError(eq, non_params)

lib/ModelingToolkitBase/src/systems/problem_utils.jl

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1873,7 +1873,7 @@ function __process_SciMLProblem(
18731873

18741874
add_initials!(sys, op)
18751875

1876-
_sys = unhack_system(sys)
1876+
_sys = reverse_all_default_reversible_transformations(sys)
18771877
obs = observed(_sys)
18781878

18791879
guesses = operating_point_preprocess(sys, guesses; name = "guesses")
@@ -2367,7 +2367,7 @@ function get_u0(sys::AbstractSystem, varmap; kwargs...)
23672367
op = build_operating_point(sys, varmap)
23682368
binds = bindings(sys)
23692369
no_override_merge_except_missing!(op, binds)
2370-
obs = observed(unhack_system(sys))
2370+
obs = observed(reverse_all_default_reversible_transformations(sys))
23712371
add_observed_equations!(op, obs)
23722372

23732373
return varmap_to_vars(op, unknowns(sys); kwargs...)
@@ -2385,7 +2385,7 @@ function get_p(sys::AbstractSystem, varmap; split = is_split(sys), kwargs...)
23852385
binds = bindings(sys)
23862386
no_override_merge_except_missing!(op, binds)
23872387
add_initials!(sys, op)
2388-
obs = observed(unhack_system(sys))
2388+
obs = observed(reverse_all_default_reversible_transformations(sys))
23892389
add_observed_equations!(op, obs)
23902390

23912391
return if split

lib/ModelingToolkitBase/src/systems/system.jl

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1688,3 +1688,103 @@ Get the `IRStructure` associated with the system.
16881688
function get_irstructure(sys::System)
16891689
return get_irstructure_tlv(sys)[]::IRStructure{SymReal}
16901690
end
1691+
1692+
# Reversible transformation API
1693+
1694+
abstract type ReversibleTransformations end
1695+
1696+
"""
1697+
$TYPEDSIGNATURES
1698+
1699+
Add `tf` to the list of reversible transformations applied to `sys`. Returns a modified
1700+
copy of `sys`. Reversible transformations must define [`reverse_transformation`](@ref).
1701+
They can also define several functions to determine when they are reversed. Transformations
1702+
are reversed in the reverse order of application.
1703+
"""
1704+
function with_reversible_transformation(sys::System, @nospecialize(tf))
1705+
prev_tfs = getmetadata(sys, ReversibleTransformations, nothing)::Union{Nothing, Vector{Any}}
1706+
new_tfs::Vector{Any} = if prev_tfs === nothing
1707+
[]
1708+
else
1709+
copy(prev_tfs)
1710+
end
1711+
push!(new_tfs, tf)
1712+
return setmetadata(sys, ReversibleTransformations, new_tfs)
1713+
end
1714+
1715+
"""
1716+
$TYPEDSIGNATURES
1717+
1718+
Filter the reversible transformations applied to `sys`. Return a new system with the
1719+
updated list of transformations.
1720+
"""
1721+
function filter_reversible_transformations(fn::F, sys::System) where {F}
1722+
prev_tfs = getmetadata(sys, ReversibleTransformations, [])::Union{Nothing, Vector{Any}}
1723+
new_tfs::Vector{Any} = if prev_tfs === nothing
1724+
[]
1725+
else
1726+
copy(prev_tfs)
1727+
end
1728+
filter!(fn, new_tfs)
1729+
return setmetadata(sys, ReversibleTransformations, new_tfs)
1730+
end
1731+
1732+
"""
1733+
function reverse_transformation(sys::System, tf)
1734+
1735+
Reverse a reversible transformation ([`with_reversible_transformation`](@ref)) applied to `sys`.
1736+
"""
1737+
function reverse_transformation end
1738+
1739+
"""
1740+
$TYPEDSIGNATURES
1741+
1742+
Return a boolean indicating whether the transformation `tf` is reversed by default
1743+
when performing operations that require reversing such transformations. Defaults to
1744+
`true` for all transformations.
1745+
"""
1746+
reverse_transformation_by_default(@nospecialize(tf)) = true
1747+
1748+
"""
1749+
$TYPEDSIGNATURES
1750+
1751+
Return a boolean indicating whether the transformation `tf` is reversed by default
1752+
when building the initialization system.
1753+
"""
1754+
function reverse_transformation_during_initialization(@nospecialize(tf))
1755+
return reverse_transformation_by_default(tf)
1756+
end
1757+
1758+
function __reverse_transformations_helper(fn::F, sys::System) where {F}
1759+
tfs = getmetadata(sys, ReversibleTransformations, [])::Vector{Any}
1760+
new_tfs = []
1761+
for tf in Iterators.reverse(tfs)
1762+
if !fn(tf)::Bool
1763+
push!(new_tfs, tf)
1764+
continue
1765+
end
1766+
sys = reverse_transformation(sys, tf)::System
1767+
end
1768+
reverse!(new_tfs)
1769+
return setmetadata(sys, ReversibleTransformations, new_tfs)::System
1770+
end
1771+
1772+
"""
1773+
$TYPEDSIGNATURES
1774+
1775+
Reverse all reversible transformations ([`with_reversible_transformation`](@ref)) applied
1776+
to `sys` for which [`reverse_transformation_by_default`](@ref) is `true`.
1777+
"""
1778+
function reverse_all_default_reversible_transformations(sys::System)
1779+
return __reverse_transformations_helper(reverse_transformation_by_default, sys)
1780+
end
1781+
1782+
"""
1783+
$TYPEDSIGNATURES
1784+
1785+
Reverse all reversible transformations ([`with_reversible_transformation`](@ref)) applied
1786+
to `sys` for which [`reverse_transformation_during_initialization`](@ref) is `true`.
1787+
"""
1788+
function reverse_transformations_for_initialization(sys::System)
1789+
return __reverse_transformations_helper(reverse_transformation_during_initialization, sys)
1790+
end

0 commit comments

Comments
 (0)