openscvx 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

openscvx/_version.py

@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.1'
+__version_tuple__ = version_tuple = (0, 1, 1)

openscvx/augmentation.py

@@ -0,0 +1,44 @@
+from typing import List
+
+from openscvx.constraints.ctcs import CTCSConstraint
+
+def sort_ctcs_constraints(constraints_ctcs: List[CTCSConstraint], N: int):        
+    idx_to_nodes: dict[int, tuple] = {}
+    next_idx = 0
+    for c in constraints_ctcs:
+        # normalize None to full horizon
+        c.nodes = c.nodes or (0, N)
+        key = c.nodes
+
+        if c.idx is not None:
+            # user supplied an identifier: ensure it always points to the same interval
+            if c.idx in idx_to_nodes:
+                if idx_to_nodes[c.idx] != key:
+                    raise ValueError(
+                        f"idx={c.idx} was first used with interval={idx_to_nodes[c.idx]}, "
+                        f"but now you gave it interval={key}"
+                    )
+            else:
+                idx_to_nodes[c.idx] = key
+
+        else:
+            # no identifier: see if this interval already has one
+            for existing_id, nodes in idx_to_nodes.items():
+                if nodes == key:
+                    c.idx = existing_id
+                    break
+            else:
+                # brand-new interval: pick the next free auto-id
+                while next_idx in idx_to_nodes:
+                    next_idx += 1
+                c.idx = next_idx
+                idx_to_nodes[next_idx] = key
+                next_idx += 1
+    
+    # Extract your intervals in ascending‐idx order
+    ordered_ids       = sorted(idx_to_nodes.keys())
+    node_intervals    = [ idx_to_nodes[i] for i in ordered_ids ]
+    id_to_position    = { ident: pos for pos, ident in enumerate(ordered_ids) }
+    num_augmented_states = len(ordered_ids)
+
+    return constraints_ctcs, node_intervals, num_augmented_states,

openscvx/config.py

