Merge pull request #2591 from SciML/initialization_equations

Add initialization_equations flattening constructor

Author: ChrisRackauckas

Project.toml

@@ -93,7 +93,7 @@ LinearAlgebra = "1"
 MLStyle = "0.4.17"
 NaNMath = "0.3, 1"
 OrderedCollections = "1"
-OrdinaryDiffEq = "6.73.0"
+OrdinaryDiffEq = "6.82.0"
 PrecompileTools = "1"
 RecursiveArrayTools = "2.3, 3"
 Reexport = "0.2, 1"

src/systems/abstractsystem.jl

@@ -889,6 +889,13 @@ function namespace_equations(sys::AbstractSystem, ivs = independent_variables(sy
     map(eq -> namespace_equation(eq, sys; ivs), eqs)
 end
 
+function namespace_initialization_equations(
+        sys::AbstractSystem, ivs = independent_variables(sys))
+    eqs = initialization_equations(sys)
+    isempty(eqs) && return Equation[]
+    map(eq -> namespace_equation(eq, sys; ivs), eqs)
+end
+
 function namespace_equation(eq::Equation,
         sys,
         n = nameof(sys);
@@ -1080,6 +1087,20 @@ function equations(sys::AbstractSystem)
     end
 end
 
+function initialization_equations(sys::AbstractSystem)
+    eqs = get_initialization_eqs(sys)
+    systems = get_systems(sys)
+    if isempty(systems)
+        return eqs
+    else
+        eqs = Equation[eqs;
+                       reduce(vcat,
+                           namespace_initialization_equations.(get_systems(sys));
+                           init = Equation[])]
+        return eqs
+    end
+end
+
 function preface(sys::AbstractSystem)
     has_preface(sys) || return nothing
     pre = get_preface(sys)

src/systems/diffeqs/abstractodesystem.jl

@@ -908,10 +908,10 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     # TODO: make it work with clocks
     # ModelingToolkit.get_tearing_state(sys) !== nothing => Requires structural_simplify first
     if sys isa ODESystem && build_initializeprob &&
-       (implicit_dae || !isempty(missingvars)) &&
-       all(isequal(Continuous()), ci.var_domain) &&
-       ModelingToolkit.get_tearing_state(sys) !== nothing &&
-       t !== nothing
+       (((implicit_dae || !isempty(missingvars)) &&
+         all(isequal(Continuous()), ci.var_domain) &&
+         ModelingToolkit.get_tearing_state(sys) !== nothing) ||
+        !isempty(initialization_equations(sys))) && t !== nothing
         if eltype(u0map) <: Number
             u0map = unknowns(sys) .=> u0map
         end

src/systems/diffeqs/odesystem.jl

@@ -362,6 +362,7 @@ function flatten(sys::ODESystem, noeqs = false)
             discrete_events = discrete_events(sys),
             defaults = defaults(sys),
             name = nameof(sys),
+            initialization_eqs = initialization_equations(sys),
             checks = false)
     end
 end

src/systems/nonlinear/initializesystem.jl

@@ -23,6 +23,7 @@ function generate_initializesystem(sys::ODESystem;
     diffmap = Dict(getfield.(eqs_diff, :lhs) .=> getfield.(eqs_diff, :rhs))
     observed_diffmap = Dict(Differential(get_iv(sys)).(getfield.((observed(sys)), :lhs)) .=>
         Differential(get_iv(sys)).(getfield.((observed(sys)), :rhs)))
+    full_diffmap = merge(diffmap, observed_diffmap)
 
     full_states = unique([sts; getfield.((observed(sys)), :lhs)])
     set_full_states = Set(full_states)
@@ -39,8 +40,7 @@ function generate_initializesystem(sys::ODESystem;
             filtered_u0 = Pair[]
             for x in u0map
                 y = get(schedule.dummy_sub, x[1], x[1])
-                y = ModelingToolkit.fixpoint_sub(y, observed_diffmap)
-                y = get(diffmap, y, y)
+                y = ModelingToolkit.fixpoint_sub(y, full_diffmap)
 
                 if y isa Symbolics.Arr
                     _y = collect(y)

test/hierarchical_initialization_eqs.jl

@@ -0,0 +1,147 @@
+using ModelingToolkit, OrdinaryDiffEq
+
+t = only(@variables(t))
+D = Differential(t)
+"""
+A simple linear resistor model
+
+![Resistor](https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTpJkiEyqh-BRx27pvVH0GLZ4MP_D1oriBwJhnZdgIq7m17z9VKUWaW9MeNQAz1rTML2ho&usqp=CAU)
+"""
+@component function Resistor(; name, R = 1.0)
+    systems = @named begin
+        p = Pin()
+        n = Pin()
+    end
+    vars = @variables begin
+        v(t), [guess = 0.0]
+        i(t), [guess = 0.0]
+    end
+    params = @parameters begin
+        R = R, [description = "Resistance of this Resistor"]
+    end
+    eqs = [v ~ p.v - n.v
+           i ~ p.i
+           p.i + n.i ~ 0
+           # Ohm's Law
+           v ~ i * R]
+    return ODESystem(eqs, t, vars, params; systems, name)
+end
+@connector Pin begin
+    v(t)
+    i(t), [connect = Flow]
+end
+@component function ConstantVoltage(; name, V = 1.0)
+    systems = @named begin
+        p = Pin()
+        n = Pin()
+    end
+    vars = @variables begin
+        v(t), [guess = 0.0]
+        i(t), [guess = 0.0]
+    end
+    params = @parameters begin
+        V = 10
+    end
+    eqs = [v ~ p.v - n.v
+           i ~ p.i
+           p.i + n.i ~ 0
+           v ~ V]
+    return ODESystem(eqs, t, vars, params; systems, name)
+end
+
+@component function Capacitor(; name, C = 1.0)
+    systems = @named begin
+        p = Pin()
+        n = Pin()
+    end
+    vars = @variables begin
+        v(t), [guess = 0.0]
+        i(t), [guess = 0.0]
+    end
+    params = @parameters begin
+        C = C
+    end
+    initialization_eqs = [
+        v ~ 0
+    ]
+    eqs = [v ~ p.v - n.v
+           i ~ p.i
+           p.i + n.i ~ 0
+           C * D(v) ~ i]
+    return ODESystem(eqs, t, vars, params; systems, name, initialization_eqs)
+end
+
+@component function Ground(; name)
+    systems = @named begin
+        g = Pin()
+    end
+    eqs = [
+        g.v ~ 0
+    ]
+    return ODESystem(eqs, t, [], []; systems, name)
+end
+
+@component function Inductor(; name, L = 1.0)
+    systems = @named begin
+        p = Pin()
+        n = Pin()
+    end
+    vars = @variables begin
+        v(t), [guess = 0.0]
+        i(t), [guess = 0.0]
+    end
+    params = @parameters begin
+        (L = L)
+    end
+    eqs = [v ~ p.v - n.v
+           i ~ p.i
+           p.i + n.i ~ 0
+           L * D(i) ~ v]
+    return ODESystem(eqs, t, vars, params; systems, name)
+end
+
+"""
+This is an RLC model.  This should support markdown.  That includes
+HTML as well.
+"""
+@component function RLCModel(; name)
+    systems = @named begin
+        resistor = Resistor(R = 100)
+        capacitor = Capacitor(C = 0.001)
+        inductor = Inductor(L = 1)
+        source = ConstantVoltage(V = 30)
+        ground = Ground()
+    end
+    initialization_eqs = [
+        inductor.i ~ 0
+    ]
+    eqs = [connect(source.p, inductor.n)
+           connect(inductor.p, resistor.p, capacitor.p)
+           connect(resistor.n, ground.g, capacitor.n, source.n)]
+    return ODESystem(eqs, t, [], []; systems, name, initialization_eqs)
+end
+"""Run model RLCModel from 0 to 10"""
+function simple()
+    @mtkbuild model = RLCModel()
+    u0 = []
+    prob = ODEProblem(model, u0, (0.0, 10.0))
+    sol = solve(prob)
+end
+@test SciMLBase.successful_retcode(simple())
+
+@named model = RLCModel()
+@test length(ModelingToolkit.get_initialization_eqs(model)) == 1
+syslist = ModelingToolkit.get_systems(model)
+@test length(ModelingToolkit.get_initialization_eqs(syslist[1])) == 0
+@test length(ModelingToolkit.get_initialization_eqs(syslist[2])) == 1
+@test length(ModelingToolkit.get_initialization_eqs(syslist[3])) == 0
+@test length(ModelingToolkit.get_initialization_eqs(syslist[4])) == 0
+@test length(ModelingToolkit.get_initialization_eqs(syslist[5])) == 0
+@test length(ModelingToolkit.initialization_equations(model)) == 2
+
+u0 = []
+prob = ODEProblem(structural_simplify(model), u0, (0.0, 10.0))
+sol = solve(prob, Rodas5P())
+@test length(sol[end]) == 2
+@test length(equations(prob.f.initializeprob.f.sys)) == 0
+@test length(unknowns(prob.f.initializeprob.f.sys)) == 0

test/runtests.jl

@@ -38,6 +38,7 @@ end
             @safetestset "DDESystem Test" include("dde.jl")
             @safetestset "NonlinearSystem Test" include("nonlinearsystem.jl")
             @safetestset "InitializationSystem Test" include("initializationsystem.jl")
+            @safetestset "Hierarchical Initialization Equations" include("hierarchical_initialization_eqs.jl")
             @safetestset "PDE Construction Test" include("pde.jl")
             @safetestset "JumpSystem Test" include("jumpsystem.jl")
             @safetestset "Constraints Test" include("constraints.jl")