openscvx 0.3.2.dev170__py3-none-any.whl

openscvx/plotting/viser/plotly_integration.py

@@ -0,0 +1,333 @@
+"""Plotly integration for viser - animated 2D plots synchronized with 3D visualization.
+
+This module provides utilities for embedding plotly figures in viser's GUI with
+animated markers that synchronize with the 3D animation timeline.
+"""
+
+import numpy as np
+import plotly.graph_objects as go
+import viser
+
+from openscvx.algorithms import OptimizationResults
+
+
+def add_animated_plotly_vline(
+    server: viser.ViserServer,
+    fig: go.Figure,
+    time_array: np.ndarray,
+    use_trajectory_indexing: bool = True,
+    line_color: str = "red",
+    line_width: int = 2,
+    line_dash: str = "dash",
+    annotation_text: str = "Current",
+    annotation_position: str = "top",
+    folder_name: str | None = None,
+    aspect: float = 1.5,
+) -> tuple:
+    """Add a plotly figure to viser GUI with an animated vertical line.
+
+    This function takes any plotly figure and adds an animated vertical line that
+    synchronizes with viser's 3D animation timeline. The line shows the current
+    time position as the animation plays.
+
+    This is more generic than add_animated_plotly_marker() as it works for any
+    number of traces without needing to specify y-data for each.
+
+    Args:
+        server: ViserServer instance
+        fig: Plotly figure to display
+        time_array: Time values corresponding to animation frames (N,).
+            This should match the time array passed to add_animation_controls().
+        use_trajectory_indexing: If True, frame_idx maps directly to time indices.
+            If False, searches for nearest time value (use for node-only data).
+        line_color: Color of the vertical line
+        line_width: Width of the vertical line in pixels
+        line_dash: Dash style - "solid", "dash", "dot", "dashdot"
+        annotation_text: Text to show on the line
+        annotation_position: Position of annotation - "top", "bottom", "top left", etc.
+        folder_name: Optional GUI folder name to organize plots
+        aspect: Aspect ratio for plot display (width/height)
+
+    Returns:
+        Tuple of (plot_handle, update_callback)
+
+    Example::
+
+        from openscvx.plotting import plot_control, viser
+
+        # Create any plotly figure
+        fig = plot_control(results, "thrust_force")
+
+        # Add to viser with animated vertical line
+        _, update_plot = viser.add_animated_plotly_vline(
+            server, fig,
+            time_array=results.trajectory["time"].flatten(),
+        )
+
+        # Add to animation callbacks
+        update_callbacks.append(update_plot)
+    """
+    # Detect number of subplots in the figure
+    # Count unique xaxis/yaxis references in the layout
+    n_subplots = 1
+    if hasattr(fig, "_grid_ref") and fig._grid_ref is not None:
+        # Figure created with make_subplots - use grid dimensions
+        n_rows = len(fig._grid_ref)
+        n_cols = len(fig._grid_ref[0]) if n_rows > 0 else 1
+        n_subplots = n_rows * n_cols
+
+    # Track which shapes are our vlines (before adding new ones)
+    n_existing_shapes = len(fig.layout.shapes) if fig.layout.shapes else 0
+
+    # Add vertical line to each subplot
+    if n_subplots == 1:
+        # Single plot - add one vline
+        fig.add_vline(
+            x=time_array[0],
+            line_dash=line_dash,
+            line_color=line_color,
+            line_width=line_width,
+            annotation_text=annotation_text,
+            annotation_position=annotation_position,
+        )
+    else:
+        # Multiple subplots - add vline to each
+        # Determine grid layout
+        n_rows = len(fig._grid_ref)
+        n_cols = len(fig._grid_ref[0]) if n_rows > 0 else 1
+
+        for row_idx in range(n_rows):
+            for col_idx in range(n_cols):
+                # Add vline to this subplot
+                # Only add annotation to first subplot to avoid clutter
+                show_annotation = row_idx == 0 and col_idx == 0
+                fig.add_vline(
+                    x=time_array[0],
+                    line_dash=line_dash,
+                    line_color=line_color,
+                    line_width=line_width,
+                    annotation_text=annotation_text if show_annotation else None,
+                    annotation_position=annotation_position if show_annotation else None,
+                    row=row_idx + 1,
+                    col=col_idx + 1,
+                )
+
+    # Track indices of shapes we added
+    n_new_shapes = len(fig.layout.shapes) - n_existing_shapes
+    vline_shape_indices = list(range(n_existing_shapes, n_existing_shapes + n_new_shapes))
+
+    # Add to viser GUI
+    if folder_name:
+        with server.gui.add_folder(folder_name):
+            plot_handle = server.gui.add_plotly(figure=fig, aspect=aspect)
+    else:
+        plot_handle = server.gui.add_plotly(figure=fig, aspect=aspect)
+
+    # Create update callback
+    def update(frame_idx: int) -> None:
+        """Update vertical line position based on current frame."""
+        if use_trajectory_indexing:
+            # Direct indexing: frame_idx corresponds to time index
+            idx = min(frame_idx, len(time_array) - 1)
+        else:
+            # Search for nearest time (for node-only data)
+            current_time = time_array[frame_idx]
+            idx = min(frame_idx, len(time_array) - 1)
+
+        # Update all vertical line positions
+        current_time = time_array[idx]
+        for shape_idx in vline_shape_indices:
+            fig.layout.shapes[shape_idx].x0 = current_time
+            fig.layout.shapes[shape_idx].x1 = current_time
+
+        # Trigger viser update
+        plot_handle.figure = fig
+
+    return plot_handle, update
+
+
+def add_animated_plotly_marker(
+    server: viser.ViserServer,
+    fig: go.Figure,
+    time_array: np.ndarray,
+    marker_x_data: np.ndarray,
+    marker_y_data: np.ndarray,
+    use_trajectory_indexing: bool = True,
+    marker_name: str = "Current",
+    marker_color: str = "red",
+    marker_size: int = 12,
+    folder_name: str | None = None,
+    aspect: float = 1.5,
+) -> tuple:
+    """Add a plotly figure to viser GUI with an animated time marker.
+
+    This function takes any plotly figure and adds an animated marker that
+    synchronizes with viser's 3D animation timeline. The marker shows the
+    current position on the plot as the animation plays.
+
+    Args:
+        server: ViserServer instance
+        fig: Plotly figure to display
+        time_array: Time values corresponding to animation frames (N,).
+            This should match the time array passed to add_animation_controls().
+        marker_x_data: X-axis values for marker position (N,)
+        marker_y_data: Y-axis values for marker position (N,)
+        use_trajectory_indexing: If True, frame_idx maps directly to data indices.
+            If False, searches for nearest time value (use for node-only data).
+        marker_name: Legend name for the marker trace
+        marker_color: Color of the animated marker
+        marker_size: Size of the animated marker in points
+        folder_name: Optional GUI folder name to organize plots
+        aspect: Aspect ratio for plot display (width/height)
+
+    Returns:
+        Tuple of (plot_handle, update_callback)
+
+    Example::
+
+        from openscvx.plotting import plot_vector_norm, viser
+
+        # Create any plotly figure
+        fig = plot_vector_norm(results, "thrust")
+        thrust_norms = np.linalg.norm(results.trajectory["thrust"], axis=1)
+
+        # Add to viser with animated marker
+        _, update_plot = viser.add_animated_plotly_marker(
+            server, fig,
+            time_array=results.trajectory["time"].flatten(),
+            marker_x_data=results.trajectory["time"].flatten(),
+            marker_y_data=thrust_norms,
+        )
+
+        # Add to animation callbacks
+        update_callbacks.append(update_plot)
+    """
+    # Add marker trace to figure
+    marker_trace = go.Scatter(
+        x=[marker_x_data[0]],
+        y=[marker_y_data[0]],
+        mode="markers",
+        marker={"color": marker_color, "size": marker_size, "symbol": "circle"},
+        name=marker_name,
+    )
+    fig.add_trace(marker_trace)
+    marker_trace_idx = len(fig.data) - 1
+
+    # Add to viser GUI
+    if folder_name:
+        with server.gui.add_folder(folder_name):
+            plot_handle = server.gui.add_plotly(figure=fig, aspect=aspect)
+    else:
+        plot_handle = server.gui.add_plotly(figure=fig, aspect=aspect)
+
+    # Create update callback
+    def update(frame_idx: int) -> None:
+        """Update marker position based on current frame."""
+        if use_trajectory_indexing:
+            # Direct indexing: frame_idx corresponds to data index
+            idx = min(frame_idx, len(marker_y_data) - 1)
+        else:
+            # Search for nearest time (for node-only data)
+            current_time = time_array[frame_idx]
+            idx = min(np.searchsorted(marker_x_data, current_time), len(marker_y_data) - 1)
+
+        # Update marker position
+        fig.data[marker_trace_idx].x = [marker_x_data[idx]]
+        fig.data[marker_trace_idx].y = [marker_y_data[idx]]
+
+        # Trigger viser update
+        plot_handle.figure = fig
+
+    return plot_handle, update
+
+
+def add_animated_vector_norm_plot(
+    server: viser.ViserServer,
+    results: OptimizationResults,
+    var_name: str,
+    bounds: tuple[float, float] | None = None,
+    title: str | None = None,
+    folder_name: str | None = None,
+    aspect: float = 1.5,
+    marker_color: str = "red",
+    marker_size: int = 12,
+) -> tuple:
+    """Add animated norm plot for a state or control variable.
+
+    Convenience wrapper around add_animated_plotly_marker() that uses
+    the existing plot_vector_norm() function to create the base plot.
+
+    Args:
+        server: ViserServer instance
+        results: Optimization results containing variable data
+        var_name: Name of the state or control variable to plot
+        bounds: Optional (min, max) bounds to display on plot
+        title: Optional custom title for the plot (defaults to "‖{var_name}‖₂")
+        folder_name: Optional GUI folder name to organize plots
+        aspect: Aspect ratio for plot display (width/height)
+        marker_color: Color of the animated marker
+        marker_size: Size of the animated marker in points
+
+    Returns:
+        Tuple of (plot_handle, update_callback), or (None, None) if variable not found
+
+    Example::
+
+        from openscvx.plotting import viser
+
+        # Add animated thrust norm plot
+        _, update_thrust = viser.add_animated_vector_norm_plot(
+            server, results, "thrust",
+            title="Thrust Magnitude",
+            bounds=(0.0, max_thrust),
+            folder_name="Control Plots"
+        )
+        if update_thrust is not None:
+            update_callbacks.append(update_thrust)
+    """
+    from openscvx.plotting import plot_vector_norm
+
+    # Check if variable exists in results
+    has_in_trajectory = bool(results.trajectory) and var_name in results.trajectory
+    has_in_nodes = var_name in results.nodes
+
+    if not (has_in_trajectory or has_in_nodes):
+        import warnings
+
+        warnings.warn(f"Variable '{var_name}' not found in results, skipping plot")
+        return None, None
+
+    # Create figure using existing plotting function
+    fig = plot_vector_norm(results, var_name, bounds=bounds)
+
+    # Update title if custom title provided
+    if title is not None:
+        fig.update_layout(title_text=title)
+
+    # Determine data source and compute norms
+    if has_in_trajectory:
+        time_data = results.trajectory["time"].flatten()
+        var_data = results.trajectory[var_name]
+        use_trajectory_indexing = True
+    else:
+        time_data = results.nodes["time"].flatten()
+        var_data = results.nodes[var_name]
+        use_trajectory_indexing = False
+
+    # Compute norms
+    norm_data = np.linalg.norm(var_data, axis=1) if var_data.ndim > 1 else np.abs(var_data)
+
+    # Add animated marker
+    return add_animated_plotly_marker(
+        server,
+        fig,
+        time_array=time_data,
+        marker_x_data=time_data,
+        marker_y_data=norm_data,
+        use_trajectory_indexing=use_trajectory_indexing,
+        marker_name="Current",
+        marker_color=marker_color,
+        marker_size=marker_size,
+        folder_name=folder_name,
+        aspect=aspect,
+    )

