Fix out-of-bounds state indexing in code gen

Author: YingboMa

src/systems/diffeqs/abstractodesystem.jl

@@ -68,16 +68,19 @@ function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = param
 end
 
 function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
-    dvs = vcat(dvs, map(eq->eq.lhs, observed(sys)))
     # optimization
-    dvs′ = makesym.(value.(dvs), states=dvs)
-    ps′ = makesym.(value.(ps), states=dvs)
+    obsvars = map(eq->eq.lhs, observed(sys))
+    fulldvs = [dvs; obsvars]
+    fulldvs′ = makesym.(value.(fulldvs))
 
-    sub = Dict(dvs .=> dvs′)
+    sub = Dict(fulldvs .=> fulldvs′)
     # substitute x(t) by just x
     rhss = [substitute(deq.rhs, sub) for deq ∈ equations(sys)]
     obss = [makesym(value(eq.lhs)) ~ substitute(eq.rhs, sub) for eq ∈ observed(sys)]
 
+    dvs′ = fulldvs′[1:length(dvs)]
+    ps′ = makesym.(value.(ps), states=())
+
     # TODO: add an optional check on the ordering of observed equations
     return build_function(Let(obss, rhss), dvs′, ps′, sys.iv;
                           conv = ODEToExpr(sys),kwargs...)

src/utils.jl

@@ -222,7 +222,7 @@ Symbol("z⦗t⦘")
 """
 function tosymbol(t::Term; states=nothing, escape=true)
     if operation(t) isa Sym
-        if states !== nothing && !(any(isequal(t), states))
+        if states !== nothing && !(t in states)
             return nameof(operation(t))
         end
         op = nameof(operation(t))