Skip to content

fix: remove CSE hack, fix unscalarized variables in initializeprobpmap #3697

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jun 7, 2025
Merged

fix: remove CSE hack, fix unscalarized variables in initializeprobpmap #3697

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c58b67f
refactor: remove CSE hack
AayushSabharwal Jun 5, 2025
aff913f
test: test array unknowns occurring unscalarized in initializeprobpmap
AayushSabharwal Jun 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 11 additions & 86 deletions src/structural_transformation/symbolics_tearing.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -929,7 +929,7 @@ Update the system equations, unknowns, and observables after simplification.
"""
function update_simplified_system!(
state::TearingState, neweqs, solved_eqs, dummy_sub, var_sccs, extra_unknowns;
cse_hack = true, array_hack = true, D = nothing, iv = nothing)
array_hack = true, D = nothing, iv = nothing)
@unpack fullvars, structure = state
@unpack solvable_graph, var_to_diff, eq_to_diff, graph = structure
diff_to_var = invview(var_to_diff)
Expand Down Expand Up @@ -978,8 +978,7 @@ function update_simplified_system!(
end
@set! sys.unknowns = unknowns

obs = cse_and_array_hacks(
sys, obs, unknowns, neweqs; cse = cse_hack, array = array_hack)
obs = tearing_hacks(sys, obs, unknowns, neweqs; array = array_hack)

@set! sys.eqs = neweqs
@set! sys.observed = obs
Expand Down Expand Up @@ -1035,7 +1034,7 @@ differential variables.
according to `full_var_eq_matching`.
"""
function tearing_reassemble(state::TearingState, var_eq_matching::Matching,
full_var_eq_matching::Matching, var_sccs::Vector{Vector{Int}}; simplify = false, mm, cse_hack = true,
full_var_eq_matching::Matching, var_sccs::Vector{Vector{Int}}; simplify = false, mm,
array_hack = true, fully_determined = true)
extra_eqs_vars = get_extra_eqs_vars(
state, var_eq_matching, full_var_eq_matching, fully_determined)
Expand Down Expand Up @@ -1074,7 +1073,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching::Matching,
# var_eq_matching and full_var_eq_matching are now invalidated

sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_sccs,
extra_unknowns; cse_hack, array_hack, iv, D)
extra_unknowns; array_hack, iv, D)

@set! state.sys = sys
@set! sys.tearing_state = state
Expand Down Expand Up @@ -1223,60 +1222,22 @@ function get_extra_eqs_vars(
end

"""
# HACK 1

Since we don't support array equations, any equation of the sort `x[1:n] ~ f(...)[1:n]`
gets turned into `x[1] ~ f(...)[1], x[2] ~ f(...)[2]`. Repeatedly calling `f` gets
_very_ expensive. this hack performs a limited form of CSE specifically for this case to
avoid the unnecessary cost. This and the below hack are implemented simultaneously

# HACK 2
# HACK

Add equations for array observed variables. If `p[i] ~ (...)` are equations, add an
equation `p ~ [p[1], p[2], ...]` allow topsort to reorder them only add the new equation
if all `p[i]` are present and the unscalarized form is used in any equation (observed or
not) we first count the number of times the scalarized form of each observed variable
occurs in observed equations (and unknowns if it's split).
"""
function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = true)
# HACK 1
# mapping of rhs to temporary CSE variable
# `f(...) => tmpvar` in above example
rhs_to_tempvar = Dict()

# HACK 2
function tearing_hacks(sys, obs, unknowns, neweqs; array = true)
# map of array observed variable (unscalarized) to number of its
# scalarized terms that appear in observed equations
arr_obs_occurrences = Dict()
for (i, eq) in enumerate(obs)
lhs = eq.lhs
rhs = eq.rhs

