solvers

Convex subproblem solvers for trajectory optimization.

This module provides implementations of convex subproblem solvers used within SCvx algorithms. At each iteration of a successive convexification algorithm, the non-convex problem is approximated by a convex subproblem, which is then solved using one of these solver backends.

All solvers inherit from :class:ConvexSolver, enabling pluggable solver implementations and custom backends:

class ConvexSolver(ABC):
    @abstractmethod
    def create_variables(self, N, x_unified, u_unified, jax_constraints) -> None:
        '''Create backend-specific optimization variables (called once).'''
        ...

    @abstractmethod
    def initialize(self, lowered, settings) -> None:
        '''Build the convex subproblem structure (called once).'''
        ...

    @abstractmethod
    def solve(self, state, params, settings) -> Any:
        '''Update parameters and solve (called each iteration).'''
        ...

The Penalized Trust-Region (PTR) subproblem ships with three concrete backends:

• :class:CVXPyPTRSolver — DCP graph via CVXPy, dispatched to any of its supported conic solvers (QOCO, CLARABEL, ...). Optional code generation via cvxpygen for improved per-iteration performance.
• :class:QPAXPTRSolver — flat (Q, q, A, b, G, h) assembled as JAX arrays and solved with qpax.solve_qp. Aimed at end-to-end JAX differentiability of the SCP loop (follow-up work). Requires the qpax optional extra.
• :class:MoreauPTRSolver — sparse conic program (SOCP-capable) assembled as CSR JAX arrays and solved with moreau.jax.Solver. Uses SOC epigraphs for the PTR penalties instead of QPAX-style slack expansion; warm-starts between iterations. Requires the moreau optional extra.

All three share the abstract :class:PTRSolver contract.

Note

Solvers own their optimization variables (e.g., CVXPySolver.ocp_vars). The lowering process calls solver.create_variables() before constraint lowering, then solver.initialize() after. See :mod:openscvx.solvers.base for the interface details.

CVXPyPTRSolver

Bases: PTRSolver

CVXPy-backed implementation of the PTR convex subproblem.

Builds the subproblem as a DCP program through CVXPy and dispatches it to one of CVXPy's supported conic solvers (QOCO by default, CLARABEL, etc.). Optional code generation via cvxpygen is available for improved per-iteration performance.

The solver builds the problem structure once during initialize(), using CVXPy Parameters for values that change each iteration. The solve() method then solves and returns a structured PTRSolveResult. The JAX-pure entry point :meth:iteration_callback wraps that same solve in :func:jax.pure_callback (vmap_method="sequential" — CVXPy cannot batch) so the backend composes with jax.jit / jax.vmap alongside the JAX-native QPAX and Moreau backends.

The cost and constraint formulations are defined in the cost() and constraints() methods, which can be overridden in subclasses to customize the convex subproblem. For example::

class MyPTRSolver(CVXPyPTRSolver):
    def cost(self, settings, lowered):
        c = super().cost(settings, lowered)
        c += my_extra_term(self._ocp_vars)
        return c

Example

Using CVXPyPTRSolver with the SCP framework::

solver = CVXPyPTRSolver()
solver.create_variables(N, x_unified, u_unified, jax_constraints)
solver.initialize(lowered, settings)

# Each iteration (parameter updates done by algorithm):
result = solver.solve()
x_sol = result.x  # Unscaled state trajectory

Parameters:

Name	Type	Description	Default
cvx_solver	str	CVXPY solver backend name. Defaults to "QOCO".	'QOCO'
solver_args	Optional[dict]	Keyword arguments forwarded to the CVXPY solver (e.g. tolerances). Defaults to {"abstol": 1e-6, "reltol": 1e-9, "enforce_dpp": True}.	None
cvxpygen	bool	Enable CVXPy code generation for faster solves. Defaults to False. Warning Enabling cvxpygen currently disables sparse parameter declarations. cvxpygen does not yet support the N-D sparsity indices used by OpenSCvx's tiled parameters, so all parameters are created as dense when code generation is active. This may increase the generated solver's memory footprint and compile time but does not affect solution correctness.	False
cvxpygen_override	bool	Overwrite existing generated solver directory without prompting. Defaults to False.	False

Attributes:

Name	Type	Description
ocp_vars	CVXPyVariables	The CVXPy variables and parameters (available after create_variables())

Source code in openscvx/solvers/cvxpy_ptr_solver.py

