Merge pull request #1489 from SciML/myb/obs

Make observed function building resilient wrt namespacing

Author: YingboMa

src/structural_transformation/codegen.jl

@@ -393,20 +393,40 @@ function build_observed_function(
 
     required_algvars = Set(intersect(algvars, vars))
     obs = observed(sys)
-    observed_idx = Dict(map(x->x.lhs, obs) .=> 1:length(obs))
-    # FIXME: this is a rather rough estimate of dependencies.
-    maxidx = 0
+    observed_idx = Dict(x.lhs => i for (i, x) in enumerate(obs))
+    namespaced_to_obs = Dict(states(sys, x.lhs) => x.lhs for x in obs)
+    namespaced_to_sts = Dict(states(sys, x) => x for x in states(sys))
     sts = Set(states(sys))
+
+    # FIXME: This is a rather rough estimate of dependencies. We assume
+    # the expression depends on everything before the `maxidx`.
+    subs = Dict()
+    maxidx = 0
     for (i, s) in enumerate(dep_vars)
         idx = get(observed_idx, s, nothing)
-        if idx === nothing
-            if !(s in sts)
+        if idx !== nothing
+            idx > maxidx && (maxidx = idx)
+        else
+            s′ = get(namespaced_to_obs, s, nothing)
+            if s′ !== nothing
+                subs[s] = s′
+                s = s′
+                idx = get(observed_idx, s, nothing)
+            end
+            if idx !== nothing
+                idx > maxidx && (maxidx = idx)
+            elseif !(s in sts)
+                s′ = get(namespaced_to_sts, s, nothing)
+                if s′ !== nothing
+                    subs[s] = s′
+                    continue
+                end
                 throw(ArgumentError("$s is either an observed nor a state variable."))
             end
             continue
         end
-        idx > maxidx && (maxidx = idx)
     end
+    ts = map(t->substitute(t, subs), ts)
     vs = Set()
     for idx in 1:maxidx
         vars!(vs, obs[idx].rhs)

src/systems/diffeqs/abstractodesystem.jl

@@ -357,7 +357,7 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
     end
 
     M = calculate_massmatrix(sys)
-    
+
     _M = if sparse && !(u0 === nothing || M === I)
       SparseArrays.sparse(M)
     elseif u0 === nothing || M === I

src/systems/diffeqs/odesystem.jl

@@ -279,21 +279,39 @@ function build_explicit_observed_function(
 
     obs = observed(sys)
     sts = Set(states(sys))
-    observed_idx = Dict(map(x->x.lhs, obs) .=> 1:length(obs))
+    observed_idx = Dict(x.lhs => i for (i, x) in enumerate(obs))
+    namespaced_to_obs = Dict(states(sys, x.lhs) => x.lhs for x in obs)
+    namespaced_to_sts = Dict(states(sys, x) => x for x in states(sys))
 
     # FIXME: This is a rather rough estimate of dependencies. We assume
     # the expression depends on everything before the `maxidx`.
+    subs = Dict()
     maxidx = 0
     for (i, s) in enumerate(dep_vars)
         idx = get(observed_idx, s, nothing)
-        if idx === nothing
-            if !(s in sts)
+        if idx !== nothing
+            idx > maxidx && (maxidx = idx)
+        else
+            s′ = get(namespaced_to_obs, s, nothing)
+            if s′ !== nothing
+                subs[s] = s′
+                s = s′
+                idx = get(observed_idx, s, nothing)
+            end
+            if idx !== nothing
+                idx > maxidx && (maxidx = idx)
+            elseif !(s in sts)
+                s′ = get(namespaced_to_sts, s, nothing)
+                if s′ !== nothing
+                    subs[s] = s′
+                    continue
+                end
                 throw(ArgumentError("$s is either an observed nor a state variable."))
             end
             continue
         end
-        idx > maxidx && (maxidx = idx)
     end
+    ts = map(t->substitute(t, subs), ts)
     obsexprs = map(eq -> eq.lhs←eq.rhs, obs[1:maxidx])
 
     dvs = DestructuredArgs(states(sys), inbounds=!checkbounds)

test/components.jl

@@ -19,6 +19,15 @@ function check_contract(sys)
     end
 end
 
+function check_rc_sol(sol)
+    @test sol[rc_model.resistor.p.i] == sol[resistor.p.i] == sol[capacitor.p.i]
+    @test sol[rc_model.resistor.n.i] == sol[resistor.n.i] == -sol[capacitor.p.i]
+    @test sol[rc_model.capacitor.n.i] ==sol[capacitor.n.i] == -sol[capacitor.p.i]
+    @test iszero(sol[rc_model.ground.g.i])
+    @test iszero(sol[rc_model.ground.g.v])
+    @test sol[rc_model.resistor.v] == sol[resistor.v] == sol[source.p.v] - sol[capacitor.p.v]
+end
+
 include("../examples/rc_model.jl")
 
 sys = structural_simplify(rc_model)
@@ -31,13 +40,10 @@ u0 = [
      ]
 prob = ODEProblem(sys, u0, (0, 10.0))
 sol = solve(prob, Rodas4())
-
-@test sol[resistor.p.i] == sol[capacitor.p.i]
-@test sol[resistor.n.i] == -sol[capacitor.p.i]
-@test sol[capacitor.n.i] == -sol[capacitor.p.i]
-@test iszero(sol[ground.g.i])
-@test iszero(sol[ground.g.v])
-@test sol[resistor.v] == sol[source.p.v] - sol[capacitor.p.v]
+check_rc_sol(sol)
+prob = ODAEProblem(sys, u0, (0, 10.0))
+sol = solve(prob, Rodas4())
+check_rc_sol(sol)
 
 # Outer/inner connections
 function rc_component(;name)