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

[AayushSabharwal]

Close #4159

[baggepinnen]

I'm testing this PR from ControlSystemsMTK.jl and hit a dispatch ambiguity in _build_op_from_solution when calling with a vector of time points.

The two methods are:

# Method 1 (line 28): scalar
_build_op_from_solution(op::LinearizationOpPoint)

# Method 2 (line 44): vector
_build_op_from_solution(op::LinearizationOpPoint{S, <:AbstractVector}) where {S}

Method 1 implicitly constrains S <: AbstractODESolution (from the struct definition), while method 2 uses an unconstrained where {S}. Julia sees these as ambiguous for LinearizationOpPoint{ODESolution{...}, Vector{Float64}} since neither method is strictly more specific than the other.

[baggepinnen]

I've been exercising this branch through ControlSystemsMTK.trajectory_ss (linearizing a gain-scheduled closed loop at ~800 points along a trajectory) and hit two performance issues and one correctness issue in the multi-timepoint path. Details below with file/line references against this branch.

1. Operating-point values are boxed into SymbolicT, forcing per-point recompilation

_build_op_from_solution builds the op as a Dict{SymbolicT, SymbolicT}:

# src/linearization.jl
result = Dict{SymbolicT, SymbolicT}()          # line 33 (scalar method)
param_vals = Dict{SymbolicT, SymbolicT}()      # line 52 (vector method)
...
result[sts[i]] = u[i]                          # u[i]::Float64 gets widened to a symbolic constant
result[p] = getp(op.sol, p)(op.sol)            # numeric, also widened

Because the value type is SymbolicT, every numeric value is converted to a symbolic constant on insertion. The per-point linearize(sys, lin_fun::LinearizationFunction; t, op, ...) then takes the symbolic branch for every op entry:

# src/linearization.jl, ~line 911-915
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)   # taken for every entry, because every value looks symbolic
    end
    ...

getu(prob, v) builds an observed-function RuntimeGeneratedFunction keyed on the value v, and since each time point has distinct values, fresh codegen runs at essentially every point. Measured on a 7-state closed loop, 21 op entries, 800 time points:

• One-time linearization_function: ~0.09 s (fine)
• The 800-point loop: 273 s, 92% compilation time, ~340 ms/point
• Replaying the same 40 operating points: 268 ms/pt cold → 33.8 ms/pt warm — i.e. the cost tracks the values, confirming value-dependent recompilation rather than warmup.
• A trivial 3-entry model does not trigger this and runs at ~0.4 ms/pt.

