openscvx 0.5.3.dev19__tar.gz → 0.5.3.dev20__tar.gz

{openscvx-0.5.3.dev19/openscvx.egg-info → openscvx-0.5.3.dev20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev19
+Version: 0.5.3.dev20
 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.dev19 → openscvx-0.5.3.dev20}/docs/UnderTheHood/batching_jit_grad.md

@@ -49,6 +49,81 @@ 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.
 
+## Batched, disk-cached solves with `solve_batched()`
+
+`jax.vmap(solve_jax)` and `solve_batched()` both run `B` SCP solves over
+stacked boundary conditions; they differ in **who owns the batch axis**, and
+that one difference is the whole point.
+
+* `jax.vmap(solve_jax)` — *the caller* owns the axis. Maximally idiomatic,
+  composes with `grad` / `scan`, but **in-process JIT only**: `jax.export` has
+  no outer `vmap` rule (`call_exported` cannot be batched), so an exported
+  artifact can never be re-vmapped, and every fresh process re-traces and
+  re-compiles the entire loop.
+* `solve_batched(x0_stack, xf_stack)` — *the library* owns the axis. The vmap
+  is applied **inside** the method, before any export, so the whole loop lowers
+  to ordinary batched XLA ops and serializes cleanly. Under
+  `save_compiled=True` it is written to the solver cache on the first run and
+  deserialized on later processes — skipping the XLA compile that
+  `jax.vmap(solve_jax)` pays on every launch.
+
+```python
+problem = ox.Problem(..., solver={"backend": "qpax"})
+problem.settings.sim.save_compiled = True
+problem.initialize()
+
+# First process: traces, exports, writes <cache>/compiled_solve_batched_<hash>.jax
+# Later processes: deserialize-and-.call, no compile.
+batched = problem.solve_batched(x0_stack, xf_stack, params)
+# batched.x.shape == (B, N, n_states) — same pytree jax.vmap(solve_jax) returns.
+```
+
+Pick `solve_batched` when **cross-process cold-start dominates** — short-lived
+processes (CI sweeps, notebook restarts, deployed workers) that otherwise pay
+the compile every launch. Pick `jax.vmap(solve_jax)` for everything in-program,
+where you own the transforms and want `grad` / `scan` composition. With
+`save_compiled=False`, `solve_batched` is just an in-process batched wrapper and
+the two produce identical results.
+
+The exported artifact is closed and deployable, not composable: it is the
+`call_exported`-has-no-outer-`vmap`-rule wall again. Use `solve_jax` for
+user-driven `jax.vmap` / `jax.grad`; `solve_batched` for a cached batched solve.
+
+### CVXPy cannot be exported
+
+Export requires every op in the loop — the backend solve included — to lower to
+serializable XLA. CVXPy's `iteration_callback` is a `jax.pure_callback`, and
+`jax.export` cannot serialize host callbacks. So `PTRSolver.exportable` is
+`False` for CVXPy (and `True` for QPAX / Moreau), and
+`solve_batched(save_compiled=True)` on the CVXPy backend **raises a teaching
+error** pointing at QPAX / Moreau rather than silently degrading to an uncached
+in-process solve. With `save_compiled=False`, CVXPy still runs `B` sequential
+in-process solves.
+
+### Cache invalidation is correctness, not just freshness
+
+The exported loop closes over far more than the symbolic problem does. Caching
+on the symbolic hash alone would silently deserialize a **stale artifact** when
+any of that extra state changes — a wrong-answer bug, not a cache miss. So
+`get_solve_batched_cache_path` (`openscvx/utils/caching.py`) extends the
+symbolic hash with everything additionally baked into the loop, each runtime
+object contributing through its own `_hash_into` (the same protocol the symbolic
+layer uses):
+
+* the convex backend class + its `solver_args`,
+* the algorithm + autotuner + convergence thresholds + initial weights,
+* the discretizer scheme (class + hold type + ODE solver),
+* the state/control scaling matrices `inv_S_x` / `inv_S_u`,
+* the fixed batch size `B` (the artifact is exported at one `B`), and
+* the JAX version (exported artifacts are not guaranteed to survive an
+  incompatible jax bump).
+
+Change any of these — backend, a tolerance, `k_max`, a state bound — and you get
+a different cache path, so a stale artifact is never reused.
+
+`examples/abstract/brachistochrone_batched.py` shows both paths side by side;
+re-running it is itself the cross-process demo.
+
 ## What `solve_jax()` returns
 
 `solve_jax()` returns the same `OptimizationResults` type as `solve()`,

