Add linearize function

Co-authored-by: Fredrik Bagge Carlson <[email protected]>

Author: YingboMa

Project.toml

@@ -17,6 +17,7 @@ Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
@@ -55,6 +56,7 @@ DiffRules = "0.1, 1.0"
 Distributions = "0.23, 0.24, 0.25"
 DocStringExtensions = "0.7, 0.8, 0.9"
 DomainSets = "0.5"
+ForwardDiff = "0.10.3"
 Graphs = "1.5.2"
 IfElse = "0.1"
 JuliaFormatter = "1"

src/ModelingToolkit.jl

@@ -4,7 +4,7 @@ $(DocStringExtensions.README)
 module ModelingToolkit
 using DocStringExtensions
 using AbstractTrees
-using DiffEqBase, SciMLBase, Reexport
+using DiffEqBase, SciMLBase, ForwardDiff, Reexport
 using Distributed
 using StaticArrays, LinearAlgebra, SparseArrays, LabelledArrays
 using InteractiveUtils

src/inputoutput.jl

@@ -61,20 +61,20 @@ unbound_outputs(sys) = filter(x -> !is_bound(sys, x), outputs(sys))
 Determine whether or not input/output variable `u` is "bound" within the system, i.e., if it's to be considered internal to `sys`.
 A variable/signal is considered bound if it appears in an equation together with variables from other subsystems.
 The typical usecase for this function is to determine whether the input to an IO component is connected to another component,
-or if it remains an external input that the user has to supply before simulating the system. 
+or if it remains an external input that the user has to supply before simulating the system.
 
 See also [`bound_inputs`](@ref), [`unbound_inputs`](@ref), [`bound_outputs`](@ref), [`unbound_outputs`](@ref)
 """
 function is_bound(sys, u, stack = [])
     #=
-    For observed quantities, we check if a variable is connected to something that is bound to something further out. 
+    For observed quantities, we check if a variable is connected to something that is bound to something further out.
     In the following scenario
     julia> observed(syss)
         2-element Vector{Equation}:
         sys₊y(tv) ~ sys₊x(tv)
         y(tv) ~ sys₊x(tv)
     sys₊y(t) is bound to the outer y(t) through the variable sys₊x(t) and should thus return is_bound(sys₊y(t)) = true.
-    When asking is_bound(sys₊y(t)), we know that we are looking through observed equations and can thus ask 
+    When asking is_bound(sys₊y(t)), we know that we are looking through observed equations and can thus ask
     if var is bound, if it is, then sys₊y(t) is also bound. This can lead to an infinite recursion, so we maintain a stack of variables we have previously asked about to be able to break cycles
     =#
     u ∈ Set(stack) && return false # Cycle detected
@@ -241,7 +241,7 @@ function toparam(sys, ctrls::AbstractVector)
     ODESystem(eqs, name = nameof(sys))
 end
 
-function inputs_to_parameters!(state::TransformationState)
+function inputs_to_parameters!(state::TransformationState, check_bound = true)
     @unpack structure, fullvars, sys = state
     @unpack var_to_diff, graph, solvable_graph = structure
     @assert solvable_graph === nothing
@@ -254,7 +254,7 @@ function inputs_to_parameters!(state::TransformationState)
     input_to_parameters = Dict()
     new_fullvars = []
     for (i, v) in enumerate(fullvars)
-        if isinput(v) && !is_bound(sys, v)
+        if isinput(v) && !(check_bound && is_bound(sys, v))
             if var_to_diff[i] !== nothing
                 error("Input $(fullvars[i]) is differentiated!")
             end
@@ -296,9 +296,12 @@ function inputs_to_parameters!(state::TransformationState)
 
     @set! sys.eqs = map(Base.Fix2(substitute, input_to_parameters), equations(sys))
     @set! sys.states = setdiff(states(sys), keys(input_to_parameters))
-    @set! sys.ps = [parameters(sys); new_parameters]
+    ps = parameters(sys)
+    @set! sys.ps = [ps; new_parameters]
 
     @set! state.sys = sys
     @set! state.fullvars = new_fullvars
     @set! state.structure = structure
+    base_params = length(ps)
+    return state, (base_params + 1):(base_params + length(new_parameters)) # (1:length(new_parameters)) .+ base_params
 end

src/systems/abstractsystem.jl

@@ -950,10 +950,10 @@ types during tearing.
 """
 function structural_simplify(sys::AbstractSystem; simplify = false, kwargs...)
     sys = expand_connections(sys)
