add utilities and tests for disturbance modeling

src/systems/analysis_points.jl

@@ -643,3 +643,38 @@ function linearization_function(sys::AbstractSystem,
 
     return linearization_function(system_modifier(sys), input_vars, output_vars; kwargs...)
 end
+
+# 
+
+"""
+    generate_control_function(sys::ModelingToolkit.AbstractODESystem, input_ap_name::Union{Symbol, Vector{Symbol}, AnalysisPoint, Vector{AnalysisPoint}}, dist_ap_name::Union{Symbol, Vector{Symbol}, AnalysisPoint, Vector{AnalysisPoint}}; system_modifier = identity, kwargs)
+
+When called with analysis points as input arguments, we assume that all analysis points corresponds to connections that should be opened (broken). The use case for this is to get rid of input signal blocks, such as `Step` or `Sine`, since these are useful for simulation but are not needed when using the plant model in a controller or state estimator.
+"""
+function generate_control_function(
+        sys::ModelingToolkit.AbstractODESystem, input_ap_name::Union{
+            Symbol, Vector{Symbol}, AnalysisPoint, Vector{AnalysisPoint}},
+        dist_ap_name::Union{
+            Nothing, Symbol, Vector{Symbol}, AnalysisPoint, Vector{AnalysisPoint}} = nothing;
+        system_modifier = identity,
+        kwargs...)
+    input_ap_name = canonicalize_ap(input_ap_name)
+    u = []
+    for input_ap in input_ap_name
+        sys, (du, _) = open_loop(sys, input_ap)
+        push!(u, du)
+    end
+    if dist_ap_name === nothing
+        return ModelingToolkit.generate_control_function(system_modifier(sys), u; kwargs...)
+    end
+
+    dist_ap_name = canonicalize_ap(dist_ap_name)
+    d = []
+    for dist_ap in dist_ap_name
+        sys, (du, _) = open_loop(sys, dist_ap)
+        push!(d, du)
+    end
+
+    ModelingToolkit.generate_control_function(system_modifier(sys), u, d; kwargs...)
+end
+

src/systems/diffeqs/odesystem.jl

@@ -449,6 +449,8 @@ an array of inputs `inputs` is given, and `param_only` is false for a time-depen
 """
 function build_explicit_observed_function(sys, ts;
         inputs = nothing,
+        disturbance_inputs = nothing,
+        disturbance_argument = false,
         expression = false,
         eval_expression = false,
         eval_module = @__MODULE__,
@@ -575,12 +577,15 @@ function build_explicit_observed_function(sys, ts;
     if inputs !== nothing
         ps = setdiff(ps, inputs) # Inputs have been converted to parameters by io_preprocessing, remove those from the parameter list
     end
+    if disturbance_inputs !== nothing
+        ps = setdiff(ps, disturbance_inputs)
+    end
     _ps = ps
     if ps isa Tuple
         ps = DestructuredArgs.(unwrap.(ps), inbounds = !checkbounds)
     elseif has_index_cache(sys) && get_index_cache(sys) !== nothing
         ps = DestructuredArgs.(reorder_parameters(get_index_cache(sys), unwrap.(ps)))
-        if isempty(ps) && inputs !== nothing
+        if isempty(ps) && (inputs !== nothing || disturbance_inputs !== nothing)
             ps = (:EMPTY,)
         end
     else
@@ -601,13 +606,20 @@ function build_explicit_observed_function(sys, ts;
         args = param_only ? [ipts, ps..., ivs...] : [dvs..., ipts, ps..., ivs...]
         p_start += 1
     end
+    if disturbance_argument
+        disturbance_inputs = unwrap.(disturbance_inputs)
+        dist_inputs = DestructuredArgs(disturbance_inputs, inbounds = !checkbounds)
+        args = [args; dist_inputs]
+    end
     pre = get_postprocess_fbody(sys)
 
     array_wrapper = if param_only
-        wrap_array_vars(sys, ts; ps = _ps, dvs = nothing, inputs, history = is_dde(sys)) .∘
+        wrap_array_vars(sys, ts; ps = _ps, dvs = nothing, inputs,
+            history = is_dde(sys), extra_args = (disturbance_inputs,)) .∘
         wrap_parameter_dependencies(sys, isscalar)
     else
-        wrap_array_vars(sys, ts; ps = _ps, inputs, history = is_dde(sys)) .∘
+        wrap_array_vars(sys, ts; ps = _ps, inputs, history = is_dde(sys),
+            extra_args = (disturbance_inputs,)) .∘
         wrap_parameter_dependencies(sys, isscalar)
     end
     mtkparams_wrapper = wrap_mtkparameters(sys, isscalar, p_start)

test/downstream/test_disturbance_model.jl

@@ -0,0 +1,215 @@
+#=
+This file implements and tests a typical workflow for state estimation with disturbance models
+The primary subject of the tests is the analysis-point features and the
+analysis-point specific method for `generate_control_function`.
+=#
+using ModelingToolkit, OrdinaryDiffEq, LinearAlgebra, Test
+using ModelingToolkitStandardLibrary.Mechanical.Rotational
+using ModelingToolkitStandardLibrary.Blocks
+using ModelingToolkit: connect
+# using Plots
+
+using ModelingToolkit: t_nounits as t, D_nounits as D
+
+indexof(sym, syms) = findfirst(isequal(sym), syms)
+
+## Build the system model ======================================================
+@mtkmodel SystemModel begin
+    @parameters begin
+        m1 = 1
+        m2 = 1
+        k = 10 # Spring stiffness
+        c = 3  # Damping coefficient
+    end
+    @components begin
+        inertia1 = Inertia(; J = m1, phi = 0, w = 0)
+        inertia2 = Inertia(; J = m2, phi = 0, w = 0)
+        spring = Spring(; c = k)
+        damper = Damper(; d = c)
+        torque = Torque(use_support = false)
+    end
+    @equations begin
+        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)
+    end
+end
+
+@mtkmodel ModelWithInputs begin
+    @components begin
+        input_signal = Blocks.Sine(frequency = 1, amplitude = 1)
+        disturbance_signal1 = Blocks.Constant(k = 0)
+        disturbance_signal2 = Blocks.Constant(k = 0)
+        disturbance_torque1 = Torque(use_support = false)
+        disturbance_torque2 = Torque(use_support = false)
+        system_model = SystemModel()
+    end
+    @equations begin
+        connect(input_signal.output, :u, system_model.torque.tau)
+        connect(disturbance_signal1.output, :d1, disturbance_torque1.tau)
+        connect(disturbance_signal2.output, :d2, disturbance_torque2.tau)
+        connect(disturbance_torque1.flange, system_model.inertia1.flange_b)
+        connect(disturbance_torque2.flange, system_model.inertia2.flange_b)
+    end
+end
+
+@named model = ModelWithInputs() # Model with load disturbance
+ssys = structural_simplify(model)
+prob = ODEProblem(ssys, [], (0.0, 10.0))
+sol = solve(prob, Tsit5())
+plot(sol)
+
+##
+using ControlSystemsBase
+# cmodel = complete(model)
+# P = cmodel.system_model
+# lsys = named_ss(
+#     model, [:u, :d1], [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w])
+
+##
+# If we now want to add a disturbance model, we cannot do that since we have already connected a constant to the disturbance input. We have also already used the name `d` for an analysis point, but that might not be an issue since we create an outer model and get a new namespace.
+
+s = tf("s")
+dist(; name) = ODESystem(1 / s; name)
+
+@mtkmodel SystemModelWithDisturbanceModel begin
+    @components begin
+        input_signal = Blocks.Sine(frequency = 1, amplitude = 1)
+        disturbance_signal1 = Blocks.Constant(k = 0)
+        disturbance_signal2 = Blocks.Constant(k = 0)
+        disturbance_torque1 = Torque(use_support = false)
+        disturbance_torque2 = Torque(use_support = false)
+        disturbance_model = dist()
+        system_model = SystemModel()
+    end
+    @equations begin
+        connect(input_signal.output, :u, system_model.torque.tau)
+        connect(disturbance_signal1.output, :d1, disturbance_model.input)
+        connect(disturbance_model.output, disturbance_torque1.tau)
+        connect(disturbance_signal2.output, :d2, disturbance_torque2.tau)
+        connect(disturbance_torque1.flange, system_model.inertia1.flange_b)
+        connect(disturbance_torque2.flange, system_model.inertia2.flange_b)
+    end
+end
+
+@named model_with_disturbance = SystemModelWithDisturbanceModel()
+# ssys = structural_simplify(open_loop(model_with_disturbance, :d)) # Open loop worked, but it's a bit awkward that we have to use it here
+# lsys2 = named_ss(model_with_disturbance, [:u, :d1],
+# [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w])
+ssys = structural_simplify(model_with_disturbance)
+prob = ODEProblem(ssys, [], (0.0, 10.0))
+sol = solve(prob, Tsit5())
+@test SciMLBase.successful_retcode(sol)
+# plot(sol)
+
+## 
+# Now we only have an integrating disturbance affecting inertia1, what if we want both integrating and direct Gaussian? We'd need a "PI controller" disturbancemodel. If we add the disturbance model (s+1)/s we get the integrating and non-integrating noises being correlated which is fine, it reduces the dimensions of the sigma point by 1.
+
+dist3(; name) = ODESystem(ss(1 + 10 / s, balance = false); name)
+
+@mtkmodel SystemModelWithDisturbanceModel begin
+    @components begin
+        input_signal = Blocks.Sine(frequency = 1, amplitude = 1)
+        disturbance_signal1 = Blocks.Constant(k = 0)
+        disturbance_signal2 = Blocks.Constant(k = 0)
+        disturbance_torque1 = Torque(use_support = false)
+        disturbance_torque2 = Torque(use_support = false)
+        disturbance_model = dist3()
+        system_model = SystemModel()
+
+        y = Blocks.Add()
+        angle_sensor = AngleSensor()
+        output_disturbance = Blocks.Constant(k = 0)
+    end
+    @equations begin
+        connect(input_signal.output, :u, system_model.torque.tau)
+        connect(disturbance_signal1.output, :d1, disturbance_model.input)
+        connect(disturbance_model.output, disturbance_torque1.tau)
+        connect(disturbance_signal2.output, :d2, disturbance_torque2.tau)
+        connect(disturbance_torque1.flange, system_model.inertia1.flange_b)
+        connect(disturbance_torque2.flange, system_model.inertia2.flange_b)
+
+        connect(system_model.inertia1.flange_b, angle_sensor.flange)
+        connect(angle_sensor.phi, y.input1)
+        connect(output_disturbance.output, :dy, y.input2)
+    end
+end
+
+@named model_with_disturbance = SystemModelWithDisturbanceModel()
+# ssys = structural_simplify(open_loop(model_with_disturbance, :d)) # Open loop worked, but it's a bit awkward that we have to use it here
+# lsys3 = named_ss(model_with_disturbance, [:u, :d1],
+#     [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w])
+ssys = structural_simplify(model_with_disturbance)
+prob = ODEProblem(ssys, [], (0.0, 10.0))
+sol = solve(prob, Tsit5())
+@test SciMLBase.successful_retcode(sol)
+# plot(sol)
+
+## Generate function for an augmented Unscented Kalman Filter =====================
+# temp = open_loop(model_with_disturbance, :d)
+outputs = [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w]
+(f_oop1, f_ip), x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
+    model_with_disturbance, [:u], [:d1, :d2, :dy], split = false)
+
+(f_oop2, f_ip2), x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
+    model_with_disturbance, [:u], [:d1, :d2, :dy],
+    disturbance_argument = true, split = false)
+
+measurement = ModelingToolkit.build_explicit_observed_function(
+    io_sys, outputs, inputs = ModelingToolkit.inputs(io_sys)[1:1])
+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)
+
+op = ModelingToolkit.inputs(io_sys) .=> 0
+x0, p = ModelingToolkit.get_u0_p(io_sys, op, op)
+x = zeros(5)
+u = zeros(1)
+d = zeros(3)
+@test f_oop2(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]
+
+# Add to the integrating disturbance input
+d = [1, 0, 0]
+@test sort(f_oop2(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]
+
+d = [0, 1, 0]
+@test sort(f_oop2(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]
+
+d = [0, 0, 1]
+@test sort(f_oop2(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
+
+## Further downstream tests that the functions generated above actually have the properties required to use for state estimation
+# 
+# using LowLevelParticleFilters, SeeToDee
+# Ts = 0.001
+# discrete_dynamics = SeeToDee.Rk4(f_oop2, Ts)
+# nx = length(x_sym)
+# nu = 1
+# nw = 2
+# ny = length(outputs)
+# R1 = Diagonal([1e-5, 1e-5])
+# R2 = 0.1 * I(ny)
+# op = ModelingToolkit.inputs(io_sys) .=> 0
+# x0, p = ModelingToolkit.get_u0_p(io_sys, op, op)
+# d0 = LowLevelParticleFilters.SimpleMvNormal(x0, 10.0I(nx))
+# measurement_model = UKFMeasurementModel{Float64, false, false}(measurement, R2; nx, ny)
+# kf = UnscentedKalmanFilter{false, false, true, false}(
+#     discrete_dynamics, measurement_model, R1, d0; nu, Ts, p)
+
+# tvec = 0:Ts:sol.t[end]
+# u = vcat.(Array(sol(tvec, idxs = P.torque.tau.u)))
+# y = collect.(eachcol(Array(sol(tvec, idxs = outputs)) .+ 1e-2 .* randn.()))
+
+# inds = 1:5805
+# res = forward_trajectory(kf, u, y)
+
+# plot(res, size = (1000, 1000));
+# plot!(sol, idxs = x_sym, sp = (1:nx)', l = :dash);

test/runtests.jl

@@ -112,6 +112,7 @@ end
         @safetestset "Linearization Dummy Derivative Tests" include("downstream/linearization_dd.jl")
         @safetestset "Inverse Models Test" include("downstream/inversemodel.jl")
         @safetestset "Analysis Points Test" include("downstream/analysis_points.jl")
+        @safetestset "Analysis Points Test" include("downstream/test_disturbance_model.jl")
     end
 
     if GROUP == "All" || GROUP == "Extensions"