openscvx 0.1.2__py3-none-any.whl → 0.2.1.dev0__py3-none-any.whl

openscvx/discretization.py

@@ -17,7 +17,30 @@ def dVdt(
     n_u: int,
     N: int,
     dis_type: str,
+    **params
 ) -> jnp.ndarray:
+    """Compute the time derivative of the augmented state vector.
+    
+    This function computes the time derivative of the augmented state vector V,
+    which includes the state, state transition matrix, and control influence matrix.
+    
+    Args:
+        tau (float): Current normalized time in [0,1].
+        V (jnp.ndarray): Augmented state vector.
+        u_cur (np.ndarray): Control input at current node.
+        u_next (np.ndarray): Control input at next node.
+        state_dot (callable): Function computing state derivatives.
+        A (callable): Function computing state Jacobian.
+        B (callable): Function computing control Jacobian.
+        n_x (int): Number of states.
+        n_u (int): Number of controls.
+        N (int): Number of nodes in trajectory.
+        dis_type (str): Discretization type ("ZOH" or "FOH").
+        **params: Additional parameters passed to state_dot, A, and B.
+        
+    Returns:
+        jnp.ndarray: Time derivative of augmented state vector.
+    """
     # Define the nodes
     nodes = jnp.arange(0, N-1)
 
