increase flexibility in disturbance-argument codegen

docs/src/tutorials/disturbance_modeling.md

@@ -119,7 +119,15 @@ y &= g(x, u, p, t)
 \end{aligned}
 ```
 
-To make use of the model defined above for state estimation, we may want to generate a Julia function for the dynamics ``f`` and the output equations ``g`` that we can plug into, e.g., a nonlinear version of a Kalman filter or a particle filter, etc. MTK contains utilities to do this, namely, [`ModelingToolkit.generate_control_function`](@ref) and [`ModelingToolkit.build_explicit_observed_function`](@ref) (described in more details in ["Input output"](@ref inputoutput)). These functions take keyword arguments `disturbance_inputs` and `disturbance_argument`, that indicate which variables in the model are considered part of ``w``, and whether or not these variables are to be added as function arguments to ``f``, i.e., whether we have ``f(x, u, p, t)`` or ``f(x, u, p, t, w)``. If we do not include the disturbance inputs as function arguments, MTK will assume that the ``w`` variables are all zero, but any dynamics associated with these variables, such as disturbance models, will be included in the generated function. This allows a state estimator to estimate the state of the disturbance model, provided that this state is [observable](https://en.wikipedia.org/wiki/Observability) from the measured outputs of the system.
+To make use of the model defined above for state estimation, we may want to generate a Julia function for the dynamics ``f`` and the output equations ``g`` that we can plug into, e.g., a nonlinear version of a Kalman filter or a particle filter, etc. MTK contains utilities to do this, namely, [`ModelingToolkit.generate_control_function`](@ref) and [`ModelingToolkit.build_explicit_observed_function`](@ref) (described in more details in ["Input output"](@ref inputoutput)).
+
+These functions support two types of disturbance inputs:
+
+- **Unknown disturbances** (`disturbance_inputs`): Variables that are part of ``w`` but NOT added as function arguments to ``f``. MTK assumes these variables are zero, but any dynamics associated with them (such as disturbance models) are included in the generated function. This allows a state estimator to estimate the state of the disturbance model, provided that this state is [observable](https://en.wikipedia.org/wiki/Observability) from measured outputs.
+
+- **Known disturbances** (`known_disturbance_inputs`): Variables that are part of ``w`` AND added as function arguments to ``f``, resulting in ``f(x, u, p, t, w)``. Use this when disturbances can be measured or are otherwise known.
+
+You can mix and match: some disturbances can be unknown while others are known. For example, `generate_control_function(sys, inputs; disturbance_inputs=[w1], known_disturbance_inputs=[w2, w3])` generates a function `f(x, u, p, t, [w2, w3])` where `w1` is set to zero but its dynamics are preserved, while `w2` and `w3` must be provided as arguments.
 
 Below, we demonstrate
 
@@ -187,7 +195,7 @@ outputs = [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w]
 (f_oop, f_ip), x_sym,
 p_sym,
 io_sys = ModelingToolkit.generate_control_function(
-    model_with_disturbance, inputs, disturbance_inputs; disturbance_argument = true)
+    model_with_disturbance, inputs; known_disturbance_inputs = disturbance_inputs)
 
 g = ModelingToolkit.build_explicit_observed_function(
     io_sys, outputs; inputs)
@@ -196,7 +204,7 @@ op = ModelingToolkit.inputs(io_sys) .=> 0
 x0 = ModelingToolkit.get_u0(io_sys, op)
 p = MTKParameters(io_sys, op)
 u = zeros(1) # Control input
-w = zeros(length(disturbance_inputs)) # Disturbance input
+w = zeros(length(disturbance_inputs)) # Disturbance input (known disturbances are provided as arguments)
 @test f_oop(x0, u, p, t, w) == zeros(5)
 @test g(x0, u, p, 0.0) == [0, 0, 0, 0]
 

src/inputoutput.jl

@@ -156,28 +156,35 @@ has_var(ex, x) = x ∈ Set(get_variables(ex))
 """
     (f_oop, f_ip), x_sym, p_sym, io_sys = generate_control_function(
             sys::System,
-            inputs             = unbound_inputs(sys),
-            disturbance_inputs = disturbances(sys);
-            implicit_dae       = false,
-            simplify           = false,
-            split              = true,
+            inputs                   = unbound_inputs(sys),
+            disturbance_inputs       = disturbances(sys);
+            known_disturbance_inputs = nothing,
+            implicit_dae             = false,
+            simplify                 = false,
+            split                    = true,
         )
 
 For a system `sys` with inputs (as determined by [`unbound_inputs`](@ref) or user specified), generate functions with additional input argument `u`
 
 The returned functions are the out-of-place (`f_oop`) and in-place (`f_ip`) forms:
 ```