# HACK 1
if cse && is_getindexed_array(rhs)
rhs_arr = arguments(rhs)[1]
iscall(rhs_arr) && operation(rhs_arr) isa Symbolics.Operator && continue
if !haskey(rhs_to_tempvar, rhs_arr)
tempvar = gensym(Symbol(lhs))
N = length(rhs_arr)
tempvar = unwrap(Symbolics.variable(
tempvar; T = Symbolics.symtype(rhs_arr)))
tempvar = setmetadata(
tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr))
tempeq = tempvar ~ rhs_arr
rhs_to_tempvar[rhs_arr] = tempvar
push!(obs, tempeq)
end

# getindex_wrapper is used because `observed2graph` treats `x` and `x[i]` as different,
# so it doesn't find a dependency between this equation and `tempvar ~ rhs_arr`
# which fails the topological sort
neweq = lhs ~ getindex_wrapper(
rhs_to_tempvar[rhs_arr], Tuple(arguments(rhs)[2:end]))
obs[i] = neweq
end
# end HACK 1

array || continue
iscall(lhs) || continue
operation(lhs) === getindex || continue
Expand All @@ -1287,31 +1248,6 @@ function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = tru
continue
end

# Also do CSE for `equations(sys)`
if cse
for (i, eq) in enumerate(neweqs)
(; lhs, rhs) = eq
is_getindexed_array(rhs) || continue
rhs_arr = arguments(rhs)[1]
if !haskey(rhs_to_tempvar, rhs_arr)
tempvar = gensym(Symbol(lhs))
N = length(rhs_arr)
tempvar = unwrap(Symbolics.variable(
tempvar; T = Symbolics.symtype(rhs_arr)))
tempvar = setmetadata(
tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr))
tempeq = tempvar ~ rhs_arr
rhs_to_tempvar[rhs_arr] = tempvar
push!(obs, tempeq)
end
# don't need getindex_wrapper, but do it anyway to know that this
# hack took place
neweq = lhs ~ getindex_wrapper(
rhs_to_tempvar[rhs_arr], Tuple(arguments(rhs)[2:end]))
neweqs[i] = neweq
end
end

# 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.
Expand Down Expand Up @@ -1346,18 +1282,7 @@ function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = tru
return obs
end

function is_getindexed_array(rhs)
(!ModelingToolkit.isvariable(rhs) || ModelingToolkit.iscalledparameter(rhs)) &&
iscall(rhs) && operation(rhs) === getindex &&
Symbolics.shape(rhs) != Symbolics.Unknown()
end

# PART OF HACK 1
getindex_wrapper(x, i) = x[i...]

@register_symbolic getindex_wrapper(x::AbstractArray, i::Tuple{Vararg{Int}})

