diff --git a/src/structural_transformation/symbolics_tearing.jl b/src/structural_transformation/symbolics_tearing.jl index 84cee928cd..629ffe1636 100644 --- a/src/structural_transformation/symbolics_tearing.jl +++ b/src/structural_transformation/symbolics_tearing.jl @@ -574,35 +574,35 @@ function tearing_reassemble(state::TearingState, var_eq_matching, # TODO: compute the dependency correctly so that we don't have to do this obs = [fast_substitute(observed(sys), obs_sub); subeqs] - # HACK: Substitute non-scalarized symbolic arrays of observed variables - # E.g. if `p[1] ~ (...)` and `p[2] ~ (...)` then substitute `p => [p[1], p[2]]` in all equations - # ideally, we want to support equations such as `p ~ [p[1], p[2]]` which will then be handled - # by the topological sorting and dependency identification pieces - obs_arr_subs = Dict() + # 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 + handled_obs_arr = Set() + obs_arr_eqs = Equation[] for eq in obs lhs = eq.lhs iscall(lhs) || continue operation(lhs) === getindex || continue Symbolics.shape(lhs) !== Symbolics.Unknown() || continue arg1 = arguments(lhs)[1] - haskey(obs_arr_subs, arg1) && continue - obs_arr_subs[arg1] = [arg1[i] for i in eachindex(arg1)] # e.g. p => [p[1], p[2]] - index_first = eachindex(arg1)[1] - + arg1 in handled_obs_arr && continue + # firstindex returns 1 for multidimensional array symbolics + firstind = first(eachindex(arg1)) + scal = [arg1[i] for i in eachindex(arg1)] # respect non-1-indexed arrays # TODO: get rid of this hack together with the above hack, then remove OffsetArrays dependency - obs_arr_subs[arg1] = Origin(index_first)(obs_arr_subs[arg1]) - end - for i in eachindex(neweqs) - neweqs[i] = fast_substitute(neweqs[i], obs_arr_subs; operator = Symbolics.Operator) - end - for i in eachindex(obs) - obs[i] = fast_substitute(obs[i], obs_arr_subs; operator = Symbolics.Operator) - end - for i in eachindex(subeqs) - subeqs[i] = fast_substitute(subeqs[i], obs_arr_subs; operator = Symbolics.Operator) - end + # `change_origin` is required because `Origin(firstind)(scal)` makes codegen + # try to `create_array(OffsetArray{...}, ...)` which errors. + # `term(Origin(firstind), scal)` doesn't retain the `symtype` and `size` + # of `scal`. + push!(obs_arr_eqs, arg1 ~ change_origin(Origin(firstind), scal)) + push!(handled_obs_arr, arg1) + end + append!(obs, obs_arr_eqs) + append!(subeqs, obs_arr_eqs) + # need to re-sort subeqs + subeqs = ModelingToolkit.topsort_equations(subeqs, [eq.lhs for eq in subeqs]) @set! sys.eqs = neweqs @set! sys.observed = obs @@ -629,6 +629,16 @@ function tearing_reassemble(state::TearingState, var_eq_matching, return invalidate_cache!(sys) end +function change_origin(origin, arr) + return origin(arr) +end + +@register_array_symbolic change_origin(origin::Origin, arr::AbstractArray) begin + size = size(arr) + eltype = eltype(arr) + ndims = ndims(arr) +end + function tearing(state::TearingState; kwargs...) state.structure.solvable_graph === nothing && find_solvables!(state; kwargs...) complete!(state.structure) diff --git a/test/structural_transformation/utils.jl b/test/structural_transformation/utils.jl index 8644d96945..c8142f1dd7 100644 --- a/test/structural_transformation/utils.jl +++ b/test/structural_transformation/utils.jl @@ -40,3 +40,16 @@ end @test ModelingToolkit.đť‘ neighbors(g, 1) == [2] @test ModelingToolkit.đť‘‘neighbors(g, 2) == [1] end + +@testset "array observed used unscalarized in another observed" begin + @variables x(t) y(t)[1:2] z(t)[1:2] + @parameters foo(::AbstractVector)[1:2] + _tmp_fn(x) = 2x + @mtkbuild sys = ODESystem([D(x) ~ z[1] + z[2], y[1] ~ 2t, y[2] ~ 3t, z ~ foo(y)], t) + @test length(equations(sys)) == 1 + @test length(observed(sys)) == 6 + @test any(eq -> isequal(eq.lhs, y), observed(sys)) + @test any(eq -> isequal(eq.lhs, z), observed(sys)) + prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0), [foo => _tmp_fn]) + @test_nowarn prob.f(prob.u0, prob.p, 0.0) +end