add observed to var_to_name Dict (#1343)

Author: isaacsas

src/systems/diffeqs/odesystem.jl

@@ -141,6 +141,7 @@ function ODESystem(
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, dvs′)
     process_variables!(var_to_name, defaults, ps′)
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     tgrad = RefValue(Vector{Num}(undef, 0))
     jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))

src/systems/diffeqs/sdesystem.jl

@@ -128,6 +128,7 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, dvs′)
     process_variables!(var_to_name, defaults, ps′)
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     tgrad = RefValue(Vector{Num}(undef, 0))
     jac = RefValue{Any}(Matrix{Num}(undef, 0, 0))

src/systems/discrete_system/discrete_system.jl

@@ -98,7 +98,8 @@ function DiscreteSystem(
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, dvs′)
     process_variables!(var_to_name, defaults, ps′)
-
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
+    
     sysnames = nameof.(systems)
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))

src/systems/jumps/jumpsystem.jl

@@ -101,6 +101,7 @@ function JumpSystem(eqs, iv, states, ps;
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, states)
     process_variables!(var_to_name, defaults, ps)
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     JumpSystem{typeof(ap)}(ap, value(iv), states, ps, var_to_name, observed, name, systems, defaults, connector_type, checks = checks)
 end

src/systems/nonlinear/nonlinearsystem.jl

@@ -98,6 +98,7 @@ function NonlinearSystem(eqs, states, ps;
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, states)
     process_variables!(var_to_name, defaults, ps)
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     NonlinearSystem(eqs, states, ps, var_to_name, observed, jac, name, systems, defaults, nothing, connector_type, checks = checks)
 end

src/systems/optimization/optimizationsystem.jl

@@ -77,6 +77,7 @@ function OptimizationSystem(op, states, ps;
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, states)
     process_variables!(var_to_name, defaults, ps)
+    isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     OptimizationSystem(
                        value(op), states, ps, var_to_name,

test/discretesystem.jl

@@ -112,3 +112,12 @@ linearized_eqs = [
     y(t - 2.0) ~ y(t)
 ]
 @test all(eqs2 .== linearized_eqs)
+
+# observed variable handling
+@variables t x(t) RHS(t)
+@parameters τ   
+@named fol = DiscreteSystem([D(x) ~ (1 - x)/τ]; observed=[RHS ~ (1 - x)/τ])
+@test isequal(RHS, @nonamespace fol.RHS)
+RHS2 = RHS
+@unpack RHS = fol
+@test isequal(RHS, RHS2)

test/jumpsystem.jl

@@ -175,3 +175,11 @@ function paffect!(integrator)
 end
 sol = solve(jprob, SSAStepper(), tstops=[1000.0], callback=DiscreteCallback(pcondit,paffect!))
 @test sol[1,end] == 100
+
+# observed variable handling
+@variables OBS(t)
+@named js5 = JumpSystem([maj1,maj2], t, [S], [β,γ]; observed=[OBS ~ 2*S])
+OBS2 = OBS
+@test isequal(OBS2, @nonamespace js5.OBS)
+@unpack OBS = js5
+@test isequal(OBS2, OBS)

test/nonlinearsystem.jl

@@ -156,3 +156,12 @@ end
     @test isequal(union(Set(states(sys1)), Set(states(sys2))), Set(states(sys3)))
     @test isequal(union(Set(equations(sys1)), Set(equations(sys2))), Set(equations(sys3)))
 end
+
+# observed variable handling
+@variables t x(t) RHS(t)
+@parameters τ   
+@named fol = NonlinearSystem([0  ~ (1 - x)/τ], [x], [τ]; observed=[RHS ~ (1 - x)/τ])
+@test isequal(RHS, @nonamespace fol.RHS)
+RHS2 = RHS
+@unpack RHS = fol
+@test isequal(RHS, RHS2)

test/odesystem.jl

@@ -504,3 +504,14 @@ prob = ODEProblem(outersys, [sys.x=>1.0; collect(sys.ms).=>1:3], (0, 1.0))
 @variables t x(t)
 @named sys = ODESystem([D(x) ~ x/x], t)
 @test equations(alias_elimination(sys)) == [D(x) ~ 1]
+
+# observed variable handling
+@variables t x(t) RHS(t)
+@parameters τ   
+D = Differential(t) 
+@named fol = ODESystem([D(x)  ~ (1 - x)/τ]; observed=[RHS ~ (1 - x)/τ])
+@test isequal(RHS, @nonamespace fol.RHS)
+RHS2 = RHS
+@unpack RHS = fol
+@test isequal(RHS, RHS2)
+