diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl index f4e29346ff..8e7b2ac97e 100644 --- a/src/systems/diffeqs/abstractodesystem.jl +++ b/src/systems/diffeqs/abstractodesystem.jl @@ -1230,7 +1230,7 @@ function flatten_equations(eqs) error("LHS ($(eq.lhs)) and RHS ($(eq.rhs)) must either both be array expressions or both scalar") size(eq.lhs) == size(eq.rhs) || error("Size of LHS ($(eq.lhs)) and RHS ($(eq.rhs)) must match: got $(size(eq.lhs)) and $(size(eq.rhs))") - return collect(eq.lhs) .~ collect(eq.rhs) + return vec(collect(eq.lhs) .~ collect(eq.rhs)) else eq end diff --git a/src/systems/index_cache.jl b/src/systems/index_cache.jl index 4fc2f18bfd..623623da42 100644 --- a/src/systems/index_cache.jl +++ b/src/systems/index_cache.jl @@ -404,6 +404,7 @@ function SymbolicIndexingInterface.timeseries_parameter_index(ic::IndexCache, sy sym = get(ic.symbol_to_variable, sym, nothing) sym === nothing && return nothing end + sym = unwrap(sym) idx = check_index_map(ic.discrete_idx, sym) idx === nothing || return ParameterTimeseriesIndex(idx.clock_idx, (idx.buffer_idx, idx.idx_in_clock)) @@ -411,7 +412,8 @@ function SymbolicIndexingInterface.timeseries_parameter_index(ic::IndexCache, sy args = arguments(sym) idx = timeseries_parameter_index(ic, args[1]) idx === nothing && return nothing - ParameterIndex(idx.portion, (idx.idx..., args[2:end]...), idx.validate_size) + return ParameterTimeseriesIndex( + idx.timeseries_idx, (idx.parameter_idx..., args[2:end]...)) end function check_index_map(idxmap, sym) diff --git a/test/if_lifting.jl b/test/if_lifting.jl index d702355506..b72b468bf1 100644 --- a/test/if_lifting.jl +++ b/test/if_lifting.jl @@ -21,7 +21,9 @@ using ModelingToolkit: t_nounits as t, D_nounits as D, IfLifting, no_if_lift @test operation(only(equations(ss2)).rhs) === ifelse discvar = only(parameters(ss2)) - prob2 = ODEProblem(ss2, [x => 0.0], (0.0, 5.0)) + prob1 = ODEProblem(ss1, [ss1.x => 0.0], (0.0, 5.0)) + sol1 = solve(prob1, Tsit5()) + prob2 = ODEProblem(ss2, [ss2.x => 0.0], (0.0, 5.0)) sol2 = solve(prob2, Tsit5()) @test count(isapprox(pi), sol2.t) == 2 @test any(isapprox(pi), sol2.discretes[1].t) diff --git a/test/symbolic_indexing_interface.jl b/test/symbolic_indexing_interface.jl index e46c197dde..4a7cd926b4 100644 --- a/test/symbolic_indexing_interface.jl +++ b/test/symbolic_indexing_interface.jl @@ -224,3 +224,14 @@ end end @test isempty(get_all_timeseries_indexes(sys, a)) end + +@testset "`timeseries_parameter_index` on unwrapped scalarized timeseries parameter" begin + @variables x(t)[1:2] + @parameters p(t)[1:2, 1:2] + ev = [x[1] ~ 2.0] => [p ~ -ones(2, 2)] + @mtkbuild sys = ODESystem(D(x) ~ p * x, t; continuous_events = [ev]) + p = ModelingToolkit.unwrap(p) + @test timeseries_parameter_index(sys, p) === ParameterTimeseriesIndex(1, (1, 1)) + @test timeseries_parameter_index(sys, p[1, 1]) === + ParameterTimeseriesIndex(1, (1, 1, 1, 1)) +end