Trajectory linearization: avoid per-timepoint recompilation; require op for loop-opening parameters

lib/ModelingToolkitBase/test/analysis_points.jl

@@ -249,6 +249,58 @@ if @isdefined(ModelingToolkit)
             @test m1.C ≈ m2.C
             @test m1.D ≈ m2.D
         end
+        # The vector-of-time-points path (which builds the problem and setters once and
+        # updates the operating point in place) must agree with calling the scalar
+        # LinearizationOpPoint form at each time point individually.
+        for (i, ti) in enumerate(ts)
+            mats_i, _ = linearize(
+                sys, sys.plant_input, sys.plant_output;
+                op = ModelingToolkit.LinearizationOpPoint(sol, ti)
+            )
+            @test mats_vec[i].A ≈ mats_i.A
+            @test mats_vec[i].B ≈ mats_i.B
+            @test mats_vec[i].C ≈ mats_i.C
+            @test mats_vec[i].D ≈ mats_i.D
+        end
+    end
+
+    @testset "loop_openings require an operating point" begin
+        # Opening a loop turns the opened signal into a parameter whose operating-point
+        # value is not implied by the solution. It must be provided explicitly, otherwise
+        # `linearize` errors instead of silently using a stale/default value.
+        # Linearize plant_input -> P.output.u while opening the `plant_output` analysis
+        # point (so the opened AP is neither the input nor the output).
+        ssys_solve = mtkcompile(sys)
+        prob = ODEProblem(ssys_solve, [P.x => 0.0], (0.0, 1.0))
+        sol = solve(prob, Rodas5())
+        ts = [0.0, 0.5, 1.0]
+
+        lf_lo, _ = linearization_function(
+            sys, sys.plant_input, P.output.u; loop_openings = [sys.plant_output]
+        )
+        lops = lf_lo.loop_opening_params
+        @test !isempty(lops)
+
+        # Not providing the loop-opening parameter errors (both vector and scalar paths).
+        @test_throws Exception linearize(
+            sys, sys.plant_input, P.output.u;
+            op = ModelingToolkit.LinearizationOpPoint(sol, ts),
+            loop_openings = [sys.plant_output]
+        )
+        @test_throws Exception linearize(
+            sys, sys.plant_input, P.output.u;
+            op = ModelingToolkit.LinearizationOpPoint(sol, 0.0),
+            loop_openings = [sys.plant_output]
+        )
+
+        # Providing it via the `op` keyword of `LinearizationOpPoint` works.
+        op_lo = Dict(p => 0.0 for p in lops)
+        mats_lo, _, _ = linearize(
+            sys, sys.plant_input, P.output.u;
+            op = ModelingToolkit.LinearizationOpPoint(sol, ts; op = op_lo),
+            loop_openings = [sys.plant_output]
+        )
+        @test length(mats_lo) == length(ts)
     end
 
     @testset "Complicated model" begin

src/linearization.jl