@@ -0,0 +1,247 @@
+import numpy as np
+from dataclasses import dataclass, field
+from typing import Dict, List
+
+
+def get_affine_scaling_matrices(n, minimum, maximum):
+    S = np.diag(np.maximum(np.ones(n), abs(minimum - maximum) / 2))
+    c = (maximum + minimum) / 2
+    return S, c
+
+
+@dataclass
+class DiscretizationConfig:
+    dis_type: str = "FOH"
+    custom_integrator: bool = True
+    solver: str = "Tsit5"
+    args: Dict = field(default_factory=dict)
+    atol: float = 1e-3
+    rtol: float = 1e-6
+
+    """
+    Configuration class for discretization settings.
+
+    This class defines the parameters required for discretizing system dynamics.
+
+    Main arguments:
+    These are the arguments most commonly used day-to-day.
+
+    Args:
+        dis_type (str): The type of discretization to use (e.g., "FOH" for First-Order Hold). Defaults to "FOH".
+        custom_integrator (bool): This enables our custom fixed-step RK45 algorthim. This tends to be faster then Diffrax but unless your going for speed, its reccomended to stick with Diffrax for robustness and other solver options. Defaults to False.
+        solver (str): Not used if custom_integrator is enabled. Any choice of solver in Diffrax is valid, please refer here, [How to Choose a Solver](https://docs.kidger.site/diffrax/usage/how-to-choose-a-solver/). Defaults to "Tsit5".
+
+    Other arguments:
+    These arguments are less frequently used, and for most purposes you shouldn't need to understand these.
+
+    Args:
+        args (Dict): Additional arguments to pass to the solver which can be found [here](https://docs.kidger.site/diffrax/api/diffeqsolve/). Defaults to an empty dictionary.
+        atol (float): Absolute tolerance for the solver. Defaults to 1e-3.
+        rtol (float): Relative tolerance for the solver. Defaults to 1e-6.
+    """
+
+
+@dataclass
+class DevConfig:
+    profiling: bool = False
+    debug: bool = False
+    printing: bool = True
+
+    """
+    Configuration class for development settings.
+
+    This class defines the parameters used for development and debugging purposes.
+
+    Main arguments:
+    These are the arguments most commonly used day-to-day.
+
+    Args:
+        profiling (bool): Whether to enable profiling for performance analysis. Defaults to False.
+        debug (bool): Disables all precompilation so you can place breakpoints and inspect values. Defaults to False.
+    """
+
+
+@dataclass
+class ConvexSolverConfig:
+    solver: str = "QOCO"
+    solver_args: dict = field(default_factory=lambda: {"abstol": 1e-6, "reltol": 1e-9})
+    cvxpygen: bool = False
+
+    """
+    Configuration class for convex solver settings.
+
+    This class defines the parameters required for configuring a convex solver.
+
+    These are the arguments most commonly used day-to-day. Generally I have found [QOCO](https://qoco-org.github.io/qoco/index.html) to be the most performant of the CVXPY solvers for these types of problems (I do have a bias as the author is from my group) and can handle up to SOCP's. 
+    [CLARABEL](https://clarabel.org/stable/) is also a great option with feasibility checking and can handle a few more problem types.
+    [CVXPYGen](https://github.com/cvxgrp/cvxpygen) is also great if your problem isn't too large and allows. I have found qocogen to be the most performant of the CVXPYGen solvers. 
+
+    Args:
+        solver (str):  The name of the CVXPY solver to use. A list of options can be found [here](https://www.cvxpy.org/tutorial/solvers/index.html). Defaults to "QOCO".
+        solver_args (dict): Ensure you are using the correct arguments for your solver as they are not all common. Additional arguments to configure the solver, such as tolerances. 
+                            Defaults to {"abstol": 1e-6, "reltol": 1e-9}.
+        cvxpygen (bool): Whether to enable CVXPY code generation for the solver. Defaults to False.
+    """
+
+
+@dataclass
+class PropagationConfig:
+    inter_sample: int = 30
+    dt: float = 0.1
+    solver: str = "Dopri8"
+    args: Dict = field(default_factory=dict)
+    atol: float = 1e-3
+    rtol: float = 1e-6
+
+    """
+    Configuration class for propagation settings.
+
+    This class defines the parameters required for propagating the nonlinear system dynamics using the optimal control sequence.
+
+    Main arguments:
+    These are the arguments most commonly used day-to-day.
+    
+    Args:
+        dt (float): The time step for propagation. Defaults to 0.1.
+        inter_sample (int): How dense the propagation within multishot discretization should be.
+
+    Other arguments:
+    The solver should likley not to be changed as it is a high accuracy 8th order runga kutta method.
+    
+    Args:
+        solver (str): The numerical solver to use for propagation (e.g., "Dopri8"). Defaults to "Dopri8".
+        args (Dict): Additional arguments to pass to the solver. Defaults to an empty dictionary.
+        atol (float): Absolute tolerance for the solver. Defaults to 1e-3.
+        rtol (float): Relative tolerance for the solver. Defaults to 1e-6.
+    """
+
+
+@dataclass
+class SimConfig:
+    x_bar: np.ndarray
+    u_bar: np.ndarray
+    initial_state: np.ndarray
+    final_state: np.ndarray
+    max_state: np.ndarray
+    min_state: np.ndarray
+    max_control: np.ndarray
+    min_control: np.ndarray
+    total_time: float
+    idx_x_true: slice
+    idx_u_true: slice
+    idx_t: slice
+    idx_y: slice
+    idx_s: slice
+    ctcs_node_intervals: list = None
+    constraints_ctcs: List[callable] = field(
+        default_factory=list
+    )  # TODO (norrisg): clean this up, consider moving to dedicated `constraints` dataclass
+    constraints_nodal: List[callable] = field(default_factory=list)
+    n_states: int = None
+    n_controls: int = None
+    S_x: np.ndarray = None
+    inv_S_x: np.ndarray = None
+    c_x: np.ndarray = None
+    S_u: np.ndarray = None
+    inv_S_u: np.ndarray = None
+    c_u: np.ndarray = None
+
+    def __post_init__(self):
+        self.n_states = len(self.max_state)
+        self.n_controls = len(self.max_control)
+
+        assert (
+            len(self.initial_state.value) == self.n_states - (self.idx_y.stop - self.idx_y.start)
+        ), f"Initial state must have {self.n_states - (self.idx_y.stop - self.idx_y.start)} elements"
+        assert (
+            len(self.final_state.value) == self.n_states - (self.idx_y.stop - self.idx_y.start)
+        ), f"Final state must have {self.n_states - (self.idx_y.stop - self.idx_y.start)} elements"
+        assert (
+            self.max_state.shape[0] == self.n_states
+        ), f"Max state must have {self.n_states} elements"
+        assert (
+            self.min_state.shape[0] == self.n_states
+        ), f"Min state must have {self.n_states} elements"
+        assert (
+            self.max_control.shape[0] == self.n_controls
+        ), f"Max control must have {self.n_controls} elements"
+        assert (
+            self.min_control.shape[0] == self.n_controls
+        ), f"Min control must have {self.n_controls} elements"
+
+        if self.S_x is None or self.c_x is None:
+            self.S_x, self.c_x = get_affine_scaling_matrices(
+                self.n_states, self.min_state, self.max_state
+            )
+            # Use the fact that S_x is diagonal to compute the inverse
+            self.inv_S_x = np.diag(1 / np.diag(self.S_x))
+        if self.S_u is None or self.c_u is None:
+            self.S_u, self.c_u = get_affine_scaling_matrices(
+                self.n_controls, self.min_control, self.max_control
+            )
+            self.inv_S_u = np.diag(1 / np.diag(self.S_u))
+
+
+@dataclass
+class ScpConfig:
+    n: int = None
+    k_max: int = 200
+    w_tr: float = 1e0
+    lam_vc: float = 1e0
+    ep_tr: float = 1e-4
+    ep_vb: float = 1e-4
+    ep_vc: float = 1e-8
+    lam_cost: float = 0.0
+    lam_vb: float = 0.0
+    uniform_time_grid: bool = False
+    cost_drop: int = -1
+    cost_relax: float = 1.0
+    w_tr_adapt: float = 1.0
+    w_tr_max: float = None
+    w_tr_max_scaling_factor: float = None
+
+    """
+    Configuration class for Sequential Convex Programming (SCP).
+
+    This class defines the parameters used to configure the SCP solver. You will very likely need to modify
+    the weights for your problem. Please refer to my guide [here](https://haynec.github.io/openscvx/hyperparameter_tuning) for more information.
+
+    Attributes:
+        n (int): The number of discretization nodes. Defaults to `None`.
+        k_max (int): The maximum number of SCP iterations. Defaults to 200.
+        w_tr (float): The trust region weight. Defaults to 1.0.
+        lam_vc (float): The penalty weight for virtual control. Defaults to 1.0.
+        ep_tr (float): The trust region convergence tolerance. Defaults to 1e-4.
+        ep_vb (float): The boundary constraint convergence tolerance. Defaults to 1e-4.
+        ep_vc (float): The virtual constraint convergence tolerance. Defaults to 1e-8.
+        lam_cost (float): The weight for original cost. Defaults to 0.0.
+        lam_vb (float): The weight for virtual buffer. This is only used if there are nonconvex nodal constraints present. Defaults to 0.0.
+        uniform_time_grid (bool): Whether to use a uniform time grid. TODO haynec add a link to the time dilation page. Defaults to `False`.
+        cost_drop (int): The number of iterations to allow for cost stagnation before termination. Defaults to -1 (disabled).
+        cost_relax (float): The relaxation factor for cost reduction. Defaults to 1.0.
+        w_tr_adapt (float): The adaptation factor for the trust region weight. Defaults to 1.0.
+        w_tr_max (float): The maximum allowable trust region weight. Defaults to `None`.
+        w_tr_max_scaling_factor (float): The scaling factor for the maximum trust region weight. Defaults to `None`.
+    """
+
+    def __post_init__(self):
+        keys_to_scale = ["w_tr", "lam_vc", "lam_cost", "lam_vb"]
+        scale = max(getattr(self, key) for key in keys_to_scale)
+        for key in keys_to_scale:
+            setattr(self, key, getattr(self, key) / scale)
+
+        if self.w_tr_max_scaling_factor is not None and self.w_tr_max is None:
+            self.w_tr_max = self.w_tr_max_scaling_factor * self.w_tr
+
+
+@dataclass
+class Config:
+    sim: SimConfig
+    scp: ScpConfig
+    cvx: ConvexSolverConfig
+    dis: DiscretizationConfig
+    prp: PropagationConfig
+    dev: DevConfig
+
+    def __post_init__(self):
+        pass

