feat: [AI] allow passing an operating point to linearize using a solution

Author: AayushSabharwal

lib/ModelingToolkitBase/test/analysis_points.jl

@@ -204,6 +204,24 @@ if @isdefined(ModelingToolkit)
         @test matrices.D[] == 0
     end
 
+    @testset "LinearizationOpPoint" begin
+        # sys here is the two-analysis-point system (plant_input + plant_output)
+        # Simulate to obtain a solution, then verify that linearizing at t=0 via
+        # LinearizationOpPoint gives the same result as the default operating point.
+        ssys_solve = mtkcompile(sys)
+        prob = ODEProblem(ssys_solve, [P.x => 0.0], (0.0, 1.0))
+        sol = solve(prob, Rodas5())
+        matrices_ref, _ = linearize(sys, sys.plant_input, sys.plant_output)
+        matrices_op, _ = linearize(
+            sys, sys.plant_input, sys.plant_output;
+            op = ModelingToolkit.LinearizationOpPoint(sol, 0.0)
+        )
+        @test matrices_op.A ≈ matrices_ref.A
+        @test matrices_op.B ≈ matrices_ref.B
+        @test matrices_op.C ≈ matrices_ref.C
+        @test matrices_op.D ≈ matrices_ref.D
+    end
+
     @testset "Complicated model" begin
         # Parameters
         m1 = 1

src/ModelingToolkit.jl

@@ -158,7 +158,7 @@ end
 export SemilinearODEFunction, SemilinearODEProblem
 export alias_elimination
 export linearize, linearization_function,
-    LinearizationProblem, linearization_ap_transform
+    LinearizationProblem, LinearizationOpPoint, linearization_ap_transform
 export solve
 export map_variables_to_equations, substitute_component
 

src/linearization.jl

@@ -1,3 +1,42 @@
+"""
+    $(TYPEDEF)
+
+Wraps an `ODESolution` and a time point `t`. When passed as `op` to [`linearize`](@ref),
+an operating point is constructed from the values of differential state variables and
+parameters of `sol` evaluated at time `t`. Algebraic variables are not set and will be
+determined by the initialization algorithm.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
+struct LinearizationOpPoint{S <: SciMLBase.AbstractODESolution, T}
+    """
+    The solution to extract operating point values from.
+    """
+    sol::S
+    """
+    The time point at which to evaluate the solution.
+    """
+    t::T
+end
+
+function _build_op_from_solution(op::LinearizationOpPoint)
+    sol_sys = MTKBase.indp_to_system(op.sol)
+    eqs = equations(sol_sys)
+    sts = unknowns(sol_sys)
+    u = op.sol(op.t)
+    result = Dict{SymbolicT, SymbolicT}()
+    for (i, eq) in enumerate(eqs)
+        isdiffeq(eq) || continue
+        result[sts[i]] = u[i]
+    end
+    for p in parameters(sol_sys)
+        result[p] = getp(op.sol, p)(op.sol)
+    end
+    return result
+end
+
 """
     lin_fun, simplified_sys = linearization_function(sys::AbstractSystem, inputs, outputs; simplify = false, initialize = true, initialization_solver_alg = nothing, kwargs...)
 
@@ -770,6 +809,10 @@ function linearize(
         op = Dict(), allow_input_derivatives = false,
         p = DiffEqBase.NullParameters()
     )
+    if op isa LinearizationOpPoint
+        t = op.t
+        op = _build_op_from_solution(op)
+    end
     prob = LinearizationProblem(lin_fun, t)
     op = as_atomic_dict_with_defaults(Dict{SymbolicT, SymbolicT}(op), COMMON_NOTHING)
     evaluate_varmap!(op, keys(op))
@@ -797,6 +840,10 @@ function linearize(
         zero_dummy_der = false,
         kwargs...
     )
+    if op isa LinearizationOpPoint
+        t = op.t
+        op = _build_op_from_solution(op)
+    end
     lin_fun,
         ssys = linearization_function(
         sys,