pass variables/states as named-tuples into affect

Author: dd

src/systems/callbacks.jl

@@ -23,10 +23,11 @@ function FunctionalAffect(f, sts, pars, ctx = nothing)
     # sts & pars contain either pairs: resistor.R => R, or Syms: R
     vs = [x isa Pair ? x.first : x for x in sts]
     vs_syms = [x isa Pair ? Symbol(x.second) : getname(x) for x in sts]
+    length(vs_syms) == length(unique(vs_syms)) || error("Variables are not unique")
 
     ps = [x isa Pair ? x.first : x for x in pars]
     ps_syms = [x isa Pair ? Symbol(x.second) : getname(x) for x in pars]
-    length(vs_syms) + length(ps_syms) == length(unique(vcat(vs_syms, ps_syms))) || error("All symbols for variables & parameters must be unique.")
+    length(ps_syms) == length(unique(ps_syms)) || error("Parameters are not unique")
 
     FunctionalAffect(f, vs, vs_syms, ps, ps_syms, ctx)
 end
@@ -146,10 +147,10 @@ end
 
 struct SymbolicDiscreteCallback
     # condition can be one of:
-    # TODO: Iterative
     #   Δt::Real - Periodic with period Δt
     #   Δts::Vector{Real} - events trigger in this times (Preset)
     #   condition::Vector{Equation} - event triggered when condition is true
+    # TODO: Iterative
     condition
     affects
 
@@ -411,18 +412,14 @@ function compile_user_affect(affect::FunctionalAffect, sys, dvs, ps; kwargs...)
 
     # HACK: filter out eliminated symbols. Not clear this is the right thing to do
     # (MTK should keep these symbols)
-    v = filter(x -> !isnothing(x[1]), collect(zip(v_inds, states_syms(affect))))
-    v_inds = [x[1] for x in v]
-    v_syms = [x[2] for x in v]
-    p = filter(x -> !isnothing(x[1]), collect(zip(p_inds, parameters_syms(affect))))
-    p_inds = [x[1] for x in p]
-    p_syms = [x[2] for x in p]
-
-    kwargs = zip(vcat(v_syms, p_syms), vcat(v_inds, p_inds))
+    u = filter(x -> !isnothing(x[2]), collect(zip(states_syms(affect), v_inds)))
+    p = filter(x -> !isnothing(x[2]), collect(zip(parameters_syms(affect), p_inds)))
+    u = NamedTuple(u)
+    p = NamedTuple(p)
 
-    let kwargs=kwargs, user_affect=func(affect), ctx = context(affect)
+    let u=u, p=p, user_affect=func(affect), ctx = context(affect)
         function (integ)
-            user_affect(integ, ctx; kwargs...)
+            user_affect(integ, u, p, ctx)
         end
     end
 end

test/funcaffect.jl

@@ -6,7 +6,7 @@ D = Differential(t)
 
 eqs = [ D(u) ~ -u ]
 
-affect1!(integ, ctx; u) = integ.u[u] += 10
+affect1!(integ, u, p, ctx) = integ.u[u.u] += 10
 
 @named sys = ODESystem(eqs, t, [u], [], discrete_events=[[4.0, ]=>(affect1!, [u], [], nothing)])
 prob = ODEProblem(sys, [u=> 10.0], (0, 10.0))
@@ -15,8 +15,8 @@ i4 = findfirst(==(4.0), sol[:t])
 @test sol.u[i4+1][1] > 10.0
 
 # context
-function affect2!(integ, ctx; u)
-    integ.u[u] += ctx[1]
+function affect2!(integ, u, p, ctx)
+    integ.u[u.u] += ctx[1]
     ctx[1] *= 2
 end
 ctx1 = [10.0, ]
@@ -30,9 +30,9 @@ i8 = findfirst(==(8.0), sol[:t])
 @test ctx1[1] == 40.0
 
 # parameter
-function affect3!(integ, ctx; u, a)
-    integ.u[u] += integ.p[a]
-    integ.p[a] *= 2
+function affect3!(integ, u, p, ctx)
+    integ.u[u.u] += integ.p[p.a]
+    integ.p[p.a] *= 2
 end
 
 @parameters a = 10.0
@@ -46,9 +46,9 @@ i8 = findfirst(==(8.0), sol[:t])
 @test sol.u[i8+1][1] > 20.0
 
 # rename parameter
-function affect3!(integ, ctx; u, b)
-    integ.u[u] += integ.p[b]
-    integ.p[b] *= 2
+function affect3!(integ, u, p, ctx)
+    integ.u[u.u] += integ.p[p.b]
+    integ.p[p.b] *= 2
 end
 
 @named sys = ODESystem(eqs, t, [u], [a], discrete_events=[[4.0, 8.0]=>(affect3!, [u], [a=> :b], nothing)])
@@ -61,8 +61,21 @@ i8 = findfirst(==(8.0), sol[:t])
 @test sol.u[i8+1][1] > 20.0
 
 # same name
-@test_throws ErrorException ODESystem(eqs, t, [u], [a], discrete_events=[[4.0, 8.0]=>(affect3!, [u], [a=> :u], nothing)]; name=:sys)
+@variables v(t)
+@test_throws ErrorException ODESystem(eqs, t, [u], [a], discrete_events=[[4.0, 8.0]=>(affect3!, [u, v => :u], [a], nothing)]; name=:sys)
 
+@named resistor = ODESystem(D(v) ~ v, t, [v], [])
 
+# nested namespace
+ctx = [0]
+function affect4!(integ, u, p, ctx)
+    ctx[1] += 1
+    @test u.resistor₊v == 1
+end
+s1 = compose(ODESystem(Equation[], t, [], [], name=:s1, discrete_events=1.0=>(affect4!, [resistor.v], [], ctx)), resistor)
+s2 = structural_simplify(s1)
+prob = ODEProblem(s2, [resistor.v=> 10.0], (0, 2.01))
+sol = solve(prob, Tsit5())
+@test ctx[1] == 2