main.py

#!/usr/bin/env python3
import math
import os
import numpy as np
from scipy.signal import butter, filtfilt
from qt_plot import plot_all_data

def get_input_filenamess():
    for i in os.listdir('input'):
        if i.endswith('.csv'):
            yield os.path.join('input', i)

def get_output_filename_for_input(input_filename):
    name = os.path.basename(input_filename).rsplit('.', 1)[0]
    return os.path.join('output', f'{name}_output.csv')

def load_from_csv(filename):
    # The CSVs have inconsistent delimiters and columns.
    # - handle both ; and , delimiters
    # - search for a set of expected headers and normalise them

    expected_column_aliases = {
        'time': ['timestamp'],
        'velocity': ['barbell velocity','barbell_velocity'],
        'displacement': ['barbell position','barbell_position']
    }

    # Read the first line to determine the delimiter
    with open(filename, 'r') as f:
        first_line = f.readline()
        delimiter = ';' if ';' in first_line else ','

    # Load CSV data into a matrix with headers
    data = np.genfromtxt(filename, delimiter=delimiter, names=True)

    # Normalize column names
    normalized_data = {}
    for standard_name, aliases in expected_column_aliases.items():
        allowable_names = [alias for alias in aliases]
        allowable_names.append(standard_name)

        matched = False

        for allowable_name in allowable_names:
            if allowable_name in data.dtype.names:
                normalized_data[standard_name] = data[allowable_name]
                matched = True
                break

        if not matched:
            raise ValueError(f'Expected column for {standard_name} not found in {filename}. Available columns: {data.dtype.names}')
    
    return normalized_data

def convert_timestamp_to_seconds(timestamps):
    return timestamps / 1000.0 # assuming timestamps are in milliseconds


class AnalysisResult:
    
    def __init__(
            self, 
            start_index,
            free_fall_index, 
            displacement_at_free_fall, 
            max_displacement_index,
            max_displacement,
            mean_velocity,
            mean_acceleration
        ):
        self.start_index = start_index
        self.end_index = max_displacement_index

        self.free_fall_index = free_fall_index
        self.displacement_at_free_fall = displacement_at_free_fall

        self.max_displacement = max_displacement
        self.max_displacement_index = max_displacement_index

        self.mean_velocity = mean_velocity
        self.mean_acceleration = mean_acceleration

def remove_position_offset(data):
    # Remove any initial offset in displacement
    initial_displacement = data['displacement'][0]
    data['displacement'] = data['displacement'] - initial_displacement
    return data;

def do_analysis(data) -> AnalysisResult:
    # Find the peak displacement over the data, then work backwards until velocity is 0 on either side
    max_displacement_index = np.argmax(data['displacement'])
    max_displacement = data['displacement'][max_displacement_index]

    # Find start index
    # this is the first index before peak displacement which has velocity > 0
    # velocity might be 0 at the peak, so start looking just before the peak
    i = max_displacement_index - 1
    while True:
        if i > 0 and data['velocity'][i-1] > 0:
            i -= 1
        else:
            break
    start_index = i

    # Find free fall
    # Find the first point (after the start) where acceleration goes below -9.81 m/s², or 'G'
    while data['acceleration'][i] > -9.81 and i < len(data['acceleration']) - 1:
        i += 1
    free_fall_index = i
    displacement_at_free_fall = data['displacement'][free_fall_index]

    # Determine the mean velocity and acceleration from start to free fall
    mean_velocity = np.mean(data['velocity'][start_index:free_fall_index])
    mean_acceleration = np.mean(data['acceleration'][start_index:free_fall_index])

    return AnalysisResult(
        start_index, 
        free_fall_index, 
        displacement_at_free_fall, 
        max_displacement_index, 
        max_displacement,
        mean_velocity,
        mean_acceleration
        )


def write_to_output_csv(output_filename, data, analysis_result: AnalysisResult):
    
    # add a column which is 1 when then index is within the analysis window, else 0
    analysis_window = np.zeros(len(data['time']), dtype=int)
    analysis_window[analysis_result.start_index:analysis_result.end_index + 1] = 1
    data['analysis_window'] = analysis_window
    data['max_displacement'] = np.full(len(data['time']), analysis_result.max_displacement)
    data['displacement_at_free_fall'] = np.full(len(data['time']), analysis_result.displacement_at_free_fall)

    # convert the dict to a structured array for saving
    structured_array = np.zeros(len(data['time']), dtype=[(key, 'f8') for key in data.keys()])
    format = []
    for key in data.keys():
        structured_array[key] = data[key]
        if key in ['analysis_window']:
            format.append('%d')
        elif key in ['time']:
            format.append('%.3f')
        else:
            format.append('%.10f')
    np.savetxt(output_filename, structured_array, delimiter=',', header=','.join(structured_array.dtype.names), comments='',
               fmt=format)

def write_summary_csv(all_data, summary_filename):

    # The summary CSV has columns:
    # name, max_displacement, displacement_at_free_fall, mean_velocity, mean_acceleration

    summary_data = np.zeros(len(all_data),
                            dtype=[
                                ('name', 'U100'), 
                                ('max_displacement', 'f8'), 
                                ('displacement_at_free_fall', 'f8'),
                                ('mean_velocity', 'f8'),
                                ('mean_acceleration', 'f8')
                            ])

    for i, entry in enumerate(all_data):
        summary_data[i] = (
                entry['name'], 
                entry['max_displacement'], 
                entry['displacement_at_free_fall'], 
                entry['mean_velocity'], 
                entry['mean_acceleration']
            )
    
    np.savetxt(summary_filename, summary_data, delimiter=',', header=','.join(summary_data.dtype.names), comments='',
               fmt=['%s', '%.3f', '%.3f', '%.3f', '%.3f'])


def butterworth_filter(data, cutoff_freq, sampling_rate, order=4):
    nyquist_freq = sampling_rate / 2.0
    normalized_cutoff = cutoff_freq / nyquist_freq
    b, a = butter(order, normalized_cutoff, btype='low')
    filtered_data = filtfilt(b, a, data)
    return filtered_data


def main():

    all_data = []
    os.makedirs('output', exist_ok=True)

    for input_file in get_input_filenamess():

        try:
            output_file = get_output_filename_for_input(input_file)

            data = load_from_csv(input_file)
            data['time'] = convert_timestamp_to_seconds(data['time'])

            data['velocity_raw'] = data['velocity']

            # Apply a 4th order butterworth filter to the velocity with
            # a cutoff frequency of 6 Hz to reduce noise
            sampling_rate = 200.0 # Hz
            cutoff_freq = 6.0  # Hz
            data['velocity'] = butterworth_filter(
                data['velocity'], 
                cutoff_freq, 
                sampling_rate, 
                order=4 )

            # Derive acceleration from velocity
            data['acceleration'] = np.gradient(data['velocity'], data['time'])

            data = remove_position_offset(data)
            result = do_analysis(data)

            all_data.append({
                "name": os.path.basename(input_file),
                "data": data,
                "start": result.start_index,
                "end": result.end_index,
                "max_displacement_index": result.max_displacement_index,
                "max_displacement": result.max_displacement,
                "free_fall_index": result.free_fall_index,
                "displacement_at_free_fall": result.displacement_at_free_fall,
                "mean_velocity": result.mean_velocity,
                "mean_acceleration": result.mean_acceleration
            })

            write_to_output_csv(output_file, data, result)
            
        except Exception as e:
            print(f"ERROR processing {input_file}: {e}")

    write_summary_csv(all_data, 'output/summary.csv')
    
    plot_all_data(all_data)


if __name__ == "__main__":
    main()