class CVXPyPTRSolver(PTRSolver):
    """CVXPy-backed implementation of the PTR convex subproblem.

    Builds the subproblem as a DCP program through CVXPy and dispatches it to
    one of CVXPy's supported conic solvers (QOCO by default, CLARABEL, etc.).
    Optional code generation via cvxpygen is available for improved per-iteration
    performance.

    The solver builds the problem structure once during ``initialize()``, using
    CVXPy Parameters for values that change each iteration. The ``solve()``
    method then solves and returns a structured ``PTRSolveResult``. The
    JAX-pure entry point :meth:`iteration_callback` wraps that same solve in
    :func:`jax.pure_callback` (``vmap_method="sequential"`` — CVXPy cannot
    batch) so the backend composes with ``jax.jit`` / ``jax.vmap`` alongside
    the JAX-native QPAX and Moreau backends.

    The cost and constraint formulations are defined in the ``cost()`` and
    ``constraints()`` methods, which can be overridden in subclasses to
    customize the convex subproblem. For example::

        class MyPTRSolver(CVXPyPTRSolver):
            def cost(self, settings, lowered):
                c = super().cost(settings, lowered)
                c += my_extra_term(self._ocp_vars)
                return c

    Example:
        Using CVXPyPTRSolver with the SCP framework::

            solver = CVXPyPTRSolver()
            solver.create_variables(N, x_unified, u_unified, jax_constraints)
            solver.initialize(lowered, settings)

            # Each iteration (parameter updates done by algorithm):
            result = solver.solve()
            x_sol = result.x  # Unscaled state trajectory

    Args:
        cvx_solver: CVXPY solver backend name. Defaults to ``"QOCO"``.
        solver_args: Keyword arguments forwarded to the CVXPY solver
            (e.g. tolerances). Defaults to
            ``{"abstol": 1e-6, "reltol": 1e-9, "enforce_dpp": True}``.
        cvxpygen: Enable CVXPy code generation for faster solves.
            Defaults to ``False``.

            !!! warning
                Enabling cvxpygen currently disables sparse parameter
                declarations. cvxpygen does not yet support the N-D sparsity
                indices used by OpenSCvx's tiled parameters, so all parameters
                are created as dense when code generation is active. This may
                increase the generated solver's memory footprint and compile
                time but does not affect solution correctness.
        cvxpygen_override: Overwrite existing generated solver directory
            without prompting. Defaults to ``False``.

    Attributes:
        ocp_vars: The CVXPy variables and parameters (available after create_variables())
    """

    def __init__(
        self,
        cvx_solver: str = "QOCO",
        solver_args: Optional[dict] = None,
        cvxpygen: bool = False,
        cvxpygen_override: bool = False,
    ):
        """Initialize CVXPyPTRSolver with solver configuration.

        Call create_variables() then initialize() to build the problem structure.
        """
        self.cvx_solver = cvx_solver
        self.solver_args = (
            solver_args
            if solver_args is not None
            else {"abstol": 1e-06, "reltol": 1e-09, "enforce_dpp": True, "warm_start": False}
        )
        self.cvxpygen = cvxpygen
        self.cvxpygen_override = cvxpygen_override

        self._ocp_vars: "CVXPyVariables" = None
        self._problem: cp.Problem = None
        self._solve_fn: callable = None
        # Stashed at create_variables / initialize so iteration_callback can
        # close over them without retracing them off the ``lowered`` /
        # ``settings`` arguments each iteration.
        self._jax_constraints: Optional["LoweredJaxConstraints"] = None
        self._settings: Optional[Config] = None
        self._slice_cont = slice(0, 0)
        self._slice_imp = slice(0, 0)

    @property
    def exportable(self) -> bool:
        """CVXPy is not ``jax.export``-serializable — see :attr:`PTRSolver.exportable`.

        The JAX-pure :meth:`iteration_callback` wraps the host CVXPy solve in a
        :func:`jax.pure_callback`, and ``jax.export`` cannot serialize host
        callbacks. ``solve_batched(save_compiled=True)`` therefore refuses this
        backend with a teaching error pointing at QPAX / Moreau rather than
        silently falling back to an uncached in-process solve.
        """
        return False

    @property
    def ocp_vars(self) -> "CVXPyVariables":
        """The CVXPy variables and parameters.

        Returns:
            The CVXPyVariables dataclass, or None if create_variables() not called.
        """
        return self._ocp_vars

    def create_variables(
        self,
        N: int,
        x_unified: "UnifiedState",
        u_unified: "UnifiedControl",
        jax_constraints: "LoweredJaxConstraints",
        dynamics_sparsity: Optional[tuple] = None,
        constraint_sparsity: Optional[list] = None,
    ) -> None:
        """Create CVXPy optimization variables.

        Creates all CVXPy Variable and Parameter objects needed for the optimal
        control problem. This includes state/control variables, dynamics parameters,
        constraint linearization parameters, and scaling matrices.

        Args:
            N: Number of discretization nodes
            x_unified: Unified state interface with dimensions and scaling bounds
            u_unified: Unified control interface with dimensions and scaling bounds
            jax_constraints: Lowered JAX constraints (for sizing linearization params)
            dynamics_sparsity: Optional tuple ``(A_d, B_d, C_d)`` of boolean
                ndarrays giving the discrete-time Jacobian sparsity patterns.
                ``A_d`` has shape ``(n_x, n_x)``; ``B_d`` and ``C_d`` have
                shape ``(n_x, n_u)``.
            constraint_sparsity: Optional list of ``(x_mask, u_mask)`` boolean
                1-D arrays, one per nodal constraint.
        """
        from openscvx.symbolic.lower import _tile_sparsity, create_cvxpy_variables

        n_states = len(x_unified.max)
        n_controls = len(u_unified.max)
        slice_cont = u_unified.slice_continuous
        slice_imp = u_unified.slice_impulsive
        n_controls_cont = int(slice_cont.stop - slice_cont.start)
        n_controls_imp = int(slice_imp.stop - slice_imp.start)
        if n_controls_cont + n_controls_imp != n_controls:
            raise ValueError(
                "Unified control slices are inconsistent with control dimension. "
                f"continuous={n_controls_cont}, impulsive={n_controls_imp}, total={n_controls}."
            )
        self._slice_cont = slice_cont
        self._slice_imp = slice_imp

        S_x, c_x = self._scaling(x_unified)
        S_u, c_u = self._scaling(u_unified)

        # Convert boolean sparsity patterns to CVXPY index format
        A_d_sp = B_d_sp = C_d_sp = None
        if dynamics_sparsity is not None:
            A_d_pat, B_d_pat, C_d_pat = dynamics_sparsity
            A_d_sp = _tile_sparsity(A_d_pat, N - 1)
            B_d_sp = _tile_sparsity(B_d_pat, N - 1)
            C_d_sp = _tile_sparsity(C_d_pat, N - 1)

        # TODO: (griffin-norris) Remove once cvxpygen supports N-D sparsity
        # indices. cvxpygen's handle_sparsity() assumes 2-D (rows, cols) but
        # our tiled parameters produce 3-D indices (slices, rows, cols).
        # Dropping sparsity here is safe — it only affects codegen performance.
        if self.cvxpygen:
            A_d_sp = B_d_sp = C_d_sp = None
            constraint_sparsity = None

        # Create all CVXPy variables for the OCP
        self._ocp_vars = create_cvxpy_variables(
            N=N,
            n_states=n_states,
            n_controls=n_controls,
            S_x=S_x,
            c_x=c_x,
            S_u=S_u,
            c_u=c_u,
            n_nodal_constraints=len(jax_constraints.nodal),
            n_cross_node_constraints=len(jax_constraints.cross_node),
            A_d_sparsity=A_d_sp,
            B_d_sparsity=B_d_sp,
            C_d_sparsity=C_d_sp,
            constraint_sparsity=constraint_sparsity,
        )

        # Stash the constraint catalog so iteration_callback can unstack the
        # JAX-pure ``SubproblemData`` back into the per-constraint dict layout
        # ``update_constraint_linearizations`` expects.
        self._jax_constraints = jax_constraints

    def lower_convex_constraints(self, constraints, parameters=None):
        """Lower user ``.convex()`` constraints into CVXPy constraint objects.

        Delegates to :func:`openscvx.symbolic.lower.lower_cvxpy_constraints`,
        feeding it the unscaled-state and unscaled-control CVXPy expressions
        built by :meth:`create_variables`.
        """
        from openscvx.symbolic.lower import lower_cvxpy_constraints

        if self._ocp_vars is None:
            raise RuntimeError(
                "CVXPyPTRSolver.lower_convex_constraints() called before "
                "create_variables(); the CVXPy variables it needs don't "
                "exist yet."
            )
        return lower_cvxpy_constraints(
            constraints,
            self._ocp_vars.x_nonscaled,
            self._ocp_vars.u_nonscaled,
            parameters,
        )

    def initialize(
        self,
        lowered: "LoweredProblem",
        settings: "Config",
    ) -> None:
        """Build the CVXPy optimal control problem.

        Constructs the complete optimization problem by calling ``cost()`` and
        ``constraints()`` to build the objective and constraint formulations,
        then assembles them into a CVXPy Problem.

        If cvxpygen is enabled, generates compiled solver code for improved
        performance.

        Note:
            ``create_variables()`` must be called before this method.

        Args:
            lowered: Lowered problem containing:
                - ``cvxpy_constraints``: Lowered convex constraints
                - ``jax_constraints``: JAX constraint functions (for structure)
            settings: Problem configuration (node count, scaling, etc.)

        Raises:
            RuntimeError: If create_variables() has not been called.
        """
        if self._ocp_vars is None:
            raise RuntimeError(
                "CVXPyPTRSolver.initialize() called before create_variables(). "
                "Call create_variables() first to create optimization variables."
            )

        objective = self.cost(settings, lowered)
        constr = self.constraints(settings, lowered)
        prob = cp.Problem(cp.Minimize(objective), constr)

        if self.cvxpygen:
            if not CVXPYGEN_AVAILABLE:
                raise ImportError(
                    "cvxpygen is required for code generation but not installed. "
                    "Install it with: pip install openscvx[cvxpygen] or pip install cvxpygen"
                )
            # Check to see if solver directory exists
            if not os.path.exists("solver"):
                cpg.generate_code(prob, solver=self.cvx_solver, code_dir="solver", wrapper=True)
            else:
                # Prompt the use to indicate if they wish to overwrite the solver
                # directory or use the existing compiled solver
                if self.cvxpygen_override:
                    cpg.generate_code(
                        prob,
                        solver=self.cvx_solver,
                        code_dir="solver",
                        wrapper=True,
                    )
                else:
                    overwrite = input("Solver directory already exists. Overwrite? (y/n): ")
                    if overwrite.lower() == "y":
                        cpg.generate_code(
                            prob,
                            solver=self.cvx_solver,
                            code_dir="solver",
                            wrapper=True,
                        )

        self._problem = prob
        self._settings = settings
        self._setup_solve_function()

    def cost(
        self,
        settings: "Config",
        lowered: "LoweredProblem",
    ) -> cp.Expression:
        """Build the cost expression for the convex subproblem.

        Constructs the PTR objective function including:

        - Boundary condition costs (Minimize/Maximize state components)
        - Trust region penalty (deviation from linearization point)
        - Virtual control penalty (dynamics defect relaxation)
        - Virtual buffer penalty (nonconvex constraint violation relaxation)

        Override this method in subclasses to customize the cost formulation.
        Use ``super().cost(settings, lowered)`` to include the standard PTR
        cost terms and add to them.

        Args:
            settings: Configuration object with solver settings
            lowered: Lowered problem containing constraint structure

        Returns:
            CVXPy expression representing the total cost to minimize.
        """
        ocp_vars = self._ocp_vars
        jax_constraints = lowered.jax_constraints

        lam_prox = ocp_vars.lam_prox
        prox_c = ocp_vars.prox_c
        prox_cc = ocp_vars.prox_cc
        lam_cost = ocp_vars.lam_cost
        lam_vc = ocp_vars.lam_vc
        lam_vb_nodal = ocp_vars.lam_vb_nodal
        lam_vb_cross = ocp_vars.lam_vb_cross
        x = ocp_vars.x
        u = ocp_vars.u
        nu = ocp_vars.nu
        nu_vb = ocp_vars.nu_vb
        nu_vb_cross = ocp_vars.nu_vb_cross

        cost = cp.sum(lam_cost) * 0
        cost += cp.sum(lam_vb_nodal) * 0
        cost += cp.sum(lam_vb_cross) * 0

        # Boundary condition cost terms (use scaled x for numerical conditioning)
        for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
            if settings.sim.x.initial_type[i] == "Minimize":
                cost += lam_cost[i] * x[0][i]
            if settings.sim.x.final_type[i] == "Minimize":
                cost += lam_cost[i] * x[-1][i]
            if settings.sim.x.initial_type[i] == "Maximize":
                cost -= lam_cost[i] * x[0][i]
            if settings.sim.x.final_type[i] == "Maximize":
                cost -= lam_cost[i] * x[-1][i]

        # Trust-region cost in expanded form:
        #   sum_i [lam_i * z_i^2 + prox_c_i * z_i + prox_cc_i], z_i = [x_i, u_i].
        cost += sum(
            cp.sum(cp.multiply(lam_prox[i], cp.square(cp.hstack((x[i], u[i])))))
            + cp.sum(cp.multiply(prox_c[i], cp.hstack((x[i], u[i]))))
            + prox_cc[i]
            for i in range(settings.sim.n)
        )

        # Virtual Control Slack
        cost += sum(cp.sum(lam_vc[i - 1] * cp.abs(nu[i - 1])) for i in range(1, settings.sim.n))

        # Virtual buffer penalty for nodal constraints (per-node weighting).
        # Equalities penalize |nu_vb|, inequalities penalize pos(nu_vb).
        idx_ncvx = 0
        if jax_constraints.nodal:
            for constraint in jax_constraints.nodal:
                pen = cp.abs if constraint.is_equality else cp.pos
                cost += lam_vb_nodal[:, idx_ncvx] @ pen(nu_vb[idx_ncvx])
                idx_ncvx += 1

        # Virtual slack penalty for cross-node constraints
        idx_cross = 0
        if jax_constraints.cross_node:
            for constraint in jax_constraints.cross_node:
                pen = cp.abs if constraint.is_equality else cp.pos
                cost += lam_vb_cross[idx_cross] * pen(nu_vb_cross[idx_cross])
                idx_cross += 1

        return cost

    def constraints(
        self,
        settings: "Config",
        lowered: "LoweredProblem",
    ) -> list:
        """Build the constraint list for the convex subproblem.

        Constructs all PTR constraints including:

        - Linearized nodal constraints (from JAX-lowered nonconvex constraints)
        - Linearized cross-node constraints
        - Convex constraints (already lowered to CVXPy)
        - Boundary conditions (fixed initial/terminal states)
        - Uniform time grid constraints
        - State and control deviation definitions
        - Linearized dynamics
        - State and control box constraints
        - CTCS constraints

        Override this method in subclasses to customize the constraint
        formulation. Use ``super().constraints(settings, lowered)`` to include
        the standard PTR constraints and extend them.

        Args:
            settings: Configuration object with solver settings
            lowered: Lowered problem containing lowered constraints

        Returns:
            List of CVXPy constraints.
        """
        ocp_vars = self._ocp_vars
        jax_constraints = lowered.jax_constraints
        cvxpy_constraints = lowered.cvxpy_constraints

        x_init = ocp_vars.x_init
        x_term = ocp_vars.x_term
        A_d = ocp_vars.A_d
        B_d = ocp_vars.B_d
        C_d = ocp_vars.C_d
        E_d = ocp_vars.E_d
        dyn_bias = ocp_vars.dyn_bias
        x0_imp_bias = ocp_vars.x0_imp_bias
        nu = ocp_vars.nu
        g = ocp_vars.g
        grad_g_x = ocp_vars.grad_g_x
        grad_g_u = ocp_vars.grad_g_u
        nu_vb = ocp_vars.nu_vb
        g_cross = ocp_vars.g_cross
        grad_g_X_cross = ocp_vars.grad_g_X_cross
        grad_g_U_cross = ocp_vars.grad_g_U_cross
        nu_vb_cross = ocp_vars.nu_vb_cross
        inv_S_x = ocp_vars.inv_S_x
        c_x = ocp_vars.c_x
        inv_S_u = ocp_vars.inv_S_u
        c_u = ocp_vars.c_u
        x_nonscaled = ocp_vars.x_nonscaled
        u_nonscaled = ocp_vars.u_nonscaled
        slice_cont = settings.sim.u.slice_continuous
        slice_imp = settings.sim.u.slice_impulsive
        has_impulsive = bool(slice_imp.stop > slice_imp.start)

        constr = []

        # Linearized nodal constraints (from JAX-lowered non-convex)
        idx_ncvx = 0
        if jax_constraints.nodal:
            for constraint in jax_constraints.nodal:
                # nodes should already be validated and normalized in preprocessing
                nodes = constraint.nodes
                for node in nodes:
                    residual = (
                        g[idx_ncvx][node]
                        + grad_g_x[idx_ncvx][node] @ x_nonscaled[node]
                        + grad_g_u[idx_ncvx][node] @ u_nonscaled[node]
                    )
                    constr += [residual == nu_vb[idx_ncvx][node]]
                idx_ncvx += 1

        # Linearized cross-node constraints (from JAX-lowered non-convex)
        idx_cross = 0
        if jax_constraints.cross_node:
            for constraint in jax_constraints.cross_node:
                # Linearization in affine form:
                # g_tilde + Σ_k(∇g_X[k]·X[k] + ∇g_U[k]·U[k]) == nu_vb
                # Sum over all trajectory nodes to couple multiple nodes
                residual = g_cross[idx_cross]
                for k in range(settings.sim.n):
                    # Contribution from state at node k
                    residual += grad_g_X_cross[idx_cross][k, :] @ x_nonscaled[k]
                    # Contribution from control at node k
                    residual += grad_g_U_cross[idx_cross][k, :] @ u_nonscaled[k]
                # Add constraint: residual == slack variable
                constr += [residual == nu_vb_cross[idx_cross]]
                idx_cross += 1

        # Convex constraints (already lowered to CVXPy)
        if cvxpy_constraints.constraints:
            constr += cvxpy_constraints.constraints

        # Boundary conditions (Fix)
        for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
            if settings.sim.x.initial_type[i] == "Fix":
                if has_impulsive:
                    constr += [
                        x_nonscaled[0][i]
                        == x0_imp_bias[i] + E_d[0][i, slice_imp] @ u_nonscaled[0][slice_imp]
                    ]
                else:
                    constr += [x_nonscaled[0][i] == x_init[i]]  # Initial Boundary Conditions
            if settings.sim.x.final_type[i] == "Fix":
                constr += [x_nonscaled[-1][i] == x_term[i]]  # Final Boundary Conditions

        if settings.sim._uniform_time_grid:
            S_u_inv_td = inv_S_u[settings.sim.time_dilation_slice, settings.sim.time_dilation_slice]
            c_u_td = c_u[settings.sim.time_dilation_slice]
            constr += [
                S_u_inv_td @ (u_nonscaled[i][settings.sim.time_dilation_slice] - c_u_td)
                == S_u_inv_td @ (u_nonscaled[i - 1][settings.sim.time_dilation_slice] - c_u_td)
                for i in range(1, settings.sim.n)
            ]

        constr += [
            inv_S_x @ (x_nonscaled[i] - c_x)
            == inv_S_x
            @ (
                A_d[i - 1] @ x_nonscaled[i - 1]
                + B_d[i - 1][:, slice_cont] @ u_nonscaled[i - 1][slice_cont]
                + C_d[i - 1][:, slice_cont] @ u_nonscaled[i][slice_cont]
                + (E_d[i][:, slice_imp] @ u_nonscaled[i][slice_imp] if has_impulsive else 0)
                + dyn_bias[i - 1]
                - c_x
            )
            + nu[i - 1]
            for i in range(1, settings.sim.n)
        ]  # Dynamics Constraint

        constr += [
            inv_S_u @ (u_nonscaled[i] - c_u) <= inv_S_u @ (settings.sim.u.max - c_u)
            for i in range(settings.sim.n)
        ]
        constr += [
            inv_S_u @ (u_nonscaled[i] - c_u) >= inv_S_u @ (settings.sim.u.min - c_u)
            for i in range(settings.sim.n)
        ]  # Control Constraints

        # TODO: (norrisg) formalize this
        constr += [
            inv_S_x @ (x_nonscaled[i][:] - c_x) <= inv_S_x @ (settings.sim.x.max - c_x)
            for i in range(settings.sim.n)
        ]
        constr += [
            inv_S_x @ (x_nonscaled[i][:] - c_x) >= inv_S_x @ (settings.sim.x.min - c_x)
            for i in range(settings.sim.n)
        ]  # State Constraints (Also implemented in CTCS but included for numerical stability)

        for idx, nodes in zip(
            np.arange(settings.sim.ctcs_slice.start, settings.sim.ctcs_slice.stop),
            settings.sim.ctcs_node_intervals,
        ):
            start_idx = 1 if nodes[0] == 0 else nodes[0]
            constr += [
                cp.abs(x_nonscaled[i][idx] - x_nonscaled[i - 1][idx]) <= settings.sim.x.max[idx]
                for i in range(start_idx, nodes[1])
            ]
            constr += [x_nonscaled[0][idx] == 0]

        return constr

    def _setup_solve_function(self) -> None:
        """Configure the solve function based on solver settings.

        Sets up either cvxpygen-based solving or standard CVXPy solving
        based on the configuration.
        """
        if self.cvxpygen:
            cpg_solver_path = os.path.join("solver", "cpg_solver.py")
            if not os.path.isfile(cpg_solver_path):
                raise ImportError(
                    "cvxpygen solver not found. Make sure cvxpygen is installed and code "
                    "generation has been run. Install with: uv pip install openscvx[cvxpygen]"
                )
            try:
                from solver.cpg_solver import cpg_solve
            except ImportError as exc:
                raise ImportError(
                    "cvxpygen solver not found. Make sure cvxpygen is installed and code "
                    "generation has been run. Install with: uv pip install openscvx[cvxpygen]"
                ) from exc
            # cvxpygen v1.0+ registers CPG during generate_code(); re-register here so
            # reusing an existing solver/ directory (without regenerating) still works.
            self._problem.register_solve("CPG", cpg_solve)
            solver_args = self.solver_args
            self._solve_fn = lambda: self._problem.solve(method="CPG", **solver_args)
        else:
            solver = self.cvx_solver
            solver_args = dict(self.solver_args)

            def _solve_with_dpp_fallback():
                try:
                    return self._problem.solve(solver=solver, **solver_args)
                except cp.error.DPPError:
                    fallback_args = dict(solver_args)
                    fallback_args.pop("enforce_dpp", None)
                    fallback_args["ignore_dpp"] = True
                    return self._problem.solve(solver=solver, **fallback_args)

            self._solve_fn = _solve_with_dpp_fallback

    def update_dynamics_linearization(
        self,
        x_bar: np.ndarray,
        u_bar: np.ndarray,
        A_d: np.ndarray,
        B_d: np.ndarray,
        C_d: np.ndarray,
        x_prop: np.ndarray,
        x_prop_plus: np.ndarray | None = None,
        D_d: np.ndarray | None = None,
        E_d: np.ndarray | None = None,
    ) -> None:
        """Update dynamics linearization point and matrices.

        Sets the current linearization point (previous iterate) and the
        discretized dynamics matrices for the convex subproblem.

        Args:
            x_bar: Previous state trajectory, shape (N, n_states)
            u_bar: Previous control trajectory, shape (N, n_controls)
            A_d: Discretized state Jacobian, shape (N-1, n_states, n_states)
            B_d: Discretized control Jacobian (current node), shape (N-1, n_states, n_controls)
            C_d: Discretized control Jacobian (next node), shape (N-1, n_states, n_controls)
            x_prop: Propagated state from continuous dynamics, shape (N-1, n_states)
            x_prop_plus: Optional impulsive/discrete propagated state, shape (N, n_states)
            D_d: Optional impulsive/discrete Jacobian wrt state, shape (N, n_states, n_states)
            E_d: Optional impulsive/discrete Jacobian wrt control, shape (N, n_states, n_controls)
        """
        x_bar_arr = np.asarray(x_bar)
        u_bar_arr = np.asarray(u_bar)
        self._ocp_vars.x_bar.value = x_bar_arr
        self._ocp_vars.u_bar.value = u_bar_arr

        A_eff = np.asarray(A_d)
        B_eff = np.asarray(B_d)
        C_eff = np.asarray(C_d)

        if D_d is not None:
            D_arr = np.asarray(D_d)
            # Temporary DPP-safe workaround: absorb D_d into A/B/C numerically
            # so the CVXPY graph has only Parameter@Variable products.
            if D_arr.ndim == 3 and D_arr.shape[0] == A_eff.shape[0] + 1:
                D_steps = D_arr[1:]
            elif D_arr.ndim == 3 and D_arr.shape[0] == A_eff.shape[0]:
                D_steps = D_arr
            else:
                raise ValueError(
                    "Unexpected D_d shape for dynamics update: "
                    f"{D_arr.shape}, expected "
                    f"{(A_eff.shape[0] + 1, A_eff.shape[1], A_eff.shape[2])} "
                    f"or {(A_eff.shape[0], A_eff.shape[1], A_eff.shape[2])}."
                )

            A_eff = np.einsum("kij,kjl->kil", D_steps, A_eff)
            B_eff = np.einsum("kij,kjl->kil", D_steps, B_eff)
            C_eff = np.einsum("kij,kjl->kil", D_steps, C_eff)

        self._set_param("A_d", A_eff)
        self._set_param("B_d", B_eff)
        self._set_param("C_d", C_eff)
        x_prop_arr = np.asarray(x_prop)
        x_prop_plus_arr = np.asarray(x_prop_plus) if x_prop_plus is not None else None
        E_arr = None
        if E_d is not None:
            E_arr = np.asarray(E_d)
            self._ocp_vars.E_d.value = E_arr

        self.dynamics_biases(
            A_d=A_eff,
            B_d=B_eff,
            C_d=C_eff,
            x_bar=x_bar_arr,
            u_bar=u_bar_arr,
            x_prop=x_prop_arr,
            x_prop_plus=x_prop_plus_arr,
            E_d=E_arr,
        )

    def dynamics_biases(
        self,
        A_d: np.ndarray,
        B_d: np.ndarray,
        C_d: np.ndarray,
        x_bar: np.ndarray,
        u_bar: np.ndarray,
        x_prop: np.ndarray,
        x_prop_plus: np.ndarray | None = None,
        E_d: np.ndarray | None = None,
    ) -> None:
        """Update affine dynamics bias parameters for the current linearization.

        Computes and stores:
            - ``dyn_bias``: per-step affine bias in linearized dynamics
            - ``x0_imp_bias``: initial-node impulsive affine bias
        """
        has_impulsive = bool(self._slice_imp.stop > self._slice_imp.start)

        dyn_bias = np.zeros((A_d.shape[0], A_d.shape[1]), dtype=float)
        for k in range(A_d.shape[0]):
            i = k + 1
            base = x_prop_plus[i] if has_impulsive else x_prop[k]
            bias_k = (
                base
                - A_d[k] @ x_bar[k]
                - B_d[k][:, self._slice_cont] @ u_bar[k][self._slice_cont]
                - C_d[k][:, self._slice_cont] @ u_bar[i][self._slice_cont]
                - (E_d[i][:, self._slice_imp] @ u_bar[i][self._slice_imp] if has_impulsive else 0)
            )
            dyn_bias[k] = bias_k
        self._ocp_vars.dyn_bias.value = dyn_bias

        if has_impulsive:
            x0_imp_bias = x_prop_plus[0] - E_d[0][:, self._slice_imp] @ u_bar[0][self._slice_imp]
        else:
            x0_imp_bias = np.zeros(A_d.shape[1], dtype=float)
        self._ocp_vars.x0_imp_bias.value = x0_imp_bias

    def update_constraint_linearizations(
        self,
        nodal: List[dict] = None,
        cross_node: List[dict] = None,
    ) -> None:
        """Update linearized constraint values and gradients.

        Sets constraint function values and gradients at the current
        linearization point for both nodal and cross-node constraints.

        Args:
            nodal: List of dicts for nodal constraints, each containing:
                - ``g``: Constraint value at linearization point
                - ``grad_g_x``: Gradient w.r.t. state
                - ``grad_g_u``: Gradient w.r.t. control
            cross_node: List of dicts for cross-node constraints, each containing:
                - ``g``: Constraint value at linearization point
                - ``grad_g_X``: Gradient w.r.t. full state trajectory
                - ``grad_g_U``: Gradient w.r.t. full control trajectory
        """
        if nodal:
            x_bar = self._ocp_vars.x_bar.value
            u_bar = self._ocp_vars.u_bar.value
            x_bar_arr = np.asarray(x_bar)
            u_bar_arr = np.asarray(u_bar)
            for g_id, constraint_data in enumerate(nodal):
                g_arr = np.asarray(constraint_data["g"])
                grad_x_arr = np.asarray(constraint_data["grad_g_x"])
                grad_u_arr = np.asarray(constraint_data["grad_g_u"])

                # Convert to affine form in (X, U):
                # g_tilde + grad_x·X + grad_u·U, where
                # g_tilde = g(X_bar, U_bar) - grad_x·X_bar - grad_u·U_bar.
                g_tilde = (
                    g_arr
                    - np.sum(grad_x_arr * x_bar_arr, axis=1)
                    - np.sum(grad_u_arr * u_bar_arr, axis=1)
                )

                self._set_param(f"g_{g_id}", g_tilde)
                self._set_param(f"grad_g_x_{g_id}", grad_x_arr)
                self._set_param(f"grad_g_u_{g_id}", grad_u_arr)

        if cross_node:
            x_bar = self._ocp_vars.x_bar.value
            u_bar = self._ocp_vars.u_bar.value
            x_bar_arr = np.asarray(x_bar)
            u_bar_arr = np.asarray(u_bar)
            for g_id, constraint_data in enumerate(cross_node):
                g_val = np.asarray(constraint_data["g"])
                grad_X_arr = np.asarray(constraint_data["grad_g_X"])
                grad_U_arr = np.asarray(constraint_data["grad_g_U"])

                g_tilde = g_val - np.sum(grad_X_arr * x_bar_arr) - np.sum(grad_U_arr * u_bar_arr)

                self._set_param(f"g_cross_{g_id}", g_tilde)
                self._set_param(f"grad_g_X_cross_{g_id}", grad_X_arr)
                self._set_param(f"grad_g_U_cross_{g_id}", grad_U_arr)

    def update_penalties(
        self,
        lam_prox: np.ndarray,
        lam_cost: Union[float, np.ndarray],
        lam_vc: np.ndarray,
        lam_vb_nodal: np.ndarray,
        lam_vb_cross: np.ndarray,
    ) -> None:
        """Update SCP penalty weights.

        Sets the penalty weights that balance competing objectives in the
        PTR convex subproblem.

        Args:
            lam_prox: Trust region weights, shape ``(N, n_states + n_controls)``.
            lam_cost: Cost function weight. Scalar or array of shape
                ``(n_states,)`` for per-state weighting.
            lam_vc: Virtual control penalty weights, shape (N-1, n_states)
            lam_vb_nodal: Virtual buffer penalty weights for nodal constraints,
                shape ``(N, n_nodal_constraints)``.
            lam_vb_cross: Virtual buffer penalty weights for cross-node
                constraints, shape ``(n_cross_node_constraints,)``.
        """
        lam_prox_arr = np.asarray(lam_prox)
        self._set_param("lam_prox", lam_prox_arr)
        self._set_param("lam_cost", lam_cost)
        self._set_param("lam_vc", lam_vc)
        self._set_param("lam_vb_nodal", lam_vb_nodal)
        self._set_param("lam_vb_cross", lam_vb_cross)

    def proximal_cost_terms(
        self,
        lam_prox: np.ndarray,
        x_bar: np.ndarray,
        u_bar: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        """Compute linear and constant proximal cost terms."""
        x_ref = (self._ocp_vars.inv_S_x @ (x_bar.T - self._ocp_vars.c_x[:, None])).T
        u_ref = (self._ocp_vars.inv_S_u @ (u_bar.T - self._ocp_vars.c_u[:, None])).T
        z_ref = np.hstack((x_ref, u_ref))
        prox_c = -2.0 * lam_prox * z_ref
        prox_cc = np.sum(lam_prox * np.square(z_ref), axis=1)
        return prox_c, prox_cc

    def update_proximal_terms(self) -> None:
        """Update proximal expansion parameters from current references and weights."""
        lam_prox_arr = np.asarray(self._ocp_vars.lam_prox.value)
        x_bar = np.asarray(self._ocp_vars.x_bar.value)
        u_bar = np.asarray(self._ocp_vars.u_bar.value)
        prox_c, prox_cc = self.proximal_cost_terms(lam_prox_arr, x_bar, u_bar)
        self._set_param("prox_c", prox_c)
        self._set_param("prox_cc", prox_cc)

    def update_boundary_conditions(
        self,
        x_init: np.ndarray = None,
        x_term: np.ndarray = None,
    ) -> None:
        """Update boundary condition parameters.

        Sets initial and/or terminal state constraints. Only sets parameters
        that exist in the problem (some problems may not have both).

        Args:
            x_init: Initial state vector, shape (n_states,). Optional.
            x_term: Terminal state vector, shape (n_states,). Optional.
        """
        # No-op before initialize() — the CVXPy problem (and its param_dict)
        # isn't built yet. Callers like Problem._sync_boundary_conditions
        # may invoke this both before and after initialize().
        if self._problem is None:
            return
        if x_init is not None and "x_init" in self._problem.param_dict:
            self._set_param("x_init", x_init)
        if x_term is not None and "x_term" in self._problem.param_dict:
            self._set_param("x_term", x_term)

    def get_stats(self) -> dict:
        """Get solver statistics for diagnostics and printing.

        Returns:
            Dict containing:
                - ``n_variables``: Total number of optimization variables
                - ``n_parameters``: Total number of parameters
                - ``n_constraints``: Total number of constraints
        """
        if self._problem is None:
            return {"n_variables": 0, "n_parameters": 0, "n_constraints": 0}

        return {
            "n_variables": sum(var.size for var in self._problem.variables()),
            "n_parameters": sum(param.size for param in self._problem.parameters()),
            "n_constraints": sum(constraint.size for constraint in self._problem.constraints),
        }

    def _set_param(self, name: str, value: np.ndarray) -> None:
        """Set a CVXPy parameter with helpful error messages on failure.

        Args:
            name: The parameter name in problem.param_dict
            value: The value to assign

        Raises:
            ValueError: If the value is not real, with diagnostic information.
        """
        try:
            param = self._problem.param_dict[name]
            value_arr = np.asarray(value)

            # Ensure the value shape matches the parameter shape exactly
            # This is critical for Python 3.11+ where NumPy/CVXPy are stricter about shapes
            if hasattr(param, "shape") and param.shape is not None:
                expected_shape = param.shape
                if value_arr.shape != expected_shape:
                    # Try to reshape if sizes match
                    if value_arr.size == np.prod(expected_shape):
                        value_arr = value_arr.reshape(expected_shape)
                    else:
                        # If sizes don't match, try squeezing extra dimensions first
                        value_arr = np.squeeze(value_arr)
                        if value_arr.shape != expected_shape and value_arr.size == np.prod(
                            expected_shape
                        ):
                            value_arr = value_arr.reshape(expected_shape)
                        elif value_arr.shape != expected_shape:
                            raise ValueError(
                                f"Parameter '{name}' shape mismatch: expected {expected_shape}, "
                                f"got {value.shape} (after squeezing: {value_arr.shape})"
                            )

            param.value = value_arr
        except ValueError as e:
            if "must be real" in str(e):
                arr = np.asarray(value)
                nan_mask = ~np.isfinite(arr)
                nan_indices = np.argwhere(nan_mask)

                index_value_strs = [
                    f"  {tuple(int(i) for i in idx)} -> {arr[tuple(idx)]}"
                    for idx in nan_indices[:20]
                ]
                if len(nan_indices) > 20:
                    index_value_strs.append(f"  ... and {len(nan_indices) - 20} more")

                arr_str = np.array2string(arr, threshold=200, edgeitems=3, max_line_width=120)
                msg = (
                    f"Parameter '{name}' with shape {arr.shape} contains "
                    f"{len(nan_indices)} non-real value(s):\n"
                    + "\n".join(index_value_strs)
                    + f"\n\n{name} = {arr_str}"
                )
                raise ValueError(msg) from e
            raise

    def solve(self) -> PTRSolveResult:
        """Solve the convex subproblem and return structured results.

        Call ``update_dynamics_linearization()``, ``update_constraint_linearizations()``,
        and ``update_penalties()`` before calling this method.

        Returns:
            PTRSolveResult containing unscaled trajectories, slack variables,
            cost, and solver status.

        Raises:
            RuntimeError: If initialize() has not been called.
        """
        if self._problem is None:
            raise RuntimeError(
                "CVXPyPTRSolver.solve() called before initialize(). "
                "Call initialize() first to build the problem structure."
            )

        self._solve_fn()

        # Get scaling matrices
        S_x = self._ocp_vars.S_x
        c_x = self._ocp_vars.c_x
        S_u = self._ocp_vars.S_u
        c_u = self._ocp_vars.c_u

        # Unscale state and control trajectories
        x_scaled = self._problem.var_dict["x"].value  # (N, n_states)
        u_scaled = self._problem.var_dict["u"].value  # (N, n_controls)
        x = (S_x @ x_scaled.T + np.expand_dims(c_x, axis=1)).T
        u = (S_u @ u_scaled.T + np.expand_dims(c_u, axis=1)).T

        # Get virtual control slack
        nu = self._problem.var_dict["nu"].value

        # Get nodal constraint violation slacks
        nu_vb = [var.value for var in self._ocp_vars.nu_vb]

        # Get cross-node constraint violation slacks
        nu_vb_cross = [var.value for var in self._ocp_vars.nu_vb_cross]

        return PTRSolveResult(
            x=x,
            u=u,
            nu=nu,
            nu_vb=nu_vb,
            nu_vb_cross=nu_vb_cross,
            cost=self._problem.value,
            status=self._problem.status,
        )

    def iteration_callback(self) -> Callable[..., SubproblemSolution]:
        """JAX-pure ``(state, SubproblemData) -> SubproblemSolution``.

        Wraps the existing NumPy ``solve()`` path in :func:`jax.pure_callback`
        with ``vmap_method="sequential"``. CVXPy can't trace — it builds a DCP
        graph and reaches into a backend solver through Python attribute
        mutation — so the only JAX-friendly surface it can present is a host
        callback. Under :func:`jax.jit` the callback fires once per call; under
        :func:`jax.vmap` it fires ``B`` times serially (CVXPy can't ingest a
        batched parameter set).

        The closure unstacks :class:`SubproblemData` back to the per-constraint
        dict layout :meth:`update_constraint_linearizations` expects, runs the
        four ``update_*`` methods + :meth:`solve`, and packs the result into a
        :class:`SubproblemSolution` whose ``nu_vb`` is collapsed to a stacked
        ``(N, n_nodal)`` array (the list-of-arrays form on
        :attr:`PTRSolveResult.nu_vb` is preserved for the NumPy path).

        ``status_code`` is mapped via :func:`status_str_to_code`; CVXPy's
        ``"optimal_inaccurate"`` / ``"solver_error"`` / ... labels collapse to
        :attr:`StatusCode.UNKNOWN` — the SCP trust-region check is the
        authoritative convergence gate.

        The returned callable takes ``(state, data)``: ``state`` is the
        :class:`AlgorithmState` pytree, accepted for cross-backend signature
        uniformity but unused (CVXPy exposes no warm-start through
        ``update_*``); ``data`` is the :class:`SubproblemData` pytree carrying
        the per-iteration linearization arrays, penalty weights, and boundary
        conditions.
        """
        if self._problem is None or self._settings is None or self._jax_constraints is None:
            raise RuntimeError(
                "CVXPyPTRSolver.iteration_callback() requires initialize() to "
                "have been called first."
            )

        settings = self._settings
        jax_constraints = self._jax_constraints
        n_x = settings.sim.n_states
        n_u = settings.sim.n_controls
        N = settings.sim.n
        n_nodal = len(jax_constraints.nodal)
        n_cross = len(jax_constraints.cross_node)
        slice_imp = settings.sim.u.slice_impulsive
        has_impulsive = bool(slice_imp.stop > slice_imp.start)
        f = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32

        result_struct = SubproblemSolution(
            x=jax.ShapeDtypeStruct((N, n_x), f),
            u=jax.ShapeDtypeStruct((N, n_u), f),
            nu=jax.ShapeDtypeStruct((N - 1, n_x), f),
            nu_vb=jax.ShapeDtypeStruct((N, n_nodal), f),
            nu_vb_cross=jax.ShapeDtypeStruct((n_cross,), f),
            cost=jax.ShapeDtypeStruct((), f),
            status_code=jax.ShapeDtypeStruct((), jnp.int32),
        )

        def host_solve(data: SubproblemData) -> SubproblemSolution:
            # SubproblemData always carries (N, n_x, n_x) D_d / (N, n_x, n_u) E_d
            # / (N, n_x) x_prop_plus arrays — zero-filled when the problem has
            # no impulsive component. ``update_dynamics_linearization`` treats
            # a non-None D_d as "absorb into A/B/C", which would zero them out
            # on the no-impulsive path; pass None in that case to keep A/B/C
            # intact. Same logic for x_prop_plus / E_d.
            x_prop_plus = np.asarray(data.x_prop_plus) if has_impulsive else None
            D_d = np.asarray(data.D_d) if has_impulsive else None
            E_d = np.asarray(data.E_d) if has_impulsive else None
            self.update_dynamics_linearization(
                x_bar=np.asarray(data.x_bar),
                u_bar=np.asarray(data.u_bar),
                A_d=np.asarray(data.A_d),
                B_d=np.asarray(data.B_d),
                C_d=np.asarray(data.C_d),
                x_prop=np.asarray(data.x_prop),
                x_prop_plus=x_prop_plus,
                D_d=D_d,
                E_d=E_d,
            )
            nodal = _unstack_nodal(data, jax_constraints)
            cross_node = _unstack_cross(data, jax_constraints)
            self.update_constraint_linearizations(
                nodal=nodal or None,
                cross_node=cross_node or None,
            )
            self.update_penalties(
                lam_prox=np.asarray(data.lam_prox),
                lam_cost=np.asarray(data.lam_cost),
                lam_vc=np.asarray(data.lam_vc),
                lam_vb_nodal=np.asarray(data.lam_vb_nodal),
                lam_vb_cross=np.asarray(data.lam_vb_cross),
            )
            # Set the proximal-term parameter (``prox_c``) from the freshly
            # updated penalties; the legacy ``_subproblem`` called this between
            # ``update_penalties`` and ``solve``, and without it ``prox_c`` has
            # no value on the first solve of a process.
            self.update_proximal_terms()
            self.update_boundary_conditions(
                x_init=_unmask_bc(data.x_init),
                x_term=_unmask_bc(data.x_term),
            )

            try:
                result = self.solve()
            except cp.error.SolverError:
                # Infeasible / numerical failure. Surface it as a non-OPTIMAL
                # status on a finite (zero) solution so the SCP loop fails
                # loudly at the step() boundary, instead of the exception
                # propagating through pure_callback as an opaque CpuCallback
                # error.
                return SubproblemSolution(
                    x=jnp.zeros((N, n_x), dtype=f),
                    u=jnp.zeros((N, n_u), dtype=f),
                    nu=jnp.zeros((N - 1, n_x), dtype=f),
                    nu_vb=jnp.zeros((N, n_nodal), dtype=f),
                    nu_vb_cross=jnp.zeros((n_cross,), dtype=f),
                    cost=jnp.asarray(0.0, dtype=f),
                    status_code=jnp.asarray(int(StatusCode.INFEASIBLE), dtype=jnp.int32),
                )

            # ``optimal_inaccurate`` is a usable solution — the legacy path
            # consumed it without error — so treat it as OPTIMAL; the step()
            # status gate would otherwise reject the UNKNOWN it maps to.
            code = (
                StatusCode.OPTIMAL
                if (result.status == "optimal_inaccurate")
                or (result.status == "solved")
                or (result.status == "1 (for description visit https://qoco-org.github.io/qoco/)")
                else status_str_to_code(result.status)
            )
            return SubproblemSolution(
                x=jnp.asarray(result.x, dtype=f),
                u=jnp.asarray(result.u, dtype=f),
                nu=jnp.asarray(result.nu, dtype=f),
                # nu_vb collapses from PTRSolveResult's list-of-(N,) layout to
                # the (N, n_nodal) stacked array SubproblemSolution declares
                # for pure_callback's ShapeDtypeStruct contract.
                nu_vb=(
                    jnp.zeros((N, 0), dtype=f)
                    if n_nodal == 0
                    else jnp.asarray(
                        np.stack([np.asarray(a) for a in result.nu_vb], axis=-1),
                        dtype=f,
                    )
                ),
                nu_vb_cross=jnp.asarray(
                    np.asarray(result.nu_vb_cross, dtype=float).reshape(n_cross),
                    dtype=f,
                ),
                cost=jnp.asarray(result.cost, dtype=f),
                status_code=jnp.asarray(int(code), dtype=jnp.int32),
            )

        def step(state, data: SubproblemData) -> SubproblemSolution:
            del state  # CVXPy exposes no warm-start through ``update_*``.
            return jax.pure_callback(
                host_solve,
                result_struct,
                data,
                vmap_method="sequential",
            )

        return step

    def citation(self) -> List[str]:
        """Return BibTeX citations for CVXPy.

        Returns:
            List containing BibTeX entries for CVXPy and DCCP papers.
        """
        return [
            r"""@article{diamond2016cvxpy,
  title={CVXPY: A Python-embedded modeling language for convex optimization},
  author={Diamond, Steven and Boyd, Stephen},
  journal={Journal of Machine Learning Research},
  volume={17},
  number={83},
  pages={1--5},
  year={2016}
}""",
            r"""@article{agrawal2018rewriting,
  title={A rewriting system for convex optimization problems},
  author={Agrawal, Akshay and Verschueren, Robin and Diamond, Steven and Boyd, Stephen},
  journal={Journal of Control and Decision},
  volume={5},
  number={1},
  pages={42--60},
  year={2018},
  publisher={Taylor \& Francis}
}""",
        ]

exportable: bool property

CVXPy is not jax.export-serializable — see :attr:PTRSolver.exportable.

The JAX-pure :meth:iteration_callback wraps the host CVXPy solve in a :func:jax.pure_callback, and jax.export cannot serialize host callbacks. solve_batched(save_compiled=True) therefore refuses this backend with a teaching error pointing at QPAX / Moreau rather than silently falling back to an uncached in-process solve.

ocp_vars: CVXPyVariables property

The CVXPy variables and parameters.

Returns:

Type	Description
CVXPyVariables	The CVXPyVariables dataclass, or None if create_variables() not called.

__init__(cvx_solver: str = 'QOCO', solver_args: Optional[dict] = None, cvxpygen: bool = False, cvxpygen_override: bool = False)

Initialize CVXPyPTRSolver with solver configuration.

Call create_variables() then initialize() to build the problem structure.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def __init__(
    self,
    cvx_solver: str = "QOCO",
    solver_args: Optional[dict] = None,
    cvxpygen: bool = False,
    cvxpygen_override: bool = False,
):
    """Initialize CVXPyPTRSolver with solver configuration.

    Call create_variables() then initialize() to build the problem structure.
    """
    self.cvx_solver = cvx_solver
    self.solver_args = (
        solver_args
        if solver_args is not None
        else {"abstol": 1e-06, "reltol": 1e-09, "enforce_dpp": True, "warm_start": False}
    )
    self.cvxpygen = cvxpygen
    self.cvxpygen_override = cvxpygen_override

    self._ocp_vars: "CVXPyVariables" = None
    self._problem: cp.Problem = None
    self._solve_fn: callable = None
    # Stashed at create_variables / initialize so iteration_callback can
    # close over them without retracing them off the ``lowered`` /
    # ``settings`` arguments each iteration.
    self._jax_constraints: Optional["LoweredJaxConstraints"] = None
    self._settings: Optional[Config] = None
    self._slice_cont = slice(0, 0)
    self._slice_imp = slice(0, 0)

citation() -> List[str]

Return BibTeX citations for CVXPy.

Returns:

Type	Description
List[str]	List containing BibTeX entries for CVXPy and DCCP papers.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

    def citation(self) -> List[str]:
        """Return BibTeX citations for CVXPy.

        Returns:
            List containing BibTeX entries for CVXPy and DCCP papers.
        """
        return [
            r"""@article{diamond2016cvxpy,
  title={CVXPY: A Python-embedded modeling language for convex optimization},
  author={Diamond, Steven and Boyd, Stephen},
  journal={Journal of Machine Learning Research},
  volume={17},
  number={83},
  pages={1--5},
  year={2016}
}""",
            r"""@article{agrawal2018rewriting,
  title={A rewriting system for convex optimization problems},
  author={Agrawal, Akshay and Verschueren, Robin and Diamond, Steven and Boyd, Stephen},
  journal={Journal of Control and Decision},
  volume={5},
  number={1},
  pages={42--60},
  year={2018},
  publisher={Taylor \& Francis}
}""",
        ]

constraints(settings: Config, lowered: LoweredProblem) -> list

Build the constraint list for the convex subproblem.

Constructs all PTR constraints including:

• Linearized nodal constraints (from JAX-lowered nonconvex constraints)
• Linearized cross-node constraints
• Convex constraints (already lowered to CVXPy)
• Boundary conditions (fixed initial/terminal states)
• Uniform time grid constraints
• State and control deviation definitions
• Linearized dynamics
• State and control box constraints
• CTCS constraints

Override this method in subclasses to customize the constraint formulation. Use super().constraints(settings, lowered) to include the standard PTR constraints and extend them.

Parameters:

Name	Type	Description	Default
settings	Config	Configuration object with solver settings	required
lowered	LoweredProblem	Lowered problem containing lowered constraints	required

Returns:

Type	Description
list	List of CVXPy constraints.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def constraints(
    self,
    settings: "Config",
    lowered: "LoweredProblem",
) -> list:
    """Build the constraint list for the convex subproblem.

    Constructs all PTR constraints including:

    - Linearized nodal constraints (from JAX-lowered nonconvex constraints)
    - Linearized cross-node constraints
    - Convex constraints (already lowered to CVXPy)
    - Boundary conditions (fixed initial/terminal states)
    - Uniform time grid constraints
    - State and control deviation definitions
    - Linearized dynamics
    - State and control box constraints
    - CTCS constraints

    Override this method in subclasses to customize the constraint
    formulation. Use ``super().constraints(settings, lowered)`` to include
    the standard PTR constraints and extend them.

    Args:
        settings: Configuration object with solver settings
        lowered: Lowered problem containing lowered constraints

    Returns:
        List of CVXPy constraints.
    """
    ocp_vars = self._ocp_vars
    jax_constraints = lowered.jax_constraints
    cvxpy_constraints = lowered.cvxpy_constraints

    x_init = ocp_vars.x_init
    x_term = ocp_vars.x_term
    A_d = ocp_vars.A_d
    B_d = ocp_vars.B_d
    C_d = ocp_vars.C_d
    E_d = ocp_vars.E_d
    dyn_bias = ocp_vars.dyn_bias
    x0_imp_bias = ocp_vars.x0_imp_bias
    nu = ocp_vars.nu
    g = ocp_vars.g
    grad_g_x = ocp_vars.grad_g_x
    grad_g_u = ocp_vars.grad_g_u
    nu_vb = ocp_vars.nu_vb
    g_cross = ocp_vars.g_cross
    grad_g_X_cross = ocp_vars.grad_g_X_cross
    grad_g_U_cross = ocp_vars.grad_g_U_cross
    nu_vb_cross = ocp_vars.nu_vb_cross
    inv_S_x = ocp_vars.inv_S_x
    c_x = ocp_vars.c_x
    inv_S_u = ocp_vars.inv_S_u
    c_u = ocp_vars.c_u
    x_nonscaled = ocp_vars.x_nonscaled
    u_nonscaled = ocp_vars.u_nonscaled
    slice_cont = settings.sim.u.slice_continuous
    slice_imp = settings.sim.u.slice_impulsive
    has_impulsive = bool(slice_imp.stop > slice_imp.start)

    constr = []

    # Linearized nodal constraints (from JAX-lowered non-convex)
    idx_ncvx = 0
    if jax_constraints.nodal:
        for constraint in jax_constraints.nodal:
            # nodes should already be validated and normalized in preprocessing
            nodes = constraint.nodes
            for node in nodes:
                residual = (
                    g[idx_ncvx][node]
                    + grad_g_x[idx_ncvx][node] @ x_nonscaled[node]
                    + grad_g_u[idx_ncvx][node] @ u_nonscaled[node]
                )
                constr += [residual == nu_vb[idx_ncvx][node]]
            idx_ncvx += 1

    # Linearized cross-node constraints (from JAX-lowered non-convex)
    idx_cross = 0
    if jax_constraints.cross_node:
        for constraint in jax_constraints.cross_node:
            # Linearization in affine form:
            # g_tilde + Σ_k(∇g_X[k]·X[k] + ∇g_U[k]·U[k]) == nu_vb
            # Sum over all trajectory nodes to couple multiple nodes
            residual = g_cross[idx_cross]
            for k in range(settings.sim.n):
                # Contribution from state at node k
                residual += grad_g_X_cross[idx_cross][k, :] @ x_nonscaled[k]
                # Contribution from control at node k
                residual += grad_g_U_cross[idx_cross][k, :] @ u_nonscaled[k]
            # Add constraint: residual == slack variable
            constr += [residual == nu_vb_cross[idx_cross]]
            idx_cross += 1

    # Convex constraints (already lowered to CVXPy)
    if cvxpy_constraints.constraints:
        constr += cvxpy_constraints.constraints

    # Boundary conditions (Fix)
    for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
        if settings.sim.x.initial_type[i] == "Fix":
            if has_impulsive:
                constr += [
                    x_nonscaled[0][i]
                    == x0_imp_bias[i] + E_d[0][i, slice_imp] @ u_nonscaled[0][slice_imp]
                ]
            else:
                constr += [x_nonscaled[0][i] == x_init[i]]  # Initial Boundary Conditions
        if settings.sim.x.final_type[i] == "Fix":
            constr += [x_nonscaled[-1][i] == x_term[i]]  # Final Boundary Conditions

    if settings.sim._uniform_time_grid:
        S_u_inv_td = inv_S_u[settings.sim.time_dilation_slice, settings.sim.time_dilation_slice]
        c_u_td = c_u[settings.sim.time_dilation_slice]
        constr += [
            S_u_inv_td @ (u_nonscaled[i][settings.sim.time_dilation_slice] - c_u_td)
            == S_u_inv_td @ (u_nonscaled[i - 1][settings.sim.time_dilation_slice] - c_u_td)
            for i in range(1, settings.sim.n)
        ]

    constr += [
        inv_S_x @ (x_nonscaled[i] - c_x)
        == inv_S_x
        @ (
            A_d[i - 1] @ x_nonscaled[i - 1]
            + B_d[i - 1][:, slice_cont] @ u_nonscaled[i - 1][slice_cont]
            + C_d[i - 1][:, slice_cont] @ u_nonscaled[i][slice_cont]
            + (E_d[i][:, slice_imp] @ u_nonscaled[i][slice_imp] if has_impulsive else 0)
            + dyn_bias[i - 1]
            - c_x
        )
        + nu[i - 1]
        for i in range(1, settings.sim.n)
    ]  # Dynamics Constraint

    constr += [
        inv_S_u @ (u_nonscaled[i] - c_u) <= inv_S_u @ (settings.sim.u.max - c_u)
        for i in range(settings.sim.n)
    ]
    constr += [
        inv_S_u @ (u_nonscaled[i] - c_u) >= inv_S_u @ (settings.sim.u.min - c_u)
        for i in range(settings.sim.n)
    ]  # Control Constraints

    # TODO: (norrisg) formalize this
    constr += [
        inv_S_x @ (x_nonscaled[i][:] - c_x) <= inv_S_x @ (settings.sim.x.max - c_x)
        for i in range(settings.sim.n)
    ]
    constr += [
        inv_S_x @ (x_nonscaled[i][:] - c_x) >= inv_S_x @ (settings.sim.x.min - c_x)
        for i in range(settings.sim.n)
    ]  # State Constraints (Also implemented in CTCS but included for numerical stability)

    for idx, nodes in zip(
        np.arange(settings.sim.ctcs_slice.start, settings.sim.ctcs_slice.stop),
        settings.sim.ctcs_node_intervals,
    ):
        start_idx = 1 if nodes[0] == 0 else nodes[0]
        constr += [
            cp.abs(x_nonscaled[i][idx] - x_nonscaled[i - 1][idx]) <= settings.sim.x.max[idx]
            for i in range(start_idx, nodes[1])
        ]
        constr += [x_nonscaled[0][idx] == 0]

    return constr

cost(settings: Config, lowered: LoweredProblem) -> cp.Expression

Build the cost expression for the convex subproblem.

Constructs the PTR objective function including:

• Boundary condition costs (Minimize/Maximize state components)
• Trust region penalty (deviation from linearization point)
• Virtual control penalty (dynamics defect relaxation)
• Virtual buffer penalty (nonconvex constraint violation relaxation)

Override this method in subclasses to customize the cost formulation. Use super().cost(settings, lowered) to include the standard PTR cost terms and add to them.

Parameters:

Name	Type	Description	Default
settings	Config	Configuration object with solver settings	required
lowered	LoweredProblem	Lowered problem containing constraint structure	required

Returns:

Type	Description
Expression	CVXPy expression representing the total cost to minimize.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def cost(
    self,
    settings: "Config",
    lowered: "LoweredProblem",
) -> cp.Expression:
    """Build the cost expression for the convex subproblem.

    Constructs the PTR objective function including:

    - Boundary condition costs (Minimize/Maximize state components)
    - Trust region penalty (deviation from linearization point)
    - Virtual control penalty (dynamics defect relaxation)
    - Virtual buffer penalty (nonconvex constraint violation relaxation)

    Override this method in subclasses to customize the cost formulation.
    Use ``super().cost(settings, lowered)`` to include the standard PTR
    cost terms and add to them.

    Args:
        settings: Configuration object with solver settings
        lowered: Lowered problem containing constraint structure

    Returns:
        CVXPy expression representing the total cost to minimize.
    """
    ocp_vars = self._ocp_vars
    jax_constraints = lowered.jax_constraints

    lam_prox = ocp_vars.lam_prox
    prox_c = ocp_vars.prox_c
    prox_cc = ocp_vars.prox_cc
    lam_cost = ocp_vars.lam_cost
    lam_vc = ocp_vars.lam_vc
    lam_vb_nodal = ocp_vars.lam_vb_nodal
    lam_vb_cross = ocp_vars.lam_vb_cross
    x = ocp_vars.x
    u = ocp_vars.u
    nu = ocp_vars.nu
    nu_vb = ocp_vars.nu_vb
    nu_vb_cross = ocp_vars.nu_vb_cross

    cost = cp.sum(lam_cost) * 0
    cost += cp.sum(lam_vb_nodal) * 0
    cost += cp.sum(lam_vb_cross) * 0

    # Boundary condition cost terms (use scaled x for numerical conditioning)
    for i in range(settings.sim.true_state_slice.start, settings.sim.true_state_slice.stop):
        if settings.sim.x.initial_type[i] == "Minimize":
            cost += lam_cost[i] * x[0][i]
        if settings.sim.x.final_type[i] == "Minimize":
            cost += lam_cost[i] * x[-1][i]
        if settings.sim.x.initial_type[i] == "Maximize":
            cost -= lam_cost[i] * x[0][i]
        if settings.sim.x.final_type[i] == "Maximize":
            cost -= lam_cost[i] * x[-1][i]

    # Trust-region cost in expanded form:
    #   sum_i [lam_i * z_i^2 + prox_c_i * z_i + prox_cc_i], z_i = [x_i, u_i].
    cost += sum(
        cp.sum(cp.multiply(lam_prox[i], cp.square(cp.hstack((x[i], u[i])))))
        + cp.sum(cp.multiply(prox_c[i], cp.hstack((x[i], u[i]))))
        + prox_cc[i]
        for i in range(settings.sim.n)
    )

    # Virtual Control Slack
    cost += sum(cp.sum(lam_vc[i - 1] * cp.abs(nu[i - 1])) for i in range(1, settings.sim.n))

    # Virtual buffer penalty for nodal constraints (per-node weighting).
    # Equalities penalize |nu_vb|, inequalities penalize pos(nu_vb).
    idx_ncvx = 0
    if jax_constraints.nodal:
        for constraint in jax_constraints.nodal:
            pen = cp.abs if constraint.is_equality else cp.pos
            cost += lam_vb_nodal[:, idx_ncvx] @ pen(nu_vb[idx_ncvx])
            idx_ncvx += 1

    # Virtual slack penalty for cross-node constraints
    idx_cross = 0
    if jax_constraints.cross_node:
        for constraint in jax_constraints.cross_node:
            pen = cp.abs if constraint.is_equality else cp.pos
            cost += lam_vb_cross[idx_cross] * pen(nu_vb_cross[idx_cross])
            idx_cross += 1

    return cost

create_variables(N: int, x_unified: UnifiedState, u_unified: UnifiedControl, jax_constraints: LoweredJaxConstraints, dynamics_sparsity: Optional[tuple] = None, constraint_sparsity: Optional[list] = None) -> None

Create CVXPy optimization variables.

Creates all CVXPy Variable and Parameter objects needed for the optimal control problem. This includes state/control variables, dynamics parameters, constraint linearization parameters, and scaling matrices.

Parameters:

Name	Type	Description	Default
N	int	Number of discretization nodes	required
x_unified	UnifiedState	Unified state interface with dimensions and scaling bounds	required
u_unified	UnifiedControl	Unified control interface with dimensions and scaling bounds	required
jax_constraints	LoweredJaxConstraints	Lowered JAX constraints (for sizing linearization params)	required
dynamics_sparsity	Optional[tuple]	Optional tuple (A_d, B_d, C_d) of boolean ndarrays giving the discrete-time Jacobian sparsity patterns. A_d has shape (n_x, n_x); B_d and C_d have shape (n_x, n_u).	None
constraint_sparsity	Optional[list]	Optional list of (x_mask, u_mask) boolean 1-D arrays, one per nodal constraint.	None

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def create_variables(
    self,
    N: int,
    x_unified: "UnifiedState",
    u_unified: "UnifiedControl",
    jax_constraints: "LoweredJaxConstraints",
    dynamics_sparsity: Optional[tuple] = None,
    constraint_sparsity: Optional[list] = None,
) -> None:
    """Create CVXPy optimization variables.

    Creates all CVXPy Variable and Parameter objects needed for the optimal
    control problem. This includes state/control variables, dynamics parameters,
    constraint linearization parameters, and scaling matrices.

    Args:
        N: Number of discretization nodes
        x_unified: Unified state interface with dimensions and scaling bounds
        u_unified: Unified control interface with dimensions and scaling bounds
        jax_constraints: Lowered JAX constraints (for sizing linearization params)
        dynamics_sparsity: Optional tuple ``(A_d, B_d, C_d)`` of boolean
            ndarrays giving the discrete-time Jacobian sparsity patterns.
            ``A_d`` has shape ``(n_x, n_x)``; ``B_d`` and ``C_d`` have
            shape ``(n_x, n_u)``.
        constraint_sparsity: Optional list of ``(x_mask, u_mask)`` boolean
            1-D arrays, one per nodal constraint.
    """
    from openscvx.symbolic.lower import _tile_sparsity, create_cvxpy_variables

    n_states = len(x_unified.max)
    n_controls = len(u_unified.max)
    slice_cont = u_unified.slice_continuous
    slice_imp = u_unified.slice_impulsive
    n_controls_cont = int(slice_cont.stop - slice_cont.start)
    n_controls_imp = int(slice_imp.stop - slice_imp.start)
    if n_controls_cont + n_controls_imp != n_controls:
        raise ValueError(
            "Unified control slices are inconsistent with control dimension. "
            f"continuous={n_controls_cont}, impulsive={n_controls_imp}, total={n_controls}."
        )
    self._slice_cont = slice_cont
    self._slice_imp = slice_imp

    S_x, c_x = self._scaling(x_unified)
    S_u, c_u = self._scaling(u_unified)

    # Convert boolean sparsity patterns to CVXPY index format
    A_d_sp = B_d_sp = C_d_sp = None
    if dynamics_sparsity is not None:
        A_d_pat, B_d_pat, C_d_pat = dynamics_sparsity
        A_d_sp = _tile_sparsity(A_d_pat, N - 1)
        B_d_sp = _tile_sparsity(B_d_pat, N - 1)
        C_d_sp = _tile_sparsity(C_d_pat, N - 1)

    # TODO: (griffin-norris) Remove once cvxpygen supports N-D sparsity
    # indices. cvxpygen's handle_sparsity() assumes 2-D (rows, cols) but
    # our tiled parameters produce 3-D indices (slices, rows, cols).
    # Dropping sparsity here is safe — it only affects codegen performance.
    if self.cvxpygen:
        A_d_sp = B_d_sp = C_d_sp = None
        constraint_sparsity = None

    # Create all CVXPy variables for the OCP
    self._ocp_vars = create_cvxpy_variables(
        N=N,
        n_states=n_states,
        n_controls=n_controls,
        S_x=S_x,
        c_x=c_x,
        S_u=S_u,
        c_u=c_u,
        n_nodal_constraints=len(jax_constraints.nodal),
        n_cross_node_constraints=len(jax_constraints.cross_node),
        A_d_sparsity=A_d_sp,
        B_d_sparsity=B_d_sp,
        C_d_sparsity=C_d_sp,
        constraint_sparsity=constraint_sparsity,
    )

    # Stash the constraint catalog so iteration_callback can unstack the
    # JAX-pure ``SubproblemData`` back into the per-constraint dict layout
    # ``update_constraint_linearizations`` expects.
    self._jax_constraints = jax_constraints

dynamics_biases(A_d: np.ndarray, B_d: np.ndarray, C_d: np.ndarray, x_bar: np.ndarray, u_bar: np.ndarray, x_prop: np.ndarray, x_prop_plus: np.ndarray | None = None, E_d: np.ndarray | None = None) -> None

Update affine dynamics bias parameters for the current linearization.

Computes and stores

• dyn_bias: per-step affine bias in linearized dynamics
• x0_imp_bias: initial-node impulsive affine bias

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def dynamics_biases(
    self,
    A_d: np.ndarray,
    B_d: np.ndarray,
    C_d: np.ndarray,
    x_bar: np.ndarray,
    u_bar: np.ndarray,
    x_prop: np.ndarray,
    x_prop_plus: np.ndarray | None = None,
    E_d: np.ndarray | None = None,
) -> None:
    """Update affine dynamics bias parameters for the current linearization.

    Computes and stores:
        - ``dyn_bias``: per-step affine bias in linearized dynamics
        - ``x0_imp_bias``: initial-node impulsive affine bias
    """
    has_impulsive = bool(self._slice_imp.stop > self._slice_imp.start)

    dyn_bias = np.zeros((A_d.shape[0], A_d.shape[1]), dtype=float)
    for k in range(A_d.shape[0]):
        i = k + 1
        base = x_prop_plus[i] if has_impulsive else x_prop[k]
        bias_k = (
            base
            - A_d[k] @ x_bar[k]
            - B_d[k][:, self._slice_cont] @ u_bar[k][self._slice_cont]
            - C_d[k][:, self._slice_cont] @ u_bar[i][self._slice_cont]
            - (E_d[i][:, self._slice_imp] @ u_bar[i][self._slice_imp] if has_impulsive else 0)
        )
        dyn_bias[k] = bias_k
    self._ocp_vars.dyn_bias.value = dyn_bias

    if has_impulsive:
        x0_imp_bias = x_prop_plus[0] - E_d[0][:, self._slice_imp] @ u_bar[0][self._slice_imp]
    else:
        x0_imp_bias = np.zeros(A_d.shape[1], dtype=float)
    self._ocp_vars.x0_imp_bias.value = x0_imp_bias

get_stats() -> dict

Get solver statistics for diagnostics and printing.

Returns:

Type	Description
dict	Dict containing: - n_variables: Total number of optimization variables - n_parameters: Total number of parameters - n_constraints: Total number of constraints

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def get_stats(self) -> dict:
    """Get solver statistics for diagnostics and printing.

    Returns:
        Dict containing:
            - ``n_variables``: Total number of optimization variables
            - ``n_parameters``: Total number of parameters
            - ``n_constraints``: Total number of constraints
    """
    if self._problem is None:
        return {"n_variables": 0, "n_parameters": 0, "n_constraints": 0}

    return {
        "n_variables": sum(var.size for var in self._problem.variables()),
        "n_parameters": sum(param.size for param in self._problem.parameters()),
        "n_constraints": sum(constraint.size for constraint in self._problem.constraints),
    }

initialize(lowered: LoweredProblem, settings: Config) -> None

Build the CVXPy optimal control problem.

Constructs the complete optimization problem by calling cost() and constraints() to build the objective and constraint formulations, then assembles them into a CVXPy Problem.

If cvxpygen is enabled, generates compiled solver code for improved performance.

Note

create_variables() must be called before this method.

Parameters:

Name	Type	Description	Default
lowered	LoweredProblem	Lowered problem containing: - cvxpy_constraints: Lowered convex constraints - jax_constraints: JAX constraint functions (for structure)	required
settings	Config	Problem configuration (node count, scaling, etc.)	required

Raises:

Type	Description
RuntimeError	If create_variables() has not been called.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def initialize(
    self,
    lowered: "LoweredProblem",
    settings: "Config",
) -> None:
    """Build the CVXPy optimal control problem.

    Constructs the complete optimization problem by calling ``cost()`` and
    ``constraints()`` to build the objective and constraint formulations,
    then assembles them into a CVXPy Problem.

    If cvxpygen is enabled, generates compiled solver code for improved
    performance.

    Note:
        ``create_variables()`` must be called before this method.

    Args:
        lowered: Lowered problem containing:
            - ``cvxpy_constraints``: Lowered convex constraints
            - ``jax_constraints``: JAX constraint functions (for structure)
        settings: Problem configuration (node count, scaling, etc.)

    Raises:
        RuntimeError: If create_variables() has not been called.
    """
    if self._ocp_vars is None:
        raise RuntimeError(
            "CVXPyPTRSolver.initialize() called before create_variables(). "
            "Call create_variables() first to create optimization variables."
        )

    objective = self.cost(settings, lowered)
    constr = self.constraints(settings, lowered)
    prob = cp.Problem(cp.Minimize(objective), constr)

    if self.cvxpygen:
        if not CVXPYGEN_AVAILABLE:
            raise ImportError(
                "cvxpygen is required for code generation but not installed. "
                "Install it with: pip install openscvx[cvxpygen] or pip install cvxpygen"
            )
        # Check to see if solver directory exists
        if not os.path.exists("solver"):
            cpg.generate_code(prob, solver=self.cvx_solver, code_dir="solver", wrapper=True)
        else:
            # Prompt the use to indicate if they wish to overwrite the solver
            # directory or use the existing compiled solver
            if self.cvxpygen_override:
                cpg.generate_code(
                    prob,
                    solver=self.cvx_solver,
                    code_dir="solver",
                    wrapper=True,
                )
            else:
                overwrite = input("Solver directory already exists. Overwrite? (y/n): ")
                if overwrite.lower() == "y":
                    cpg.generate_code(
                        prob,
                        solver=self.cvx_solver,
                        code_dir="solver",
                        wrapper=True,
                    )

    self._problem = prob
    self._settings = settings
    self._setup_solve_function()

iteration_callback() -> Callable[..., SubproblemSolution]

JAX-pure (state, SubproblemData) -> SubproblemSolution.

Wraps the existing NumPy solve() path in :func:jax.pure_callback with vmap_method="sequential". CVXPy can't trace — it builds a DCP graph and reaches into a backend solver through Python attribute mutation — so the only JAX-friendly surface it can present is a host callback. Under :func:jax.jit the callback fires once per call; under :func:jax.vmap it fires B times serially (CVXPy can't ingest a batched parameter set).

