fix: fix callback codegen, observed eqs with non-scalarized symbolic arrays

Author: AayushSabharwal

src/structural_transformation/symbolics_tearing.jl

@@ -551,7 +551,43 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     end
 
     sys = state.sys
+
+    obs_sub = dummy_sub
+    for eq in neweqs
+        isdiffeq(eq) || continue
+        obs_sub[eq.lhs] = eq.rhs
+    end
+    # 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()
+
+    for eq in obs
+        lhs = eq.lhs
+        istree(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)]
+    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
+
     @set! sys.eqs = neweqs
+    @set! sys.observed = obs
+
     unknowns = Any[v
                    for (i, v) in enumerate(fullvars)
                    if diff_to_var[i] === nothing && ispresent(i)]
@@ -563,15 +599,6 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     @set! sys.unknowns = unknowns
     @set! sys.substitutions = Substitutions(subeqs, deps)
 
-    obs_sub = dummy_sub
-    for eq in equations(sys)
-        isdiffeq(eq) || continue
-        obs_sub[eq.lhs] = eq.rhs
-    end
-    # TODO: compute the dependency correctly so that we don't have to do this
-    obs = [fast_substitute(observed(sys), obs_sub); subeqs]
-    @set! sys.observed = obs
-
     # Only makes sense for time-dependent
     # TODO: generalize to SDE
     if sys isa ODESystem

src/systems/callbacks.jl

@@ -388,6 +388,7 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
             return (args...) -> () # We don't do anything in the callback, we're just after the event
         end
     else
+        eqs = flatten_equations(eqs)
         rhss = map(x -> x.rhs, eqs)
         outvar = :u
         if outputidxs === nothing
@@ -457,7 +458,7 @@ end
 
 function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = unknowns(sys),
         ps = full_parameters(sys); kwargs...)
-    eqs = map(cb -> cb.eqs, cbs)
+    eqs = map(cb -> flatten_equations(cb.eqs), cbs)
     num_eqs = length.(eqs)
     (isempty(eqs) || sum(num_eqs) == 0) && return nothing
     # fuse equations to create VectorContinuousCallback
@@ -471,12 +472,8 @@ function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = unknow
     rhss = map(x -> x.rhs, eqs)
     root_eq_vars = unique(collect(Iterators.flatten(map(ModelingToolkit.vars, rhss))))
 
-    u = map(x -> time_varying_as_func(value(x), sys), dvs)
-    p = map.(x -> time_varying_as_func(value(x), sys), reorder_parameters(sys, ps))
-    t = get_iv(sys)
-    pre = get_preprocess_constants(rhss)
-    rf_oop, rf_ip = build_function(rhss, u, p..., t; expression = Val{false},
-        postprocess_fbody = pre, kwargs...)
+    rf_oop, rf_ip = generate_custom_function(sys, rhss, dvs, ps; expression = Val{false},
+        kwargs...)
 
     affect_functions = map(cbs) do cb # Keep affect function separate
         eq_aff = affects(cb)
@@ -487,16 +484,16 @@ function generate_rootfinding_callback(cbs, sys::AbstractODESystem, dvs = unknow
         cond = function (u, t, integ)
             if DiffEqBase.isinplace(integ.sol.prob)
                 tmp, = DiffEqBase.get_tmp_cache(integ)
-                rf_ip(tmp, u, parameter_values(integ)..., t)
+                rf_ip(tmp, u, parameter_values(integ), t)
                 tmp[1]
             else
-                rf_oop(u, parameter_values(integ)..., t)
+                rf_oop(u, parameter_values(integ), t)
             end
         end
         ContinuousCallback(cond, affect_functions[])
     else
         cond = function (out, u, t, integ)
-            rf_ip(out, u, parameter_values(integ)..., t)
+            rf_ip(out, u, parameter_values(integ), t)
         end
 
         # since there may be different number of conditions and affects,

test/odesystem.jl

@@ -995,3 +995,30 @@ let # Issue https://github.com/SciML/ModelingToolkit.jl/issues/2322
     sol = solve(prob, Rodas4())
     @test sol(1)[]≈0.6065307685451087 rtol=1e-4
 end
+
+# Issue#2599
+@variables x(t) y(t)
+eqs = [D(x) ~ x * t, y ~ 2x]
+@mtkbuild sys = ODESystem(eqs, t; continuous_events = [[y ~ 3] => [x ~ 2]])
+prob = ODEProblem(sys, [x => 1.0], (0.0, 10.0))
+@test_nowarn solve(prob, Tsit5())
+
+# Issue#2383
+@variables x(t)[1:3]
+@parameters p[1:3, 1:3]
+eqs = [
+    D(x) ~ p * x
+]
+@mtkbuild sys = ODESystem(eqs, t; continuous_events = [[norm(x) ~ 3.0] => [x ~ ones(3)]])
+# array affect equations used to not work
+prob1 = @test_nowarn ODEProblem(sys, [x => ones(3)], (0.0, 10.0), [p => ones(3, 3)])
+sol1 = @test_nowarn solve(prob1, Tsit5())
+
+# array condition equations also used to not work
+@mtkbuild sys = ODESystem(
+    eqs, t; continuous_events = [[x ~ sqrt(3) * ones(3)] => [x ~ ones(3)]])
+# array affect equations used to not work
+prob2 = @test_nowarn ODEProblem(sys, [x => ones(3)], (0.0, 10.0), [p => ones(3, 3)])
+sol2 = @test_nowarn solve(prob, Tsit5())
+
+@test sol1 ≈ sol2