""" 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