openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/lowered/__init__.py

@@ -0,0 +1,30 @@
+"""Lowered problem dataclasses.
+
+This module contains dataclasses representing the outputs of the lowering phase,
+where symbolic expressions are converted to executable JAX and CVXPy code.
+"""
+
+from openscvx.lowered.cvxpy_constraints import LoweredCvxpyConstraints
+from openscvx.lowered.cvxpy_variables import CVXPyVariables
+from openscvx.lowered.dynamics import Dynamics
+from openscvx.lowered.jax_constraints import (
+    LoweredCrossNodeConstraint,
+    LoweredJaxConstraints,
+    LoweredNodalConstraint,
+)
+from openscvx.lowered.parameters import ParameterDict
+from openscvx.lowered.problem import LoweredProblem
+from openscvx.lowered.unified import UnifiedControl, UnifiedState
+
+__all__ = [
+    "LoweredProblem",
+    "LoweredJaxConstraints",
+    "LoweredCvxpyConstraints",
+    "LoweredNodalConstraint",
+    "LoweredCrossNodeConstraint",
+    "CVXPyVariables",
+    "ParameterDict",
+    "Dynamics",
+    "UnifiedState",
+    "UnifiedControl",
+]

openscvx/lowered/cvxpy_constraints.py

@@ -0,0 +1,23 @@
+"""CVXPy-lowered constraint dataclass."""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import cvxpy as cp
+
+
+@dataclass
+class LoweredCvxpyConstraints:
+    """CVXPy-lowered convex constraints.
+
+    Contains constraints that have been lowered to CVXPy constraint objects.
+    These are added directly to the optimal control problem without
+    linearization.
+
+    Attributes:
+        constraints: List of CVXPy constraint objects (cp.Constraint).
+            Includes both nodal and cross-node convex constraints.
+    """
+
+    constraints: list["cp.Constraint"] = field(default_factory=list)

openscvx/lowered/cvxpy_variables.py

