Merge pull request #837 from SciML/myb/obsfix

Handle dependencies of observed variables properly

Author: YingboMa

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "5.9.0"
+version = "5.9.1"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

docs/src/tutorials/acausal_components.md

@@ -12,7 +12,7 @@ equalities before solving. Let's see this in action.
 ## Copy-Paste Example
 
 ```julia
-using ModelingToolkit, Plots
+using ModelingToolkit, Plots, DifferentialEquations
 
 @parameters t
 
@@ -23,15 +23,15 @@ function Pin(;name)
 end
 
 function Ground(;name)
-    g = Pin(:g)
+    @named g = Pin()
     eqs = [g.v ~ 0]
     ODESystem(eqs, t, [], [], systems=[g], name=name)
 end
 
 function Resistor(;name, R = 1.0)
     val = R
-    p = Pin(:p)
-    n = Pin(:n)
+    @named p = Pin()
+    @named n = Pin()
     @variables v(t)
     @parameters R
     eqs = [
@@ -44,8 +44,8 @@ end
 
 function Capacitor(; name, C = 1.0)
     val = C
-    p = Pin(:p)
-    n = Pin(:n)
+    @named p = Pin()
+    @named n = Pin()
     @variables v(t)
     @parameters C
     D = Differential(t)
@@ -59,8 +59,8 @@ end
 
 function ConstantVoltage(;name, V = 1.0)
     val = V
-    p = Pin(:p)
-    n = Pin(:n)
+    @named p = Pin()
+    @named n = Pin()
     @parameters V
     eqs = [
            V ~ p.v - n.v
@@ -135,7 +135,7 @@ names to correspond to duplicates of this topology with unique variables.
 One can then construct a `Pin` like:
 
 ```julia
-Pin(:mypin1)
+Pin(name=:mypin1)
 ```
 
 or equivalently using the `@named` helper macro:
@@ -151,7 +151,7 @@ that the voltage in such a `Pin` is equal to zero. This gives:
 
 ```julia
 function Ground(;name)
-    g = Pin(:g)
+    @named g = Pin()
     eqs = [g.v ~ 0]
     ODESystem(eqs, t, [], [], systems=[g], name=name)
 end
@@ -167,8 +167,8 @@ zero. This leads to our resistor equations:
 ```julia
 function Resistor(;name, R = 1.0)
     val = R
-    p = Pin(:p)
-    n = Pin(:n)
+    @named p = Pin()
+    @named n = Pin()
     @variables v(t)
     @parameters R
     eqs = [
@@ -190,8 +190,8 @@ Using our knowledge of circuits we similarly construct the Capacitor:
 ```julia
 function Capacitor(; name, C = 1.0)
     val = C
-    p = Pin(:p)
-    n = Pin(:n)
+    @named p = Pin()
+    @named n = Pin()
     @variables v(t)
     @parameters C
     D = Differential(t)
@@ -212,8 +212,8 @@ model this as:
 ```julia
 function ConstantVoltage(;name, V = 1.0)
     val = V
-    p = Pin(:p)
-    n = Pin(:n)
+    @named p = Pin()
+    @named n = Pin()
     @parameters V
     eqs = [
            V ~ p.v - n.v
@@ -406,10 +406,15 @@ observed(sys)
  capacitor₊n₊v(t) ~ 0.0
  source₊n₊v(t) ~ 0.0
  ground₊g₊i(t) ~ 0.0
- ⋮
+ source₊n₊i(t) ~ capacitor₊p₊i(t)
+ source₊p₊i(t) ~ -capacitor₊p₊i(t)
+ capacitor₊n₊i(t) ~ -capacitor₊p₊i(t)
+ resistor₊n₊i(t) ~ -capacitor₊p₊i(t)
  ground₊g₊v(t) ~ 0.0
  source₊p₊v(t) ~ source₊V
  capacitor₊p₊v(t) ~ capacitor₊v(t)
+ resistor₊p₊v(t) ~ source₊p₊v(t)
+ resistor₊n₊v(t) ~ capacitor₊p₊v(t)
  resistor₊v(t) ~ -((capacitor₊p₊v(t)) - (source₊p₊v(t)))
 ```
 