openscvx/discretization.py

@@ -0,0 +1,169 @@
+import jax.numpy as jnp
+import numpy as np
+
+from openscvx.integrators import solve_ivp_rk45, solve_ivp_diffrax
+
+
+def dVdt(
+    tau: float,
+    V: jnp.ndarray,
+    u_cur: np.ndarray,
+    u_next: np.ndarray,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    n_x: int,
+    n_u: int,
+    N: int,
+    dis_type: str,
+) -> jnp.ndarray:
+    # Define the nodes
+    nodes = jnp.arange(0, N-1)
+
+    # Define indices for slicing the augmented state vector
+    i0 = 0
+    i1 = n_x
+    i2 = i1 + n_x * n_x
+    i3 = i2 + n_x * n_u
+    i4 = i3 + n_x * n_u
+    i5 = i4 + n_x
+
+    # Unflatten V
+    V = V.reshape(-1, i5)
+
+    # Compute the interpolation factor based on the discretization type
+    if dis_type == "ZOH":
+        beta = 0.0
+    elif dis_type == "FOH":
+        beta = (tau) * N
+    alpha = 1 - beta
+
+    # Interpolate the control input
+    u = u_cur + beta * (u_next - u_cur)
+    s = u[:, -1]
+
+    # Initialize the augmented Jacobians
+    dfdx = jnp.zeros((V.shape[0], n_x, n_x))
+    dfdu = jnp.zeros((V.shape[0], n_x, n_u))
+
+    # Ensure x_seq and u have the same batch size
+    x = V[:, :n_x]
+    u = u[: x.shape[0]]
+
+    # Compute the nonlinear propagation term
+    f = state_dot(x, u[:, :-1], nodes)
+    F = s[:, None] * f
+
+    # Evaluate the State Jacobian
+    dfdx = A(x, u[:, :-1], nodes)
+    sdfdx = s[:, None, None] * dfdx
+
+    # Evaluate the Control Jacobian
+    dfdu_veh = B(x, u[:, :-1], nodes)
+    dfdu = dfdu.at[:, :, :-1].set(s[:, None, None] * dfdu_veh)
+    dfdu = dfdu.at[:, :, -1].set(f)
+
+    # Compute the defect
+    z = F - jnp.einsum("ijk,ik->ij", sdfdx, x) - jnp.einsum("ijk,ik->ij", dfdu, u)
+
+    # Stack up the results into the augmented state vector
+    # fmt: off
+    dVdt = jnp.zeros_like(V)
+    dVdt = dVdt.at[:, i0:i1].set(F)
+    dVdt = dVdt.at[:, i1:i2].set(jnp.matmul(sdfdx, V[:, i1:i2].reshape(-1, n_x, n_x)).reshape(-1, n_x * n_x))
+    dVdt = dVdt.at[:, i2:i3].set((jnp.matmul(sdfdx, V[:, i2:i3].reshape(-1, n_x, n_u)) + dfdu * alpha).reshape(-1, n_x * n_u))
+    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()
+
+
+def calculate_discretization(
+    x,
+    u,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    n_x: int,
+    n_u: int,
+    N: int,
+    custom_integrator: bool,
+    debug: bool,
+    solver: str,
+    rtol,
+    atol,
+    dis_type: str,
+):
+
+    # Define indices for slicing the augmented state vector
+    i0 = 0
+    i1 = n_x
+    i2 = i1 + n_x * n_x
+    i3 = i2 + n_x * n_u
+    i4 = i3 + n_x * n_u
+    i5 = i4 + n_x
+
+    # 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(
+        jnp.eye(n_x).reshape(1, -1).repeat(N - 1, axis=0)
+    )
+
+    # choose integrator
+    if custom_integrator:
+        # fmt: off
+        sol = solve_ivp_rk45(
+            lambda t,y,*a: dVdt(t, y, *a),
+            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),
+            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),
+            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,
+            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
+    z_bar = Vend[:, i4:i5]
+    # fmt: on
+
+    return A_bar, B_bar, C_bar, z_bar, Vmulti
+
+
+def get_discretization_solver(state_dot, A, B, params):
+    return lambda x, u: calculate_discretization(
+        x=x,
+        u=u,
+        state_dot=state_dot,
+        A=A,
+        B=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,
+    )