-f_oop : (x,u,p,t)      -> rhs
-f_ip  : (xout,x,u,p,t) -> nothing
+f_oop : (x,u,p,t)      -> rhs         # basic form
+f_oop : (x,u,p,t,w)    -> rhs         # with known_disturbance_inputs
+f_ip  : (xout,x,u,p,t) -> nothing     # basic form
+f_ip  : (xout,x,u,p,t,w) -> nothing   # with known_disturbance_inputs
 ```
 
 The return values also include the chosen state-realization (the remaining unknowns) `x_sym` and parameters, in the order they appear as arguments to `f`.
 
-If `disturbance_inputs` is an array of variables, the generated dynamics function will preserve any state and dynamics associated with disturbance inputs, but the disturbance inputs themselves will (by default) not be included as inputs to the generated function. The use case for this is to generate dynamics for state observers that estimate the influence of unmeasured disturbances, and thus require unknown variables for the disturbance model, but without disturbance inputs since the disturbances are not available for measurement. To add an input argument corresponding to the disturbance inputs, either include the disturbance inputs among the control inputs, or set `disturbance_argument=true`, in which case an additional input argument `w` is added to the generated function `(x,u,p,t,w)->rhs`.
+# Disturbance Handling
+
+- `disturbance_inputs`: Unknown disturbance inputs. The generated dynamics will preserve any state and dynamics associated with these disturbances, but the disturbance inputs themselves will not be included as function arguments. This is useful for state observers that estimate unmeasured disturbances.
+
+- `known_disturbance_inputs`: Known disturbance inputs. The generated dynamics will preserve state and dynamics, and the disturbance inputs will be added as an additional input argument `w` to the generated function: `(x,u,p,t,w)->rhs`.
 
 # Example
 
-```
+```julia
 using ModelingToolkit: generate_control_function, varmap_to_vars, defaults
 f, x_sym, ps = generate_control_function(sys, expression=Val{false}, simplify=false)
 p = varmap_to_vars(defaults(sys), ps)
@@ -188,6 +195,7 @@ f[1](x, inputs, p, t)
 """
 function generate_control_function(sys::AbstractSystem, inputs = unbound_inputs(sys),
         disturbance_inputs = disturbances(sys);
+        known_disturbance_inputs = nothing,
         disturbance_argument = false,
         implicit_dae = false,
         simplify = false,
@@ -196,30 +204,55 @@ function generate_control_function(sys::AbstractSystem, inputs = unbound_inputs(
         split = true,
         kwargs...)
     isempty(inputs) && @warn("No unbound inputs were found in system.")
+
+    # Handle backward compatibility for disturbance_argument
+    if disturbance_argument
+        Base.depwarn("The `disturbance_argument` keyword argument is deprecated. Use `known_disturbance_inputs` instead. " *
+                     "For `disturbance_argument=true`, pass `known_disturbance_inputs=disturbance_inputs, disturbance_inputs=nothing`. " *
+                     "For `disturbance_argument=false`, use `disturbance_inputs` as before.",
+                     :generate_control_function)
+        if known_disturbance_inputs !== nothing
+            error("Cannot specify both `disturbance_argument=true` and `known_disturbance_inputs`")
+        end
+        known_disturbance_inputs = disturbance_inputs
+        disturbance_inputs = nothing
+    end
+
+    # Collect all disturbance inputs for mtkcompile
+    all_disturbances = vcat(
+        disturbance_inputs === nothing ? [] : disturbance_inputs,
+        known_disturbance_inputs === nothing ? [] : known_disturbance_inputs
+    )
+
     if !isscheduled(sys)
-        sys = mtkcompile(sys; inputs, disturbance_inputs, split)
+        sys = mtkcompile(sys; inputs, disturbance_inputs=all_disturbances, split)
     end
-    if disturbance_inputs !== nothing
-        # add to inputs for the purposes of io processing
-        inputs = [inputs; disturbance_inputs]
+
+    # Add all disturbances to inputs for the purposes of io processing
+    if !isempty(all_disturbances)
+        inputs = [inputs; all_disturbances]
     end
 
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
     ps = setdiff(ps, inputs)
+
+    # Remove unknown disturbances from inputs (we don't want them as actual inputs to the dynamics)
     if disturbance_inputs !== nothing
-        # remove from inputs since we do not want them as actual inputs to the dynamics
         inputs = setdiff(inputs, disturbance_inputs)
-        # ps = [ps; disturbance_inputs]
     end
+
     inputs = map(value, inputs)
-    disturbance_inputs = unwrap.(disturbance_inputs)
+
+    # Prepare disturbance arrays for substitution and function arguments
+    unknown_disturbances = disturbance_inputs === nothing ? [] : unwrap.(disturbance_inputs)
+    known_disturbances = known_disturbance_inputs === nothing ? [] : unwrap.(known_disturbance_inputs)
 
     eqs = [eq for eq in full_equations(sys)]
 
-    if disturbance_inputs !== nothing && !disturbance_argument
-        # Set all disturbance *inputs* to zero (we just want to keep the disturbance state)
-        subs = Dict(disturbance_inputs .=> 0)
+    # Set unknown disturbance inputs to zero (we just want to keep the disturbance state)
+    if !isempty(unknown_disturbances)
+        subs = Dict(unknown_disturbances .=> 0)
         eqs = [eq.lhs ~ substitute(eq.rhs, subs) for eq in eqs]
     end
     check_operator_variables(eqs, Differential)
@@ -231,8 +264,9 @@ function generate_control_function(sys::AbstractSystem, inputs = unbound_inputs(
     p = reorder_parameters(sys, ps)
     t = get_iv(sys)
 
-    if disturbance_argument
-        args = (dvs, inputs, p..., t, disturbance_inputs)
+    # Construct args with known disturbances if provided
+    if !isempty(known_disturbances)
+        args = (dvs, inputs, p..., t, known_disturbances)
     else
         args = (dvs, inputs, p..., t)
     end
@@ -245,7 +279,8 @@ function generate_control_function(sys::AbstractSystem, inputs = unbound_inputs(
     f = eval_or_rgf.(f; eval_expression, eval_module)
     f = GeneratedFunctionWrapper{(
         3 + implicit_dae, length(args) - length(p) + 1, is_split(sys))}(f...)
-    ps = setdiff(parameters(sys), inputs, disturbance_inputs)
+    # Return parameters excluding both control inputs and all disturbances
+    ps = setdiff(parameters(sys), inputs, all_disturbances)
     (; f = (f, f), dvs, ps, io_sys = sys)
 end
 

src/systems/codegen.jl

@@ -982,6 +982,8 @@ Generates a function that computes the observed value(s) `ts` in the system `sys
 - `output_type = Array` the type of the array generated by a out-of-place vector-valued function
 - `param_only = false` if true, only allow the generated function to access system parameters
 - `inputs = nothing` additinoal symbolic variables that should be provided to the generated function