The closure unstacks :class:SubproblemData back to the per-constraint dict layout :meth:update_constraint_linearizations expects, runs the four update_* methods + :meth:solve, and packs the result into a :class:SubproblemSolution whose nu_vb is collapsed to a stacked (N, n_nodal) array (the list-of-arrays form on :attr:PTRSolveResult.nu_vb is preserved for the NumPy path).

status_code is mapped via :func:status_str_to_code; CVXPy's "optimal_inaccurate" / "solver_error" / ... labels collapse to :attr:StatusCode.UNKNOWN — the SCP trust-region check is the authoritative convergence gate.

The returned callable takes (state, data): state is the :class:AlgorithmState pytree, accepted for cross-backend signature uniformity but unused (CVXPy exposes no warm-start through update_*); data is the :class:SubproblemData pytree carrying the per-iteration linearization arrays, penalty weights, and boundary conditions.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def iteration_callback(self) -> Callable[..., SubproblemSolution]:
    """JAX-pure ``(state, SubproblemData) -> SubproblemSolution``.

    Wraps the existing NumPy ``solve()`` path in :func:`jax.pure_callback`
    with ``vmap_method="sequential"``. CVXPy can't trace — it builds a DCP
    graph and reaches into a backend solver through Python attribute
    mutation — so the only JAX-friendly surface it can present is a host
    callback. Under :func:`jax.jit` the callback fires once per call; under
    :func:`jax.vmap` it fires ``B`` times serially (CVXPy can't ingest a
    batched parameter set).

    The closure unstacks :class:`SubproblemData` back to the per-constraint
    dict layout :meth:`update_constraint_linearizations` expects, runs the
    four ``update_*`` methods + :meth:`solve`, and packs the result into a
    :class:`SubproblemSolution` whose ``nu_vb`` is collapsed to a stacked
    ``(N, n_nodal)`` array (the list-of-arrays form on
    :attr:`PTRSolveResult.nu_vb` is preserved for the NumPy path).

    ``status_code`` is mapped via :func:`status_str_to_code`; CVXPy's
    ``"optimal_inaccurate"`` / ``"solver_error"`` / ... labels collapse to
    :attr:`StatusCode.UNKNOWN` — the SCP trust-region check is the
    authoritative convergence gate.

    The returned callable takes ``(state, data)``: ``state`` is the
    :class:`AlgorithmState` pytree, accepted for cross-backend signature
    uniformity but unused (CVXPy exposes no warm-start through
    ``update_*``); ``data`` is the :class:`SubproblemData` pytree carrying
    the per-iteration linearization arrays, penalty weights, and boundary
    conditions.
    """
    if self._problem is None or self._settings is None or self._jax_constraints is None:
        raise RuntimeError(
            "CVXPyPTRSolver.iteration_callback() requires initialize() to "
            "have been called first."
        )

    settings = self._settings
    jax_constraints = self._jax_constraints
    n_x = settings.sim.n_states
    n_u = settings.sim.n_controls
    N = settings.sim.n
    n_nodal = len(jax_constraints.nodal)
    n_cross = len(jax_constraints.cross_node)
    slice_imp = settings.sim.u.slice_impulsive
    has_impulsive = bool(slice_imp.stop > slice_imp.start)
    f = jnp.float64 if jax.config.read("jax_enable_x64") else jnp.float32

    result_struct = SubproblemSolution(
        x=jax.ShapeDtypeStruct((N, n_x), f),
        u=jax.ShapeDtypeStruct((N, n_u), f),
        nu=jax.ShapeDtypeStruct((N - 1, n_x), f),
        nu_vb=jax.ShapeDtypeStruct((N, n_nodal), f),
        nu_vb_cross=jax.ShapeDtypeStruct((n_cross,), f),
        cost=jax.ShapeDtypeStruct((), f),
        status_code=jax.ShapeDtypeStruct((), jnp.int32),
    )

    def host_solve(data: SubproblemData) -> SubproblemSolution:
        # SubproblemData always carries (N, n_x, n_x) D_d / (N, n_x, n_u) E_d
        # / (N, n_x) x_prop_plus arrays — zero-filled when the problem has
        # no impulsive component. ``update_dynamics_linearization`` treats
        # a non-None D_d as "absorb into A/B/C", which would zero them out
        # on the no-impulsive path; pass None in that case to keep A/B/C
        # intact. Same logic for x_prop_plus / E_d.
        x_prop_plus = np.asarray(data.x_prop_plus) if has_impulsive else None
        D_d = np.asarray(data.D_d) if has_impulsive else None
        E_d = np.asarray(data.E_d) if has_impulsive else None
        self.update_dynamics_linearization(
            x_bar=np.asarray(data.x_bar),
            u_bar=np.asarray(data.u_bar),
            A_d=np.asarray(data.A_d),
            B_d=np.asarray(data.B_d),
            C_d=np.asarray(data.C_d),
            x_prop=np.asarray(data.x_prop),
            x_prop_plus=x_prop_plus,
            D_d=D_d,
            E_d=E_d,
        )
        nodal = _unstack_nodal(data, jax_constraints)
        cross_node = _unstack_cross(data, jax_constraints)
        self.update_constraint_linearizations(
            nodal=nodal or None,
            cross_node=cross_node or None,
        )
        self.update_penalties(
            lam_prox=np.asarray(data.lam_prox),
            lam_cost=np.asarray(data.lam_cost),
            lam_vc=np.asarray(data.lam_vc),
            lam_vb_nodal=np.asarray(data.lam_vb_nodal),
            lam_vb_cross=np.asarray(data.lam_vb_cross),
        )
        # Set the proximal-term parameter (``prox_c``) from the freshly
        # updated penalties; the legacy ``_subproblem`` called this between
        # ``update_penalties`` and ``solve``, and without it ``prox_c`` has
        # no value on the first solve of a process.
        self.update_proximal_terms()
        self.update_boundary_conditions(
            x_init=_unmask_bc(data.x_init),
            x_term=_unmask_bc(data.x_term),
        )

        try:
            result = self.solve()
        except cp.error.SolverError:
            # Infeasible / numerical failure. Surface it as a non-OPTIMAL
            # status on a finite (zero) solution so the SCP loop fails
            # loudly at the step() boundary, instead of the exception
            # propagating through pure_callback as an opaque CpuCallback
            # error.
            return SubproblemSolution(
                x=jnp.zeros((N, n_x), dtype=f),
                u=jnp.zeros((N, n_u), dtype=f),
                nu=jnp.zeros((N - 1, n_x), dtype=f),
                nu_vb=jnp.zeros((N, n_nodal), dtype=f),
                nu_vb_cross=jnp.zeros((n_cross,), dtype=f),
                cost=jnp.asarray(0.0, dtype=f),
                status_code=jnp.asarray(int(StatusCode.INFEASIBLE), dtype=jnp.int32),
            )

        # ``optimal_inaccurate`` is a usable solution — the legacy path
        # consumed it without error — so treat it as OPTIMAL; the step()
        # status gate would otherwise reject the UNKNOWN it maps to.
        code = (
            StatusCode.OPTIMAL
            if (result.status == "optimal_inaccurate")
            or (result.status == "solved")
            or (result.status == "1 (for description visit https://qoco-org.github.io/qoco/)")
            else status_str_to_code(result.status)
        )
        return SubproblemSolution(
            x=jnp.asarray(result.x, dtype=f),
            u=jnp.asarray(result.u, dtype=f),
            nu=jnp.asarray(result.nu, dtype=f),
            # nu_vb collapses from PTRSolveResult's list-of-(N,) layout to
            # the (N, n_nodal) stacked array SubproblemSolution declares
            # for pure_callback's ShapeDtypeStruct contract.
            nu_vb=(
                jnp.zeros((N, 0), dtype=f)
                if n_nodal == 0
                else jnp.asarray(
                    np.stack([np.asarray(a) for a in result.nu_vb], axis=-1),
                    dtype=f,
                )
            ),
            nu_vb_cross=jnp.asarray(
                np.asarray(result.nu_vb_cross, dtype=float).reshape(n_cross),
                dtype=f,
            ),
            cost=jnp.asarray(result.cost, dtype=f),
            status_code=jnp.asarray(int(code), dtype=jnp.int32),
        )

    def step(state, data: SubproblemData) -> SubproblemSolution:
        del state  # CVXPy exposes no warm-start through ``update_*``.
        return jax.pure_callback(
            host_solve,
            result_struct,
            data,
            vmap_method="sequential",
        )

    return step