{openscvx-0.5.3.dev19 → openscvx-0.5.3.dev20}/examples/abstract/brachistochrone_batched.py

@@ -1,20 +1,30 @@
-"""Batched brachistochrone via ``jax.vmap(problem.solve_jax)``.
+"""Batched brachistochrone — two ways to solve four problems at once.
 
 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``.
+starting x-coordinate — and contrasts the two batching entry points:
+
+* ``jax.vmap(problem.solve_jax)`` — the caller owns the batch axis. Maximally
+  idiomatic, composes with ``grad`` / ``scan``, but in-process JIT only: every
+  fresh process re-traces and re-compiles the whole SCP loop.
+* ``problem.solve_batched(x0_stack, xf_stack)`` — the library owns the batch
+  axis, so the whole vmapped loop is a single function that ``jax.export`` can
+  serialize. Under ``save_compiled=True`` it is written to the solver cache on
+  the first run and deserialized on later runs, skipping that compile. Reach
+  for it when cross-process cold-start dominates (CI sweeps, short-lived
+  workers); reach for ``jax.vmap(solve_jax)`` when you stay inside one program.
+
+Run with ``python examples/abstract/brachistochrone_batched.py``. Re-running is
+itself the cross-process demo: the second invocation deserializes the cached
+``solve_batched`` artifact instead of compiling it.
 
 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.
+  for thread safety — see :meth:`solve_jax`'s docstring. That same host
+  callback is why ``solve_batched(save_compiled=True)`` refuses CVXPy: it
+  cannot be exported.
 * 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.
@@ -113,14 +123,33 @@ if __name__ == "__main__":
     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.
+    # --- caller owns the batch axis: jax.vmap(solve_jax) ---------------------
+    # Bare ``jax.vmap`` — composes like any JAX transform, in-process only.
     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"jax.vmap(solve_jax) 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)}")
+
+    # --- library owns the batch axis: solve_batched, disk-cached -------------
+    # The same batched solve, but ``solve_batched`` owns the vmap so the whole
+    # loop is one exportable artifact. Under ``save_compiled=True`` it lands in
+    # the solver cache on the first run; re-run this script (a fresh process) to
+    # see it deserialized instead of recompiled. ``xf`` is shared across the
+    # batch, so broadcast the default terminal pin to the same leading axis.
+    export_problem = build_problem()
+    export_problem.settings.sim.save_compiled = True
+    export_problem.initialize()
+
+    x_term_stack = jnp.broadcast_to(export_problem.state.x_term_pin, x_initial_stack.shape)
+    batched = export_problem.solve_batched(x_initial_stack, x_term_stack)
+
+    print(f"\nsolve_batched (save_compiled=True) over {x_initial_stack.shape[0]} ICs:")
+    print(f"  result.x.shape:   {batched.x.shape}")
+    print(f"  result.t_final:   {np.asarray(batched.t_final).reshape(-1)}")
+    print(f"  result.converged: {np.asarray(batched.converged)}")

openscvx-0.5.3.dev20/examples/drone/2d_obstacle_avoidance_batched_ic.py

