openscvx 0.5.3.dev30__tar.gz → 0.5.3.dev31__tar.gz

{openscvx-0.5.3.dev30/openscvx.egg-info → openscvx-0.5.3.dev31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev30
+Version: 0.5.3.dev31
 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.dev30 → openscvx-0.5.3.dev31}/examples/abstract/flappy_bird.py

@@ -193,34 +193,17 @@ plotting_dict = {
 
 def _extract_multishot_position_segments(results) -> list[np.ndarray]:
     """Extract per-segment propagated states from SCP discretization history."""
-    v_history = getattr(results, "discretization_history", None) or []
-    if not v_history:
+    prop = results.multishot_propagation()
+    if prop is None:
         return []
-
-    position_slice = None
-    for state in results._states:
+    pos_slice = None
+    for state in prop.states:
         if state.name == "position":
-            position_slice = state._slice
+            pos_slice = state._slice
             break
-    if position_slice is None:
+    if pos_slice is None:
         return []
-
-    n_x = results.x.shape[1]
-    n_u = results.u.shape[1]
-    i4 = n_x + n_x * n_x + 2 * n_x * n_u
-
-    v = np.asarray(v_history[-1])
-    n_segments = v.shape[0] // i4
-    segments = []
-    for seg_idx in range(n_segments):
-        seg_start = seg_idx * i4
-        seg_end = seg_start + i4
-        seg_positions = []
-        for t_idx in range(v.shape[1]):
-            state = v[seg_start:seg_end, t_idx][:n_x]
-            seg_positions.append(state[position_slice])
-        segments.append(np.asarray(seg_positions))
-    return segments
+    return [seg[:, pos_slice] for seg in prop.segments()]
 
 
 def _build_singleshot_propagated_trajectory(

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/drone/logo.py

@@ -388,7 +388,7 @@ if __name__ == "__main__":
 
     results.update_plotting_data(**plotting_dict)
 
-    # Define the plane the logo actually lies in (from three points on the logo)
+    # Define the plane the logo actually lies in
     p0 = np.asarray(get_kp_pose(0.0))
     p1 = np.asarray(get_kp_pose(0.33))
     p2 = np.asarray(get_kp_pose(0.67))

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/frax/panda_frax_pick_place.py

@@ -248,43 +248,6 @@ def _compute_panda_keypoints(
     return keypoints, ee_pos
 
 
-def _q_from_V_multishot(
-    V: np.ndarray,
-    n_x: int,
-    n_u: int,
-    q_slice: slice,
-    t_nodes: np.ndarray,
-) -> tuple[np.ndarray | None, np.ndarray | None]:
-    """Unpack joint angles from the SCP multi-shoot V matrix.
-
-    V has shape ``((N-1) * i4, n_substeps)`` where
-    ``i4 = n_x + n_x² + 2·n_x·n_u``.  The first ``n_x`` rows of each
-    ``i4``-row block are the integrated state at that substep.
-
-    Returns ``(q_traj, t_traj)`` or ``(None, None)`` if extraction fails.
-    """
-    i4 = n_x + n_x * n_x + 2 * n_x * n_u
-    n_rows, n_sub = V.shape
-    if i4 <= 0 or n_rows % i4 != 0 or n_sub < 1:
-        return None, None
-    n_seg = n_rows // i4
-    if n_seg != len(t_nodes) - 1:
-        return None, None
-    q_rows: list[np.ndarray] = []
-    t_rows: list[float] = []
-    for seg in range(n_seg):
-        t0, t1 = float(t_nodes[seg]), float(t_nodes[seg + 1])
-        j0 = 0 if seg == 0 else 1  # skip duplicated segment-start sample
-        for j in range(j0, n_sub):
-            alpha = j / (n_sub - 1) if n_sub > 1 else 0.0
-            x_vec = np.asarray(V[seg * i4 : seg * i4 + n_x, j], dtype=np.float64)
-            q_rows.append(x_vec[q_slice])
-            t_rows.append((1.0 - alpha) * t0 + alpha * t1)
-    if not q_rows:
-        return None, None
-    return np.stack(q_rows), np.asarray(t_rows, dtype=np.float64)
-
-
 def visualize(results, robot) -> None:
     """Animate the optimised trajectory in Viser (CAD mesh + markers).
 
@@ -307,22 +270,10 @@ def visualize(results, robot) -> None:
     q_traj = np.asarray(results.trajectory["q"])
 
     # ── Prefer multishot V over post-process single propagation ──────────────
-    _dh = getattr(results, "discretization_history", None) or []
-    if _dh:
-        _V = np.asarray(_dh[-1], dtype=np.float64)
-        _n_x = results.x.shape[1]
-        _n_u = results.u.shape[1]
-        _t_nodes_raw = results.nodes.get("time", None)
-        if _t_nodes_raw is None:
-            _t_nodes = np.linspace(0.0, float(t_vec[-1]), len(results.nodes["q"]))
-        else:
-            _t_nodes = np.asarray(_t_nodes_raw).flatten()
-        _q_ms, _t_ms = _q_from_V_multishot(_V, _n_x, _n_u, q.slice, _t_nodes)
-        if _q_ms is not None:
-            print(f"[viser] Multishot V: {len(_q_ms)} frames across {len(_t_nodes) - 1} segments.")
-            q_traj, t_vec = _q_ms, _t_ms
-        else:
-            print("[viser] Multishot V extraction failed; using post-process trajectory.")
+    prop = results.multishot_propagation()
+    if prop is not None:
+        q_traj, t_vec = prop.state("q")
+        print(f"[viser] Multishot V: {len(q_traj)} frames across {prop.n_segments} segments.")
 
     n_frames = len(q_traj)
     n_segs = robot.num_joints + 1

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/frax/panda_frax_viewplanning.py

@@ -267,43 +267,6 @@ def _compute_panda_keypoints(
     return keypoints, ee_pos, ee_quats
 
 
-def _q_from_V_multishot(
-    V: np.ndarray,
-    n_x: int,
-    n_u: int,
-    q_slice: slice,
-    t_nodes: np.ndarray,
-) -> tuple[np.ndarray | None, np.ndarray | None]:
-    """Unpack joint angles from the SCP multi-shoot V matrix.
-
-    V has shape ``((N-1) * i4, n_substeps)`` where
-    ``i4 = n_x + n_x² + 2·n_x·n_u``.  The first ``n_x`` rows of each
-    ``i4``-row block are the integrated state at that substep.
-
-    Returns ``(q_traj, t_traj)`` or ``(None, None)`` if extraction fails.
-    """
-    i4 = n_x + n_x * n_x + 2 * n_x * n_u
-    n_rows, n_sub = V.shape
-    if i4 <= 0 or n_rows % i4 != 0 or n_sub < 1:
-        return None, None
-    n_seg = n_rows // i4
-    if n_seg != len(t_nodes) - 1:
-        return None, None
-    q_rows: list[np.ndarray] = []
-    t_rows: list[float] = []
-    for seg in range(n_seg):
-        t0, t1 = float(t_nodes[seg]), float(t_nodes[seg + 1])
-        j0 = 0 if seg == 0 else 1  # skip duplicated segment-start sample
-        for j in range(j0, n_sub):
-            alpha = j / (n_sub - 1) if n_sub > 1 else 0.0
-            x_vec = np.asarray(V[seg * i4 : seg * i4 + n_x, j], dtype=np.float64)
-            q_rows.append(x_vec[q_slice])
-            t_rows.append((1.0 - alpha) * t0 + alpha * t1)
-    if not q_rows:
-        return None, None
-    return np.stack(q_rows), np.asarray(t_rows, dtype=np.float64)
-
-
 def visualize(results, robot) -> None:
     """Animate the trajectory in Viser: CAD mesh + viewcone + markers.
 
@@ -327,22 +290,10 @@ def visualize(results, robot) -> None:
     q_traj = np.asarray(results.trajectory["q"])
 
     # ── Prefer multishot V over post-process single propagation ──────────────
-    _dh = getattr(results, "discretization_history", None) or []
-    if _dh:
-        _V = np.asarray(_dh[-1], dtype=np.float64)
-        _n_x = results.x.shape[1]
-        _n_u = results.u.shape[1]
-        _t_nodes_raw = results.nodes.get("time", None)
-        if _t_nodes_raw is None:
-            _t_nodes = np.linspace(0.0, float(t_vec[-1]), len(results.nodes["q"]))
-        else:
-            _t_nodes = np.asarray(_t_nodes_raw).flatten()
-        _q_ms, _t_ms = _q_from_V_multishot(_V, _n_x, _n_u, q.slice, _t_nodes)
-        if _q_ms is not None:
-            print(f"[viser] Multishot V: {len(_q_ms)} frames across {len(_t_nodes) - 1} segments.")
-            q_traj, t_vec = _q_ms, _t_ms
-        else:
-            print("[viser] Multishot V extraction failed; using post-process trajectory.")
+    prop = results.multishot_propagation()
+    if prop is not None:
+        q_traj, t_vec = prop.state("q")
+        print(f"[viser] Multishot V: {len(q_traj)} frames across {prop.n_segments} segments.")
 
     n_frames = len(q_traj)
     n_segs = robot.num_joints + 1

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/frax/panda_frax_waypoint.py

@@ -197,43 +197,6 @@ def _compute_panda_keypoints(q_traj: np.ndarray, robot) -> tuple[np.ndarray, np.
     return keypoints, ee_pos
 
 
-def _q_from_V_multishot(
-    V: np.ndarray,
-    n_x: int,
-    n_u: int,
-    q_slice: slice,
-    t_nodes: np.ndarray,
-) -> tuple[np.ndarray | None, np.ndarray | None]:
-    """Unpack joint angles from the SCP multi-shoot V matrix.
-
-    V has shape ``((N-1) * i4, n_substeps)`` where
-    ``i4 = n_x + n_x² + 2·n_x·n_u``.  The first ``n_x`` rows of each
-    ``i4``-row block are the integrated state at that substep.
-
-    Returns ``(q_traj, t_traj)`` or ``(None, None)`` if extraction fails.
-    """
-    i4 = n_x + n_x * n_x + 2 * n_x * n_u
-    n_rows, n_sub = V.shape
-    if i4 <= 0 or n_rows % i4 != 0 or n_sub < 1:
-        return None, None
-    n_seg = n_rows // i4
-    if n_seg != len(t_nodes) - 1:
-        return None, None
-    q_rows: list[np.ndarray] = []
-    t_rows: list[float] = []
-    for seg in range(n_seg):
-        t0, t1 = float(t_nodes[seg]), float(t_nodes[seg + 1])
-        j0 = 0 if seg == 0 else 1  # skip duplicated segment-start sample
-        for j in range(j0, n_sub):
-            alpha = j / (n_sub - 1) if n_sub > 1 else 0.0
-            x_vec = np.asarray(V[seg * i4 : seg * i4 + n_x, j], dtype=np.float64)
-            q_rows.append(x_vec[q_slice])
-            t_rows.append((1.0 - alpha) * t0 + alpha * t1)
-    if not q_rows:
-        return None, None
-    return np.stack(q_rows), np.asarray(t_rows, dtype=np.float64)
-
-
 def visualize(
     results,
     robot,
@@ -265,31 +228,18 @@ def visualize(
     tau_traj = np.asarray(results.trajectory["tau"])
 
     # ── Prefer multishot V over post-process single propagation ──────────────
-    _dh = getattr(results, "discretization_history", None) or []
-    if _dh:
-        _V = np.asarray(_dh[-1], dtype=np.float64)
-        _n_x = results.x.shape[1]
-        _n_u = results.u.shape[1]
-        _t_nodes_raw = results.nodes.get("time", None)
-        if _t_nodes_raw is None:
-            _t_nodes = np.linspace(0.0, float(t_vec[-1]), len(results.nodes["q"]))
-        else:
-            _t_nodes = np.asarray(_t_nodes_raw).flatten()
-        _q_ms, _t_ms = _q_from_V_multishot(_V, _n_x, _n_u, q.slice, _t_nodes)
-        if _q_ms is not None:
-            print(f"[viser] Multishot V: {len(_q_ms)} frames across {len(_t_nodes) - 1} segments.")
-            q_traj, t_vec = _q_ms, _t_ms
-            # Torque is ZOH: replicate each node's tau value across its substeps.
-            _tau_nodes = np.asarray(results.nodes["tau"])
-            _n_sub = _V.shape[1]
-            _tau_rows: list[np.ndarray] = []
-            for _seg in range(len(_t_nodes) - 1):
-                _j0 = 0 if _seg == 0 else 1
-                for _j in range(_j0, _n_sub):
-                    _tau_rows.append(_tau_nodes[_seg])
-            tau_traj = np.stack(_tau_rows)
-        else:
-            print("[viser] Multishot V extraction failed; using post-process trajectory.")
+    prop = results.multishot_propagation()
+    if prop is not None:
+        q_traj, t_vec = prop.state("q")
+        print(f"[viser] Multishot V: {len(q_traj)} frames across {prop.n_segments} segments.")
+        # Torque is ZOH: replicate each node's tau value across its substeps.
+        _tau_nodes = np.asarray(results.nodes["tau"])
+        _tau_rows: list[np.ndarray] = []
+        for _seg in range(prop.n_segments):
+            _j0 = 0 if _seg == 0 else 1
+            for _j in range(_j0, prop.n_substeps):
+                _tau_rows.append(_tau_nodes[_seg])
+        tau_traj = np.stack(_tau_rows)
 
     keypoints, ee_pos = _compute_panda_keypoints(q_traj, robot)
     n_segs = robot.num_joints + 1

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/mjx/double_cartpole_mjx.py

@@ -212,44 +212,6 @@ def simulate_mujoco(results) -> dict:
     return {"time": rec_t, "qpos": rec_q, "qvel": rec_qd}
 
 
-def qpos_from_V_multishot(
-    V: np.ndarray,
-    *,
-    n_q: int,
-    n_v: int,
-    n_u: int,
-    t_nodes: np.ndarray,
-) -> tuple[np.ndarray | None, np.ndarray | None]:
-    """Unpack generalized coordinates from the SCP multi-shoot matrix ``V``."""
-    if V.size == 0:
-        return None, None
-    n_x = n_q + n_v
-    i4 = n_x + n_x * n_x + 2 * n_x * n_u
-    n_rows, n_sub = V.shape
-    if i4 <= 0 or n_rows % i4 != 0 or n_sub < 1:
-        return None, None
-    n_seg = n_rows // i4
-    if n_seg != len(t_nodes) - 1:
-        return None, None
-
-    q_rows: list[np.ndarray] = []
-    t_rows: list[float] = []
-    for seg in range(n_seg):
-        t0 = float(t_nodes[seg])
-        t1 = float(t_nodes[seg + 1])
-        j0 = 0 if seg == 0 else 1
-        for j in range(j0, n_sub):
-            alpha = j / (n_sub - 1) if n_sub > 1 else 0.0
-            t_s = (1.0 - alpha) * t0 + alpha * t1
-            row0 = seg * i4
-            x_vec = np.asarray(V[row0 : row0 + n_x, j], dtype=np.float64).ravel()
-            q_rows.append(x_vec[:n_q])
-            t_rows.append(t_s)
-    if not q_rows:
-        return None, None
-    return np.stack(q_rows, axis=0), np.asarray(t_rows, dtype=np.float64)
-
-
 def visualize(results, sim: dict | None = None) -> None:
     """Animate the double-link cartpole in a Viser 3D scene.
 
@@ -288,21 +250,9 @@ def visualize(results, sim: dict | None = None) -> None:
     fk_nodes = [fk_joints(q_nodes[i]) for i in range(len(q_nodes))]
 
     # Multi-shoot integrated trajectory from ``V``
-    _dh = getattr(results, "discretization_history", None) or []
-    V_multishot = _dh[-1] if len(_dh) > 0 else None
-    q_ms_v, t_ms_v = (
-        qpos_from_V_multishot(
-            np.asarray(V_multishot, dtype=np.float64),
-            n_q=n_q,
-            n_v=n_v,
-            n_u=n_u,
-            t_nodes=t_nodes,
-        )
-        if V_multishot is not None
-        else (None, None)
-    )
-
-    if q_ms_v is not None and t_ms_v is not None:
+    prop = results.multishot_propagation(t_nodes=t_nodes)
+    if prop is not None:
+        q_ms_v, t_ms_v = prop.state("qpos")
         fk_multishot_anim = [fk_joints(q_ms_v[i]) for i in range(len(q_ms_v))]
         t_multishot_lookup = t_ms_v
     else:

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/mjx/triple_cartpole_3d_mjx.py

@@ -59,7 +59,6 @@ except ImportError:
 import openscvx as ox
 from examples.plotting_viser import (
     create_snapshot_plotting_server,
-    extract_multishoot_trajectory,
 )
 from openscvx import Problem
 
@@ -252,62 +251,6 @@ def _foh_controls_at_times(
     return u_out
 
 
-def extract_multishoot_qpos_chronological(
-    V_multi_shoot: np.ndarray,
-    *,
-    n_x: int,
-    n_u: int,
-    n_q: int,
-    t_nodes: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None:
-    """Unpack propagated states from ``V`` in time order (SCP / realtime layout).
-
-    ``V`` has shape ``((N-1) * segment_size, n_substeps)``.  Each segment block
-    stores the integrated state in its first ``n_x`` rows at every substep column.
-    This matches :func:`openscvx.plotting.viser.scp.extract_propagation_positions`
-    and the realtime ``extract_multishoot_trajectory`` helpers — **not** nodal
-    linear interpolation of ``results.nodes``.
-    """
-    V_multi_shoot = np.asarray(V_multi_shoot, dtype=np.float64)
-    if V_multi_shoot.size == 0:
-        return None
-
-    segment_size = n_x + n_x * n_x + 2 * n_x * n_u
-    n_rows, n_sub = V_multi_shoot.shape
-    if segment_size <= 0 or n_rows % segment_size != 0 or n_sub < 1:
-        return None
-    n_seg = n_rows // segment_size
-
-    t_nodes = np.asarray(t_nodes, dtype=np.float64).ravel()
-    if t_nodes.size != n_seg + 1:
-        if t_nodes.size < 2:
-            return None
-        t_nodes = np.linspace(float(t_nodes[0]), float(t_nodes[-1]), n_seg + 1)
-
-    q_rows: list[np.ndarray] = []
-    qd_rows: list[np.ndarray] = []
-    t_rows: list[float] = []
-    for seg_idx in range(n_seg):
-        seg_start = seg_idx * segment_size
-        t0, t1 = float(t_nodes[seg_idx]), float(t_nodes[seg_idx + 1])
-        j0 = 0 if seg_idx == 0 else 1
-        for t_idx in range(j0, n_sub):
-            alpha = t_idx / (n_sub - 1) if n_sub > 1 else 0.0
-            state = np.asarray(
-                V_multi_shoot[seg_start : seg_start + n_x, t_idx], dtype=np.float64
-            ).ravel()
-            q_rows.append(state[:n_q])
-            qd_rows.append(state[n_q:n_x])
-            t_rows.append((1.0 - alpha) * t0 + alpha * t1)
-    if not q_rows:
-        return None
-
-    t_ms = np.asarray(t_rows, dtype=np.float64)
-    q_ms = np.stack(q_rows, axis=0)
-    qd_ms = np.stack(qd_rows, axis=0)
-    return q_ms, qd_ms, t_ms, t_nodes
-
-
 def segment_tip_paths_from_V(
     V_multi_shoot: np.ndarray,
     *,
@@ -362,23 +305,11 @@ def visualize(results) -> None:
     else:
         t_nodes = np.asarray(t_nodes).flatten()
 
-    _dh = getattr(results, "discretization_history", None) or []
-    V_multishot = np.asarray(_dh[-1], dtype=np.float64) if len(_dh) > 0 else None
-    ms_traj = (
-        extract_multishoot_qpos_chronological(
-            V_multishot,
-            n_x=n_x,
-            n_u=n_u,
-            n_q=n_q,
-            t_nodes=t_nodes,
-        )
-        if V_multishot is not None
-        else None
-    )
-
-    using_multishot = ms_traj is not None
+    prop = results.multishot_propagation(t_nodes=t_nodes)
+    using_multishot = prop is not None
     if using_multishot:
-        q_angle, _, t_play, t_nodes_ms = ms_traj
+        q_angle, t_play = prop.state("qpos")
+        t_nodes_ms = prop.t_nodes
         u_play = _foh_controls_at_times(t_play, u_nodes, t_nodes_ms)
         fk_multishot_anim = [fk_joints(q_angle[i]) for i in range(len(q_angle))]
         tip_pos = np.array([fk[4] for fk in fk_multishot_anim], dtype=np.float64)
@@ -386,12 +317,13 @@ def visualize(results) -> None:
             tip_vel = np.gradient(tip_pos, t_play, axis=0)
         else:
             tip_vel = np.zeros_like(tip_pos)
-        n_seg = V_multishot.shape[0] // (n_x + n_x * n_x + 2 * n_x * n_u)
+        V_multishot = prop.V
         print(
             f"[viser] Multi-shoot V: {len(t_play)} propagated samples "
-            f"({V_multishot.shape[1]} cols × {n_seg} segments)."
+            f"({prop.n_substeps} cols × {prop.n_segments} segments)."
         )
     else:
+        V_multishot = None
         print(
             "[viser] WARNING: could not decode discretization_history V; "
             "falling back to nodal linear interpolation."
@@ -481,13 +413,7 @@ def visualize(results) -> None:
     tip_multishot_segs = None
     if using_multishot and V_multishot is not None:
         # Faint cloud of every integrated state (realtime-style unpack).
-        all_qpos, _ = extract_multishoot_trajectory(
-            V_multishot,
-            n_x,
-            n_u,
-            position_slice=slice(0, n_q),
-            velocity_slice=None,
-        )
+        all_qpos, _ = prop.state("qpos")
         if len(all_qpos) > 0:
             tip_cloud = np.array(
                 [fk_joints(all_qpos[i])[4] for i in range(len(all_qpos))],

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/mjx/triple_cartpole_mjx.py

@@ -37,7 +37,6 @@ except ImportError:
 import openscvx as ox
 from examples.plotting_viser import (
     create_snapshot_plotting_server,
-    extract_multishoot_trajectory,
 )
 from openscvx import Problem
 
@@ -175,55 +174,6 @@ def _foh_controls_at_times(
     return u_out
 
 
-def extract_multishoot_qpos_chronological(
-    V_multi_shoot: np.ndarray,
-    *,
-    n_x: int,
-    n_u: int,
-    n_q: int,
-    t_nodes: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray] | None:
-    """Unpack propagated states from ``V`` in time order (SCP / realtime layout)."""
-    V_multi_shoot = np.asarray(V_multi_shoot, dtype=np.float64)
-    if V_multi_shoot.size == 0:
-        return None
-
-    segment_size = n_x + n_x * n_x + 2 * n_x * n_u
-    n_rows, n_sub = V_multi_shoot.shape
-    if segment_size <= 0 or n_rows % segment_size != 0 or n_sub < 1:
-        return None
-    n_seg = n_rows // segment_size
-
-    t_nodes = np.asarray(t_nodes, dtype=np.float64).ravel()
-    if t_nodes.size != n_seg + 1:
-        if t_nodes.size < 2:
-            return None
-        t_nodes = np.linspace(float(t_nodes[0]), float(t_nodes[-1]), n_seg + 1)
-
-    q_rows: list[np.ndarray] = []
-    qd_rows: list[np.ndarray] = []
-    t_rows: list[float] = []
-    for seg_idx in range(n_seg):
-        seg_start = seg_idx * segment_size
-        t0, t1 = float(t_nodes[seg_idx]), float(t_nodes[seg_idx + 1])
-        j0 = 0 if seg_idx == 0 else 1
-        for t_idx in range(j0, n_sub):
-            alpha = t_idx / (n_sub - 1) if n_sub > 1 else 0.0
-            state = np.asarray(
-                V_multi_shoot[seg_start : seg_start + n_x, t_idx], dtype=np.float64
-            ).ravel()
-            q_rows.append(state[:n_q])
-            qd_rows.append(state[n_q:n_x])
-            t_rows.append((1.0 - alpha) * t0 + alpha * t1)
-    if not q_rows:
-        return None
-
-    t_ms = np.asarray(t_rows, dtype=np.float64)
-    q_ms = np.stack(q_rows, axis=0)
-    qd_ms = np.stack(qd_rows, axis=0)
-    return q_ms, qd_ms, t_ms, t_nodes
-
-
 def segment_tip_paths_from_V(
     V_multi_shoot: np.ndarray,
     *,
@@ -277,23 +227,11 @@ def visualize(results) -> None:
     else:
         t_nodes = np.asarray(t_nodes).flatten()
 
-    _dh = getattr(results, "discretization_history", None) or []
-    V_multishot = np.asarray(_dh[-1], dtype=np.float64) if len(_dh) > 0 else None
-    ms_traj = (
-        extract_multishoot_qpos_chronological(
-            V_multishot,
-            n_x=n_x,
-            n_u=n_u,
-            n_q=n_q,
-            t_nodes=t_nodes,
-        )
-        if V_multishot is not None
-        else None
-    )
-
-    using_multishot = ms_traj is not None
+    prop = results.multishot_propagation(t_nodes=t_nodes)
+    using_multishot = prop is not None
     if using_multishot:
-        q_angle, _, t_play, t_nodes_ms = ms_traj
+        q_angle, t_play = prop.state("qpos")
+        t_nodes_ms = prop.t_nodes
         u_play = _foh_controls_at_times(t_play, u_nodes, t_nodes_ms)
         fk_multishot_anim = [fk_joints(q_angle[i]) for i in range(len(q_angle))]
         tip_pos = np.array([fk[4] for fk in fk_multishot_anim], dtype=np.float64)
@@ -301,12 +239,13 @@ def visualize(results) -> None:
             tip_vel = np.gradient(tip_pos, t_play, axis=0)
         else:
             tip_vel = np.zeros_like(tip_pos)
-        n_seg = V_multishot.shape[0] // (n_x + n_x * n_x + 2 * n_x * n_u)
+        V_multishot = prop.V
         print(
             f"[viser] Multi-shoot V: {len(t_play)} propagated samples "
-            f"({V_multishot.shape[1]} cols × {n_seg} segments)."
+            f"({prop.n_substeps} cols × {prop.n_segments} segments)."
         )
     else:
+        V_multishot = None
         print(
             "[viser] WARNING: could not decode discretization_history V; "
             "falling back to nodal linear interpolation."
@@ -363,13 +302,7 @@ def visualize(results) -> None:
     cart_multishot_segs = None
     tip_multishot_segs = None
     if using_multishot and V_multishot is not None:
-        all_qpos, _ = extract_multishoot_trajectory(
-            V_multishot,
-            n_x,
-            n_u,
-            position_slice=slice(0, n_q),
-            velocity_slice=None,
-        )
+        all_qpos, _ = prop.state("qpos")
         if len(all_qpos) > 0:
             tip_cloud = np.array(
                 [fk_joints(all_qpos[i])[4] for i in range(len(all_qpos))],

{openscvx-0.5.3.dev30 → openscvx-0.5.3.dev31}/examples/plotting_viser.py

@@ -1756,12 +1756,10 @@ def extract_multishoot_trajectory(
 ) -> tuple[np.ndarray, np.ndarray | None]:
     """Extract position and velocity trajectories from multi-shoot data.
 
-    The multi-shoot format stores propagation segments with state, STM, and
-    sensitivity data packed together. This function extracts the state
-    components from each segment.
+    Uses time-ordered chronological stitching with deduplicated segment boundaries.
 
     Args:
-        V_multi_shoot: Multi-shoot data array of shape (n_segments * segment_size, n_nodes)
+        V_multi_shoot: Multi-shoot data array of shape (n_segments * segment_size, n_substeps)
         n_x: Number of states
         n_u: Number of controls
         position_slice: Slice for extracting position from state (default: first 3)
@@ -1772,30 +1770,19 @@ def extract_multishoot_trajectory(
         Tuple of (positions, velocities) as float32 arrays.
         positions: Shape (n_total_points, 3)
         velocities: Shape (n_total_points, 3) or None if velocity_slice is None
-
-    Note:
-        The segment size is computed as: n_x + n_x² + 2*n_x*n_u
-        This accounts for: state (n_x) + STM (n_x²) + sensitivities (2*n_x*n_u)
     """
-    # Segment size: state + STM + control sensitivities
-    segment_size = n_x + n_x * n_x + 2 * n_x * n_u
-
-    all_pos_segments = []
-    all_vel_segments = [] if velocity_slice is not None else None
-
-    for i_node in range(V_multi_shoot.shape[1]):
-        node_data = V_multi_shoot[:, i_node]
-        segments_for_node = node_data.reshape(-1, segment_size)
-        pos_segments = segments_for_node[:, position_slice]
-        all_pos_segments.append(pos_segments)
-
-        if velocity_slice is not None:
-            vel_segments = segments_for_node[:, velocity_slice]
-            all_vel_segments.append(vel_segments)
-
-    positions = np.vstack(all_pos_segments).astype(np.float32)
-    velocities = np.vstack(all_vel_segments).astype(np.float32) if all_vel_segments else None
-
+    from openscvx.algorithms.history import unpack_multishot_V
+
+    n_segments = V_multi_shoot.shape[0] // (n_x + n_x * n_x + 2 * n_x * n_u)
+    placeholder_t = np.linspace(0.0, 1.0, n_segments + 1)
+    prop = unpack_multishot_V(V_multi_shoot, n_x=n_x, n_u=n_u, t_nodes=placeholder_t, states=())
+    positions, _ = prop.slice_states(position_slice)
+    positions = positions.astype(np.float32)
+    if velocity_slice is not None:
+        velocities, _ = prop.slice_states(velocity_slice)
+        velocities = velocities.astype(np.float32)
+    else:
+        velocities = None
     return positions, velocities