@@ -0,0 +1,124 @@
+"""CVXPy variables and parameters dataclass for the optimal control problem."""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List
+
+import numpy as np
+
+if TYPE_CHECKING:
+    import cvxpy as cp
+
+
+@dataclass
+class CVXPyVariables:
+    """CVXPy variables and parameters for the optimal control problem.
+
+    This dataclass holds all CVXPy Variable and Parameter objects needed to
+    construct and solve the optimal control problem. It replaces the previous
+    untyped dictionary approach with a typed, self-documenting structure.
+
+    The variables are organized into logical groups:
+        - SCP weights: Parameters controlling trust region and penalty weights
+        - State: Variables and parameters for the state trajectory
+        - Control: Variables and parameters for the control trajectory
+        - Dynamics: Parameters for the discretized dynamics constraints
+        - Nodal constraints: Parameters for linearized non-convex nodal constraints
+        - Cross-node constraints: Parameters for linearized cross-node constraints
+        - Scaling: Affine scaling matrices and offset vectors
+        - Scaled expressions: CVXPy expressions for scaled state/control at each node
+
+    Attributes:
+        w_tr: Trust region weight parameter (scalar, nonneg)
+        lam_cost: Cost function weight parameter (scalar, nonneg)
+        lam_vc: Virtual control penalty weights (N-1 x n_states, nonneg)
+        lam_vb: Virtual buffer penalty weight (scalar, nonneg)
+
+        x: State variable (N x n_states)
+        dx: State error variable (N x n_states)
+        x_bar: Previous SCP state parameter (N x n_states)
+        x_init: Initial state parameter (n_states,)
+        x_term: Terminal state parameter (n_states,)
+
+        u: Control variable (N x n_controls)
+        du: Control error variable (N x n_controls)
+        u_bar: Previous SCP control parameter (N x n_controls)
+
+        A_d: Discretized state Jacobian parameter (N-1 x n_states*n_states)
+        B_d: Discretized control Jacobian parameter (N-1 x n_states*n_controls)
+        C_d: Discretized control Jacobian (next node) parameter
+        x_prop: Propagated state parameter (N-1 x n_states)
+        nu: Virtual control variable (N-1 x n_states)
+
+        g: List of constraint value parameters (one per nodal constraint)
+        grad_g_x: List of state gradient parameters (one per nodal constraint)
+        grad_g_u: List of control gradient parameters (one per nodal constraint)
+        nu_vb: List of virtual buffer variables (one per nodal constraint)
+
+        g_cross: List of cross-node constraint value parameters
+        grad_g_X_cross: List of trajectory state gradient parameters
+        grad_g_U_cross: List of trajectory control gradient parameters
+        nu_vb_cross: List of cross-node virtual buffer variables
+
+        S_x: State scaling matrix (n_states x n_states)
+        inv_S_x: Inverse state scaling matrix
+        c_x: State offset vector (n_states,)
+        S_u: Control scaling matrix (n_controls x n_controls)
+        inv_S_u: Inverse control scaling matrix
+        c_u: Control offset vector (n_controls,)
+
+        x_nonscaled: List of scaled state expressions at each node
+        u_nonscaled: List of scaled control expressions at each node
+        dx_nonscaled: List of scaled state error expressions at each node
+        du_nonscaled: List of scaled control error expressions at each node
+    """
+
+    # SCP weight parameters
+    w_tr: "cp.Parameter"
+    lam_cost: "cp.Parameter"
+    lam_vc: "cp.Parameter"
+    lam_vb: "cp.Parameter"
+
+    # State variables and parameters
+    x: "cp.Variable"
+    dx: "cp.Variable"
+    x_bar: "cp.Parameter"
+    x_init: "cp.Parameter"
+    x_term: "cp.Parameter"
+
+    # Control variables and parameters
+    u: "cp.Variable"
+    du: "cp.Variable"
+    u_bar: "cp.Parameter"
+
+    # Dynamics discretization parameters
+    A_d: "cp.Parameter"
+    B_d: "cp.Parameter"
+    C_d: "cp.Parameter"
+    x_prop: "cp.Parameter"
+    nu: "cp.Variable"
+
+    # Nodal constraint linearization (lists, one per constraint)
+    g: List["cp.Parameter"] = field(default_factory=list)
+    grad_g_x: List["cp.Parameter"] = field(default_factory=list)
+    grad_g_u: List["cp.Parameter"] = field(default_factory=list)
+    nu_vb: List["cp.Variable"] = field(default_factory=list)
+
+    # Cross-node constraint linearization (lists, one per constraint)
+    g_cross: List["cp.Parameter"] = field(default_factory=list)
+    grad_g_X_cross: List["cp.Parameter"] = field(default_factory=list)
+    grad_g_U_cross: List["cp.Parameter"] = field(default_factory=list)
+    nu_vb_cross: List["cp.Variable"] = field(default_factory=list)
+
+    # Scaling matrices and offsets (numpy arrays)
+    S_x: np.ndarray = field(default_factory=lambda: np.array([]))
+    inv_S_x: np.ndarray = field(default_factory=lambda: np.array([]))
+    c_x: np.ndarray = field(default_factory=lambda: np.array([]))
+    S_u: np.ndarray = field(default_factory=lambda: np.array([]))
+    inv_S_u: np.ndarray = field(default_factory=lambda: np.array([]))
+    c_u: np.ndarray = field(default_factory=lambda: np.array([]))
+
+    # Scaled CVXPy expressions at each node (lists of length N)
+    x_nonscaled: List = field(default_factory=list)
+    u_nonscaled: List = field(default_factory=list)
+    dx_nonscaled: List = field(default_factory=list)
+    du_nonscaled: List = field(default_factory=list)

