"""
This file is the implementation of the environment plot to visualize the environment objects.
Author: Ruihua Han
"""
from __future__ import annotations
import glob
import os
import re
import shutil
from collections.abc import Iterable
from math import cos, sin
from typing import Any
import imageio
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import matplotlib.transforms as mtransforms
import mpl_toolkits.mplot3d.art3d as art3d
import numpy as np
from matplotlib.lines import Line2D
from matplotlib.patches import (
Arrow,
Circle,
Ellipse,
PathPatch,
Polygon,
Rectangle,
Wedge,
)
from matplotlib.path import Path
from mpl_toolkits.mplot3d import Axes3D
from shapely.geometry import MultiPolygon as ShapelyMultiPolygon
from shapely.geometry import Polygon as ShapelyPolygon
from shapely.geometry.polygon import orient
from irsim.config import palette_param
from irsim.config.path_param import path_manager as pm
from irsim.util.util import points_to_xy_list, traj_to_xy_list
_FMT_COLOR = re.compile(r"C\d+|[bgrcmykw]")
def _with_default_color(
fmt: str, kwargs: dict[str, Any], default: str
) -> dict[str, Any]:
"""Return ``kwargs`` with ``color=default`` unless a color is already given.
A color counts as given when ``kwargs`` has ``color`` or ``c``, or when
the Matplotlib format string ``fmt`` names one: a single letter such as
``r``, a ``C0``-style cycle color, or a full color name.
"""
if "color" in kwargs or "c" in kwargs:
return kwargs
if _FMT_COLOR.search(fmt) or mcolors.is_color_like(fmt):
return kwargs
return {**kwargs, "color": default}
[docs]
class EnvPlot:
"""
EnvPlot class for visualizing the environment.
Args:
world: The world object with ranges and a ``plot_parse`` dictionary that
configures plotting (e.g., saved_figure, figure_pixels, show_title,
no_axis, tight).
objects (list, optional): Initial objects to draw. Default is ``[]``.
kwargs: Plot overrides that update ``world.plot_parse`` at runtime
(e.g., ``saved_figure``, ``figure_pixels``, ``show_title``,
``no_axis``, ``tight``).
"""
def __init__(
self,
world: Any,
objects: list[Any] | None = None,
**kwargs: Any,
) -> None:
"""
Initialize the EnvPlot instance.
Sets up the matplotlib figure, configures plotting parameters,
and initializes the plot with world data and objects. Any provided
``kwargs`` override values in ``world.plot_parse``.
"""
world.plot_parse.update(kwargs)
# default saved figure kwargs
self.saved_figure_kwargs: dict[str, Any] = {
"dpi": 100,
"bbox_inches": "tight",
}
self.saved_figure_kwargs.update(world.plot_parse.get("saved_figure", {}))
figure_pixels = world.plot_parse.get("figure_pixels", [1000, 800])
if objects is None:
objects = []
# Create matplotlib figure and axes
self.fig, self.ax = plt.subplots(
figsize=(
figure_pixels[0] / self.saved_figure_kwargs["dpi"],
figure_pixels[1] / self.saved_figure_kwargs["dpi"],
),
dpi=self.saved_figure_kwargs["dpi"],
)
# Default size for saved animation frames; refined on the first saved
# frame (see save_figure) so every frame is pinned to it
# (bbox_inches="tight" otherwise recomputes it from the window).
self._save_size = self.fig.get_size_inches().copy()
self.viewpoint = world.plot_parse.get("viewpoint", None)
# Initialize dynamic plotting lists
self.dyna_line_list: list[Any] = []
self.dyna_point_list: list[Any] = []
self.dyna_quiver_list: list[Any] = []
# Grid-map overlay artist (managed by draw_grid_map), reused across
# reset/reload so artists don't stack.
self._grid_im: Any = None
# Set in _init_plot; pre-set so the first call has a defined previous world.
self.world: Any = None
# Initialize the plot with world data
self._init_plot(world, objects, **kwargs)
def _init_plot(
self,
world: Any,
objects: list[Any],
**kwargs: Any,
) -> None:
"""
Initialize or re-initialize drawing on the existing figure/axes.
Safe to call when the world and objects change (e.g., YAML reload);
it reuses the current figure without creating a new window.
Args:
world: World instance providing ranges and grid map.
objects (list): List of objects to initialize and draw.
**kwargs: Plot overrides merged into ``world.plot_parse``.
"""
# Clear the previous world's fog overlay before swapping worlds, so a
# reload's new FogMap doesn't orphan the old artist on the reused axes.
if self.world is not None and self.world.fog_map is not None:
self.world.fog_map.plot_clear()
world.plot_parse.update(kwargs)
self.world = world
self.saved_ani_kwargs: dict[str, Any] = {}
# Frozen crop box for animation frames (computed lazily on first save).
self._save_bbox: Any = None
self.title = world.plot_parse.get("title", None)
self.show_title = world.plot_parse.get("show_title", True)
self.saved_figure_kwargs.update(world.plot_parse.get("saved_figure", {}))
if isinstance(self.ax, Axes3D):
self.ax.set_box_aspect([1, 1, 1])
else:
self.ax.set_aspect("equal")
self.ax.set_xlim(self.world.x_range)
self.ax.set_ylim(self.world.y_range)
# Centre the camera before the first frame is drawn: the limits above
# frame the whole world, so a follow viewpoint would otherwise show one
# frame at the world centre and only jump to its target on the next step.
self.set_ax_viewpoint(objects)
self.ax.set_xlabel("x [m]")
self.ax.set_ylabel("y [m]")
self.init_objects_plot(objects)
self.draw_grid_map(world.grid_map)
self.draw_fog(world.fog_map)
if world.plot_parse.get("no_axis", False):
plt.axis("off")
if world.plot_parse.get("tight", True):
self.fig.tight_layout()
[docs]
def step(
self, mode: str = "dynamic", objects: list[Any] | None = None, **kwargs: Any
) -> None:
"""Advance the plot by one step for the given objects.
Args:
mode (str): Which objects to update: "dynamic", "static", or "all".
objects (list | None): The objects to update/draw. Defaults to empty list.
**kwargs: Extra drawing options passed through to objects' plot methods.
"""
if objects is None:
objects = []
if self.show_title:
self.update_title()
self.set_ax_viewpoint(objects)
if isinstance(self.ax, Axes3D):
self.clear_components(mode, objects)
self.draw_components(mode, objects, **kwargs)
else:
self.clear_components(mode)
self.step_objects_plot(mode, objects, **kwargs)
self.step_world_plot()
[docs]
def init_objects_plot(self, objects: list[Any], **kwargs: Any) -> None:
"""Initialize plot state for provided objects, then render once.
Args:
objects (list): Objects to be initialized on the axes.
**kwargs: Extra drawing options passed to initialization/plot.
"""
if self.show_title:
self.update_title()
[obj._init_plot(self.ax, **kwargs) for obj in objects]
self.step_objects_plot("all", objects, **kwargs)
[docs]
def step_objects_plot(
self, mode: str = "dynamic", objects: list[Any] | None = None, **kwargs: Any
) -> None:
"""
Update the plot for the objects by transform based on the object's original geometry.
"""
if objects is None:
objects = []
if mode == "dynamic":
# A sensor on a static object still watches a changing scene, so
# the object is refreshed whenever it carries sensors.
[
obj._step_plot(**kwargs)
for obj in objects
if not obj.static or obj.sensors
]
elif mode == "static":
[obj._step_plot(**kwargs) for obj in objects if obj.static]
elif mode == "all":
[obj._step_plot(**kwargs) for obj in objects]
else:
self.logger.error("Error: Invalid draw mode")
[docs]
def draw_components(
self, mode: str = "all", objects: list[Any] | None = None, **kwargs: Any
) -> None:
"""
Draw the components in the environment with global axis.
Args:
mode (str): 'static', 'dynamic', or 'all' to specify which objects to draw.
objects (list): List of objects to draw.
kwargs: Additional plotting options.
"""
if objects is None:
objects = []
if mode == "static":
for obj in objects:
if obj.static:
obj.plot(self.ax, **kwargs)
elif mode == "dynamic":
for obj in objects:
if not obj.static:
obj.plot(self.ax, **kwargs)
elif mode == "all":
for obj in objects:
obj.plot(self.ax, **kwargs)
else:
self.logger.error("Error: Invalid draw mode")
[docs]
def clear_components(
self, mode: str = "all", objects: list[Any] | None = None
) -> None:
"""
Clear the components in the environment.
Args:
mode (str): 'static', 'dynamic', or 'all' to specify which objects to clear.
objects (list): List of objects to clear.
"""
if objects is None:
objects = []
if mode == "dynamic":
for obj in objects:
if not obj.static:
obj.plot_clear()
self._clear_dynamic_artists()
elif mode == "static":
for obj in objects:
if obj.static:
obj.plot_clear()
elif mode == "all":
for obj in objects:
obj.plot_clear(all=True)
self._clear_dynamic_artists()
else:
self.logger.error("Error: Invalid clear mode")
def _clear_dynamic_artists(self) -> None:
"""Remove the per-step lines, points, and quivers, then forget them."""
for line in self.dyna_line_list:
line.pop(0).remove()
for points in self.dyna_point_list:
points.remove()
for quiver in self.dyna_quiver_list:
quiver.remove()
self.dyna_line_list = []
self.dyna_point_list = []
self.dyna_quiver_list = []
[docs]
def draw_grid_map(self, grid_map: Any | None = None, **kwargs: Any) -> None:
"""
Draw the grid map on the plot.
Args:
grid_map (optional): The grid map to draw.
"""
if grid_map is None:
# Reloading into a world without an obstacle map: drop the old grid.
if self._grid_im is not None:
self._grid_im.remove()
self._grid_im = None
return
# Reuse the existing grid image on repeated _init_plot (e.g. on every
# env.reset() / reload()) instead of stacking new imshow artists, which
# leaks memory and slows interactive redraws.
if self._grid_im is None:
self._grid_im = self.ax.imshow(
grid_map.T,
cmap="Greys",
origin="lower",
extent=self.x_range + self.y_range,
zorder=0,
)
else:
self._grid_im.set_data(grid_map.T)
# extent may change if a reload swaps in a differently sized world
self._grid_im.set_extent(self.x_range + self.y_range)
if isinstance(self.ax, Axes3D) and self.logger is not None:
self.logger.warning("Map will not show in 3D plot")
[docs]
def draw_fog(
self, fog_map: Any | None = None, zorder: int = 2, **kwargs: Any
) -> None:
"""Initialize the fog-of-map overlay.
The :class:`~irsim.world.map.FogMap` owns its artist and refreshes it via
:meth:`step_world_plot`; this just hands it the axes. Skipped in 3D. On
``env.reset()`` the same ``FogMap`` is reused and refreshed; on reload the
previous overlay is cleared in :meth:`_init_plot`.
Args:
fog_map: A :class:`~irsim.world.map.FogMap`, or ``None`` to skip.
zorder: Drawing layer for the overlay.
"""
if fog_map is None or isinstance(self.ax, Axes3D):
return
fog_map._init_plot(self.ax, zorder=zorder, **kwargs)
[docs]
def step_world_plot(self) -> None:
"""Refresh per-step world-level overlays (not tied to a single object).
Currently steps the fog-of-map overlay; extend here for other
world-level visuals updated every step. ``_step_plot`` is a no-op when
there is no artist (fog disabled or 3D), so no extra guard is needed.
"""
fog_map = self.world.fog_map
if fog_map is not None:
fog_map._step_plot()
[docs]
def draw_trajectory(
self,
traj: list[Any] | np.ndarray,
traj_type: str = "-",
label: str = "trajectory",
show_direction: bool = False,
refresh: bool = False,
**kwargs: Any,
) -> None:
"""
Draw a trajectory on the plot.
Args:
traj (list or np.ndarray): List of points or array of points [x, y, theta].
traj_type (str): Matplotlib format string for the line (e.g. '-',
'r--'). Unless it names a color, or ``color`` is passed, the
line uses ``palette_param.path``.
See https://matplotlib.org/3.1.1/api/_as_gen/matplotlib.pyplot.plot.html for details.
label (str): Label for the trajectory.
show_direction (bool): Whether to show the direction of the trajectory.
refresh (bool): Whether to refresh the plot.
kwargs: Additional plotting options for ax.plot()
"""
path_x_list, path_y_list = traj_to_xy_list(traj)
kwargs = _with_default_color(traj_type, kwargs, palette_param.path)
line = self.ax.plot(path_x_list, path_y_list, traj_type, label=label, **kwargs)
if show_direction:
if isinstance(traj, list):
u_list = [cos(p[2, 0]) for p in traj]
v_list = [sin(p[2, 0]) for p in traj]
elif isinstance(traj, np.ndarray):
u_list = [cos(p[2]) for p in traj.T]
v_list = [sin(p[2]) for p in traj.T]
if isinstance(self.ax, Axes3D):
path_z_list = [0] * len(path_x_list)
w_list = [0] * len(u_list)
self.ax.quiver(
path_x_list, path_y_list, path_z_list, u_list, v_list, w_list
)
else:
self.ax.quiver(path_x_list, path_y_list, u_list, v_list, width=0.003)
if refresh:
self.dyna_line_list.append(line)
[docs]
def draw_points(
self,
points: list[Any] | np.ndarray | None,
s: int = 10,
c: str | None = None,
refresh: bool = True,
**kwargs: Any,
) -> None:
"""
Draw points on the plot.
Args:
points (list): List of points, each point as [x, y] or (2, 1) array
or (np.array): points array: (2, N), NL number of points
s (int): Size of the points.
c (str): Color of the points.
refresh (bool): Whether to refresh the plot.
kwargs: Additional plotting options.
"""
if points is None:
return
x_coordinates, y_coordinates = points_to_xy_list(points)
if c is None:
c = palette_param.marker
points_plot = self.ax.scatter(x_coordinates, y_coordinates, s, c, **kwargs)
if refresh:
self.dyna_point_list.append(points_plot)
[docs]
def draw_quiver(
self,
point: np.ndarray | None,
refresh: bool = False,
color: str | None = None,
**kwargs: Any,
) -> None:
"""
Draw a quiver plot on the plot.
Args:
points (4*1 np.ndarray): List of points, each point as [x, y, u, v]. u, v are the components of the vector.
kwargs: Additional plotting options.
"""
if point is None:
return
if color is None:
color = palette_param.quiver
ax_point = self.ax.scatter(point[0], point[1], color=color)
ax_quiver = self.ax.quiver(
point[0],
point[1], # starting positions
point[2],
point[3], # vector components (direction)
color=color,
**kwargs,
)
if refresh:
self.dyna_quiver_list.append(ax_quiver)
self.dyna_point_list.append(ax_point)
[docs]
def draw_quivers(
self,
points: Iterable[np.ndarray],
refresh: bool = False,
color: str | None = None,
**kwargs: Any,
) -> None:
"""
Draw a series of quiver plot on the plot.
Args:
points (list or np.ndarray): List of points, each point as [x, y, u, v]. u, v are the components of the vector.
"""
for point in points:
self.draw_quiver(point, refresh, color=color, **kwargs)
[docs]
def draw_box(
self, vertices: np.ndarray, refresh: bool = False, color: str = "-"
) -> None:
"""
Draw a box by the vertices.
Args:
vertices (np.ndarray): 2xN array of vertices.
refresh (bool): Whether to refresh the plot.
color (str): Matplotlib format string for the outline; unless it
names a color, ``palette_param.path``
is used.
"""
temp_vertex = np.c_[vertices, vertices[0:2, 0]]
box_line = self.ax.plot(
temp_vertex[0, :],
temp_vertex[1, :],
color,
**_with_default_color(color, {}, palette_param.path),
)
if refresh:
self.dyna_line_list.append(box_line)
[docs]
def update_title(self) -> None:
"""Update the figure title with current time/status or a custom title."""
if self.title is not None:
self.ax.set_title(self.title, pad=3)
else:
self.ax.set_title(
f"Simulation Time: {self.world.time:.2f}s, Status: {self.world.status}",
pad=3,
)
def _init_save_bbox(self) -> Any:
"""Tight crop box of the initial figure, captured once and reused.
Reusing one crop box (with the figure pinned to its initial size) keeps
every animation frame the same size, regardless of how the window was
resized or zoomed. The mp4 encoder's even-size requirement is handled at
write time (ffmpeg pad), not here, since bbox->pixel rounding can't
guarantee an even pixel count.
"""
self.fig.set_size_inches(self._save_size, forward=False)
self.fig.canvas.draw()
pad = self.saved_figure_kwargs.get("pad_inches", 0.1)
return self.fig.get_tightbbox(self.fig.canvas.get_renderer()).padded(pad)
[docs]
def save_animate(
self,
ani_name: str = "animation",
suffix: str = ".gif",
last_frame_duration: int = 1,
rm_fig_path: bool = True,
**kwargs: Any,
) -> None:
"""
Save the animation by streaming frames to avoid loading all images into memory.
Args:
ani_name (str): Name of the animation. Default is 'animation'.
last_frame_duration (int): Duration of the last frame for the gif. Default is 1 second.
suffix (str): Suffix of the animation file. Default is '.gif'.
rm_fig_path (bool): Whether to remove the figure path after saving. Default is True.
kwargs: Additional arguments for saving the animation.
For GIF: See pillow plugin documentation.
For video: See ffmpeg/pyav plugin documentation.
"""
self.saved_ani_kwargs.update(kwargs)
self.logger.info("Start to create animation")
ap = pm.ani_path
fp = pm.ani_buffer_path
if not os.path.exists(ap):
os.makedirs(ap)
images = list(glob.glob(fp + "/*.png"))
images.sort()
if not images:
self.logger.warning("No images found to create animation")
return
full_name = ap + "/" + ani_name + suffix
num_images = len(images)
if suffix == ".gif":
# GIF frame timing in milliseconds
frame_duration_ms = 100
# Per-frame durations in milliseconds
durations_ms = [frame_duration_ms] * num_images
# Extend the last frame duration (input argument is seconds)
last_frame_duration_ms = int(last_frame_duration * 1000)
if num_images > 0 and last_frame_duration_ms > frame_duration_ms:
durations_ms[-1] = last_frame_duration_ms
gif_kwargs = self.saved_ani_kwargs.copy()
# Convert milliseconds to seconds for imageio writer
with imageio.get_writer(
full_name,
mode="I",
loop=0,
duration=durations_ms,
**gif_kwargs,
) as writer:
for image_path in images:
frame = imageio.imread(str(image_path))
writer.append_data(frame)
else:
# Video format (e.g., .mp4) - stream frames to encoder
video_kwargs = self.saved_ani_kwargs.copy()
fps = video_kwargs.pop("fps", 10)
# H.264 needs even dimensions; the tight crop is often odd. Let ffmpeg
# pad (not resize) each frame up to even, so the video isn't distorted
# and the encoder doesn't reject odd sizes. macro_block_size=1 stops
# imageio from doing its own resize first.
video_kwargs.setdefault("macro_block_size", 1)
video_kwargs.setdefault(
"output_params", ["-vf", "pad=ceil(iw/2)*2:ceil(ih/2)*2:color=white"]
)
# Use get_writer for memory-efficient streaming writes
with imageio.get_writer(full_name, fps=fps, **video_kwargs) as writer:
for image_path in images:
frame = imageio.imread(str(image_path))
writer.append_data(frame)
self.logger.info(f"{ani_name} created successfully, saved in {ap}")
if rm_fig_path:
shutil.rmtree(fp)
[docs]
def set_ax_viewpoint(self, objects: list[Any] | None = None) -> None:
"""
Set the viewpoint of the plot windows by the viewpoint parameter.
Args:
objects (list): List of objects to search for the viewpoint.
"""
if self.viewpoint is None:
point = None
elif isinstance(self.viewpoint, str):
if objects is None:
target = None
else:
target = next(
(
obj
for obj in objects
if getattr(obj, "name", None) == self.viewpoint
),
None,
)
point = target.state[0:2].flatten().tolist() if target is not None else None
elif isinstance(self.viewpoint, list):
point = self.viewpoint
else:
point = None
if point is None:
pass
else:
# Preserve current zoom span if available; otherwise use world size
cur_xlim = self.ax.get_xlim()
cur_ylim = self.ax.get_ylim()
cur_width = cur_xlim[1] - cur_xlim[0]
cur_height = cur_ylim[1] - cur_ylim[0]
view_point_x1 = point[0] - cur_width / 2.0
view_point_x2 = point[0] + cur_width / 2.0
view_point_y1 = point[1] - cur_height / 2.0
view_point_y2 = point[1] + cur_height / 2.0
self.ax.set_xlim([view_point_x1, view_point_x2])
self.ax.set_ylim([view_point_y1, view_point_y2])
[docs]
def show(self) -> None:
"""
Display the plot.
"""
plt.show()
[docs]
def close(self) -> None:
"""
Close this plot's figure.
"""
plt.close(self.fig)
@property
def _world_param(self):
"""Access world_param via world's instance if available."""
if hasattr(self.world, "_wp"):
return self.world._wp
from irsim.config import world_param
return world_param
@property
def _env_param(self):
"""Access env_param, fallback to global."""
from irsim.config import env_param
return env_param
@property
def logger(self):
"""Environment logger used by plotting helpers."""
return self._env_param.logger
@property
def x_range(self):
"""World x-axis range used by the plot."""
return self.world.x_range
@property
def y_range(self):
"""World y-axis range used by the plot."""
return self.world.y_range
[docs]
def linewidth_from_data_units(
linewidth: float, axis: Any, reference: str = "y"
) -> float:
"""
Convert a linewidth in data units to linewidth in points.
Parameters
----------
linewidth: float
Linewidth in data units of the respective reference-axis
axis: matplotlib axis
The axis which is used to extract the relevant transformation
data (data limits and size must not change afterwards)
reference: string
The axis that is taken as a reference for the data width.
Possible values: 'x' and 'y'. Defaults to 'y'.
Returns
-------
linewidth: float
Linewidth in points
"""
fig = axis.get_figure()
if reference == "x":
length = fig.bbox_inches.width * axis.get_position().width
value_range = np.diff(axis.get_xlim()).item()
elif reference == "y":
length = fig.bbox_inches.height * axis.get_position().height
value_range = np.diff(axis.get_ylim()).item()
# Convert length to points
length *= 72
# Scale linewidth to value range
return linewidth * (length / value_range)
[docs]
def geometry_to_path(geometry: Any) -> Path:
"""Convert a Shapely polygonal geometry to a Matplotlib path.
Polygon interiors are preserved as oppositely oriented subpaths, allowing
one ``PathPatch`` to render Polygon and MultiPolygon geometries including
holes.
Args:
geometry: Shapely Polygon or MultiPolygon in local coordinates.
Returns:
matplotlib.path.Path: Compound path containing every polygon ring.
Raises:
ValueError: If the geometry is empty or not polygonal.
"""
if isinstance(geometry, ShapelyPolygon):
polygons = [geometry]
elif isinstance(geometry, ShapelyMultiPolygon):
polygons = list(geometry.geoms)
else:
raise ValueError("compound rendering requires a Polygon or MultiPolygon")
paths = []
for polygon in polygons:
polygon = orient(polygon, sign=1.0)
for ring in (polygon.exterior, *polygon.interiors):
vertices = np.asarray(ring.coords, dtype=float)
codes = np.full(len(vertices), Path.LINETO, dtype=np.uint8)
codes[0] = Path.MOVETO
codes[-1] = Path.CLOSEPOLY
paths.append(Path(vertices, codes))
if not paths:
raise ValueError("compound rendering requires a non-empty geometry")
return Path.make_compound_path(*paths)
[docs]
def draw_patch(
ax: Any,
shape: str,
state: np.ndarray | None = None,
radius: float | None = None,
vertices: np.ndarray | None = None,
color: str | None = None,
zorder: int | None = None,
linestyle: str | None = None,
geometry: Any | None = None,
**kwargs: Any,
) -> Any:
"""
Draw a geometric element (patch or line) on the given axes.
Supported shapes and expected inputs (refer to irsim.world.object_base plotting patterns):
- circle: use ``state`` (x,y,theta), ``radius``, and optional ``center`` (body-frame
offset); created in the body frame and transformed, matching the collision geometry
- rectangle: prefer body-frame ``vertices`` (2xN) else use ``width``/``height``
with an optional ``center``; transformed by ``state``
- polygon: use body-frame ``vertices`` (2xN); transformed by ``state``
- compound: use a local-frame Shapely Polygon or MultiPolygon as ``geometry``
- ellipse: use ``width``/``height`` with ``state`` transform
- wedge: use ``radius`` and either ``theta1``/``theta2`` (deg) or ``fov`` (rad); transformed by ``state``
- arrow: use ``state`` for position/orientation or provide ``theta``; supports ``arrow_length`` and ``arrow_width``
- line|linestring: use ``vertices`` (2xN) in world coordinates; ``state`` is not
applied to them
Coordinate frames:
- A patch is built in the object's body frame and placed by ``state``, which
rotates by theta and then translates by (x, y). Pass body-frame vertices -
:py:attr:`~irsim.world.object_base.ObjectBase.original_vertices` - together
with the pose, as ``ObjectPlot`` does; passing vertices that already carry
the pose applies it twice. Pass ``state=None`` for vertices that are
already in world coordinates.
- A line is added with its vertices as data coordinates and receives no
transform, so its vertices are always world coordinates.
Styling:
- color/edgecolor/facecolor, alpha, zorder, linestyle are respected when applicable.
Returns the created matplotlib artist (patch, line, or list from ax.plot).
"""
state = state if state is not None else np.zeros((3, 1))
# Normalize some common styling kwargs
facecolor = kwargs.pop("facecolor", None)
edgecolor = kwargs.pop("edgecolor", None)
alpha = kwargs.pop("alpha", None)
fill = kwargs.pop("fill", None)
if shape in ("line", "linestring"):
# Lines are artists rather than patches, so they are styled directly.
return _draw_line(ax, vertices, color, alpha, zorder, linestyle, kwargs)
if shape == "arrow":
# The arrow is built in world coordinates and has its own color and
# zorder fallbacks, so it is styled without a state transform.
created_element = ax.add_patch(_arrow_patch(state, kwargs))
set_patch_property(
created_element,
ax,
state=None, # world coords already applied
color=color if color is not None else kwargs.pop("arrow_color", None),
alpha=alpha,
zorder=zorder if zorder is not None else kwargs.pop("arrow_zorder", None),
fill=fill,
)
else:
builder = _PATCH_BUILDERS.get(shape)
if builder is None:
raise ValueError(f"Unsupported shape type: {shape}")
created_element = ax.add_patch(
builder(radius=radius, vertices=vertices, geometry=geometry, **kwargs)
)
set_patch_property(
created_element,
ax,
state=state,
color=color,
facecolor=facecolor,
edgecolor=edgecolor,
alpha=alpha,
zorder=zorder,
linestyle=linestyle,
fill=fill,
)
# 3D conversion for patches if needed
if isinstance(ax, Axes3D) and created_element is not None:
art3d.patch_2d_to_3d(created_element, z=kwargs.pop("z", 0), zdir="z")
return created_element
def _circle_patch(radius: float | None = None, center: Any = None, **_: Any) -> Circle:
"""Build a circle at the body-frame center, ready for the state transform."""
if radius is None:
raise ValueError("circle requires radius")
if center is None:
xy = (0.0, 0.0)
else:
c = np.asarray(center, dtype=float).flatten()
xy = (c[0], c[1])
return Circle(xy, radius)
def _rectangle_patch(
vertices: np.ndarray | None = None,
width: float | None = None,
height: float | None = None,
center: Any = None,
**_: Any,
) -> Any:
"""Build a rectangle from its body-frame vertices, or from width and height.
``center`` is the body-frame center of the box, as it is for a circle. It
defaults to the origin, which matches the collision geometry of a plain
rectangle; an Ackermann box, offset from the origin by half a wheelbase,
passes its own centroid so the two still coincide.
"""
if vertices is not None:
return Polygon(vertices.T)
if width is None or height is None:
raise ValueError("rectangle requires either vertices or width/height")
if center is None:
cx, cy = 0.0, 0.0
else:
c = np.asarray(center, dtype=float).flatten()
cx, cy = float(c[0]), float(c[1])
return Rectangle((cx - width / 2.0, cy - height / 2.0), width, height)
def _polygon_patch(vertices: np.ndarray | None = None, **_: Any) -> Polygon:
"""Build a polygon from its body-frame vertices, placed later by ``state``."""
if vertices is None:
raise ValueError("polygon requires vertices (2xN)")
return Polygon(vertices.T)
def _compound_patch(geometry: Any = None, **_: Any) -> PathPatch:
"""Build one patch for a Polygon or MultiPolygon, holes included."""
if geometry is None:
raise ValueError("compound requires a Polygon or MultiPolygon geometry")
return PathPatch(geometry_to_path(geometry))
def _ellipse_patch(
width: float | None = None, height: float | None = None, **_: Any
) -> Ellipse:
"""Build a body-frame ellipse, ready for the state transform."""
if width is None or height is None:
raise ValueError("ellipse requires width and height")
return Ellipse((0.0, 0.0), width, height, angle=0.0)
def _wedge_patch(
radius: float | None = None,
theta1: float | None = None,
theta2: float | None = None,
fov: float | None = None,
**_: Any,
) -> Wedge:
"""Build a field-of-view wedge from explicit angles, or from an fov."""
if radius is None:
raise ValueError("wedge requires radius")
if theta1 is None or theta2 is None:
# build from fov in radians centered at 0
fov = np.pi if fov is None else fov
theta1 = -180 * fov / (2 * np.pi)
theta2 = 180 * fov / (2 * np.pi)
return Wedge((0.0, 0.0), radius, theta1, theta2)
# Shapes drawn as a patch: each builder takes what it needs and ignores the rest,
# after which draw_patch adds and styles the patch the same way for all of them.
_PATCH_BUILDERS = {
"circle": _circle_patch,
"rectangle": _rectangle_patch,
"polygon": _polygon_patch,
"compound": _compound_patch,
"ellipse": _ellipse_patch,
"wedge": _wedge_patch,
}
def _arrow_patch(state: np.ndarray, kwargs: dict[str, Any]) -> Arrow:
"""Build a heading arrow in world coordinates, from the pose in ``state``."""
arrow_length = kwargs.pop("arrow_length", 0.4)
arrow_width = kwargs.pop("arrow_width", 0.6)
# Orientation: use provided theta or state[2]
theta = kwargs.pop("theta", float(state[2, 0]) if state.shape[0] >= 3 else 0.0)
return Arrow(
float(state[0, 0]),
float(state[1, 0]),
float(arrow_length * cos(theta)),
float(arrow_length * sin(theta)),
width=arrow_width,
)
def _draw_line(
ax: Any,
vertices: np.ndarray | None,
color: str | None,
alpha: float | None,
zorder: int | None,
linestyle: str | None,
kwargs: dict[str, Any],
) -> Any:
"""Add a 2D or 3D line and style it directly.
The vertices are used as data coordinates: a line carries no patch
transform, so ``state`` never reaches it and the vertices must be in world
coordinates. ``fill`` does not apply to a line and is ignored.
"""
if vertices is None:
raise ValueError("line/linestring requires vertices (2xN)")
if isinstance(ax, Axes3D):
zs = kwargs.pop("z", None)
if zs is None:
# One height per vertex: a shorter default breaks the draw.
zs = np.zeros(vertices.shape[1])
line3d = art3d.Line3D(vertices[0, :], vertices[1, :], zs=zs)
if color is not None:
line3d.set_color(color)
if alpha is not None:
line3d.set_alpha(alpha)
if zorder is not None:
line3d.set_zorder(zorder)
ax.add_line(line3d)
return line3d
line2d = Line2D(vertices[0, :], vertices[1, :])
if linestyle is not None:
line2d.set_linestyle(linestyle)
if color is not None:
line2d.set_color(color)
if alpha is not None:
line2d.set_alpha(alpha)
if zorder is not None:
line2d.set_zorder(zorder)
ax.add_line(line2d)
return line2d
[docs]
def set_patch_property(
element: Any,
ax: Any,
state: np.ndarray | None = None,
**kwargs: Any,
) -> None:
"""
Apply transform and style properties to a patch/artist, matching ObjectBase semantics.
- If ``state`` provided with at least 3 rows, apply rotation by theta and translation by (x, y).
- Apply color/facecolor/edgecolor, alpha, zorder, linestyle when supported.
"""
# Transform by state (x, y, theta)
if state is not None and state.shape[0] >= 3:
x = float(state[0, 0])
y = float(state[1, 0])
theta = float(state[2, 0])
transform = mtransforms.Affine2D().rotate(theta).translate(x, y) + ax.transData
if hasattr(element, "set_transform"):
element.set_transform(transform)
# Styling
color = kwargs.get("color")
facecolor = kwargs.get("facecolor")
edgecolor = kwargs.get("edgecolor")
alpha = kwargs.get("alpha")
zorder = kwargs.get("zorder")
linestyle = kwargs.get("linestyle")
linewidth = kwargs.get("linewidth")
fill = kwargs.get("fill")
if color is not None and hasattr(element, "set_color"):
element.set_color(color)
if facecolor is not None and hasattr(element, "set_facecolor"):
element.set_facecolor(facecolor)
if edgecolor is not None and hasattr(element, "set_edgecolor"):
element.set_edgecolor(edgecolor)
if alpha is not None and hasattr(element, "set_alpha"):
element.set_alpha(alpha)
if zorder is not None and hasattr(element, "set_zorder"):
element.set_zorder(zorder)
if linestyle is not None and hasattr(element, "set_linestyle"):
element.set_linestyle(linestyle)
if linewidth is not None and hasattr(element, "set_linewidth"):
element.set_linewidth(linewidth)
if fill is not None and hasattr(element, "set_fill"):
element.set_fill(fill)
