openscvx 0.5.2.dev16__tar.gz → 0.5.3.dev12__tar.gz

{openscvx-0.5.2.dev16/openscvx.egg-info → openscvx-0.5.3.dev12}/PKG-INFO

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

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/docs/UnderTheHood/lowering_architecture.md

@@ -171,6 +171,16 @@ of the abstract `PTRSolver` base. Three backends ship today:
 | QPAX | `QPAXPTRSolver` | `solver={"backend": "qpax"}` | JAX-native QP via `qpax.solve_qp`. Flat `(Q, q, A, b, G, h)` assembly. Supports box / dynamics (continuous and impulsive) / CTCS / boundary-Fix; **rejects** user `.convex()` and cross-node at `initialize()` with a clear "use `CVXPyPTRSolver`" message. Enables a path toward an end-to-end JAX-differentiable SCP loop in follow-up work. |
 | Moreau | `MoreauPTRSolver` | `solver={"backend": "moreau"}` | JAX-native conic solver (`moreau.jax.Solver`). Sparse CSR assembly; SOC epigraphs for the L1 / pos PTR penalties (fewer variables and rows than QPAX). Warm-starts between SCP iterations. Same supported subset as QPAX (continuous and impulsive dynamics). Paves the way for user `.convex()` SOC support in a follow-up. |
 
+Each backend exposes two per-iteration entry points: the historical NumPy
+`update_*` + `solve()` stages used by today's Python-side SCP loop, and a
+JAX-pure `iteration_callback()` built once at `initialize()`. The callback
+returns a `(state, SubproblemData) -> SubproblemSolution` callable; all
+three backends emit the same input/output pytree shape so the callable
+composes with `jax.jit` and `jax.vmap` (CVXPy via `jax.pure_callback` with
+`vmap_method="sequential"` — its batched solves are serial). The downstream
+batchable-`Problem.solve()` work wraps this callback in `lax.while_loop` to
+keep the entire SCP iteration on the JAX boundary.
+
 Picking a backend at construction time:
 
 ```python

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/docs/UsersGuide/03_obstacle_avoidance_vmap.md

@@ -373,7 +373,7 @@ while running:
     step_result = problem.step()
 
     # Extract control to apply
-    control = problem.state.V_history[-1]  # Access current solution
+    control = problem.history.V_history[-1]  # Access current solution
 ```
 
 Parameter updates take effect immediately on the next `step()` call without any recompilation.

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/examples/realtime/3DoF_pdg_realtime.py

@@ -455,8 +455,8 @@ def create_realtime_server(optimization_problem) -> viser.ViserServer:
                 results.update(get_print_queue_data(optimization_problem))
                 metrics_text.content = format_metrics_markdown(results)
 
-                if optimization_problem.state.V_history:
-                    V_multi_shoot = np.asarray(optimization_problem.state.V_history[-1])
+                if optimization_problem.history.V_history:
+                    V_multi_shoot = np.asarray(optimization_problem.history.V_history[-1])
                     update_trajectory(V_multi_shoot)
 
                 time.sleep(0.05)

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/examples/realtime/6DoF_pdg_realtime.py

@@ -613,8 +613,8 @@ def create_realtime_server(optimization_problem) -> viser.ViserServer:
                 results.update(get_print_queue_data(optimization_problem))
                 metrics_text.content = format_metrics_markdown(results)
 
-                if optimization_problem.state.V_history:
-                    V_multi_shoot = np.asarray(optimization_problem.state.V_history[-1])
+                if optimization_problem.history.V_history:
+                    V_multi_shoot = np.asarray(optimization_problem.history.V_history[-1])
                     update_trajectory(V_multi_shoot)
 
                 time.sleep(0.05)

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/examples/realtime/cinema_vp_realtime.py