openscvx/lowered/dynamics.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+from typing import Callable, Optional
+
+import jax.numpy as jnp
+
+
+@dataclass
+class Dynamics:
+    """Dataclass to hold a system dynamics function and its Jacobians.
+
+    This dataclass is used internally by openscvx to store the compiled dynamics
+    function and its gradients after symbolic expressions are lowered to JAX.
+    Users typically don't instantiate this class directly.
+
+    Attributes:
+        f (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
+            Function defining the continuous time nonlinear system dynamics
+            as x_dot = f(x, u, ...params).
+            - x: 1D array (state at a single node), shape (n_x,)
+            - u: 1D array (control at a single node), shape (n_u,)
+            - Additional parameters: passed as keyword arguments with names
+              matching the parameter name plus an underscore (e.g., g_ for
+              Parameter('g')).
+            If you use vectorized integration or batch evaluation, x and u
+            may be 2D arrays (N, n_x) and (N, n_u).
+        A (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            Jacobian of ``f`` w.r.t. ``x``.
+        B (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            Jacobian of ``f`` w.r.t. ``u``.
+    """
+
+    f: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
+    A: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    B: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None

openscvx/lowered/jax_constraints.py

@@ -0,0 +1,133 @@
+"""JAX-lowered constraint dataclass."""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Callable, List, Optional
+
+import jax.numpy as jnp
+
+if TYPE_CHECKING:
+    from openscvx.symbolic.expr import CTCS
+
+
+@dataclass
+class LoweredNodalConstraint:
+    """
+    Dataclass to hold a lowered symbolic constraint function and its jacobians.
+
+    This is a simplified drop-in replacement for NodalConstraint that holds
+    only the essential lowered JAX functions and their jacobians, without
+    the complexity of convex/vectorized flags or post-initialization logic.
+
+    Designed for use with symbolic expressions that have been lowered to JAX
+    and will be linearized for sequential convex programming.
+
+    Args:
+        func (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]):
+            The lowered constraint function g(x, u, ...params) that returns
+            constraint residuals. Should follow g(x, u) <= 0 convention.
+            - x: 1D array (state at a single node), shape (n_x,)
+            - u: 1D array (control at a single node), shape (n_u,)
+            - Additional parameters: passed as keyword arguments
+
+        grad_g_x (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            Jacobian of g w.r.t. x. If None, should be computed using jax.jacfwd.
+
+        grad_g_u (Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]]):
+            Jacobian of g w.r.t. u. If None, should be computed using jax.jacfwd.
+
+        nodes (Optional[List[int]]): List of node indices where this constraint applies.
+            Set after lowering from NodalConstraint.
+    """
+
+    func: Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]
+    grad_g_x: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    grad_g_u: Optional[Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]] = None
+    nodes: Optional[List[int]] = None
+
+
+@dataclass
+class LoweredCrossNodeConstraint:
+    """Lowered cross-node constraint with trajectory-level evaluation.
+
+    Unlike regular LoweredNodalConstraint which operates on single-node vectors
+    and is vmapped across the trajectory, LoweredCrossNodeConstraint operates
+    on full trajectory arrays to relate multiple nodes simultaneously.
+
+    This is necessary for constraints like:
+    - Rate limits: x[k] - x[k-1] <= max_rate
+    - Multi-step dependencies: x[k] = 2*x[k-1] - x[k-2]
+    - Periodic boundaries: x[0] = x[N-1]
+
+    The function signatures differ from LoweredNodalConstraint:
+    - Regular: f(x, u, node, params) -> scalar (vmapped to handle (N, n_x))
+    - Cross-node: f(X, U, params) -> scalar (single constraint with fixed node indices)
+
+    Attributes:
+        func: Function (X, U, params) -> scalar residual
+            where X: (N, n_x), U: (N, n_u)
+            Returns constraint residual following g(X, U) <= 0 convention
+            The constraint references fixed trajectory nodes (e.g., X[5] - X[4])
+        grad_g_X: Function (X, U, params) -> (N, n_x) Jacobian wrt full state trajectory
+            This is typically sparse - most constraints only couple nearby nodes
+        grad_g_U: Function (X, U, params) -> (N, n_u) Jacobian wrt full control trajectory
+            Often zero or very sparse for cross-node state constraints
+
+    Example:
+        For rate constraint x[5] - x[4] <= r:
+
+            func(X, U, params) -> scalar residual
+            grad_g_X(X, U, params) -> (N, n_x) sparse Jacobian
+                where grad_g_X[5, :] = ∂g/∂x[5] (derivative wrt node 5)
+                and grad_g_X[4, :] = ∂g/∂x[4] (derivative wrt node 4)
+                all other entries are zero
+
+    Performance Note - Dense Jacobian Storage:
+        The Jacobian matrices grad_g_X and grad_g_U are stored as DENSE arrays with
+        shape (N, n_x) and (N, n_u), but most cross-node constraints only couple a
+        small number of nearby nodes, making these matrices extremely sparse.
+
+        For example, a rate limit constraint x[k] - x[k-1] <= r only has non-zero
+        Jacobian entries at positions [k, :] and [k-1, :]. All other N-2 rows are
+        zero but still stored in memory.
+
+        Memory impact for large problems:
+        - A single constraint with N=100 nodes, n_x=10 states requires ~8KB for
+          grad_g_X (compared to ~160 bytes if sparse with 2 non-zero rows)
+        - Multiple cross-node constraints multiply this overhead
+        - May cause issues for N > 1000 with many constraints
+
+        Performance impact:
+        - Slower autodiff (computes many zero gradients)
+        - Inefficient constraint linearization in the SCP solver
+        - Potential GPU memory limitations for very large problems
+
+        The current implementation prioritizes simplicity and compatibility with
+        JAX's autodiff over memory efficiency. Future versions may support sparse
+        Jacobian formats (COO, CSR, or custom sparse representations) if this
+        becomes a bottleneck in practice.
+    """
+
+    func: Callable[[jnp.ndarray, jnp.ndarray, dict], jnp.ndarray]
+    grad_g_X: Callable[[jnp.ndarray, jnp.ndarray, dict], jnp.ndarray]
+    grad_g_U: Callable[[jnp.ndarray, jnp.ndarray, dict], jnp.ndarray]
+
+
+@dataclass
+class LoweredJaxConstraints:
+    """JAX-lowered non-convex constraints with gradient functions.
+
+    Contains constraints that have been lowered to JAX callable functions
+    with automatically computed gradients. These are used for linearization
+    in the SCP (Sequential Convex Programming) loop.
+
+    Attributes:
+        nodal: List of LoweredNodalConstraint objects. Each has `func`,
+            `grad_g_x`, `grad_g_u` callables and `nodes` list.
+        cross_node: List of LoweredCrossNodeConstraint objects. Each has
+            `func`, `grad_g_X`, `grad_g_U` for trajectory-level constraints.
+        ctcs: CTCS constraints (unchanged from input, not lowered here).
+    """
+
+    nodal: list[LoweredNodalConstraint] = field(default_factory=list)
+    cross_node: list[LoweredCrossNodeConstraint] = field(default_factory=list)
+    ctcs: list["CTCS"] = field(default_factory=list)

