openscvx 2.dev6__tar.gz → 2.dev8__tar.gz

{openscvx-2.dev6/openscvx.egg-info → openscvx-2.dev8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 2.dev6
+Version: 2.dev8
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License
@@ -26,6 +26,8 @@ Requires-Dist: pydantic>=2.0
 Provides-Extra: cvxpygen
 Requires-Dist: cvxpygen; extra == "cvxpygen"
 Requires-Dist: qocogen; extra == "cvxpygen"
+Provides-Extra: qpax
+Requires-Dist: qpax>=0.1.1; extra == "qpax"
 Provides-Extra: stl
 Requires-Dist: stljax; extra == "stl"
 Provides-Extra: lie

{openscvx-2.dev6 → openscvx-2.dev8}/docs/UnderTheHood/lowering_architecture.md

@@ -160,9 +160,40 @@ JAX lowering has no dependency on:
 
 This means JAX-lowered dynamics and constraints could be used with alternative solvers.
 
+## Convex Subproblem Backends
+
+The convex subproblem at each SCP iteration is solved by a concrete subclass
+of the abstract `PTRSolver` base. Two backends ship today:
+
+| Backend | Class | Selector | Notes |
+|---------|-------|----------|-------|
+| CVXPy (default) | `CVXPyPTRSolver` | `solver={"backend": "cvxpy"}` (default) | DCP graph via CVXPy, dispatched to QOCO / CLARABEL / etc. Supports user `.convex()` constraints, cross-node constraints, CTCS, and impulsive controls. Optional cvxpygen code generation. |
+| QPAX | `QPAXPTRSolver` | `solver={"backend": "qpax"}` | JAX-native QP via `qpax.solve_qp`. Flat `(Q, q, A, b, G, h)` assembly. Supports box / dynamics / CTCS / boundary-Fix; **rejects** user `.convex()`, cross-node, and impulsive at `initialize()` with a clear "use `CVXPyPTRSolver`" message. Enables a path toward an end-to-end JAX-differentiable SCP loop in follow-up work. |
+
+Picking a backend at construction time:
+
+```python
+import openscvx as ox
+
+# Default — same as today's CVXPyPTRSolver.
+problem = ox.Problem(...)
+
+# Explicit CVXPy with a different inner solver.
+problem = ox.Problem(..., solver={"backend": "cvxpy", "cvx_solver": "CLARABEL"})
+
+# JAX-native QPAX (requires the `qpax` extra: pip install openscvx[qpax]).
+problem = ox.Problem(..., solver={"backend": "qpax"})
+```
+
+QPAX consumes the global JAX dtype — pass `float_dtype="float64"` to `Problem`
+if you need tight inner-solver tolerances; the default `float32` is enough for
+many problems but caps the QP's conditioning.
+
 ## Further Reading
 
 - `openscvx/symbolic/lower.py` — Main lowering implementation
 - `openscvx/lowered/` — Dataclass definitions
 - `openscvx/symbolic/lowerers/jax.py` — JAX visitor implementation
 - `openscvx/symbolic/lowerers/cvxpy.py` — CVXPy visitor implementation
+- `openscvx/solvers/ptr_solver.py` — Abstract `PTRSolver` contract
+- `openscvx/solvers/cvxpy_ptr_solver.py` / `qpax_ptr_solver.py` — Concrete backends

{openscvx-2.dev6 → openscvx-2.dev8}/examples/rocket/6DoF_pdg.py

@@ -218,11 +218,6 @@ problem = Problem(
     discretizer={
         "diffrax_kwargs": {"stepsize_controller": dfx.StepTo(np.linspace(0.0, 1 / (n - 1), 3))}
     },
-    solver={
-        "cvxpygen": True,
-        "cvx_solver": "qocogen",
-        "solver_args": {},
-    },
 )
 
 problem.settings.dev.printing = False

{openscvx-2.dev6 → openscvx-2.dev8}/openscvx/__init__.py

@@ -29,7 +29,10 @@ from openscvx.expert import ByofSpec
 from openscvx.integrations import DynamicsAdapter, MjxDynamics
 from openscvx.loader import load_dict, load_json, load_yaml
 from openscvx.problem import Problem
-from openscvx.solvers import PTRSolver
+from openscvx.solvers import CVXPyPTRSolver, PTRSolver
+
+# QPAXPTRSolver is exposed lazily via __getattr__ below to keep `import qpax`
+# off the hot import path for users who don't install the optional extra.
 from openscvx.symbolic.expr import (
     CTCS,
     Abs,
@@ -91,6 +94,16 @@ from openscvx.utils.cache import clear_cache, get_cache_dir, get_cache_size
 
 load_results = OptimizationResults.load
 
+
+def __getattr__(name: str):
+    """Lazy export for backends that depend on optional packages."""
+    if name == "QPAXPTRSolver":
+        from openscvx.solvers.qpax_ptr_solver import QPAXPTRSolver
+
+        return QPAXPTRSolver
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
 __all__ = [
     # Main Trajectory Optimization Entrypoint
     "Problem",
@@ -188,6 +201,8 @@ __all__ = [
     "VectorizeDiscretizeLinearize",
     # Convex Solver
     "PTRSolver",
+    "CVXPyPTRSolver",
+    "QPAXPTRSolver",
     # Algorithm & Autotuning
     "PenalizedTrustRegion",
     "AugmentedLagrangian",

{openscvx-2.dev6 → openscvx-2.dev8}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '2.dev6'
-__version_tuple__ = version_tuple = (2, 'dev6')
+__version__ = version = '2.dev8'
+__version_tuple__ = version_tuple = (2, 'dev8')
 
-__commit_id__ = commit_id = 'gcabefce07'
+__commit_id__ = commit_id = 'g69d91aabb'

{openscvx-2.dev6 → openscvx-2.dev8}/openscvx/loader.py

@@ -40,7 +40,7 @@ from pydantic import BaseModel, ConfigDict
 from openscvx.algorithms import PenalizedTrustRegionConfig
 from openscvx.config import SettingsSpec
 from openscvx.discretization import DiscretizerSpec
-from openscvx.solvers import SolverSpec
+from openscvx.solvers import PTRSolverSpec
 from openscvx.symbolic.expr.control import ControlSpec
 from openscvx.symbolic.expr.expr import Expr
 from openscvx.symbolic.expr.parameter import ParameterSpec
@@ -66,7 +66,7 @@ class ProblemSpec(BaseModel):
     constraints: List[str] = []
     algorithm: Optional[PenalizedTrustRegionConfig] = None
     discretizer: Optional[DiscretizerSpec] = None
-    solver: Optional[SolverSpec] = None
+    solver: Optional[PTRSolverSpec] = None
     settings: Optional[SettingsSpec] = None
     states_prop: Optional[List[StateSpec]] = None
     dynamics_prop: Optional[Dict[str, Any]] = None

{openscvx-2.dev6 → openscvx-2.dev8}/openscvx/lowered/cvxpy_constraints.py

@@ -17,7 +17,10 @@ class LoweredCvxpyConstraints:
 
     Attributes:
         constraints: List of CVXPy constraint objects (cp.Constraint).
-            Includes both nodal and cross-node convex constraints.
+            Includes both nodal and cross-node convex constraints. Empty
+            for backends that don't accept ``.convex()`` constraints — the
+            refusal happens earlier, in
+            :meth:`openscvx.solvers.base.ConvexSolver.lower_convex_constraints`.
     """
 
     constraints: list["cp.Constraint"] = field(default_factory=list)

{openscvx-2.dev6 → openscvx-2.dev8}/openscvx/problem.py

@@ -201,9 +201,11 @@ class Problem:
                     discretizer=ox.LinearizeDiscretize(dis_type="ZOH", ode_solver="Dopri8")
             solver: Convex subproblem solver configuration. Accepts:
 
-                - ``None`` — uses ``PTRSolver()`` with defaults (QOCO backend).
+                - ``None`` — uses ``CVXPyPTRSolver()`` with defaults (QOCO backend).
                 - A ``ConvexSolver`` instance — used directly.
-                - A ``dict`` — passed as kwargs to ``PTRSolver()``.
+                - A ``dict`` — validated as ``PTRSolverSpec``; the ``backend``
+                  field (``"cvxpy"`` or ``"qpax"``) selects the concrete
+                  backend.
 
                 Examples::
 
@@ -216,8 +218,12 @@ class Problem:
                     # Enable cvxpygen code generation
                     solver={"cvxpygen": True}
 
+                    # JAX-native QPAX backend (no cvx_solver / cvxpygen fields)
+                    solver={"backend": "qpax"}
+
                     # Instance
-                    solver=ox.PTRSolver(cvx_solver="CLARABEL")
+                    solver=ox.CVXPyPTRSolver(cvx_solver="CLARABEL")
+                    solver=ox.QPAXPTRSolver()
             byof (ByofSpec, optional): Expert mode only. Raw JAX functions to
                 bypass symbolic layer. See :class:`openscvx.expert.ByofSpec` for
                 detailed documentation.
@@ -393,9 +399,9 @@ class Problem:
     def solver(self) -> ConvexSolver:
         """Access the convex subproblem solver instance.
 
-        Attributes such as ``cvx_solver``, ``solver_args``, ``cvxpygen``, and
-        ``cvxpygen_override`` can be modified freely before ``initialize``
-        is called::
+        Backend-specific attributes (e.g. ``cvx_solver``, ``solver_args``,
+        ``cvxpygen``, ``cvxpygen_override`` on :class:`CVXPyPTRSolver`) can
+        be modified freely before ``initialize`` is called::
 
             problem.solver.solver_args = {"abstol": 1e-6, "reltol": 1e-9}
             problem.solver.cvxpygen = True
@@ -407,7 +413,7 @@ class Problem:
             will have no effect on subsequent solves.
 
         Returns:
-            The solver instance (e.g., PTRSolver).
+            The solver instance — a concrete :class:`PTRSolver` subclass.
         """
         return self._solver
 
@@ -510,12 +516,12 @@ class Problem:
                 self._lowered.x_prop_unified.final[state._slice] = state.final
                 self._lowered.x_prop_unified.final_type[state._slice] = state.final_type
 
-        # Update CVXPy solver parameters (only if solver is initialized)
-        if self._solver._problem is not None:
-            self._solver.update_boundary_conditions(
-                x_init=self._lowered.x_unified.initial,
-                x_term=self._lowered.x_unified.final,
-            )
+        # Push to the solver — both backends short-circuit on a pre-initialize
+        # call, so this is safe to invoke from any lifecycle point.
+        self._solver.update_boundary_conditions(
+            x_init=self._lowered.x_unified.initial,
+            x_term=self._lowered.x_unified.final,
+        )
 
     def _sync_guesses(self):
         """Sync trajectory guesses from State/Control objects to lowered representation.

openscvx-2.dev8/openscvx/solvers/__init__.py

@@ -0,0 +1,93 @@
+"""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:
+
+```python
+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 two 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).
+
+Both 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.
+"""
+
+import warnings
+from typing import Any
+
+from .base import ConvexSolver, PTRSolverSpec
+from .cvxpy_ptr_solver import CVXPyPTRSolver
+from .ptr_solver import PTRSolver, PTRSolveResult
+
+
+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)
+
+
+def __getattr__(name: str):
+    """Deprecated alias: ``SolverSpec`` → :class:`PTRSolverSpec`."""
+    if name == "SolverSpec":
+        warnings.warn(
+            "openscvx.solvers.SolverSpec is deprecated; use PTRSolverSpec.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return PTRSolverSpec
+    if name == "QPAXPTRSolver":
+        # Lazy import so users without the qpax extra don't pay a hard
+        # ImportError just for `from openscvx.solvers import QPAXPTRSolver`
+        # — the import error gets deferred to instantiation time, where the
+        # error message points at the install command.
+        from .qpax_ptr_solver import QPAXPTRSolver
+
+        return QPAXPTRSolver
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+__all__ = [
+    # Base classes
+    "ConvexSolver",
+    "PTRSolver",
+    "PTRSolveResult",
+    # PTR backends
+    "CVXPyPTRSolver",
+    "QPAXPTRSolver",
+    # Config
+    "PTRSolverSpec",
+    "resolve_solver_config",
+]

{openscvx-2.dev6 → openscvx-2.dev8}/openscvx/solvers/base.py

@@ -5,40 +5,29 @@ must follow for use within successive convexification algorithms.
 
 !!! note
 
-    Solvers own their optimization variables via ``create_variables()``.
-    Convex constraint lowering remains in ``lower.py`` but uses the solver's
-    variables.
-
-    When adding non-CVXPy backends, there are two approaches:
-
-    1. **Solver owns the lowerer**: The solver implements a
-       ``lower_convex_constraints()`` method containing the lowering logic.
-
-    2. **Solver determines the lowerer**: The solver references which lowerer
-       to use, but the lowering logic stays in ``lower.py``. Example:
-
-       ```python
-       # In solver
-       @property
-       def lowerer(self):
-           from openscvx.symbolic.lower import lower_cvxpy_constraints
-           return lower_cvxpy_constraints
-
-       # In lower_symbolic_problem()
-       lowered_constraints = solver.lowerer(constraints, solver.variables, parameters)
-       ```
+    Solvers own both their optimization variables (``create_variables()``) and
+    the lowering of any user ``.convex()`` constraints
+    (``lower_convex_constraints()``). The default ``lower_convex_constraints``
+    refuses user ``.convex()`` constraints with a clear error — backends that
+    accept them override it. This keeps ``openscvx.symbolic.lower``
+    backend-agnostic: it never branches on solver type, it just delegates.
+
+    See :class:`openscvx.solvers.ptr_solver.PTRSolver` for the PTR-specific
+    interface every PTR backend implements.
 """
 
+import warnings
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple
 
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, model_validator
 
 if TYPE_CHECKING:
     from openscvx.config import Config
     from openscvx.lowered import LoweredProblem
     from openscvx.lowered.jax_constraints import LoweredJaxConstraints
     from openscvx.lowered.unified import UnifiedControl, UnifiedState
+    from openscvx.symbolic.constraint_set import ConstraintSet
 
 
 class ConvexSolver(ABC):
@@ -91,10 +80,6 @@ class ConvexSolver(ABC):
                     return MyResult(...)
     """
 
-    #: Backend solver name (e.g., ``"QOCO"``, ``"CLARABEL"``).  Subclasses
-    #: must set this in ``__init__``.
-    cvx_solver: str
-
     @abstractmethod
     def create_variables(
         self,
@@ -128,6 +113,50 @@ class ConvexSolver(ABC):
         """
         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,
@@ -241,11 +270,17 @@ class ConvexSolver(ABC):
 # Pydantic spec for dict / YAML validation
 # =============================================================================
 
-_SOLVER_MAP: Dict[str, type] = {}  # populated by __init__.py after all classes are imported
 
+class PTRSolverSpec(BaseModel):
+    """Validates PTR solver configuration from dict/YAML input.
 
-class SolverSpec(BaseModel):
-    """Validates 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`) or ``"qpax"``
+    (:class:`openscvx.solvers.qpax_ptr_solver.QPAXPTRSolver`).
+
+    ``cvx_solver``, ``cvxpygen``, and ``cvxpygen_override`` are CVXPy-only;
+    setting them under ``backend="qpax"`` is a configuration error.
 
     !!! warning
         Enabling ``cvxpygen`` currently disables sparse parameter declarations.
@@ -255,15 +290,61 @@ class SolverSpec(BaseModel):
     """
 
     type: Literal["PTRSolver"] = "PTRSolver"
-    cvx_solver: str = "QOCO"
+    backend: Literal["cvxpy", "qpax"] = "cvxpy"
+    cvx_solver: Optional[str] = None
     solver_args: Optional[Dict[str, Any]] = None
     cvxpygen: bool = False
     cvxpygen_override: bool = False
 
     model_config = ConfigDict(extra="forbid")
 
+    @model_validator(mode="after")
+    def _check_backend_fields(self):
+        if self.backend == "qpax":
+            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'."
+                )
+        return self
+
     def build(self) -> ConvexSolver:
-        cls = _SOLVER_MAP.get(self.type)
-        if cls is None:
-            raise ValueError(f"Unknown solver {self.type!r}; expected one of {sorted(_SOLVER_MAP)}")
-        return cls(**self.model_dump(exclude={"type"}, exclude_unset=True))
+        # Local imports keep CVXPy / qpax 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,
+            )
+        from .qpax_ptr_solver import QPAXPTRSolver
+
+        return QPAXPTRSolver(solver_args=self.solver_args)
+
+
+def __getattr__(name: str):
+    """Deprecated alias: ``SolverSpec`` → :class:`PTRSolverSpec`.
+
+    Kept for one release so existing dict/YAML configs and tests that import
+    ``SolverSpec`` continue to work. Emit a ``DeprecationWarning`` on access.
+    """
+    if name == "SolverSpec":
+        warnings.warn(
+            "openscvx.solvers.base.SolverSpec is deprecated; use PTRSolverSpec.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return PTRSolverSpec
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")