Visualization

visualization

Visualization: SE3 trajectories and invariant plots.

Functions

plot_invariants

plot_invariants(linear_inv, angular_inv, ax=None, title='DHB invariants')

Plot linear and angular invariant time series.

Source code in src/dhb_xr/visualization/plot.py

def plot_invariants(
    linear_inv: np.ndarray,
    angular_inv: np.ndarray,
    ax: Optional[Any] = None,
    title: str = "DHB invariants",
) -> Any:
    """Plot linear and angular invariant time series."""
    if not HAS_MPL:
        raise ImportError("matplotlib required for plot_invariants")
    n = linear_inv.shape[0]
    k_lin, k_ang = linear_inv.shape[1], angular_inv.shape[1]
    if ax is None:
        fig, axes = plt.subplots(2, 1, sharex=True)
        ax_lin, ax_ang = axes[0], axes[1]
    else:
        ax_lin, ax_ang = ax
    t = np.arange(n)
    for j in range(k_lin):
        ax_lin.plot(t, linear_inv[:, j], label=f"lin_{j}")
    ax_lin.set_ylabel("linear")
    ax_lin.legend(loc="right", fontsize=8)
    ax_lin.set_title(title)
    for j in range(k_ang):
        ax_ang.plot(t, angular_inv[:, j], label=f"ang_{j}")
    ax_ang.set_ylabel("angular")
    ax_ang.legend(loc="right", fontsize=8)
    ax_ang.set_xlabel("step")
    return (ax_lin, ax_ang)

plot_se3_trajectory

plot_se3_trajectory(
    positions,
    quaternions=None,
    ax=None,
    title="SE(3) trajectory",
    show_orientation=False,
    vis_type="arrow",
    box_size=(0.03, 0.03, 0.03),
    color="b",
    alpha=0.7,
    num_frames=8,
    label=None,
)

Plot 3D position trajectory with optional orientation visualization.

Parameters:

Name	Type	Description	Default
positions	ndarray	(N, 3) position array	required
quaternions	Optional[ndarray]	(N, 4) quaternion array (wxyz), optional	None
ax	Optional[Any]	Matplotlib 3D axes, created if None	None
title	str	Plot title	'SE(3) trajectory'
show_orientation	bool	If True, show orientation frames or cubes	False
vis_type	str	"arrow" for coordinate frames, "cube" for oriented boxes	'arrow'
box_size	tuple	Size of cubes (if vis_type="cube")	(0.03, 0.03, 0.03)
color	str	Trajectory line color	'b'
alpha	float	Transparency for cubes	0.7
num_frames	int	Number of orientation samples to show along trajectory	8
label	Optional[str]	Legend label for trajectory	None

Returns:

Type	Description
Any	Matplotlib axes

Source code in src/dhb_xr/visualization/plot.py

def plot_se3_trajectory(
    positions: np.ndarray,
    quaternions: Optional[np.ndarray] = None,
    ax: Optional[Any] = None,
    title: str = "SE(3) trajectory",
    show_orientation: bool = False,
    vis_type: str = "arrow",
    box_size: tuple = (0.03, 0.03, 0.03),
    color: str = "b",
    alpha: float = 0.7,
    num_frames: int = 8,
    label: Optional[str] = None,
) -> Any:
    """
    Plot 3D position trajectory with optional orientation visualization.

    Args:
        positions: (N, 3) position array
        quaternions: (N, 4) quaternion array (wxyz), optional
        ax: Matplotlib 3D axes, created if None
        title: Plot title
        show_orientation: If True, show orientation frames or cubes
        vis_type: "arrow" for coordinate frames, "cube" for oriented boxes
        box_size: Size of cubes (if vis_type="cube")
        color: Trajectory line color
        alpha: Transparency for cubes
        num_frames: Number of orientation samples to show along trajectory
        label: Legend label for trajectory

    Returns:
        Matplotlib axes
    """
    if not HAS_MPL:
        raise ImportError("matplotlib required for plot_se3_trajectory")
    positions = np.asarray(positions)
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], "-", color=color, label=label or "position", alpha=0.8)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.set_title(title)

    if show_orientation and quaternions is not None:
        from dhb_xr.core import geometry as geom
        n = len(positions)
        indices = np.linspace(0, n - 1, num_frames, dtype=int)
        for i in indices:
            R = geom.quat_to_rot(quaternions[i])
            p = positions[i]
            if vis_type == "cube":
                draw_box(ax, p, R, size=box_size, color=color, alpha=alpha)
            else:  # arrow
                draw_frame(ax, p, R, length=0.05)
    return ax