openscvx/lowered/parameters.py

@@ -0,0 +1,54 @@
+"""Parameter dictionary that syncs between JAX and CVXPy."""
+
+import numpy as np
+
+
+class ParameterDict(dict):
+    """Dictionary that syncs to both internal _parameters dict and CVXPy parameters.
+
+    This allows users to naturally update parameters like:
+        problem.parameters["obs_radius"] = 2.0
+
+    Changes automatically propagate to:
+    1. Internal _parameters dict (plain dict for JAX)
+    2. CVXPy parameters (for optimization)
+    """
+
+    def __init__(self, problem, internal_dict, *args, **kwargs):
+        self._problem = problem
+        self._internal_dict = internal_dict  # Reference to plain dict for JAX
+        super().__init__()
+        # Initialize with float enforcement by using __setitem__
+        if args:
+            other = args[0]
+            if hasattr(other, "items"):
+                for key, value in other.items():
+                    self[key] = value
+            else:
+                for key, value in other:
+                    self[key] = value
+        for key, value in kwargs.items():
+            self[key] = value
+
+    def __setitem__(self, key, value):
+        # Enforce float dtype to prevent int/float mismatch bugs
+        value = np.asarray(value, dtype=float)
+        super().__setitem__(key, value)
+        # Sync to internal dict for JAX
+        self._internal_dict[key] = value
+        # Sync to CVXPy if it exists
+        lowered = getattr(self._problem, "_lowered", None)
+        if lowered is not None and key in lowered.cvxpy_params:
+            lowered.cvxpy_params[key].value = value
+
+    def update(self, other=None, **kwargs):
+        """Update multiple parameters and sync to internal dict and CVXPy."""
+        if other is not None:
+            if hasattr(other, "items"):
+                for key, value in other.items():
+                    self[key] = value
+            else:
+                for key, value in other:
+                    self[key] = value
+        for key, value in kwargs.items():
+            self[key] = value

