using integrator interface

Author: dd

src/systems/callbacks.jl

@@ -23,12 +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(ps_syms) == length(unique(ps_syms)) || error("Parameters are not unique.")
-    
+    length(vs_syms) + length(ps_syms) == length(unique(vcat(vs_syms, ps_syms))) || error("All symbols for variables & parameters must be unique.")
+
     FunctionalAffect(f, vs, vs_syms, ps, ps_syms, ctx)
 end
 
@@ -405,20 +404,16 @@ function compile_user_affect(affect::FunctionalAffect, sys, dvs, ps; kwargs...)
     # (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 = Tuple([x[2] 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 = Tuple([x[2] for x in p])
-
-    let v_inds=v_inds, p_inds=p_inds, v_syms=v_syms, p_syms=p_syms, user_affect=func(affect), ctx = context(affect)
-        function (integ)
-            uv = @views integ.u[v_inds]
-            pv = @views integ.p[p_inds]
+    p_syms = [x[2] for x in p]
 
-            u = LArray{v_syms}(uv)
-            p = LArray{p_syms}(pv)
+    kwargs = zip(vcat(v_syms, p_syms), vcat(v_inds, p_inds))
 
-            user_affect(integ.t, u, p, ctx)
+    let kwargs=kwargs, user_affect=func(affect), ctx = context(affect)
+        function (integ)
+            user_affect(integ, ctx; kwargs...)
         end
     end
 end

test/funcaffect.jl

@@ -0,0 +1,33 @@
+using ModelingToolkit, Test, DifferentialEquations
+
+@parameters t a b
+@variables u(t)
+D = Differential(t)
+
+eqs = [ D(u) ~ -u ]
+
+affect1!(integ, ctx; u) = integ.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))
+sol = solve(prob, Tsit5())
+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]
+    ctx[1] *= 2
+end
+ctx1 = [10.0, ]
+@named sys = ODESystem(eqs, t, [u], [], discrete_events=[[4.0, 8.0]=>(affect2!, [u], [], ctx1)])
+prob = ODEProblem(sys, [u=> 10.0], (0, 10.0))
+sol = solve(prob, Tsit5())
+i4 = findfirst(==(4.0), sol[:t])
+@test sol.u[i4+1][1] > 10.0
+i8 = findfirst(==(8.0), sol[:t])
+@test sol.u[i8+1][1] > 20.0
+@test ctx1[1] == 40.0
+
+
+