# PART OF HACK 2
# PART OF HACK
function change_origin(origin, arr)
if all(isone, Tuple(origin))
return arr
Expand Down Expand Up @@ -1385,11 +1310,11 @@ new residual equations after tearing. End users are encouraged to call [`mtkcomp
instead, which calls this function internally.
"""
function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing,
simplify = false, cse_hack = true, array_hack = true, fully_determined = true, kwargs...)
simplify = false, array_hack = true, fully_determined = true, kwargs...)
var_eq_matching, full_var_eq_matching, var_sccs, can_eliminate = tearing(state)
invalidate_cache!(tearing_reassemble(
state, var_eq_matching, full_var_eq_matching, var_sccs; mm,
simplify, cse_hack, array_hack, fully_determined))
simplify, array_hack, fully_determined))
end

"""
Expand All @@ -1399,7 +1324,7 @@ Perform index reduction and use the dummy derivative technique to ensure that
the system is balanced.
"""
function dummy_derivative(sys, state = TearingState(sys); simplify = false,
mm = nothing, cse_hack = true, array_hack = true, fully_determined = true, kwargs...)
mm = nothing, array_hack = true, fully_determined = true, kwargs...)
jac = let state = state
(eqs, vars) -> begin
symeqs = EquationsView(state)[eqs]
Expand All @@ -1425,5 +1350,5 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false,
state, jac; state_priority,
kwargs...)
tearing_reassemble(state, var_eq_matching, full_var_eq_matching, var_sccs;
simplify, mm, cse_hack, array_hack, fully_determined)
simplify, mm, array_hack, fully_determined)
end
14 changes: 0 additions & 14 deletions src/systems/nonlinear/initializesystem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -780,20 +780,6 @@ function unhack_observed(obseqs::Vector{Equation}, eqs::Vector{Equation})
push!(rm_idxs, i)
continue
end
if operation(eq.rhs) == StructuralTransformations.getindex_wrapper
var, idxs = arguments(eq.rhs)
subs[eq.rhs] = var[idxs...]
push!(tempvars, var)
end
end

for (i, eq) in enumerate(eqs)
iscall(eq.rhs) || continue
if operation(eq.rhs) == StructuralTransformations.getindex_wrapper
var, idxs = arguments(eq.rhs)
subs[eq.rhs] = var[idxs...]
push!(tempvars, var)
end
end

for (i, eq) in enumerate(obseqs)
Expand Down
31 changes: 31 additions & 0 deletions test/code_generation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,3 +79,34 @@ end
@test SciMLBase.successful_retcode(sol)
end
end

@testset "scalarized array observed calling same function multiple times" begin
@variables x(t) y(t)[1:2]
@parameters foo(::Real)[1:2]
val = Ref(0)
function _tmp_fn2(x)
val[] += 1
return [x, 2x]
end
@mtkcompile sys = System([D(x) ~ y[1] + y[2], y ~ foo(x)], t)
@test length(equations(sys)) == 1
@test length(ModelingToolkit.observed(sys)) == 3
prob = ODEProblem(sys, [x => 1.0, foo => _tmp_fn2], (0.0, 1.0))
val[] = 0
@test_nowarn prob.f(prob.u0, prob.p, 0.0)
@test val[] == 1

@testset "CSE in equations(sys)" begin
val[] = 0
@variables z(t)[1:2]
@mtkcompile sys = System(
[D(y) ~ foo(x), D(x) ~ sum(y), zeros(2) ~ foo(prod(z))], t)
@test length(equations(sys)) == 5
@test length(ModelingToolkit.observed(sys)) == 0
prob = ODEProblem(
sys, [y => ones(2), z => 2ones(2), x => 3.0, foo => _tmp_fn2], (0.0, 1.0))
val[] = 0
@test_nowarn prob.f(prob.u0, prob.p, 0.0)
@test val[] == 2
end
end
21 changes: 21 additions & 0 deletions test/initializationsystem.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1643,3 +1643,24 @@ end
@test !SciMLBase.isinplace(prob)
@test !SciMLBase.isinplace(prob.f.initialization_data.initializeprob)
end

@testset "Array unknowns occurring unscalarized in initializeprobpmap" begin
@variables begin
u(t)[1:2] = 0.9ones(2)
x(t)[1:2], [guess = 0.01ones(2)]
o(t)[1:2]
end
@parameters p[1:4] = [2.0, 1.875, 2.0, 1.875]

eqs = [D(u[1]) ~ p[1] * u[1] - p[2] * u[1] * u[2] + x[1] + 0.1
D(u[2]) ~ p[4] * u[1] * u[2] - p[3] * u[2] - x[2]
o[1] ~ sum(p) * sum(u)
o[2] ~ sum(p) * sum(x)
x[1] ~ 0.01exp(-1)
x[2] ~ 0.01cos(t)]

@mtkbuild sys = ODESystem(eqs, t)
prob = ODEProblem(sys, [], (0.0, 1.0))
sol = solve(prob, Tsit5())
@test SciMLBase.successful_retcode(sol)
end
73 changes: 5 additions & 68 deletions test/structural_transformation/utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,7 @@ end
@mtkcompile sys = System(
[D(x) ~ z[1] + z[2] + foo(z)[1], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t)
@test length(equations(sys)) == 1
@test length(observed(sys)) == 7
@test length(observed(sys)) == 6
@test any(obs -> isequal(obs, y), observables(sys))
@test any(obs -> isequal(obs, z), observables(sys))
prob = ODEProblem(sys, [x => 1.0, foo => _tmp_fn], (0.0, 1.0))
Expand All @@ -62,76 +62,20 @@ end
@test length(unknowns(isys)) == 5
@test length(equations(isys)) == 4
@test !any(equations(isys)) do eq
iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.getindex_wrapper,
StructuralTransformations.change_origin]
iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.change_origin]
end
end

@testset "scalarized array observed calling same function multiple times" begin
@variables x(t) y(t)[1:2]
@parameters foo(::Real)[1:2]
val = Ref(0)
function _tmp_fn2(x)
val[] += 1
return [x, 2x]
end
@mtkcompile sys = System([D(x) ~ y[1] + y[2], y ~ foo(x)], t)
@test length(equations(sys)) == 1
@test length(observed(sys)) == 4
prob = ODEProblem(sys, [x => 1.0, foo => _tmp_fn2], (0.0, 1.0))
val[] = 0
@test_nowarn prob.f(prob.u0, prob.p, 0.0)
@test val[] == 1

isys = ModelingToolkit.generate_initializesystem(sys)
@test length(unknowns(isys)) == 3
@test length(equations(isys)) == 2
@test !any(equations(isys)) do eq
iscall(eq.rhs) && operation(eq.rhs) in [StructuralTransformations.getindex_wrapper,
StructuralTransformations.change_origin]
end

@testset "CSE hack in equations(sys)" begin
val[] = 0
@variables z(t)[1:2]
@mtkcompile sys = System(
[D(y) ~ foo(x), D(x) ~ sum(y), zeros(2) ~ foo(prod(z))], t)
@test length(equations(sys)) == 5
@test length(observed(sys)) == 2
prob = ODEProblem(
sys, [y => ones(2), z => 2ones(2), x => 3.0, foo => _tmp_fn2], (0.0, 1.0))
val[] = 0
@test_nowarn prob.f(prob.u0, prob.p, 0.0)
@test val[] == 2

isys = ModelingToolkit.generate_initializesystem(sys)
@test length(unknowns(isys)) == 5
@test length(equations(isys)) == 2
@test !any(equations(isys)) do eq
iscall(eq.rhs) &&
operation(eq.rhs) in [StructuralTransformations.getindex_wrapper,
StructuralTransformations.change_origin]
end
end
end

@testset "array and cse hacks can be disabled" begin
@testset "array hack can be disabled" begin
@testset "fully_determined = true" begin
@variables x(t) y(t)[1:2] z(t)[1:2]
@parameters foo(::AbstractVector)[1:2]
_tmp_fn(x) = 2x
@named sys = System(
[D(x) ~ z[1] + z[2] + foo(z)[1], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t)

sys1 = mtkcompile(sys; cse_hack = false)
@test length(observed(sys1)) == 6
@test !any(observed(sys1)) do eq
iscall(eq.rhs) &&
operation(eq.rhs) == StructuralTransformations.getindex_wrapper
end

sys2 = mtkcompile(sys; array_hack = false)
@test length(observed(sys2)) == 5
@test length(observed(sys2)) == 4
@test !any(observed(sys2)) do eq
iscall(eq.rhs) && operation(eq.rhs) == StructuralTransformations.change_origin
end
Expand All @@ -144,15 +88,8 @@ end
@named sys = System(
[D(x) ~ z[1] + z[2] + foo(z)[1] + w, y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t)

sys1 = mtkcompile(sys; cse_hack = false, fully_determined = false)
@test length(observed(sys1)) == 6
@test !any(observed(sys1)) do eq
iscall(eq.rhs) &&
operation(eq.rhs) == StructuralTransformations.getindex_wrapper
end

sys2 = mtkcompile(sys; array_hack = false, fully_determined = false)
@test length(observed(sys2)) == 5
@test length(observed(sys2)) == 4
@test !any(observed(sys2)) do eq
iscall(eq.rhs) && operation(eq.rhs) == StructuralTransformations.change_origin
end
Expand Down
Loading