Overview

Utilities for plotting and visualizing DHB trajectories and invariants.

Main Functions

plot_se3_trajectory

plot

Plot SE(3) trajectories and invariant sequences with cubes and coordinate frames.

Functions

draw_box

draw_box(ax, center, rotation_matrix, size=(0.05, 0.05, 0.05), color='blue', alpha=0.3)

Draw a 3D oriented box (cuboid) at the given center with the given rotation.

Parameters:

Name	Type	Description	Default
ax		Matplotlib 3D axes	required
center	ndarray	(3,) position of box center	required
rotation_matrix	ndarray	(3, 3) rotation matrix	required
size	tuple	(width, height, depth) of box	(0.05, 0.05, 0.05)
color	str	Face color	'blue'
alpha	float	Transparency	0.3

Source code in src/dhb_xr/visualization/plot.py

def draw_box(
    ax,
    center: np.ndarray,
    rotation_matrix: np.ndarray,
    size: tuple = (0.05, 0.05, 0.05),
    color: str = "blue",
    alpha: float = 0.3,
) -> None:
    """
    Draw a 3D oriented box (cuboid) at the given center with the given rotation.

    Args:
        ax: Matplotlib 3D axes
        center: (3,) position of box center
        rotation_matrix: (3, 3) rotation matrix
        size: (width, height, depth) of box
        color: Face color
        alpha: Transparency
    """
    # Cuboid vertices before rotation (unit cube centered at origin)
    vertices = np.array([
        [0.5, 0.5, -0.5], [-0.5, 0.5, -0.5], [-0.5, -0.5, -0.5], [0.5, -0.5, -0.5],  # Bottom
        [0.5, 0.5, 0.5], [-0.5, 0.5, 0.5], [-0.5, -0.5, 0.5], [0.5, -0.5, 0.5]  # Top
    ], dtype=np.float64)

    # Scale and rotate vertices
    vertices = vertices * np.array(size)
    vertices = np.dot(vertices, rotation_matrix.T)
    vertices += center

    # Define the indices of vertices for each face of the cuboid
    faces = [
        [vertices[0], vertices[1], vertices[2], vertices[3]],
        [vertices[4], vertices[5], vertices[6], vertices[7]],
        [vertices[0], vertices[3], vertices[7], vertices[4]],
        [vertices[1], vertices[2], vertices[6], vertices[5]],
        [vertices[0], vertices[1], vertices[5], vertices[4]],
        [vertices[2], vertices[3], vertices[7], vertices[6]]
    ]

    # Create a Poly3DCollection object
    face_collection = Poly3DCollection(faces, facecolors=color, linewidths=1, edgecolors='k', alpha=alpha)
    ax.add_collection3d(face_collection)

draw_frame

draw_frame(ax, position, rotation_matrix, length=0.05, linewidth=2.0)

Draw a coordinate frame (x=red, y=green, z=blue arrows) at the given pose.

Parameters:

Name	Type	Description	Default
ax		Matplotlib 3D axes	required
position	ndarray	(3,) position of frame origin	required
rotation_matrix	ndarray	(3, 3) rotation matrix	required
length	float	Length of arrows	0.05
linewidth	float	Line width	2.0

Source code in src/dhb_xr/visualization/plot.py

def draw_frame(
    ax,
    position: np.ndarray,
    rotation_matrix: np.ndarray,
    length: float = 0.05,
    linewidth: float = 2.0,
) -> None:
    """
    Draw a coordinate frame (x=red, y=green, z=blue arrows) at the given pose.

    Args:
        ax: Matplotlib 3D axes
        position: (3,) position of frame origin
        rotation_matrix: (3, 3) rotation matrix
        length: Length of arrows
        linewidth: Line width
    """
    for j, c in enumerate(AXIS_COLORS):
        ax.quiver(
            position[0], position[1], position[2],
            rotation_matrix[0, j], rotation_matrix[1, j], rotation_matrix[2, j],
            color=c, length=length, normalize=True, linewidth=linewidth
        )

