add outputidxs

Author: isaacsas

src/build_function.jl

@@ -129,7 +129,8 @@ end
 function _build_function(target::JuliaTarget, rhss, args...;
                          conv = simplified_expr, expression = Val{true},
                          checkbounds = false, constructor=nothing,
-                         linenumbers = false, multithread=false, headerfun=addheader)
+                         linenumbers = false, multithread=false, 
+                         headerfun=addheader, outputidxs=nothing)
     argnames = [gensym(:MTKArg) for i in 1:length(args)]
     arg_pairs = map(vars_to_pairs,zip(argnames,args))
     ls = reduce(vcat,first.(arg_pairs))
@@ -139,6 +140,8 @@ function _build_function(target::JuliaTarget, rhss, args...;
     fname = gensym(:ModelingToolkitFunction)
     fargs = Expr(:tuple,argnames...)
 
+
+    oidx = isnothing(outputidxs) ? (i -> i) : (i -> outputidxs[i])
     X = gensym(:MTIIPVar)
 	if eltype(eltype(rhss)) <: AbstractArray # Array of arrays of arrays
 		ip_sys_exprs = reduce(vcat,[vec(reduce(vcat,[vec([:($X[$i][$j][$k] = $(conv(rhs))) for (k, rhs) ∈ enumerate(rhsel2)]) for (j, rhsel2) ∈ enumerate(rhsel)],init=Expr[])) for (i,rhsel) ∈ enumerate(rhss)],init=Expr[])
@@ -151,7 +154,7 @@ function _build_function(target::JuliaTarget, rhss, args...;
     elseif rhss isa SparseMatrixCSC
         ip_sys_exprs = [:($X.nzval[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss.nzval)]
     else
-        ip_sys_exprs = [:($X[$i] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
+        ip_sys_exprs = [:($X[$(oidx(i))] = $(conv(rhs))) for (i, rhs) ∈ enumerate(rhss)]
     end
 
     ip_let_expr = Expr(:let, var_eqs, build_expr(:block, ip_sys_exprs))

src/systems/jumps/jumpsystem.jl

@@ -20,20 +20,25 @@ generate_rate_function(js, rate) = build_function(rate, states(js), parameters(j
                                                   independent_variable(js), 
                                                   expression=Val{false})
 
-generate_affect_function(js, affect) = build_function(affect, states(js), 
+generate_affect_function(js, affect, outputidxs) = build_function(affect, states(js), 
                                                       parameters(js), 
                                                       independent_variable(js),
                                                       expression=Val{false},
-                                                      headerfun=add_integrator_header)[2]
-function assemble_vrj(js, vrj) 
+                                                      headerfun=add_integrator_header, 
+                                                      outputidxs=outputidxs)[2]
+function assemble_vrj(js, vrj, statetoid) 
     rate   = generate_rate_function(js, vrj.rate)
-    affect = generate_affect_function(js, vrj.affect!)
+    outputvars = (convert(Variable,affect.lhs) for affect in vrj.affect!)
+    outputidxs = ((statetoid[var] for var in outputvars)...,)
+    affect = generate_affect_function(js, vrj.affect!, outputidxs)
     VariableRateJump(rate, affect)
 end
 
-function assemble_crj(js, crj)
+function assemble_crj(js, crj, statetoid)
     rate   = generate_rate_function(js, crj.rate)
-    affect = generate_affect_function(js, crj.affect!)
+    outputvars = (convert(Variable,affect.lhs) for affect in crj.affect!)
+    outputidxs = ((statetoid[var] for var in outputvars)...,)
+    affect = generate_affect_function(js, crj.affect!, outputidxs)
     ConstantRateJump(rate, affect)
 end
 
@@ -47,16 +52,17 @@ Generates a JumpProblem from a JumpSystem.
 function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
     vrjs = Vector{VariableRateJump}()
     crjs = Vector{ConstantRateJump}()
+    statetoid = Dict(convert(Variable,state) => i for (i,state) in enumerate(states(js)))
     for j in equations(js)
         if j isa ConstantRateJump
-            push!(crjs, assemble_crj(js, j))
+            push!(crjs, assemble_crj(js, j, statetoid))
         elseif j isa VariableRateJump
-            push!(vrjs, assemble_vrj(js, j))
+            push!(vrjs, assemble_vrj(js, j, statetoid))
         else
             (j isa MassActionJump) && error("Generation of JumpProblems with MassActionJumps is not yet supported.")
         end
     end
     ((prob isa DiscreteProblem) && !isempty(vrjs)) && error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")    
-    jset = JumpSet(Tuple(vrjs...), Tuple(crjs...), nothing, nothing)
+    jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, nothing)
     JumpProblem(prob, aggregator, jset)
 end

test/jumpsystem.jl

@@ -11,8 +11,9 @@ affect₂ = [I ~ I - 1, R ~ R + 1]
 j₁      = ConstantRateJump(rate₁,affect₁)
 j₂      = VariableRateJump(rate₂,affect₂)
 js      = JumpSystem([j₁,j₂], t, [S,I,R], [β,γ])
-mtjump1 = MT.assemble_crj(js, j₁)
-mtjump2 = MT.assemble_vrj(js, j₂)
+statetoid = Dict(convert(Variable,state) => i for (i,state) in enumerate(states(js)))
+mtjump1 = MT.assemble_crj(js, j₁, statetoid)
+mtjump2 = MT.assemble_vrj(js, j₂, statetoid)
 
 # doc version
 rate1(u,p,t) = (0.1/1000.0)*u[1]*u[2]
@@ -58,22 +59,23 @@ u₀ = [999,1,0]; p = (0.1/1000,0.01); tspan = (0.,250.)
 dprob = DiscreteProblem(u₀,tspan,p)
 jprob = JumpProblem(js2, dprob, Direct())
 sol = solve(jprob, SSAStepper())
+plot(sol)
 
 # test the MT JumpProblem rates/affects are correct
-rate2(u,p,t) = 0.01u[2]
-jump2 = ConstantRateJump(rate2,affect2!)
-mtjumps = jprob.discrete_jump_aggregation
-@test abs(mtjumps.rates[1](u,p,tf) - jump1.rate(u,p,tf)) < 10*eps()
-@test abs(mtjumps.rates[2](u,p,tf) - jump2.rate(u,p,tf)) < 10*eps()
-mtjumps.affects![1](mtintegrator)
-jump1.affect!(integrator)
-@test all(integrator.u .== mtintegrator.u) 
-mtintegrator.u .= u; integrator.u .= u
-mtjumps.affects![2](mtintegrator)
-jump2.affect!(integrator)
-@test all(integrator.u .== mtintegrator.u)
+# rate2(u,p,t) = 0.01u[2]
+# jump2 = ConstantRateJump(rate2,affect2!)
+# mtjumps = jprob.discrete_jump_aggregation
+# @test abs(mtjumps.rates[1](u,p,tf) - jump1.rate(u,p,tf)) < 10*eps()
+# @test abs(mtjumps.rates[2](u,p,tf) - jump2.rate(u,p,tf)) < 10*eps()
+# mtjumps.affects![1](mtintegrator)
+# jump1.affect!(integrator)
+# @test all(integrator.u .== mtintegrator.u) 
+# mtintegrator.u .= u; integrator.u .= u
+# mtjumps.affects![2](mtintegrator)
+# jump2.affect!(integrator)
+# @test all(integrator.u .== mtintegrator.u)
 
-# # direct vers
+# # # direct vers
 # p = (0.1/1000,0.01)
 # prob = DiscreteProblem([999,1,0],(0.0,250.0),p)
 # r1(u,p,t) = (0.1/1000.0)*u[1]*u[2]
@@ -90,7 +92,7 @@ jump2.affect!(integrator)
 # j2 = ConstantRateJump(r2,a2!)
 # jset = JumpSet((),(j1,j2),nothing,nothing)
 # jprob = JumpProblem(prob,Direct(),jset)
-# sol = solve(jprob, SSAStepper())
+# sol2 = solve(jprob, SSAStepper())
 
 # using Plots
 # plot(sol)