handle variables marked as disturbance when building controlled dynamics

Author: baggepinnen

src/inputoutput.jl

@@ -188,7 +188,8 @@ t = 0
 f[1](x, inputs, p, t)
 ```
 """
-function generate_control_function(sys::AbstractODESystem, inputs = unbound_inputs(sys), disturbance_inputs = nothing;
+function generate_control_function(sys::AbstractODESystem, inputs = unbound_inputs(sys),
+                                   disturbance_inputs = disturbances(sys);
                                    implicit_dae = false,
                                    simplify = false,
                                    kwargs...)
@@ -329,7 +330,7 @@ The structure represents a model of a disturbance, along with the input variable
 - `model::M`: A model of the disturbance. This is typically an `ODESystem`, but type that implements [`ModelingToolkit.get_disturbance_system`](@ref)`(dist::DisturbanceModel) -> ::ODESystem` is supported.
 """
 struct DisturbanceModel{M}
-    input
+    input::Any
     model::M
 end
 
@@ -387,18 +388,17 @@ dist = ModelingToolkit.DisturbanceModel(model.torque.tau.u, dmodel)
 `f_oop` will have an extra state corresponding to the integrator in the disturbance model. This state will not be affected by any input, but will affect the dynamics from where it enters, in this case it will affect additively from `model.torque.tau.u`.
 """
 function add_input_disturbance(sys, dist::DisturbanceModel)
-    t = ModelingToolkit.get_iv(sys)
-    @variables d(t)=0 
-    @variables u(t)=0 [input=true]
+    t = get_iv(sys)
+    @variables d(t)=0 [disturbance = true]
+    @variables u(t)=0 [input = true]
     dsys = get_disturbance_system(dist)
 
-    eqs = [
-        dsys.input.u[1] ~ d
-        dist.input ~ u + dsys.output.u[1]
-    ]
+    eqs = [dsys.input.u[1] ~ d
+           dist.input ~ u + dsys.output.u[1]]
+
+    augmented_sys = ODESystem(eqs, t, systems = [sys, dsys], name = gensym(:outer))
 
-    augmented_sys = ODESystem(eqs, t, systems=[sys, dsys], name=gensym(:outer))
-    
-    (f_oop, f_ip), dvs, p = ModelingToolkit.generate_control_function(augmented_sys, [u], [d])
+    (f_oop, f_ip), dvs, p = generate_control_function(augmented_sys, [u],
+                                                      [d])
     (f_oop, f_ip), augmented_sys, dvs, p
-end
+end

src/variables.jl

@@ -158,6 +158,10 @@ function isdisturbance(x)
     Symbolics.getmetadata(x, VariableDisturbance, false)
 end
 
+function disturbances(sys)
+    [filter(isdisturbance, states(sys)); filter(isdisturbance, parameters(sys))]
+end
+
 ## Tunable =====================================================================
 struct VariableTunable end
 Symbolics.option_to_metadata_type(::Val{:tunable}) = VariableTunable

test/input_output_handling.jl

@@ -196,8 +196,6 @@ eqs = [Differential(t)(x) ~ u]
 @named sys = ODESystem(eqs, t)
 @test_nowarn structural_simplify(sys)
 
-
-
 #=
 ## Disturbance input handling
 We test that the generated disturbance dynamics is correct by calling the dynamics in two different points that differ in the disturbance state, and check that we get the same result as when we call the linearized dynamics in the same two points. The true system is linear so the linearized dynamics are exact.
@@ -222,15 +220,21 @@ c = 10   # Damping coefficient
 @named damper = Rotational.Damper(; d = c)
 @named torque = Rotational.Torque()
 
-function SystemModel(u=nothing; name=:model)
-    eqs = [
-        connect(torque.flange, inertia1.flange_a)
-        connect(inertia1.flange_b, spring.flange_a, damper.flange_a)
-        connect(inertia2.flange_a, spring.flange_b, damper.flange_b)
-    ]
+function SystemModel(u = nothing; name = :model)
+    eqs = [connect(torque.flange, inertia1.flange_a)
+           connect(inertia1.flange_b, spring.flange_a, damper.flange_a)
+           connect(inertia2.flange_a, spring.flange_b, damper.flange_b)]
     if u !== nothing
         push!(eqs, connect(torque.tau, u.output))
-        return @named model = ODESystem(eqs, t; systems = [torque, inertia1, inertia2, spring, damper, u])
+        return @named model = ODESystem(eqs, t;
+                                        systems = [
+                                            torque,
+                                            inertia1,
+                                            inertia2,
+                                            spring,
+                                            damper,
+                                            u,
+                                        ])
     end
     ODESystem(eqs, t; systems = [torque, inertia1, inertia2, spring, damper], name)
 end
@@ -259,7 +263,5 @@ pn = ModelingToolkit.varmap_to_vars(def, p)
 xp0 = f_oop(x0, u, pn, 0)
 xp1 = f_oop(x1, u, pn, 0)
 
-@test xp0 ≈ matrices.A*x0 + matrices.B*[u; 0]
-@test xp1 ≈ matrices.A*x1 + matrices.B*[u; 0]
-
-
+@test xp0 ≈ matrices.A * x0 + matrices.B * [u; 0]
+@test xp1 ≈ matrices.A * x1 + matrices.B * [u; 0]