plot_invariants

plot_invariants(linear_inv, angular_inv, ax=None, title='DHB invariants')

Plot linear and angular invariant time series.

Source code in src/dhb_xr/visualization/plot.py

def plot_invariants(
    linear_inv: np.ndarray,
    angular_inv: np.ndarray,
    ax: Optional[Any] = None,
    title: str = "DHB invariants",
) -> Any:
    """Plot linear and angular invariant time series."""
    if not HAS_MPL:
        raise ImportError("matplotlib required for plot_invariants")
    n = linear_inv.shape[0]
    k_lin, k_ang = linear_inv.shape[1], angular_inv.shape[1]
    if ax is None:
        fig, axes = plt.subplots(2, 1, sharex=True)
        ax_lin, ax_ang = axes[0], axes[1]
    else:
        ax_lin, ax_ang = ax
    t = np.arange(n)
    for j in range(k_lin):
        ax_lin.plot(t, linear_inv[:, j], label=f"lin_{j}")
    ax_lin.set_ylabel("linear")
    ax_lin.legend(loc="right", fontsize=8)
    ax_lin.set_title(title)
    for j in range(k_ang):
        ax_ang.plot(t, angular_inv[:, j], label=f"ang_{j}")
    ax_ang.set_ylabel("angular")
    ax_ang.legend(loc="right", fontsize=8)
    ax_ang.set_xlabel("step")
    return (ax_lin, ax_ang)

plot_se3_trajectories

plot_se3_trajectories(
    trajectories,
    ax=None,
    show_orientation=True,
    vis_type="cube",
    box_size_scale=0.05,
    num_frames=6,
    title="SE(3) Trajectories",
    show_legend=True,
)

Plot multiple SE(3) trajectories with orientation visualization.

Parameters:

Name	Type	Description	Default
trajectories	Union[List[Dict], Dict[str, Dict]]	List of dicts with 'positions' and 'quaternions', or dict of name -> trajectory	required
ax	Optional[Any]	Matplotlib 3D axes	None
show_orientation	bool	Show orientation cubes/frames	True
vis_type	str	"cube" or "arrow"	'cube'
box_size_scale	float	Size of cubes	0.05
num_frames	int	Number of orientation samples per trajectory	6
title	str	Plot title	'SE(3) Trajectories'
show_legend	bool	Show legend	True

Returns:

Type	Description
Any	Matplotlib axes

Source code in src/dhb_xr/visualization/plot.py

def plot_se3_trajectories(
    trajectories: Union[List[Dict], Dict[str, Dict]],
    ax: Optional[Any] = None,
    show_orientation: bool = True,
    vis_type: str = "cube",
    box_size_scale: float = 0.05,
    num_frames: int = 6,
    title: str = "SE(3) Trajectories",
    show_legend: bool = True,
) -> Any:
    """
    Plot multiple SE(3) trajectories with orientation visualization.

    Args:
        trajectories: List of dicts with 'positions' and 'quaternions', or dict of name -> trajectory
        ax: Matplotlib 3D axes
        show_orientation: Show orientation cubes/frames
        vis_type: "cube" or "arrow"
        box_size_scale: Size of cubes
        num_frames: Number of orientation samples per trajectory
        title: Plot title
        show_legend: Show legend

    Returns:
        Matplotlib axes
    """
    if not HAS_MPL:
        raise ImportError("matplotlib required for plot_se3_trajectories")

    # Normalize input to dict
    if isinstance(trajectories, list):
        trajectories = {f"traj_{i}": t for i, t in enumerate(trajectories)}

    if ax is None:
        fig = plt.figure(figsize=(10, 8))
        ax = fig.add_subplot(111, projection="3d")

    colors = ["blue", "green", "red", "orange", "purple", "cyan", "magenta", "brown"]
    box_size = (box_size_scale, box_size_scale, box_size_scale)

    for idx, (name, traj) in enumerate(trajectories.items()):
        positions = np.asarray(traj["positions"])
        quaternions = np.asarray(traj.get("quaternions"))
        color = colors[idx % len(colors)]

        ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], color=color, label=name, linewidth=2)

        if show_orientation and quaternions is not None:
            from dhb_xr.core import geometry as geom
            n = len(positions)
            indices = np.linspace(0, n - 1, num_frames, dtype=int)
            for i in indices:
                R = geom.quat_to_rot(quaternions[i])
                p = positions[i]
                if vis_type == "cube":
                    draw_box(ax, p, R, size=box_size, color=color, alpha=0.5)
                else:
                    draw_frame(ax, p, R, length=0.05)

    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.set_title(title)
    if show_legend:
        ax.legend()

    return ax

