bennettldavid
c6b08a089d
- 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)
287 lines
9.1 KiB
Python
287 lines
9.1 KiB
Python
"""
|
|
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) |