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	ee058fa...	master / ee058fa...
ODEProblem	5.38 ± 0.19 ms	5.54 ± 0.17 ms	0.971 ± 0.046
init	0.0701 ± 0.027 ms	0.0883 ± 0.028 ms	0.794 ± 0.39
large_parameter_init/ODEProblem	22.7 ± 3.9 ms	24 ± 4.3 ms	0.944 ± 0.23
large_parameter_init/init	0.0899 ± 0.035 ms	0.109 ± 0.042 ms	0.828 ± 0.45
mtkcompile	9.42 ± 0.29 ms	9.67 ± 0.51 ms	0.975 ± 0.059
sparse_analytical_jacobian/ODEProblem	25.8 ± 0.77 ms	27.1 ± 0.94 ms	0.95 ± 0.043
sparse_analytical_jacobian/f_iip	0.071 ± 0.01 μs	0.07 ± 0 μs	1.01 ± 0.14
sparse_analytical_jacobian/f_oop	0.386 ± 0.011 ms	0.384 ± 0.012 ms	1 ± 0.043
time_to_load	5.84 ± 0.096 s	5.9 ± 0.066 s	0.99 ± 0.02

Memory benchmarks

	master	ee058fa...	master / ee058fa...
ODEProblem	0.0376 M allocs: 1.92 MB	0.0376 M allocs: 1.94 MB	0.992
init	0.417 k allocs: 0.0698 MB	0.417 k allocs: 0.0698 MB	1
large_parameter_init/ODEProblem	0.325 M allocs: 11.4 MB	0.325 M allocs: 11.4 MB	1
large_parameter_init/init	0.605 k allocs: 0.172 MB	0.605 k allocs: 0.172 MB	1
mtkcompile	0.0618 M allocs: 3.45 MB	0.0618 M allocs: 3.43 MB	1
sparse_analytical_jacobian/ODEProblem	0.193 M allocs: 7.93 MB	0.193 M allocs: 7.93 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.5 kB	0.153 k allocs: 14.5 kB	1

[github-actions]

Benchmark Results (Julia v1)

Time benchmarks

	master	ee058fa...	master / ee058fa...
ODEProblem	5.35 ± 0.34 ms	5.33 ± 0.33 ms	1 ± 0.089
init	0.043 ± 0.021 ms	0.0511 ± 0.021 ms	0.841 ± 0.53
large_parameter_init/ODEProblem	25.9 ± 4.9 ms	27.2 ± 5.9 ms	0.951 ± 0.27
large_parameter_init/init	0.0669 ± 0.041 ms	0.0679 ± 0.028 ms	0.985 ± 0.73
mtkcompile	8.96 ± 0.71 ms	8.47 ± 0.59 ms	1.06 ± 0.11
sparse_analytical_jacobian/ODEProblem	24.1 ± 3.3 ms	24.4 ± 3.4 ms	0.985 ± 0.19
sparse_analytical_jacobian/f_iip	0.09 ± 0.001 μs	0.1 ± 0.001 μs	0.9 ± 0.013
sparse_analytical_jacobian/f_oop	0.131 ± 0.015 ms	0.13 ± 0.014 ms	1 ± 0.16
time_to_load	5.79 ± 0.028 s	5.93 ± 0.093 s	0.976 ± 0.016

Memory benchmarks

	master	ee058fa...	master / ee058fa...
ODEProblem	0.0375 M allocs: 1.68 MB	0.0375 M allocs: 1.7 MB	0.991
init	0.416 k allocs: 0.0518 MB	0.416 k allocs: 0.0518 MB	1
large_parameter_init/ODEProblem	0.352 M allocs: 12.8 MB	0.352 M allocs: 12.8 MB	1
large_parameter_init/init	0.799 k allocs: 0.152 MB	0.799 k allocs: 0.152 MB	1
mtkcompile	0.0598 M allocs: 2.76 MB	0.0598 M allocs: 2.76 MB	1
sparse_analytical_jacobian/ODEProblem	0.186 M allocs: 6.89 MB	0.186 M allocs: 6.89 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.145 k allocs: 11 kB	0.145 k allocs: 11 kB	1

[baggepinnen]

what's the status of this PR? Mike has started to think about ops in dyad so I'd like the low level stuff to be available and working to start trying things out. Last time I tried, there were some performance problems that I hoe were addressed by #4679, anything else needed?

[AayushSabharwal]

I'll prioritize it. It needs CI to pass, there are some failures in linearization.