@@ -428,5 +433,5 @@ sol[resistor.v]
 or we can plot the timeseries of the resistor's voltage:
 
 ```julia
-plot(sol,vars=(resistor.v,))
+plot(sol, vars=[resistor.v])
 ```

src/structural_transformation/StructuralTransformations.jl

@@ -16,7 +16,7 @@ using ModelingToolkit: ODESystem, var_from_nested_derivative, Differential,
                        states, equations, vars, Symbolic, diff2term, value,
                        operation, arguments, Sym, Term, simplify, solve_for,
                        isdiffeq, isdifferential,
-                       get_structure, default_u0, default_p
+                       get_structure, get_reduced_states, default_u0, default_p
 
 using ModelingToolkit.BipartiteGraphs
 using ModelingToolkit.SystemStructures

src/structural_transformation/codegen.jl

@@ -272,9 +272,14 @@ function build_observed_function(
     required_algvars = Set(intersect(algvars, syms_set))
     obs = observed(sys)
     observed_idx = Dict(map(x->x.lhs, obs) .=> 1:length(obs))
-    for sym in syms
-        idx = get(observed_idx, sym, nothing)
+    # FIXME: this is a rather rough estimate of dependencies.
+    maxidx = 0
+    for (i, s) in enumerate(syms)
+        idx = get(observed_idx, s, nothing)
         idx === nothing && continue
+        idx > maxidx && (maxidx = idx)
+    end
+    for idx in 1:maxidx
         vs = vars(obs[idx].rhs)
         union!(required_algvars, intersect(algvars, vs))
     end
@@ -308,7 +313,10 @@ function build_observed_function(
         ],
         [],
         Let(
-            collect(Iterators.flatten(solves)),
+            [
+             collect(Iterators.flatten(solves))
+             map(eq -> eq.lhs←eq.rhs, obs[1:maxidx])
+            ],
             isscalar ? output[1] : MakeArray(output, output_type)
            )
     ) |> Code.toexpr

src/structural_transformation/tearing.jl

@@ -180,6 +180,7 @@ function tearing_reassemble(sys; simplify=false)
     @set! sys.structure.algeqs = newalgeqs
     @set! sys.eqs = neweqs
     @set! sys.states = newstates
+    @set! sys.reduced_states = [get_reduced_states(sys); solvars]
     @set! sys.observed = vcat(observed(sys), obseqs)
     return sys
 end

src/systems/abstractsystem.jl

@@ -154,6 +154,7 @@ for prop in [
              :inequality_constraints
              :controls
              :loss
+             :reduced_states
             ]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
@@ -467,6 +468,13 @@ end
 """
 $(SIGNATURES)
 
-Structurally simplify algebraic equations in a system.
+Structurally simplify algebraic equations in a system and compute the
+topological sort of the observed equations.
 """
-structural_simplify(sys::AbstractSystem) = tearing(alias_elimination(sys))
+function structural_simplify(sys::AbstractSystem)
+    sys = tearing(alias_elimination(sys))
+    s = structure(sys)
+    fullstates = [get_reduced_states(sys); states(sys)]
+    @set! sys.observed = topsort_equations(observed(sys), fullstates)
+    return sys
+end

src/systems/alias_elimination.jl

@@ -3,12 +3,13 @@ using SymbolicUtils: Rewriters
 const KEEP = typemin(Int)
 
 function alias_elimination(sys)