+- `disturbance_inputs = nothing` symbolic variables representing unknown disturbance inputs (removed from parameters, not added as function arguments)
+- `known_disturbance_inputs = nothing` symbolic variables representing known disturbance inputs (removed from parameters, added as function arguments)
 - `checkbounds = true` checks bounds if true when destructuring parameters
 - `op = Operator` sets the recursion terminator for the walk done by `vars` to identify the variables that appear in `ts`. See the documentation for `vars` for more detail.
 - `throw = true` if true, throw an error when generating a function for `ts` that reference variables that do not exist.
@@ -1005,16 +1007,18 @@ The signatures will be of the form `g(...)` with arguments:
 
 - `output` for in-place functions
 - `unknowns` if `param_only` is `false`
-- `inputs` if `inputs` is an array of symbolic inputs that should be available in `ts` 
+- `inputs` if `inputs` is an array of symbolic inputs that should be available in `ts`
 - `p...` unconditionally; note that in the case of `MTKParameters` more than one parameters argument may be present, so it must be splatted
 - `t` if the system is time-dependent; for example systems of nonlinear equations will not have `t`
+- `known_disturbance_inputs` if provided; these are disturbance inputs that are known and provided as arguments
 
-For example, a function `g(op, unknowns, p..., inputs, t)` will be the in-place function generated if `return_inplace` is true, `ts` is a vector, 
-an array of inputs `inputs` is given, and `param_only` is false for a time-dependent system.
+For example, a function `g(op, unknowns, p..., inputs, t, known_disturbances)` will be the in-place function generated if `return_inplace` is true, `ts` is a vector,
+an array of inputs `inputs` is given, `known_disturbance_inputs` is provided, and `param_only` is false for a time-dependent system.
 """
 function build_explicit_observed_function(sys, ts;
         inputs = nothing,
         disturbance_inputs = nothing,
+        known_disturbance_inputs = nothing,
         disturbance_argument = false,
         expression = false,
         eval_expression = false,
@@ -1095,22 +1099,36 @@ function build_explicit_observed_function(sys, ts;
         ps = setdiff(ps, inputs) # Inputs have been converted to parameters by io_preprocessing, remove those from the parameter list
         inputs = (inputs,)
     end
+    # Handle backward compatibility for disturbance_argument
+    if disturbance_argument
+        Base.depwarn("The `disturbance_argument` keyword argument is deprecated. Use `known_disturbance_inputs` instead. " *
+                     "For `disturbance_argument=true`, pass `known_disturbance_inputs=disturbance_inputs, disturbance_inputs=nothing`. " *
+                     "For `disturbance_argument=false`, use `disturbance_inputs` as before.",
+                     :build_explicit_observed_function)
+        if known_disturbance_inputs !== nothing
+            error("Cannot specify both `disturbance_argument=true` and `known_disturbance_inputs`")
+        end
+        known_disturbance_inputs = disturbance_inputs
+        disturbance_inputs = nothing
+    end
+
+    # Remove disturbance inputs from parameters (both known and unknown)
     if disturbance_inputs !== nothing
-        # Disturbance inputs may or may not be included as inputs, depending on disturbance_argument
         ps = setdiff(ps, disturbance_inputs)
     end
-    if disturbance_argument
-        disturbance_inputs = (disturbance_inputs,)
+    if known_disturbance_inputs !== nothing
+        ps = setdiff(ps, known_disturbance_inputs)
+        known_disturbance_inputs = (known_disturbance_inputs,)
     else
-        disturbance_inputs = ()
+        known_disturbance_inputs = ()
     end
     ps = reorder_parameters(sys, ps)
     iv = if is_time_dependent(sys)
         (get_iv(sys),)
     else
         ()
     end
-    args = (dvs..., inputs..., ps..., iv..., disturbance_inputs...)
+    args = (dvs..., inputs..., ps..., iv..., known_disturbance_inputs...)
     p_start = length(dvs) + length(inputs) + 1
     p_end = length(dvs) + length(inputs) + length(ps)
     fns = build_function_wrapper(

test/downstream/test_disturbance_model.jl

@@ -167,6 +167,10 @@ measurement2 = ModelingToolkit.build_explicit_observed_function(
     io_sys, [io_sys.y.output.u], inputs = ModelingToolkit.inputs(io_sys)[1:1],
     disturbance_inputs = ModelingToolkit.inputs(io_sys)[2:end],
     disturbance_argument = true)
+# Test new interface with known_disturbance_inputs (equivalent to above)
+measurement3 = ModelingToolkit.build_explicit_observed_function(
+    io_sys, [io_sys.y.output.u], inputs = ModelingToolkit.inputs(io_sys)[1:1],
+    known_disturbance_inputs = ModelingToolkit.inputs(io_sys)[2:end])
 
 op = ModelingToolkit.inputs(io_sys) .=> 0
 x0 = ModelingToolkit.get_u0(io_sys, op)
@@ -177,21 +181,25 @@ d = zeros(3)
 @test f[1](x, u, p, t, d) == zeros(5)
 @test measurement(x, u, p, 0.0) == [0, 0, 0, 0]
 @test measurement2(x, u, p, 0.0, d) == [0]
+@test measurement3(x, u, p, 0.0, d) == [0]  # Test new interface
 
 # Add to the integrating disturbance input
 d = [1, 0, 0]
 @test sort(f[1](x, u, p, 0.0, d)) == [0, 0, 0, 1, 1] # Affects disturbance state and one velocity
 @test measurement2(x, u, p, 0.0, d) == [0]
+@test measurement3(x, u, p, 0.0, d) == [0]  # Test new interface
 
 d = [0, 1, 0]
 @test sort(f[1](x, u, p, 0.0, d)) == [0, 0, 0, 0, 1] # Affects one velocity
 @test measurement(x, u, p, 0.0) == [0, 0, 0, 0]
 @test measurement2(x, u, p, 0.0, d) == [0]
+@test measurement3(x, u, p, 0.0, d) == [0]  # Test new interface
 
 d = [0, 0, 1]
 @test sort(f[1](x, u, p, 0.0, d)) == [0, 0, 0, 0, 0] # Affects nothing
 @test measurement(x, u, p, 0.0) == [0, 0, 0, 0]
 @test measurement2(x, u, p, 0.0, d) == [1] # We have now disturbed the output
+@test measurement3(x, u, p, 0.0, d) == [1]  # Test new interface
 
 ## Further downstream tests that the functions generated above actually have the properties required to use for state estimation
 # 

test/input_output_handling.jl

@@ -216,6 +216,53 @@ end
         u = [rand()]
         d = [rand()]
         @test f[1](x, u, p, t, d) ≈ -x + u + [d[]^2]
+
+        ## Test new known_disturbance_inputs parameter (equivalent to disturbance_argument=true)
+        @variables x(t)=0 u(t)=0 [input=true] d(t)=0
+        eqs = [
+            D(x) ~ -x + u + d^2
+        ]
+
+        @named sys = System(eqs, t)
+        f, dvs,
+        ps,
+        io_sys = ModelingToolkit.generate_control_function(
+            sys, [u];
+            known_disturbance_inputs = [d],
+            simplify, split)
+
+        @test isequal(dvs[], x)
+        @test isempty(ps)
+
+        p = [rand()]
+        x = [rand()]
+        u = [rand()]
+        d = [rand()]
+        @test f[1](x, u, p, t, d) ≈ -x + u + [d[]^2]
+
+        ## Test mixed known and unknown disturbances
+        @variables x(t)=0 u(t)=0 [input=true] d1(t)=0 d2(t)=0
+        eqs = [
+            D(x) ~ -x + u + d1^2 + d2^3
+        ]
+
+        @named sys = System(eqs, t)
+        f, dvs,
+        ps,
+        io_sys = ModelingToolkit.generate_control_function(
+            sys, [u], [d1];           # d1 is unknown (set to zero)
+            known_disturbance_inputs = [d2],     # d2 is known (function argument)
+            simplify, split)
+
+        @test isequal(dvs[], x)
+        @test isempty(ps)
+
+        p = [rand()]
+        x = [rand()]
+        u = [rand()]
+        d2_val = [rand()]
+        # d1 is set to zero, so only u and d2 affect the dynamics
+        @test f[1](x, u, p, t, d2_val) ≈ -x + u + [d2_val[]^3]
     end
 end