Fix state extraction

Co-authored-by: "Yingbo Ma" <[email protected]>

Author: shashi

src/systems/abstractsystem.jl

@@ -157,7 +157,7 @@ end
 renamespace(namespace,name) = Symbol(namespace,:₊,name)
 
 function namespace_variables(sys::AbstractSystem)
-    [rename(x,renamespace(sys.name,x.name)) for x in states(sys)]
+    [Sym{symtype(x.op)}(renamespace(sys.name,x.op.name))(x.args...) for x in states(sys)]
 end
 
 function namespace_parameters(sys::AbstractSystem)
@@ -190,7 +190,11 @@ end
 namespace_expr(O,name,ivname) = O
 
 independent_variable(sys::AbstractSystem) = sys.iv
-states(sys::AbstractSystem) = unique(isempty(sys.systems) ? setdiff(sys.states, convert.(Variable,sys.pins)) : [sys.states;reduce(vcat,namespace_variables.(sys.systems))])
+function states(sys::AbstractSystem)
+    unique(isempty(sys.systems) ?
+           setdiff(sys.states, value.(sys.pins)) :
+           [sys.states;reduce(vcat,namespace_variables.(sys.systems))])
+end
 parameters(sys::AbstractSystem) = isempty(sys.systems) ? sys.ps : [sys.ps;reduce(vcat,namespace_parameters.(sys.systems))]
 pins(sys::AbstractSystem) = isempty(sys.systems) ? sys.pins : [sys.pins;reduce(vcat,namespace_pins.(sys.systems))]
 function observed(sys::AbstractSystem)

src/systems/diffeqs/abstractodesystem.jl

@@ -271,8 +271,14 @@ function DiffEqBase.ODEProblem{iip}(sys::AbstractODESystem,u0map,tspan,
                                     kwargs...) where iip
     dvs = states(sys)
     ps = parameters(sys)
-    u0 = varmap_to_vars(u0map,dvs)
-    p = varmap_to_vars(parammap,ps)
+    u0map′ = [lower_varname(value(k), sys.iv) => value(v) for (k, v) in u0map]
+    parammap′ = [value(k) => value(v) for (k, v) in parammap]
+    u0 = varmap_to_vars(u0map′,dvs)
+    if !(parammap isa DiffEqBase.NullParameters)
+        p = varmap_to_vars(parammap′,ps)
+    else
+        p = ps
+    end
     f = ODEFunction{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,checkbounds=checkbounds,
                         linenumbers=linenumbers,parallel=parallel,simplify=simplify,
                         sparse=sparse,eval_expression=eval_expression,kwargs...)
@@ -306,10 +312,17 @@ function ODEProblemExpr{iip}(sys::AbstractODESystem,u0map,tspan,
                                     simplify = true,
                                     linenumbers = false, parallel=SerialForm(),
                                     kwargs...) where iip
+
     dvs = states(sys)
     ps = parameters(sys)
-    u0 = varmap_to_vars(u0map,dvs)
-    p = varmap_to_vars(parammap,ps)
+    u0map′ = [lower_varname(value(k), sys.iv) => value(v) for (k, v) in u0map]
+    parammap′ = [value(k) => value(v) for (k, v) in parammap]
+    u0 = varmap_to_vars(u0map′,dvs)
+    if !(parammap isa DiffEqBase.NullParameters)
+        p = varmap_to_vars(parammap′,ps)
+    else
+        p = ps
+    end
     f = ODEFunctionExpr{iip}(sys,dvs,ps,u0;tgrad=tgrad,jac=jac,checkbounds=checkbounds,
                         linenumbers=linenumbers,parallel=parallel,
                         simplify=simplify,

src/systems/diffeqs/first_order_transform.jl

@@ -5,6 +5,12 @@ function lower_varname(var::Term, idv, order)
     return Sym{symtype(var.op)}(name)(var.args[1])
 end
 
+function lower_varname(t::Term, iv)
+    var, order = var_from_nested_derivative(t)
+    lower_varname(var, iv, order)
+end
+lower_varname(t::Sym, iv) = t
+
 function flatten_differential(O::Term)
     @assert is_derivative(O) "invalid differential: $O"
     is_derivative(O.args[1]) || return (O.args[1], O.op.x, 1)

test/runtests.jl

@@ -7,7 +7,7 @@ using SafeTestsets, Test
 @safetestset "Direct Usage Test" begin include("direct.jl") end
 @safetestset "System Linearity Test" begin include("linearity.jl") end
 @safetestset "Build Function Test" begin include("build_function.jl") end
-#@safetestset "ODESystem Test" begin include("odesystem.jl") end
+@safetestset "ODESystem Test" begin include("odesystem.jl") end
 @safetestset "LabelledArrays Test" begin include("labelledarrays.jl") end
 @safetestset "Mass Matrix Test" begin include("mass_matrix.jl") end
 @safetestset "SteadyStateSystem Test" begin include("steadystatesystems.jl") end