-    sys = flatten(sys)
-    s = get_structure(sys)
-    if !(s isa SystemStructure)
-        sys = initialize_system_structure(sys)
-        s = structure(sys)
-    end
+    # FIXME: update `structure` too
+    #sys = flatten(sys)
+    #s = get_structure(sys)
+    #if !(s isa SystemStructure)
+    sys = initialize_system_structure(sys)
+    s = structure(sys)
+    #end
     is_linear_equations, eadj, cadj = find_linear_equations(sys)
 
     v_eliminated, v_types, n_null_vars, degenerate_equations, linear_equations = alias_eliminate_graph(
@@ -18,9 +19,12 @@ function alias_elimination(sys)
     s = structure(sys)
     @unpack fullvars, graph = s
 
+    n_reduced_states = length(v_eliminated) - n_null_vars
+    reduced_states = similar(v_eliminated, Any, n_reduced_states)
     subs = OrderedDict()
-    if length(v_eliminated) - n_null_vars > 0
-        for v in v_eliminated[n_null_vars+1:end]
+    if n_reduced_states > 0
+        for (i, v) in enumerate(@view v_eliminated[n_null_vars+1:end])
+            reduced_states[i] = fullvars[v]
             subs[fullvars[v]] = iszeroterm(v_types, v) ? 0.0 :
                                 isalias(v_types, v) ? fullvars[alias(v_types, v)] :
                                 -fullvars[negalias(v_types, v)]
@@ -63,6 +67,7 @@ function alias_elimination(sys)
 
     @set! sys.eqs = eqs
     @set! sys.states = newstates
+    @set! sys.reduced_states = [get_reduced_states(sys); reduced_states]
     @set! sys.observed = [get_observed(sys); [lhs ~ rhs for (lhs, rhs) in pairs(subs)]]
     @set! sys.structure = nothing
     return sys

src/systems/diffeqs/odesystem.jl

@@ -74,6 +74,7 @@ struct ODESystem <: AbstractODESystem
     structure: structural information of the system
     """
     structure::Any
+    reduced_states::Vector
 end
 
 function ODESystem(
@@ -101,7 +102,7 @@ function ODESystem(
     if length(unique(sysnames)) != length(sysnames)
         throw(ArgumentError("System names must be unique."))
     end
-    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, default_u0, default_p, nothing)
+    ODESystem(deqs, iv′, dvs′, ps′, observed, tgrad, jac, Wfact, Wfact_t, name, systems, default_u0, default_p, nothing, [])
 end
 
 var_from_nested_derivative(x, i=0) = (missing, missing)
@@ -235,7 +236,7 @@ ODESystem(eq::Equation, args...; kwargs...) = ODESystem([eq], args...; kwargs...
 $(SIGNATURES)
 
 Build the observed function assuming the observed equations are all explicit,
-i.e. there are no cycles or dependencies.
+i.e. there are no cycles.
 """
 function build_explicit_observed_function(
         sys, syms;
@@ -249,7 +250,14 @@ function build_explicit_observed_function(
 
     obs = observed(sys)
     observed_idx = Dict(map(x->x.lhs, obs) .=> 1:length(obs))
-    output = map(sym->obs[observed_idx[sym]].rhs, syms)
+    output = similar(syms, Any)
+    # FIXME: this is a rather rough estimate of dependencies.
+    maxidx = 0
+    for (i, s) in enumerate(syms)
+        idx = observed_idx[s]
+        idx > maxidx && (maxidx = idx)
+        output[i] = obs[idx].rhs
+    end
 
     ex = Func(
         [
@@ -258,7 +266,10 @@ function build_explicit_observed_function(
          independent_variable(sys)
         ],
         [],
-        isscalar ? output[1] : MakeArray(output, output_type)
+        Let(
+            map(eq -> eq.lhs←eq.rhs, obs[1:maxidx]),
+            isscalar ? output[1] : MakeArray(output, output_type)
+           )
     ) |> toexpr
 
     expression ? ex : @RuntimeGeneratedFunction(ex)

src/systems/nonlinear/nonlinearsystem.jl

@@ -48,6 +48,7 @@ struct NonlinearSystem <: AbstractSystem
     structure: structural information of the system
     """
     structure::Any
+    reduced_states::Any
 end
 
 function NonlinearSystem(eqs, states, ps;
@@ -60,7 +61,7 @@ function NonlinearSystem(eqs, states, ps;
     default_p isa Dict || (default_p = Dict(default_p))
     default_u0 = Dict(value(k) => value(default_u0[k]) for k in keys(default_u0))
     default_p = Dict(value(k) => value(default_p[k]) for k in keys(default_p))
-    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, default_u0, default_p, nothing)
+    NonlinearSystem(eqs, value.(states), value.(ps), observed, name, systems, default_u0, default_p, nothing, [])
 end
 
 function calculate_jacobian(sys::NonlinearSystem;sparse=false,simplify=false)

test/components.jl

@@ -100,14 +100,16 @@ sol = solve(prob, Rodas4())
 @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]
 
 prob = ODAEProblem(sys, u0, (0, 10.0))
-sol = solve(prob, Rodas4())
+sol = solve(prob, Tsit5())
 
 @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]
 #using Plots
 #plot(sol)