Suggested fix: keep the operating-point values numeric (don't widen to SymbolicT) so the getu branch is skipped for plain numbers.

2. __linearize_multiple_op_barrier reconstructs the problem and setters every iteration

The vector path correctly builds lin_fun once, but then:

# src/linearization.jl, ~line 929-939
function __linearize_multiple_op_barrier(ssys, lin_fun; ops, ts, allow_input_derivatives)
    for (op, t) in zip(ops, ts)
        res, xpt = linearize(ssys, lin_fun; op, t, allow_input_derivatives)  # high-level call
        ...

Each high-level linearize(ssys, lin_fun; op, t) constructs a fresh LinearizationProblem(lin_fun, t) and rebuilds the setu/getu accessors from scratch. The LinearizationProblem docstring (~line 543) already says it "can be symbolically indexed to efficiently update the operating point for multiple linearizations in a loop" with the .t field mutated — which is exactly the fast path the loop should use: build the LinearizationProblem once, cache the setters for the op keys, update values + mutate .t, then solve. Combined with (1) this should remove essentially all of the per-point compilation.

3. Opened-loop parameters are silently frozen across the trajectory (correctness)

When loop_openings is used, handle_loop_openings / Break(ap, true) turns the opened signal into a parameter of the linearized system that does not exist in the solution. _build_op_from_solution only extracts the solution system's states and parameters, so these opened-loop parameters never get a per-time-point value — they stay frozen at whatever value the function was built with (ops[1], i.e. t[1]) for the entire trajectory.

This produces silently wrong results. Concretely: linearizing a gain-scheduled controller r -> u with loop_openings=[y, v] vs loop_openings=[u] gives identical systems at a single time point, but over 800 points they diverge by ~19% in frequency response, because the opened scheduling signal is pinned to its t[1] value rather than tracking the trajectory.

Rather than silently using a stale value, I think linearize should error when opened-loop parameters are not provided in the op, so the user is forced to specify their values explicitly (e.g. set the opened signals to 0). That keeps the semantics explicit and avoids the silent-wrong-answer failure mode.

[github-actions]

Benchmark Results (Julia vlts)

Time benchmarks

	master	a0d4896...	master / a0d4896...
ODEProblem	6.3 ± 0.93 ms	5.64 ± 0.47 ms	1.12 ± 0.19
init	0.0695 ± 0.01 ms	0.0649 ± 0.02 ms	1.07 ± 0.36
large_parameter_init/ODEProblem	28.9 ± 7.8 ms	27.2 ± 4.1 ms	1.06 ± 0.33
large_parameter_init/init	0.0882 ± 0.043 ms	0.068 ± 0.021 ms	1.3 ± 0.75
mtkcompile	9.08 ± 0.96 ms	8.34 ± 0.77 ms	1.09 ± 0.15
sparse_analytical_jacobian/ODEProblem	25.7 ± 2.3 ms	24.4 ± 2.5 ms	1.06 ± 0.15
sparse_analytical_jacobian/f_iip	0.06 ± 0 μs	0.06 ± 0 μs	1 ± 0
sparse_analytical_jacobian/f_oop	0.3 ± 0.012 ms	0.297 ± 0.01 ms	1.01 ± 0.052
time_to_load	5.04 ± 0.094 s	4.69 ± 0.026 s	1.07 ± 0.021

Memory benchmarks

	master	a0d4896...	master / a0d4896...
ODEProblem	31.5 k allocs: 1.93 MB	31.5 k allocs: 1.95 MB	0.992
init	0.416 k allocs: 0.0697 MB	0.416 k allocs: 0.0697 MB	1
large_parameter_init/ODEProblem	0.316 M allocs: 11.5 MB	0.316 M allocs: 11.5 MB	1
large_parameter_init/init	0.604 k allocs: 0.172 MB	0.604 k allocs: 0.172 MB	1
mtkcompile	0.0596 M allocs: 3.45 MB	0.0596 M allocs: 3.45 MB	1
sparse_analytical_jacobian/ODEProblem	0.196 M allocs: 8.31 MB	0.196 M allocs: 8.31 MB	1
sparse_analytical_jacobian/f_iip	0 allocs: 0 B	0 allocs: 0 B
sparse_analytical_jacobian/f_oop	0.634 k allocs: 19.6 kB	0.634 k allocs: 19.6 kB	1
time_to_load	0.153 k allocs: 14.6 kB	0.153 k allocs: 14.6 kB	1

[github-actions]

Benchmark Results (Julia v1)

Time benchmarks

	master	a0d4896...	master / a0d4896...
ODEProblem	6.42 ± 1.2 ms	5.09 ± 0.47 ms	1.26 ± 0.26
init	0.0542 ± 0.0082 ms	0.0526 ± 0.0079 ms	1.03 ± 0.22
large_parameter_init/ODEProblem	30.7 ± 5.5 ms	26.1 ± 5.8 ms	1.18 ± 0.34
large_parameter_init/init	0.0592 ± 0.027 ms	0.053 ± 0.015 ms	1.12 ± 0.6
mtkcompile	8.9 ± 1.3 ms	7.25 ± 0.58 ms	1.23 ± 0.2
sparse_analytical_jacobian/ODEProblem	26.3 ± 3.9 ms	21.8 ± 3.3 ms	1.21 ± 0.26
sparse_analytical_jacobian/f_iip	0.06 ± 0.01 μs	0.06 ± 0.01 μs	1 ± 0.24
sparse_analytical_jacobian/f_oop	0.0949 ± 0.012 ms	0.0975 ± 0.013 ms	0.973 ± 0.18
time_to_load	4.89 ± 0.056 s	4.69 ± 0.023 s	1.04 ± 0.013

Memory benchmarks

	master	a0d4896...	master / a0d4896...
ODEProblem	30.1 k allocs: 1.75 MB	30.1 k allocs: 1.75 MB	1
init	0.415 k allocs: 0.0517 MB	0.415 k allocs: 0.0517 MB	1
large_parameter_init/ODEProblem	0.329 M allocs: 12.5 MB	0.328 M allocs: 12.4 MB	1.01
large_parameter_init/init	0.798 k allocs: 0.151 MB	0.798 k allocs: 0.151 MB	1
mtkcompile	0.0561 M allocs: 2.75 MB	0.0561 M allocs: 2.75 MB	1
sparse_analytical_jacobian/ODEProblem	0.189 M allocs: 7.29 MB	0.189 M allocs: 7.28 MB	1
sparse_analytical_jacobian/f_iip	0 allocs: 0 B	0 allocs: 0 B
sparse_analytical_jacobian/f_oop	0.848 k allocs: 27 kB	0.848 k allocs: 27 kB	1
time_to_load	0.146 k allocs: 11.4 kB	0.146 k allocs: 11.4 kB	1