-    sys = alias_elimination(sys)
     state = TearingState(sys)
-    state = inputs_to_parameters!(state)
-    sys = state.sys
+    state, = inputs_to_parameters!(state)
+    sys = alias_elimination!(state)
+    state = TearingState(sys)
     check_consistency(state)
     if sys isa ODESystem
         sys = dae_order_lowering(dummy_derivative(sys, state))
@@ -967,6 +967,87 @@ function structural_simplify(sys::AbstractSystem; simplify = false, kwargs...)
     return sys
 end
 
+export linearize
+# TODO: The order of the states and equations should match so that the Jacobians are ẋ = Ax
+function linearize(sys::AbstractSystem, inputs, outputs; simplify = false, kwargs...)
+    sys = expand_connections(sys)
+    state = TearingState(sys)
+    markio!(state, inputs, outputs)
+    state, input_idxs = inputs_to_parameters!(state, false)
+    sys = alias_elimination!(state)
+    state = TearingState(sys)
+    check_consistency(state)
+    if sys isa ODESystem
+        sys = dae_order_lowering(dummy_derivative(sys, state))
+    end
+    state = TearingState(sys)
+    find_solvables!(state; kwargs...)
+    sys = tearing_reassemble(state, tearing(state), simplify = simplify)
+    fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
+    @set! sys.observed = topsort_equations(observed(sys), fullstates)
+    invalidate_cache!(sys)
+
+    eqs = equations(sys)
+    check_operator_variables(eqs, Differential)
+    # Sort equations and states such that diff.eqs. match differential states and the rest are algebraic
+    diffstates = collect_operator_variables(sys, Differential)
+    eqs = sort(eqs, by = e -> !isoperator(e.lhs, Differential),
+               alg = Base.Sort.DEFAULT_STABLE)
+    @set! sys.eqs = eqs
+    diffstates = [arguments(e.lhs)[1] for e in eqs[1:length(diffstates)]]
+    sts = [diffstates; setdiff(states(sys), diffstates)]
+    @set! sys.states = sts
+
+    diff_idxs = 1:length(diffstates)
+    alge_idxs = (length(diffstates) + 1):length(sts)
+    fun = ODEFunction(sys)
+    lin_fun = let fun = fun,
+        h = ModelingToolkit.build_explicit_observed_function(sys, outputs)
+
+        (u, p, t) -> begin
+            uf = SciMLBase.UJacobianWrapper(fun, t, p)
+            fg_xz = ForwardDiff.jacobian(uf, u)
+            pf = SciMLBase.ParamJacobianWrapper(fun, t, u)
+            # TODO: this is very inefficient, p contains all parameters of the system
+            fg_u = ForwardDiff.jacobian(pf, p)[:, input_idxs]
+            h_xz = ForwardDiff.jacobian(xz -> h(xz, p, t), u)
+            h_u = ForwardDiff.jacobian(p -> h(u, p, t), p)[:, input_idxs]
+            (f_x = fg_xz[diff_idxs, diff_idxs],
+             f_z = fg_xz[diff_idxs, alge_idxs],
+             g_x = fg_xz[alge_idxs, diff_idxs],
+             g_z = fg_xz[alge_idxs, alge_idxs],
+             f_u = fg_u[diff_idxs, :],
+             g_u = fg_u[alge_idxs, :],
+             h_x = h_xz[:, diff_idxs],
+             h_z = h_xz[:, alge_idxs],
+             h_u = h_u)
+        end
+    end
+    return sys, lin_fun
+end
+
+function markio!(state::TearingState, inputs, outputs)
+    fullvars = state.fullvars
+    inputset = Set(inputs)
+    outputset = Set(outputs)
+    for (i, v) in enumerate(fullvars)
+        if v in inputset
+            v = setmetadata(v, ModelingToolkit.VariableInput, true)
+            v = setmetadata(v, ModelingToolkit.VariableOutput, false)
+            fullvars[i] = v
+        elseif v in outputset
+            v = setmetadata(v, ModelingToolkit.VariableInput, false)
+            v = setmetadata(v, ModelingToolkit.VariableOutput, true)
+            fullvars[i] = v
+        else
+            v = setmetadata(v, ModelingToolkit.VariableInput, false)
+            v = setmetadata(v, ModelingToolkit.VariableOutput, false)
+            fullvars[i] = v
+        end
+    end
+    state
+end
+
 @latexrecipe function f(sys::AbstractSystem)
     return latexify(equations(sys))
 end

