--- Final Solver (used for performance test) ---
import numpy as np
from scipy import signal
from scipy.linalg import expm
from numba import njit
from typing import Any, Dict, List


@njit
def _simulate_foh(M, u, C, D, n, dt):
    """Numba-accelerated FOH simulation for uniform time steps."""
    n_aug = n + 2
    n_steps = len(u) - 1
    y = np.empty(len(u), dtype=np.float64)
    z = np.zeros(n_aug, dtype=np.float64)
    
    # Initial output (zero initial state)
    if n == 1:
        y[0] = C[0] * z[0] + D * u[0]
    else:
        y[0] = np.dot(C, z[:n]) + D * u[0]
    
    for k in range(n_steps):
        uk = u[k]
        ukp1 = u[k+1]
        s = (ukp1 - uk) / dt
        
        z[n] = uk
        z[n+1] = s
        
        z[:] = np.dot(M, z)
        
        if n == 1:
            y[k+1] = C[0] * z[0] + D * ukp1
        else:
            y[k+1] = np.dot(C, z[:n]) + D * ukp1
    
    return y


@njit
def _simulate_foh_nonuniform(A, B, C, D, u, t, dt_arr, n):
    """Numba-accelerated FOH simulation for non-uniform time steps.
    We cannot call expm from numba, so we pass precomputed M for each step from Python.
    This function is placeholder; we handle non-uniform in Python with per-step expm.
    """
    pass


class Solver:
    def __init__(self):
        # Pre-import and compile the Numba function with a warmup call
        self._simulate_foh = _simulate_foh
        # Warmup to trigger JIT compilation (compilation time doesn't count)
        M_tmp = np.eye(3, dtype=np.float64)
        u_tmp = np.ones(10, dtype=np.float64)
        C_tmp = np.ones(1, dtype=np.float64)
        D_tmp = 0.0
        _ = self._simulate_foh(M_tmp, u_tmp, C_tmp, D_tmp, 1, 0.1)
    
    def solve(self, problem: Dict[str, Any], **kwargs) -> Dict[str, List[float]]:
        num = np.asarray(problem["num"], dtype=float)
        den = np.asarray(problem["den"], dtype=float)
        u = np.asarray(problem["u"], dtype=float)
        t = np.asarray(problem["t"], dtype=float)
        
        # Convert to state-space (controllable canonical form)
        A, B, C, D = signal.tf2ss(num, den)
        n = A.shape[0]
        B = B.ravel()
        C = C.ravel()
        D = D.ravel()[0]
        
        dt_arr = np.diff(t)
        
        # Build augmented system matrix template (without time scaling)
        n_aug = n + 2
        A_aug = np.zeros((n_aug, n_aug), dtype=float)
        A_aug[:n, :n] = A
        A_aug[:n, n] = B
        A_aug[n, n+1] = 1.0
        
        # Check if time steps are uniform
        if len(dt_arr) == 0:
            # Single time point
            return {"yout": [float(D * u[0])]}
        
        first_dt = dt_arr[0]
        uniform = np.allclose(dt_arr, first_dt)
        
        if uniform:
            # Single matrix exponential for all steps
            M = expm(A_aug * first_dt)
            y = self._simulate_foh(M, u, C, D, n, first_dt)
        else:
            # Non-uniform: compute expm for each step
            n_steps = len(t) - 1
            y = np.empty(len(t))
            z = np.zeros(n_aug)
            
            if n == 1:
                y[0] = C[0] * z[0] + D * u[0]
            else:
                y[0] = np.dot(C, z[:n]) + D * u[0]
            
            for k in range(n_steps):
                dt_k = dt_arr[k]
                uk = u[k]
                ukp1 = u[k+1]
                s = (ukp1 - uk) / dt_k
                
                M_k = expm(A_aug * dt_k)
                
                z[n] = uk
                z[n+1] = s
                z[:] = np.dot(M_k, z)
                
                if n == 1:
                    y[k+1] = C[0] * z[0] + D * ukp1
                else:
                    y[k+1] = np.dot(C, z[:n]) + D * ukp1
        
        return {"yout": y.tolist()}
--- End Solver ---
Running performance test...
============================= test session starts ==============================
platform linux -- Python 3.12.13, pytest-9.0.2, pluggy-1.6.0
rootdir: /tests
plugins: anyio-4.13.0, jaxtyping-0.3.9
collected 3 items

../tests/test_outputs.py .
--- Performance Summary ---
Validity: True
Total Baseline Time: 9.2229s
Total Solver Time:   0.2909s
Raw Speedup:         31.7068 x
Final Reward (Score): 31.7068
---------------------------
..

==================================== PASSES ====================================
=========================== short test summary info ============================
PASSED ../tests/test_outputs.py::test_solver_exists
PASSED ../tests/test_outputs.py::test_solver_validity
PASSED ../tests/test_outputs.py::test_solver_speedup
======================== 3 passed in 121.99s (0:02:01) =========================
