diff --git a/README.md b/README.md index f54cc743..a630a009 100755 --- a/README.md +++ b/README.md @@ -142,7 +142,7 @@ python benchmarks/qp_solver_comparison.py #### Multi-robot 3D coordination -Cinematic 3D simulation rendering with Manim. Multi-robot reach-avoid in 3D, rendered via CBFKit's Manim backend. Shows the visualization stack scales from quick matplotlib plots to publication-quality 3D animations. +Cinematic 3D simulation rendering with Manim. Multi-robot reach-avoid in 3D, rendered via CBFKit's Manim backend. Shows the visualization stack scales from quick matplotlib plots to publication-quality 3D animations. The same backend also renders 2D `CBFAnimator` scenes: pass `backend="manim"` (or `"manim-"`) to any `CBFAnimator` and `save("out.mp4")` produces a publication-quality video (`.gif` also supported).

Manim 3D render of multi-robot reach-avoid

diff --git a/examples/unicycle/reach_goal/manim_2d_animation.py b/examples/unicycle/reach_goal/manim_2d_animation.py new file mode 100644 index 00000000..3fb11092 --- /dev/null +++ b/examples/unicycle/reach_goal/manim_2d_animation.py @@ -0,0 +1,116 @@ +"""Tutorial: render a CBF-filtered unicycle reach-avoid run with the Manim 2D backend. + +Runs the README quick-start simulation (unicycle drives to a goal while a CBF +safety filter keeps it clear of an obstacle), then renders the trajectory with +``CBFAnimator(backend="manim-medium")``. This is the script that produced +``media/showcase/manim_2d_animator.gif``. + +Requires the ``manim`` extra (``pip install cbfkit[manim]``) plus ffmpeg, and +on macOS the cairo/pango libraries (``brew install ffmpeg cairo pango``). +""" + +import os + +import jax.numpy as jnp +from jax import jit + +from cbfkit.certificates import concatenate_certificates, rectify_relative_degree +from cbfkit.certificates.barrier_functions import ellipsoidal_barrier_factory +from cbfkit.certificates.conditions.barrier_conditions.zeroing_barriers import linear_class_k +from cbfkit.controllers.cbf_clf import vanilla_cbf_clf_qp_controller +from cbfkit.estimators import naive +from cbfkit.integration import runge_kutta_4 +from cbfkit.sensors import perfect +from cbfkit.simulation import simulator +from cbfkit.systems.unicycle.models.olfatisaber2002approximate.dynamics import ( + approx_unicycle_dynamics, +) +from cbfkit.utils.animators.deps import _HAS_MANIM + +# Simulation Parameters +initial_state = jnp.array([0.0, 0.0, 0.0]) +actuation_limits = jnp.array([5.0, jnp.pi]) +goal = jnp.array([4.0, 0.0]) +obstacle_center = jnp.array([2.0, 0.5, 0.0]) +obstacle_radii = jnp.array([0.5, 0.5]) +dt = 1e-2 +num_steps = 500 if not os.getenv("CBFKIT_TEST_MODE") else 50 + +# Dynamics +dynamics = approx_unicycle_dynamics(lam=1.0) # state: [x, y, theta] + + +# Nominal controller - drives toward the goal, ignorant of the obstacle +@jit +def nominal_controller(t, state, key, data): + x, y, th = state + heading = jnp.arctan2(goal[1] - y, goal[0] - x) + return ( + jnp.array( + [ + jnp.linalg.norm(jnp.array([x - goal[0], y - goal[1]])), # speed + jnp.arctan2(jnp.sin(heading - th), jnp.cos(heading - th)), # steering + ] + ), + {}, + ) + + +# CBF barrier around the obstacle +cbf_factory, _, _ = ellipsoidal_barrier_factory( + system_position_indices=(0, 1), + obstacle_position_indices=(0, 1), + ellipsoid_axis_indices=(0, 1), +) +barrier = rectify_relative_degree( + function=cbf_factory(obstacle_center, obstacle_radii), + system_dynamics=dynamics, + state_dim=3, + form="exponential", +)(certificate_conditions=linear_class_k(10.0)) + +# Safety-filtered controller +controller = vanilla_cbf_clf_qp_controller( + control_limits=actuation_limits, + dynamics_func=dynamics, + barriers=concatenate_certificates(barrier), +) + +results = simulator.execute( + x0=initial_state, + dt=dt, + num_steps=num_steps, + dynamics=dynamics, + integrator=runge_kutta_4, + nominal_controller=nominal_controller, + controller=controller, + sensor=perfect, + estimator=naive, +) +print(f"Final position: ({results.states[-1, 0]:.2f}, {results.states[-1, 1]:.2f})") + +# Render with the Manim 2D backend (skipped in test mode / without manim) +if os.getenv("CBFKIT_TEST_MODE"): + print("CBFKIT_TEST_MODE set: skipping Manim render.") +elif not _HAS_MANIM: + print("Manim not found. Rendering disabled. Install 'cbfkit[manim]' to enable it.") +else: + import numpy as np + + from cbfkit.utils.animator import CBFAnimator + + anim = CBFAnimator( + np.asarray(results.states), + dt=dt, + backend="manim-medium", # or "manim" / "manim-low" for quicker renders + title="CBF Safety Filter: Unicycle Reach-Avoid", + aspect="equal", + x_lim=(-0.7, 4.7), + y_lim=(-1.2, 1.8), + ) + anim.add_goal(np.asarray(goal), radius=0.25, color="g", label="Goal") + anim.add_obstacle(np.asarray(obstacle_center[:2]), radius=float(obstacle_radii[0]), alpha=0.3) + anim.add_agent(x_idx=0, y_idx=1, body_radius=0.12, body_color="tab:blue", trail=True) + anim.show_time() + out = anim.save("manim_2d_unicycle_reach_avoid.mp4") + print(f"Animation saved to {out}") diff --git a/media/showcase/manim_2d_animator.gif b/media/showcase/manim_2d_animator.gif new file mode 100644 index 00000000..e476bbfc Binary files /dev/null and b/media/showcase/manim_2d_animator.gif differ diff --git a/mypy.ini b/mypy.ini index 5a0cf492..0d498ace 100644 --- a/mypy.ini +++ b/mypy.ini @@ -40,6 +40,9 @@ ignore_missing_imports = True [mypy-matplotlib.*] ignore_missing_imports = True +[mypy-manim.*] +ignore_missing_imports = True + [mypy-scipy.*] ignore_missing_imports = True diff --git a/src/cbfkit/utils/animators/__init__.py b/src/cbfkit/utils/animators/__init__.py index ced6aeb0..0a4f19e0 100644 --- a/src/cbfkit/utils/animators/__init__.py +++ b/src/cbfkit/utils/animators/__init__.py @@ -9,12 +9,13 @@ * ``"matplotlib"`` — generates MP4 / GIF animations via ``FuncAnimation``. * ``"plotly"`` (default) — generates interactive HTML files with play/pause controls and a timeline slider. Requires ``pip install cbfkit[plotly]``. -* ``"manim"`` — high-quality 3D animations (MP4) via Manim. Currently only - available for 3D multi-robot scenes via :func:`visualize_3d_multi_robot`. - Requires ``pip install cbfkit[manim]``. +* ``"manim"`` / ``"manim-"`` — high-quality MP4 / GIF animations via + Manim, for both :class:`CBFAnimator` 2D scenes and 3D multi-robot scenes + (:func:`visualize_3d_multi_robot`). Requires ``pip install cbfkit[manim]``. """ from .animator import CBFAnimator +from .manim_backend import CBFAnimator2DScene from .config import AnimationConfig, DEFAULT_CONFIG from .deps import ( _HAS_MANIM, @@ -29,6 +30,7 @@ __all__ = [ "AnimationConfig", "CBFAnimator", + "CBFAnimator2DScene", "DEFAULT_CONFIG", "save_animation", "_HAS_MATPLOTLIB", diff --git a/src/cbfkit/utils/animators/animator.py b/src/cbfkit/utils/animators/animator.py index 33001b91..36e44e90 100644 --- a/src/cbfkit/utils/animators/animator.py +++ b/src/cbfkit/utils/animators/animator.py @@ -6,11 +6,12 @@ from .config import AnimationConfig from .deps import _require_manim, _require_matplotlib, _require_plotly +from .manim_backend import _ManimMixin from .matplotlib_backend import _MatplotlibMixin from .plotly_backend import _PlotlyMixin -class CBFAnimator(_MatplotlibMixin, _PlotlyMixin): +class CBFAnimator(_MatplotlibMixin, _PlotlyMixin, _ManimMixin): """Declarative 2D trajectory animator for CBFKit simulations. Build an animation by chaining ``add_*`` calls, then :meth:`save` or @@ -31,7 +32,9 @@ class CBFAnimator(_MatplotlibMixin, _PlotlyMixin): aspect : str or None Axis aspect ratio (e.g. ``"equal"``). *None* keeps matplotlib default. backend : str - ``"plotly"`` (default) or ``"matplotlib"``. + ``"plotly"`` (default), ``"matplotlib"``, or ``"manim"`` / + ``"manim-"`` with quality one of ``low``, ``medium``, + ``high``, ``production`` (requires ``pip install cbfkit[manim]``). config : AnimationConfig, optional Animation parameters. Uses :data:`DEFAULT_CONFIG` when *None*. """ @@ -54,15 +57,16 @@ def __init__( ) self._backend = backend + self._manim_quality: Optional[str] = None if backend == "matplotlib": _require_matplotlib() elif backend.startswith("manim"): - # The 2D Manim backend is unimplemented whether or not Manim is - # installed, so surface that before requiring the optional dependency. - raise NotImplementedError( - "Manim 2D backend not yet implemented. " - "Use visualize_3d_multi_robot(backend='manim') for 3D scenes." - ) + # Validate the quality suffix before requiring the optional + # dependency so a typo'd backend string fails loudly either way. + from cbfkit.utils.visualization import _parse_manim_backend + + self._manim_quality = _parse_manim_backend(backend) + _require_manim() else: _require_plotly() @@ -396,23 +400,30 @@ def _compute_prediction(self, spec: dict, frame: int): def build(self): """Create the figure and all static / dynamic artists. - Returns ``(fig, ax)`` for the matplotlib backend, or the Plotly - ``Figure`` for the plotly backend. + Returns ``(fig, ax)`` for the matplotlib backend, the Plotly + ``Figure`` for the plotly backend, or the configured + :class:`~cbfkit.utils.animators.manim_backend.CBFAnimator2DScene` + class for the manim backend. """ if self._backend == "plotly": return self._build_plotly() + if self._backend.startswith("manim"): + return self._build_manim() return self._build_matplotlib() def animate(self): """Build (if needed) and create the animation object. - Returns a matplotlib ``FuncAnimation`` or a Plotly ``Figure`` - (which already contains the animation frames). + Returns a matplotlib ``FuncAnimation``, a Plotly ``Figure`` + (which already contains the animation frames), or the configured + Scene class for the manim backend (rendering happens in :meth:`save`). """ if self._backend == "plotly": if self._fig is None: self._build_plotly() return self._fig + if self._backend.startswith("manim"): + return self._build_manim() return self._animate_matplotlib() def save(self, path: str, config: Optional[AnimationConfig] = None) -> str: @@ -420,11 +431,15 @@ def save(self, path: str, config: Optional[AnimationConfig] = None) -> str: * matplotlib: saves MP4 (ffmpeg) or GIF (pillow fallback). * plotly: saves an interactive ``.html`` file. + * manim: renders MP4 (or GIF if *path* ends in ``.gif``) at the + quality encoded in the backend string (``manim-``). Returns the absolute path of the saved file. """ if self._backend == "plotly": return self._save_plotly(path) + if self._backend.startswith("manim"): + return self._save_manim(path) return self._save_matplotlib(path, config) def show(self): @@ -432,9 +447,12 @@ def show(self): * matplotlib: opens a matplotlib window. * plotly: opens the default web browser. + * manim: renders and opens the video in the default player. """ if self._backend == "plotly": return self._show_plotly() + if self._backend.startswith("manim"): + return self._show_manim() return self._show_matplotlib() # -- properties --------------------------------------------------------- diff --git a/src/cbfkit/utils/animators/manim_backend.py b/src/cbfkit/utils/animators/manim_backend.py new file mode 100644 index 00000000..c3e57727 --- /dev/null +++ b/src/cbfkit/utils/animators/manim_backend.py @@ -0,0 +1,467 @@ +"""Manim 2D backend mixin for CBFAnimator. + +Renders the declarative element API (goals, obstacles, trajectories, agents, +predictions, time overlay) as a high-quality Manim video. Mirrors the +architecture of :mod:`cbfkit.utils.visualizations.manim_3d_multi_robot`: +data is injected into a Scene subclass via class attributes, then rendered. + +Usage (standalone test with synthetic data): + manim -pql src/cbfkit/utils/animators/manim_backend.py CBFAnimator2DScene + +Usage (from library): + from cbfkit.utils.animator import CBFAnimator + CBFAnimator(states, backend="manim-medium").add_agent(0, 1).save("out.mp4") +""" + +from __future__ import annotations + +import os +import shutil + +import numpy as np + +from .deps import _require_manim + +try: + from manim import ( + UL, + UP, + Circle, + DashedLine, + DashedVMobject, + Dot, + Ellipse, + Line, + Scene, + Text, + TracedPath, # noqa: F401 (re-exported for parity with the 3D module) + ValueTracker, + VGroup, + rate_functions, + tempconfig, + ) + + _MANIM_AVAILABLE = True +except ImportError: + _MANIM_AVAILABLE = False + # Fallback base class so this module still imports when Manim is absent; + # actual rendering is gated by _require_manim() in _ManimMixin methods. + Scene = object # type: ignore[assignment,misc] + + +# --------------------------------------------------------------------------- +# Colour handling: translate matplotlib-style colour strings to hex +# --------------------------------------------------------------------------- +_TAB10_HEX = { + "tab:blue": "#1f77b4", + "tab:orange": "#ff7f0e", + "tab:green": "#2ca02c", + "tab:red": "#d62728", + "tab:purple": "#9467bd", + "tab:brown": "#8c564b", + "tab:pink": "#e377c2", + "tab:gray": "#7f7f7f", + "tab:olive": "#bcbd22", + "tab:cyan": "#17becf", +} + +_NAMED_HEX = { + "r": "#d62728", + "red": "#d62728", + "g": "#2ca02c", + "green": "#2ca02c", + "b": "#1f77b4", + "blue": "#1f77b4", + "k": "#000000", + "black": "#000000", + "w": "#ffffff", + "white": "#ffffff", + "y": "#bcbd22", + "yellow": "#bcbd22", + "c": "#17becf", + "cyan": "#17becf", + "m": "#e377c2", + "magenta": "#e377c2", + "orange": "#ff7f0e", + "purple": "#9467bd", + "brown": "#8c564b", + "pink": "#e377c2", + "gray": "#7f7f7f", + "grey": "#7f7f7f", +} + + +def _to_hex(color: str) -> str: + """Best-effort conversion of a matplotlib-style colour string to hex.""" + if not isinstance(color, str) or color.startswith("#"): + return color + if color in _TAB10_HEX: + return _TAB10_HEX[color] + return _NAMED_HEX.get(color.lower(), color) + + +# --------------------------------------------------------------------------- +# Core Manim Scene +# --------------------------------------------------------------------------- +class CBFAnimator2DScene(Scene): + """Render a 2D CBFAnimator element set as a Manim animation. + + Pass data via class attributes before calling ``scene.render()``, or let + the scene generate synthetic demo data when run standalone. + """ + + # --- class-level defaults (overridden by _ManimMixin._build_manim) ----- + states: np.ndarray | None = None + dt: float = 0.1 + x_lim: tuple[float, float] = (-4, 4) + y_lim: tuple[float, float] = (-4, 4) + title: str = "System Behavior" + aspect: str | None = None + goals: list[dict] = [] + obstacles: list[dict] = [] + trajectories: list[dict] = [] + agents: list[dict] = [] + predictions: list[dict] = [] + show_time_overlay: bool = False + # Bound CBFAnimator._compute_prediction (set by _build_manim) + compute_prediction = None + + # ----------------------------------------------------------------------- + def construct(self): + if self.states is None: + self._synthetic_demo() + + states = np.asarray(self.states) + n_frames = len(states) + + # White canvas so matplotlib-convention colours (black obstacles, + # dark trajectories) stay visible, matching the other backends. + self.camera.background_color = "#ffffff" + + axes, unit_x, unit_y = self._build_axes() + self.add(axes) + + if self.title: + self.add(Text(self.title, font_size=28, color="#000000").to_edge(UP, buff=0.2)) + + def c2p(x: float, y: float): + return axes.c2p(float(x), float(y)) + + # --- static goals --------------------------------------------------- + for g in self.goals: + pos, color = g["position"], _to_hex(g["color"]) + dot = Dot(c2p(pos[0], pos[1]), radius=0.06, color=color) + ring = Circle(radius=g["radius"] * unit_x, color=color, stroke_width=2) + ring.move_to(c2p(pos[0], pos[1])) + self.add(dot, DashedVMobject(ring, num_dashes=24)) + + # --- static obstacles ------------------------------------------------- + for obs in self.obstacles: + c, color = obs["center"], _to_hex(obs["color"]) + if obs["ellipse_radii"] is not None: + rx, ry = obs["ellipse_radii"] + mob = Ellipse(width=2 * rx * unit_x, height=2 * ry * unit_y, color=color) + elif obs["radius"] is not None: + mob = Circle(radius=obs["radius"] * unit_x, color=color) + else: + continue + mob.set_fill(color, opacity=obs["alpha"]).set_stroke(color, width=2) + mob.move_to(c2p(c[0], c[1])) + self.add(mob) + + # --- progress tracker ------------------------------------------------- + progress = ValueTracker(0) + + def current_frame() -> int: + return min(int(progress.get_value()), n_frames - 1) + + def _reveal_updater(segments_or_dots): + """Reveal pre-built per-frame mobjects up to the current frame.""" + + def updater(_mob): + frame = current_frame() + for k, seg in enumerate(segments_or_dots): + seg.set_opacity(seg.target_opacity if k < frame else 0.0) + + return updater + + def _build_reveal_group(xs, ys, color, *, style, stroke_width, opacity, dashed=False): + """Pre-build per-frame segments (line) or dots (scatter), initially hidden.""" + group = VGroup() + if style == "scatter": + for k in range(len(xs)): + d = Dot(c2p(xs[k], ys[k]), radius=0.03, color=color) + d.target_opacity = opacity + d.set_opacity(0) + group.add(d) + else: + line_cls = DashedLine if dashed else Line + for k in range(len(xs) - 1): + seg = line_cls( + c2p(xs[k], ys[k]), + c2p(xs[k + 1], ys[k + 1]), + stroke_width=stroke_width, + color=color, + ) + seg.target_opacity = opacity + seg.set_opacity(0) + group.add(seg) + group.add_updater(_reveal_updater(list(group))) + return group + + # --- animated trajectories -------------------------------------------- + for spec in self.trajectories: + data = spec["data"] if spec["data"] is not None else states + group = _build_reveal_group( + data[:, spec["x_idx"]], + data[:, spec["y_idx"]], + _to_hex(spec["color"]), + style=spec["style"], + stroke_width=2 * spec["linewidth"], + opacity=spec["alpha"], + ) + self.add(group) + + # --- agents (body + optional zone + optional trail) -------------------- + for spec in self.agents: + src = spec["data"] if spec["data"] is not None else states + x_idx, y_idx = spec["x_idx"], spec["y_idx"] + color = _to_hex(spec["body_color"]) + start = c2p(src[0, x_idx], src[0, y_idx]) + + body = Circle(radius=spec["body_radius"] * unit_x, color=color) + body.set_fill(color, opacity=spec["body_alpha"]).set_stroke(width=0) + body.move_to(start) + + def _make_follow(_src, _xi, _yi): + def updater(mob): + frame = current_frame() + mob.move_to(c2p(_src[frame, _xi], _src[frame, _yi])) + + return updater + + body.add_updater(_make_follow(src, x_idx, y_idx)) + self.add(body) + + if spec["zone_radius"] is not None: + zone_color = _to_hex(spec["zone_color"]) + zone = Circle(radius=spec["zone_radius"] * unit_x, color=zone_color) + zone.set_fill(zone_color, opacity=spec["zone_alpha"]).set_stroke(width=0) + zone.move_to(start) + zone.add_updater(_make_follow(src, x_idx, y_idx)) + self.add(zone) + + if spec["trail"]: + trail = _build_reveal_group( + src[:, x_idx], + src[:, y_idx], + _to_hex(spec["trail_color"]), + style="line", + stroke_width=3, + opacity=spec["trail_alpha"], + dashed=spec["trail_style"] == "--", + ) + self.add(trail) + + # --- per-frame predictions --------------------------------------------- + for spec in self.predictions: + self.add(self._prediction_group(spec, c2p, current_frame)) + + # --- time overlay -------------------------------------------------------- + if self.show_time_overlay: + time_text = Text("t = 0.0 s", font_size=20, color="#000000").to_corner(UL, buff=0.3) + + def _time_updater(mob): + mob.become( + Text( + f"t = {current_frame() * self.dt:.1f} s", font_size=20, color="#000000" + ).to_corner(UL, buff=0.3) + ) + + time_text.add_updater(_time_updater) + self.add(time_text) + + # --- play: real-time playback, clamped to a sane render length ----------- + playback_seconds = min(30.0, max(2.0, n_frames * self.dt)) + self.play( + progress.animate(run_time=playback_seconds, rate_func=rate_functions.linear).set_value( + n_frames - 1 + ) + ) + self.wait(0.5) + + # ----------------------------------------------------------------------- + def _build_axes(self): + """Create 2D axes fitting the frame; returns (axes, unit_x, unit_y). + + ``unit_x`` / ``unit_y`` convert data units to scene units so element + radii can be drawn to scale. With ``aspect="equal"`` both units match. + """ + from manim import Axes # local import keeps module importable sans manim + + x_span = float(self.x_lim[1] - self.x_lim[0]) + y_span = float(self.y_lim[1] - self.y_lim[0]) + max_w, max_h = 12.0, 6.0 + + if self.aspect == "equal": + unit = min(max_w / x_span, max_h / y_span) + x_len, y_len = x_span * unit, y_span * unit + else: + x_len, y_len = max_w, max_h + + axes = Axes( + x_range=[self.x_lim[0], self.x_lim[1], max(x_span / 8, 1e-6)], + y_range=[self.y_lim[0], self.y_lim[1], max(y_span / 8, 1e-6)], + x_length=x_len, + y_length=y_len, + tips=False, + axis_config={ + "include_ticks": True, + "tick_size": 0.04, + "stroke_width": 1.5, + "color": "#555555", + }, + ) + return axes, x_len / x_span, y_len / y_span + + # ----------------------------------------------------------------------- + def _prediction_group(self, spec: dict, c2p, current_frame): + """A VGroup redrawn every frame from the shared prediction computer.""" + color = _to_hex(spec["color"]) + group = VGroup() + + def updater(mob): + frame = current_frame() + px, py = self.compute_prediction(spec, frame) + segs = VGroup() + n = len(px) - 1 + for k in range(n): + seg = Line( + c2p(px[k], py[k]), + c2p(px[k + 1], py[k + 1]), + stroke_width=2 * spec["linewidth"], + color=color, + ) + alpha = spec["alpha"] * (1.0 - k / n) if spec["fade"] else spec["alpha"] + seg.set_opacity(alpha) + segs.add(seg) + mob.become(segs) if n >= 1 else mob.become(VGroup()) + + group.add_updater(updater) + return group + + # ----------------------------------------------------------------------- + def _synthetic_demo(self): + """Simple goal-reaching demo so the scene renders standalone.""" + n = 100 + t = np.linspace(0, 1, n) + states = np.stack([-3 + 6 * t, 1.5 * np.sin(2 * np.pi * t) * (1 - t)], axis=1) + type(self).states = states + type(self).goals = [{"position": (3.0, 0.0), "radius": 0.3, "color": "g", "label": "Goal"}] + type(self).obstacles = [ + {"center": (0.0, 0.5), "radius": 0.5, "ellipse_radii": None, "color": "k", "alpha": 0.3} + ] + type(self).agents = [ + { + "x_idx": 0, + "y_idx": 1, + "data": None, + "body_radius": 0.2, + "body_color": "blue", + "body_alpha": 0.8, + "zone_radius": None, + "zone_color": "blue", + "zone_alpha": 0.15, + "trail": True, + "trail_color": "blue", + "trail_alpha": 0.5, + "trail_style": "-", + "label": "Agent", + "zorder": 5, + } + ] + type(self).show_time_overlay = True + + +# --------------------------------------------------------------------------- +# Mixin for CBFAnimator +# --------------------------------------------------------------------------- +class _ManimMixin: + """Manim-specific build / save / show methods. + + Mixed into :class:`~cbfkit.utils.animators.animator.CBFAnimator`. + Expects the host class to provide the element descriptor lists, limits, + ``_dt``, ``_title``, ``_aspect``, ``_manim_quality``, and + ``_compute_prediction``. + """ + + def _build_manim(self): + """Inject animator state into :class:`CBFAnimator2DScene`; return it.""" + _require_manim() + CBFAnimator2DScene.states = np.asarray(self._states) + CBFAnimator2DScene.dt = self._dt + CBFAnimator2DScene.x_lim = self._x_lim + CBFAnimator2DScene.y_lim = self._y_lim + CBFAnimator2DScene.title = self._title + CBFAnimator2DScene.aspect = self._aspect + CBFAnimator2DScene.goals = self._goals + CBFAnimator2DScene.obstacles = self._obstacles + CBFAnimator2DScene.trajectories = self._trajectories + CBFAnimator2DScene.agents = self._agents + CBFAnimator2DScene.predictions = self._predictions + CBFAnimator2DScene.show_time_overlay = self._show_time + CBFAnimator2DScene.compute_prediction = self._compute_prediction + return CBFAnimator2DScene + + def _save_manim(self, path: str) -> str: + """Render to *path* (``.mp4``, or ``.gif`` via Manim's gif format).""" + import tempfile + + scene_cls = self._build_manim() + + overrides = { + "quality": self._manim_quality, + # Caching keys partial movie files by animation hash; with a fresh + # media dir per render, stale in-process cache entries would point + # at deleted files and crash a second render, so disable it. + "disable_caching": True, + } + if path.lower().endswith(".gif"): + overrides["format"] = "gif" + # Render intermediates in a temp dir so the caller's cwd stays clean; + # the finished video is copied to *path* below. tempconfig scopes the + # global Manim config so repeated renders in one process stay isolated. + with tempfile.TemporaryDirectory(prefix="cbfkit_manim_") as tmp_media: + overrides["media_dir"] = tmp_media + with tempconfig(overrides): + scene = scene_cls() + scene.render() + # For MP4 the file writer reports the exact output path; for GIF + # it still reports the .mp4 name, so fall back to globbing the + # media dir for the rendered file with the requested extension. + rendered = str(scene.renderer.file_writer.movie_file_path) + if not os.path.exists(rendered): + import glob + + ext = os.path.splitext(path)[1] or ".mp4" + candidates = glob.glob( + os.path.join(tmp_media, "videos", "**", f"*{ext}"), recursive=True + ) + if not candidates: + raise FileNotFoundError( + f"Manim did not produce a {ext} file under {tmp_media!r}." + ) + rendered = candidates[0] + + out_dir = os.path.dirname(path) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + shutil.copy2(rendered, path) + return os.path.abspath(path) + + def _show_manim(self): + """Render and open the result in the default player (Manim preview).""" + scene_cls = self._build_manim() + with tempconfig({"quality": self._manim_quality, "preview": True, "disable_caching": True}): + scene = scene_cls() + scene.render() diff --git a/tests/test_utils/test_animator.py b/tests/test_utils/test_animator.py index 994851db..031abe83 100644 --- a/tests/test_utils/test_animator.py +++ b/tests/test_utils/test_animator.py @@ -78,7 +78,9 @@ def test_build_creates_figure(self, simple_states): plt.close(fig) def test_build_sets_limits_and_labels(self, simple_states): - a = CBFAnimator(simple_states, x_lim=(-5, 5), y_lim=(-3, 3), title="Test", backend="matplotlib") + a = CBFAnimator( + simple_states, x_lim=(-5, 5), y_lim=(-3, 3), title="Test", backend="matplotlib" + ) fig, ax = a.build() assert ax.get_xlim() == (-5, 5) assert ax.get_ylim() == (-3, 3) @@ -259,7 +261,9 @@ def test_save_creates_file(self, simple_states, tmp_path): def test_save_gif_fallback(self, simple_states, tmp_path): """Force GIF fallback by using a config and verifying some file is created.""" - a = CBFAnimator(simple_states, dt=0.1, config=AnimationConfig(fps=5, dpi=50), backend="matplotlib") + a = CBFAnimator( + simple_states, dt=0.1, config=AnimationConfig(fps=5, dpi=50), backend="matplotlib" + ) a.add_trajectory(x_idx=0, y_idx=1) path = str(tmp_path / "test_anim.mp4") result = a.save(path) @@ -423,7 +427,86 @@ def test_require_manim_raises_when_missing(self, monkeypatch): class TestManimBackend: - def test_manim_backend_accepted_but_not_implemented(self, simple_states): - """backend='manim' is a valid value but raises NotImplementedError for 2D.""" - with pytest.raises(NotImplementedError, match="Manim 2D backend not yet implemented"): + """2D Manim backend: construction, dispatch, and guards. + + Dependency-free tests monkeypatch ``_HAS_MANIM`` / the render path so + they run in CI (where the ``manim`` extra is not installed); the real + render smoke test skips without manim. + """ + + def test_manim_backend_raises_import_error_when_missing(self, simple_states, monkeypatch): + from cbfkit.utils.animators import deps + + monkeypatch.setattr(deps, "_HAS_MANIM", False) + with pytest.raises(ImportError, match=r"cbfkit\[manim\]"): CBFAnimator(simple_states, backend="manim") + + def test_manim_invalid_quality_raises_value_error(self, simple_states): + # Quality is validated before the optional dependency is required, + # so this fails loudly whether or not manim is installed. + with pytest.raises(ValueError, match="Unknown Manim backend"): + CBFAnimator(simple_states, backend="manim-ultra") + + @pytest.mark.parametrize( + "backend,quality", + [ + ("manim", "low_quality"), + ("manim-low", "low_quality"), + ("manim-medium", "medium_quality"), + ("manim-high", "high_quality"), + ("manim-production", "production_quality"), + ], + ) + def test_manim_quality_parsing(self, simple_states, backend, quality, monkeypatch): + from cbfkit.utils.animators import deps + + monkeypatch.setattr(deps, "_HAS_MANIM", True) + a = CBFAnimator(simple_states, backend=backend) + assert a._manim_quality == quality + + def test_manim_save_dispatches_to_backend(self, simple_states, monkeypatch, tmp_path): + from cbfkit.utils.animators import deps + + monkeypatch.setattr(deps, "_HAS_MANIM", True) + a = CBFAnimator(simple_states, backend="manim-medium") + + recorded = {} + + def fake_save(path): + recorded["path"] = path + recorded["quality"] = a._manim_quality + return path + + monkeypatch.setattr(a, "_save_manim", fake_save) + out = a.save(str(tmp_path / "anim.mp4")) + assert recorded["path"] == out + assert recorded["quality"] == "medium_quality" + + def test_manim_real_render_smoke(self, simple_states, tmp_path): + from cbfkit.utils.animators.deps import _HAS_MANIM + + if not _HAS_MANIM: + pytest.skip("manim not installed") + + a = CBFAnimator(simple_states[:10], dt=0.1, backend="manim", title="Test", aspect="equal") + a.add_goal((1.0, 0.0), radius=0.2, color="g") + a.add_obstacle((0.0, 0.5), radius=0.3, color="k") + a.add_agent(x_idx=0, y_idx=1, body_radius=0.1, body_color="blue") + a.show_time() + out = a.save(str(tmp_path / "manim_smoke.mp4")) + assert os.path.exists(out) + assert os.path.getsize(out) > 0 + + def test_manim_repeated_renders_same_process(self, simple_states, tmp_path): + """Regression: a second render in the same process must not reuse + cached partial-movie files from the first render's deleted temp dir.""" + from cbfkit.utils.animators.deps import _HAS_MANIM + + if not _HAS_MANIM: + pytest.skip("manim not installed") + + for i in range(2): + a = CBFAnimator(simple_states[:8], dt=0.1, backend="manim") + a.add_agent(x_idx=0, y_idx=1, body_radius=0.1) + out = a.save(str(tmp_path / f"repeat_{i}.mp4")) + assert os.path.getsize(out) > 0