openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/plotting/viser/scp.py

@@ -0,0 +1,459 @@
+"""SCP iteration visualization components for viser.
+
+This module contains functions for visualizing the successive convex programming
+(SCP) optimization process, showing how the solution evolves across iterations.
+"""
+
+import threading
+import time
+from typing import Callable
+
+import matplotlib.pyplot as plt
+import numpy as np
+import viser
+
+# Type alias for update callbacks
+UpdateCallback = Callable[[int], None]
+
+
+def add_scp_iteration_nodes(
+    server: viser.ViserServer,
+    positions: list[np.ndarray],
+    colors: list[tuple[int, int, int]] | None = None,
+    point_size: float = 0.3,
+    cmap_name: str = "viridis",
+) -> tuple[list[viser.PointCloudHandle], UpdateCallback]:
+    """Add animated optimization nodes that update per SCP iteration.
+
+    Pre-buffers point clouds for all iterations and toggles visibility for performance.
+    This avoids transmitting point data on every frame update.
+
+    Args:
+        server: ViserServer instance
+        positions: List of position arrays per iteration, each shape (N, 3)
+        colors: Optional list of RGB colors per iteration. If None, uses viridis colormap.
+        point_size: Size of node markers
+        cmap_name: Matplotlib colormap name (default: "viridis")
+
+    Returns:
+        Tuple of (list of point_handles, update_callback)
+    """
+    n_iterations = len(positions)
+
+    # Default: use viridis colormap
+    if colors is None:
+        cmap = plt.get_cmap(cmap_name)
+        colors = []
+        for i in range(n_iterations):
+            t = i / max(n_iterations - 1, 1)
+            rgb = cmap(t)[:3]
+            colors.append(tuple(int(c * 255) for c in rgb))
+
+    # Convert colors to numpy arrays for viser compatibility
+    colors_np = [np.array([c[0], c[1], c[2]], dtype=np.uint8) for c in colors]
+
+    # Pre-create point clouds for all iterations (only first visible initially)
+    handles = []
+    for i in range(n_iterations):
+        pos = np.asarray(positions[i], dtype=np.float32)
+        handle = server.scene.add_point_cloud(
+            f"/scp/nodes/iter_{i}",
+            points=pos,
+            colors=colors_np[i],
+            point_size=point_size,
+            visible=(i == 0),
+        )
+        handles.append(handle)
+
+    # Track current visible iteration to minimize visibility toggles
+    state = {"current_idx": 0}
+
+    def update(iter_idx: int) -> None:
+        idx = min(iter_idx, n_iterations - 1)
+        if idx != state["current_idx"]:
+            handles[state["current_idx"]].visible = False
+            handles[idx].visible = True
+            state["current_idx"] = idx
+
+    return handles, update
+
+
+def add_scp_iteration_attitudes(
+    server: viser.ViserServer,
+    positions: list[np.ndarray],
+    attitudes: list[np.ndarray] | None,
+    axes_length: float = 1.5,
+    axes_radius: float = 0.03,
+    stride: int = 1,
+) -> tuple[list[viser.FrameHandle], UpdateCallback | None]:
+    """Add animated attitude frames at each node that update per SCP iteration.
+
+    Args:
+        server: ViserServer instance
+        positions: List of position arrays per iteration, each shape (N, 3)
+        attitudes: List of quaternion arrays per iteration, each shape (N, 4) in wxyz format.
+            If None, returns empty list and None callback.
+        axes_length: Length of coordinate frame axes
+        axes_radius: Radius of axes cylinders
+        stride: Show attitude frame every `stride` nodes (1 = all nodes)
+
+    Returns:
+        Tuple of (list of frame handles, update_callback)
+    """
+    if attitudes is None:
+        return [], None
+
+    n_iterations = len(positions)
+    n_nodes = len(positions[0])
+
+    # Create frame handles for nodes at stride intervals
+    node_indices = list(range(0, n_nodes, stride))
+    handles = []
+
+    for i, node_idx in enumerate(node_indices):
+        handle = server.scene.add_frame(
+            f"/scp/attitudes/frame_{i}",
+            wxyz=attitudes[0][node_idx],
+            position=positions[0][node_idx],
+            axes_length=axes_length,
+            axes_radius=axes_radius,
+        )
+        handles.append(handle)
+
+    def update(iter_idx: int) -> None:
+        idx = min(iter_idx, n_iterations - 1)
+        pos = positions[idx]
+        att = attitudes[idx]
+
+        for i, node_idx in enumerate(node_indices):
+            # Handle case where number of nodes changes between iterations
+            if node_idx < len(pos) and node_idx < len(att):
+                handles[i].position = pos[node_idx]
+                handles[i].wxyz = att[node_idx]
+
+    return handles, update
+
+
+def add_scp_ghost_iterations(
+    server: viser.ViserServer,
+    positions: list[np.ndarray],
+    point_size: float = 0.15,
+    cmap_name: str = "viridis",
+) -> tuple[list[viser.PointCloudHandle], UpdateCallback]:
+    """Add ghost trails showing all previous SCP iterations.
+
+    Pre-buffers point clouds for all iterations and toggles visibility for performance.
+    Shows all previous iterations with viridis coloring to visualize convergence.
+
+    Args:
+        server: ViserServer instance
+        positions: List of position arrays per iteration, each shape (N, 3)
+        point_size: Size of ghost points
+        cmap_name: Matplotlib colormap name for ghost colors
+
+    Returns:
+        Tuple of (list of handles, update_callback)
+    """
+    n_iterations = len(positions)
+    cmap = plt.get_cmap(cmap_name)
+
+    # Pre-create point clouds for all iterations with their colors
+    # (all initially hidden, shown progressively as ghosts)
+    handles = []
+    for i in range(n_iterations):
+        t = i / max(n_iterations - 1, 1)
+        rgb = cmap(t)[:3]
+        color = np.array([int(c * 255) for c in rgb], dtype=np.uint8)
+        pos = np.asarray(positions[i], dtype=np.float32)
+
+        handle = server.scene.add_point_cloud(
+            f"/scp/ghosts/iter_{i}",
+            points=pos,
+            colors=color,
+            point_size=point_size,
+            visible=False,  # All start hidden
+        )
+        handles.append(handle)
+
+    # Track which iterations are currently visible as ghosts
+    state = {"visible_up_to": -1}
+
+    def update(iter_idx: int) -> None:
+        idx = min(iter_idx, n_iterations - 1)
+        # Ghosts are iterations 0 through idx-1 (everything before current)
+        new_visible_up_to = idx - 1
+
+        if new_visible_up_to != state["visible_up_to"]:
+            # Show/hide only the iterations that changed
+            if new_visible_up_to > state["visible_up_to"]:
+                # Show newly visible ghosts
+                for i in range(state["visible_up_to"] + 1, new_visible_up_to + 1):
+                    handles[i].visible = True
+            else:
+                # Hide ghosts that should no longer be visible
+                for i in range(new_visible_up_to + 1, state["visible_up_to"] + 1):
+                    handles[i].visible = False
+            state["visible_up_to"] = new_visible_up_to
+
+    return handles, update
+
+
+def extract_propagation_positions(
+    discretization_history: list[np.ndarray],
+    n_x: int,
+    n_u: int,
+    position_slice: slice,
+    scene_scale: float = 1.0,
+) -> list[list[np.ndarray]]:
+    """Extract 3D position trajectories from discretization history.
+
+    The discretization history contains the multi-shot integration results.
+    Each V matrix has shape (flattened_size, n_timesteps) where:
+    - flattened_size = (N-1) * i4
+    - i4 = n_x + n_x*n_x + 2*n_x*n_u (state + STM + control influence matrices)
+    - n_timesteps = number of integration substeps
+
+    Args:
+        discretization_history: List of V matrices from each SCP iteration
+        n_x: Number of states
+        n_u: Number of controls
+        position_slice: Slice for extracting position from state vector
+        scene_scale: Divide positions by this factor for visualization
+
+    Returns:
+        List of propagation trajectories per iteration.
+        Each iteration contains a list of (n_substeps, 3) arrays, one per segment.
+    """
+    if not discretization_history:
+        return []
+
+    i4 = n_x + n_x * n_x + 2 * n_x * n_u
+    propagations = []
+
+    for V in discretization_history:
+        # V shape: (flattened_size, n_timesteps)
+        n_timesteps = V.shape[1]
+        n_segments = V.shape[0] // i4  # N-1 segments
+
+        iteration_segments = []
+        for seg_idx in range(n_segments):
+            # Extract this segment's data across all timesteps
+            seg_start = seg_idx * i4
+            seg_end = seg_start + i4
+
+            # For each timestep, extract the position from the state
+            segment_positions = []
+            for t_idx in range(n_timesteps):
+                # Get full state at this segment and timestep
+                state = V[seg_start:seg_end, t_idx][:n_x]
+                # Extract position components
+                pos = state[position_slice] / scene_scale
+                segment_positions.append(pos)
+
+            iteration_segments.append(np.array(segment_positions, dtype=np.float32))
+
+        propagations.append(iteration_segments)
+
+    return propagations
+
+
+def add_scp_propagation_lines(
+    server: viser.ViserServer,
+    propagations: list[list[np.ndarray]],
+    line_width: float = 2.0,
+    cmap_name: str = "viridis",
+) -> tuple[list, UpdateCallback]:
+    """Add animated nonlinear propagation lines that update per SCP iteration.
+
+    Shows the actual integrated trajectory between optimization nodes,
+    revealing defects (gaps) in early iterations that close as SCP converges.
+    All iterations up to the current one are shown with viridis coloring,
+    similar to ghost iterations for nodes.
+
+    Args:
+        server: ViserServer instance
+        propagations: List of propagation trajectories per iteration from
+            extract_propagation_positions(). Each iteration contains a list
+            of (n_substeps, 3) position arrays, one per segment.
+        line_width: Width of propagation lines
+        cmap_name: Matplotlib colormap name (default: "viridis")
+
+    Returns:
+        Tuple of (list of line handles, update_callback)
+    """
+    if not propagations:
+        return [], lambda _: None
+
+    n_iterations = len(propagations)
+    n_segments = len(propagations[0])
+    cmap = plt.get_cmap(cmap_name)
+
+    # Pre-compute colors for each iteration
+    iteration_colors = []
+    for i in range(n_iterations):
+        t = i / max(n_iterations - 1, 1)
+        rgb = cmap(t)[:3]
+        iteration_colors.append(np.array([int(c * 255) for c in rgb], dtype=np.uint8))
+
+    # Create line handles for each (iteration, segment) pair
+    # Structure: handles[iter_idx][seg_idx]
+    all_handles = []
+
+    for iter_idx in range(n_iterations):
+        iter_handles = []
+        color = iteration_colors[iter_idx]
+
+        for seg_idx in range(n_segments):
+            seg_pos = propagations[iter_idx][seg_idx]  # Shape (n_substeps, 3)
+
+            if len(seg_pos) < 2:
+                iter_handles.append(None)
+                continue
+
+            # Create line segments connecting consecutive substeps
+            segments = np.array(
+                [[seg_pos[i], seg_pos[i + 1]] for i in range(len(seg_pos) - 1)],
+                dtype=np.float32,
+            )
+
+            handle = server.scene.add_line_segments(
+                f"/scp/propagation/iter_{iter_idx}/segment_{seg_idx}",
+                points=segments,
+                colors=color,
+                line_width=line_width,
+                visible=(iter_idx == 0),  # Only first iteration visible initially
+            )
+            iter_handles.append(handle)
+
+        all_handles.append(iter_handles)
+
+    def update(iter_idx: int) -> None:
+        idx = min(iter_idx, n_iterations - 1)
+
+        # Show all iterations up to and including current, hide the rest
+        for i in range(n_iterations):
+            should_show = i <= idx
+            for handle in all_handles[i]:
+                if handle is not None:
+                    handle.visible = should_show
+
+    return all_handles, update
+
+
+def add_scp_animation_controls(
+    server: viser.ViserServer,
+    n_iterations: int,
+    update_callbacks: list[UpdateCallback],
+    autoplay: bool = False,
+    frame_duration_ms: int = 500,
+    folder_name: str = "SCP Animation",
+) -> None:
+    """Add GUI controls for stepping through SCP iterations.
+
+    Creates play/pause button, step buttons, iteration slider, and speed control.
+
+    Args:
+        server: ViserServer instance
+        n_iterations: Total number of SCP iterations
+        update_callbacks: List of update functions to call each iteration
+        autoplay: Whether to start playing automatically
+        frame_duration_ms: Default milliseconds per iteration frame
+        folder_name: Name for the GUI folder
+    """
+    # Filter out None callbacks
+    callbacks = [cb for cb in update_callbacks if cb is not None]
+
+    def update_all(iter_idx: int) -> None:
+        """Update all visualization components."""
+        for callback in callbacks:
+            callback(iter_idx)
+
+    # --- GUI Controls ---
+    with server.gui.add_folder(folder_name):
+        play_button = server.gui.add_button("Play")
+        with server.gui.add_folder("Step Controls", expand_by_default=False):
+            prev_button = server.gui.add_button("< Previous")
+            next_button = server.gui.add_button("Next >")
+        iter_slider = server.gui.add_slider(
+            "Iteration",
+            min=0,
+            max=n_iterations - 1,
+            step=1,
+            initial_value=0,
+        )
+        speed_slider = server.gui.add_slider(
+            "Speed (ms/iter)",
+            min=50,
+            max=2000,
+            step=50,
+            initial_value=frame_duration_ms,
+        )
+        loop_checkbox = server.gui.add_checkbox("Loop", initial_value=True)
+
+    # Animation state
+    state = {"playing": autoplay, "iteration": 0, "needs_update": True}
+
+    @play_button.on_click
+    def _(_) -> None:
+        state["playing"] = not state["playing"]
+        state["needs_update"] = True  # Trigger immediate update on play
+        play_button.name = "Pause" if state["playing"] else "Play"
+
+    @prev_button.on_click
+    def _(_) -> None:
+        if state["iteration"] > 0:
+            state["iteration"] -= 1
+            iter_slider.value = state["iteration"]
+            update_all(state["iteration"])
+
+    @next_button.on_click
+    def _(_) -> None:
+        if state["iteration"] < n_iterations - 1:
+            state["iteration"] += 1
+            iter_slider.value = state["iteration"]
+            update_all(state["iteration"])
+
+    @iter_slider.on_update
+    def _(_) -> None:
+        if not state["playing"]:
+            state["iteration"] = int(iter_slider.value)
+            update_all(state["iteration"])
+
+    def animation_loop() -> None:
+        """Background thread for SCP iteration playback."""
+        last_update = time.time()
+        while True:
+            time.sleep(0.016)  # ~60 fps check rate
+
+            # Handle immediate update requests (e.g., on play button click)
+            if state["needs_update"]:
+                state["needs_update"] = False
+                last_update = time.time()
+                update_all(state["iteration"])
+                continue
+
+            if state["playing"]:
+                current_time = time.time()
+                elapsed_ms = (current_time - last_update) * 1000
+
+                if elapsed_ms >= speed_slider.value:
+                    last_update = current_time
+                    state["iteration"] += 1
+
+                    if state["iteration"] >= n_iterations:
+                        if loop_checkbox.value:
+                            state["iteration"] = 0
+                        else:
+                            state["iteration"] = n_iterations - 1
+                            state["playing"] = False
+                            play_button.name = "Play"
+
+                    iter_slider.value = state["iteration"]
+                    update_all(state["iteration"])
+
+    # Start animation thread
+    thread = threading.Thread(target=animation_loop, daemon=True)
+    thread.start()
+
+    # Initial update to ensure first frame is fully rendered
+    update_all(0)

