openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/__init__.py

@@ -0,0 +1,123 @@
+import os
+
+# Set Equinox error handling to return NaN instead of crashing
+os.environ["EQX_ON_ERROR"] = "nan"
+
+# Cache management
+# Core symbolic expressions - flat namespace for most common functions
+import openscvx.symbolic.expr.lie as lie
+import openscvx.symbolic.expr.linalg as linalg
+import openscvx.symbolic.expr.spatial as spatial
+import openscvx.symbolic.expr.stl as stl
+from openscvx.expert import ByofSpec
+from openscvx.problem import Problem
+from openscvx.symbolic.expr import (
+    CTCS,
+    Abs,
+    Add,
+    Block,
+    Concat,
+    Constant,
+    Constraint,
+    Control,
+    Cos,
+    Diag,
+    Div,
+    Equality,
+    Exp,
+    Expr,
+    Fixed,
+    Free,
+    Hstack,
+    Index,
+    Inequality,
+    Leaf,
+    Log,
+    LogSumExp,
+    MatMul,
+    Max,
+    Maximize,
+    Minimize,
+    Mul,
+    Neg,
+    NodalConstraint,
+    Parameter,
+    Power,
+    Sin,
+    Sqrt,
+    Stack,
+    State,
+    Sub,
+    Sum,
+    Tan,
+    Variable,
+    Vstack,
+    ctcs,
+)
+from openscvx.symbolic.time import Time
+from openscvx.utils.cache import clear_cache, get_cache_dir, get_cache_size
+
+__all__ = [
+    # Main Trajectory Optimization Entrypoint
+    "Problem",
+    # Cache management
+    "get_cache_dir",
+    "clear_cache",
+    "get_cache_size",
+    # Time configuration
+    "Time",
+    # Core base classes
+    "Expr",
+    "Leaf",
+    "Parameter",
+    "Variable",
+    "State",
+    "Control",
+    # Boundary condition helpers
+    "Free",
+    "Fixed",
+    "Minimize",
+    "Maximize",
+    # Basic arithmetic operations
+    "Add",
+    "Sub",
+    "Mul",
+    "Div",
+    "MatMul",
+    "Neg",
+    "Power",
+    "Sum",
+    # Array operations
+    "Index",
+    "Concat",
+    "Stack",
+    "Hstack",
+    "Vstack",
+    "Block",
+    "Diag",
+    "Constant",
+    # Mathematical functions
+    "Sin",
+    "Cos",
+    "Tan",
+    "Sqrt",
+    "Abs",
+    "Exp",
+    "Log",
+    "LogSumExp",
+    "Max",
+    # Constraints
+    "Constraint",
+    "Equality",
+    "Inequality",
+    "NodalConstraint",
+    "CTCS",
+    "ctcs",
+    # Submodules
+    "stl",
+    "spatial",
+    "linalg",
+    "lie",
+    # Expert mode types
+    "ByofSpec",
+]

openscvx/_version.py

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

openscvx/algorithms/__init__.py

@@ -0,0 +1,92 @@
+"""Successive convexification algorithms for trajectory optimization.
+
+This module provides implementations of SCvx (Successive Convexification) algorithms
+for solving non-convex trajectory optimization problems through iterative convex
+approximation.
+
+All algorithms inherit from :class:`Algorithm`, enabling pluggable algorithm
+implementations and custom SCvx variants:
+
+```python
+class Algorithm(ABC):
+    @abstractmethod
+    def initialize(self, ocp, discretization_solver, jax_constraints,
+                   solve_ocp, emitter, params, settings) -> None:
+        '''Store compiled infrastructure and warm-start solvers.'''
+        ...
+
+    @abstractmethod
+    def step(self, state, params, settings) -> bool:
+        '''Execute one iteration using stored infrastructure.'''
+        ...
+```
+
+Immutable components (ocp, discretization_solver, jax_constraints, etc.) are stored
+during ``initialize()``. Mutable configuration (params, settings) is passed per-step
+to support runtime parameter updates and tolerance tuning.
+
+:class:`AlgorithmState` holds mutable state during SCP iterations. Algorithms
+that require additional state can subclass it:
+
+```python
+@dataclass
+class MyAlgorithmState(AlgorithmState):
+    my_custom_field: float = 0.0
+```
+
+Note:
+    ``AlgorithmState`` currently combines iteration metrics (costs, weights),
+    trajectory history, and discretization data. A future refactor may separate
+    these concerns into distinct classes for clearer data flow:
+
+    ```python
+    @dataclass
+    class AlgorithmState:
+        # Mutable iteration state
+        k: int
+        J_tr: float
+        J_vb: float
+        J_vc: float
+        w_tr: float
+        lam_cost: float
+        lam_vc: ...
+        lam_vb: float
+
+    @dataclass
+    class TrajectoryHistory:
+        # Accumulated trajectory solutions
+        X: List[np.ndarray]
+        U: List[np.ndarray]
+
+        @property
+        def x(self): return self.X[-1]
+
+        @property
+        def u(self): return self.U[-1]
+
+    @dataclass
+    class DebugHistory:
+        # Optional diagnostic data (discretization matrices, etc.)
+        V_history: List[np.ndarray]
+        VC_history: List[np.ndarray]
+        TR_history: List[np.ndarray]
+    ```
+
+Current Implementations:
+
+- :class:`PenalizedTrustRegion`: Penalized Trust Region (PTR) algorithm
+"""
+
+from .base import Algorithm, AlgorithmState
+from .optimization_results import OptimizationResults
+from .penalized_trust_region import PenalizedTrustRegion
+
+__all__ = [
+    # Base class
+    "Algorithm",
+    "AlgorithmState",
+    # Core results
+    "OptimizationResults",
+    # PTR algorithm
+    "PenalizedTrustRegion",
+]

