refactor: use new structural_simplify in linearization

Author: vyudu

src/linearization.jl

@@ -58,9 +58,8 @@ 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...)
+    ssys = structural_simplify(sys; inputs, outputs, simplify, kwargs...)
+    diff_idxs, alge_idxs = eq_idxs(ssys)
     if zero_dummy_der
         dummyder = setdiff(unknowns(ssys), unknowns(sys))
         defs = Dict(x => 0.0 for x in dummyder)
@@ -87,9 +86,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 +112,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,12 +126,24 @@ 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
 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
+
 """
     $(TYPEDEF)
 
@@ -192,7 +203,7 @@ A callable struct which linearizes a system.
 $(TYPEDFIELDS)
 """
 struct LinearizationFunction{
-    DI <: AbstractVector{Int}, AI <: AbstractVector{Int}, II, P <: ODEProblem,
+    DI <: AbstractVector{Int}, AI <: AbstractVector{Int}, I, P <: ODEProblem,
     H, C, J1, J2, J3, J4, IA <: SciMLBase.DAEInitializationAlgorithm, IK}
     """
     The indexes of differential equations in the linearized system.
@@ -206,7 +217,7 @@ struct LinearizationFunction{
     The indexes of parameters in the linearized system which represent
     input variables.
     """
-    input_idxs::II
+    inputs::I
     """
     The number of unknowns in the linearized system.
     """
@@ -281,6 +292,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 +306,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],
@@ -482,9 +494,8 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
         outputs; simplify = false, 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...)
+    sys = structural_simplify(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

@@ -2486,21 +2486,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