openscvx/plotting/viser/server.py

@@ -0,0 +1,112 @@
+"""Viser server setup utilities."""
+
+import matplotlib.pyplot as plt
+import numpy as np
+import viser
+from viser.theme import TitlebarButton, TitlebarConfig, TitlebarImage
+
+
+def compute_velocity_colors(vel: np.ndarray, cmap_name: str = "viridis") -> np.ndarray:
+    """Compute RGB colors based on velocity magnitude.
+
+    Args:
+        vel: Velocity array of shape (N, 3)
+        cmap_name: Matplotlib colormap name
+
+    Returns:
+        Array of RGB colors with shape (N, 3), values in [0, 255]
+    """
+    vel_norms = np.linalg.norm(vel, axis=1)
+    vel_range = vel_norms.max() - vel_norms.min()
+    if vel_range < 1e-8:
+        vel_normalized = np.zeros_like(vel_norms)
+    else:
+        vel_normalized = (vel_norms - vel_norms.min()) / vel_range
+
+    cmap = plt.get_cmap(cmap_name)
+    colors = np.array([[int(c * 255) for c in cmap(v)[:3]] for v in vel_normalized])
+    return colors
+
+
+def compute_grid_size(pos: np.ndarray, padding: float = 1.2) -> float:
+    """Compute grid size based on trajectory extent.
+
+    Args:
+        pos: Position array of shape (N, 3)
+        padding: Padding factor (1.2 = 20% padding)
+
+    Returns:
+        Grid size (width and height)
+    """
+    max_x = np.abs(pos[:, 0]).max()
+    max_y = np.abs(pos[:, 1]).max()
+    return max(max_x, max_y) * 2 * padding
+
+
+def create_server(
+    pos: np.ndarray,
+    dark_mode: bool = True,
+) -> viser.ViserServer:
+    """Create a viser server with basic scene setup.
+
+    Args:
+        pos: Position array for computing grid size
+        dark_mode: Whether to use dark theme
+
+    Returns:
+        ViserServer instance with grid and origin frame
+    """
+    server = viser.ViserServer()
+
+    # Configure theme with OpenSCvx branding
+    # TitlebarButton and TitlebarConfig are TypedDict classes (create as plain dicts)
+    buttons = (
+        TitlebarButton(
+            text="Getting Started",
+            icon="Description",
+            href="https://haynec.github.io/OpenSCvx/latest/getting-started/",
+        ),
+        TitlebarButton(
+            text="Docs",
+            icon="Description",
+            href="https://haynec.github.io/OpenSCvx/",
+        ),
+        TitlebarButton(
+            text="GitHub",
+            icon="GitHub",
+            href="https://github.com/haynec/OpenSCvx",
+        ),
+    )
+
+    # Add OpenSCvx logo to titlebar (loaded from GitHub)
+    logo_url = (
+        "https://raw.githubusercontent.com/haynec/OpenSCvx/main/figures/openscvx_logo_square.png"
+    )
+    image = TitlebarImage(
+        image_url_light=logo_url,
+        image_url_dark=logo_url,  # Use same logo for both themes
+        image_alt="OpenSCvx",
+        href="https://github.com/haynec/OpenSCvx",
+    )
+
+    titlebar_config = TitlebarConfig(buttons=buttons, image=image)
+
+    server.gui.configure_theme(
+        titlebar_content=titlebar_config,
+        dark_mode=dark_mode,
+    )
+
+    grid_size = compute_grid_size(pos)
+    server.scene.add_grid(
+        "/grid",
+        width=grid_size,
+        height=grid_size,
+        position=np.array([0.0, 0.0, 0.0]),
+    )
+    server.scene.add_frame(
+        "/origin",
+        wxyz=(1.0, 0.0, 0.0, 0.0),
+        position=(0.0, 0.0, 0.0),
+    )
+
+    return server