@@ -10,11 +10,17 @@ When `t` is an `AbstractVector`, [`linearize`](@ref) calls [`linearization_funct
 once and evaluates the linearization at each time point, returning vectors of matrices and
 extras.
 
+The `op` keyword argument provides additional operating-point values that are merged into
+the solution-derived operating point at every time point (taking precedence). This is how
+values for variables that are not present in `sol` are supplied — in particular the
+parameters created by `loop_openings`, e.g.
+`LinearizationOpPoint(sol, t; op = Dict(opened_signal => 0))`.
+
 # Fields
 
 $(TYPEDFIELDS)
 """
-struct LinearizationOpPoint{S <: SciMLBase.AbstractODESolution, T}
+struct LinearizationOpPoint{S <: SciMLBase.AbstractODESolution, T, D <: AbstractDict}
     """
     The solution to extract operating point values from.
     """
@@ -23,6 +29,15 @@ struct LinearizationOpPoint{S <: SciMLBase.AbstractODESolution, T}
     The time point (or vector of time points) at which to evaluate the solution.
     """
     t::T
+    """
+    Additional operating-point values merged into the solution-derived operating point at
+    each time point (takes precedence over solution-derived values).
+    """
+    op::D
+end
+
+function LinearizationOpPoint(sol::SciMLBase.AbstractODESolution, t; op = Dict{SymbolicT, SymbolicT}())
+    return LinearizationOpPoint(sol, t, op)
 end
 
 function _build_op_from_solution(op::LinearizationOpPoint)
@@ -38,6 +53,9 @@ function _build_op_from_solution(op::LinearizationOpPoint)
     for p in parameters(sol_sys)
         result[p] = getp(op.sol, p)(op.sol)
     end
+    for (k, v) in op.op
+        result[unwrap(k)] = v
+    end
     return result
 end
 
@@ -54,12 +72,14 @@ function _build_op_from_solution(op::LinearizationOpPoint{S, <:AbstractVector})
         param_vals[p] = getp(op.sol, p)(op.sol)
     end
     # Interpolate once per time point to build the per-point operating-point dict.
+    extra_op = Dict{SymbolicT, SymbolicT}(unwrap(k) => v for (k, v) in op.op)
     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
+        merge!(result, extra_op)
         result
     end
 end
@@ -124,6 +144,7 @@ function linearization_function(
         warn_empty_op = true,
         t = 0.0,
         ignore_system_initial_conditions = false,
+        loop_opening_params = SymbolicT[],
         kwargs...
     )
     op = Dict(op)
@@ -243,7 +264,7 @@ function linearization_function(
     lin_fun = LinearizationFunction(
         diff_idxs, alge_idxs, input_getter, length(inputs), length(unknowns(sys)),
         prob, h, u0 === nothing ? nothing : similar(u0, T), uf_jac, h_jac, pf_jac,
-        hp_jac, initializealg, initialization_kwargs
+        hp_jac, initializealg, initialization_kwargs, collect(SymbolicT, loop_opening_params)
     )
     return lin_fun, sys
 end
@@ -409,6 +430,12 @@ struct LinearizationFunction{
     Keyword arguments to be passed to `SciMLBase.get_initial_values`.
     """
     initialize_kwargs::IK
+    """
+    Variables turned into parameters by `loop_openings`. Their operating-point values are
+    not implied by the rest of the system, so they must be provided explicitly in the `op`
+    passed to `linearize`; otherwise an error is thrown.
+    """
+    loop_opening_params::Vector{SymbolicT}
 end
 
 SymbolicIndexingInterface.symbolic_container(f::LinearizationFunction) = f.prob
@@ -908,11 +935,10 @@ function linearize(
     prob = LinearizationProblem(lin_fun, t)
     op = as_atomic_dict_with_defaults(Dict{SymbolicT, SymbolicT}(op), COMMON_NOTHING)
     evaluate_varmap!(op, keys(op))
+    _check_loop_opening_op(lin_fun.loop_opening_params, op)
     for (k, v) in op
         isequal(v, COMMON_NOTHING) && continue
-        if symbolic_type(v) != NotSymbolic() || is_array_of_symbolics(v)
-            v = getu(prob, v)(prob)
-        end
+        v = _resolve_op_value(prob, v)
         if is_parameter(prob, Initial(k))
             setu(prob, Initial(k))(prob, v)
         else
@@ -926,12 +952,75 @@ function linearize(
     return solve(prob; allow_input_derivatives)
 end
 
+# Resolve an operating-point value to something that can be passed to a setter.
+# A numeric value stored in a `Dict{SymbolicT, SymbolicT}` is wrapped as a symbolic
+# constant; unwrap it to a plain number/array. Routing a constant through `getu` would
+# build a fresh `RuntimeGeneratedFunction` (with the value baked into the generated code)
+# on every call, which dominates runtime when linearizing along a trajectory. Only
+# genuinely symbolic values (expressions referencing the problem state) need `getu`.
+function _resolve_op_value(prob, v)
+    if SU.isconst(v)
+        return SU.unwrap_const(v)
+    elseif symbolic_type(v) != NotSymbolic() || is_array_of_symbolics(v)
+        return getu(prob, v)(prob)
+    else
+        return v
+    end
+end
+
+# Variables turned into parameters by `loop_openings` have no operating-point value implied
+# by the rest of the system, so they must be supplied explicitly in `op`. Error (rather than
+# silently using a stale/default value) if any of them is missing.
+function _check_loop_opening_op(loop_opening_params, op)
+    isempty(loop_opening_params) && return nothing
+    missing_params = SymbolicT[]
+    for p in loop_opening_params
+        v = get(op, p, COMMON_NOTHING)
+        isequal(v, COMMON_NOTHING) && push!(missing_params, p)
+    end
+    isempty(missing_params) && return nothing
+    params_str = join(string.(missing_params), ", ")
+    error("""
+          The operating point does not provide values for the loop-opening parameter(s): \
+          $(params_str). When `loop_openings` is used, the opened signals become \
+          parameters whose operating-point values are not implied by the rest of the \
+          system, so they must be provided explicitly in `op` (e.g. set to zero). When \
+          linearizing along a trajectory with `LinearizationOpPoint`, pass them via its \
+          `op` keyword argument: `LinearizationOpPoint(sol, t; op = Dict(signal => value))`.
+          """)
+end
+
 function __linearize_multiple_op_barrier(ssys, lin_fun; ops, ts, allow_input_derivatives)
     T = eltype(lin_fun.prob.u0)
     results = @NamedTuple{A::Matrix{T}, B::Matrix{T}, C::Matrix{T}, D::Matrix{T}}[]
     xpts = @NamedTuple{x::typeof(lin_fun.prob.u0), p::typeof(lin_fun.prob.p), t::typeof(lin_fun.prob.tspan[1])}[]
+    isempty(ops) && return results, xpts
+
+    # Build the linearization problem once and reuse it across all time points, mutating
+    # only the operating point and `.t`. This avoids reconstructing the problem and
+    # rebuilding the symbolic setters on every iteration. The set of operating-point keys
+    # is identical across time points (only the values change), so resolve the first op to
+    # obtain the keys and build the setters once.
+    prob = LinearizationProblem(lin_fun, ts[1])
+    op1 = as_atomic_dict_with_defaults(Dict{SymbolicT, SymbolicT}(ops[1]), COMMON_NOTHING)
+    evaluate_varmap!(op1, keys(op1))
+    # The op keys are identical across time points, so checking the first one suffices.
+    _check_loop_opening_op(lin_fun.loop_opening_params, op1)
+    op_keys = collect(keys(op1))
+    setters = map(op_keys) do k
+        is_parameter(prob, Initial(k)) ? setu(prob, Initial(k)) : setu(prob, k)
+    end
+
     for (op, t) in zip(ops, ts)
-        res, xpt = linearize(ssys, lin_fun; op, t, allow_input_derivatives)::Tuple{eltype(results), eltype(xpts)}
+        op = as_atomic_dict_with_defaults(Dict{SymbolicT, SymbolicT}(op), COMMON_NOTHING)
+        evaluate_varmap!(op, keys(op))
+        for (setter, k) in zip(setters, op_keys)
+            v = get(op, k, COMMON_NOTHING)
+            isequal(v, COMMON_NOTHING) && continue
+            setter(prob, _resolve_op_value(prob, v))
+        end
+        prob.t = t
+        res, xpt = solve(prob; allow_input_derivatives)::Tuple{eltype(results), eltype(xpts)}
         push!(results, res)
         push!(xpts, xpt)
     end

src/systems/analysis_points.jl

@@ -1,13 +1,18 @@
 """
     $(TYPEDSIGNATURES)
 
-Given a list of analysis points, break the connection for each and set the output to zero.
+Given a list of analysis points, break the connection for each. The output of each broken
+analysis point is turned into a parameter of the system. Returns the modified system and a
+vector of the variables that were turned into parameters (the "loop-opening parameters"),
+whose operating-point values must be supplied by the user when linearizing.
 """
 function handle_loop_openings(sys::AbstractSystem, aps)
+    loop_opening_params = SymbolicT[]
     for ap in canonicalize_ap(sys, aps)
         sys, (d_vs,) = apply_transformation(Break(ap, true), sys)
+        append!(loop_opening_params, d_vs)
     end
-    return sys
+    return sys, loop_opening_params
 end
 
 const DOC_LOOP_OPENINGS = """
@@ -41,7 +46,7 @@ function get_linear_analysis_function(
     )
     dus = SymbolicT[]
     us = SymbolicT[]
-    sys = handle_loop_openings(sys, loop_openings)
+    sys, loop_opening_params = handle_loop_openings(sys, loop_openings)
     aps = canonicalize_ap(sys, aps)
     for ap in aps
         sys, (du, u) = apply_transformation(transform(ap), sys)
@@ -58,7 +63,7 @@ function get_linear_analysis_function(
             append!(us, u)
         end
     end
-    return linearization_function(system_modifier(sys), dus, us; kwargs...)
+    return linearization_function(system_modifier(sys), dus, us; loop_opening_params, kwargs...)
 end
 """
     $(TYPEDSIGNATURES)
@@ -135,6 +140,8 @@ Returns
 - `sys`: The transformed system.
 - `input_vars`: A vector of input variables corresponding to the input analysis points.
 - `output_vars`: A vector of output variables corresponding to the output analysis points.
+- `loop_opening_params`: A vector of the variables that were turned into parameters by the
+  `loop_openings` (their operating-point values must be supplied when linearizing).
 """
 function linearization_ap_transform(
         sys,
@@ -166,20 +173,22 @@ function linearization_ap_transform(
         end
         push!(output_vars, output_var)
     end
-    sys = handle_loop_openings(sys, map(AnalysisPoint, collect(loop_openings)))
-    return sys, input_vars, output_vars
+    sys, loop_opening_params = handle_loop_openings(sys, map(AnalysisPoint, collect(loop_openings)))
+    return sys, input_vars, output_vars, loop_opening_params
 end
 
 function linearization_function(
         sys::AbstractSystem,
         inputs::Union{Symbol, Vector{Symbol}, AnalysisPoint, Vector{AnalysisPoint}},
         outputs; loop_openings = [], system_modifier = identity, kwargs...
     )
-    sys, input_vars,
-        output_vars = linearization_ap_transform(
+    sys, input_vars, output_vars,
+        loop_opening_params = linearization_ap_transform(
         sys, inputs, outputs, loop_openings
     )
-    return linearization_function(system_modifier(sys), input_vars, output_vars; kwargs...)
+    return linearization_function(
+        system_modifier(sys), input_vars, output_vars; loop_opening_params, kwargs...
+    )
 end
 
 @doc """