Tests (half working) for linearize in ModelingToolkit

linearize_tests.jl

using ModelingToolkit
using ModelingToolkit: renamespace
using LinearAlgebra
using ModelingToolkit: isinput, isoutput, is_bound, isdifferential
using Symbolics: jacobian, substitute

"""
    linearize(sys::ODESystem; numeric=false, x0=nothing)

Linaerize `sys` around operating point `x0`. `x0` can be `nothing` for a symbolic linearization, or a `Dict` that maps states and parameters to values.

Returns a named tuple with fields `A,B,x,u`
representing a linear state-space system
```
ẋ = Ax + Bu
```
where `x` is the deviation from the operating point `x0`.

If `numeric = true`, symbolics parameters are replaced with their default values and the matrices `A,B` will have element type `Float64`, only pass this option if all parameters and states have default values or have values in `x0`.

NOTE: This function is experimental and temporary, in anticipation of linearization support in ModelingToolkit.
"""
function linearize(sys::ODESystem; x0=nothing, numeric=false)
    sysr = initialize_system_structure(alias_elimination(sys))
    # check_consistency(sys) # currently not consistent due to u treated as states
    if sysr isa ODESystem
        sys = dae_index_lowering(sysr)
    end
    sysr = tearing(sys, simplify=false)
    # x = [map(eq->eq.lhs, observed(sys)); states(sys)]
    u = filter(u->isinput(u) && !is_bound(sysr, u), states(sysr))
    # TODO: current problem is that u above contains also internal inputs, i.e., already bound inputs. 
    x = filter(x->!isinput(x) && !isoutput(x), states(sysr))
    dynamics = [e.rhs for e in equations(sysr) if isdifferential(e.lhs)]
    A = jacobian(dynamics, x)
    B = jacobian(dynamics, u) 

    if numeric || x0 !== nothing
        num(x::Num) = x.val
        # sym2val is the operaing point we linearize around
        sym2val = ModelingToolkit.defaults(sysr)
        if x0 isa AbstractDict
            sym2val = merge(sym2val, x0) # NOTE: it's important to have x0 last for x0 to take precedence
        end
        A = substitute.(A, Ref(sym2val)) .|> num .|> Float64
        B = substitute.(B, Ref(sym2val)) .|> num .|> Float64
    end
    size(A,1) == size(A,2) || @warn "size(A,1) == size(A,2), This can happen if there are algebraic variables left in the system. `linearize` is very experimental and the provided system is unfortunately not supported."
    size(A,1) == size(B,1) || @warn "size(A,1) == size(B,1), This can happen if there are algebraic variables left in the system. `linearize` is very experimental and the provided system is unfortunately not supported."
    x = states(sysr, states(sysr)) # redo this to get the namespace right for later use
    u = filter(isinput, x)
    x = filter(!isinput, x)
    (; A, B, x, u)
end



@variables t
Dₜ = Differential(t)
## Tests


A = [0 1; 0 0]
B = [0; 1;;]
C = [1 0]
D = 0

# This system is already linear so linearize should return the same matrices
function StateSpace(A, B, C, D=0; x0=zeros(size(A,1)), name)
    nx = size(A,1)
    nu = size(B,2)
    ny = size(C,1)
    if B isa AbstractVector
        B = reshape(B, length(B), 1)
    end
    if D == 0
        D = zeros(ny, nu)
    end
    @variables x[1:nx](t)=x0 u[1:nu](t)=0 [input=true] y[1:ny](t)=C*x0 [output=true]
    x = collect(x) # https://github.com/JuliaSymbolics/Symbolics.jl/issues/379
    u = collect(u)
    y = collect(y)
    # @parameters A=A B=B C=C D=D # This is buggy
    eqs = [
        Dₜ.(x) .~ A*x .+ B*u
        y      .~ C*x .+ D*u
    ]
    ODESystem(eqs, t, name=name)
end
@named sys = StateSpace(A,B,C,D)
lin = linearize(sys, numeric=true)

@test lin.A == A
@test lin.B == B
@test lin.C == C
@test lin.D == D


## Test saturation
function Saturation(; y_max, y_min=y_max > 0 ? -y_max : -Inf, name)
    if !ModelingToolkit.isvariable(y_max)
        y_max ≥ y_min || throw(ArgumentError("y_min must be smaller than y_max"))
    end
    @variables u(t)=0 [input=true] y(t)=0 [output=true]
    @parameters y_max=y_max y_min=y_min
    ie = ModelingToolkit.IfElse.ifelse
    eqs = [
        # The equation below is equivalent to y ~ clamp(u, y_min, y_max)
        y ~ ie(u > y_max, y_max, ie( (y_min < u) & (u < y_max), u, y_min))
    ]
    ODESystem(eqs, t, name=name)
end

@named sat = Saturation(; y_max=1)
# inside the linear region, the function is identity
lin = linearize(sat, numeric=true)
@test isempty(lin.A) # there are no differential variables in this system
@test isempty(lin.B) # there are no differential variables in this system
@test isempty(lin.C) # there are no differential variables in this system
@test lin.D[] == 1


## Add differential equations
@variables x(t)=0 u(t)=0 [input=true] 
@named saturated_integrator = ODESystem([Dₜ(x) ~ u, x~sat.u], t, systems=[sat])

lin = linearize(saturated_integrator, numeric=true)
@test lin.A == [1;;]
@test lin.B == [1;;]
@test lin.C == [1;;]
@test lin.D == [0;;]

# in the test below, we linearize around a point outside the saturation's linear region, where it becomes a constant equation. Ideally, this constant should be made known in the object returned from linearize
lin = linearize(saturated_integrator, numeric=true, x0=Dict(sat.u => 2))
@test isempty(lin.A) # there are no differential variables left in the system
@test isempty(lin.B) # there are no differential variables left in the system
@test isempty(lin.C) # there are no differential variables left in the system
@test lin.D[] == 0
@test lin.output_offset == [y_max] # the output is the constant Y_max since the saturation has saturated