lower_convex_constraints(constraints, parameters=None)

Lower user .convex() constraints into CVXPy constraint objects.

Delegates to :func:openscvx.symbolic.lower.lower_cvxpy_constraints, feeding it the unscaled-state and unscaled-control CVXPy expressions built by :meth:create_variables.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def lower_convex_constraints(self, constraints, parameters=None):
    """Lower user ``.convex()`` constraints into CVXPy constraint objects.

    Delegates to :func:`openscvx.symbolic.lower.lower_cvxpy_constraints`,
    feeding it the unscaled-state and unscaled-control CVXPy expressions
    built by :meth:`create_variables`.
    """
    from openscvx.symbolic.lower import lower_cvxpy_constraints

    if self._ocp_vars is None:
        raise RuntimeError(
            "CVXPyPTRSolver.lower_convex_constraints() called before "
            "create_variables(); the CVXPy variables it needs don't "
            "exist yet."
        )
    return lower_cvxpy_constraints(
        constraints,
        self._ocp_vars.x_nonscaled,
        self._ocp_vars.u_nonscaled,
        parameters,
    )

proximal_cost_terms(lam_prox: np.ndarray, x_bar: np.ndarray, u_bar: np.ndarray) -> tuple[np.ndarray, np.ndarray]

Compute linear and constant proximal cost terms.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def proximal_cost_terms(
    self,
    lam_prox: np.ndarray,
    x_bar: np.ndarray,
    u_bar: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Compute linear and constant proximal cost terms."""
    x_ref = (self._ocp_vars.inv_S_x @ (x_bar.T - self._ocp_vars.c_x[:, None])).T
    u_ref = (self._ocp_vars.inv_S_u @ (u_bar.T - self._ocp_vars.c_u[:, None])).T
    z_ref = np.hstack((x_ref, u_ref))
    prox_c = -2.0 * lam_prox * z_ref
    prox_cc = np.sum(lam_prox * np.square(z_ref), axis=1)
    return prox_c, prox_cc

solve() -> PTRSolveResult

Solve the convex subproblem and return structured results.

Call update_dynamics_linearization(), update_constraint_linearizations(), and update_penalties() before calling this method.

Returns:

Type	Description
PTRSolveResult	PTRSolveResult containing unscaled trajectories, slack variables,
PTRSolveResult	cost, and solver status.

Raises:

Type	Description
RuntimeError	If initialize() has not been called.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def solve(self) -> PTRSolveResult:
    """Solve the convex subproblem and return structured results.

    Call ``update_dynamics_linearization()``, ``update_constraint_linearizations()``,
    and ``update_penalties()`` before calling this method.

    Returns:
        PTRSolveResult containing unscaled trajectories, slack variables,
        cost, and solver status.

    Raises:
        RuntimeError: If initialize() has not been called.
    """
    if self._problem is None:
        raise RuntimeError(
            "CVXPyPTRSolver.solve() called before initialize(). "
            "Call initialize() first to build the problem structure."
        )

    self._solve_fn()

    # Get scaling matrices
    S_x = self._ocp_vars.S_x
    c_x = self._ocp_vars.c_x
    S_u = self._ocp_vars.S_u
    c_u = self._ocp_vars.c_u

    # Unscale state and control trajectories
    x_scaled = self._problem.var_dict["x"].value  # (N, n_states)
    u_scaled = self._problem.var_dict["u"].value  # (N, n_controls)
    x = (S_x @ x_scaled.T + np.expand_dims(c_x, axis=1)).T
    u = (S_u @ u_scaled.T + np.expand_dims(c_u, axis=1)).T

    # Get virtual control slack
    nu = self._problem.var_dict["nu"].value

    # Get nodal constraint violation slacks
    nu_vb = [var.value for var in self._ocp_vars.nu_vb]

    # Get cross-node constraint violation slacks
    nu_vb_cross = [var.value for var in self._ocp_vars.nu_vb_cross]

    return PTRSolveResult(
        x=x,
        u=u,
        nu=nu,
        nu_vb=nu_vb,
        nu_vb_cross=nu_vb_cross,
        cost=self._problem.value,
        status=self._problem.status,
    )

update_boundary_conditions(x_init: np.ndarray = None, x_term: np.ndarray = None) -> None

Update boundary condition parameters.

Sets initial and/or terminal state constraints. Only sets parameters that exist in the problem (some problems may not have both).

Parameters:

Name	Type	Description	Default
x_init	ndarray	Initial state vector, shape (n_states,). Optional.	None
x_term	ndarray	Terminal state vector, shape (n_states,). Optional.	None

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def update_boundary_conditions(
    self,
    x_init: np.ndarray = None,
    x_term: np.ndarray = None,
) -> None:
    """Update boundary condition parameters.

    Sets initial and/or terminal state constraints. Only sets parameters
    that exist in the problem (some problems may not have both).

    Args:
        x_init: Initial state vector, shape (n_states,). Optional.
        x_term: Terminal state vector, shape (n_states,). Optional.
    """
    # No-op before initialize() — the CVXPy problem (and its param_dict)
    # isn't built yet. Callers like Problem._sync_boundary_conditions
    # may invoke this both before and after initialize().
    if self._problem is None:
        return
    if x_init is not None and "x_init" in self._problem.param_dict:
        self._set_param("x_init", x_init)
    if x_term is not None and "x_term" in self._problem.param_dict:
        self._set_param("x_term", x_term)

update_constraint_linearizations(nodal: List[dict] = None, cross_node: List[dict] = None) -> None

Update linearized constraint values and gradients.

Sets constraint function values and gradients at the current linearization point for both nodal and cross-node constraints.

Parameters:

Name	Type	Description	Default
nodal	List[dict]	List of dicts for nodal constraints, each containing: - g: Constraint value at linearization point - grad_g_x: Gradient w.r.t. state - grad_g_u: Gradient w.r.t. control	None
cross_node	List[dict]	List of dicts for cross-node constraints, each containing: - g: Constraint value at linearization point - grad_g_X: Gradient w.r.t. full state trajectory - grad_g_U: Gradient w.r.t. full control trajectory	None

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def update_constraint_linearizations(
    self,
    nodal: List[dict] = None,
    cross_node: List[dict] = None,
) -> None:
    """Update linearized constraint values and gradients.

    Sets constraint function values and gradients at the current
    linearization point for both nodal and cross-node constraints.

    Args:
        nodal: List of dicts for nodal constraints, each containing:
            - ``g``: Constraint value at linearization point
            - ``grad_g_x``: Gradient w.r.t. state
            - ``grad_g_u``: Gradient w.r.t. control
        cross_node: List of dicts for cross-node constraints, each containing:
            - ``g``: Constraint value at linearization point
            - ``grad_g_X``: Gradient w.r.t. full state trajectory
            - ``grad_g_U``: Gradient w.r.t. full control trajectory
    """
    if nodal:
        x_bar = self._ocp_vars.x_bar.value
        u_bar = self._ocp_vars.u_bar.value
        x_bar_arr = np.asarray(x_bar)
        u_bar_arr = np.asarray(u_bar)
        for g_id, constraint_data in enumerate(nodal):
            g_arr = np.asarray(constraint_data["g"])
            grad_x_arr = np.asarray(constraint_data["grad_g_x"])
            grad_u_arr = np.asarray(constraint_data["grad_g_u"])

            # Convert to affine form in (X, U):
            # g_tilde + grad_x·X + grad_u·U, where
            # g_tilde = g(X_bar, U_bar) - grad_x·X_bar - grad_u·U_bar.
            g_tilde = (
                g_arr
                - np.sum(grad_x_arr * x_bar_arr, axis=1)
                - np.sum(grad_u_arr * u_bar_arr, axis=1)
            )

            self._set_param(f"g_{g_id}", g_tilde)
            self._set_param(f"grad_g_x_{g_id}", grad_x_arr)
            self._set_param(f"grad_g_u_{g_id}", grad_u_arr)

    if cross_node:
        x_bar = self._ocp_vars.x_bar.value
        u_bar = self._ocp_vars.u_bar.value
        x_bar_arr = np.asarray(x_bar)
        u_bar_arr = np.asarray(u_bar)
        for g_id, constraint_data in enumerate(cross_node):
            g_val = np.asarray(constraint_data["g"])
            grad_X_arr = np.asarray(constraint_data["grad_g_X"])
            grad_U_arr = np.asarray(constraint_data["grad_g_U"])

            g_tilde = g_val - np.sum(grad_X_arr * x_bar_arr) - np.sum(grad_U_arr * u_bar_arr)

            self._set_param(f"g_cross_{g_id}", g_tilde)
            self._set_param(f"grad_g_X_cross_{g_id}", grad_X_arr)
            self._set_param(f"grad_g_U_cross_{g_id}", grad_U_arr)

update_dynamics_linearization(x_bar: np.ndarray, u_bar: np.ndarray, A_d: np.ndarray, B_d: np.ndarray, C_d: np.ndarray, x_prop: np.ndarray, x_prop_plus: np.ndarray | None = None, D_d: np.ndarray | None = None, E_d: np.ndarray | None = None) -> None

Update dynamics linearization point and matrices.

Sets the current linearization point (previous iterate) and the discretized dynamics matrices for the convex subproblem.

Parameters:

Name	Type	Description	Default
x_bar	ndarray	Previous state trajectory, shape (N, n_states)	required
u_bar	ndarray	Previous control trajectory, shape (N, n_controls)	required
A_d	ndarray	Discretized state Jacobian, shape (N-1, n_states, n_states)	required
B_d	ndarray	Discretized control Jacobian (current node), shape (N-1, n_states, n_controls)	required
C_d	ndarray	Discretized control Jacobian (next node), shape (N-1, n_states, n_controls)	required
x_prop	ndarray	Propagated state from continuous dynamics, shape (N-1, n_states)	required
x_prop_plus	ndarray | None	Optional impulsive/discrete propagated state, shape (N, n_states)	None
D_d	ndarray | None	Optional impulsive/discrete Jacobian wrt state, shape (N, n_states, n_states)	None
E_d	ndarray | None	Optional impulsive/discrete Jacobian wrt control, shape (N, n_states, n_controls)	None

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def update_dynamics_linearization(
    self,
    x_bar: np.ndarray,
    u_bar: np.ndarray,
    A_d: np.ndarray,
    B_d: np.ndarray,
    C_d: np.ndarray,
    x_prop: np.ndarray,
    x_prop_plus: np.ndarray | None = None,
    D_d: np.ndarray | None = None,
    E_d: np.ndarray | None = None,
) -> None:
    """Update dynamics linearization point and matrices.

    Sets the current linearization point (previous iterate) and the
    discretized dynamics matrices for the convex subproblem.

    Args:
        x_bar: Previous state trajectory, shape (N, n_states)
        u_bar: Previous control trajectory, shape (N, n_controls)
        A_d: Discretized state Jacobian, shape (N-1, n_states, n_states)
        B_d: Discretized control Jacobian (current node), shape (N-1, n_states, n_controls)
        C_d: Discretized control Jacobian (next node), shape (N-1, n_states, n_controls)
        x_prop: Propagated state from continuous dynamics, shape (N-1, n_states)
        x_prop_plus: Optional impulsive/discrete propagated state, shape (N, n_states)
        D_d: Optional impulsive/discrete Jacobian wrt state, shape (N, n_states, n_states)
        E_d: Optional impulsive/discrete Jacobian wrt control, shape (N, n_states, n_controls)
    """
    x_bar_arr = np.asarray(x_bar)
    u_bar_arr = np.asarray(u_bar)
    self._ocp_vars.x_bar.value = x_bar_arr
    self._ocp_vars.u_bar.value = u_bar_arr

    A_eff = np.asarray(A_d)
    B_eff = np.asarray(B_d)
    C_eff = np.asarray(C_d)

    if D_d is not None:
        D_arr = np.asarray(D_d)
        # Temporary DPP-safe workaround: absorb D_d into A/B/C numerically
        # so the CVXPY graph has only Parameter@Variable products.
        if D_arr.ndim == 3 and D_arr.shape[0] == A_eff.shape[0] + 1:
            D_steps = D_arr[1:]
        elif D_arr.ndim == 3 and D_arr.shape[0] == A_eff.shape[0]:
            D_steps = D_arr
        else:
            raise ValueError(
                "Unexpected D_d shape for dynamics update: "
                f"{D_arr.shape}, expected "
                f"{(A_eff.shape[0] + 1, A_eff.shape[1], A_eff.shape[2])} "
                f"or {(A_eff.shape[0], A_eff.shape[1], A_eff.shape[2])}."
            )

        A_eff = np.einsum("kij,kjl->kil", D_steps, A_eff)
        B_eff = np.einsum("kij,kjl->kil", D_steps, B_eff)
        C_eff = np.einsum("kij,kjl->kil", D_steps, C_eff)

    self._set_param("A_d", A_eff)
    self._set_param("B_d", B_eff)
    self._set_param("C_d", C_eff)
    x_prop_arr = np.asarray(x_prop)
    x_prop_plus_arr = np.asarray(x_prop_plus) if x_prop_plus is not None else None
    E_arr = None
    if E_d is not None:
        E_arr = np.asarray(E_d)
        self._ocp_vars.E_d.value = E_arr

    self.dynamics_biases(
        A_d=A_eff,
        B_d=B_eff,
        C_d=C_eff,
        x_bar=x_bar_arr,
        u_bar=u_bar_arr,
        x_prop=x_prop_arr,
        x_prop_plus=x_prop_plus_arr,
        E_d=E_arr,
    )

update_penalties(lam_prox: np.ndarray, lam_cost: Union[float, np.ndarray], lam_vc: np.ndarray, lam_vb_nodal: np.ndarray, lam_vb_cross: np.ndarray) -> None

Update SCP penalty weights.

Sets the penalty weights that balance competing objectives in the PTR convex subproblem.

Parameters:

Name	Type	Description	Default
lam_prox	ndarray	Trust region weights, shape (N, n_states + n_controls).	required
lam_cost	Union[float, ndarray]	Cost function weight. Scalar or array of shape (n_states,) for per-state weighting.	required
lam_vc	ndarray	Virtual control penalty weights, shape (N-1, n_states)	required
lam_vb_nodal	ndarray	Virtual buffer penalty weights for nodal constraints, shape (N, n_nodal_constraints).	required
lam_vb_cross	ndarray	Virtual buffer penalty weights for cross-node constraints, shape (n_cross_node_constraints,).	required

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def update_penalties(
    self,
    lam_prox: np.ndarray,
    lam_cost: Union[float, np.ndarray],
    lam_vc: np.ndarray,
    lam_vb_nodal: np.ndarray,
    lam_vb_cross: np.ndarray,
) -> None:
    """Update SCP penalty weights.

    Sets the penalty weights that balance competing objectives in the
    PTR convex subproblem.

    Args:
        lam_prox: Trust region weights, shape ``(N, n_states + n_controls)``.
        lam_cost: Cost function weight. Scalar or array of shape
            ``(n_states,)`` for per-state weighting.
        lam_vc: Virtual control penalty weights, shape (N-1, n_states)
        lam_vb_nodal: Virtual buffer penalty weights for nodal constraints,
            shape ``(N, n_nodal_constraints)``.
        lam_vb_cross: Virtual buffer penalty weights for cross-node
            constraints, shape ``(n_cross_node_constraints,)``.
    """
    lam_prox_arr = np.asarray(lam_prox)
    self._set_param("lam_prox", lam_prox_arr)
    self._set_param("lam_cost", lam_cost)
    self._set_param("lam_vc", lam_vc)
    self._set_param("lam_vb_nodal", lam_vb_nodal)
    self._set_param("lam_vb_cross", lam_vb_cross)

update_proximal_terms() -> None

Update proximal expansion parameters from current references and weights.

Source code in openscvx/solvers/cvxpy_ptr_solver.py

def update_proximal_terms(self) -> None:
    """Update proximal expansion parameters from current references and weights."""
    lam_prox_arr = np.asarray(self._ocp_vars.lam_prox.value)
    x_bar = np.asarray(self._ocp_vars.x_bar.value)
    u_bar = np.asarray(self._ocp_vars.u_bar.value)
    prox_c, prox_cc = self.proximal_cost_terms(lam_prox_arr, x_bar, u_bar)
    self._set_param("prox_c", prox_c)
    self._set_param("prox_cc", prox_cc)

ConvexSolver

Bases: ABC

Abstract base class for convex subproblem solvers.

This class defines the interface for solvers that handle the convex subproblems generated at each iteration of a successive convexification algorithm.

Subclasses must implement all abstract methods below.

The solver lifecycle has two phases:

Setup (called once):

• create_variables: Create backend-specific variables
• initialize: Build the problem structure using lowered constraints

Per-iteration (called each SCP iteration):

• update_dynamics_linearization: Set linearization point and dynamics matrices
• update_constraint_linearizations: Set constraint values and gradients
• update_penalties: Set penalty weights
• solve: Solve and return results

Example

Implementing a custom solver::

class MySolver(ConvexSolver):
    def create_variables(self, N, x_unified, u_unified, jax_constraints):
        self._vars = create_my_variables(N, x_unified, ...)

    def initialize(self, lowered, settings):
        self._prob = build_my_problem(self._vars, lowered, settings)

    def update_dynamics_linearization(self, **kwargs):
        # Set x_bar, u_bar, A_d, B_d, etc.
        ...

    def update_constraint_linearizations(self, **kwargs):
        # Set constraint function values and gradients
        ...

    def update_penalties(self, **kwargs):
        # Set lam_prox, lam_cost, lam_vc, lam_vb_nodal, lam_vb_cross
        ...

    def solve(self):
        self._prob.solve()
        return MyResult(...)

Source code in openscvx/solvers/base.py

class ConvexSolver(ABC):
    """Abstract base class for convex subproblem solvers.

    This class defines the interface for solvers that handle the convex
    subproblems generated at each iteration of a successive convexification
    algorithm.

    Subclasses must implement all abstract methods below.

    The solver lifecycle has two phases:

    **Setup (called once):**

    - create_variables: Create backend-specific variables
    - initialize: Build the problem structure using lowered constraints

    **Per-iteration (called each SCP iteration):**

    - update_dynamics_linearization: Set linearization point and dynamics matrices
    - update_constraint_linearizations: Set constraint values and gradients
    - update_penalties: Set penalty weights
    - solve: Solve and return results

    Example:
        Implementing a custom solver::

            class MySolver(ConvexSolver):
                def create_variables(self, N, x_unified, u_unified, jax_constraints):
                    self._vars = create_my_variables(N, x_unified, ...)

                def initialize(self, lowered, settings):
                    self._prob = build_my_problem(self._vars, lowered, settings)

                def update_dynamics_linearization(self, **kwargs):
                    # Set x_bar, u_bar, A_d, B_d, etc.
                    ...

                def update_constraint_linearizations(self, **kwargs):
                    # Set constraint function values and gradients
                    ...

                def update_penalties(self, **kwargs):
                    # Set lam_prox, lam_cost, lam_vc, lam_vb_nodal, lam_vb_cross
                    ...

                def solve(self):
                    self._prob.solve()
                    return MyResult(...)
    """

    @abstractmethod
    def create_variables(
        self,
        N: int,
        x_unified: "UnifiedState",
        u_unified: "UnifiedControl",
        jax_constraints: "LoweredJaxConstraints",
        dynamics_sparsity: Optional[tuple] = None,
        constraint_sparsity: Optional[list] = None,
    ) -> None:
        """Create backend-specific optimization variables.

        This method creates the optimization variables (decision variables and
        parameters) for this solver's backend. Called once during problem setup,
        before constraint lowering.

        The solver should store its variables on ``self`` for use in subsequent
        ``initialize()`` and ``solve()`` calls.

        Args:
            N: Number of discretization nodes
            x_unified: Unified state interface with dimensions and scaling bounds
            u_unified: Unified control interface with dimensions and scaling bounds
            jax_constraints: Lowered JAX constraints (for sizing linearization params)
            dynamics_sparsity: Optional tuple ``(A_d, B_d, C_d)`` of boolean
                ndarrays giving the discrete-time Jacobian sparsity patterns.
                ``A_d`` has shape ``(n_x, n_x)``; ``B_d`` and ``C_d`` have
                shape ``(n_x, n_u)``.
            constraint_sparsity: Optional list of ``(x_mask, u_mask)`` boolean
                1-D arrays, one per nodal constraint.
        """
        raise NotImplementedError

    def lower_convex_constraints(
        self,
        constraints: "ConstraintSet",
        parameters: Optional[Dict[str, Any]] = None,
    ) -> Tuple[List[Any], Dict[str, Any]]:
        """Lower user ``.convex()`` constraints into this backend's form.

        Called once by :func:`openscvx.symbolic.lower.lower_symbolic_problem`
        after ``create_variables()`` and before ``initialize()``.

        The default implementation refuses any user ``.convex()``
        constraints — appropriate for backends like
        :class:`openscvx.solvers.qpax_ptr_solver.QPAXPTRSolver` that don't
        accept second-order-cone constraints. Backends that do accept them
        (e.g. :class:`openscvx.solvers.cvxpy_ptr_solver.CVXPyPTRSolver`)
        override this to invoke their backend-specific lowerer.

        Args:
            constraints: Categorized symbolic constraints. Only the
                ``nodal_convex`` / ``cross_node_convex`` lists matter here;
                non-convex constraints go through the JAX lowering pipeline.
            parameters: Optional dict of symbolic ``Parameter`` objects
                referenced by the constraints. May be ``None``.

        Returns:
            ``(lowered_list, parameter_map)``. The first is a list of
            backend-specific constraint objects (e.g. ``cp.Constraint``);
            the second maps parameter names to backend-specific parameter
            objects. Both are empty for the default refusal path.

        Raises:
            NotImplementedError: if the user defined any ``.convex()``
                constraints and this backend doesn't override.
        """
        n = len(constraints.nodal_convex) + len(constraints.cross_node_convex)
        if n:
            raise NotImplementedError(
                f"{type(self).__name__} does not support user-defined "
                f".convex() constraints ({n} defined). Drop the .convex() "
                "constraint or switch to a backend that supports them "
                "(e.g. openscvx.CVXPyPTRSolver)."
            )
        return [], {}

    @abstractmethod
    def initialize(
        self,
        lowered: "LoweredProblem",
        settings: "Config",
    ) -> None:
        """Build the convex subproblem structure.

        This method constructs the optimization problem once, using CVXPy
        Parameters (or equivalent) for values that change each iteration.
        Called once during problem setup, not at each SCP iteration.

        The solver should store its problem representation on ``self`` for use
        in subsequent ``solve()`` calls.

        Args:
            lowered: Lowered problem containing:
                - ``cvxpy_constraints``: Lowered convex constraints
                - ``jax_constraints``: JAX constraint functions
                - ``x_unified``, ``u_unified``: State/control interfaces
            settings: Configuration object with solver settings
        """
        raise NotImplementedError

    @abstractmethod
    def update_dynamics_linearization(self, **kwargs) -> None:
        """Update dynamics linearization point and matrices.

        Called at each SCP iteration before ``solve()`` to set the current
        linearization point and discretized dynamics matrices.

        The specific parameters depend on the solver implementation.
        See concrete solver classes for expected arguments.
        """
        raise NotImplementedError

    @abstractmethod
    def update_constraint_linearizations(self, **kwargs) -> None:
        """Update linearized constraint values and gradients.

        Called at each SCP iteration before ``solve()`` to set constraint
        function values and gradients at the current linearization point.

        The specific parameters depend on the solver implementation.
        See concrete solver classes for expected arguments.
        """
        raise NotImplementedError

    @abstractmethod
    def update_penalties(self, **kwargs) -> None:
        """Update SCP penalty weights.

        Called at each SCP iteration before ``solve()`` to set the current
        penalty weights for trust region, virtual control, and virtual buffer.

        The specific parameters depend on the solver implementation.
        See concrete solver classes for expected arguments.
        """
        raise NotImplementedError

    def update_proximal_terms(self) -> None:
        """Update optional proximal-expansion parameters.

        Called after ``update_penalties()`` and before ``solve()``. Solvers
        that formulate trust-region penalties via expanded quadratic forms can
        override this hook to refresh any derived coefficients.

        Default implementation is a no-op.
        """
        return None

    @abstractmethod
    def update_boundary_conditions(self, **kwargs) -> None:
        """Update boundary condition parameters.

        Called once during algorithm initialization to set initial and terminal
        state constraints.

        The specific parameters depend on the solver implementation.
        See concrete solver classes for expected arguments.
        """
        raise NotImplementedError

    @abstractmethod
    def get_stats(self) -> dict:
        """Get solver statistics for diagnostics and printing.

        Returns:
            Dict containing solver statistics. Expected keys:
                - ``n_variables``: Total number of optimization variables
                - ``n_parameters``: Total number of parameters
                - ``n_constraints``: Total number of constraints
        """
        raise NotImplementedError

    @abstractmethod
    def solve(self) -> Any:
        """Solve the convex subproblem and return results.

        Called at each SCP iteration after updating linearization and penalties.
        Returns a solver-specific result object containing the solution.

        Returns:
            Solver-specific result object (e.g., ``PTRSolveResult`` for PTR).
        """
        raise NotImplementedError

    @abstractmethod
    def citation(self) -> List[str]:
        """Return BibTeX citations for this solver.

        Implementations should return a list of BibTeX entry strings for the
        papers that should be cited when using this solver.

        Returns:
            List of BibTeX citation strings.
        """
        raise NotImplementedError

citation() -> List[str] abstractmethod

Return BibTeX citations for this solver.

Implementations should return a list of BibTeX entry strings for the papers that should be cited when using this solver.

Returns:

Type	Description
List[str]	List of BibTeX citation strings.

Source code in openscvx/solvers/base.py

@abstractmethod
def citation(self) -> List[str]:
    """Return BibTeX citations for this solver.

    Implementations should return a list of BibTeX entry strings for the
    papers that should be cited when using this solver.

    Returns:
        List of BibTeX citation strings.
    """
    raise NotImplementedError

create_variables(N: int, x_unified: UnifiedState, u_unified: UnifiedControl, jax_constraints: LoweredJaxConstraints, dynamics_sparsity: Optional[tuple] = None, constraint_sparsity: Optional[list] = None) -> None abstractmethod

Create backend-specific optimization variables.

This method creates the optimization variables (decision variables and parameters) for this solver's backend. Called once during problem setup, before constraint lowering.

The solver should store its variables on self for use in subsequent initialize() and solve() calls.

Parameters:

Name	Type	Description	Default
N	int	Number of discretization nodes	required
x_unified	UnifiedState	Unified state interface with dimensions and scaling bounds	required
u_unified	UnifiedControl	Unified control interface with dimensions and scaling bounds	required
jax_constraints	LoweredJaxConstraints	Lowered JAX constraints (for sizing linearization params)	required
dynamics_sparsity	Optional[tuple]	Optional tuple (A_d, B_d, C_d) of boolean ndarrays giving the discrete-time Jacobian sparsity patterns. A_d has shape (n_x, n_x); B_d and C_d have shape (n_x, n_u).	None
constraint_sparsity	Optional[list]	Optional list of (x_mask, u_mask) boolean 1-D arrays, one per nodal constraint.	None

Source code in openscvx/solvers/base.py

@abstractmethod
def create_variables(
    self,
    N: int,
    x_unified: "UnifiedState",
    u_unified: "UnifiedControl",
    jax_constraints: "LoweredJaxConstraints",
    dynamics_sparsity: Optional[tuple] = None,
    constraint_sparsity: Optional[list] = None,
) -> None:
    """Create backend-specific optimization variables.

    This method creates the optimization variables (decision variables and
    parameters) for this solver's backend. Called once during problem setup,
    before constraint lowering.

    The solver should store its variables on ``self`` for use in subsequent
    ``initialize()`` and ``solve()`` calls.

    Args:
        N: Number of discretization nodes
        x_unified: Unified state interface with dimensions and scaling bounds
        u_unified: Unified control interface with dimensions and scaling bounds
        jax_constraints: Lowered JAX constraints (for sizing linearization params)
        dynamics_sparsity: Optional tuple ``(A_d, B_d, C_d)`` of boolean
            ndarrays giving the discrete-time Jacobian sparsity patterns.
            ``A_d`` has shape ``(n_x, n_x)``; ``B_d`` and ``C_d`` have
            shape ``(n_x, n_u)``.
        constraint_sparsity: Optional list of ``(x_mask, u_mask)`` boolean
            1-D arrays, one per nodal constraint.
    """
    raise NotImplementedError

get_stats() -> dict abstractmethod

Get solver statistics for diagnostics and printing.

Returns:

Type	Description
dict	Dict containing solver statistics. Expected keys: - n_variables: Total number of optimization variables - n_parameters: Total number of parameters - n_constraints: Total number of constraints

Source code in openscvx/solvers/base.py

@abstractmethod
def get_stats(self) -> dict:
    """Get solver statistics for diagnostics and printing.

    Returns:
        Dict containing solver statistics. Expected keys:
            - ``n_variables``: Total number of optimization variables
            - ``n_parameters``: Total number of parameters
            - ``n_constraints``: Total number of constraints
    """
    raise NotImplementedError

initialize(lowered: LoweredProblem, settings: Config) -> None abstractmethod

Build the convex subproblem structure.

This method constructs the optimization problem once, using CVXPy Parameters (or equivalent) for values that change each iteration. Called once during problem setup, not at each SCP iteration.

The solver should store its problem representation on self for use in subsequent solve() calls.

Parameters:

Name	Type	Description	Default
lowered	LoweredProblem	Lowered problem containing: - cvxpy_constraints: Lowered convex constraints - jax_constraints: JAX constraint functions - x_unified, u_unified: State/control interfaces	required
settings	Config	Configuration object with solver settings	required

Source code in openscvx/solvers/base.py

@abstractmethod
def initialize(
    self,
    lowered: "LoweredProblem",
    settings: "Config",
) -> None:
    """Build the convex subproblem structure.

    This method constructs the optimization problem once, using CVXPy
    Parameters (or equivalent) for values that change each iteration.
    Called once during problem setup, not at each SCP iteration.

    The solver should store its problem representation on ``self`` for use
    in subsequent ``solve()`` calls.

    Args:
        lowered: Lowered problem containing:
            - ``cvxpy_constraints``: Lowered convex constraints
            - ``jax_constraints``: JAX constraint functions
            - ``x_unified``, ``u_unified``: State/control interfaces
        settings: Configuration object with solver settings
    """
    raise NotImplementedError

lower_convex_constraints(constraints: ConstraintSet, parameters: Optional[Dict[str, Any]] = None) -> Tuple[List[Any], Dict[str, Any]]

Lower user .convex() constraints into this backend's form.

Called once by :func:openscvx.symbolic.lower.lower_symbolic_problem after create_variables() and before initialize().

The default implementation refuses any user .convex() constraints — appropriate for backends like :class:openscvx.solvers.qpax_ptr_solver.QPAXPTRSolver that don't accept second-order-cone constraints. Backends that do accept them (e.g. :class:openscvx.solvers.cvxpy_ptr_solver.CVXPyPTRSolver) override this to invoke their backend-specific lowerer.

Parameters:

Name	Type	Description	Default
constraints	ConstraintSet	Categorized symbolic constraints. Only the nodal_convex / cross_node_convex lists matter here; non-convex constraints go through the JAX lowering pipeline.	required
parameters	Optional[Dict[str, Any]]	Optional dict of symbolic Parameter objects referenced by the constraints. May be None.	None

Returns:

Type	Description
List[Any]	(lowered_list, parameter_map). The first is a list of
Dict[str, Any]	backend-specific constraint objects (e.g. cp.Constraint);
Tuple[List[Any], Dict[str, Any]]	the second maps parameter names to backend-specific parameter
Tuple[List[Any], Dict[str, Any]]	objects. Both are empty for the default refusal path.

Raises:

Type	Description
NotImplementedError	if the user defined any .convex() constraints and this backend doesn't override.

Source code in openscvx/solvers/base.py

def lower_convex_constraints(
    self,
    constraints: "ConstraintSet",
    parameters: Optional[Dict[str, Any]] = None,
) -> Tuple[List[Any], Dict[str, Any]]:
    """Lower user ``.convex()`` constraints into this backend's form.

    Called once by :func:`openscvx.symbolic.lower.lower_symbolic_problem`
    after ``create_variables()`` and before ``initialize()``.

    The default implementation refuses any user ``.convex()``
    constraints — appropriate for backends like
    :class:`openscvx.solvers.qpax_ptr_solver.QPAXPTRSolver` that don't
    accept second-order-cone constraints. Backends that do accept them
    (e.g. :class:`openscvx.solvers.cvxpy_ptr_solver.CVXPyPTRSolver`)
    override this to invoke their backend-specific lowerer.

    Args:
        constraints: Categorized symbolic constraints. Only the
            ``nodal_convex`` / ``cross_node_convex`` lists matter here;
            non-convex constraints go through the JAX lowering pipeline.
        parameters: Optional dict of symbolic ``Parameter`` objects
            referenced by the constraints. May be ``None``.

    Returns:
        ``(lowered_list, parameter_map)``. The first is a list of
        backend-specific constraint objects (e.g. ``cp.Constraint``);
        the second maps parameter names to backend-specific parameter
        objects. Both are empty for the default refusal path.

    Raises:
        NotImplementedError: if the user defined any ``.convex()``
            constraints and this backend doesn't override.
    """
    n = len(constraints.nodal_convex) + len(constraints.cross_node_convex)
    if n:
        raise NotImplementedError(
            f"{type(self).__name__} does not support user-defined "
            f".convex() constraints ({n} defined). Drop the .convex() "
            "constraint or switch to a backend that supports them "
            "(e.g. openscvx.CVXPyPTRSolver)."
        )
    return [], {}

solve() -> Any abstractmethod

Solve the convex subproblem and return results.

Called at each SCP iteration after updating linearization and penalties. Returns a solver-specific result object containing the solution.

Returns:

Type	Description
Any	Solver-specific result object (e.g., PTRSolveResult for PTR).

Source code in openscvx/solvers/base.py

@abstractmethod
def solve(self) -> Any:
    """Solve the convex subproblem and return results.

    Called at each SCP iteration after updating linearization and penalties.
    Returns a solver-specific result object containing the solution.

    Returns:
        Solver-specific result object (e.g., ``PTRSolveResult`` for PTR).
    """
    raise NotImplementedError

update_boundary_conditions(**kwargs) -> None abstractmethod

Update boundary condition parameters.

Called once during algorithm initialization to set initial and terminal state constraints.

The specific parameters depend on the solver implementation. See concrete solver classes for expected arguments.

Source code in openscvx/solvers/base.py

@abstractmethod
def update_boundary_conditions(self, **kwargs) -> None:
    """Update boundary condition parameters.

    Called once during algorithm initialization to set initial and terminal
    state constraints.

    The specific parameters depend on the solver implementation.
    See concrete solver classes for expected arguments.
    """
    raise NotImplementedError

update_constraint_linearizations(**kwargs) -> None abstractmethod

Update linearized constraint values and gradients.

Called at each SCP iteration before solve() to set constraint function values and gradients at the current linearization point.

The specific parameters depend on the solver implementation. See concrete solver classes for expected arguments.

Source code in openscvx/solvers/base.py

@abstractmethod
def update_constraint_linearizations(self, **kwargs) -> None:
    """Update linearized constraint values and gradients.

    Called at each SCP iteration before ``solve()`` to set constraint
    function values and gradients at the current linearization point.

    The specific parameters depend on the solver implementation.
    See concrete solver classes for expected arguments.
    """
    raise NotImplementedError

update_dynamics_linearization(**kwargs) -> None abstractmethod

Update dynamics linearization point and matrices.

Called at each SCP iteration before solve() to set the current linearization point and discretized dynamics matrices.

The specific parameters depend on the solver implementation. See concrete solver classes for expected arguments.

Source code in openscvx/solvers/base.py

@abstractmethod
def update_dynamics_linearization(self, **kwargs) -> None:
    """Update dynamics linearization point and matrices.

    Called at each SCP iteration before ``solve()`` to set the current
    linearization point and discretized dynamics matrices.

    The specific parameters depend on the solver implementation.
    See concrete solver classes for expected arguments.
    """
    raise NotImplementedError

update_penalties(**kwargs) -> None abstractmethod

Update SCP penalty weights.

Called at each SCP iteration before solve() to set the current penalty weights for trust region, virtual control, and virtual buffer.

The specific parameters depend on the solver implementation. See concrete solver classes for expected arguments.

Source code in openscvx/solvers/base.py

@abstractmethod
def update_penalties(self, **kwargs) -> None:
    """Update SCP penalty weights.

    Called at each SCP iteration before ``solve()`` to set the current
    penalty weights for trust region, virtual control, and virtual buffer.

    The specific parameters depend on the solver implementation.
    See concrete solver classes for expected arguments.
    """
    raise NotImplementedError

update_proximal_terms() -> None

Update optional proximal-expansion parameters.

Called after update_penalties() and before solve(). Solvers that formulate trust-region penalties via expanded quadratic forms can override this hook to refresh any derived coefficients.

Default implementation is a no-op.

Source code in openscvx/solvers/base.py

def update_proximal_terms(self) -> None:
    """Update optional proximal-expansion parameters.

    Called after ``update_penalties()`` and before ``solve()``. Solvers
    that formulate trust-region penalties via expanded quadratic forms can
    override this hook to refresh any derived coefficients.

    Default implementation is a no-op.
    """
    return None

PTRSolveResult dataclass

Result from solving a PTR convex subproblem.

Contains the solution trajectories and slack variables from a single SCP iteration. All trajectories are unscaled (physical units).

Attributes:

Name	Type	Description
x	ndarray	State trajectory, shape (N, n_states). Unscaled.
u	ndarray	Control trajectory, shape (N, n_controls). Unscaled.
nu	ndarray	Virtual control slack for dynamics defects, shape (N-1, n_states).
nu_vb	List[ndarray]	Nonconvex nodal constraint violation slacks. List of arrays, one per nodal constraint.
nu_vb_cross	List[float]	Cross-node constraint violation slacks. List of scalars, one per cross-node constraint.
cost	float	Optimal objective value.
status	str	Solver status string (e.g., "optimal", "infeasible").

Source code in openscvx/solvers/ptr_solver.py

@dataclass
class PTRSolveResult:
    """Result from solving a PTR convex subproblem.

    Contains the solution trajectories and slack variables from a single
    SCP iteration. All trajectories are unscaled (physical units).

    Attributes:
        x: State trajectory, shape (N, n_states). Unscaled.
        u: Control trajectory, shape (N, n_controls). Unscaled.
        nu: Virtual control slack for dynamics defects, shape (N-1, n_states).
        nu_vb: Nonconvex nodal constraint violation slacks. List of arrays,
            one per nodal constraint.
        nu_vb_cross: Cross-node constraint violation slacks. List of scalars,
            one per cross-node constraint.
        cost: Optimal objective value.
        status: Solver status string (e.g., "optimal", "infeasible").
    """

    x: np.ndarray
    u: np.ndarray
    nu: np.ndarray
    nu_vb: List[np.ndarray]
    nu_vb_cross: List[float]
    cost: float
    status: str

PTRSolver

Bases: ConvexSolver

Abstract base class for Penalized Trust-Region convex subproblem solvers.

Defines the contract every PTR backend must satisfy: per-iteration entry points for updating the linearization point, constraint gradients, penalty weights, and boundary conditions, plus a solve() returning a :class:PTRSolveResult. The set of variables (state, control, virtual control nu, per-constraint virtual buffer nu_vb, cross-node slack nu_vb_cross) is fixed by the PTR formulation; how each backend realizes those variables is an implementation detail.

Subclasses must additionally implement :meth:create_variables and :meth:initialize from :class:ConvexSolver.

Attributes:

Name	Type	Description
SUPPORTED_CONE_TYPES	FrozenSet[Type[ConeConstraint]]	The set of :class:~openscvx.solvers.cones.ConeConstraint subclasses this backend can assemble. An empty frozenset (the default) means the backend rejects all user .convex() constraints. Concrete subclasses declare their capabilities by overriding this attribute.

Source code in openscvx/solvers/ptr_solver.py

class PTRSolver(ConvexSolver):
    """Abstract base class for Penalized Trust-Region convex subproblem solvers.

    Defines the contract every PTR backend must satisfy: per-iteration entry
    points for updating the linearization point, constraint gradients, penalty
    weights, and boundary conditions, plus a ``solve()`` returning a
    :class:`PTRSolveResult`. The set of variables (state, control, virtual
    control ``nu``, per-constraint virtual buffer ``nu_vb``, cross-node slack
    ``nu_vb_cross``) is fixed by the PTR formulation; how each backend
    realizes those variables is an implementation detail.

    Subclasses must additionally implement :meth:`create_variables` and
    :meth:`initialize` from :class:`ConvexSolver`.

    Attributes:
        SUPPORTED_CONE_TYPES: The set of
            :class:`~openscvx.solvers.cones.ConeConstraint` subclasses this
            backend can assemble.  An empty frozenset (the default) means the
            backend rejects **all** user ``.convex()`` constraints.  Concrete
            subclasses declare their capabilities by overriding this attribute.
    """

    # Subclasses declare which cone types they support.
    SUPPORTED_CONE_TYPES: ClassVar[FrozenSet[Type[ConeConstraint]]] = frozenset()

    # Registry populated via __init_subclass__ so _raise_unsupported_cone_error
    # can suggest which backends do support a given cone type.
    _solver_registry: ClassVar[List[Type["PTRSolver"]]] = []

    #: Backend-specific solve options (tolerances, iteration caps, ...). Each
    #: concrete backend assigns its own in ``__init__``; declared here with a
    #: ``None`` default so the base :meth:`_hash_into` contract is self-contained
    #: — a subclass that omits it hashes as an empty mapping rather than raising.
    solver_args: Optional[dict] = None

    def __init_subclass__(cls, **kwargs: Any) -> None:
        super().__init_subclass__(**kwargs)
        PTRSolver._solver_registry.append(cls)

    # ------------------------------------------------------------------
    # Principled convex-constraint lowering (shared by QPAX and Moreau)
    # ------------------------------------------------------------------

    def lower_convex_constraints(
        self,
        constraints: "ConstraintSet",
        parameters: Optional[Dict[str, Any]] = None,
    ) -> Tuple[List[Any], Dict[str, Any]]:
        """Lower user ``.convex()`` constraints into cone constraints.

        Separates auto-generated impulsive zero-pin equalities from genuine
        user constraints, canonicalises the latter via
        :func:`~openscvx.symbolic.canonicalize.canonicalize_nodal_constraint`,
        and verifies each cone type against :attr:`SUPPORTED_CONE_TYPES`.

        Cross-node convex constraints are not yet supported by the JAX
        backends and raise :class:`NotImplementedError`.

        Raises:
            NotImplementedError: For cross-node convex constraints or
                unsupported cone types (with a message listing which backends
                do support the cone type).
            ValueError: If a constraint is not affine in state/control.
        """
        from openscvx.symbolic.canonicalize import canonicalize_nodal_constraint

        if constraints.cross_node_convex:
            raise NotImplementedError(
                f"{type(self).__name__} does not support cross-node .convex() "
                f"constraints ({len(constraints.cross_node_convex)} defined). "
                "Use CVXPyPTRSolver."
            )

        pins, user_nodal = self._partition_nodal_convex(constraints)
        self._impulsive_pins = pins

        user_cones: List[ConeConstraint] = []
        for nc in user_nodal:
            cones = canonicalize_nodal_constraint(nc)
            for cone in cones:
                if type(cone) not in self.SUPPORTED_CONE_TYPES:
                    self._raise_unsupported_cone_error(cone)
                user_cones.append(cone)

        self._user_cone_constraints: List[ConeConstraint] = user_cones
        self._parameters_dict: Dict[str, Any] = parameters or {}
        return [], {}

    @staticmethod
    def _partition_nodal_convex(
        constraints: "ConstraintSet",
    ) -> Tuple[List[Tuple[List[int], slice]], List]:
        """Separate auto-generated impulsive zero-pin constraints from genuine
        user-defined convex constraints.

        Auto-generated pins match the pattern ``Control == 0`` injected by
        :func:`~openscvx.symbolic.lower._augment_impulsive_constraints`.  All
        other ``nodal_convex`` entries are returned as user-defined constraints.

        Returns:
            tuple: ``(pins, user_nodal)`` where *pins* is a list of
            ``(nodes, ctrl_slice)`` pairs and *user_nodal* is the list of
            :class:`~openscvx.symbolic.expr.constraint.NodalConstraint`
            objects that need full canonicalization.
        """
        from openscvx.symbolic.expr.constraint import Equality, NodalConstraint
        from openscvx.symbolic.expr.control import Control
        from openscvx.symbolic.expr.expr import Constant

        pins: List[Tuple[List[int], slice]] = []
        user_nodal: List = []

        for entry in constraints.nodal_convex:
            if (
                isinstance(entry, NodalConstraint)
                and isinstance(entry.constraint, Equality)
                and isinstance(entry.constraint.rhs, Constant)
                and np.all(np.asarray(entry.constraint.rhs.value) == 0)
                and isinstance(entry.constraint.lhs, Control)
                and entry.constraint.lhs._slice is not None
            ):
                pins.append(([int(k) for k in entry.nodes], entry.constraint.lhs._slice))
            else:
                user_nodal.append(entry)

        return pins, user_nodal

    def _raise_unsupported_cone_error(self, cone: ConeConstraint) -> None:
        """Raise :class:`NotImplementedError` naming which backends support *cone*.

        Scans :attr:`_solver_registry` to find alternatives.
        """
        cone_cls = type(cone)
        alternatives = [
            cls.__name__
            for cls in PTRSolver._solver_registry
            if cone_cls in cls.SUPPORTED_CONE_TYPES
        ]
        msg = f"{type(self).__name__} does not support {cone_cls.__name__} constraints."
        if alternatives:
            msg += f"  Backends that do: {', '.join(alternatives)}."
        else:
            msg += (
                "  No registered JAX backend supports this cone type; use CVXPyPTRSolver instead."
            )
        raise NotImplementedError(msg)

    @abstractmethod
    def update_dynamics_linearization(
        self,
        x_bar: np.ndarray,
        u_bar: np.ndarray,
        A_d: np.ndarray,
        B_d: np.ndarray,
        C_d: np.ndarray,
        x_prop: np.ndarray,
        x_prop_plus: np.ndarray | None = None,
        D_d: np.ndarray | None = None,
        E_d: np.ndarray | None = None,
    ) -> None:
        """Update dynamics linearization point and discrete-time matrices.

        Args:
            x_bar: Previous state trajectory, shape (N, n_states).
            u_bar: Previous control trajectory, shape (N, n_controls).
            A_d: Discretized state Jacobian, shape (N-1, n_states, n_states).
            B_d: Discretized control Jacobian (current node),
                shape (N-1, n_states, n_controls).
            C_d: Discretized control Jacobian (next node),
                shape (N-1, n_states, n_controls).
            x_prop: Propagated state from continuous dynamics,
                shape (N-1, n_states).
            x_prop_plus: Optional impulsive/discrete propagated state,
                shape (N, n_states).
            D_d: Optional impulsive/discrete Jacobian wrt state,
                shape (N, n_states, n_states).
            E_d: Optional impulsive/discrete Jacobian wrt control,
                shape (N, n_states, n_controls).
        """
        raise NotImplementedError

    @abstractmethod
    def update_constraint_linearizations(
        self,
        nodal: List[dict] = None,
        cross_node: List[dict] = None,
    ) -> None:
        """Update linearized constraint values and gradients.

        Args:
            nodal: List of dicts, one per nodal constraint, each containing:
                ``g`` (value at linearization point),
                ``grad_g_x`` (gradient w.r.t. state),
                ``grad_g_u`` (gradient w.r.t. control).
            cross_node: List of dicts, one per cross-node constraint, each
                containing: ``g``, ``grad_g_X`` (gradient w.r.t. full state
                trajectory), ``grad_g_U`` (gradient w.r.t. full control
                trajectory).
        """
        raise NotImplementedError

    @abstractmethod
    def update_penalties(
        self,
        lam_prox: np.ndarray,
        lam_cost: Union[float, np.ndarray],
        lam_vc: np.ndarray,
        lam_vb_nodal: np.ndarray,
        lam_vb_cross: np.ndarray,
    ) -> None:
        """Update SCP penalty weights.

        Args:
            lam_prox: Trust region weights, shape (N, n_states + n_controls).
            lam_cost: Cost function weight. Scalar or shape (n_states,).
            lam_vc: Virtual control penalty weights, shape (N-1, n_states).
            lam_vb_nodal: Virtual buffer penalty weights for nodal constraints,
                shape (N, n_nodal_constraints).
            lam_vb_cross: Virtual buffer penalty weights for cross-node
                constraints, shape (n_cross_node_constraints,).
        """
        raise NotImplementedError

    @abstractmethod
    def update_boundary_conditions(
        self,
        x_init: np.ndarray = None,
        x_term: np.ndarray = None,
    ) -> None:
        """Update initial and/or terminal state parameters.

        Args:
            x_init: Initial state vector, shape (n_states,). Optional.
            x_term: Terminal state vector, shape (n_states,). Optional.
        """
        raise NotImplementedError

    @abstractmethod
    def solve(self) -> PTRSolveResult:
        """Solve the convex subproblem and return a :class:`PTRSolveResult`.

        Call the four ``update_*`` methods first to set the linearization
        point, constraint gradients, penalties, and boundary conditions.
        """
        raise NotImplementedError

    @abstractmethod
    def iteration_callback(self) -> Callable[..., SubproblemSolution]:
        """Return a JAX-friendly ``(state, SubproblemData) -> SubproblemSolution``.

        Built once at :meth:`initialize`; called once per SCP iteration. The
        callable consumes the algorithm-state pytree plus an assembled
        :class:`SubproblemData`, performs one ``assemble → solve → unpack``
        cycle on the JAX boundary, and returns a :class:`SubproblemSolution`.

        Replaces the four ``update_*`` methods + :meth:`solve` for the
        JAX-pure SCP path. All backends share the same input/output pytree
        shape so the SCP loop downstream can wrap the callback in
        ``lax.while_loop`` and compose with ``jax.jit`` / ``jax.vmap``.
        """
        raise NotImplementedError

    @property
    def exportable(self) -> bool:
        """Whether this backend's solve survives ``jax.export`` serialization.

        :meth:`~openscvx.problem.Problem.solve_batched` exports the whole
        vmapped SCP loop to disk under ``save_compiled=True``, which requires
        every op in the loop — the backend solve included — to lower to
        serializable XLA. The pure-JAX backends (QPAX, Moreau) qualify, so the
        base default is ``True``. CVXPy's :meth:`iteration_callback` is a
        :func:`jax.pure_callback`, whose host call ``jax.export`` cannot
        serialize, so it overrides this to ``False`` and ``solve_batched``
        raises a teaching error rather than silently degrading to an in-process
        solve.
        """
        return True

    def _hash_into(self, hasher: "hashlib._Hash") -> None:
        """Contribute this backend's identity to the ``solve_batched`` cache key.

        The exported batched loop bakes in the backend class and its
        ``solver_args`` (tolerances, iteration caps, ...), neither of which the
        symbolic problem hash covers. Mirrors the symbolic ``_hash_into``
        protocol (see :meth:`~openscvx.symbolic.expr.expr.Expr._hash_into`) so a
        backend that adds a field affecting the artifact updates the hash right
        next to that field.
        """
        from openscvx.utils.caching import hash_value_into

        hasher.update(type(self).__name__.encode())
        for key in sorted(self.solver_args or {}):
            hasher.update(str(key).encode())
            hash_value_into(hasher, self.solver_args[key])

    @staticmethod
    def _extract_impulsive_pins(
        constraints: "ConstraintSet",
    ) -> Optional[List[Tuple[List[int], slice]]]:
        """Recognize auto-generated impulsive zero-pin constraints.

        :func:`openscvx.symbolic.lower._augment_impulsive_constraints` injects
        a ``Control == 0`` equality at every non-impulse node for each
        impulsive control. CVXPy lowers these alongside user ``.convex()``
        constraints, but JAX backends that otherwise refuse user
        ``.convex()`` constraints still need to honor them.

        This helper detects that exact shape — a ``NodalConstraint`` wrapping
        ``Equality(Control, Constant(0))`` over a list of nodes — and
        returns the implied ``(nodes, slice)`` pin list. If any
        ``nodal_convex`` entry doesn't match the auto-augmentation shape,
        returns ``None`` so the caller can fall back to the default
        refusal.

        Cross-node convex constraints are never produced by the
        auto-augmentation, so any presence aborts recognition.
        """
        from openscvx.symbolic.expr.constraint import Equality, NodalConstraint
        from openscvx.symbolic.expr.control import Control
        from openscvx.symbolic.expr.expr import Constant

        if constraints.cross_node_convex:
            return None

        pins: List[Tuple[List[int], slice]] = []
        for entry in constraints.nodal_convex:
            if not isinstance(entry, NodalConstraint):
                return None
            inner = entry.constraint
            if not isinstance(inner, Equality):
                return None
            rhs = inner.rhs
            if not isinstance(rhs, Constant) or not np.all(np.asarray(rhs.value) == 0):
                return None
            lhs = inner.lhs
            if not isinstance(lhs, Control) or lhs._slice is None:
                return None
            pins.append(([int(k) for k in entry.nodes], lhs._slice))
        return pins

    @staticmethod
    def _scaling(unified: Union["UnifiedState", "UnifiedControl"]) -> Tuple[np.ndarray, np.ndarray]:
        """Compute the affine scaling matrices ``(S, c)`` for a unified
        state or control interface.

        The PTR formulation works on scaled variables ``z = S⁻¹ (x - c)``;
        ``S`` and ``c`` are chosen so the scaled variables sit roughly on
        ``[-1, 1]``. Bounds come from ``scaling_min``/``scaling_max`` when
        provided, otherwise from ``min``/``max``. Shared by every backend.
        """
        from openscvx.config import get_affine_scaling_matrices

        n = len(unified.max)
        lower = np.array(
            unified.scaling_min if unified.scaling_min is not None else unified.min,
            dtype=float,
        )
        upper = np.array(
            unified.scaling_max if unified.scaling_max is not None else unified.max,
            dtype=float,
        )
        return get_affine_scaling_matrices(n, lower, upper)

exportable: bool property

Whether this backend's solve survives jax.export serialization.

:meth:~openscvx.problem.Problem.solve_batched exports the whole vmapped SCP loop to disk under save_compiled=True, which requires every op in the loop — the backend solve included — to lower to serializable XLA. The pure-JAX backends (QPAX, Moreau) qualify, so the base default is True. CVXPy's :meth:iteration_callback is a :func:jax.pure_callback, whose host call jax.export cannot serialize, so it overrides this to False and solve_batched raises a teaching error rather than silently degrading to an in-process solve.

iteration_callback() -> Callable[..., SubproblemSolution] abstractmethod

Return a JAX-friendly (state, SubproblemData) -> SubproblemSolution.

Built once at :meth:initialize; called once per SCP iteration. The callable consumes the algorithm-state pytree plus an assembled :class:SubproblemData, performs one assemble → solve → unpack cycle on the JAX boundary, and returns a :class:SubproblemSolution.

Replaces the four update_* methods + :meth:solve for the JAX-pure SCP path. All backends share the same input/output pytree shape so the SCP loop downstream can wrap the callback in lax.while_loop and compose with jax.jit / jax.vmap.

Source code in openscvx/solvers/ptr_solver.py

@abstractmethod
def iteration_callback(self) -> Callable[..., SubproblemSolution]:
    """Return a JAX-friendly ``(state, SubproblemData) -> SubproblemSolution``.

    Built once at :meth:`initialize`; called once per SCP iteration. The
    callable consumes the algorithm-state pytree plus an assembled
    :class:`SubproblemData`, performs one ``assemble → solve → unpack``
    cycle on the JAX boundary, and returns a :class:`SubproblemSolution`.

    Replaces the four ``update_*`` methods + :meth:`solve` for the
    JAX-pure SCP path. All backends share the same input/output pytree
    shape so the SCP loop downstream can wrap the callback in
    ``lax.while_loop`` and compose with ``jax.jit`` / ``jax.vmap``.
    """
    raise NotImplementedError

lower_convex_constraints(constraints: ConstraintSet, parameters: Optional[Dict[str, Any]] = None) -> Tuple[List[Any], Dict[str, Any]]

Lower user .convex() constraints into cone constraints.

Separates auto-generated impulsive zero-pin equalities from genuine user constraints, canonicalises the latter via :func:~openscvx.symbolic.canonicalize.canonicalize_nodal_constraint, and verifies each cone type against :attr:SUPPORTED_CONE_TYPES.

Cross-node convex constraints are not yet supported by the JAX backends and raise :class:NotImplementedError.

Raises:

Type	Description
NotImplementedError	For cross-node convex constraints or unsupported cone types (with a message listing which backends do support the cone type).
ValueError	If a constraint is not affine in state/control.

Source code in openscvx/solvers/ptr_solver.py

def lower_convex_constraints(
    self,
    constraints: "ConstraintSet",
    parameters: Optional[Dict[str, Any]] = None,
) -> Tuple[List[Any], Dict[str, Any]]:
    """Lower user ``.convex()`` constraints into cone constraints.

    Separates auto-generated impulsive zero-pin equalities from genuine
    user constraints, canonicalises the latter via
    :func:`~openscvx.symbolic.canonicalize.canonicalize_nodal_constraint`,
    and verifies each cone type against :attr:`SUPPORTED_CONE_TYPES`.

    Cross-node convex constraints are not yet supported by the JAX
    backends and raise :class:`NotImplementedError`.

    Raises:
        NotImplementedError: For cross-node convex constraints or
            unsupported cone types (with a message listing which backends
            do support the cone type).
        ValueError: If a constraint is not affine in state/control.
    """
    from openscvx.symbolic.canonicalize import canonicalize_nodal_constraint

    if constraints.cross_node_convex:
        raise NotImplementedError(
            f"{type(self).__name__} does not support cross-node .convex() "
            f"constraints ({len(constraints.cross_node_convex)} defined). "
            "Use CVXPyPTRSolver."
        )

    pins, user_nodal = self._partition_nodal_convex(constraints)
    self._impulsive_pins = pins

    user_cones: List[ConeConstraint] = []
    for nc in user_nodal:
        cones = canonicalize_nodal_constraint(nc)
        for cone in cones:
            if type(cone) not in self.SUPPORTED_CONE_TYPES:
                self._raise_unsupported_cone_error(cone)
            user_cones.append(cone)

    self._user_cone_constraints: List[ConeConstraint] = user_cones
    self._parameters_dict: Dict[str, Any] = parameters or {}
    return [], {}

solve() -> PTRSolveResult abstractmethod

Solve the convex subproblem and return a :class:PTRSolveResult.

Call the four update_* methods first to set the linearization point, constraint gradients, penalties, and boundary conditions.

Source code in openscvx/solvers/ptr_solver.py

@abstractmethod
def solve(self) -> PTRSolveResult:
    """Solve the convex subproblem and return a :class:`PTRSolveResult`.

    Call the four ``update_*`` methods first to set the linearization
    point, constraint gradients, penalties, and boundary conditions.
    """
    raise NotImplementedError

update_boundary_conditions(x_init: np.ndarray = None, x_term: np.ndarray = None) -> None abstractmethod

Update initial and/or terminal state parameters.

Parameters:

Name	Type	Description	Default
x_init	ndarray	Initial state vector, shape (n_states,). Optional.	None
x_term	ndarray	Terminal state vector, shape (n_states,). Optional.	None

Source code in openscvx/solvers/ptr_solver.py

@abstractmethod
def update_boundary_conditions(
    self,
    x_init: np.ndarray = None,
    x_term: np.ndarray = None,
) -> None:
    """Update initial and/or terminal state parameters.

    Args:
        x_init: Initial state vector, shape (n_states,). Optional.
        x_term: Terminal state vector, shape (n_states,). Optional.
    """
    raise NotImplementedError

update_constraint_linearizations(nodal: List[dict] = None, cross_node: List[dict] = None) -> None abstractmethod

Update linearized constraint values and gradients.

Parameters:

Name	Type	Description	Default
nodal	List[dict]	List of dicts, one per nodal constraint, each containing: g (value at linearization point), grad_g_x (gradient w.r.t. state), grad_g_u (gradient w.r.t. control).	None
cross_node	List[dict]	List of dicts, one per cross-node constraint, each containing: g, grad_g_X (gradient w.r.t. full state trajectory), grad_g_U (gradient w.r.t. full control trajectory).	None

Source code in openscvx/solvers/ptr_solver.py

@abstractmethod
def update_constraint_linearizations(
    self,
    nodal: List[dict] = None,
    cross_node: List[dict] = None,
) -> None:
    """Update linearized constraint values and gradients.

    Args:
        nodal: List of dicts, one per nodal constraint, each containing:
            ``g`` (value at linearization point),
            ``grad_g_x`` (gradient w.r.t. state),
            ``grad_g_u`` (gradient w.r.t. control).
        cross_node: List of dicts, one per cross-node constraint, each
            containing: ``g``, ``grad_g_X`` (gradient w.r.t. full state
            trajectory), ``grad_g_U`` (gradient w.r.t. full control
            trajectory).
    """
    raise NotImplementedError

update_dynamics_linearization(x_bar: np.ndarray, u_bar: np.ndarray, A_d: np.ndarray, B_d: np.ndarray, C_d: np.ndarray, x_prop: np.ndarray, x_prop_plus: np.ndarray | None = None, D_d: np.ndarray | None = None, E_d: np.ndarray | None = None) -> None abstractmethod

Update dynamics linearization point and discrete-time matrices.

Parameters:

Name	Type	Description	Default
x_bar	ndarray	Previous state trajectory, shape (N, n_states).	required
u_bar	ndarray	Previous control trajectory, shape (N, n_controls).	required
A_d	ndarray	Discretized state Jacobian, shape (N-1, n_states, n_states).	required
B_d	ndarray	Discretized control Jacobian (current node), shape (N-1, n_states, n_controls).	required
C_d	ndarray	Discretized control Jacobian (next node), shape (N-1, n_states, n_controls).	required
x_prop	ndarray	Propagated state from continuous dynamics, shape (N-1, n_states).	required
x_prop_plus	ndarray | None	Optional impulsive/discrete propagated state, shape (N, n_states).	None
D_d	ndarray | None	Optional impulsive/discrete Jacobian wrt state, shape (N, n_states, n_states).	None
E_d	ndarray | None	Optional impulsive/discrete Jacobian wrt control, shape (N, n_states, n_controls).	None

Source code in openscvx/solvers/ptr_solver.py

@abstractmethod
def update_dynamics_linearization(
    self,
    x_bar: np.ndarray,
    u_bar: np.ndarray,
    A_d: np.ndarray,
    B_d: np.ndarray,
    C_d: np.ndarray,
    x_prop: np.ndarray,
    x_prop_plus: np.ndarray | None = None,
    D_d: np.ndarray | None = None,
    E_d: np.ndarray | None = None,
) -> None:
    """Update dynamics linearization point and discrete-time matrices.

    Args:
        x_bar: Previous state trajectory, shape (N, n_states).
        u_bar: Previous control trajectory, shape (N, n_controls).
        A_d: Discretized state Jacobian, shape (N-1, n_states, n_states).
        B_d: Discretized control Jacobian (current node),
            shape (N-1, n_states, n_controls).
        C_d: Discretized control Jacobian (next node),
            shape (N-1, n_states, n_controls).
        x_prop: Propagated state from continuous dynamics,
            shape (N-1, n_states).
        x_prop_plus: Optional impulsive/discrete propagated state,
            shape (N, n_states).
        D_d: Optional impulsive/discrete Jacobian wrt state,
            shape (N, n_states, n_states).
        E_d: Optional impulsive/discrete Jacobian wrt control,
            shape (N, n_states, n_controls).
    """
    raise NotImplementedError

update_penalties(lam_prox: np.ndarray, lam_cost: Union[float, np.ndarray], lam_vc: np.ndarray, lam_vb_nodal: np.ndarray, lam_vb_cross: np.ndarray) -> None abstractmethod

Update SCP penalty weights.

Parameters:

Name	Type	Description	Default
lam_prox	ndarray	Trust region weights, shape (N, n_states + n_controls).	required
lam_cost	Union[float, ndarray]	Cost function weight. Scalar or shape (n_states,).	required
lam_vc	ndarray	Virtual control penalty weights, shape (N-1, n_states).	required
lam_vb_nodal	ndarray	Virtual buffer penalty weights for nodal constraints, shape (N, n_nodal_constraints).	required
lam_vb_cross	ndarray	Virtual buffer penalty weights for cross-node constraints, shape (n_cross_node_constraints,).	required

Source code in openscvx/solvers/ptr_solver.py

@abstractmethod
def update_penalties(
    self,
    lam_prox: np.ndarray,
    lam_cost: Union[float, np.ndarray],
    lam_vc: np.ndarray,
    lam_vb_nodal: np.ndarray,
    lam_vb_cross: np.ndarray,
) -> None:
    """Update SCP penalty weights.

    Args:
        lam_prox: Trust region weights, shape (N, n_states + n_controls).
        lam_cost: Cost function weight. Scalar or shape (n_states,).
        lam_vc: Virtual control penalty weights, shape (N-1, n_states).
        lam_vb_nodal: Virtual buffer penalty weights for nodal constraints,
            shape (N, n_nodal_constraints).
        lam_vb_cross: Virtual buffer penalty weights for cross-node
            constraints, shape (n_cross_node_constraints,).
    """
    raise NotImplementedError

PTRSolverSpec

Bases: BaseModel

Validates PTR solver configuration from dict/YAML input.

The backend discriminator selects which concrete PTR backend to build: "cvxpy" (the default, :class:openscvx.solvers.cvxpy_ptr_solver.CVXPyPTRSolver), "qpax" (:class:openscvx.solvers.qpax_ptr_solver.QPAXPTRSolver), or "moreau" (:class:openscvx.solvers.moreau_ptr_solver.MoreauPTRSolver).

Backend-specific fields mirror the constructor parameters of each class. Fields that don't belong to the active backend are a configuration error and raise ValidationError.

CVXPy-only

cvx_solver, cvxpygen, cvxpygen_override.

QPAX-only

solver_tol.

Moreau-only

verbose, device, tol_gap_abs, tol_feas.

Shared between QPAX and Moreau

max_iter.

All backends

solver_args (escape hatch for settings not covered by named fields).

Warning

Enabling cvxpygen currently disables sparse parameter declarations. cvxpygen does not yet support the N-D sparsity indices used by OpenSCvx's tiled parameters, so all parameters are created as dense when code generation is active.

Source code in openscvx/solvers/base.py

class PTRSolverSpec(BaseModel):
    """Validates PTR solver configuration from dict/YAML input.

    The ``backend`` discriminator selects which concrete PTR backend to build:
    ``"cvxpy"`` (the default,
    :class:`openscvx.solvers.cvxpy_ptr_solver.CVXPyPTRSolver`),
    ``"qpax"`` (:class:`openscvx.solvers.qpax_ptr_solver.QPAXPTRSolver`),
    or ``"moreau"``
    (:class:`openscvx.solvers.moreau_ptr_solver.MoreauPTRSolver`).

    Backend-specific fields mirror the constructor parameters of each class.
    Fields that don't belong to the active backend are a configuration error
    and raise ``ValidationError``.

    CVXPy-only:
        ``cvx_solver``, ``cvxpygen``, ``cvxpygen_override``.

    QPAX-only:
        ``solver_tol``.

    Moreau-only:
        ``verbose``, ``device``, ``tol_gap_abs``, ``tol_feas``.

    Shared between QPAX and Moreau:
        ``max_iter``.

    All backends:
        ``solver_args`` (escape hatch for settings not covered by named fields).

    !!! warning
        Enabling ``cvxpygen`` currently disables sparse parameter declarations.
        cvxpygen does not yet support the N-D sparsity indices used by
        OpenSCvx's tiled parameters, so all parameters are created as dense
        when code generation is active.
    """

    type: Literal["PTRSolver"] = "PTRSolver"
    backend: Literal["cvxpy", "qpax", "moreau"] = "cvxpy"

    # CVXPy-only
    cvx_solver: Optional[str] = None
    cvxpygen: bool = False
    cvxpygen_override: bool = False

    # QPAX-specific
    solver_tol: Optional[float] = None

    # Shared between QPAX and Moreau
    max_iter: Optional[int] = None

    # Moreau-specific
    verbose: Optional[bool] = None
    device: Optional[str] = None
    tol_gap_abs: Optional[float] = None
    tol_feas: Optional[float] = None

    # Escape hatch for all backends
    solver_args: Optional[Dict[str, Any]] = None

    model_config = ConfigDict(extra="forbid")

    @model_validator(mode="after")
    def _check_backend_fields(self):
        if self.backend in ("qpax", "moreau"):
            offenders = [
                name
                for name, value in (
                    ("cvx_solver", self.cvx_solver),
                    ("cvxpygen", self.cvxpygen),
                    ("cvxpygen_override", self.cvxpygen_override),
                )
                if value
            ]
            if offenders:
                raise ValueError(
                    f"{offenders} only valid for backend='cvxpy'; "
                    "remove these fields or set backend='cvxpy'."
                )
        if self.backend == "cvxpy":
            offenders = [
                name
                for name, value in (
                    ("solver_tol", self.solver_tol),
                    ("max_iter", self.max_iter),
                    ("verbose", self.verbose),
                    ("device", self.device),
                    ("tol_gap_abs", self.tol_gap_abs),
                    ("tol_feas", self.tol_feas),
                )
                if value is not None
            ]
            if offenders:
                raise ValueError(
                    f"{offenders} not valid for backend='cvxpy'; pass these "
                    "via solver_args or switch to the qpax / moreau backend."
                )
        if self.backend == "qpax":
            offenders = [
                name
                for name, value in (
                    ("verbose", self.verbose),
                    ("device", self.device),
                    ("tol_gap_abs", self.tol_gap_abs),
                    ("tol_feas", self.tol_feas),
                )
                if value is not None
            ]
            if offenders:
                raise ValueError(
                    f"{offenders} only valid for backend='moreau'; "
                    "remove these fields or set backend='moreau'."
                )
        if self.backend == "moreau" and self.solver_tol is not None:
            raise ValueError(
                "solver_tol is only valid for backend='qpax'; "
                "remove this field or set backend='qpax'."
            )
        return self

    def build(self) -> ConvexSolver:
        # Local imports keep CVXPy / qpax / moreau out of the import path until
        # the corresponding backend is actually requested.
        if self.backend == "cvxpy":
            from .cvxpy_ptr_solver import CVXPyPTRSolver

            return CVXPyPTRSolver(
                cvx_solver=self.cvx_solver or "QOCO",
                solver_args=self.solver_args,
                cvxpygen=self.cvxpygen,
                cvxpygen_override=self.cvxpygen_override,
            )
        if self.backend == "moreau":
            from .moreau_ptr_solver import MoreauPTRSolver

            kwargs: Dict[str, Any] = {}
            if self.max_iter is not None:
                kwargs["max_iter"] = self.max_iter
            if self.verbose is not None:
                kwargs["verbose"] = self.verbose
            if self.device is not None:
                kwargs["device"] = self.device
            if self.tol_gap_abs is not None:
                kwargs["tol_gap_abs"] = self.tol_gap_abs
            if self.tol_feas is not None:
                kwargs["tol_feas"] = self.tol_feas
            if self.solver_args is not None:
                kwargs["solver_args"] = self.solver_args
            return MoreauPTRSolver(**kwargs)

        from .qpax_ptr_solver import QPAXPTRSolver

        kwargs = {}
        if self.solver_tol is not None:
            kwargs["solver_tol"] = self.solver_tol
        if self.max_iter is not None:
            kwargs["max_iter"] = self.max_iter
        if self.solver_args is not None:
            kwargs["solver_args"] = self.solver_args
        return QPAXPTRSolver(**kwargs)

resolve_solver_config(val: Any) -> PTRSolverSpec

Validate a dict / Spec into a :class:PTRSolverSpec instance.

Source code in openscvx/solvers/__init__.py

def resolve_solver_config(val: Any) -> PTRSolverSpec:
    """Validate a dict / Spec into a :class:`PTRSolverSpec` instance."""
    if isinstance(val, PTRSolverSpec):
        return val
    return PTRSolverSpec.model_validate(val)