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

lib/ModelingToolkitBase/src/systems/analysis_points.jl

@@ -459,6 +459,7 @@ function get_analysis_variable(var, name, iv; perturb = true)
         pvar = unwrap(only(@variables $name(iv)::T))
         default = zero(T)
     end
+    pvar = SU.setmetadata(pvar, AnalysisVariable, true)
     return pvar, default
 end
 

lib/ModelingToolkitBase/src/variables.jl

@@ -45,6 +45,12 @@ struct VariableMisc end
 struct VariableUnshifted end
 struct VariableShift end
 struct VariableTimeDomain end
+"""
+    $TYPEDEF
+
+Metadata key used to mark variables introduced by analysis point transformations.
+"""
+struct AnalysisVariable end
 
 Symbolics.option_to_metadata_type(::Val{:unit}) = VariableUnit
 Symbolics.option_to_metadata_type(::Val{:connect}) = VariableConnectType

lib/ModelingToolkitBase/test/analysis_points.jl

@@ -204,6 +204,53 @@ 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
+
+        # Vector of time points: linearization_function is built once and reused.
+        ts = [0.0, 0.5, 1.0]
+        mats_vec, _, extras_vec = linearize(
+            sys, sys.plant_input, sys.plant_output;
+            op = ModelingToolkit.LinearizationOpPoint(sol, ts)
+        )
+        @test length(mats_vec) == 3
+        @test length(extras_vec) == 3
+        # The system is linear so all operating points yield the same A,B,C,D.
+        for mats_t in mats_vec
+            @test mats_t.A ≈ matrices_ref.A
+            @test mats_t.B ≈ matrices_ref.B
+            @test mats_t.C ≈ matrices_ref.C
+            @test mats_t.D ≈ matrices_ref.D
+        end
+        # Two-arg form: linearize(ssys, lin_fun; op=LinearizationOpPoint(sol, ts))
+        lin_fun, ssys_lin = linearization_function(sys, sys.plant_input, sys.plant_output)
+        mats_vec2, extras_vec2 = linearize(
+            ssys_lin, lin_fun;
+            op = ModelingToolkit.LinearizationOpPoint(sol, ts)
+        )
+        @test length(mats_vec2) == 3
+        for (m1, m2) in zip(mats_vec, mats_vec2)
+            @test m1.A ≈ m2.A
+            @test m1.B ≈ m2.B
+            @test m1.C ≈ m2.C
+            @test m1.D ≈ m2.D
+        end
+    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