feat: [AI] also support passing a vector of time points

Author: AayushSabharwal

lib/ModelingToolkitBase/test/analysis_points.jl

@@ -220,6 +220,35 @@ if @isdefined(ModelingToolkit)
         @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

src/linearization.jl

@@ -1,10 +1,14 @@
 """
     $(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.
+Wraps an `ODESolution` and a time point (or vector of time points) `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 `t`. Algebraic
+variables are not set and will be determined by the initialization algorithm.
+
+When `t` is an `AbstractVector`, [`linearize`](@ref) calls [`linearization_function`](@ref)
+once and evaluates the linearization at each time point, returning vectors of matrices and
+extras.
 
 # Fields
 
@@ -16,7 +20,7 @@ struct LinearizationOpPoint{S <: SciMLBase.AbstractODESolution, T}
     """
     sol::S
     """
-    The time point at which to evaluate the solution.
+    The time point (or vector of time points) at which to evaluate the solution.
     """
     t::T
 end
@@ -37,6 +41,29 @@ function _build_op_from_solution(op::LinearizationOpPoint)
     return result
 end
 
+function _build_op_from_solution(op::LinearizationOpPoint{S, <:AbstractVector}) where {S}
+    sol_sys = MTKBase.indp_to_system(op.sol)
+    eqs = equations(sol_sys)
+    sts = unknowns(sol_sys)
+    # Find differential equation indices and extract parameters once — both are
+    # time-independent, so we only do this work once regardless of how many time
+    # points are requested.
+    diff_idxs = findall(isdiffeq, eqs)
+    param_vals = Dict{SymbolicT, SymbolicT}()
+    for p in parameters(sol_sys)
+        param_vals[p] = getp(op.sol, p)(op.sol)
+    end
+    # Interpolate once per time point to build the per-point operating-point dict.
+    return map(op.t) do ti
+        u = op.sol(ti)
+        result = copy(param_vals)
+        for i in diff_idxs
+            result[sts[i]] = u[i]
+        end
+        result
+    end
+end
+
 """
     lin_fun, simplified_sys = linearization_function(sys::AbstractSystem, inputs, outputs; simplify = false, initialize = true, initialization_solver_alg = nothing, kwargs...)
 
@@ -809,6 +836,13 @@ function linearize(
         op = Dict(), allow_input_derivatives = false,
         p = DiffEqBase.NullParameters()
     )
+    if op isa LinearizationOpPoint && op.t isa AbstractVector
+        ops = _build_op_from_solution(op)
+        results = map(zip(ops, op.t)) do (op_i, ti)
+            linearize(sys, lin_fun; t = ti, op = op_i, allow_input_derivatives, p)
+        end
+        return first.(results), last.(results)
+    end
     if op isa LinearizationOpPoint
         t = op.t
         op = _build_op_from_solution(op)
@@ -840,6 +874,21 @@ function linearize(
         zero_dummy_der = false,
         kwargs...
     )
+    if op isa LinearizationOpPoint && op.t isa AbstractVector
+        ops = _build_op_from_solution(op)
+        ts = op.t
+        # Build the linearization function once using the first operating point, then
+        # reuse it for all subsequent time points — this avoids redundant `mtkcompile`
+        # and Jacobian preparation work.
+        lin_fun, ssys = linearization_function(
+            sys, inputs, outputs;
+            zero_dummy_der, op = ops[1], t = ts[1], kwargs...
+        )
+        results = map(zip(ops, ts)) do (op_i, ti)
+            linearize(ssys, lin_fun; op = op_i, t = ti, allow_input_derivatives)
+        end
+        return first.(results), ssys, last.(results)
+    end
     if op isa LinearizationOpPoint
         t = op.t
         op = _build_op_from_solution(op)