openscvx/algorithms/autotuning.py

@@ -0,0 +1,24 @@
+"""Autotuning functions for SCP (Successive Convex Programming) parameters."""
+
+from typing import TYPE_CHECKING
+
+from openscvx.config import Config
+
+if TYPE_CHECKING:
+    from .base import AlgorithmState
+
+
+def update_scp_weights(state: "AlgorithmState", settings: Config, scp_k: int):
+    """Update SCP weights and cost parameters based on iteration number.
+
+    Args:
+        state: Solver state containing current weight values (mutated in place)
+        settings: Configuration object containing adaptation parameters
+        scp_k: Current SCP iteration number
+    """
+    # Update trust region weight in state
+    state.w_tr = min(state.w_tr * settings.scp.w_tr_adapt, settings.scp.w_tr_max)
+
+    # Update cost relaxation parameter after cost_drop iterations
+    if scp_k > settings.scp.cost_drop:
+        state.lam_cost = state.lam_cost * settings.scp.cost_relax

openscvx/algorithms/base.py

@@ -0,0 +1,351 @@
+"""Base class for successive convexification algorithms.
+
+This module defines the abstract interface that all SCP algorithm implementations
+must follow, along with the AlgorithmState dataclass that holds mutable state
+during SCP iterations.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, Union
+
+import numpy as np
+
+if TYPE_CHECKING:
+    import cvxpy as cp
+
+    from openscvx.config import Config
+    from openscvx.lowered.jax_constraints import LoweredJaxConstraints
+
+
+@dataclass
+class AlgorithmState:
+    """Mutable state for SCP iterations.
+
+    This dataclass holds all state that changes during the solve process.
+    It stores only the evolving trajectory arrays, not the full State/Control
+    objects which contain immutable configuration metadata.
+
+    Trajectory arrays are stored in history lists, with the current guess
+    accessed via properties that return the latest entry.
+
+    A fresh instance is created for each solve, enabling easy reset functionality.
+
+    Attributes:
+        k: Current iteration number (starts at 1)
+        J_tr: Current trust region cost
+        J_vb: Current virtual buffer cost
+        J_vc: Current virtual control cost
+        w_tr: Current trust region weight (may adapt during solve)
+        lam_cost: Current cost weight (may relax during solve)
+        lam_vc: Current virtual control penalty weight
+        lam_vb: Current virtual buffer penalty weight
+        n_x: Number of states (for unpacking V vectors)
+        n_u: Number of controls (for unpacking V vectors)
+        N: Number of trajectory nodes (for unpacking V vectors)
+        X: List of state trajectory iterates
+        U: List of control trajectory iterates
+        V_history: List of discretization history
+    """
+
+    k: int
+    J_tr: float
+    J_vb: float
+    J_vc: float
+    w_tr: float
+    lam_cost: float
+    lam_vc: Union[float, np.ndarray]
+    lam_vb: float
+    n_x: int
+    n_u: int
+    N: int
+    X: List[np.ndarray] = field(default_factory=list)
+    U: List[np.ndarray] = field(default_factory=list)
+    V_history: List[np.ndarray] = field(default_factory=list)
+    VC_history: List[np.ndarray] = field(default_factory=list)
+    TR_history: List[np.ndarray] = field(default_factory=list)
+
+    @property
+    def x(self) -> np.ndarray:
+        """Get current state trajectory array.
+
+        Returns:
+            Current state trajectory guess (latest entry in history), shape (N, n_states)
+        """
+        return self.X[-1]
+
+    @property
+    def u(self) -> np.ndarray:
+        """Get current control trajectory array.
+
+        Returns:
+            Current control trajectory guess (latest entry in history), shape (N, n_controls)
+        """
+        return self.U[-1]
+
+    @property
+    def x_prop(self) -> np.ndarray:
+        """Extract propagated state trajectory from latest V.
+
+        Returns:
+            Propagated state trajectory x_prop with shape (N-1, n_x), or None if no V_history
+
+        Example:
+            After running an iteration, access the propagated states::
+
+                problem.step()
+                x_prop = problem.state.x_prop  # Shape (N-1, n_x)
+        """
+        if not self.V_history:
+            return None
+
+        # V_history contains Vmulti from discretization
+        # Shape: (flattened_size, n_timesteps) where flattened_size = (N-1) * i4
+        V = self.V_history[-1]
+
+        # Take final timestep and reshape to (N-1, i4)
+        i4 = self.n_x + self.n_x * self.n_x + 2 * self.n_x * self.n_u
+        V_final = V[:, -1].reshape(-1, i4)
+
+        # Extract propagated state (first n_x elements of each row)
+        return V_final[:, : self.n_x]
+
+    @property
+    def A_d(self) -> np.ndarray:
+        """Extract discretized state transition matrix from latest V.
+
+        Returns:
+            Discretized state Jacobian A_d with shape (N-1, n_x, n_x), or None if no V_history
+
+        Example:
+            After running an iteration, access linearization matrices::
+
+                problem.step()
+                A_d = problem.state.A_d  # Shape (N-1, n_x, n_x)
+        """
+        if not self.V_history:
+            return None
+
+        # Extract indices for unpacking V vector
+        i1 = self.n_x
+        i2 = i1 + self.n_x * self.n_x
+
+        # V_history contains Vmulti from discretization
+        # Shape: (flattened_size, n_timesteps) where flattened_size = (N-1) * i4
+        V = self.V_history[-1]
+
+        # Take final timestep and reshape to (N-1, i4)
+        i4 = self.n_x + self.n_x * self.n_x + 2 * self.n_x * self.n_u
+        V_final = V[:, -1].reshape(-1, i4)
+
+        # Extract and reshape A_d matrix
+        return V_final[:, i1:i2].reshape(self.N - 1, self.n_x, self.n_x)
+
+    @property
+    def B_d(self) -> np.ndarray:
+        """Extract discretized control influence matrix (current node) from latest V.
+
+        Returns:
+            Discretized control Jacobian B_d with shape (N-1, n_x, n_u), or None if no V_history
+
+        Example:
+            After running an iteration, access linearization matrices::
+
+                problem.step()
+                B_d = problem.state.B_d  # Shape (N-1, n_x, n_u)
+        """
+        if not self.V_history:
+            return None
+
+        # Extract indices for unpacking V vector
+        i1 = self.n_x
+        i2 = i1 + self.n_x * self.n_x
+        i3 = i2 + self.n_x * self.n_u
+
+        # V_history contains Vmulti from discretization
+        V = self.V_history[-1]
+
+        # Take final timestep and reshape to (N-1, i4)
+        i4 = self.n_x + self.n_x * self.n_x + 2 * self.n_x * self.n_u
+        V_final = V[:, -1].reshape(-1, i4)
+
+        # Extract and reshape B_d matrix
+        return V_final[:, i2:i3].reshape(self.N - 1, self.n_x, self.n_u)
+
+    @property
+    def C_d(self) -> np.ndarray:
+        """Extract discretized control influence matrix (next node) from latest V.
+
+        Returns:
+            Discretized control Jacobian C_d with shape (N-1, n_x, n_u), or None if no V_history
+
+        Example:
+            After running an iteration, access linearization matrices::
+
+                problem.step()
+                C_d = problem.state.C_d  # Shape (N-1, n_x, n_u)
+        """
+        if not self.V_history:
+            return None
+
+        # Extract indices for unpacking V vector
+        i2 = self.n_x + self.n_x * self.n_x
+        i3 = i2 + self.n_x * self.n_u
+        i4 = i3 + self.n_x * self.n_u
+
+        # V_history contains Vmulti from discretization
+        V = self.V_history[-1]
+
+        # Take final timestep and reshape to (N-1, i4)
+        V_final = V[:, -1].reshape(-1, i4)
+
+        # Extract and reshape C_d matrix
+        return V_final[:, i3:i4].reshape(self.N - 1, self.n_x, self.n_u)
+
+    @classmethod
+    def from_settings(cls, settings: "Config") -> "AlgorithmState":
+        """Create initial algorithm state from configuration.
+
+        Copies only the trajectory arrays from settings, leaving all metadata
+        (bounds, boundary conditions, etc.) in the original settings object.
+
+        Args:
+            settings: Configuration object containing initial guesses and SCP parameters
+
+        Returns:
+            Fresh AlgorithmState initialized from settings with copied arrays
+        """
+        return cls(
+            k=1,
+            J_tr=1e2,
+            J_vb=1e2,
+            J_vc=1e2,
+            w_tr=settings.scp.w_tr,
+            lam_cost=settings.scp.lam_cost,
+            lam_vc=settings.scp.lam_vc,
+            lam_vb=settings.scp.lam_vb,
+            n_x=settings.sim.n_states,
+            n_u=settings.sim.n_controls,
+            N=settings.scp.n,
+            X=[settings.sim.x.guess.copy()],
+            U=[settings.sim.u.guess.copy()],
+            V_history=[],
+            VC_history=[],
+            TR_history=[],
+        )
+
+
+class Algorithm(ABC):
+    """Abstract base class for successive convexification algorithms.
+
+    This class defines the interface for SCP algorithms used in trajectory
+    optimization. Implementations should remain minimal and functional,
+    delegating state management to the AlgorithmState dataclass.
+
+    The two core methods mirror the SCP workflow:
+
+    - initialize: Store compiled infrastructure and warm-start solvers
+    - step: Execute one convex subproblem iteration
+
+    Immutable components (ocp, discretization_solver, jax_constraints, etc.) are
+    stored during initialize(). Mutable configuration (params, settings) is passed
+    per-step to support runtime parameter updates and tolerance tuning.
+
+    !!! tip "Statefullness"
+        Avoid storing mutable iteration state (costs, weights, trajectories) on
+        ``self``. All iteration state should live in :class:`AlgorithmState` or
+        a subclass thereof, passed explicitly to ``step()``. This keeps algorithm
+        classes stateless w.r.t. iteration, making data flow explicit and staying
+        close to functional programming principles where possible.
+
+    Example:
+        Implementing a custom algorithm::
+
+            class MyAlgorithm(Algorithm):
+                def initialize(self, ocp, discretization_solver,
+                               jax_constraints, solve_ocp, emitter,
+                               params, settings):
+                    # Store compiled infrastructure
+                    self._ocp = ocp
+                    self._discretization_solver = discretization_solver
+                    self._jax_constraints = jax_constraints
+                    self._solve_ocp = solve_ocp
+                    self._emitter = emitter
+                    # Warm-start with initial params/settings...
+
+                def step(self, state, params, settings):
+                    # Run one iteration using self._* and per-step params/settings
+                    return converged
+    """
+
+    @abstractmethod
+    def initialize(
+        self,
+        ocp: "cp.Problem",
+        discretization_solver: callable,
+        jax_constraints: "LoweredJaxConstraints",
+        solve_ocp: callable,
+        emitter: callable,
+        params: dict,
+        settings: "Config",
+    ) -> None:
+        """Initialize the algorithm and store compiled infrastructure.
+
+        This method stores immutable components and performs any setup required
+        before the SCP loop begins (e.g., warm-starting solvers). The params and
+        settings are passed for warm-start but may change between steps.
+
+        Args:
+            ocp: The CVXPy optimal control problem
+            discretization_solver: Compiled discretization solver function
+            jax_constraints: JIT-compiled JAX constraint functions
+            solve_ocp: Callable that solves the OCP (captures solver config)
+            emitter: Callback for emitting iteration progress data
+            params: Problem parameters dictionary (for warm-start only)
+            settings: Configuration object (for warm-start only)
+        """
+        ...
+
+    @abstractmethod
+    def step(
+        self,
+        state: AlgorithmState,
+        params: dict,
+        settings: "Config",
+    ) -> bool:
+        """Execute one iteration of the SCP algorithm.
+
+        This method solves a single convex subproblem, updates the algorithm
+        state in place, and returns whether convergence criteria are met.
+
+        Uses stored infrastructure (ocp, discretization_solver, etc.) with
+        per-step params and settings to support runtime modifications.
+
+        Args:
+            state: Mutable algorithm state (modified in place)
+            params: Problem parameters dictionary (may change between steps)
+            settings: Configuration object (may change between steps)
+
+        Returns:
+            True if convergence criteria are satisfied, False otherwise.
+        """
+        ...
+
+    @abstractmethod
+    def citation(self) -> List[str]:
+        """Return BibTeX citations for this algorithm.
+
+        Implementations should return a list of BibTeX entry strings for the
+        papers that should be cited when using this algorithm.
+
+        Returns:
+            List of BibTeX citation strings.
+
+        Example:
+            Getting citations for an algorithm::
+
+                algorithm = PenalizedTrustRegion()
+                for bibtex in algorithm.citation():
+                    print(bibtex)
+        """
+        ...