"""
Animation tools for MechanicsDSL
Provides animation capabilities for mechanical systems and fluid simulations.
"""
from typing import List, Optional, Tuple
import numpy as np
try:
import matplotlib.animation as animation
import matplotlib.pyplot as plt
HAS_MATPLOTLIB = True
except ImportError:
HAS_MATPLOTLIB = False
from ..utils import config, logger
[docs]
class Animator:
"""
Animation handler for mechanical system simulations.
Supports:
- Pendulum animations (single, double, multi-body)
- Particle trajectory animations
- Fluid particle visualizations
"""
def __init__(self, trail_length: int = None, fps: int = None):
if not HAS_MATPLOTLIB:
raise ImportError("matplotlib required for animations")
self.trail_length = trail_length or config.trail_length
self.fps = fps or config.animation_fps
self.fig: Optional[plt.Figure] = None
self.ax: Optional[plt.Axes] = None
self.animation: Optional[animation.FuncAnimation] = None
[docs]
def animate_pendulum(
self, solution: dict, length: float = 1.0, title: str = "Pendulum"
) -> animation.FuncAnimation:
"""
Create pendulum animation from simulation solution.
Args:
solution: Simulation result with 't' and 'y' arrays
length: Pendulum length for visualization
title: Animation title
Returns:
matplotlib FuncAnimation object
"""
t = solution["t"]
y = solution["y"]
# Assume first coordinate is angle
theta = y[0]
self.setup_figure(
xlim=(-1.5 * length, 1.5 * length), ylim=(-1.5 * length, 0.5 * length), title=title
)
# Create elements
(line,) = self.ax.plot([], [], "o-", lw=2, color="blue")
(trail,) = self.ax.plot([], [], "-", lw=1, alpha=0.5, color="red")
time_text = self.ax.text(0.02, 0.95, "", transform=self.ax.transAxes)
# Trail buffer
trail_x = []
trail_y = []
def init():
line.set_data([], [])
trail.set_data([], [])
time_text.set_text("")
return line, trail, time_text
def update(frame):
x = length * np.sin(theta[frame])
y_pos = -length * np.cos(theta[frame])
line.set_data([0, x], [0, y_pos])
# Update trail
trail_x.append(x)
trail_y.append(y_pos)
if len(trail_x) > self.trail_length:
trail_x.pop(0)
trail_y.pop(0)
trail.set_data(trail_x, trail_y)
time_text.set_text(f"t = {t[frame]:.2f}s")
return line, trail, time_text
interval = 1000 / self.fps
self.animation = animation.FuncAnimation(
self.fig, update, frames=len(t), init_func=init, blit=True, interval=interval
)
return self.animation
[docs]
def animate(self, solution: dict, parameters: dict = None, system_name: str = "system"):
"""
Generic animation dispatcher (backward compatibility with MechanicsVisualizer).
Args:
solution: Simulation result dictionary
parameters: Physical parameters (optional)
system_name: Name of the system
Returns:
matplotlib FuncAnimation object
"""
if not solution or not solution.get("success"):
return None
coords = solution.get("coordinates", [])
name = (system_name or "").lower()
length = parameters.get("l", 1.0) if parameters else 1.0
# Dispatch to appropriate animation
if "pendulum" in name or any("theta" in str(c) for c in coords):
return self.animate_pendulum(solution, length=length, title=system_name)
else:
return self.animate_pendulum(solution, length=length, title=system_name)
[docs]
def animate_particles(
self, positions: List[Tuple[np.ndarray, np.ndarray]], title: str = "Particles"
) -> animation.FuncAnimation:
"""
Animate particle positions over time.
Args:
positions: List of (x_array, y_array) for each frame
title: Animation title
Returns:
matplotlib FuncAnimation object
"""
# Determine bounds
all_x = np.concatenate([p[0] for p in positions])
all_y = np.concatenate([p[1] for p in positions])
margin = 0.1 * max(all_x.max() - all_x.min(), all_y.max() - all_y.min())
self.setup_figure(
xlim=(all_x.min() - margin, all_x.max() + margin),
ylim=(all_y.min() - margin, all_y.max() + margin),
title=title,
)
scatter = self.ax.scatter([], [], c="blue", s=10)
def update(frame):
x, y = positions[frame]
scatter.set_offsets(np.c_[x, y])
return (scatter,)
interval = 1000 / self.fps
self.animation = animation.FuncAnimation(
self.fig, update, frames=len(positions), blit=True, interval=interval
)
return self.animation
[docs]
def save(self, filename: str, dpi: int = 100) -> bool:
"""
Save animation to file.
Args:
filename: Output filename (mp4, gif, etc.)
dpi: Resolution
Returns:
True if successful
"""
if self.animation is None:
logger.error("No animation to save")
return False
try:
writer = None
if filename.endswith(".gif"):
writer = "pillow"
elif filename.endswith(".mp4"):
writer = "ffmpeg"
self.animation.save(filename, writer=writer, dpi=dpi)
logger.info(f"Animation saved to {filename}")
return True
except Exception as e:
logger.error(f"Failed to save animation: {e}")
return False