openscvx 0.5.3.dev15__tar.gz → 0.5.3.dev17__tar.gz

{openscvx-0.5.3.dev15/openscvx.egg-info → openscvx-0.5.3.dev17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev15
+Version: 0.5.3.dev17
 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.3.dev17/docs/UnderTheHood/batching_jit_grad.md

@@ -0,0 +1,134 @@
+# Batching, JIT, and Grad with `Problem.solve_jax()`
+
+`Problem` exposes two solve entry points. `solve()` is the familiar
+Python-loop driver — real-time prints, wall-clock `time_limit`, `continuous`
+mode, populated per-iteration history. `solve_jax()` is its JAX-pure
+sibling: it drives the same fused `iteration_fn` body inside a
+`lax.while_loop` and returns an `OptimizationResults`
+(`openscvx/algorithms/optimization_results.py`) pytree, so it composes with
+`jax.vmap`, `jax.jit`, and `jax.grad`.
+
+```python
+problem = ox.Problem(..., solver={"backend": "qpax"})
+problem.initialize()
+
+# Familiar interactive solve — unchanged: prints, history, time limit.
+result = problem.solve()
+
+# JAX-pure single solve — silent, no per-iteration history, batchable.
+result = problem.solve_jax()
+
+# Batched — standard jax.vmap, no library-specific API.
+batched = jax.vmap(problem.solve_jax, in_axes=(0, 0, None))(
+    x_initial_stack, x_final_stack, params
+)
+
+# JIT-compile once, solve forever (e.g. MPC inner loop).
+fast_solve = jax.jit(problem.solve_jax)
+```
+
+A worked example lives at `examples/abstract/brachistochrone_batched.py`
+— four brachistochrone problems with shifted starting x-coordinates,
+solved in parallel via `jax.vmap(problem.solve_jax)`.
+
+## When to choose `solve()` vs. `solve_jax()`
+
+| Use case | Entry point |
+|---|---|
+| Interactive use, real-time prints, plotting per iteration | `solve()` |
+| Wall-clock `time_limit` or `continuous=True` | `solve()` |
+| Per-iteration trajectories / weights / diagnostics | `solve()` |
+| Batched solve over many initial conditions or parameters | `solve_jax()` |
+| JIT-compile once, solve forever (MPC inner loop, scenario sweeps) | `solve_jax()` |
+| `jax.grad` through the solver (best-effort, untested) | `solve_jax()` |
+
+The split exists to make user intent explicit — a single dispatched
+`solve()` would have to infer from arguments which path the caller wants,
+and silent routing failures (like `continuous=True` quietly skipping the
+Python loop) are exactly the failure class the two-method design exists to
+prevent. See `plans/jax-pure-solve.md`'s Decision Log for the longer
+discussion.
+
+## What `solve_jax()` returns
+
+`solve_jax()` returns the same `OptimizationResults` type as `solve()`,
+but built via
+`OptimizationResults.from_final_state(state, problem=...)` instead of
+`from_history(history, final_state, ...)`. The differences:
+
+* **Per-iteration history is empty.** `X = [state.x]` and `U = [state.u]`
+  are single-element lists so `result.x` / `result.u` continue to return
+  the final iterate; every `*_history` field is `[]`. List growth doesn't
+  fit inside `lax.while_loop`. If you need per-iteration trajectories, use
+  `solve()` or wrap `solve_jax()` in `lax.scan` manually.
+* **Post-process fields stay `None`.** `t_full`, `x_full`, `u_full`,
+  `cost`, `ctcs_violation` are populated only by `post_process()`. They're
+  outside the JAX pytree's children surface, so a batched `solve_jax`
+  result doesn't force every consumer to handle `None` leaves. If you want
+  to post-process, call `post_process()` per batch element after the
+  batched solve.
+* **`converged` is a `jnp.bool_`** under the JAX-pure path (a `(B,)` array
+  under vmap), not a Python `bool` like `solve()` returns. Most uses
+  (`if result.converged: ...`) work either way.
+
+## `solve_jax()` arguments
+
+```python
+problem.solve_jax(
+    x_initial=None,      # boundary-condition pin (full unified state vector,
+                         #     ``jnp.nan`` at non-Fix entries — see
+                         #     ``AlgorithmState.from_settings``); falls back
+                         #     to the default from settings
+    x_final=None,        # terminal pin, same conventions
+    parameters=None,     # parameters dict for this solve; falls back to
+                         #     ``self._parameters``
+    *,
+    max_iters=None,      # SCP iteration cap; non-default rebuilds the
+                         #     cached ``lax.while_loop`` closure (one extra
+                         #     trace; subsequent calls at the same cap hit
+                         #     the cache)
+)
+```
+
+Positional kwargs (rather than a single `inputs` pytree) keep
+`jax.vmap(problem.solve_jax, in_axes=(0, 0, None))` ergonomic;
+multi-argument gradient uses `jax.grad(..., argnums=(0, 1))`.
+
+## Caveats
+
+### CVXPy under `vmap` is sequential
+
+The CVXPy backend's `iteration_callback`
+(`openscvx/solvers/cvxpy_ptr_solver.py`) host-calls CVXPy through
+`jax.pure_callback` with `vmap_method="sequential"`. Host
+CVXPy is not thread-safe, so a `jax.vmap(problem.solve_jax)` over the
+CVXPy backend runs `B` sequential CVXPy solves. The QPAX and Moreau
+backends are pure JAX end-to-end and run in parallel under vmap.
+
+### Moreau warm-start is bypassed
+
+The Moreau backend's `_warm_start` carry is host-side mutable state that
+`lax.while_loop` doesn't thread. Both `solve()` and `solve_jax()` cold-start
+the inner Moreau solve every SCP iteration (see the Moreau module
+docstring and `plans/jax-pure-solve.md`'s Decision Log 2026-05-27).
+Restoring warm-start to the SCP loop requires threading `(x, z, s)`
+through an `AlgorithmState.moreau_carry` field — a Future Extension.
+
+### `jax.grad` is best-effort, untested
+
+The QPAX backend's `solve_qp` is not differentiable (its convergence flag
+costs reverse-mode autodiff). The CVXPy backend's `pure_callback` is
+non-differentiable by default. End-to-end gradient validation through
+`solve_jax` is a follow-up — `jax.grad(loss_fn)(x0)` will run without
+erroring under QPAX (`solve_qp_primal` would be `custom_vjp`, but the
+backend currently uses `solve_qp`), but the gradient's correctness is not
+yet pinned down by a test against finite differences.
+
+### Convergence under `vmap` with mixed-rate elements
+
+`lax.while_loop` continues while *any* batch element still needs
+iterations — converged elements would otherwise keep receiving body calls
+while their peers iterate. `make_solve_loop`'s body selects the unchanged
+state for converged elements (`jax.tree.map(jnp.where, ...)`), so a
+batched solve agrees with the single-problem `solve_jax` on each element.
+The cost is per-iteration work × slowest-element iteration count.

{openscvx-0.5.3.dev15 → openscvx-0.5.3.dev17}/docs/UnderTheHood/lowering_architecture.md

@@ -168,18 +168,33 @@ of the abstract `PTRSolver` base. Three 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 (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. |
+| 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. The `iteration_callback` uses `qpax.solve_qp`, whose convergence flag maps to a `status_code` so a diverged solve fails loudly (chosen over the differentiable `solve_qp_primal`, which exposes no such flag). |
 | 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
+`update_*` + `solve()` stages (kept for direct / interactive solver use), 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.
+returns a `(state, SubproblemData) -> SubproblemSolution` callable; all three
+backends emit the same input/output pytree shape (CVXPy host-calls via
+`jax.pure_callback`).
+
+`Problem.initialize()` fuses that callback with discretization, constraint
+linearization, the SCP metrics, and the autotuner into one JAX function —
+`iteration_fn`, built by `make_scp_iteration`
+(`openscvx/algorithms/scvx/iteration.py`). A single JIT'd `iteration_fn` call
+per iteration replaces the former per-step NumPy↔JAX stitching, and the Python
+SCP loop now drives that one fused body. Forward propagation stays a separate
+JAX export, used only by `post_process()` — so a compiled problem holds **one
+`iteration_fn` plus one propagation export**, not three discretization /
+propagation exports stitched together in Python. Because `iteration_fn` is a
+registered pytree in and out, it also composes with `jax.jit`; `make_solve_loop`
+wraps it in `lax.while_loop` to drive the **JAX-pure** `Problem.solve_jax()`
+entry point, which composes with `jax.vmap`, `jax.jit`, and `jax.grad` (see
+`docs/UnderTheHood/batching_jit_grad.md`). The familiar
+`Problem.solve()` keeps its Python-loop shape — real-time prints, wall-clock
+`time_limit`, populated per-iteration history — and is the right entry
+point for interactive use. `solve_jax()` is the right entry point for
+batched / JIT'd / differentiable workflows.
 
 Picking a backend at construction time:
 

openscvx-0.5.3.dev17/examples/abstract/brachistochrone_batched.py

@@ -0,0 +1,126 @@
+"""Batched brachistochrone via ``jax.vmap(problem.solve_jax)``.
+
+Solves four brachistochrone problems in parallel — each with a different
+starting x-coordinate — using the JAX-pure
+:meth:`~openscvx.problem.Problem.solve_jax` entry point. ``solve_jax``
+returns an :class:`~openscvx.algorithms.optimization_results.OptimizationResults`
+pytree, so ``jax.vmap`` over the boundary-condition pins is the standard JAX
+composition pattern; no library-specific batching API is needed.
+
+Run with ``python examples/abstract/brachistochrone_batched.py``.
+
+Notes:
+
+* Under the QPAX backend used here, vmap'd subproblem solves run in
+  parallel (no host callback). The CVXPy backend is single-threaded under
+  vmap because :func:`jax.pure_callback` uses ``vmap_method="sequential"``
+  for thread safety — see :meth:`solve_jax`'s docstring.
+* The Moreau warm-start is bypassed on the JAX-pure path (its carry is
+  host-side state that ``lax.while_loop`` doesn't thread). Use QPAX or
+  CVXPy when you want batched solves.
+"""
+
+import os
+import sys
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+grandparent_dir = os.path.dirname(os.path.dirname(current_dir))
+sys.path.append(grandparent_dir)
+
+import openscvx as ox
+from openscvx import Problem
+
+
+def build_problem(n: int = 8):
+    g = 9.81
+
+    position = ox.State("position", shape=(2,))
+    position.max = np.array([10.0, 10.0])
+    position.min = np.array([0.0, 0.0])
+    position.initial = np.array([0.0, 10.0])
+    position.final = [10.0, 5.0]
+
+    velocity = ox.State("velocity", shape=(1,))
+    velocity.max = np.array([10.0])
+    velocity.min = np.array([0.0])
+    velocity.initial = np.array([0.0])
+    velocity.final = [("free", 10.0)]
+
+    theta = ox.Control("theta", shape=(1,))
+    theta.max = np.array([100.5 * jnp.pi / 180])
+    theta.min = np.array([0.0])
+    theta.guess = np.linspace(5 * jnp.pi / 180, 100.5 * jnp.pi / 180, n).reshape(-1, 1)
+
+    states = [position, velocity]
+    controls = [theta]
+
+    dynamics = {
+        "position": ox.Concat(
+            velocity[0] * ox.Sin(theta[0]),
+            -velocity[0] * ox.Cos(theta[0]),
+        ),
+        "velocity": g * ox.Cos(theta[0]),
+    }
+
+    constraint_exprs = []
+    for state in states:
+        constraint_exprs.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
+
+    time = ox.Time(
+        initial=0.0,
+        final=("minimize", 2.0),
+        min=0.0,
+        max=2.0,
+        uniform_time_grid=True,
+    )
+
+    return Problem(
+        dynamics=dynamics,
+        states=states,
+        controls=controls,
+        time=time,
+        constraints=constraint_exprs,
+        N=n,
+        float_dtype="float64",
+        algorithm={
+            "autotuner": "ConstantProximalWeight",
+            "lam_prox": 1e0,
+            "lam_cost": 6e-1,
+            "ep_tr": 1e-5,
+            "ep_vb": 1e-5,
+            "ep_vc": 1e-9,
+            "k_max": 30,
+        },
+        solver={"backend": "qpax"},
+    )
+
+
+if __name__ == "__main__":
+    problem = build_problem()
+    problem.initialize()
+
+    # The default boundary-condition pin (the full unified state vector with
+    # ``jnp.nan`` at non-Fix entries — see
+    # :meth:`AlgorithmState.from_settings`).
+    default_pin = problem.state.x_init_pin
+
+    # Stack four ICs by varying the starting x-coordinate (component 0 of
+    # the unified state vector is the x-position).
+    shifts = jnp.array([0.0, 0.3, 0.6, 0.9])
+    x_initial_stack = jnp.stack([default_pin.at[0].set(default_pin[0] + s) for s in shifts])
+
+    # Bare ``jax.vmap`` — the library exposes no batching wrapper of its own.
+    batched_solve = jax.vmap(problem.solve_jax, in_axes=(0, None, None))
+    results = batched_solve(x_initial_stack, None, None)
+
+    # ``results`` is a batched ``OptimizationResults`` pytree — every child
+    # (``X``, ``U``, ``t_final``, ``converged``, ...) has a leading batch axis.
+    print(f"Batched solve over {x_initial_stack.shape[0]} initial conditions:")
+    print(f"  result.x.shape:   {results.x.shape}")
+    print(f"  result.u.shape:   {results.u.shape}")
+    print(f"  result.t_final:   {np.asarray(results.t_final).reshape(-1)}")
+    print(f"  result.converged: {np.asarray(results.converged)}")

{openscvx-0.5.3.dev15 → openscvx-0.5.3.dev17}/mkdocs.yml

@@ -197,4 +197,5 @@ nav:
       - 'Foundations/discretization.md'
     - Under the Hood:
       - 'UnderTheHood/vectorization_and_vmapping.md'
-      - 'UnderTheHood/lowering_architecture.md'
+      - 'UnderTheHood/lowering_architecture.md'
+      - 'UnderTheHood/batching_jit_grad.md'

{openscvx-0.5.3.dev15 → openscvx-0.5.3.dev17}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.3.dev15'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev15')
+__version__ = version = '0.5.3.dev17'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev17')
 
-__commit_id__ = commit_id = 'ge47b6d19c'
+__commit_id__ = commit_id = 'gb2143958f'

{openscvx-0.5.3.dev15 → openscvx-0.5.3.dev17}/openscvx/algorithms/base.py

@@ -29,7 +29,6 @@ from openscvx.utils.printing import Column
 if TYPE_CHECKING:
     from openscvx.config import Config
     from openscvx.lowered.jax_constraints import LoweredJaxConstraints
-    from openscvx.solvers import ConvexSolver
 
     from .weights import Weights
 
@@ -373,6 +372,16 @@ class AlgorithmState:
         actual_reduction: Actual reduction in ``J_nonlin`` for this iter.
         acceptance_ratio: ``actual_reduction / predicted_reduction``.
         adaptive_state_code: :class:`AdaptiveStateCode` value as ``int32``.
+        x_init_pin: Initial-state boundary condition, shape ``(n_states,)``.
+            ``jnp.nan`` where the state is not pinned at t0 (``initial_type``
+            is not ``"Fix"``). Carried on the pytree — rather than read from
+            ``settings`` — so the fused SCP iteration body assembles the
+            subproblem's initial boundary rows as a pure function of state,
+            and a future ``jax.vmap`` over problems can batch boundary
+            conditions per element.
+        x_term_pin: Terminal-state boundary condition, shape ``(n_states,)``.
+            ``jnp.nan`` where the state is not pinned at tf (``final_type`` is
+            not ``"Fix"``).
     """
 
     x: jnp.ndarray
@@ -394,6 +403,8 @@ class AlgorithmState:
     actual_reduction: jnp.ndarray
     acceptance_ratio: jnp.ndarray
     adaptive_state_code: jnp.ndarray
+    x_init_pin: jnp.ndarray
+    x_term_pin: jnp.ndarray
 
     # Field order is the source of truth for tree_flatten / tree_unflatten;
     # keep _FIELDS in sync with the dataclass field declarations above.
@@ -417,6 +428,8 @@ class AlgorithmState:
         "actual_reduction",
         "acceptance_ratio",
         "adaptive_state_code",
+        "x_init_pin",
+        "x_term_pin",
     )
 
     def replace(self, **changes) -> "AlgorithmState":
@@ -475,6 +488,17 @@ class AlgorithmState:
         else:
             lam_cost_init = np.full(n_states, weights.lam_cost)
 
+        # Boundary-condition pins: the physical value where the state is fixed
+        # at t0 / tf, ``nan`` elsewhere. The subproblem only reads these at
+        # ``"Fix"`` entries, so the sentinel marks "unpinned" without poisoning
+        # any value the solver consumes.
+        x_initial = np.asarray(settings.sim.x.initial, dtype=float).reshape(-1)
+        x_final = np.asarray(settings.sim.x.final, dtype=float).reshape(-1)
+        init_fixed = np.asarray(settings.sim.x.initial_type) == "Fix"
+        final_fixed = np.asarray(settings.sim.x.final_type) == "Fix"
+        x_init_pin = np.where(init_fixed, x_initial, np.nan)
+        x_term_pin = np.where(final_fixed, x_final, np.nan)
+
         return cls(
             x=put(jnp.asarray(settings.sim.x.guess, dtype=f)),
             u=put(jnp.asarray(settings.sim.u.guess, dtype=f)),
@@ -495,6 +519,8 @@ class AlgorithmState:
             actual_reduction=put(jnp.asarray(0.0, dtype=f)),
             acceptance_ratio=put(jnp.asarray(0.0, dtype=f)),
             adaptive_state_code=put(jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=i)),
+            x_init_pin=put(jnp.asarray(x_init_pin, dtype=f)),
+            x_term_pin=put(jnp.asarray(x_term_pin, dtype=f)),
         )
 
 
@@ -798,15 +824,21 @@ class Algorithm(ABC):
     @abstractmethod
     def initialize(
         self,
-        solver: "ConvexSolver",
-        discretization_solver: callable,
-        jax_constraints: "LoweredJaxConstraints",
+        iteration_fn: Callable,
         emitter: callable,
-        params: dict,
+        jax_constraints: "LoweredJaxConstraints",
         settings: "Config",
-        discretization_solver_impulsive: Optional[Callable] = None,
     ) -> None:
-        """Initialize the algorithm and store compiled infrastructure."""
+        """Store the fused SCP iteration body and per-iteration infrastructure.
+
+        Args:
+            iteration_fn: The JAX-pure ``(state, params) -> (next_state,
+                diagnostics)`` body built by
+                :func:`~openscvx.algorithms.scvx.iteration.make_scp_iteration`.
+            emitter: Per-iteration diagnostics sink (printing queue / no-op).
+            jax_constraints: Lowered JAX constraints the body operates over.
+            settings: Problem configuration.
+        """
         raise NotImplementedError
 
     @abstractmethod