@@ -0,0 +1,326 @@
+"""Batched 2D obstacle avoidance: solve over B random initial positions.
+
+Same planar double-integrator + CTCS obstacle-avoidance problem as
+``2d_obstacle_avoidance.py``, but ``solve_batched`` runs all B solves in one
+compiled dispatch (Moreau backend).  The initial position is a ``Parameter``
+pinned via ``(position == initial_position).convex().at([0])``.  Pass
+``parameters={"initial_position": ic_batch}`` (shape ``(B, 2)``) to pin each
+IC at node 0.  Per-batch ``x_guess`` stacks
+(a position linspace from each IC to the shared target) seed the SCP iterate
+independently.  ``post_process_batched`` then propagates
+every converged solution through the full nonlinear dynamics.
+
+All B trajectories are visualised together with Plotly — converged runs in
+blue, diverged in orange — on top of the obstacle field.
+"""
+
+import os
+import sys
+
+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
+
+# ── Batch configuration ────────────────────────────────────────────────────────
+N_BATCH = 200  # number of initial conditions to solve simultaneously
+N = 40  # SCP discretisation nodes
+TOTAL_TIME = 120.0
+
+# ── Obstacle field (identical to 2d_obstacle_avoidance.py) ────────────────────
+obstacle_radius_min, obstacle_radius_max = 1.0, 2.5
+
+np.random.seed(42)
+obstacle_centers_list = []
+n_rows, n_cols = 20, 20
+for i in range(n_rows):
+    for j in range(n_cols):
+        x = -6.0 + i * 6.0
+        y = -7.5 + j * 5.0
+        x += np.random.uniform(-1.0, 1.0)
+        y += np.random.uniform(-1.0, 1.0)
+        obstacle_centers_list.append([x, y])
+
+obstacle_centers = np.array(obstacle_centers_list)  # (n_obs, 2)
+obstacle_radii = np.random.uniform(
+    obstacle_radius_min, obstacle_radius_max, size=len(obstacle_centers_list)
+)
+
+# ── States ─────────────────────────────────────────────────────────────────────
+_default_ic = np.array([-10.0, -10.0])  # used for guess; actual IC set by ZeroConeConstraint
+
+position = ox.State("position", shape=(2,))
+position.max = np.array([150.0, 150.0])
+position.min = np.array([-15.0, -15.0])
+position.initial = [ox.Free(-10.0), ox.Free(-10.0)]
+position.final = np.array([100.0, 100.0])
+
+velocity = ox.State("velocity", shape=(2,))
+velocity.max = np.array([10.0, 10.0])
+velocity.min = np.array([-10.0, -10.0])
+velocity.initial = np.array([0.0, 0.0])
+velocity.final = [("free", 0.0), ("free", 0.0)]
+
+# ── Controls ───────────────────────────────────────────────────────────────────
+force = ox.Control("force", shape=(2,))
+a_max = 20.0
+force.max = np.array([a_max, a_max])
+force.min = np.array([-a_max, -a_max])
+
+m = 1.0
+
+# ── Batched parameter: initial position ───────────────────────────────────────
+initial_position = ox.Parameter("initial_position", shape=(2,), value=_default_ic)
+
+# ── Dynamics ───────────────────────────────────────────────────────────────────
+dynamics = {
+    "position": velocity,
+    "velocity": (1.0 / m) * force,
+}
+
+# ── Constraints ────────────────────────────────────────────────────────────────
+states = [position, velocity]
+controls = [force]
+constraints = []
+
+for state in states:
+    constraints.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
+
+constraints.extend([force <= force.max, force.min <= force])
+
+obstacle_avoidance = ox.ctcs(
+    obstacle_radii
+    <= ox.Vmap(
+        lambda obs_center: ox.linalg.Norm(position - obs_center),
+        batch=obstacle_centers,
+    )
+)
+constraints.append(obstacle_avoidance)
+
+# ── Constraints ── initial position pinned via convex equality at node 0 ───────
+constraints.append((position == initial_position).convex().at([0]))
+
+# ── Initial guesses ────────────────────────────────────────────────────────────
+position.guess = np.linspace(_default_ic, position.final, N)
+velocity.guess = np.zeros((N, 2))
+force.guess = np.zeros((N, 2))
+
+# ── Time ───────────────────────────────────────────────────────────────────────
+time_var = ox.Time(
+    initial=0.0,
+    final=("minimize", TOTAL_TIME),
+    min=0.0,
+    max=TOTAL_TIME,
+)
+
+# ── Problem ────────────────────────────────────────────────────────────────────
+# Moreau backend: JAX-native, exportable, required for solve_batched.
+problem = Problem(
+    dynamics=dynamics,
+    states=states,
+    controls=controls,
+    time=time_var,
+    constraints=constraints,
+    N=N,
+    algorithm={
+        "autotuner": ox.RampProximalWeight(ramp_factor=1.1, lam_prox_max=1e3),
+        "lam_cost": 1e-2,
+        "lam_vc": 1e1,
+    },
+    solver=ox.MoreauPTRSolver(),
+    float_dtype="float64",
+)
+
+
+# ── Helper: sample valid initial positions ─────────────────────────────────────
+def sample_initial_positions(
+    B: int,
+    seed: int = 0,
+    x_range: tuple[float, float] = (-14.0, 20.0),
+    y_range: tuple[float, float] = (-14.0, 20.0),
+    safety_margin: float = 0.5,
+) -> np.ndarray:
+    """Sample B positions in [x_range × y_range] that lie outside all obstacles.
+
+    The default range covers the lower-left quadrant of the obstacle field,
+    giving varied starting positions from far-corner approaches to positions
+    already adjacent to the first few rows of obstacles.
+
+    Args:
+        B: Number of positions to sample.
+        seed: RNG seed.
+        x_range: (min, max) for the x coordinate.
+        y_range: (min, max) for the y coordinate.
+        safety_margin: Extra clearance beyond the obstacle radius.
+
+    Returns:
+        Array of shape ``(B, 2)``.
+    """
+    rng = np.random.default_rng(seed)
+    positions: list[np.ndarray] = []
+    while len(positions) < B:
+        p = rng.uniform(
+            [x_range[0], y_range[0]],
+            [x_range[1], y_range[1]],
+        )
+        dists = np.linalg.norm(obstacle_centers - p, axis=1)
+        if np.all(dists > obstacle_radii + safety_margin):
+            positions.append(p)
+    return np.array(positions)
+
+
+def build_batched_x_guess(problem, ic_batch: np.ndarray) -> np.ndarray:
+    """Build per-batch state guesses with position linspace IC → target."""
+    base_x = np.asarray(problem.state.x)  # (N, n_x)
+    pos_sl = position._slice
+    final_pos = np.asarray(position.final)
+    B = ic_batch.shape[0]
+    x_guess_batch = np.broadcast_to(base_x, (B,) + base_x.shape).copy()
+    t = np.linspace(0.0, 1.0, N)
+    x_guess_batch[:, :, pos_sl] = (
+        ic_batch[:, None, :] * (1.0 - t[None, :, None])
+        + final_pos[None, None, :] * t[None, :, None]
+    )
+    return x_guess_batch
+
+
+# ── Visualisation ──────────────────────────────────────────────────────────────
+def plot_batched_2d_trajectories(results, ic_batch: np.ndarray):
+    """Plotly figure overlaying all B propagated trajectories on the obstacle field.
+
+    Args:
+        results: :class:`~openscvx.algorithms.OptimizationResults` returned by
+            :meth:`~openscvx.problem.Problem.post_process_batched`.  Must have
+            ``trajectory["position"]`` of shape ``(B, T, 2)``.
+        ic_batch: ``(B, 2)`` array of initial positions used for the batch.
+    """
+    import plotly.graph_objects as go
+
+    converged = np.asarray(results.converged, dtype=bool).reshape(-1)
+    B = len(converged)
+
+    # Extract propagated positions: (B, T, 2)
+    pos_all = np.asarray(results.trajectory["position"], dtype=np.float64)
+    if pos_all.ndim == 2:
+        # Fallback: single-batch dim was squeezed
+        pos_all = pos_all[np.newaxis]
+
+    fig = go.Figure()
+
+    # ── Background obstacles ──────────────────────────────────────────────────
+    for center, radius in zip(obstacle_centers, obstacle_radii):
+        xc, yc = center
+        fig.add_shape(
+            type="circle",
+            x0=xc - radius,
+            y0=yc - radius,
+            x1=xc + radius,
+            y1=yc + radius,
+            fillcolor="rgba(220, 180, 180, 0.30)",
+            line={"color": "rgba(200, 120, 120, 0.45)", "width": 1},
+            layer="below",
+        )
+
+    # ── Trajectories ──────────────────────────────────────────────────────────
+    # Draw diverged runs first so converged runs render on top.
+    for b in range(B):
+        ok = bool(converged[b])
+        color = "rgba(30, 130, 220, 0.65)" if ok else "rgba(220, 100, 30, 0.50)"
+        pos = pos_all[b]  # (T, 2)
+        fig.add_trace(
+            go.Scatter(
+                x=pos[:, 0],
+                y=pos[:, 1],
+                mode="lines",
+                line={"color": color, "width": 1.5 if ok else 1.0},
+                showlegend=False,
+                hovertemplate=f"run {b}  {'✓' if ok else '✗'}<extra></extra>",
+            )
+        )
+
+    # ── Start markers (one per batch element) ─────────────────────────────────
+    starts_ok = ic_batch[converged]
+    starts_bad = ic_batch[~converged]
+    if len(starts_ok):
+        fig.add_trace(
+            go.Scatter(
+                x=starts_ok[:, 0],
+                y=starts_ok[:, 1],
+                mode="markers",
+                name="Start (converged)",
+                marker={"color": "rgba(20, 100, 200, 0.9)", "size": 7, "symbol": "circle"},
+            )
+        )
+    if len(starts_bad):
+        fig.add_trace(
+            go.Scatter(
+                x=starts_bad[:, 0],
+                y=starts_bad[:, 1],
+                mode="markers",
+                name="Start (diverged)",
+                marker={"color": "rgba(200, 80, 20, 0.9)", "size": 7, "symbol": "circle-open"},
+            )
+        )
+
+    # ── Target ────────────────────────────────────────────────────────────────
+    tgt = position.final
+    fig.add_trace(
+        go.Scatter(
+            x=[tgt[0]],
+            y=[tgt[1]],
+            mode="markers",
+            name="Target",
+            marker={"color": "black", "size": 12, "symbol": "star"},
+        )
+    )
+
+    # ── Layout ────────────────────────────────────────────────────────────────
+    n_ok = int(converged.sum())
+    fig.update_layout(
+        title=f"Batched 2-D obstacle avoidance  —  {B} ICs  |  {n_ok}/{B} converged",
+        xaxis_title="x  [m]",
+        yaxis_title="y  [m]",
+        yaxis={"scaleanchor": "x", "scaleratio": 1},
+        template="simple_white",
+        paper_bgcolor="white",
+        plot_bgcolor="white",
+        autosize=False,
+        width=720,
+        height=720,
+        margin={"l": 60, "r": 30, "t": 50, "b": 60},
+    )
+    return fig
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    print(f"Initializing 2-D obstacle avoidance problem (Moreau, B={N_BATCH}) ...")
+    problem.initialize()
+
+    # Sample B valid initial positions
+    ic_batch = sample_initial_positions(N_BATCH)
+    print(f"Sampled {N_BATCH} initial positions in x∈[-14,20], y∈[-14,20]")
+    print(f"  ic_batch.shape = {ic_batch.shape}")
+
+    # ── Compile + solve ───────────────────────────────────────────────────────
+    # B is inferred from ic_batch.shape[0].  solve_batched vmaps over
+    # initial_position (shape (B, 2)) and the per-batch x_guess stack.
+    x_guess_batch = build_batched_x_guess(problem, ic_batch)
+    print("Compiling and running solve_batched …")
+    results = problem.solve_batched(
+        parameters={"initial_position": ic_batch},
+        x_guess=x_guess_batch,
+    )
+
+    # ── Post-process ──────────────────────────────────────────────────────────
+    print("Post-processing (nonlinear propagation) …")
+    results = problem.post_process_batched(results)
+
+    # ── Plot ──────────────────────────────────────────────────────────────────
+    print("Plotting …")
+    fig = plot_batched_2d_trajectories(results, ic_batch)
+    fig.show()