openscvx/lowered/problem.py

@@ -0,0 +1,70 @@
+"""LoweredProblem dataclass - container for all lowering outputs."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict
+
+from openscvx.lowered.cvxpy_constraints import LoweredCvxpyConstraints
+from openscvx.lowered.cvxpy_variables import CVXPyVariables
+from openscvx.lowered.dynamics import Dynamics
+from openscvx.lowered.jax_constraints import LoweredJaxConstraints
+from openscvx.lowered.unified import UnifiedControl, UnifiedState
+
+if TYPE_CHECKING:
+    import cvxpy as cp
+
+
+@dataclass
+class LoweredProblem:
+    """Container for all outputs from symbolic problem lowering.
+
+    This dataclass holds all the results of lowering symbolic expressions
+    to executable JAX and CVXPy code. It provides a clean, typed interface
+    for accessing the various components needed for optimization.
+
+    Attributes:
+        dynamics: Optimization dynamics with fields f, A, B (JAX functions)
+        dynamics_prop: Propagation dynamics with fields f, A, B
+        jax_constraints: Non-convex constraints lowered to JAX with gradients
+        cvxpy_constraints: Convex constraints lowered to CVXPy
+        x_unified: Aggregated optimization state interface
+        u_unified: Aggregated optimization control interface
+        x_prop_unified: Aggregated propagation state interface
+        ocp_vars: Typed CVXPy variables and parameters for OCP construction
+        cvxpy_params: Dict mapping user parameter names to CVXPy Parameter objects
+
+    Example:
+        After lowering a symbolic problem::
+
+            lowered = lower_symbolic_problem(
+                dynamics_aug=dynamics,
+                states_aug=states,
+                controls_aug=controls,
+                constraints=constraint_set,
+                parameters=params,
+                N=50,
+            )
+
+            # Access components
+            dx_dt = lowered.dynamics.f(x, u, node, params)
+            jacobian_A = lowered.dynamics.A(x, u, node, params)
+
+            # Use CVXPy objects
+            ocp = OptimalControlProblem(settings, lowered)
+    """
+
+    # JAX dynamics
+    dynamics: Dynamics
+    dynamics_prop: Dynamics
+
+    # Lowered constraints (separate types for JAX vs CVXPy)
+    jax_constraints: LoweredJaxConstraints
+    cvxpy_constraints: LoweredCvxpyConstraints
+
+    # Unified interfaces
+    x_unified: UnifiedState
+    u_unified: UnifiedControl
+    x_prop_unified: UnifiedState
+
+    # CVXPy objects
+    ocp_vars: CVXPyVariables
+    cvxpy_params: Dict[str, "cp.Parameter"]