Provided modular architecture for animated bouncing ball analysis
- Organized project into src directory with subpackages (analysis, data, visualization, utils) - Added comprehensive README with project overview and structure - Implemented data loading, bounce detection, and visualization modules - Created example scripts and Jupyter notebook for project usage - Added requirements.txt for dependency management - Included output files for different ball types (golf, lacrosse, metal)
This commit is contained in:
287
src/visualization/animation.py
Normal file
287
src/visualization/animation.py
Normal file
@@ -0,0 +1,287 @@
|
||||
"""
|
||||
Animation module for visualizing bouncing ball data.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib.animation as animation
|
||||
from matplotlib.patches import Circle
|
||||
import pandas as pd
|
||||
from IPython.display import HTML
|
||||
|
||||
|
||||
def create_ball_animation(df, ball_radius=0.5, fps=30, duration=None, title="Bouncing Ball Animation"):
|
||||
"""
|
||||
Create an animation of a bouncing ball from position data.
|
||||
|
||||
Parameters:
|
||||
df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
|
||||
ball_radius (float): Radius of the ball in inches
|
||||
fps (int): Frames per second for the animation
|
||||
duration (float): Duration of the animation in seconds (if None, uses the full data)
|
||||
title (str): Title for the animation
|
||||
|
||||
Returns:
|
||||
matplotlib.animation.Animation: The animation object
|
||||
"""
|
||||
# Extract time and position data
|
||||
time_data = df['Time'].values
|
||||
position_data = df['Position'].values
|
||||
|
||||
# Determine animation duration
|
||||
if duration is None:
|
||||
duration = time_data[-1]
|
||||
|
||||
# Calculate the number of frames based on fps and duration
|
||||
num_frames = int(fps * duration)
|
||||
|
||||
# Create figure and axis
|
||||
fig, ax = plt.subplots(figsize=(10, 6))
|
||||
|
||||
# Set axis limits
|
||||
max_height = np.max(position_data) + 2 * ball_radius
|
||||
ax.set_xlim(-5, 5)
|
||||
ax.set_ylim(0, max_height)
|
||||
|
||||
# Set labels and title
|
||||
ax.set_xlabel('X Position (inches)')
|
||||
ax.set_ylabel('Y Position (inches)')
|
||||
ax.set_title(title)
|
||||
|
||||
# Add a horizontal line at y=0 to represent the ground
|
||||
ax.axhline(y=0, color='black', linestyle='-', linewidth=2)
|
||||
|
||||
# Create the ball (circle patch)
|
||||
ball = Circle((0, position_data[0]), ball_radius, color='red', zorder=2)
|
||||
ax.add_patch(ball)
|
||||
|
||||
# Create a line to show the ball's path
|
||||
line, = ax.plot([], [], 'b-', alpha=0.3, linewidth=1)
|
||||
|
||||
# Initialize empty lists for the path
|
||||
path_x = []
|
||||
path_y = []
|
||||
|
||||
def init():
|
||||
"""Initialize the animation"""
|
||||
ball.center = (0, position_data[0])
|
||||
line.set_data([], [])
|
||||
return ball, line
|
||||
|
||||
def animate(i):
|
||||
"""Update the animation for frame i"""
|
||||
# Calculate the current time based on the frame number
|
||||
current_time = i * duration / num_frames
|
||||
|
||||
# Find the closest time index in our data
|
||||
idx = np.abs(time_data - current_time).argmin()
|
||||
|
||||
# Update ball position
|
||||
y_pos = position_data[idx]
|
||||
ball.center = (0, y_pos)
|
||||
|
||||
# Update path
|
||||
path_x.append(0)
|
||||
path_y.append(y_pos)
|
||||
line.set_data(path_x, path_y)
|
||||
|
||||
return ball, line
|
||||
|
||||
# Create the animation
|
||||
anim = animation.FuncAnimation(
|
||||
fig, animate, init_func=init, frames=num_frames,
|
||||
interval=1000/fps, blit=True
|
||||
)
|
||||
|
||||
plt.close() # Prevent the static plot from displaying
|
||||
|
||||
return anim
|
||||
|
||||
|
||||
def create_ball_animation_with_model(df, peak_indices, cor, initial_height,
|
||||
ball_radius=0.5, fps=30, duration=None,
|
||||
title="Bouncing Ball Animation with Model",
|
||||
g=386.09):
|
||||
"""
|
||||
Create an animation of a bouncing ball from position data with an ideal model comparison.
|
||||
|
||||
Parameters:
|
||||
df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
|
||||
peak_indices (numpy.ndarray): Indices of detected bounces
|
||||
cor (float): Coefficient of Restitution
|
||||
initial_height (float): Initial height of the ball in inches
|
||||
ball_radius (float): Radius of the ball in inches
|
||||
fps (int): Frames per second for the animation
|
||||
duration (float): Duration of the animation in seconds (if None, uses the full data)
|
||||
title (str): Title for the animation
|
||||
g (float): Acceleration due to gravity (in inches/s²)
|
||||
|
||||
Returns:
|
||||
matplotlib.animation.Animation: The animation object
|
||||
"""
|
||||
# Extract time and position data
|
||||
time_data = df['Time'].values
|
||||
position_data = df['Position'].values
|
||||
|
||||
# Determine animation duration
|
||||
if duration is None:
|
||||
duration = time_data[-1]
|
||||
|
||||
# Calculate the number of frames based on fps and duration
|
||||
num_frames = int(fps * duration)
|
||||
|
||||
# Create figure and axis
|
||||
fig, ax = plt.subplots(figsize=(10, 6))
|
||||
|
||||
# Set axis limits
|
||||
max_height = np.max(position_data) + 2 * ball_radius
|
||||
ax.set_xlim(-5, 5)
|
||||
ax.set_ylim(0, max_height)
|
||||
|
||||
# Set labels and title
|
||||
ax.set_xlabel('X Position (inches)')
|
||||
ax.set_ylabel('Y Position (inches)')
|
||||
ax.set_title(title)
|
||||
|
||||
# Add a horizontal line at y=0 to represent the ground
|
||||
ax.axhline(y=0, color='black', linestyle='-', linewidth=2)
|
||||
|
||||
# Create the real data ball (circle patch)
|
||||
real_ball = Circle((0, position_data[0]), ball_radius, color='red', zorder=2, label='Actual')
|
||||
ax.add_patch(real_ball)
|
||||
|
||||
# Create the model ball (circle patch)
|
||||
model_ball = Circle((2, initial_height), ball_radius, color='blue', zorder=2, alpha=0.7, label='Model')
|
||||
ax.add_patch(model_ball)
|
||||
|
||||
# Create lines to show the balls' paths
|
||||
real_line, = ax.plot([], [], 'r-', alpha=0.3, linewidth=1)
|
||||
model_line, = ax.plot([], [], 'b-', alpha=0.3, linewidth=1)
|
||||
|
||||
# Add legend
|
||||
ax.legend(handles=[
|
||||
plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='red', markersize=10, label='Actual'),
|
||||
plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='blue', markersize=10, label='Model')
|
||||
])
|
||||
|
||||
# Initialize empty lists for the paths
|
||||
real_path_x = []
|
||||
real_path_y = []
|
||||
model_path_x = []
|
||||
model_path_y = []
|
||||
|
||||
# Generate model data
|
||||
dt = 0.001 # Time step for simulation
|
||||
model_time = []
|
||||
model_position = []
|
||||
|
||||
# Initialize simulation parameters
|
||||
y = initial_height
|
||||
v_y = 0.0 # Starting from rest
|
||||
t = 0.0
|
||||
|
||||
# Run simulation for the duration
|
||||
while t <= duration:
|
||||
model_time.append(t)
|
||||
model_position.append(y)
|
||||
|
||||
# Kinematic update
|
||||
v_y -= g * dt
|
||||
y += v_y * dt
|
||||
t += dt
|
||||
|
||||
# Bounce condition
|
||||
if y <= 0:
|
||||
y = 0 # Ground impact
|
||||
v_y = -v_y * cor # Apply COR
|
||||
|
||||
# Convert to numpy arrays for faster indexing
|
||||
model_time = np.array(model_time)
|
||||
model_position = np.array(model_position)
|
||||
|
||||
def init():
|
||||
"""Initialize the animation"""
|
||||
real_ball.center = (0, position_data[0])
|
||||
model_ball.center = (2, initial_height)
|
||||
real_line.set_data([], [])
|
||||
model_line.set_data([], [])
|
||||
return real_ball, model_ball, real_line, model_line
|
||||
|
||||
def animate(i):
|
||||
"""Update the animation for frame i"""
|
||||
# Calculate the current time based on the frame number
|
||||
current_time = i * duration / num_frames
|
||||
|
||||
# Find the closest time index in our data
|
||||
real_idx = np.abs(time_data - current_time).argmin()
|
||||
model_idx = np.abs(model_time - current_time).argmin()
|
||||
|
||||
# Update ball positions
|
||||
real_y_pos = position_data[real_idx]
|
||||
model_y_pos = model_position[model_idx]
|
||||
|
||||
real_ball.center = (0, real_y_pos)
|
||||
model_ball.center = (2, model_y_pos)
|
||||
|
||||
# Update paths
|
||||
real_path_x.append(0)
|
||||
real_path_y.append(real_y_pos)
|
||||
model_path_x.append(2)
|
||||
model_path_y.append(model_y_pos)
|
||||
|
||||
real_line.set_data(real_path_x, real_path_y)
|
||||
model_line.set_data(model_path_x, model_path_y)
|
||||
|
||||
return real_ball, model_ball, real_line, model_line
|
||||
|
||||
# Create the animation
|
||||
anim = animation.FuncAnimation(
|
||||
fig, animate, init_func=init, frames=num_frames,
|
||||
interval=1000/fps, blit=True
|
||||
)
|
||||
|
||||
plt.close() # Prevent the static plot from displaying
|
||||
|
||||
return anim
|
||||
|
||||
|
||||
def display_animation(anim, html_video=True):
|
||||
"""
|
||||
Display an animation in a Jupyter notebook.
|
||||
|
||||
Parameters:
|
||||
anim (matplotlib.animation.Animation): The animation to display
|
||||
html_video (bool): If True, convert to HTML5 video, otherwise use JavaScript
|
||||
|
||||
Returns:
|
||||
IPython.display.HTML: The HTML object containing the animation
|
||||
"""
|
||||
if html_video:
|
||||
# Convert to HTML5 video
|
||||
html = anim.to_html5_video()
|
||||
return HTML(html)
|
||||
else:
|
||||
# Use JavaScript animation
|
||||
return HTML(anim.to_jshtml())
|
||||
|
||||
|
||||
def save_animation(anim, filename, fps=30, dpi=100):
|
||||
"""
|
||||
Save an animation to a file.
|
||||
|
||||
Parameters:
|
||||
anim (matplotlib.animation.Animation): The animation to save
|
||||
filename (str): The filename to save to (should end with .mp4, .gif, etc.)
|
||||
fps (int): Frames per second
|
||||
dpi (int): Resolution in dots per inch
|
||||
"""
|
||||
# Determine the writer based on the file extension
|
||||
if filename.endswith('.mp4'):
|
||||
writer = animation.FFMpegWriter(fps=fps)
|
||||
elif filename.endswith('.gif'):
|
||||
writer = animation.PillowWriter(fps=fps)
|
||||
else:
|
||||
raise ValueError("Unsupported file format. Use .mp4 or .gif")
|
||||
|
||||
# Save the animation
|
||||
anim.save(filename, writer=writer, dpi=dpi)
|
||||
Reference in New Issue
Block a user