Files
bounce/src/lab2_part2.py
bennettldavid c6b08a089d 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)
2025-03-01 16:55:29 -07:00

408 lines
13 KiB
Python

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.signal import find_peaks
sns.set_theme(style="whitegrid", palette="muted")
def read_trial(path):
"""
Reads a CSV file containing time and position data (with no header)
and returns a DataFrame with columns 'Time' and 'Position'.
"""
return pd.read_csv(path, header=None, names=['Time', 'Position'])
def detect_bounces(df,
prominence=0,
min_distance=10,
min_time_diff=0.2,
max_bounces=7,
fs=None):
"""
Detects bounce events by locating peaks in the signal.
Parameters:
- df: DataFrame with 'Time' and 'Position'
- prominence: required prominence for peaks (in the signal)
- min_distance: minimum number of data points between peaks
- min_time_diff: minimum time difference (seconds) between bounces
- max_bounces: maximum bounces to report
- fs: override sampling frequency (Hz). If None, computed from 'Time'.
Returns:
- peak_indices (array of indices),
- bounce_heights (array of position values at these bounces),
- bounce_times (array of times of these bounces),
- signal_data (the data used for detection, here the raw Position values)
"""
if fs is None:
dt = np.median(np.diff(df['Time']))
if dt <= 0:
raise ValueError(
"Time differences are non-positive. "
"Ensure your 'Time' column is in seconds."
)
fs = 1 / dt
signal_data = df['Position'].values
peaks, _ = find_peaks(signal_data, prominence=prominence, distance=min_distance)
filtered_peaks = []
filtered_times = []
group_start_time = None
group_best_idx = None
group_best_value = None
for idx in peaks:
current_time = df['Time'].iloc[idx]
current_value = signal_data[idx]
if group_start_time is None:
group_start_time = current_time
group_best_idx = idx
group_best_value = current_value
continue
if (current_time - group_start_time) < min_time_diff:
if current_value > group_best_value:
group_best_idx = idx
group_best_value = current_value
else:
filtered_peaks.append(group_best_idx)
filtered_times.append(df['Time'].iloc[group_best_idx])
group_start_time = current_time
group_best_idx = idx
group_best_value = current_value
if group_best_idx is not None:
filtered_peaks.append(group_best_idx)
filtered_times.append(df['Time'].iloc[group_best_idx])
filtered_peaks = filtered_peaks[:max_bounces]
filtered_times = filtered_times[:max_bounces]
bounce_heights = df['Position'].iloc[filtered_peaks].values
bounce_times = np.array(filtered_times)
return np.array(filtered_peaks), bounce_heights, bounce_times, signal_data
def compute_cor(bounce_heights):
"""
Computes the Coefficient of Restitution (COR) for consecutive bounces:
e = sqrt( h_{n+1} / h_n )
Returns an array of COR values.
"""
cor_values = []
for i in range(len(bounce_heights) - 1):
if bounce_heights[i] > 0:
cor_values.append(np.sqrt(bounce_heights[i+1] / bounce_heights[i]))
else:
cor_values.append(np.nan)
return np.array(cor_values)
def process_trial(path,
prominence,
min_distance=10,
min_time_diff=0.2,
max_bounces=7,
fs=None):
"""
Reads trial data, detects bounces, computes COR values, and returns:
- df (the DataFrame)
- peak_indices
- bounce_heights
- bounce_times
- cor_values
- avg_cor
- signal_data
"""
df = read_trial(path)
peak_indices, bounce_heights, bounce_times, signal_data = detect_bounces(
df, prominence, min_distance, min_time_diff, max_bounces, fs
)
cor_values = compute_cor(bounce_heights)
avg_cor = np.mean(cor_values) if cor_values.size > 0 else np.nan
return df, peak_indices, bounce_heights, bounce_times, cor_values, avg_cor, signal_data
def plot_trial(df, peak_indices, signal_data, label, ball_type,
initial_height=None, cor=None, bounces=7, dt=0.001, g=386.09):
"""
Plots raw data (Position vs. Time) for a single trial and marks the detected bounces with red 'x'.
Ensures:
- Displays exactly 7 bounces.
- X-axis starts just before the first data point and ends slightly after the last peak or data point.
- Y-axis is a tight fit to cover the full range of data.
"""
plt.figure(figsize=(10, 6))
# 1) Plot raw signal
plt.plot(df['Time'], df['Position'], linewidth=1.5, color='blue', label='Raw Data')
# 2) Mark the detected bounces
if len(peak_indices) > 0:
plt.plot(
df['Time'].iloc[peak_indices],
df['Position'].iloc[peak_indices],
'gx',
markersize=10,
label='Detected Bounces'
)
# 3) Simulated Idealized Bounce Trajectory (Ensuring Exactly 7 Bounces)
if initial_height is not None and cor is not None and not np.isnan(cor):
time_points = []
position_points = []
# Initialize simulation parameters
y = initial_height
v_y = 0.0 # Starting from rest
t = 0.0
bounce_count = 0
model_first_bounce_time = None
while bounce_count < bounces: # Ensure exactly 7 bounces
time_points.append(t)
position_points.append(y)
# Kinematic update
v_y -= g * dt
y += v_y * dt
t += dt
# Bounce condition
if y <= 0:
bounce_count += 1
y = 0 # Ground impact
v_y = -v_y * cor # Apply COR
if model_first_bounce_time is None:
model_first_bounce_time = t # Store first bounce time
# Align first model bounce with data's first detected bounce
if model_first_bounce_time is not None and len(peak_indices) > 0:
data_first_bounce_time = df['Time'].iloc[peak_indices[0]]
time_shift = data_first_bounce_time - model_first_bounce_time
time_points = [t + time_shift for t in time_points]
# Plot idealized trajectory
plt.plot(time_points, position_points, 'r--', label='Ideal Trajectory')
# 4) Formatting the plot
plt.title(f'{ball_type} Trial: {label}')
plt.xlabel('Time (s)')
plt.ylabel('Position (inches)')
plt.autoscale(True, 'x', True)
plt.grid(True)
plt.legend()
plt.show()
def process_and_plot_all(file_paths,
ball_type,
plot_trials=False,
min_distance=10,
min_time_diff=0.2,
relative_prominence=None,
low_height_threshold=8,
low_height_adjustment=0.5,
high_height_threshold=15,
high_height_adjustment=1.2,
max_bounces=7,
fs=None):
"""
Iterates over the given 'file_paths' dict:
- extracts the numeric initial height from the key label (e.g. "11 inches"),
- computes effective prominence,
- processes each trial,
- optionally plots each trial (with or without an ideal bounce overlay).
Returns a summary DataFrame with:
'Trial', 'Initial Height', 'Average COR', 'Bounce Times', 'Bounce Heights', and 'COR Values'.
"""
results = []
for init_label, path in file_paths.items():
try:
initial_height = float(init_label.split()[0])
except ValueError:
initial_height = np.nan
# Compute effective prominence
effective_prominence = relative_prominence * initial_height
if initial_height < low_height_threshold:
effective_prominence *= low_height_adjustment
elif initial_height > high_height_threshold:
effective_prominence *= high_height_adjustment
print(f"Processing trial '{init_label}' with effective prominence: {effective_prominence}")
# Process the trial
df, peak_indices, bounce_heights, bounce_times, cor_values, avg_cor, signal_data = process_trial(
path, effective_prominence,
min_distance=min_distance,
min_time_diff=min_time_diff,
max_bounces=max_bounces,
fs=fs
)
results.append({
'Trial': init_label,
'Initial Height': initial_height,
'Average COR': avg_cor,
'Bounce Times': bounce_times,
'Bounce Heights': bounce_heights,
'COR Values': cor_values
})
if plot_trials:
plot_trial(
df,
peak_indices,
signal_data,
label=init_label,
ball_type=ball_type,
initial_height=initial_height,
cor=avg_cor,
bounces=5,
dt=0.001,
g=386.09 #using inches for position
)
summary_df = pd.DataFrame(results)
summary_df.sort_values(by='Initial Height', inplace=True)
return summary_df
def plot_cor_table(cor_df, ball_type):
"""
Plots a line (and markers) of Average COR vs. Initial Height from the summary DataFrame.
"""
plt.figure(figsize=(8, 5))
sns.lineplot(x='Initial Height', y='Average COR', marker='o', data=cor_df)
plt.title(f'Average COR vs. Initial Height for {ball_type} Ball')
plt.xlabel('Initial Height (inches)')
plt.ylabel('Average COR')
plt.grid(True)
plt.tight_layout()
plt.show()
if __name__ == '__main__':
golf_file_paths = {
"11 inches": "golf_11.csv",
"12 inches": "golf_12.csv",
"13 inches": "golf_13.csv",
"14 inches": "golf_14.csv",
"15 inches": "golf_15.csv"
}
lacrosse_file_paths = {
"18 inches": "l_18.csv",
"19 inches": "l_19.csv",
"20 inches": "l_20.csv",
"21 inches": "l_21.csv",
"22 inches": "l_22.csv"
}
metal_file_paths = {
"2 inches": "metal_2.csv",
"4 inches": "metal_4.csv",
"6 inches": "metal_6.csv",
"8 inches": "metal_8.csv",
"10 inches": "metal_10.csv"
}
# General detection/plotting parameters
min_distance = 5
min_time_diff = 0.05
fs = 50
# Prominence adjustments
golf_relative_prominence = 0.15
golf_low_threshold = 12.5
golf_low_adjustment = 1.1
golf_high_threshold = 14.5
golf_high_adjustment = 1.3
lacrosse_relative_prominence = 0.05
lacrosse_low_threshold = 18.5
lacrosse_low_adjustment = 1.1
lacrosse_high_threshold = 21.5
lacrosse_high_adjustment = 1.3
metal_relative_prominence = 0.05
metal_low_threshold = 5
metal_low_adjustment = 0.8
metal_high_threshold = 8
metal_high_adjustment = 1.2
# Process & plot Golf
golf_results = process_and_plot_all(
golf_file_paths,
ball_type='Golf',
plot_trials=True,
min_distance=min_distance,
min_time_diff=min_time_diff,
relative_prominence=golf_relative_prominence,
low_height_threshold=golf_low_threshold,
low_height_adjustment=golf_low_adjustment,
high_height_threshold=golf_high_threshold,
high_height_adjustment=golf_high_adjustment,
max_bounces=7,
fs=fs
)
# Process & plot Lacrosse
lacrosse_results = process_and_plot_all(
lacrosse_file_paths,
ball_type='Lacrosse',
plot_trials=True,
min_distance=min_distance,
min_time_diff=min_time_diff,
relative_prominence=lacrosse_relative_prominence,
low_height_threshold=lacrosse_low_threshold,
low_height_adjustment=lacrosse_low_adjustment,
high_height_threshold=lacrosse_high_threshold,
high_height_adjustment=lacrosse_high_adjustment,
max_bounces=7,
fs=fs
)
# Process & plot Metal
metal_results = process_and_plot_all(
metal_file_paths,
ball_type='Metal',
plot_trials=True,
min_distance=min_distance,
min_time_diff=min_time_diff,
relative_prominence=metal_relative_prominence,
low_height_threshold=metal_low_threshold,
low_height_adjustment=metal_low_adjustment,
high_height_threshold=metal_high_threshold,
high_height_adjustment=metal_high_adjustment,
max_bounces=7,
fs=fs
)
# Print summary tables
print("Golf Ball COR Table:")
print(golf_results[['Trial', 'Initial Height', 'Average COR']])
print("\nLacrosse Ball COR Table:")
print(lacrosse_results[['Trial', 'Initial Height', 'Average COR']])
print("\nMetal Ball COR Table:")
print(metal_results[['Trial', 'Initial Height', 'Average COR']])
# Plot COR vs. initial height for each ball type
plot_cor_table(golf_results, 'Golf')
plot_cor_table(lacrosse_results, 'Lacrosse')
plot_cor_table(metal_results, 'Metal')