- 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)
408 lines
13 KiB
Python
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') |