openscvx/plotting/viser/primitives.py

@@ -0,0 +1,355 @@
+"""Static scene primitives for viser visualization.
+
+Functions for adding non-animated elements: obstacles, gates, constraint cones,
+ghost trajectories, etc. Called once during scene setup.
+"""
+
+import numpy as np
+import viser
+
+
+def _generate_ellipsoid_mesh(
+    center: np.ndarray,
+    radii: np.ndarray,
+    axes: np.ndarray | None = None,
+    subdivisions: int = 2,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate ellipsoid mesh vertices and faces via icosphere subdivision.
+
+    Args:
+        center: Center position (3,)
+        radii: Radii along each principal axis (3,)
+        axes: Rotation matrix (3, 3) defining principal axes. If None, uses identity.
+        subdivisions: Number of icosphere subdivisions (higher = smoother)
+
+    Returns:
+        Tuple of (vertices, faces) where vertices is (V, 3) and faces is (F, 3)
+    """
+    # Start with icosahedron vertices
+    phi = (1.0 + np.sqrt(5.0)) / 2.0  # Golden ratio
+    icosahedron_verts = np.array(
+        [
+            [-1, phi, 0],
+            [1, phi, 0],
+            [-1, -phi, 0],
+            [1, -phi, 0],
+            [0, -1, phi],
+            [0, 1, phi],
+            [0, -1, -phi],
+            [0, 1, -phi],
+            [phi, 0, -1],
+            [phi, 0, 1],
+            [-phi, 0, -1],
+            [-phi, 0, 1],
+        ],
+        dtype=np.float64,
+    )
+    # Normalize to unit sphere
+    icosahedron_verts /= np.linalg.norm(icosahedron_verts[0])
+
+    icosahedron_faces = np.array(
+        [
+            [0, 11, 5],
+            [0, 5, 1],
+            [0, 1, 7],
+            [0, 7, 10],
+            [0, 10, 11],
+            [1, 5, 9],
+            [5, 11, 4],
+            [11, 10, 2],
+            [10, 7, 6],
+            [7, 1, 8],
+            [3, 9, 4],
+            [3, 4, 2],
+            [3, 2, 6],
+            [3, 6, 8],
+            [3, 8, 9],
+            [4, 9, 5],
+            [2, 4, 11],
+            [6, 2, 10],
+            [8, 6, 7],
+            [9, 8, 1],
+        ],
+        dtype=np.int32,
+    )
+
+    vertices = icosahedron_verts
+    faces = icosahedron_faces
+
+    # Subdivide faces
+    for _ in range(subdivisions):
+        new_faces = []
+        midpoint_cache = {}
+
+        def get_midpoint(i1: int, i2: int) -> int:
+            """Get or create midpoint vertex between two vertices."""
+            key = (min(i1, i2), max(i1, i2))
+            if key in midpoint_cache:
+                return midpoint_cache[key]
+
+            nonlocal vertices
+            p1, p2 = vertices[i1], vertices[i2]
+            mid = (p1 + p2) / 2.0
+            mid = mid / np.linalg.norm(mid)  # Project onto unit sphere
+
+            idx = len(vertices)
+            vertices = np.vstack([vertices, mid])
+            midpoint_cache[key] = idx
+            return idx
+
+        for tri in faces:
+            v0, v1, v2 = tri
+            a = get_midpoint(v0, v1)
+            b = get_midpoint(v1, v2)
+            c = get_midpoint(v2, v0)
+            new_faces.extend([[v0, a, c], [v1, b, a], [v2, c, b], [a, b, c]])
+
+        faces = np.array(new_faces, dtype=np.int32)
+
+    # Scale by radii to create ellipsoid
+    vertices = vertices / radii
+
+    # Rotate by principal axes if provided
+    if axes is not None:
+        vertices = vertices @ axes.T
+
+    # Translate to center
+    vertices = vertices + center
+
+    return vertices.astype(np.float32), faces
+
+
+def add_ellipsoid_obstacles(
+    server: viser.ViserServer,
+    centers: list[np.ndarray],
+    radii: list[np.ndarray],
+    axes: list[np.ndarray] | None = None,
+    color: tuple[int, int, int] = (255, 100, 100),
+    opacity: float = 0.6,
+    wireframe: bool = False,
+    subdivisions: int = 2,
+) -> list:
+    """Add ellipsoidal obstacles to the scene.
+
+    Args:
+        server: ViserServer instance
+        centers: List of center positions, each shape (3,)
+        radii: List of radii along principal axes, each shape (3,)
+        axes: List of rotation matrices (3, 3) defining principal axes.
+            If None, ellipsoids are axis-aligned.
+        color: RGB color tuple
+        opacity: Opacity (0-1), only used when wireframe=False
+        wireframe: If True, render as wireframe instead of solid
+        subdivisions: Icosphere subdivisions (higher = smoother, 2 is usually good)
+
+    Returns:
+        List of mesh handles
+    """
+    handles = []
+
+    if axes is None:
+        axes = [None] * len(centers)
+
+    for i, (center, rad, ax) in enumerate(zip(centers, radii, axes)):
+        # Convert JAX arrays to numpy if needed
+        center = np.asarray(center, dtype=np.float64)
+        rad = np.asarray(rad, dtype=np.float64)
+        if ax is not None:
+            ax = np.asarray(ax, dtype=np.float64)
+
+        vertices, faces = _generate_ellipsoid_mesh(center, rad, ax, subdivisions)
+
+        handle = server.scene.add_mesh_simple(
+            f"/obstacles/ellipsoid_{i}",
+            vertices=vertices,
+            faces=faces,
+            color=color,
+            wireframe=wireframe,
+            opacity=opacity if not wireframe else 1.0,
+        )
+        handles.append(handle)
+
+    return handles
+
+
+def add_gates(
+    server: viser.ViserServer,
+    vertices: list,
+    color: tuple[int, int, int] = (255, 165, 0),
+    line_width: float = 3.0,
+) -> None:
+    """Add gate/obstacle wireframes to the scene.
+
+    Args:
+        server: ViserServer instance
+        vertices: List of vertex arrays (4 vertices for planar gate, 8 for box)
+        color: RGB color tuple
+        line_width: Line width for wireframe
+    """
+    for i, verts in enumerate(vertices):
+        verts = np.array(verts)
+        n_verts = len(verts)
+
+        if n_verts == 4:
+            # Planar gate: 4 vertices forming a closed loop
+            edges = [[0, 1], [1, 2], [2, 3], [3, 0]]
+        elif n_verts == 8:
+            # 3D box: 8 vertices
+            edges = [
+                [0, 1],
+                [1, 2],
+                [2, 3],
+                [3, 0],  # front face
+                [4, 5],
+                [5, 6],
+                [6, 7],
+                [7, 4],  # back face
+                [0, 4],
+                [1, 5],
+                [2, 6],
+                [3, 7],  # connecting edges
+            ]
+        else:
+            # Unknown format, skip
+            continue
+
+        # Shape (N, 2, 3) for N line segments
+        points = np.array([[verts[e[0]], verts[e[1]]] for e in edges])
+        server.scene.add_line_segments(
+            f"/gates/gate_{i}",
+            points=points,
+            colors=color,
+            line_width=line_width,
+        )
+
+
+def _generate_cone_mesh(
+    apex: np.ndarray,
+    height: float,
+    half_angle_deg: float,
+    n_segments: int = 32,
+) -> tuple[np.ndarray, np.ndarray]:
+    """Generate a cone mesh with apex at given position, opening upward.
+
+    Args:
+        apex: Apex position (3,) - the tip of the cone
+        height: Height of the cone (extends in +Z direction from apex)
+        half_angle_deg: Half-angle of the cone from the vertical axis in degrees
+        n_segments: Number of segments around the circumference
+
+    Returns:
+        Tuple of (vertices, faces) where vertices is (V, 3) and faces is (F, 3)
+    """
+    half_angle_rad = np.radians(half_angle_deg)
+    base_radius = height * np.tan(half_angle_rad)
+
+    # Vertices: apex + base circle points
+    vertices = [apex.copy()]  # Apex at index 0
+
+    # Base circle vertices
+    for i in range(n_segments):
+        angle = 2 * np.pi * i / n_segments
+        x = apex[0] + base_radius * np.cos(angle)
+        y = apex[1] + base_radius * np.sin(angle)
+        z = apex[2] + height
+        vertices.append([x, y, z])
+
+    # Center of base for closing the bottom
+    base_center = apex.copy()
+    base_center[2] += height
+    vertices.append(base_center)  # Index n_segments + 1
+
+    vertices = np.array(vertices, dtype=np.float32)
+
+    # Faces: triangles from apex to base edge pairs, plus base cap
+    faces = []
+
+    # Side faces (apex to each edge of base)
+    for i in range(n_segments):
+        next_i = (i + 1) % n_segments
+        # Triangle: apex, base[i], base[next_i]
+        faces.append([0, i + 1, next_i + 1])
+
+    # Base cap faces (center to each edge)
+    base_center_idx = n_segments + 1
+    for i in range(n_segments):
+        next_i = (i + 1) % n_segments
+        # Triangle: center, base[next_i], base[i] (reverse winding for outward normal)
+        faces.append([base_center_idx, next_i + 1, i + 1])
+
+    faces = np.array(faces, dtype=np.int32)
+
+    return vertices, faces
+
+
+def add_glideslope_cone(
+    server: viser.ViserServer,
+    apex: np.ndarray | tuple = (0.0, 0.0, 0.0),
+    height: float = 2000.0,
+    glideslope_angle_deg: float = 86.0,
+    color: tuple[int, int, int] = (100, 200, 100),
+    opacity: float = 0.2,
+    wireframe: bool = False,
+    n_segments: int = 32,
+) -> viser.MeshHandle:
+    """Add a glideslope constraint cone to the scene.
+
+    The glideslope constraint typically has the form:
+        ||position_xy|| <= tan(angle) * position_z
+
+    This creates a cone with apex at the landing site, opening upward.
+
+    Args:
+        server: ViserServer instance
+        apex: Apex position (landing site), default is origin
+        height: Height of the cone visualization
+        glideslope_angle_deg: Glideslope angle in degrees (measured from vertical).
+            For constraint ||r_xy|| <= tan(theta) * z, pass theta here.
+            Common values: 86 deg (very wide), 70 deg (moderate), 45 deg (steep)
+        color: RGB color tuple
+        opacity: Opacity (0-1)
+        wireframe: If True, render as wireframe
+        n_segments: Number of segments for cone smoothness
+
+    Returns:
+        Mesh handle for the cone
+    """
+    apex = np.asarray(apex, dtype=np.float32)
+
+    vertices, faces = _generate_cone_mesh(apex, height, glideslope_angle_deg, n_segments)
+
+    handle = server.scene.add_mesh_simple(
+        "/constraints/glideslope_cone",
+        vertices=vertices,
+        faces=faces,
+        color=color,
+        wireframe=wireframe,
+        opacity=opacity if not wireframe else 1.0,
+    )
+
+    return handle
+
+
+def add_ghost_trajectory(
+    server: viser.ViserServer,
+    pos: np.ndarray,
+    colors: np.ndarray,
+    opacity: float = 0.3,
+    point_size: float = 0.05,
+) -> None:
+    """Add a faint ghost trajectory showing the full path.
+
+    Args:
+        server: ViserServer instance
+        pos: Position array of shape (N, 3)
+        colors: RGB color array of shape (N, 3)
+        opacity: Opacity factor (0-1) applied to colors
+        point_size: Size of trajectory points
+    """
+    ghost_colors = (colors * opacity).astype(np.uint8)
+    server.scene.add_point_cloud(
+        "/ghost_traj",
+        points=pos,
+        colors=ghost_colors,
+        point_size=point_size,
+    )