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')