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)
260 lines
8.8 KiB
Python
260 lines
8.8 KiB
Python
"""
|
|
Bounce detection module for the bouncing ball analysis project.
|
|
"""
|
|
|
|
import numpy as np
|
|
from scipy.signal import find_peaks
|
|
import pandas as pd
|
|
|
|
# Default parameters for bounce detection
|
|
DEFAULT_PROMINENCE = 0.1
|
|
DEFAULT_MIN_DISTANCE = 10
|
|
DEFAULT_MIN_TIME_DIFF = 0.2
|
|
DEFAULT_MAX_BOUNCES = 7
|
|
|
|
# Ball-specific parameters
|
|
BALL_PARAMS = {
|
|
'Golf': {
|
|
'relative_prominence': 0.15,
|
|
'low_threshold': 12.5,
|
|
'low_adjustment': 1.1,
|
|
'high_threshold': 14.5,
|
|
'high_adjustment': 1.3
|
|
},
|
|
'Lacrosse': {
|
|
'relative_prominence': 0.05,
|
|
'low_threshold': 18.5,
|
|
'low_adjustment': 1.1,
|
|
'high_threshold': 21.5,
|
|
'high_adjustment': 1.3
|
|
},
|
|
'Metal': {
|
|
'relative_prominence': 0.05,
|
|
'low_threshold': 5,
|
|
'low_adjustment': 0.8,
|
|
'high_threshold': 8,
|
|
'high_adjustment': 1.2
|
|
}
|
|
}
|
|
|
|
def detect_bounces(df,
|
|
prominence=DEFAULT_PROMINENCE,
|
|
min_distance=DEFAULT_MIN_DISTANCE,
|
|
min_time_diff=DEFAULT_MIN_TIME_DIFF,
|
|
max_bounces=DEFAULT_MAX_BOUNCES,
|
|
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:
|
|
# Filter out duplicate time values and sort
|
|
unique_times = np.unique(df['Time'].values)
|
|
if len(unique_times) > 1:
|
|
dt = np.median(np.diff(unique_times))
|
|
if dt <= 0:
|
|
# Fallback to a default sampling rate if time differences are non-positive
|
|
fs = 1000 # 1000 Hz is a reasonable default for high-speed data
|
|
else:
|
|
fs = 1 / dt
|
|
else:
|
|
# If there's only one unique time value, use a default sampling rate
|
|
fs = 1000
|
|
|
|
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 )
|
|
|
|
Parameters:
|
|
bounce_heights (numpy.ndarray): Array of bounce heights
|
|
|
|
Returns:
|
|
numpy.ndarray: 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(df,
|
|
initial_height=None,
|
|
ball_type=None,
|
|
prominence=None,
|
|
min_distance=DEFAULT_MIN_DISTANCE,
|
|
min_time_diff=DEFAULT_MIN_TIME_DIFF,
|
|
max_bounces=DEFAULT_MAX_BOUNCES,
|
|
fs=None):
|
|
"""
|
|
Processes a trial by detecting bounces and computing COR values.
|
|
|
|
Parameters:
|
|
df (pandas.DataFrame): DataFrame with 'Time' and 'Position' columns
|
|
initial_height (float): Initial height of the ball
|
|
ball_type (str): Type of ball ('Golf', 'Lacrosse', or 'Metal')
|
|
prominence (float): Prominence for peak detection (if None, calculated from initial_height and ball_type)
|
|
min_distance (int): Minimum distance between peaks
|
|
min_time_diff (float): Minimum time difference between bounces
|
|
max_bounces (int): Maximum number of bounces to detect
|
|
fs (float): Sampling frequency (Hz)
|
|
|
|
Returns:
|
|
dict: Dictionary containing analysis results
|
|
"""
|
|
# Calculate prominence if not provided
|
|
if prominence is None and initial_height is not None and ball_type is not None:
|
|
prominence = calculate_effective_prominence(initial_height, ball_type)
|
|
elif prominence is None:
|
|
prominence = DEFAULT_PROMINENCE
|
|
|
|
# Detect bounces
|
|
peak_indices, bounce_heights, bounce_times, signal_data = detect_bounces(
|
|
df, prominence, min_distance, min_time_diff, max_bounces, fs
|
|
)
|
|
|
|
# Calculate COR values
|
|
cor_values = compute_cor(bounce_heights)
|
|
avg_cor = np.mean(cor_values) if cor_values.size > 0 else np.nan
|
|
|
|
# Return results as a dictionary
|
|
return {
|
|
'peak_indices': peak_indices,
|
|
'bounce_heights': bounce_heights,
|
|
'bounce_times': bounce_times,
|
|
'cor_values': cor_values,
|
|
'Average COR': avg_cor,
|
|
'signal_data': signal_data,
|
|
'Initial Height': initial_height,
|
|
'Num Bounces': len(peak_indices)
|
|
}
|
|
|
|
|
|
def calculate_effective_prominence(initial_height, ball_type):
|
|
"""
|
|
Calculate the effective prominence for bounce detection based on initial height and ball type.
|
|
|
|
Parameters:
|
|
initial_height (float): Initial height of the ball
|
|
ball_type (str): Type of ball ('Golf', 'Lacrosse', or 'Metal')
|
|
|
|
Returns:
|
|
float: Effective prominence value
|
|
"""
|
|
if ball_type not in BALL_PARAMS:
|
|
raise ValueError(f"Unknown ball type: {ball_type}. Must be one of: {list(BALL_PARAMS.keys())}")
|
|
|
|
params = BALL_PARAMS[ball_type]
|
|
|
|
# Calculate base prominence
|
|
effective_prominence = params['relative_prominence'] * initial_height
|
|
|
|
# Apply adjustments based on height
|
|
if initial_height < params['low_threshold']:
|
|
effective_prominence *= params['low_adjustment']
|
|
elif initial_height > params['high_threshold']:
|
|
effective_prominence *= params['high_adjustment']
|
|
|
|
return effective_prominence
|
|
|
|
|
|
def process_trials(file_paths, ball_type, min_distance=DEFAULT_MIN_DISTANCE,
|
|
min_time_diff=DEFAULT_MIN_TIME_DIFF, max_bounces=DEFAULT_MAX_BOUNCES, fs=None):
|
|
"""
|
|
Process multiple trials and return a summary DataFrame.
|
|
|
|
Parameters:
|
|
file_paths (dict): Dictionary mapping labels to file paths
|
|
ball_type (str): Type of ball ('Golf', 'Lacrosse', or 'Metal')
|
|
min_distance (int): Minimum distance between peaks
|
|
min_time_diff (float): Minimum time difference between bounces
|
|
max_bounces (int): Maximum number of bounces to detect
|
|
fs (float): Sampling frequency (Hz)
|
|
|
|
Returns:
|
|
pandas.DataFrame: Summary DataFrame with trial information
|
|
"""
|
|
from src.data.loader import load_trial
|
|
|
|
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 = calculate_effective_prominence(initial_height, ball_type)
|
|
|
|
# Load and process the trial
|
|
df = load_trial(path)
|
|
results.append(process_trial(
|
|
df, initial_height, ball_type, effective_prominence, min_distance, min_time_diff, max_bounces, fs
|
|
))
|
|
|
|
summary_df = pd.DataFrame(results)
|
|
summary_df.sort_values(by='Initial Height', inplace=True)
|
|
|
|
return summary_df |