@@ -434,8 +434,8 @@ def create_realtime_server(
                 update_metrics(results)
 
                 # Update trajectory from V_history
-                if optimization_problem.state.V_history:
-                    V_multi_shoot = np.array(optimization_problem.state.V_history[-1])
+                if optimization_problem.history.V_history:
+                    V_multi_shoot = np.array(optimization_problem.history.V_history[-1])
                     x_traj = np.array(optimization_problem.state.x)
                     update_trajectory_and_drone(V_multi_shoot, x_traj)
 

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/examples/realtime/drone_racing_realtime.py

@@ -355,8 +355,8 @@ def create_realtime_server(
                 update_gates()
 
                 # Update trajectory from V_history
-                if optimization_problem.state.V_history:
-                    V_multi_shoot = np.array(optimization_problem.state.V_history[-1])
+                if optimization_problem.history.V_history:
+                    V_multi_shoot = np.array(optimization_problem.history.V_history[-1])
                     update_trajectory(V_multi_shoot)
 
                 iteration += 1

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/examples/realtime/dubins_car_realtime.py

@@ -228,8 +228,8 @@ def create_realtime_server(optimization_problem) -> viser.ViserServer:
 
                 update_metrics(results)
 
-                if optimization_problem.state.V_history:
-                    V_multi_shoot = np.asarray(optimization_problem.state.V_history[-1])
+                if optimization_problem.history.V_history:
+                    V_multi_shoot = np.asarray(optimization_problem.history.V_history[-1])
                     update_trajectory(V_multi_shoot)
 
                 time.sleep(0.05)

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/examples/realtime/obstacle_avoidance_realtime.py

@@ -334,8 +334,8 @@ def create_realtime_server(
                 update_obstacles()
 
                 # Update trajectory from V_history
-                if optimization_problem.state.V_history:
-                    V_multi_shoot = np.array(optimization_problem.state.V_history[-1])
+                if optimization_problem.history.V_history:
+                    V_multi_shoot = np.array(optimization_problem.history.V_history[-1])
                     update_trajectory(V_multi_shoot)
 
                 iteration += 1

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.2.dev16'
-__version_tuple__ = version_tuple = (0, 5, 2, 'dev16')
+__version__ = version = '0.5.3.dev12'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev12')
 
-__commit_id__ = commit_id = 'ged65ccb87'
+__commit_id__ = commit_id = 'g9149c4b81'

{openscvx-0.5.2.dev16 → openscvx-0.5.3.dev12}/openscvx/algorithms/__init__.py

@@ -5,73 +5,27 @@ 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, solver, discretization_solver, jax_constraints,
-                   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 (solver, 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
-        lam_prox: float
-        lam_cost: float
-        lam_vc: ...
-        lam_vb_nodal: np.ndarray  # (N, n_nodal)
-        lam_vb_cross: np.ndarray  # (n_cross,)
-
-    @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]
-    ```
+implementations and custom SCvx variants. Immutable components (solver,
+discretization_solver, jax_constraints, etc.) are stored during ``initialize()``;
+mutable configuration (params, settings) is passed per-step.
+
+The iterate carry is split into two objects:
+
+* :class:`AlgorithmState` — a frozen, JAX-registered pytree holding only the
+  current iterate (``x``, ``u``, weights, propagated states, diagnostic scalars).
+  Every leaf is a ``jnp.ndarray`` so the state composes with ``jax.vmap`` /
+  ``jax.jit`` / ``jax.grad``.
+
+* :class:`AlgorithmHistory` — a CPU-side mutable record of every iteration's
+  trajectories, discretizations, weights, and diagnostics. Grown by the SCP loop
+  via ``record_iteration``; never crosses the JAX boundary.
+
+Autotuners follow the same contract: :class:`AutotuningBase` declares
+``update_weights`` as a pure functional update on the :class:`AlgorithmState`
+pytree (no mutation, no string returns, no list appends, no Python-level
+branching on iterate values). The next-iterate adaptive state is reported via
+the :class:`AdaptiveStateCode` IntEnum and converted to a human-readable string
+on the printing path with :func:`adaptive_state_code_to_str`.
 
 Current Implementations:
 
@@ -92,7 +46,15 @@ from .autotuner import (
     RampProximalWeight,
     RampProximalWeightSpec,
 )
-from .base import Algorithm, AlgorithmState, AutotuningBase, DiscretizationResult
+from .base import (
+    AdaptiveStateCode,
+    Algorithm,
+    AlgorithmHistory,
+    AlgorithmState,
+    AutotuningBase,
+    DiscretizationResult,
+    adaptive_state_code_to_str,
+)
 from .optimization_results import OptimizationResults
 from .scvx import PenalizedTrustRegion
 from .weights import Weights
@@ -172,7 +134,10 @@ class PenalizedTrustRegionConfig(BaseModel):
 __all__ = [
     # Base class
     "Algorithm",
+    "AlgorithmHistory",
     "AlgorithmState",
+    "AdaptiveStateCode",
+    "adaptive_state_code_to_str",
     "DiscretizationResult",
     "Weights",
     # Core results

openscvx-0.5.3.dev12/openscvx/algorithms/autotuner/adaptive_proximal_weight.py

@@ -0,0 +1,192 @@
+"""Adaptive proximal-weight autotuner.
+
+Same acceptance-ratio decision logic as :class:`AugmentedLagrangian` for
+``lam_prox``, but ``lam_vc`` and ``lam_vb_*`` are held constant at their
+current values.
+"""
+
+from typing import TYPE_CHECKING, Literal
+
+import jax
+import jax.numpy as jnp
+from pydantic import BaseModel, ConfigDict
+
+from openscvx.config import Config
+
+from ..base import AdaptiveStateCode, AutotuningBase
+from .augmented_lagrangian import AugmentedLagrangian
+
+if TYPE_CHECKING:
+    from openscvx.lowered import LoweredJaxConstraints
+
+    from ..base import AlgorithmState, CandidateIterate
+
+
+class AdaptiveProximalWeight(AutotuningBase):
+    """PTR-style proximal adaptation with fixed virtual-penalty weights.
+
+    Same four-bucket acceptance-ratio logic as
+    :class:`AugmentedLagrangian` for ``lam_prox``, but the virtual-control
+    and virtual-buffer weights are carried unchanged.
+
+    ``update_weights`` is a pure functional update on the
+    :class:`AlgorithmState` pytree; see the base-class contract.
+    """
+
+    COLUMNS = AugmentedLagrangian.COLUMNS
+
+    def __init__(
+        self,
+        gamma_1: float = 2.0,
+        gamma_2: float = 0.5,
+        eta_0: float = 1e-2,
+        eta_1: float = 1e-1,
+        eta_2: float = 0.8,
+        lam_prox_min: float = 1e-3,
+        lam_prox_max: float = 1e4,
+        lam_cost_drop: int = -1,
+        lam_cost_relax: float = 1.0,
+    ):
+        self.gamma_1 = gamma_1
+        self.gamma_2 = gamma_2
+        self.eta_0 = eta_0
+        self.eta_1 = eta_1
+        self.eta_2 = eta_2
+        self.lam_prox_min = lam_prox_min
+        self.lam_prox_max = lam_prox_max
+        self.lam_cost_drop = lam_cost_drop
+        self.lam_cost_relax = lam_cost_relax
+
+    def update_weights(
+        self,
+        state: "AlgorithmState",
+        candidate: "CandidateIterate",
+        nodal_constraints: "LoweredJaxConstraints",
+        settings: Config,
+        params: dict,
+    ) -> "AlgorithmState":
+        """Return the next-iterate state.
+
+        Pure functional update — see class docstring.
+        """
+        candidate_x_prop = candidate.x_prop_plus[1:]
+        nonlin_cost, nonlin_pen, nodal_pen = self.calculate_nonlinear_penalty(
+            candidate_x_prop,
+            candidate.x,
+            candidate.u,
+            state.lam_vc,
+            state.lam_vb_nodal,
+            state.lam_vb_cross,
+            state.lam_cost,
+            nodal_constraints,
+            params,
+            settings,
+        )
+        J_nonlin = nonlin_cost + nonlin_pen + nodal_pen
+
+        lam_cost_next = jnp.where(
+            state.k > self.lam_cost_drop,
+            state.lam_cost * self.lam_cost_relax,
+            state.lam_cost_init,
+        )
+
+        def first_iter(state):
+            return state.replace(
+                x=candidate.x,
+                u=candidate.u,
+                x_prop=candidate.x_prop,
+                x_prop_plus=candidate.x_prop_plus,
+                lam_cost=lam_cost_next,
+                J_nonlin=J_nonlin,
+                adaptive_state_code=jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=jnp.int32),
+            )
+
+        def later_iter(state):
+            prev_x_prop = state.x_prop_plus[1:]
+            prev_cost, prev_pen, prev_nodal_pen = self.calculate_nonlinear_penalty(
+                prev_x_prop,
+                state.x,
+                state.u,
+                state.lam_vc,
+                state.lam_vb_nodal,
+                state.lam_vb_cross,
+                state.lam_cost,
+                nodal_constraints,
+                params,
+                settings,
+            )
+            prev_J_nonlin = prev_cost + prev_pen + prev_nodal_pen
+
+            actual = prev_J_nonlin - J_nonlin
+            predicted = prev_J_nonlin - candidate.J_lin
+            safe_pred = jnp.where(predicted == 0.0, 1.0, predicted)
+            rho = jnp.where(predicted == 0.0, -jnp.inf, actual / safe_pred)
+
+            is_reject = rho < self.eta_0
+            is_accept_higher = (rho >= self.eta_0) & (rho < self.eta_1)
+            is_accept_constant = (rho >= self.eta_1) & (rho < self.eta_2)
+            accepted = ~is_reject
+
+            lp_higher = jnp.minimum(self.lam_prox_max, self.gamma_1 * state.lam_prox)
+            lp_lower = jnp.maximum(self.lam_prox_min, self.gamma_2 * state.lam_prox)
+            new_lam_prox = jnp.where(
+                is_reject | is_accept_higher,
+                lp_higher,
+                jnp.where(is_accept_constant, state.lam_prox, lp_lower),
+            )
+
+            code = jnp.where(
+                is_reject,
+                jnp.int32(AdaptiveStateCode.REJECT),
+                jnp.where(
+                    is_accept_higher,
+                    jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
+                    jnp.where(
+                        is_accept_constant,
+                        jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
+                        jnp.int32(AdaptiveStateCode.ACCEPT_LOWER),
+                    ),
+                ),
+            )
+
+            return state.replace(
+                x=jnp.where(accepted, candidate.x, state.x),
+                u=jnp.where(accepted, candidate.u, state.u),
+                x_prop=jnp.where(accepted, candidate.x_prop, state.x_prop),
+                x_prop_plus=jnp.where(accepted, candidate.x_prop_plus, state.x_prop_plus),
+                lam_prox=new_lam_prox,
+                # Virtual-control / virtual-buffer weights are held constant.
+                lam_cost=lam_cost_next,
+                J_nonlin=jnp.where(accepted, J_nonlin, state.J_nonlin),
+                predicted_reduction=predicted,
+                actual_reduction=actual,
+                acceptance_ratio=rho,
+                adaptive_state_code=code,
+            )
+
+        return jax.lax.cond(state.k == 1, first_iter, later_iter, state)
+
+
+# =============================================================================
+# Pydantic spec for dict / YAML validation
+# =============================================================================
+
+
+class AdaptiveProximalWeightSpec(BaseModel):
+    """Validates AdaptiveProximalWeight configuration from dict/YAML input."""
+
+    type: Literal["AdaptiveProximalWeight"] = "AdaptiveProximalWeight"
+    gamma_1: float = 2.0
+    gamma_2: float = 0.5
+    eta_0: float = 1e-2
+    eta_1: float = 1e-1
+    eta_2: float = 0.8
+    lam_prox_min: float = 1e-3
+    lam_prox_max: float = 1e4
+    lam_cost_drop: int = -1
+    lam_cost_relax: float = 1.0
+
+    model_config = ConfigDict(extra="forbid")
+
+    def build(self) -> AdaptiveProximalWeight:
+        return AdaptiveProximalWeight(**self.model_dump(exclude={"type"}, exclude_unset=True))