""" Generic animation utilities for WorldSmithAI. This module turns rendered world states into GIF or MP4 animations. It is built to work from a root-level Hugging Face Spaces ``app.py`` and does not assume any ``app/`` package or special folder layout. The animation layer is domain-agnostic. It does not know about farms, civilizations, research ecosystems, transport networks, power grids, fantasy worlds, or social networks. It advances a runtime world, asks ``visualization.renderer`` to render each state, and writes the resulting frames to an animation file. Gradio usage: from visualization.animation import animate_world_to_gif_path def simulate(prompt: str) -> str: world = build_world_from_prompt(prompt) return animate_world_to_gif_path(world, frame_count=60) MP4 usage: from visualization.animation import animate_world_to_mp4_path def simulate_video(prompt: str) -> str: world = build_world_from_prompt(prompt) return animate_world_to_mp4_path(world, frame_count=90) Future extensibility: - Add event overlays and behavior traces. - Add side-by-side charts in animation frames. - Add camera tracking for selected agents. - Add graph/network animation mode. - Add streaming frame generation for long simulations. - Add optional WebM writer for browser-native video output. """ from __future__ import annotations import copy import inspect import logging import tempfile from collections.abc import Callable, Mapping, MutableSequence, Sequence from dataclasses import dataclass, field from enum import Enum from numbers import Real from pathlib import Path from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar import matplotlib matplotlib.use("Agg", force=False) import matplotlib.pyplot as plt import numpy as np from matplotlib.animation import FFMpegWriter, FuncAnimation, PillowWriter, writers from matplotlib.figure import Figure from visualization.renderer import ( RenderSnapshot, RendererConfig, WorldRenderer, close_figure, figure_to_rgb_array, ) if TYPE_CHECKING: from core.world import World logger = logging.getLogger(__name__) _EPSILON = 1.0e-12 class AnimationFormat(str, Enum): """Supported animation output formats.""" GIF = "gif" MP4 = "mp4" class AnimationAdvanceMode(str, Enum): """Strategies for advancing a world between frames.""" WORLD_STEP = "world_step" SCHEDULER_STEP = "scheduler_step" AUTO = "auto" NONE = "none" @dataclass(frozen=True) class AnimationFrame: """One captured animation frame. The image array is intentionally omitted from ``to_dict`` to keep metadata small and JSON-friendly. """ index: int step: int | None image: np.ndarray = field(repr=False, compare=False) snapshot: RenderSnapshot | None = None metadata: Mapping[str, Any] = field(default_factory=dict) @property def shape(self) -> tuple[int, ...]: """Return the underlying image array shape.""" return tuple(int(value) for value in self.image.shape) def to_dict(self) -> dict[str, Any]: """Return a JSON-friendly frame metadata dictionary.""" return { "index": self.index, "step": self.step, "shape": list(self.shape), "snapshot": None if self.snapshot is None else self.snapshot.to_dict(), "metadata": copy.deepcopy(dict(self.metadata)), } @dataclass(frozen=True) class AnimationResult: """Result returned after writing an animation.""" path: str format: AnimationFormat frame_count: int fps: float frames: tuple[AnimationFrame, ...] = field(default_factory=tuple, repr=False, compare=False) metadata: Mapping[str, Any] = field(default_factory=dict) def to_dict(self, *, include_frame_metadata: bool = True) -> dict[str, Any]: """Return a JSON-friendly animation result dictionary.""" return { "path": self.path, "format": self.format.value, "frame_count": self.frame_count, "fps": self.fps, "frames": [frame.to_dict() for frame in self.frames] if include_frame_metadata else [], "metadata": copy.deepcopy(dict(self.metadata)), } @dataclass class AnimationConfig: """Configuration for ``WorldAnimator``. Defaults favor GIF output because GIFs are easy to display in a root-level Gradio ``app.py`` and do not require ffmpeg. """ frame_count: int = 60 steps_per_frame: int = 1 include_initial_frame: bool = True format: AnimationFormat | str = AnimationFormat.GIF fps: float | None = None interval_ms: int = 150 repeat: bool = True advance_mode: AnimationAdvanceMode | str = AnimationAdvanceMode.AUTO ensure_step_count_progress: bool = True renderer_config: RendererConfig = field(default_factory=RendererConfig) output_dir: str | Path | None = None filename_prefix: str = "worldsmithai_animation_" dpi: int | None = None close_frame_figures: bool = True close_animation_figure: bool = True title: str | None = None show_animation_title: bool = False metadata: Mapping[str, Any] = field(default_factory=dict) def resolved_fps(self) -> float: """Return effective frames per second.""" if self.fps is not None: return max(_EPSILON, float(self.fps)) interval = max(1, int(self.interval_ms)) return 1000.0 / float(interval) def resolved_format(self) -> AnimationFormat: """Return effective animation format.""" return normalize_animation_format(self.format) class AnimationError(RuntimeError): """Base exception raised by animation utilities.""" class AnimationWriterError(AnimationError): """Raised when animation frames cannot be written to the requested format.""" class WorldAdvanceError(AnimationError): """Raised when the world cannot be advanced between frames.""" @dataclass class WorldAnimator: """Capture and write animations for generic WorldSmithAI worlds. The animator mutates the supplied world by advancing it between frames. This is usually what a simulation animation should do. Callers that need to preserve the original world should pass in a copied world object. """ config: AnimationConfig = field(default_factory=AnimationConfig) name: ClassVar[str] = "world_animator" def capture_frames( self, world: World, *, frame_count: int | None = None, progress_callback: Callable[[int, int], None] | None = None, ) -> tuple[AnimationFrame, ...]: """Capture rendered frames while advancing the world. Args: world: Runtime world object. frame_count: Optional override for configured frame count. progress_callback: Optional callback receiving ``(frame_index, total_frames)``. This is useful for app-level progress reporting without importing Gradio here. Returns: Tuple of captured animation frames. """ total_frames = max(1, int(frame_count if frame_count is not None else self.config.frame_count)) renderer = WorldRenderer(config=self.config.renderer_config) frames: list[AnimationFrame] = [] for frame_index in range(total_frames): if frame_index == 0 and self.config.include_initial_frame: pass else: self.advance_world(world, steps=max(0, int(self.config.steps_per_frame))) frame = self._capture_single_frame( world, renderer=renderer, index=frame_index, ) frames.append(frame) if progress_callback is not None: progress_callback(frame_index + 1, total_frames) logger.debug( "Captured %s animation frame(s) for world at step %s", len(frames), _world_step(world), ) return tuple(frames) def animate( self, world: World, *, output_path: str | Path | None = None, format: AnimationFormat | str | None = None, progress_callback: Callable[[int, int], None] | None = None, ) -> AnimationResult: """Capture frames, write an animation file, and return the result. Args: world: Runtime world object. output_path: Optional explicit output path. format: Optional output format override. progress_callback: Optional frame-capture progress callback. Returns: ``AnimationResult`` with output path and frame metadata. """ frames = self.capture_frames(world, progress_callback=progress_callback) return self.write_frames( frames, output_path=output_path, format=format, ) def write_frames( self, frames: Sequence[AnimationFrame | np.ndarray], *, output_path: str | Path | None = None, format: AnimationFormat | str | None = None, ) -> AnimationResult: """Write existing frames to an animation file. Args: frames: Sequence of ``AnimationFrame`` objects or RGB/RGBA arrays. output_path: Optional output file path. format: Optional output format override. Returns: ``AnimationResult``. """ normalized_format = ( normalize_animation_format(format) if format is not None else self._format_from_path_or_config(output_path) ) normalized_frames = self._normalize_frames(frames) if not normalized_frames: raise AnimationWriterError("Cannot write animation with zero frames") path = ( str(output_path) if output_path is not None else temporary_animation_path( normalized_format, output_dir=self.config.output_dir, prefix=self.config.filename_prefix, ) ) write_animation_arrays( [frame.image for frame in normalized_frames], output_path=path, format=normalized_format, fps=self.config.resolved_fps(), interval_ms=int(self.config.interval_ms), repeat=bool(self.config.repeat), dpi=int(self.config.dpi or self.config.renderer_config.dpi), title=self.config.title, show_title=bool(self.config.show_animation_title), close_figure=bool(self.config.close_animation_figure), ) return AnimationResult( path=path, format=normalized_format, frame_count=len(normalized_frames), fps=self.config.resolved_fps(), frames=tuple(normalized_frames), metadata={ **copy.deepcopy(dict(self.config.metadata)), "steps_per_frame": int(self.config.steps_per_frame), "include_initial_frame": bool(self.config.include_initial_frame), "advance_mode": normalize_advance_mode(self.config.advance_mode).value, }, ) def animate_to_gif( self, world: World, *, output_path: str | Path | None = None, progress_callback: Callable[[int, int], None] | None = None, ) -> AnimationResult: """Create a GIF animation for a world.""" return self.animate( world, output_path=output_path, format=AnimationFormat.GIF, progress_callback=progress_callback, ) def animate_to_mp4( self, world: World, *, output_path: str | Path | None = None, progress_callback: Callable[[int, int], None] | None = None, ) -> AnimationResult: """Create an MP4 animation for a world. MP4 output requires an available ffmpeg writer. """ return self.animate( world, output_path=output_path, format=AnimationFormat.MP4, progress_callback=progress_callback, ) def advance_world(self, world: World, *, steps: int = 1) -> None: """Advance the runtime world by a number of simulation steps.""" for _ in range(max(0, int(steps))): self._advance_world_once(world) def _advance_world_once(self, world: World) -> None: """Advance the world by one step using the configured strategy.""" mode = normalize_advance_mode(self.config.advance_mode) before_step = _world_step(world) advanced = False if mode in {AnimationAdvanceMode.AUTO, AnimationAdvanceMode.WORLD_STEP}: step_method = getattr(world, "step", None) if callable(step_method): _call_step_method(step_method, world) advanced = True if not advanced and mode in {AnimationAdvanceMode.AUTO, AnimationAdvanceMode.SCHEDULER_STEP}: scheduler = getattr(world, "scheduler", None) if scheduler is not None: scheduler_step = getattr(scheduler, "step", None) if callable(scheduler_step): _call_step_method(scheduler_step, world) advanced = True else: advanced = _call_scheduler_phases(scheduler, world) if not advanced and mode is AnimationAdvanceMode.NONE: return if not advanced and mode is not AnimationAdvanceMode.AUTO: raise WorldAdvanceError(f"Could not advance world using mode {mode.value!r}") after_step = _world_step(world) if self.config.ensure_step_count_progress and before_step == after_step: _increment_world_step(world) def _capture_single_frame( self, world: World, *, renderer: WorldRenderer, index: int, ) -> AnimationFrame: """Capture a single rendered frame and metadata snapshot.""" render_result = renderer.render_result(world) image = figure_to_rgb_array(render_result.figure) if self.config.close_frame_figures: close_figure(render_result.figure) return AnimationFrame( index=index, step=_world_step(world), image=image, snapshot=render_result.snapshot, metadata={ "world_step": _world_step(world), "object_count": render_result.snapshot.object_count, }, ) def _normalize_frames( self, frames: Sequence[AnimationFrame | np.ndarray], ) -> tuple[AnimationFrame, ...]: """Normalize frame inputs to ``AnimationFrame`` objects.""" normalized: list[AnimationFrame] = [] for index, frame in enumerate(frames): if isinstance(frame, AnimationFrame): normalized.append(frame) continue array = normalize_image_array(frame) normalized.append( AnimationFrame( index=index, step=None, image=array, snapshot=None, metadata={"source": "array"}, ) ) validate_frame_shapes([frame.image for frame in normalized]) return tuple(normalized) def _format_from_path_or_config(self, output_path: str | Path | None) -> AnimationFormat: """Infer animation format from output path or config.""" if output_path is None: return self.config.resolved_format() suffix = Path(output_path).suffix.lower().lstrip(".") if suffix in {"gif", "mp4"}: return normalize_animation_format(suffix) return self.config.resolved_format() def animate_world( world: World, *, output_path: str | Path | None = None, frame_count: int = 60, steps_per_frame: int = 1, format: AnimationFormat | str = AnimationFormat.GIF, renderer_config: RendererConfig | None = None, output_dir: str | Path | None = None, ) -> AnimationResult: """Convenience function that animates a world to GIF or MP4.""" config = AnimationConfig( frame_count=frame_count, steps_per_frame=steps_per_frame, format=format, renderer_config=renderer_config or RendererConfig(), output_dir=output_dir, ) return WorldAnimator(config=config).animate(world, output_path=output_path, format=format) def animate_world_to_gif_path( world: World, *, output_path: str | Path | None = None, frame_count: int = 60, steps_per_frame: int = 1, renderer_config: RendererConfig | None = None, output_dir: str | Path | None = None, ) -> str: """Animate a world to a GIF path. This is the safest default for a root-level Gradio ``app.py`` because GIF output does not require ffmpeg. """ result = animate_world( world, output_path=output_path, frame_count=frame_count, steps_per_frame=steps_per_frame, format=AnimationFormat.GIF, renderer_config=renderer_config, output_dir=output_dir, ) return result.path def animate_world_to_mp4_path( world: World, *, output_path: str | Path | None = None, frame_count: int = 60, steps_per_frame: int = 1, renderer_config: RendererConfig | None = None, output_dir: str | Path | None = None, ) -> str: """Animate a world to an MP4 path. MP4 output requires ffmpeg to be available in the runtime environment. """ result = animate_world( world, output_path=output_path, frame_count=frame_count, steps_per_frame=steps_per_frame, format=AnimationFormat.MP4, renderer_config=renderer_config, output_dir=output_dir, ) return result.path def capture_world_frames( world: World, *, frame_count: int = 60, steps_per_frame: int = 1, renderer_config: RendererConfig | None = None, ) -> tuple[AnimationFrame, ...]: """Capture animation frames without writing a file.""" config = AnimationConfig( frame_count=frame_count, steps_per_frame=steps_per_frame, renderer_config=renderer_config or RendererConfig(), ) return WorldAnimator(config=config).capture_frames(world) def capture_world_frame_arrays( world: World, *, frame_count: int = 60, steps_per_frame: int = 1, renderer_config: RendererConfig | None = None, ) -> tuple[np.ndarray, ...]: """Capture animation frames as RGB arrays without writing a file.""" frames = capture_world_frames( world, frame_count=frame_count, steps_per_frame=steps_per_frame, renderer_config=renderer_config, ) return tuple(frame.image for frame in frames) def write_animation_arrays( frames: Sequence[np.ndarray], *, output_path: str | Path, format: AnimationFormat | str, fps: float, interval_ms: int = 150, repeat: bool = True, dpi: int = 120, title: str | None = None, show_title: bool = False, close_figure: bool = True, ) -> str: """Write image arrays to a GIF or MP4 animation. Args: frames: Sequence of RGB or RGBA image arrays. output_path: Destination path. format: Animation output format. fps: Frames per second. interval_ms: Display interval in milliseconds. repeat: Whether GIF animations should loop. dpi: Output DPI. title: Optional animation title. show_title: Whether to display the title above the image. close_figure: Whether to close the temporary animation figure. Returns: Output path as a string. """ normalized_format = normalize_animation_format(format) arrays = [normalize_image_array(frame) for frame in frames] validate_frame_shapes(arrays) if not arrays: raise AnimationWriterError("Cannot write animation with zero frames") path = Path(output_path) path.parent.mkdir(parents=True, exist_ok=True) figure, axes = plt.subplots(figsize=_figure_size_from_array(arrays[0], dpi=dpi), dpi=dpi) image_artist = axes.imshow(arrays[0]) axes.axis("off") if show_title and title: axes.set_title(title) def update(frame_index: int) -> tuple[Any, ...]: image_artist.set_data(arrays[frame_index]) return (image_artist,) animation = FuncAnimation( figure, update, frames=len(arrays), interval=int(interval_ms), blit=True, repeat=bool(repeat), ) try: if normalized_format is AnimationFormat.GIF: writer = PillowWriter(fps=max(_EPSILON, float(fps))) animation.save(str(path), writer=writer, dpi=int(dpi)) elif normalized_format is AnimationFormat.MP4: if not writers.is_available("ffmpeg"): raise AnimationWriterError( "MP4 output requires ffmpeg, but Matplotlib could not find an ffmpeg writer. " "Use GIF output or install ffmpeg in the Space environment." ) writer = FFMpegWriter(fps=max(_EPSILON, float(fps))) animation.save(str(path), writer=writer, dpi=int(dpi)) else: raise AnimationWriterError(f"Unsupported animation format: {normalized_format!r}") finally: if close_figure: plt.close(figure) return str(path) def temporary_animation_path( format: AnimationFormat | str, *, output_dir: str | Path | None = None, prefix: str = "worldsmithai_animation_", ) -> str: """Return a temporary animation path with the requested extension.""" normalized_format = normalize_animation_format(format) suffix = f".{normalized_format.value}" directory = None if output_dir is None else str(output_dir) if directory is not None: Path(directory).mkdir(parents=True, exist_ok=True) with tempfile.NamedTemporaryFile( suffix=suffix, prefix=prefix, dir=directory, delete=False, ) as handle: return handle.name def normalize_animation_format(value: AnimationFormat | str | None) -> AnimationFormat: """Normalize an animation format value.""" if value is None: return AnimationFormat.GIF if isinstance(value, AnimationFormat): return value normalized = str(value).lower().lstrip(".") return AnimationFormat(normalized) def normalize_advance_mode(value: AnimationAdvanceMode | str) -> AnimationAdvanceMode: """Normalize an advance mode value.""" if isinstance(value, AnimationAdvanceMode): return value return AnimationAdvanceMode(str(value)) def normalize_image_array(frame: np.ndarray) -> np.ndarray: """Normalize an image frame to uint8 RGB format.""" array = np.asarray(frame) if array.ndim != 3: raise AnimationWriterError( f"Animation frame must have shape (height, width, channels), got {array.shape}" ) if array.shape[2] not in {3, 4}: raise AnimationWriterError( f"Animation frame must have 3 or 4 channels, got shape {array.shape}" ) if array.shape[2] == 4: array = array[:, :, :3] if array.dtype == np.uint8: return array.copy() numeric = array.astype(float) if numeric.max(initial=0.0) <= 1.0 and numeric.min(initial=0.0) >= 0.0: numeric = numeric * 255.0 numeric = np.nan_to_num(numeric, nan=0.0, posinf=255.0, neginf=0.0) numeric = np.clip(numeric, 0.0, 255.0) return numeric.astype(np.uint8) def validate_frame_shapes(frames: Sequence[np.ndarray]) -> None: """Validate that all animation frames share one shape.""" if not frames: raise AnimationWriterError("Animation requires at least one frame") first_shape = tuple(frames[0].shape) for index, frame in enumerate(frames): shape = tuple(frame.shape) if shape != first_shape: raise AnimationWriterError( "All animation frames must have the same shape. " f"Frame 0 has {first_shape}, frame {index} has {shape}." ) def _figure_size_from_array(array: np.ndarray, *, dpi: int) -> tuple[float, float]: """Return Matplotlib figure size in inches for an image array.""" height, width = array.shape[:2] safe_dpi = max(1, int(dpi)) return max(1.0, width / safe_dpi), max(1.0, height / safe_dpi) def _world_step(world: World) -> int | None: """Return current world step if available.""" value = getattr(world, "step_count", None) if isinstance(value, Real) and not isinstance(value, bool): return int(value) return None def _increment_world_step(world: World) -> None: """Increment ``world.step_count`` if possible.""" current_step = _world_step(world) if current_step is None: try: setattr(world, "step_count", 1) except Exception: logger.debug("Could not initialize world.step_count", exc_info=True) return try: setattr(world, "step_count", current_step + 1) except Exception: logger.debug("Could not increment world.step_count", exc_info=True) def _call_step_method(method: Callable[..., Any], world: World) -> Any: """Call a step-like method with zero or one world argument.""" try: signature = inspect.signature(method) except (TypeError, ValueError): return method() required_parameters = [ parameter for parameter in signature.parameters.values() if parameter.default is inspect.Parameter.empty and parameter.kind in { inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY, } ] if not required_parameters: return method() return method(world) def _call_scheduler_phases(scheduler: Any, world: World) -> bool: """Call common scheduler phase methods when no direct step method exists.""" phase_names = ("process_events", "execute_agents", "update_resources") called_any = False for name in phase_names: method = getattr(scheduler, name, None) if not callable(method): continue _call_step_method(method, world) called_any = True return called_any ANIMATION_REGISTRY: Mapping[str, type[WorldAnimator]] = MappingProxyType( { WorldAnimator.name: WorldAnimator, } ) __all__ = [ "ANIMATION_REGISTRY", "AnimationAdvanceMode", "AnimationConfig", "AnimationError", "AnimationFormat", "AnimationFrame", "AnimationResult", "AnimationWriterError", "WorldAdvanceError", "WorldAnimator", "animate_world", "animate_world_to_gif_path", "animate_world_to_mp4_path", "capture_world_frame_arrays", "capture_world_frames", "normalize_animation_format", "normalize_image_array", "temporary_animation_path", "validate_frame_shapes", "write_animation_arrays", ]