diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl index 3f46b594df..cc582cd473 100644 --- a/src/structural_transformation/symbolics_tearing.jl +++ b/src/structural_transformation/symbolics_tearing.jl @@ -724,7 +724,7 @@ Update the system equations, unknowns, and observables after simplification. """ function update_simplified_system!( state::TearingState, neweqs, solved_eqs, dummy_sub, var_eq_matching, extra_unknowns; - cse_hack = true, array_hack = true) + array_hack = true) @unpack solvable_graph, var_to_diff, eq_to_diff, graph = state.structure diff_to_var = invview(var_to_diff) @@ -748,8 +748,7 @@ function update_simplified_system!( unknowns = [unknowns; extra_unknowns] @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) deps = Vector{Int}[i == 1 ? Int[] : collect(1:(i - 1)) for i in 1:length(solved_eqs)] @@ -793,7 +792,7 @@ appear in the system. Algebraic variables are variables that are not differential variables. """ function tearing_reassemble(state::TearingState, var_eq_matching, - full_var_eq_matching = nothing; simplify = false, mm = nothing, cse_hack = true, array_hack = true) + full_var_eq_matching = nothing; simplify = false, mm = nothing, array_hack = true) extra_vars = Int[] if full_var_eq_matching !== nothing for v in 𝑑vertices(state.structure.graph) @@ -829,7 +828,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching, state, var_eq_matching, eq_ordering, var_ordering, nelim_eq, nelim_var) sys = update_simplified_system!(state, neweqs, solved_eqs, dummy_sub, var_eq_matching, - extra_unknowns; cse_hack, array_hack) + extra_unknowns; array_hack) @set! state.sys = sys @set! sys.tearing_state = state @@ -837,14 +836,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching, 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 @@ -852,13 +844,7 @@ if all `p[i]` are present and the unscalarized form is used in any equation (obs 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() @@ -866,31 +852,6 @@ function cse_and_array_hacks(sys, obs, unknowns, neweqs; cse = true, array = tru 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 @@ -901,31 +862,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. @@ -960,18 +896,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 @@ -999,10 +924,10 @@ new residual equations after tearing. End users are encouraged to call [`structu instead, which calls this function internally. """ function tearing(sys::AbstractSystem, state = TearingState(sys); mm = nothing, - simplify = false, cse_hack = true, array_hack = true, kwargs...) + simplify = false, array_hack = true, kwargs...) var_eq_matching, full_var_eq_matching = tearing(state) invalidate_cache!(tearing_reassemble( - state, var_eq_matching, full_var_eq_matching; mm, simplify, cse_hack, array_hack)) + state, var_eq_matching, full_var_eq_matching; mm, simplify, array_hack)) end """ @@ -1024,7 +949,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, kwargs...) + mm = nothing, array_hack = true, kwargs...) jac = let state = state (eqs, vars) -> begin symeqs = EquationsView(state)[eqs] @@ -1048,5 +973,5 @@ function dummy_derivative(sys, state = TearingState(sys); simplify = false, end var_eq_matching = dummy_derivative_graph!(state, jac; state_priority, kwargs...) - tearing_reassemble(state, var_eq_matching; simplify, mm, cse_hack, array_hack) + tearing_reassemble(state, var_eq_matching; simplify, mm, array_hack) end diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl index 4b0086840a..cf38b94687 100644 --- a/src/systems/nonlinear/initializesystem.jl +++ b/src/systems/nonlinear/initializesystem.jl @@ -739,20 +739,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) diff --git a/test/code_generation.jl b/test/code_generation.jl index cf3d660b81..05b945f503 100644 --- a/test/code_generation.jl +++ b/test/code_generation.jl @@ -78,3 +78,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 + @mtkbuild sys = ODESystem([D(x) ~ y[1] + y[2], y ~ foo(x)], t) + @test length(equations(sys)) == 1 + @test length(observed(sys)) == 3 + prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [foo => _tmp_fn2]) + 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] + @mtkbuild sys = ODESystem( + [D(y) ~ foo(x), D(x) ~ sum(y), zeros(2) ~ foo(prod(z))], t) + @test length(equations(sys)) == 5 + @test length(observed(sys)) == 0 + prob = ODEProblem( + sys, [y => ones(2), z => 2ones(2), x => 3.0], (0.0, 1.0), [foo => _tmp_fn2]) + val[] = 0 + @test_nowarn prob.f(prob.u0, prob.p, 0.0) + @test val[] == 2 + end +end diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl index 7c512d37af..b1ba21b254 100644 --- a/test/initializationsystem.jl +++ b/test/initializationsystem.jl @@ -1650,3 +1650,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 diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl index 67cf2f72a0..24fb166755 100644 --- a/test/structural_transformation/utils.jl +++ b/test/structural_transformation/utils.jl @@ -51,7 +51,7 @@ end @mtkbuild sys = ODESystem( [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], (0.0, 1.0), [foo => _tmp_fn]) @@ -61,60 +61,11 @@ 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 - @mtkbuild sys = ODESystem([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], (0.0, 1.0), [foo => _tmp_fn2]) - 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] - @mtkbuild sys = ODESystem( - [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], (0.0, 1.0), [foo => _tmp_fn2]) - 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] @@ -122,15 +73,8 @@ end @named sys = ODESystem( [D(x) ~ z[1] + z[2] + foo(z)[1], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t) - sys1 = structural_simplify(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 = structural_simplify(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 @@ -143,15 +87,8 @@ end @named sys = ODESystem( [D(x) ~ z[1] + z[2] + foo(z)[1] + w, y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t) - sys1 = structural_simplify(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 = structural_simplify(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