@@ -52,15 +75,15 @@ def dVdt(
     u = u[: x.shape[0]]
 
     # Compute the nonlinear propagation term
-    f = state_dot(x, u[:, :-1], nodes)
+    f = state_dot(x, u[:, :-1], nodes, *params.items())
     F = s[:, None] * f
 
     # Evaluate the State Jacobian
-    dfdx = A(x, u[:, :-1], nodes)
+    dfdx = A(x, u[:, :-1], nodes, *params.items())
     sdfdx = s[:, None, None] * dfdx
 
     # Evaluate the Control Jacobian
-    dfdu_veh = B(x, u[:, :-1], nodes)
+    dfdu_veh = B(x, u[:, :-1], nodes, *params.items())
     dfdu = dfdu.at[:, :, :-1].set(s[:, None, None] * dfdu_veh)
     dfdu = dfdu.at[:, :, -1].set(f)
 
@@ -76,7 +99,8 @@ def dVdt(
     dVdt = dVdt.at[:, i3:i4].set((jnp.matmul(sdfdx, V[:, i3:i4].reshape(-1, n_x, n_u)) + dfdu * beta).reshape(-1, n_x * n_u))
     dVdt = dVdt.at[:, i4:i5].set((jnp.matmul(sdfdx, V[:, i4:i5].reshape(-1, n_x)[..., None]).squeeze(-1) + z).reshape(-1, n_x))
     # fmt: on
-    return dVdt.flatten()
+
+    return dVdt.reshape(-1)
 
 
 def calculate_discretization(
@@ -94,8 +118,38 @@ def calculate_discretization(
     rtol,
     atol,
     dis_type: str,
+    **kwargs
 ):
-
+    """Calculate the discretized system matrices.
+    
+    This function computes the discretized system matrices (A_bar, B_bar, C_bar)
+    and defect vector (z_bar) using numerical integration.
+    
+    Args:
+        x: State trajectory.
+        u: Control trajectory.
+        state_dot (callable): Function computing state derivatives.
+        A (callable): Function computing state Jacobian.
+        B (callable): Function computing control Jacobian.
+        n_x (int): Number of states.
+        n_u (int): Number of controls.
+        N (int): Number of nodes in trajectory.
+        custom_integrator (bool): Whether to use custom RK45 integrator.
+        debug (bool): Whether to use debug mode.
+        solver (str): Name of the solver to use.
+        rtol (float): Relative tolerance for integration.
+        atol (float): Absolute tolerance for integration.
+        dis_type (str): Discretization type ("ZOH" or "FOH").
+        **kwargs: Additional parameters passed to state_dot, A, and B.
+        
+    Returns:
+        tuple: (A_bar, B_bar, C_bar, z_bar, Vmulti) where:
+            - A_bar: Discretized state transition matrix
+            - B_bar: Discretized control influence matrix
+            - C_bar: Discretized control influence matrix for next node
+            - z_bar: Defect vector
+            - Vmulti: Full augmented state trajectory
+    """
     # Define indices for slicing the augmented state vector
     i0 = 0
     i1 = n_x
@@ -104,67 +158,91 @@ def calculate_discretization(
     i4 = i3 + n_x * n_u
     i5 = i4 + n_x
 
-    # initial augmented state
+    # Initial augmented state
     V0 = jnp.zeros((N - 1, i5))
     V0 = V0.at[:, :n_x].set(x[:-1].astype(float))
-    V0 = V0.at[:, n_x : n_x + n_x * n_x].set(
+    V0 = V0.at[:, n_x:n_x + n_x * n_x].set(
         jnp.eye(n_x).reshape(1, -1).repeat(N - 1, axis=0)
     )
 
-    # choose integrator
+    # Choose integrator
+    integrator_args = dict(
+        u_cur=u[:-1].astype(float),
+        u_next=u[1:].astype(float),
+        state_dot=state_dot,
+        A=A,
+        B=B,
+        n_x=n_x,
+        n_u=n_u,
+        N=N,
+        dis_type=dis_type,
+        **kwargs  # <-- adds parameter values with names
+    )
+
+    # Define dVdt wrapper using named arguments
+    def dVdt_wrapped(t, y):
+        return dVdt(t, y, **integrator_args)
+
+    # Choose integrator
     if custom_integrator:
-        # fmt: off
         sol = solve_ivp_rk45(
-            lambda t,y,*a: dVdt(t, y, *a),
-            1.0/(N-1),
+            dVdt_wrapped,
+            1.0 / (N - 1),
             V0.reshape(-1),
-            args=(u[:-1].astype(float), u[1:].astype(float),
-                  state_dot, A, B, n_x, n_u, N, dis_type),
+            args=(),
             is_not_compiled=debug,
         )
-        # fmt: on
     else:
-        # fmt: off
         sol = solve_ivp_diffrax(
-            lambda t,y,*a: dVdt(t, y, *a),
-            1.0/(N-1),
+            dVdt_wrapped,
+            1.0 / (N - 1),
             V0.reshape(-1),
-            args=(u[:-1].astype(float), u[1:].astype(float),
-                  state_dot, A, B, n_x, n_u, N, dis_type),
             solver_name=solver,
             rtol=rtol,
             atol=atol,
+            args=(),
             extra_kwargs=None,
         )
-        # fmt: on
 
     Vend = sol[-1].T.reshape(-1, i5)
     Vmulti = sol.T
 
-    # fmt: off
-    A_bar = Vend[:, i1:i2].reshape(N-1, n_x, n_x).transpose(1,2,0).reshape(n_x*n_x, -1, order='F').T
-    B_bar = Vend[:, i2:i3].reshape(N-1, n_x, n_u).transpose(1,2,0).reshape(n_x*n_u, -1, order='F').T
-    C_bar = Vend[:, i3:i4].reshape(N-1, n_x, n_u).transpose(1,2,0).reshape(n_x*n_u, -1, order='F').T
+    A_bar = Vend[:, i1:i2].reshape(N - 1, n_x, n_x).transpose(1, 2, 0).reshape(n_x * n_x, -1, order='F').T
+    B_bar = Vend[:, i2:i3].reshape(N - 1, n_x, n_u).transpose(1, 2, 0).reshape(n_x * n_u, -1, order='F').T
+    C_bar = Vend[:, i3:i4].reshape(N - 1, n_x, n_u).transpose(1, 2, 0).reshape(n_x * n_u, -1, order='F').T
     z_bar = Vend[:, i4:i5]
-    # fmt: on
 
     return A_bar, B_bar, C_bar, z_bar, Vmulti
 
 
-def get_discretization_solver(dyn: Dynamics, params):
-    return lambda x, u: calculate_discretization(
+def get_discretization_solver(dyn: Dynamics, settings, param_map):
+    """Create a discretization solver function.
+    
+    This function creates a solver that computes the discretized system matrices
+    using the specified dynamics and settings.
+    
+    Args:
+        dyn (Dynamics): System dynamics object.
+        settings: Configuration settings for discretization.
+        param_map (dict): Mapping of parameter names to values.
+        
+    Returns:
+        callable: A function that computes the discretized system matrices.
+    """
+    return lambda x, u, *params: calculate_discretization(
         x=x,
         u=u,
         state_dot=dyn.f,
         A=dyn.A,
         B=dyn.B,
-        n_x=params.sim.n_states,
-        n_u=params.sim.n_controls,
-        N=params.scp.n,
-        custom_integrator=params.dis.custom_integrator,
-        debug=params.dev.debug,
-        solver=params.dis.solver,
-        rtol=params.dis.rtol,
-        atol=params.dis.atol,
-        dis_type=params.dis.dis_type,
-    )
+        n_x=settings.sim.n_states,
+        n_u=settings.sim.n_controls,
+        N=settings.scp.n,
+        custom_integrator=settings.dis.custom_integrator,
+        debug=settings.dev.debug,
+        solver=settings.dis.solver,
+        rtol=settings.dis.rtol,
+        atol=settings.dis.atol,
+        dis_type=settings.dis.dis_type,
+        **dict(zip(param_map.keys(), params))  # <--- Named keyword args
+    )

openscvx/dynamics.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 import functools
 
 import jax.numpy as jnp
@@ -7,25 +7,165 @@ import jax.numpy as jnp
 
 @dataclass
 class Dynamics:
+    """
+    Dataclass to hold a system dynamics function and (optionally) its gradients.
+    This class is intended to be instantiated using the `dynamics` decorator wrapped around a function defining the system dynamics.
+    Both the dynamics and optional gradients should be composed of `jax` primitives to enable efficient computation.
+
+    Usage examples:
+
+    ```python
+    @dynamics
+    def f(x_, u_):
+        return x_ + u_
+    # f is now a Dynamics object
+    ```
+
+    ```python
+    @dynamics(A=grad_f_x, B=grad_f_u)
+    def f(x_, u_):
+        return x_ + u_
+    ```
+
+    Or, if a more lambda-function-style is desired, the function can be directly wrapped:
+
+    ```python
+    dyn = dynamics(lambda x_, u_: x_ + u_)
+    ```
+
+    ---
+    **Using Parameters in Dynamics**
+
+    You can use symbolic `Parameter` objects in your dynamics function to represent tunable or environment-dependent values. **The argument names for parameters must match the parameter name with an underscore suffix** (e.g., `I_sp_` for a parameter named `I_sp`). This is required for the parameter mapping to work correctly.
+
+    Example (3DoF rocket landing):
+
+    ```python
+    from openscvx.backend.parameter import Parameter
+    import jax.numpy as jnp
+
+    I_sp = Parameter("I_sp")
+    g = Parameter("g")
+    theta = Parameter("theta")
+
+    @dynamics
+    def rocket_dynamics(x_, u_, I_sp_, g_, theta_):
+        m = x_[6]
+        T = u_
+        r_dot = x_[3:6]
+        g_vec = jnp.array([0, 0, g_])
+        v_dot = T/m - g_vec
+        m_dot = -jnp.linalg.norm(T) / (I_sp_ * 9.807 * jnp.cos(theta_))
+        t_dot = 1
+        return jnp.hstack([r_dot, v_dot, m_dot, t_dot])
+
+    # Set parameter values before solving
+    I_sp.value = 225
+    g.value = 3.7114
+    theta.value = 27 * jnp.pi / 180
+    ```
+
+    ---
+    **Using Parameters in Nodal Constraints**
+
+    You can also use symbolic `Parameter` objects in nodal constraints. As with dynamics, the argument names for parameters in the constraint function must match the parameter name with an underscore suffix (e.g., `g_` for a parameter named `g`).
+
+    Example:
+
+    ```python
+    from openscvx.backend.parameter import Parameter
+    from openscvx.constraints import nodal
+    import jax.numpy as jnp
+
+    g = Parameter("g")
+    g.value = 3.7114
+
+    @nodal
+    def terminal_velocity_constraint(x_, u_, g_):
+        # Enforce a terminal velocity constraint using the gravity parameter
+        return x_[5] + g_ * x_[7]  # e.g., vz + g * t <= 0 at final node
+    ```
+
+    When building your problem, collect all parameters with `Parameter.get_all()` and pass them to your problem setup.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
+            Function defining the continuous time nonlinear system dynamics as x_dot = f(x, u, ...params).
+            - x: 1D array (state at a single node), shape (n_x,)
+            - u: 1D array (control at a single node), shape (n_u,)
+            - Additional parameters: passed as keyword arguments with names matching the parameter name plus an underscore (e.g., g_ for Parameter('g')).
+            If you want to use parameters, include them as extra arguments with the underscore naming convention.
+            If you use vectorized integration or batch evaluation, x and u may be 2D arrays (N, n_x) and (N, n_u).
+        A (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            Jacobian of `f` w.r.t. `x`. If not specified, will be calculated using `jax.jacfwd`.
+        B (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            Jacobian of `f` w.r.t. `u`. If not specified, will be calculated using `jax.jacfwd`.
+
+    Returns:
+        Dynamics: A dataclass bundling the system dynamics function and Jacobians.
+    """
+
     f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
     A: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
     B: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
 
+
 def dynamics(
-    _func=None,
+    _func: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
     *,
-    A: Optional[Callable] = None,
-    B: Optional[Callable] = None,):
-    """Decorator to mark a function as defining the system dynamics.
+    A: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
+    B: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None,
+) -> Union[Callable, Dynamics]:
+    """
+    Decorator that wraps a function defining the system dynamics as a `Dynamics` object.
+    You may optionally specify the system gradients w.r.t. `x`, `u` if desired, if not specified they will be calculated using `jax.jacfwd`.
+    Note: the dynamics as well as the optional gradients should be composed of `jax` primitives to enable efficient computation.
+
+    This decorator may be used with or without arguments:
 
-    Use as:
-    @dynamics(A=my_grad_f_x, B=my_grad_f_u)')
-    def my_dynamics(x,u): ...
+    ```
+    @dynamics
+    def f(x, u): ...
+    ```
+
+    or
+
+    ```
+    @dynamics(A=grad_f_x, B=grad_f_u)
+    def f(x, u): ...
+    ```
+
+    or, if a more lambda-function-style is desired, the function can be direclty wrapped
+
+    ```
+    dyn = dynamics(f(x,u))
+    dyn_lambda = dynamics(lambda x, u: ...)
+    ```
+
+    Args:
+        _func (callable, optional): The function to wrap. Populated
+            when using @dynamics with no extra args.
+        A (callable, optional): Jacobian of f wrt state x. Computed
+            via jax.jacfwd if not provided.
+        B (callable, optional): Jacobian of f wrt input u. Computed
+            via jax.jacfwd if not provided.
+
+    Returns:
+        Union[Callable, Dynamics]
+            A decorator if called without a function, or a `Dynamics` dataclass bundling system dynamics function
+            and Jacobians when applied to a function.
+
+    Examples:
+        >>> @dynamics
+        ... def f(x, u):
+        ...     return x + u
+        >>> isinstance(f, Dynamics)
+        True
     """
 
     def decorator(f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
-        # wrap so name, doc, signature stay on f
-        wrapped = functools.wraps(f)(f)
+        # Had to unwrap to ensure arguments are visible downstream. Originally wrapped so name, doc, signature stay on f
+        wrapped = f 
         return Dynamics(
             f=wrapped,
             A=A,
@@ -38,4 +178,3 @@ def dynamics(
     # if called as dynamics(func), we immediately decorate
     else:
         return decorator(_func)
-

openscvx/integrators.py

@@ -1,7 +1,52 @@
+import os
+os.environ["EQX_ON_ERROR"] = "nan"
+
+
+from typing import Callable, Any, Optional
+
+
+import jax
+import jax.numpy as jnp
+import diffrax as dfx
+from typing import Callable, Any
+
 import jax
 import jax.numpy as jnp
 import diffrax as dfx
 
+from diffrax._global_interpolation import DenseInterpolation
+from jax import tree_util
+
+# Safely check if DenseInterpolation is already registered
+try:
+    # Attempt to flatten a dummy DenseInterpolation instance
+    # Provide dummy arguments to create a valid instance
+    dummy_instance = DenseInterpolation(
+        ts=jnp.array([]),
+        ts_size=0,
+        infos=None,
+        interpolation_cls=None,
+        direction=None,
+        t0_if_trivial=0.0,
+        y0_if_trivial=jnp.array([]),
+    )
+    tree_util.tree_flatten(dummy_instance)
+except ValueError:
+    # Register DenseInterpolation as a PyTree node if not already registered
+    def dense_interpolation_flatten(obj):
+        # Flatten the internal data of DenseInterpolation
+        return (obj._data,), None
+
+    def dense_interpolation_unflatten(aux_data, children):
+        # Reconstruct DenseInterpolation from its flattened data
+        return DenseInterpolation(*children)
+
+    tree_util.register_pytree_node(
+        DenseInterpolation,
+        dense_interpolation_flatten,
+        dense_interpolation_unflatten,
+    )
+
 SOLVER_MAP = {
     "Tsit5": dfx.Tsit5,
     "Euler": dfx.Euler,
@@ -19,7 +64,30 @@ SOLVER_MAP = {
 }
 
 # fmt: off
-def rk45_step(f, t, y, h, *args):
+def rk45_step(
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    t: jnp.ndarray,
+    y: jnp.ndarray,
+    h: float,
+    *args
+) -> jnp.ndarray:
+    """
+    Perform a single RK45 (Runge-Kutta-Fehlberg) integration step.
+
+    This implements the classic Dorman-Prince coefficients for an
+    explicit 4(5) method, returning the fourth-order estimate.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]):
+            ODE right-hand side; signature f(t, y, *args) -> dy/dt.
+        t (jnp.ndarray): Current time.
+        y (jnp.ndarray): Current state vector.
+        h (float): Step size.
+        *args: Additional arguments passed to `f`.
+
+    Returns:
+        jnp.ndarray: Next state estimate at t + h.
+    """
     k1 = f(t, y, *args)
     k2 = f(t + h/4, y + h*k1/4, *args)
     k3 = f(t + 3*h/8, y + 3*h*k1/32 + 9*h*k2/32, *args)
@@ -31,14 +99,31 @@ def rk45_step(f, t, y, h, *args):
 
 
 def solve_ivp_rk45(
-    f,
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
     tau_final: float,
-    y_0,
+    y_0: jnp.ndarray,
     args,
     tau_0: float = 0.0,
     num_substeps: int = 50,
     is_not_compiled: bool = False,
 ):
+    """
+    Solve an initial-value ODE problem using fixed-step RK45 integration.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]):
+            ODE right-hand side; signature f(t, y, *args) -> dy/dt.
+        tau_final (float): Final integration time.
+        y_0 (jnp.ndarray): Initial state at tau_0.
+        args (tuple): Extra arguments to pass to `f`.
+        tau_0 (float, optional): Initial time. Defaults to 0.0.
+        num_substeps (int, optional): Number of output time points. Defaults to 50.
+        is_not_compiled (bool, optional): If True, use Python loop instead of
+            JAX `lax.fori_loop`. Defaults to False.
+
+    Returns:
+        jnp.ndarray: Array of shape (num_substeps, state_dim) with solution at each time.
+    """
     substeps = jnp.linspace(tau_0, tau_final, num_substeps)
 
     h = (tau_final - tau_0) / (len(substeps) - 1)
@@ -65,17 +150,40 @@ def solve_ivp_rk45(
 
 
 def solve_ivp_diffrax(
-    f,
-    tau_final,
-    y_0,
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    tau_final: float,
+    y_0: jnp.ndarray,
     args,
     tau_0: float = 0.0,
     num_substeps: int = 50,
-    solver_name="Dopri8",
+    solver_name: str = "Dopri8",
     rtol: float = 1e-3,
     atol: float = 1e-6,
     extra_kwargs=None,
 ):
+    """
+    Solve an initial-value ODE problem using a Diffrax adaptive solver.
+
+    Args:
+        f (Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray]): ODE right-hand side; signature f(t, y, *args) -> dy/dt.
+        tau_final (float): Final integration time.
+        y_0 (jnp.ndarray): Initial state at tau_0.
+        args (tuple): Extra arguments to pass to `f` in the solver term.
+        tau_0 (float, optional): Initial time. Defaults to 0.0.
+        num_substeps (int, optional): Number of save points between tau_0 and tau_final.
+            Defaults to 50.
+        solver_name (str, optional): Key into SOLVER_MAP for the Diffrax solver class.
+            Defaults to "Dopri8".
+        rtol (float, optional): Relative tolerance for adaptive stepping. Defaults to 1e-3.
+        atol (float, optional): Absolute tolerance for adaptive stepping. Defaults to 1e-6.
+        extra_kwargs (dict, optional): Additional keyword arguments forwarded to `diffeqsolve`.
+
+    Returns:
+        jnp.ndarray: Solution states at the requested save points, shape (num_substeps, state_dim).
+
+    Raises:
+        ValueError: If `solver_name` is not in SOLVER_MAP.
+    """
     substeps = jnp.linspace(tau_0, tau_final, num_substeps)
 
     solver_class = SOLVER_MAP.get(solver_name)
@@ -95,26 +203,31 @@ def solve_ivp_diffrax(
         args=args,
         stepsize_controller=stepsize_controller,
         saveat=dfx.SaveAt(ts=substeps),
+        progress_meter=dfx.NoProgressMeter(),
         **(extra_kwargs or {}),
     )
 
     return solution.ys
 
 
-# TODO: (norrisg) this function is basically identical to `solve_ivp_diffrax`, could combine, but requires returning solution and getting `.ys` wherever the `solve_ivp_diffrax` is called
 def solve_ivp_diffrax_prop(
-    f,
-    tau_final,
-    y_0,
+    f: Callable[[jnp.ndarray, jnp.ndarray, Any], jnp.ndarray],
+    tau_final: float,
+    y_0: jnp.ndarray,
     args,
     tau_0: float = 0.0,
     num_substeps: int = 50,
-    solver_name="Dopri8",
+    solver_name: str = "Dopri8",
     rtol: float = 1e-3,
     atol: float = 1e-6,
     extra_kwargs=None,
+    save_time: jnp.ndarray = None,
+    mask: jnp.ndarray = None,
 ):
-    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+    if save_time is None:
+        raise ValueError("save_time must be provided for export compatibility.")
+    if mask is None:
+        mask = jnp.ones_like(save_time, dtype=bool)
 
     solver_class = SOLVER_MAP.get(solver_name)
     if solver_class is None:
@@ -123,17 +236,23 @@ def solve_ivp_diffrax_prop(
 
     term = dfx.ODETerm(lambda t, y, args: f(t, y, *args))
     stepsize_controller = dfx.PIDController(rtol=rtol, atol=atol)
+
     solution = dfx.diffeqsolve(
         term,
         solver=solver,
         t0=tau_0,
         t1=tau_final,
-        dt0=(tau_final - tau_0) / (len(substeps) - 1),
+        dt0=(tau_final - tau_0) / 1,
         y0=y_0,
         args=args,
         stepsize_controller=stepsize_controller,
-        saveat=dfx.SaveAt(dense=True, ts=substeps),
+        saveat=dfx.SaveAt(dense=True),
         **(extra_kwargs or {}),
     )
 
-    return solution
+    # Evaluate all save_time points (static size), then mask them
+    all_evals = jax.vmap(solution.evaluate)(save_time)  # shape: (MAX_TAU_LEN, n_states)
+    masked_array = jnp.where(mask[:, None], all_evals, jnp.zeros_like(all_evals))
+         # shape: (variable_len, n_states)
+
+    return masked_array