refactor: require simplify system for linearization

Author: vyudu

docs/src/basics/Linearization.md

@@ -15,7 +15,7 @@ y &= Cx + Du
 \end{aligned}
 ```
 
-The `linearize` function expects the user to specify the inputs ``u`` and the outputs ``y`` using the syntax shown in the example below. The system model is *not* supposed to be simplified before calling `linearize`:
+The `linearize` function expects the user to specify the inputs ``u`` and the outputs ``y`` using the syntax shown in the example below.
 
 ## Example
 
@@ -29,7 +29,7 @@ eqs = [u ~ kp * (r - y) # P controller
        D(x) ~ -x + u    # First-order plant
        y ~ x]           # Output equation
 
-@named sys = ODESystem(eqs, t) # Do not call @mtkbuild when linearizing
+@mtkbuild sys = ODESystem(eqs, t) # Do not call @mtkbuild when linearizing
 matrices, simplified_sys = linearize(sys, [r], [y]) # Linearize from r to y
 matrices
 ```
@@ -45,10 +45,6 @@ using ModelingToolkit: inputs, outputs
 
     The model above has 4 variables but only three equations, there is no equation specifying the value of `r` since `r` is an input. This means that only unbalanced models can be linearized, or in other words, models that are balanced and can be simulated _cannot_ be linearized. To learn more about this, see [How to linearize a ModelingToolkit model (YouTube)](https://www.youtube.com/watch?v=-XOux-2XDGI&t=395s). Also see [ModelingToolkitStandardLibrary: Linear analysis](https://docs.sciml.ai/ModelingToolkit/stable/tutorials/linear_analysis/) for utilities that make linearization of completed models easier.
 
-!!! note "Un-simplified system"
-    
-    Linearization expects `sys` to be un-simplified, i.e., `structural_simplify` or `@mtkbuild` should not be called on the system before linearizing.
-
 ## Operating point
 
 The operating point to linearize around can be specified with the keyword argument `op` like this: `op = Dict(x => 1, r => 2)`. The operating point may include specification of unknown variables, input variables and parameters. For variables that are not specified in `op`, the default value specified in the model will be used if available, if no value is specified, an error is thrown.
@@ -79,14 +75,15 @@ eqs = [D(x) ~ v
        y.u ~ x]
 
 @named duffing = ODESystem(eqs, t, systems = [y, u], defaults = [u.u => 0])
+duffing = structural_simplify(duffing, inputs = [u.u], outputs = [y.u])
 
 # pass a constant value for `x`, since it is the variable we will change in operating points
-linfun, simplified_sys = linearization_function(duffing, [u.u], [y.u]; op = Dict(x => NaN));
+linfun = linearization_function(duffing, [u.u], [y.u]; op = Dict(x => NaN));
 
-println(linearize(simplified_sys, linfun; op = Dict(x => 1.0)))
-println(linearize(simplified_sys, linfun; op = Dict(x => 0.0)))
+println(linearize(duffing, linfun; op = Dict(x => 1.0)))
+println(linearize(duffing, linfun; op = Dict(x => 0.0)))
 
-@time linearize(simplified_sys, linfun; op = Dict(x => 0.0))
+@time linearize(duffing, linfun; op = Dict(x => 0.0))
 
 nothing # hide
 ```

src/linearization.jl

@@ -1,5 +1,5 @@
 """
-    lin_fun, simplified_sys = linearization_function(sys::AbstractSystem, inputs, outputs; simplify = false, initialize = true, initialization_solver_alg = TrustRegion(), kwargs...)
+    lin_fun, simplified_sys = linearization_function(sys::AbstractSystem, inputs, outputs; initialize = true, initialization_solver_alg = TrustRegion(), kwargs...)
 
 Return a function that linearizes the system `sys`. The function [`linearize`](@ref) provides a higher-level and easier to use interface.
 
@@ -22,7 +22,6 @@ The `simplified_sys` has undergone [`structural_simplify`](@ref) and had any occ
   - `sys`: An [`ODESystem`](@ref). This function will automatically apply simplification passes on `sys` and return the resulting `simplified_sys`.
   - `inputs`: A vector of variables that indicate the inputs of the linearized input-output model.
   - `outputs`: A vector of variables that indicate the outputs of the linearized input-output model.
-  - `simplify`: Apply simplification in tearing.
   - `initialize`: If true, a check is performed to ensure that the operating point is consistent (satisfies algebraic equations). If the op is not consistent, initialization is performed.
   - `initialization_solver_alg`: A NonlinearSolve algorithm to use for solving for a feasible set of state and algebraic variables that satisfies the specified operating point.
   - `autodiff`: An `ADType` supported by DifferentiationInterface.jl to use for calculating the necessary jacobians. Defaults to using `AutoForwardDiff()`
@@ -31,7 +30,7 @@ The `simplified_sys` has undergone [`structural_simplify`](@ref) and had any occ
 See also [`linearize`](@ref) which provides a higher-level interface.
 """
 function linearization_function(sys::AbstractSystem, inputs,
-        outputs; simplify = false,
+        outputs;
         initialize = true,
         initializealg = nothing,
         initialization_abstol = 1e-5,
@@ -58,16 +57,7 @@ function linearization_function(sys::AbstractSystem, inputs,
     outputs = mapreduce(vcat, outputs; init = []) do var
         symbolic_type(var) == ArraySymbolic() ? collect(var) : [var]
     end
-    ssys, diff_idxs, alge_idxs, input_idxs = io_preprocessing(sys, inputs, outputs;
-        simplify,
-        kwargs...)
-    if zero_dummy_der
-        dummyder = setdiff(unknowns(ssys), unknowns(sys))
-        defs = Dict(x => 0.0 for x in dummyder)
-        @set! ssys.defaults = merge(defs, defaults(ssys))
-        op = merge(defs, op)
-    end
-    sys = ssys
+    diff_idxs, alge_idxs = eq_idxs(sys)
 
     if initializealg === nothing
         initializealg = initialize ? OverrideInit() : NoInit()
@@ -87,9 +77,9 @@ function linearization_function(sys::AbstractSystem, inputs,
 
     p = parameter_values(prob)
     t0 = current_time(prob)
-    inputvals = [p[idx] for idx in input_idxs]
+    inputvals = [prob.ps[i] for i in inputs]
 
-    hp_fun = let fun = h, setter = setp_oop(sys, input_idxs)
+    hp_fun = let fun = h, setter = setp_oop(sys, inputs)
         function hpf(du, input, u, p, t)
             p = setter(p, input)
             fun(du, u, p, t)
@@ -113,7 +103,7 @@ function linearization_function(sys::AbstractSystem, inputs,
         # observed function is a `GeneratedFunctionWrapper` with iip component
         h_jac = PreparedJacobian{true}(h, similar(prob.u0, size(outputs)), autodiff,
             prob.u0, DI.Constant(p), DI.Constant(t0))
-        pf_fun = let fun = prob.f, setter = setp_oop(sys, input_idxs)
+        pf_fun = let fun = prob.f, setter = setp_oop(sys, inputs)
             function pff(du, input, u, p, t)
                 p = setter(p, input)
                 SciMLBase.ParamJacobianWrapper(fun, t, u)(du, p)
@@ -127,10 +117,22 @@ function linearization_function(sys::AbstractSystem, inputs,
     end
 
     lin_fun = LinearizationFunction(
-        diff_idxs, alge_idxs, input_idxs, length(unknowns(sys)),
+        diff_idxs, alge_idxs, length(unknowns(sys)),
         prob, h, u0 === nothing ? nothing : similar(u0), uf_jac, h_jac, pf_jac,
         hp_jac, initializealg, initialization_kwargs)
-    return lin_fun, sys
+    return lin_fun
+end
+
+function eq_idxs(sys::AbstractSystem)
+    eqs = equations(sys)
+    alg_start_idx = findfirst(!isdiffeq, eqs)
+    if alg_start_idx === nothing
+        alg_start_idx = length(eqs) + 1
+    end
+    diff_idxs = 1:(alg_start_idx - 1)
+    alge_idxs = alg_start_idx:length(eqs)
+
+    diff_idxs, alge_idxs
 end
 
 """
@@ -206,7 +208,7 @@ struct LinearizationFunction{
     The indexes of parameters in the linearized system which represent
     input variables.
     """
-    input_idxs::II
+    inputs::II
     """
     The number of unknowns in the linearized system.
     """
@@ -281,6 +283,7 @@ function (linfun::LinearizationFunction)(u, p, t)
     end
 
     fun = linfun.prob.f
+    input_vals = [linfun.prob.ps[i] for i in linfun.inputs]
     if u !== nothing # Handle systems without unknowns
         linfun.num_states == length(u) ||
             error("Number of unknown variables ($(linfun.num_states)) does not match the number of input unknowns ($(length(u)))")
@@ -294,15 +297,15 @@ function (linfun::LinearizationFunction)(u, p, t)
         end
         fg_xz = linfun.uf_jac(u, DI.Constant(p), DI.Constant(t))
         h_xz = linfun.h_jac(u, DI.Constant(p), DI.Constant(t))
-        fg_u = linfun.pf_jac([p[idx] for idx in linfun.input_idxs],
+        fg_u = linfun.pf_jac(input_vals,
             DI.Constant(u), DI.Constant(p), DI.Constant(t))
     else
         linfun.num_states == 0 ||
             error("Number of unknown variables (0) does not match the number of input unknowns ($(length(u)))")
         fg_xz = zeros(0, 0)
-        h_xz = fg_u = zeros(0, length(linfun.input_idxs))
+        h_xz = fg_u = zeros(0, length(linfun.inputs))
     end
-    h_u = linfun.hp_jac([p[idx] for idx in linfun.input_idxs],
+    h_u = linfun.hp_jac(input_vals,
         DI.Constant(u), DI.Constant(p), DI.Constant(t))
     (f_x = fg_xz[linfun.diff_idxs, linfun.diff_idxs],
         f_z = fg_xz[linfun.diff_idxs, linfun.alge_idxs],
@@ -461,7 +464,7 @@ function CommonSolve.solve(prob::LinearizationProblem; allow_input_derivatives =
 end
 
 """
-    (; A, B, C, D), simplified_sys = linearize_symbolic(sys::AbstractSystem, inputs, outputs; simplify = false, allow_input_derivatives = false, kwargs...)
+    (; A, B, C, D), simplified_sys = linearize_symbolic(sys::AbstractSystem, inputs, outputs; allow_input_derivatives = false, kwargs...)
 
 Similar to [`linearize`](@ref), but returns symbolic matrices `A,B,C,D` rather than numeric. While `linearize` uses ForwardDiff to perform the linearization, this function uses `Symbolics.jacobian`.
 
@@ -479,12 +482,10 @@ y &= h(x, z, u)
 where `x` are differential unknown variables, `z` algebraic variables, `u` inputs and `y` outputs.
 """
 function linearize_symbolic(sys::AbstractSystem, inputs,
-        outputs; simplify = false, allow_input_derivatives = false,
+        outputs; allow_input_derivatives = false,
         eval_expression = false, eval_module = @__MODULE__,
         kwargs...)
-    sys, diff_idxs, alge_idxs, input_idxs = io_preprocessing(
-        sys, inputs, outputs; simplify,
-        kwargs...)
+    diff_idxs, alge_idxs = eq_idxs(sys)
     sts = unknowns(sys)
     t = get_iv(sys)
     ps = parameters(sys; initial_parameters = true)

src/systems/abstractsystem.jl

@@ -2484,21 +2484,6 @@ function eliminate_constants(sys::AbstractSystem)
     return sys
 end
 
-function io_preprocessing(sys::AbstractSystem, inputs,
-        outputs; simplify = false, kwargs...)
-    sys, input_idxs = structural_simplify(sys, (inputs, outputs); simplify, kwargs...)
-
-    eqs = equations(sys)
-    alg_start_idx = findfirst(!isdiffeq, eqs)
-    if alg_start_idx === nothing
-        alg_start_idx = length(eqs) + 1
-    end
-    diff_idxs = 1:(alg_start_idx - 1)
-    alge_idxs = alg_start_idx:length(eqs)
-
-    sys, diff_idxs, alge_idxs, input_idxs
-end
-
 @latexrecipe function f(sys::AbstractSystem)
     return latexify(equations(sys))
 end