test: fix namespacing

Author: vyudu

src/systems/callbacks.jl

@@ -258,18 +258,12 @@ function make_affect(affect::Vector{Equation}; warn = true)
     aff_map = Dict(zip(p_as_unknowns, discretes))
     rev_map = Dict([v => k for (k, v) in aff_map])
     affect = Symbolics.substitute(affect, rev_map)
-    @mtkbuild affectsys = ImplicitDiscreteSystem(
-                  affect, iv, collect(union(unknowns, p_as_unknowns)), cb_params)
+    @mtkbuild affectsys = ImplicitDiscreteSystem(affect, iv, collect(union(unknowns, p_as_unknowns)), cb_params)
     params = filter(isparameter, map(x -> only(arguments(unwrap(x))), cb_params))
-    @show params
-
     for u in unknowns
         aff_map[u] = u
     end
 
-    @show unknowns
-    @show params
-
     return AffectSystem(affectsys, collect(unknowns), params, discretes, aff_map)
 end
 
@@ -494,16 +488,22 @@ function namespace_affect(affect::FunctionalAffect, s)
         context(affect))
 end
 
-namespace_affect(affect::AffectSystem, s) = AffectSystem(system(affect), renamespace.((s,), discretes(affect)))
-namespace_affects(af::Union{Nothing, Affect}, s) = af isa Affect ? namespace_affect(af, s) : nothing
+function namespace_affect(affect::AffectSystem, s) 
+    AffectSystem(renamespace(s, system(affect)),
+        renamespace.((s,), unknowns(affect)),
+        renamespace.((s,), parameters(affect)),
+        renamespace.((s,), discretes(affect)),
+        Dict([k => renamespace(s, v) for (k, v) in aff_to_sys(affect)]))
+end
+namespace_affect(af::Nothing, s) = nothing
 
 function namespace_callback(cb::SymbolicContinuousCallback, s)::SymbolicContinuousCallback
     SymbolicContinuousCallback(
         namespace_equation.(equations(cb), (s,)),
-        namespace_affects(affects(cb), s),
-        affect_neg = namespace_affects(affect_negs(cb), s),
-        initialize = namespace_affects(initialize_affects(cb), s),
-        finalize = namespace_affects(finalize_affects(cb), s),
+        namespace_affect(affects(cb), s),
+        affect_neg = namespace_affect(affect_negs(cb), s),
+        initialize = namespace_affect(initialize_affects(cb), s),
+        finalize = namespace_affect(finalize_affects(cb), s),
         rootfind = cb.rootfind)
 end
 
@@ -794,9 +794,9 @@ function compile_affect(
 
     ps = parameters(aff)
     dvs = unknowns(aff)
-    @show ps
 
     if aff isa AffectSystem
+        affsys = system(aff)
         aff_map = aff_to_sys(aff)
         sys_map = Dict([v => k for (k, v) in aff_map])
         build_initializeprob = has_alg_eqs(sys)
@@ -809,11 +809,11 @@ function compile_affect(
                 push!(pmap, pre_p => pval)
             end
             guesses = Pair[u => integrator[aff_map[u]] for u in updated_vals(aff)]
-            affprob = ImplicitDiscreteProblem(system(aff), Pair[], (0, 1), pmap; guesses, build_initializeprob)
+            affprob = ImplicitDiscreteProblem(affsys, Pair[], (0, 1), pmap; guesses, build_initializeprob)
 
             affsol = init(affprob, SimpleIDSolve())
             for u in unknowns(aff)
-                integrator[u] = affsol[u]
+                integrator[u] = affsol[sys_map[u]]
             end
             for p in discretes(aff)
                 integrator.ps[p] = affsol[sys_map[p]]
@@ -899,9 +899,9 @@ function continuous_events(sys::AbstractSystem)
     systems = get_systems(sys)
     cbs = [obs;
            reduce(vcat,
-               (map(o -> namespace_callback(o, s), continuous_events(s))
-               for s in systems),
+               (map(o -> namespace_callback(o, s), continuous_events(s)) for s in systems),
                init = SymbolicContinuousCallback[])]
+    @show cbs
     filter(!isempty, cbs)
 end
 

test/symbolic_events.jl

@@ -644,6 +644,7 @@ end
         dprob = DiscreteProblem(jsys, u0, tspan, p)
         jprob = JumpProblem(jsys, dprob, Direct(); kwargs...)
         sol = solve(jprob, SSAStepper(); tstops = tstops)
+        @show sol
         @test (sol(1.000000000001)[1] - sol(0.99999999999)[1]) == 1
         paramtotest === nothing || (@test sol.ps[paramtotest] == 1.0)
         @test sol(40.0)[1] == 0
@@ -654,7 +655,7 @@ end
     @variables A(t) B(t)
 
     cond1 = (t == t1)
-    affect1 = [A ~ Pre(A) + 1]
+    affect1 = [A ~ A + 1]
     cb1 = cond1 => affect1
     cond2 = (t == t2)
     affect2 = [k ~ 1.0]
@@ -704,7 +705,7 @@ end
     testsol(jsys6, u0, p, tspan; tstops = [1.0, 2.0], rng, paramtotest = k)
 end
 
-@testset "Oscillator" begin
+@testset "Namespacing" begin
     function oscillator_ce(k = 1.0; name)
         sts = @variables x(t)=1.0 v(t)=0.0 F(t)
         ps = @parameters k=k Θ=0.5
@@ -1152,7 +1153,7 @@ end
     f = ModelingToolkit.FunctionalAffect(
         f = (i, u, p, c) -> seen = true, sts = [], pars = [], discretes = [])
     cb1 = ModelingToolkit.SymbolicContinuousCallback(
-        [x ~ 0], Equation[], initialize = [x ~ 1.5], finalize = f)
+        [x ~ 0], nothing, initialize = [x ~ 1.5], finalize = f)
     @mtkbuild sys = ODESystem(D(x) ~ -1, t, [x], []; continuous_events = [cb1])
     prob = ODEProblem(sys, [x => 1.0], (0.0, 2), [])
     sol = solve(prob, Tsit5(); dtmax = 0.01)