openscvx/dynamics.py

@@ -0,0 +1,24 @@
+import jax
+import jax.numpy as jnp
+
+
+def get_augmented_dynamics(
+    dynamics: callable, g_funcs: list[callable], idx_x_true: slice, idx_u_true: slice
+) -> callable:
+    def dynamics_augmented(x: jnp.array, u: jnp.array, node: int) -> jnp.array:
+        x_dot = dynamics(x[idx_x_true], u[idx_u_true])
+
+        # Iterate through the g_func dictionary and stack the output each function
+        # to x_dot
+        for g in g_funcs:
+            x_dot = jnp.hstack([x_dot, g(x[idx_x_true], u[idx_u_true], node)])
+
+        return x_dot
+
+    return dynamics_augmented
+
+
+def get_jacobians(dyn: callable):
+    A = jax.jacfwd(dyn, argnums=0)
+    B = jax.jacfwd(dyn, argnums=1)
+    return A, B

openscvx/integrators.py

@@ -0,0 +1,139 @@
+import jax
+import jax.numpy as jnp
+import diffrax as dfx
+
+SOLVER_MAP = {
+    "Tsit5": dfx.Tsit5,
+    "Euler": dfx.Euler,
+    "Heun": dfx.Heun,
+    "Midpoint": dfx.Midpoint,
+    "Ralston": dfx.Ralston,
+    "Dopri5": dfx.Dopri5,
+    "Dopri8": dfx.Dopri8,
+    "Bosh3": dfx.Bosh3,
+    "ReversibleHeun": dfx.ReversibleHeun,
+    "ImplicitEuler": dfx.ImplicitEuler,
+    "KenCarp3": dfx.KenCarp3,
+    "KenCarp4": dfx.KenCarp4,
+    "KenCarp5": dfx.KenCarp5,
+}
+
+# fmt: off
+def rk45_step(f, t, y, h, *args):
+    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)
+    k4 = f(t + 12*h/13, y + 1932*h*k1/2197 - 7200*h*k2/2197 + 7296*h*k3/2197, *args)
+    k5 = f(t + h, y + 439*h*k1/216 - 8*h*k2 + 3680*h*k3/513 - 845*h*k4/4104, *args)
+    y_next = y + h * (25*k1/216 + 1408*k3/2565 + 2197*k4/4104 - k5/5)
+    return y_next
+# fmt: on
+
+
+def solve_ivp_rk45(
+    f,
+    tau_final: float,
+    y_0,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    is_not_compiled: bool = False,
+):
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    h = (tau_final - tau_0) / (len(substeps) - 1)
+    solution = jnp.zeros((len(substeps), len(y_0)))
+    solution = solution.at[0].set(y_0)
+
+    if is_not_compiled:
+        for i in range(1, len(substeps)):
+            t = tau_0 + i * h
+            solution = solution.at[i].set(rk45_step(f, t, solution[i - 1], h, *args))
+    else:
+
+        def body_fun(i, val):
+            t, y, V_result = val
+            y_next = rk45_step(f, t, y, h, *args)
+            V_result = V_result.at[i].set(y_next)
+            return (t + h, y_next, V_result)
+
+        _, _, solution = jax.lax.fori_loop(
+            1, len(substeps), body_fun, (tau_0, y_0, solution)
+        )
+
+    return solution
+
+
+def solve_ivp_diffrax(
+    f,
+    tau_final,
+    y_0,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    solver_name="Dopri8",
+    rtol: float = 1e-3,
+    atol: float = 1e-6,
+    extra_kwargs=None,
+):
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    solver_class = SOLVER_MAP.get(solver_name)
+    if solver_class is None:
+        raise ValueError(f"Unknown solver: {solver_name}")
+    solver = solver_class()
+
+    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),
+        y0=y_0,
+        args=args,
+        stepsize_controller=stepsize_controller,
+        saveat=dfx.SaveAt(ts=substeps),
+        **(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,
+    args,
+    tau_0: float = 0.0,
+    num_substeps: int = 50,
+    solver_name="Dopri8",
+    rtol: float = 1e-3,
+    atol: float = 1e-6,
+    extra_kwargs=None,
+):
+    substeps = jnp.linspace(tau_0, tau_final, num_substeps)
+
+    solver_class = SOLVER_MAP.get(solver_name)
+    if solver_class is None:
+        raise ValueError(f"Unknown solver: {solver_name}")
+    solver = solver_class()
+
+    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),
+        y0=y_0,
+        args=args,
+        stepsize_controller=stepsize_controller,
+        saveat=dfx.SaveAt(dense=True, ts=substeps),
+        **(extra_kwargs or {}),
+    )
+
+    return solution