src/systems/alias_elimination.jl

@@ -19,8 +19,9 @@ function aag_bareiss(sys::AbstractSystem)
     return aag_bareiss!(state.structure.graph, complete(state.structure.var_to_diff), mm)
 end
 
-function alias_elimination(sys)
-    state = TearingState(sys; quick_cancel = true)
+alias_elimination(sys) = alias_elimination!(TearingState(sys; quick_cancel = true))
+function alias_elimination!(state::TearingState)
+    sys = state.sys
     ag, mm = alias_eliminate_graph!(state)
     ag === nothing && return sys
 

test/linearize.jl

@@ -0,0 +1,83 @@
+using ModelingToolkit
+
+# r is an input, and y is an output.
+@variables t x(t)=0 y(t)=0 u(t)=0 r(t)=0
+@variables t x(t)=0 y(t)=0 u(t)=0 r(t)=0 [input = true]
+@parameters kp = 1
+D = Differential(t)
+
+eqs = [u ~ kp * (r - y)
+       D(x) ~ -x + u
+       y ~ x]
+
+@named sys = ODESystem(eqs, t)
+linearize(sys, [r], [y])
+
+##
+```
+
+  r ┌─────┐       ┌─────┐     ┌─────┐
+───►│     ├──────►│     │  u  │     │
+    │  F  │       │  C  ├────►│  P  │ y
+    └─────┘     ┌►│     │     │     ├─┬─►
+                │ └─────┘     └─────┘ │
+                │                     │
+                └─────────────────────┘
+```
+
+function plant(; name)
+    @variables x(t) = 1
+    @variables u(t)=0 [input = true] y(t)=0 [output = true]
+    D = Differential(t)
+    eqs = [D(x) ~ -x + u
+           y ~ x]
+    ODESystem(eqs, t; name = name)
+end
+
+function filt_(; name)
+    @variables x(t)=0 y(t)=0 [output = true]
+    @variables u(t)=0 [input = true]
+    D = Differential(t)
+    eqs = [D(x) ~ -2 * x + u
+           y ~ x]
+    ODESystem(eqs, t, name = name)
+end
+
+function controller(kp; name)
+    @variables y(t)=0 r(t)=0 [input = true] u(t)=0
+    @parameters kp = kp
+    eqs = [
+        u ~ kp * (r - y),
+    ]
+    ODESystem(eqs, t; name = name)
+end
+
+@named f = filt_()
+@named c = controller(1)
+@named p = plant()
+
+connections = [f.y ~ c.r # filtered reference to controller reference
+               c.u ~ p.u # controller output to plant input
+               p.y ~ c.y]
+
+@named cl = ODESystem(connections, t, systems = [f, c, p])
+
+lin, xs = linearize(cl, cl.f.u, cl.p.x)
+
+##
+using ModelingToolkitStandardLibrary.Blocks: LimPID
+#using ControlSystems
+k = 400;
+Ti = 0.5;
+Td = 1;
+Nd = 10;
+#s = tf("s")
+#expected_result_r = k*(1 + 1/(s*Ti)) |> ss
+#expected_result_y = k*(1 + 1/(s*Ti) - s*Td / (1 + s*Td/N)) |> ss
+@named pid = LimPID(; k, Ti, Td, Nd)
+ModelingToolkit.unbound_inputs(pid)
+
+@unpack reference, measurement, ctr_output = pid
+lin = linearize(pid, [reference.u, measurement.u], [ctr_output.u])
+lin, lin_fun = linearize(pid, [reference.u, measurement.u], [ctr_output.u]);
+lin_fun(prob.u0, prob.p, 0.0)