openscvx 0.4.1.dev153__tar.gz → 0.4.1.dev155__tar.gz

{openscvx-0.4.1.dev153/openscvx.egg-info → openscvx-0.4.1.dev155}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.4.1.dev153
+Version: 0.4.1.dev155
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

{openscvx-0.4.1.dev153 → openscvx-0.4.1.dev155}/openscvx/__init__.py

@@ -18,7 +18,7 @@ from openscvx.algorithms import (
     RampProximalWeight,
 )
 from openscvx.algorithms.optimization_results import OptimizationResults
-from openscvx.discretization import LinearizeDiscretize
+from openscvx.discretization import LinearizeDiscretize, LinearizeDiscretizeSparse
 from openscvx.expert import ByofSpec
 from openscvx.loader import load_dict, load_json, load_yaml
 from openscvx.problem import Problem
@@ -173,6 +173,7 @@ __all__ = [
     "ByofSpec",
     # Discretization
     "LinearizeDiscretize",
+    "LinearizeDiscretizeSparse",
     # Convex Solver
     "PTRSolver",
     # Algorithm & Autotuning

openscvx-0.4.1.dev155/openscvx/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.4.1.dev155'
+__version_tuple__ = version_tuple = (0, 4, 1, 'dev155')
+
+__commit_id__ = commit_id = 'gbd34ba2e8'

{openscvx-0.4.1.dev153 → openscvx-0.4.1.dev155}/openscvx/algorithms/penalized_trust_region.py

@@ -347,8 +347,8 @@ class PenalizedTrustRegion(Algorithm):
         init_state = AlgorithmState.from_settings(settings, self.weights)
 
         # Solve a dumb problem to initialize DPP and JAX jacobians
-        _, _, _, x_prop, V_multi_shoot = self._discretization_solver.call(
-            init_state.x, init_state.u.astype(float), params
+        _, _, _, x_prop, V_multi_shoot = self._invoke_solver(
+            self._discretization_solver, init_state.x, init_state.u.astype(float), params
         )
 
         init_state.add_discretization(V_multi_shoot.__array__())
@@ -399,15 +399,15 @@ class PenalizedTrustRegion(Algorithm):
         # Compute discretization before subproblem only for the first iteration
         if state.k == 1:
             t0 = time.time()
-            _, _, _, x_prop, V_multi_shoot = self._discretization_solver.call(
-                state.x, state.u.astype(float), params
+            _, _, _, x_prop, V_multi_shoot = self._invoke_solver(
+                self._discretization_solver, state.x, state.u.astype(float), params
             )
 
             u_state = state.u.astype(float)
             x0_prior = self._recover_prior_node_from_initial(settings, state.x[0])
             x_nodes_prior = np.vstack((x0_prior, np.asarray(x_prop)))
-            _, _, _, W_multi_shoot = self._discretization_solver_impulsive.call(
-                x_nodes_prior, u_state, params
+            _, _, _, W_multi_shoot = self._invoke_solver(
+                self._discretization_solver_impulsive, x_nodes_prior, u_state, params
             )
             dis_time = time.time() - t0
 
@@ -435,15 +435,15 @@ class PenalizedTrustRegion(Algorithm):
         candidate.J_lin = J_total
 
         t0 = time.time()
-        _, _, _, x_prop, V_multi_shoot = self._discretization_solver.call(
-            candidate.x, candidate.u.astype(float), params
+        _, _, _, x_prop, V_multi_shoot = self._invoke_solver(
+            self._discretization_solver, candidate.x, candidate.u.astype(float), params
         )
 
         u_candidate = candidate.u.astype(float)
         x0_prior = self._recover_prior_node_from_initial(settings, candidate.x[0])
         x_nodes_prior = np.vstack((x0_prior, np.asarray(x_prop)))
-        x_prop_plus, D_d, E_d, W_multi_shoot = self._discretization_solver_impulsive.call(
-            x_nodes_prior, u_candidate, params
+        x_prop_plus, D_d, E_d, W_multi_shoot = self._invoke_solver(
+            self._discretization_solver_impulsive, x_nodes_prior, u_candidate, params
         )
 
         dis_time = time.time() - t0

{openscvx-0.4.1.dev153 → openscvx-0.4.1.dev155}/openscvx/discretization/__init__.py

@@ -11,9 +11,9 @@ intermediate types, but the input (continuous nonlinear dynamics + reference
 trajectory) and output (discrete-time linear matrices A_d, B_d, C_d) are
 always consistent.
 
-The default implementation is :class:`LinearizeDiscretize`, which computes
-continuous-time Jacobians via JAX autodiff and integrates them alongside the
-nonlinear dynamics through an augmented state vector.
+:class:`Problem` uses :class:`LinearizeDiscretizeSparse` by default (sparse
+Jacobians and compact variational integration when sparsity patterns exist).
+:class:`LinearizeDiscretize` is the dense linearize-then-discretize scheme.
 """
 
 import inspect
@@ -25,6 +25,8 @@ from .linearize_discretize import (
     calculate_impulsive_discretization,
     get_impulsive_discretization_solver,
 )
+from .linearize_discretize_sparse import LinearizeDiscretizeSparse
+from .sparse_utils import color_columns, make_sparse_jacobian_fns
 
 # ---------------------------------------------------------------------------
 # Spec resolver — turn a dict into a Discretizer instance
@@ -32,6 +34,7 @@ from .linearize_discretize import (
 
 _DISCRETIZER_MAP = {
     "LinearizeDiscretize": LinearizeDiscretize,
+    "LinearizeDiscretizeSparse": LinearizeDiscretizeSparse,
 }
 
 
@@ -41,16 +44,21 @@ def _resolve_discretizer(val: Any) -> Discretizer:
     Accepted forms:
 
     * **instance** — already-constructed :class:`Discretizer` (pass-through).
-    * **dict** — keyword arguments passed to :class:`LinearizeDiscretize`.
-      An optional ``"type"`` key selects the class (currently only
-      ``"LinearizeDiscretize"``).
+    * **dict** — keyword arguments passed to the selected discretizer class.
+      An optional ``"type"`` key selects the class (defaults to
+      :class:`LinearizeDiscretizeSparse`).
 
     Examples::
 
-        # Dict with keyword overrides (default class)
+        # Dict with keyword overrides (default class: LinearizeDiscretizeSparse)
         _resolve_discretizer({"dis_type": "ZOH", "ode_solver": "Dopri8"})
 
-        # Dict with explicit type
+        # Configure integrator behavior (forwarded to Diffrax / diffeqsolve)
+        _resolve_discretizer(
+            {"diffrax_kwargs": {"num_substeps": 100, "max_steps": 20_000}}
+        )
+
+        # Dict with explicit dense discretizer
         _resolve_discretizer({"type": "LinearizeDiscretize", "dis_type": "ZOH"})
 
         # Instance pass-through
@@ -82,7 +90,10 @@ def _resolve_discretizer(val: Any) -> Discretizer:
 __all__ = [
     "Discretizer",
     "LinearizeDiscretize",
+    "LinearizeDiscretizeSparse",
     "_resolve_discretizer",
     "calculate_impulsive_discretization",
+    "color_columns",
     "get_impulsive_discretization_solver",
+    "make_sparse_jacobian_fns",
 ]

{openscvx-0.4.1.dev153 → openscvx-0.4.1.dev155}/openscvx/discretization/linearize_discretize.py

@@ -1,11 +1,16 @@
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 
 from openscvx.discretization.base import Discretizer
-from openscvx.integrators import solve_ivp_diffrax, solve_ivp_rk45
+from openscvx.integrators import (
+    DEFAULT_DIFFRAX_ATOL,
+    DEFAULT_DIFFRAX_RTOL,
+    solve_ivp_diffrax,
+    solve_ivp_rk45,
+)
 
 if TYPE_CHECKING:
     from openscvx.config import Config
@@ -31,29 +36,61 @@ class LinearizeDiscretize(Discretizer):
         ode_solver: Diffrax solver name. Any solver from
             `Diffrax <https://docs.kidger.site/diffrax/usage/how-to-choose-a-solver/>`_
             is valid. Defaults to ``"Tsit5"``.
-        custom_integrator: Use the built-in fixed-step RK45 integrator
-            instead of Diffrax. Faster but less robust. Defaults to ``False``.
-        atol: Absolute tolerance for the ODE solver. Defaults to ``1e-3``.
-        rtol: Relative tolerance for the ODE solver. Defaults to ``1e-6``.
-        args: Extra keyword arguments forwarded to
-            :func:`diffrax.diffeqsolve`. Defaults to ``{}``.
+        diffrax_kwargs: Preferred Diffrax keyword overrides. These map to
+            :func:`openscvx.integrators.solve_ivp_diffrax` kwargs, and unknown
+            keys are forwarded to :func:`diffrax.diffeqsolve` (e.g.
+            ``stepsize_controller``). Set ``rtol``/``atol`` here when using
+            the default PID controller. Defaults to ``{}``.
     """
 
     def __init__(
         self,
         dis_type: str = "FOH",
         ode_solver: str = "Tsit5",
-        custom_integrator: bool = False,
-        atol: float = 1e-3,
-        rtol: float = 1e-6,
-        args: Optional[dict] = None,
+        diffrax_kwargs: Optional[dict[str, Any]] = None,
     ):
         self.dis_type = dis_type
         self.ode_solver = ode_solver
-        self.custom_integrator = custom_integrator
-        self.atol = atol
-        self.rtol = rtol
-        self.args = args if args is not None else {}
+        self.diffrax_kwargs = dict(diffrax_kwargs) if diffrax_kwargs is not None else {}
+
+    def _resolve_diffrax_kwargs(self) -> dict[str, Any]:
+        """Build kwargs for :func:`solve_ivp_diffrax`.
+
+        Unknown keys from ``self.diffrax_kwargs`` are forwarded to
+        :func:`diffrax.diffeqsolve` via ``extra_kwargs`` so users can pass
+        objects like ``stepsize_controller`` directly from discretizer settings.
+        """
+        kwargs: dict[str, Any] = {
+            "solver_name": self.ode_solver,
+            "rtol": DEFAULT_DIFFRAX_RTOL,
+            "atol": DEFAULT_DIFFRAX_ATOL,
+        }
+
+        user_kwargs = dict(self.diffrax_kwargs)
+        extra_kwargs: dict[str, Any] = {}
+
+        nested_extra = user_kwargs.pop("extra_kwargs", None)
+        if nested_extra is not None:
+            extra_kwargs.update(dict(nested_extra))
+
+        direct_keys = {"tau_0", "num_substeps", "solver_name", "rtol", "atol"}
+        for key, value in user_kwargs.items():
+            if key in direct_keys:
+                kwargs[key] = value
+            else:
+                extra_kwargs[key] = value
+
+        kwargs["extra_kwargs"] = extra_kwargs
+        return kwargs
+
+    def _resolve_rk45_kwargs(self, *, is_not_compiled: bool) -> dict[str, Any]:
+        """Build kwargs for :func:`solve_ivp_rk45`."""
+        kwargs: dict[str, Any] = {"is_not_compiled": is_not_compiled}
+        direct_keys = {"tau_0", "num_substeps", "is_not_compiled"}
+        for key, value in self.diffrax_kwargs.items():
+            if key in direct_keys:
+                kwargs[key] = value
+        return kwargs
 
     def get_solver(self, dynamics: "Dynamics", settings: "Config") -> callable:
         """Create a multi-shoot discretization solver.
@@ -318,25 +355,23 @@ def _calculate_discretization(
     def dVdt_wrapped(t, y):
         return _dVdt(t, y, **integrator_args)
 
-    # Choose integrator
-    if discretizer.custom_integrator:
+    if settings.dev.debug:
+        rk45_kwargs = discretizer._resolve_rk45_kwargs(is_not_compiled=settings.dev.debug)
         sol = solve_ivp_rk45(
             dVdt_wrapped,
             1.0 / (N - 1),
             V0,
             args=(),
-            is_not_compiled=settings.dev.debug,
+            **rk45_kwargs,
         )
     else:
+        diffrax_kwargs = discretizer._resolve_diffrax_kwargs()
         sol = solve_ivp_diffrax(
             dVdt_wrapped,
             1.0 / (N - 1),
             V0,
-            solver_name=discretizer.ode_solver,
-            rtol=discretizer.rtol,
-            atol=discretizer.atol,
             args=(),
-            extra_kwargs=discretizer.args,
+            **diffrax_kwargs,
         )
 
     Vend = sol[-1].T.reshape(-1, i4)

openscvx-0.4.1.dev155/openscvx/discretization/linearize_discretize_sparse.py

@@ -0,0 +1,346 @@
+from typing import TYPE_CHECKING
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from openscvx.discretization.linearize_discretize import LinearizeDiscretize
+from openscvx.integrators import solve_ivp_diffrax, solve_ivp_rk45
+
+if TYPE_CHECKING:
+    from openscvx.config import Config
+    from openscvx.lowered import Dynamics
+
+
+class LinearizeDiscretizeSparse(LinearizeDiscretize):
+    """Sparse variant of linearize-then-discretize.
+
+    Uses graph-coloring-based sparse Jacobian computation and a compact
+    augmented state vector that only integrates the structurally nonzero
+    entries of Φ, B_d and C_d.  This reduces the ODE system dimension
+    from ``n_x + n_x² + 2·n_x·n_u`` to ``n_x + nnz_Ad + nnz_Bd + nnz_Cd``
+    per segment.
+
+    Requires ``A_c_sparsity`` and ``B_c_sparsity`` boolean arrays on the
+    :class:`Dynamics` object (set automatically when using the symbolic
+    problem interface).  Falls back to the dense
+    :class:`LinearizeDiscretize` path when sparsity patterns are
+    unavailable or fully dense.
+
+    Args:
+        dis_type: Control hold type. ``"FOH"`` or ``"ZOH"``.
+            Defaults to ``"FOH"``.
+        ode_solver: Diffrax solver name. Defaults to ``"Tsit5"``.
+        diffrax_kwargs: Diffrax keyword overrides inherited from
+            :class:`LinearizeDiscretize`. Unknown keys are forwarded to
+            :func:`diffrax.diffeqsolve` via ``extra_kwargs``.
+    """
+
+    def get_solver(self, dynamics: "Dynamics", settings: "Config") -> callable:
+        """Create a sparse multi-shoot discretization solver.
+
+        When ``dynamics.A_c_sparsity`` and ``dynamics.B_c_sparsity`` are
+        available and the pattern is not fully dense, builds a compact-V
+        integration path with graph-coloring Jacobians.  Otherwise
+        delegates to the dense parent implementation.
+
+        Args:
+            dynamics: System dynamics with optional sparsity annotations.
+            settings: Problem configuration.
+
+        Returns:
+            Callable ``(x, u, params) -> (A_d, B_d, C_d, x_prop, V)``.
+        """
+        from openscvx.symbolic.sparsity import discrete_sparsity
+
+        from .sparse_utils import make_sparse_jacobian_fns
+
+        A_c_pat = getattr(dynamics, "A_c_sparsity", None)
+        B_c_pat = getattr(dynamics, "B_c_sparsity", None)
+        has_sparsity = A_c_pat is not None and B_c_pat is not None
+
+        if not has_sparsity or A_c_pat.all():
+            return super().get_solver(dynamics, settings)
+
+        f_vmapped = jax.vmap(dynamics.f, in_axes=(0, 0, 0, None))
+        discretizer = self
+        n_x = settings.sim.n_states
+        n_u = settings.sim.n_controls
+
+        A_vmapped, B_vmapped = make_sparse_jacobian_fns(
+            dynamics.f,
+            A_c_pat,
+            B_c_pat,
+            n_x,
+            n_u,
+        )
+
+        Ad_pat, Bd_pat, Cd_pat = discrete_sparsity(
+            A_c_pat,
+            B_c_pat,
+            self.dis_type,
+        )
+        Ad_r, Ad_c = np.where(Ad_pat)
+        Bd_r, Bd_c = np.where(Bd_pat)
+        Cd_r, Cd_c = np.where(Cd_pat)
+
+        sparse_layout = (
+            jnp.array(Ad_r),
+            jnp.array(Ad_c),
+            len(Ad_r),
+            jnp.array(Bd_r),
+            jnp.array(Bd_c),
+            len(Bd_r),
+            jnp.array(Cd_r),
+            jnp.array(Cd_c),
+            len(Cd_r),
+        )
+
+        return lambda x, u, params: _calculate_discretization_sparse(
+            x=x,
+            u=u,
+            state_dot=f_vmapped,
+            A=A_vmapped,
+            B=B_vmapped,
+            settings=settings,
+            discretizer=discretizer,
+            params=params,
+            sparse_layout=sparse_layout,
+        )
+
+
+def _dVdt_sparse(
+    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,
+    S_x: np.ndarray,
+    c_x: np.ndarray,
+    S_u: np.ndarray,
+    c_u: np.ndarray,
+    inv_S_x: np.ndarray,
+    inv_S_u: np.ndarray,
+    params: dict,
+    Ad_rows: jnp.ndarray,
+    Ad_cols: jnp.ndarray,
+    nnz_Ad: int,
+    Bd_rows: jnp.ndarray,
+    Bd_cols: jnp.ndarray,
+    nnz_Bd: int,
+    Cd_rows: jnp.ndarray,
+    Cd_cols: jnp.ndarray,
+    nnz_Cd: int,
+) -> jnp.ndarray:
+    """Time derivative of the *compact* augmented state for sparse variational integration.
+
+    Instead of storing the full flattened Φ, B_d, C_d matrices (``n_x²`` and
+    ``n_x·n_u`` entries each), this function only tracks the structurally
+    nonzero entries as determined by ``discrete_sparsity``.
+
+    The compact layout per segment is::
+
+        V = [x(n_x), Φ_nz(nnz_Ad), B_d_nz(nnz_Bd), C_d_nz(nnz_Cd)]
+
+    At each evaluation the compact values are scattered into dense matrices
+    for the matmul, and the derivative is gathered back at the nonzero
+    positions.
+
+    Args:
+        tau: Normalized time in [0, 1] within the current segment.
+        V: Flattened compact augmented state, shape
+            ``((N-1) * (n_x + nnz_Ad + nnz_Bd + nnz_Cd),)``.
+        Ad_rows, Ad_cols: Row/column indices of A_d structural nonzeros.
+        nnz_Ad: Number of A_d structural nonzeros.
+        Bd_rows, Bd_cols: Row/column indices of B_d structural nonzeros.
+        nnz_Bd: Number of B_d structural nonzeros.
+        Cd_rows, Cd_cols: Row/column indices of C_d structural nonzeros.
+        nnz_Cd: Number of C_d structural nonzeros.
+        (remaining args identical to :func:`_dVdt`)
+
+    Returns:
+        Flattened time derivative of the compact augmented state.
+    """
+    nodes = jnp.arange(0, N - 1)
+
+    aug_dim = n_x + nnz_Ad + nnz_Bd + nnz_Cd
+    i_phi = n_x
+    i_bd = n_x + nnz_Ad
+    i_cd = n_x + nnz_Ad + nnz_Bd
+
+    V = V.reshape(-1, aug_dim)
+
+    if dis_type == "ZOH":
+        beta = 0.0
+    elif dis_type == "FOH":
+        beta = tau * N
+    alpha = 1 - beta
+
+    u = u_cur + beta * (u_next - u_cur)
+    x = V[:, :n_x]
+    u = u[: x.shape[0]]
+    batch = x.shape[0]
+
+    F = state_dot(x, u, nodes, params)
+    dfdx = A(x, u, nodes, params)
+    dfdu = B(x, u, nodes, params)
+
+    # Φ: scatter compact → dense, matmul, gather back
+    phi_nz = V[:, i_phi:i_bd]
+    Phi = jnp.zeros((batch, n_x, n_x)).at[:, Ad_rows, Ad_cols].set(phi_nz)
+    dPhi_nz = jnp.matmul(dfdx, Phi)[:, Ad_rows, Ad_cols]
+
+    # B_d: scatter, matmul + forcing, gather
+    bd_nz = V[:, i_bd:i_cd]
+    Bd = jnp.zeros((batch, n_x, n_u)).at[:, Bd_rows, Bd_cols].set(bd_nz)
+    dBd_nz = (jnp.matmul(dfdx, Bd) + dfdu * alpha)[:, Bd_rows, Bd_cols]
+
+    # C_d: scatter, matmul + forcing, gather
+    cd_nz = V[:, i_cd:]
+    Cd = jnp.zeros((batch, n_x, n_u)).at[:, Cd_rows, Cd_cols].set(cd_nz)
+    dCd_nz = (jnp.matmul(dfdx, Cd) + dfdu * beta)[:, Cd_rows, Cd_cols]
+
+    return jnp.concatenate([F, dPhi_nz, dBd_nz, dCd_nz], axis=-1).reshape(-1)
+
+
+def _calculate_discretization_sparse(
+    x: np.ndarray,
+    u: np.ndarray,
+    state_dot: callable,
+    A: callable,
+    B: callable,
+    settings: "Config",
+    discretizer: "LinearizeDiscretizeSparse",
+    params: dict,
+    sparse_layout: tuple,
+) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+    """Integrate the compact variational equations to produce discrete-time matrices.
+
+    Uses the structurally nonzero entries of Φ, B_d and C_d to form a
+    reduced augmented state, integrates with :func:`_dVdt_sparse`, then
+    scatters the results back into dense matrices for downstream
+    compatibility.
+
+    Args:
+        x: Reference state trajectory, shape ``(N, n_x)``.
+        u: Reference control trajectory, shape ``(N, n_u)``.
+        state_dot: Vmapped time-dilated dynamics.
+        A: Vmapped state Jacobian ``∂F/∂x``.
+        B: Vmapped control Jacobian ``∂F/∂u``.
+        settings: Problem configuration.
+        discretizer: Discretizer instance with integrator settings.
+        params: Parameters forwarded to dynamics callables.
+        sparse_layout: Tuple
+            ``(Ad_rows, Ad_cols, nnz_Ad, Bd_rows, Bd_cols, nnz_Bd,
+            Cd_rows, Cd_cols, nnz_Cd)``.
+
+    Returns:
+        Tuple ``(A_d, B_d, C_d, x_prop, V)``.
+    """
+    n_x = settings.sim.n_states
+    n_u = settings.sim.n_controls
+    N = settings.sim.n
+
+    (
+        Ad_rows,
+        Ad_cols,
+        nnz_Ad,
+        Bd_rows,
+        Bd_cols,
+        nnz_Bd,
+        Cd_rows,
+        Cd_cols,
+        nnz_Cd,
+    ) = sparse_layout
+
+    aug_dim = n_x + nnz_Ad + nnz_Bd + nnz_Cd
+
+    V0 = jnp.zeros((N - 1, aug_dim))
+    V0 = V0.at[:, :n_x].set(x[:-1].astype(float))
+    phi0_nz = (Ad_rows == Ad_cols).astype(x.dtype)
+    V0 = V0.at[:, n_x : n_x + nnz_Ad].set(jnp.broadcast_to(phi0_nz[None], (N - 1, nnz_Ad)))
+    V0 = V0.reshape(-1)
+
+    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=discretizer.dis_type,
+        S_x=settings.sim.S_x,
+        c_x=settings.sim.c_x,
+        S_u=settings.sim.S_u,
+        c_u=settings.sim.c_u,
+        inv_S_x=settings.sim.inv_S_x,
+        inv_S_u=settings.sim.inv_S_u,
+        params=params,
+        Ad_rows=Ad_rows,
+        Ad_cols=Ad_cols,
+        nnz_Ad=nnz_Ad,
+        Bd_rows=Bd_rows,
+        Bd_cols=Bd_cols,
+        nnz_Bd=nnz_Bd,
+        Cd_rows=Cd_rows,
+        Cd_cols=Cd_cols,
+        nnz_Cd=nnz_Cd,
+    )
+
+    def dVdt_wrapped(t, y):
+        return _dVdt_sparse(t, y, **integrator_args)
+
+    if settings.dev.debug:
+        rk45_kwargs = discretizer._resolve_rk45_kwargs(is_not_compiled=settings.dev.debug)
+        sol = solve_ivp_rk45(
+            dVdt_wrapped,
+            1.0 / (N - 1),
+            V0,
+            args=(),
+            **rk45_kwargs,
+        )
+    else:
+        diffrax_kwargs = discretizer._resolve_diffrax_kwargs()
+        sol = solve_ivp_diffrax(
+            dVdt_wrapped,
+            1.0 / (N - 1),
+            V0,
+            args=(),
+            **diffrax_kwargs,
+        )
+
+    Vend = sol[-1].T.reshape(-1, aug_dim)
+
+    x_prop = Vend[:, :n_x]
+    phi_nz = Vend[:, n_x : n_x + nnz_Ad]
+    bd_nz = Vend[:, n_x + nnz_Ad : n_x + nnz_Ad + nnz_Bd]
+    cd_nz = Vend[:, n_x + nnz_Ad + nnz_Bd :]
+
+    A_bar = jnp.zeros((N - 1, n_x, n_x)).at[:, Ad_rows, Ad_cols].set(phi_nz)
+    B_bar = jnp.zeros((N - 1, n_x, n_u)).at[:, Bd_rows, Bd_cols].set(bd_nz)
+    C_bar = jnp.zeros((N - 1, n_x, n_u))
+    if nnz_Cd > 0:
+        C_bar = C_bar.at[:, Cd_rows, Cd_cols].set(cd_nz)
+
+    # Reconstruct a dense-layout Vmulti so that
+    # DiscretizationResult.from_V can unpack it unchanged.
+    Vend_dense = jnp.concatenate(
+        [
+            x_prop,
+            A_bar.reshape(N - 1, n_x * n_x),
+            B_bar.reshape(N - 1, n_x * n_u),
+            C_bar.reshape(N - 1, n_x * n_u),
+        ],
+        axis=-1,
+    )
+    Vmulti = Vend_dense.reshape(-1, 1)
+
+    return A_bar, B_bar, C_bar, x_prop, Vmulti

openscvx-0.4.1.dev155/openscvx/discretization/sparse_utils/__init__.py

@@ -0,0 +1,11 @@
+"""Sparse linear algebra helpers built on ``jax.experimental.sparse``."""
+
+from .bcoo_helpers import precompute_sparse_indices, sparse_matmul_batched
+from .sparse_jacobian import color_columns, make_sparse_jacobian_fns
+
+__all__ = [
+    "color_columns",
+    "make_sparse_jacobian_fns",
+    "precompute_sparse_indices",
+    "sparse_matmul_batched",
+]