plot_se3_trajectory

plot_se3_trajectory(
    positions,
    quaternions=None,
    ax=None,
    title="SE(3) trajectory",
    show_orientation=False,
    vis_type="arrow",
    box_size=(0.03, 0.03, 0.03),
    color="b",
    alpha=0.7,
    num_frames=8,
    label=None,
)

Plot 3D position trajectory with optional orientation visualization.

Parameters:

Name	Type	Description	Default
positions	ndarray	(N, 3) position array	required
quaternions	Optional[ndarray]	(N, 4) quaternion array (wxyz), optional	None
ax	Optional[Any]	Matplotlib 3D axes, created if None	None
title	str	Plot title	'SE(3) trajectory'
show_orientation	bool	If True, show orientation frames or cubes	False
vis_type	str	"arrow" for coordinate frames, "cube" for oriented boxes	'arrow'
box_size	tuple	Size of cubes (if vis_type="cube")	(0.03, 0.03, 0.03)
color	str	Trajectory line color	'b'
alpha	float	Transparency for cubes	0.7
num_frames	int	Number of orientation samples to show along trajectory	8
label	Optional[str]	Legend label for trajectory	None

Returns:

Type	Description
Any	Matplotlib axes

Source code in src/dhb_xr/visualization/plot.py

def plot_se3_trajectory(
    positions: np.ndarray,
    quaternions: Optional[np.ndarray] = None,
    ax: Optional[Any] = None,
    title: str = "SE(3) trajectory",
    show_orientation: bool = False,
    vis_type: str = "arrow",
    box_size: tuple = (0.03, 0.03, 0.03),
    color: str = "b",
    alpha: float = 0.7,
    num_frames: int = 8,
    label: Optional[str] = None,
) -> Any:
    """
    Plot 3D position trajectory with optional orientation visualization.

    Args:
        positions: (N, 3) position array
        quaternions: (N, 4) quaternion array (wxyz), optional
        ax: Matplotlib 3D axes, created if None
        title: Plot title
        show_orientation: If True, show orientation frames or cubes
        vis_type: "arrow" for coordinate frames, "cube" for oriented boxes
        box_size: Size of cubes (if vis_type="cube")
        color: Trajectory line color
        alpha: Transparency for cubes
        num_frames: Number of orientation samples to show along trajectory
        label: Legend label for trajectory

    Returns:
        Matplotlib axes
    """
    if not HAS_MPL:
        raise ImportError("matplotlib required for plot_se3_trajectory")
    positions = np.asarray(positions)
    if ax is None:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
    ax.plot(positions[:, 0], positions[:, 1], positions[:, 2], "-", color=color, label=label or "position", alpha=0.8)
    ax.set_xlabel("x")
    ax.set_ylabel("y")
    ax.set_zlabel("z")
    ax.set_title(title)

    if show_orientation and quaternions is not None:
        from dhb_xr.core import geometry as geom
        n = len(positions)
        indices = np.linspace(0, n - 1, num_frames, dtype=int)
        for i in indices:
            R = geom.quat_to_rot(quaternions[i])
            p = positions[i]
            if vis_type == "cube":
                draw_box(ax, p, R, size=box_size, color=color, alpha=alpha)
            else:  # arrow
                draw_frame(ax, p, R, length=0.05)
    return ax

Usage Example

from dhb_xr.visualization.plot import plot_se3_trajectory
import matplotlib.pyplot as plt

# Plot trajectory
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')

plot_se3_trajectory(
    ax, positions, quaternions,
    color='blue', label='Trajectory'
)

ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()