src/functions.py - master - louintik/ML

src/functions.py
701 строка26 KB
louisa
first_commit
14 май 2026, 22:04
14 май 2026, 22:0451ac654
import pandas as pd
import numpy as np
from scipy.signal import find_peaks
from collections import defaultdict
import plotly.graph_objects as go
from scipy.integrate import simpson
from scipy.stats import pearsonr


def detect_data_start(filename, max_rows_to_check=13):
    # находит, где начинаются числовые данные в CSV
    with open(filename, 'r') as f:
        for line_number, line_content in enumerate(f):
            clean_line = line_content.strip()
            split_values = clean_line.split(',')
            if len(split_values) == 0:
                continue
            first_value = split_values[0].strip()
            try:
                numeric_value = float(
                    first_value.replace('+', '').replace('E', 'e')
                )
                return line_number
            except ValueError:
                continue
    raise ValueError("Не удалось найти строку с числовыми данными.")


def load_signal_data(filename):
    # загружает и подготавливает сигнал из файла
    start_row = detect_data_start(filename)
    df = pd.read_csv(filename, skiprows=start_row, header=None)
    df.columns = ['Time', 'CH1']
    
    signal = df['CH1'].astype(float).values + 0.02
    time = df['Time'].astype(float).values
    
    return signal, time


def find_all_peaks(signal, height_threshold=0.03):
    # обнаруживает положительные/отрицательные пики
    peaks_plus = find_peaks(signal, height=height_threshold)
    peaks_minus = find_peaks(-signal, height=height_threshold)
    
    peaks_index_all = np.concatenate((peaks_plus[0], peaks_minus[0]))
    peak_heights_all = np.concatenate((peaks_plus[1]['peak_heights'], peaks_minus[1]['peak_heights']))
    
    index_to_height = dict(zip(peaks_index_all, peak_heights_all))
    sorted_index_to_height = dict(sorted(index_to_height.items()))
    sorted_peak_index = list(sorted_index_to_height.keys())
    
    return sorted_peak_index, sorted_index_to_height


def group_peaks(sorted_peak_index, window_size=20):
    # группирует близкие пики
    if not sorted_peak_index:
        return {}
    
    window_number_peaks = {}
    number_group = 1
    start_peak_index = sorted_peak_index[0]
    
    for value in sorted_peak_index:
        if start_peak_index <= value <= start_peak_index + window_size:
            window_number_peaks[value] = number_group
        else:
            number_group += 1
            start_peak_index = value
            window_number_peaks[value] = number_group
    
    group_to_indices = defaultdict(list)
    for index, group in window_number_peaks.items():
        group_to_indices[group].append(index)
    
    # Remove groups with fewer than 2 peaks
    group_to_indices = {
        g: idxs for g, idxs in group_to_indices.items()
        if len(idxs) >= 2
    }
    
    return group_to_indices


def extract_key_points(group_to_indices, sorted_index_to_height):
    # определяет первую, среднюю (максимальную) и последнюю точки для каждой группы
    group_prefix_suffix = {}

    for group, indices in group_to_indices.items():
        indices_sorted = sorted(indices)

        # 1. Первый пик по индексу
        first_index = indices_sorted[0]

        # 2. Для слов из 3+ пиков: ищем максимальный среди средних
        #    Для слов из 2 пиков: второй пик становится "средним" (релаксационным)
        if len(indices_sorted) > 2:
            middle_indices = indices_sorted[1:-1]
            if middle_indices:
                middle_max_index = max(middle_indices, key=lambda x: sorted_index_to_height[x])
            else:
                middle_max_index = None
        elif len(indices_sorted) == 2:
            # Для слов из 2 пиков: второй пик = middle_max_index
            # Префикс будет от пика 1 до пика 2, суффикс = 0
            middle_max_index = indices_sorted[1]
        else:
            middle_max_index = None

        # 3. Последний пик по индексу
        last_index = indices_sorted[-1]

        group_prefix_suffix[group] = {
            'first_index': first_index,
            'middle_max_index': middle_max_index,
            'last_index': last_index
        }

    return group_prefix_suffix


def extract_prefix_features(signal, time, i1, i2):
    """Extract features from the prefix segment (from first to middle peak)."""
    if i2 is None or i1 >= i2:
        return {
            'prefix_duration': None,
            'prefix_max_amp': None,
            'prefix_min_amp': None,
            'prefix_mean': None,
            'prefix_std': None,
            'prefix_energy': None,
            'prefix_integral': None,
            'prefix_direction': None,
            'prefix_second_deriv_mean': None,
            'prefix_third_deriv_mean': None,
        }
    
    prefix_segment = signal[i1:i2]
    prefix_time = time[i1:i2]
    prefix_duration = time[i2] - time[i1]
    
    # Амплитуды
    prefix_max_amp = np.max(prefix_segment)
    prefix_min_amp = np.min(prefix_segment)
    prefix_mean = np.mean(prefix_segment)
    prefix_std = np.std(prefix_segment)
    prefix_energy = np.sum(prefix_segment ** 2)
    
    # Интеграл (площадь под кривой) с помощью Simpson's rule
    if len(prefix_segment) > 2:
        prefix_integral = simpson(prefix_segment, prefix_time)
    else:
        prefix_integral = None
    
    # Направление пика префикса (положительное или отрицательное)
    prefix_direction = 1 if prefix_segment[-1] > prefix_segment[0] else -1
    
    # Вторая и третья производные для анализа формы
    if len(prefix_segment) >= 4:
        # Вторая производная
        second_deriv = np.diff(prefix_segment, n=2)
        prefix_second_deriv_mean = np.mean(second_deriv)
        
        # Третья производная
        third_deriv = np.diff(prefix_segment, n=3)
        prefix_third_deriv_mean = np.mean(third_deriv)
    else:
        prefix_second_deriv_mean = None
        prefix_third_deriv_mean = None
    
    return {
        'prefix_duration': prefix_duration,
        'prefix_max_amp': prefix_max_amp,
        'prefix_min_amp': prefix_min_amp,
        'prefix_mean': prefix_mean,
        'prefix_std': prefix_std,
        'prefix_energy': prefix_energy,
        'prefix_integral': prefix_integral,
        'prefix_direction': prefix_direction,
        'prefix_second_deriv_mean': prefix_second_deriv_mean,
        'prefix_third_deriv_mean': prefix_third_deriv_mean,
    }


def extract_suffix_features(signal, time, i2, i3):
    """Extract features from the suffix segment (from middle to last peak)."""
    if i2 is None or i2 >= i3:
        return {
            'suffix_duration': None,
            'suffix_max_amp': None,
            'suffix_min_amp': None,
            'suffix_mean': None,
            'suffix_std': None,
            'suffix_energy': None,
            'suffix_integral': None,
            'suffix_direction': None,
            'suffix_second_deriv_mean': None,
            'suffix_third_deriv_mean': None,
        }
    
    suffix_segment = signal[i2:i3]
    suffix_time = time[i2:i3]
    suffix_duration = time[i3] - time[i2]
    
    # Амплитуды
    suffix_max_amp = np.max(suffix_segment)
    suffix_min_amp = np.min(suffix_segment)
    suffix_mean = np.mean(suffix_segment)
    suffix_std = np.std(suffix_segment)
    suffix_energy = np.sum(suffix_segment ** 2)
    
    # Интеграл (площадь под кривой)
    if len(suffix_segment) > 2:
        suffix_integral = simpson(suffix_segment, suffix_time)
    else:
        suffix_integral = None
    
    # Направление пика суффикса
    suffix_direction = 1 if suffix_segment[-1] > suffix_segment[0] else -1
    
    # Вторая и третья производные
    if len(suffix_segment) >= 4:
        second_deriv = np.diff(suffix_segment, n=2)
        suffix_second_deriv_mean = np.mean(second_deriv)
        
        third_deriv = np.diff(suffix_segment, n=3)
        suffix_third_deriv_mean = np.mean(third_deriv)
    else:
        suffix_second_deriv_mean = None
        suffix_third_deriv_mean = None
    
    return {
        'suffix_duration': suffix_duration,
        'suffix_max_amp': suffix_max_amp,
        'suffix_min_amp': suffix_min_amp,
        'suffix_mean': suffix_mean,
        'suffix_std': suffix_std,
        'suffix_energy': suffix_energy,
        'suffix_integral': suffix_integral,
        'suffix_direction': suffix_direction,
        'suffix_second_deriv_mean': suffix_second_deriv_mean,
        'suffix_third_deriv_mean': suffix_third_deriv_mean,
    }


def detect_press_type_from_filename(filename):
    """
    Определить тип нажатия по названию файла:
    - up - твердая поверхность
    - up-red - упругая поверхность
    - side - боковое нажатие
    """
    basename = filename.lower()
    
    if 'up-red' in basename or 'up_red' in basename:
        return 'elastic_surface'
    elif 'side' in basename:
        return 'side_press'
    elif 'up' in basename:
        return 'hard_surface'
    else:
        return 'unknown'


def calculate_word_amplitudes(signal, i1, i3):
    """
    Рассчитать максимальную и минимальную амплитуду во всём слове.
    Используем абсолютные значения для сравнения, т.к. пики могут быть 
    положительными и отрицательными.
    Возвращаем оба варианта: по модулю и обычные.
    """
    word_segment = signal[i1:i3]
    
    # Обычные амплитуды
    max_amp = np.max(word_segment)
    min_amp = np.min(word_segment)
    
    # Амплитуды по модулю (абсолютные значения)
    max_amp_abs = np.max(np.abs(word_segment))
    min_amp_abs = np.min(np.abs(word_segment))
    
    return {
        'word_max_amp': max_amp,
        'word_min_amp': min_amp,
        'word_max_amp_abs': max_amp_abs,
        'word_min_amp_abs': min_amp_abs,
    }


def calculate_word_duration(time, i1, i3):
    """
    Рассчитать длительность всего слова.
    Примечание: первый пик выделяется уже после начала сигнального слова,
    поэтому фактическая длительность может быть больше.
    """
    word_duration = time[i3] - time[i1]
    return {'word_duration': word_duration}


def calculate_peak_shape_analysis(signal, i1, i3, num_points=10):
    """
    Анализировать форму пика с помощью производных.
    Возвращает статистику второй и третьей производных для всего слова.
    """
    segment = signal[i1:i3]
    
    if len(segment) < 5:
        return {
            'second_deriv_mean': None,
            'second_deriv_std': None,
            'second_deriv_max': None,
            'third_deriv_mean': None,
            'third_deriv_std': None,
            'third_deriv_max': None,
        }
    
    # Вторая производная (кривизна)
    second_deriv = np.diff(segment, n=2)
    
    # Третья производная (скорость изменения кривизны)
    third_deriv = np.diff(segment, n=3)
    
    return {
        'second_deriv_mean': np.mean(second_deriv),
        'second_deriv_std': np.std(second_deriv),
        'second_deriv_max': np.max(np.abs(second_deriv)),
        'third_deriv_mean': np.mean(third_deriv),
        'third_deriv_std': np.std(third_deriv),
        'third_deriv_max': np.max(np.abs(third_deriv)),
    }


def calculate_signal_integrals(signal, time, i1, i2, i3):
    """
    Рассчитать интегралы (площадь под кривой) для префикса и суффикса.
    Учитывает направление: положительная или отрицательная площадь.
    """
    result = {
        'prefix_integral': None,
        'suffix_integral': None,
        'word_integral': None,
        'prefix_integral_abs': None,
        'suffix_integral_abs': None,
    }
    
    # Интеграл всего слова
    if i3 > i1 and len(signal[i1:i3]) > 2:
        result['word_integral'] = simpson(signal[i1:i3], time[i1:i3])
        result['word_integral_abs'] = simpson(np.abs(signal[i1:i3]), time[i1:i3])
    
    # Интеграл префикса
    if i2 is not None and i2 > i1 and len(signal[i1:i2]) > 2:
        result['prefix_integral'] = simpson(signal[i1:i2], time[i1:i2])
        result['prefix_integral_abs'] = simpson(np.abs(signal[i1:i2]), time[i1:i2])
    
    # Интеграл суффикса
    if i2 is not None and i3 > i2 and len(signal[i2:i3]) > 2:
        result['suffix_integral'] = simpson(signal[i2:i3], time[i2:i3])
        result['suffix_integral_abs'] = simpson(np.abs(signal[i2:i3]), time[i2:i3])
    
    return result


def detect_peak_direction(signal, i1, i2):
    """
    Определить направление пика префикса.
    Возвращает: 1 если положительное, -1 если отрицательное, 0 если не определено.
    """
    if i2 is None or i1 >= i2:
        return {'prefix_direction': 0}
    
    prefix_segment = signal[i1:i2]
    direction = 1 if prefix_segment[-1] > prefix_segment[0] else -1
    
    return {'prefix_direction': direction}


def calculate_prefix_suffix_correlation(signal, i1, i2, i3):
    """
    Рассчитать корреляцию между направлениями пиков префикса и суффикса.
    Проверяем, есть ли зависимость: если префикс положительный, 
    то суффикс отрицательный (и наоборот).
    """
    if i2 is None or i1 >= i2 or i2 >= i3:
        return {'prefix_suffix_correlation': None}
    
    prefix_segment = signal[i1:i2]
    suffix_segment = signal[i2:i3]
    
    # Направления
    prefix_direction = 1 if prefix_segment[-1] > prefix_segment[0] else -1
    suffix_direction = 1 if suffix_segment[-1] > suffix_segment[0] else -1
    
    # Корреляция направлений (1 = одинаковые, -1 = противоположные)
    correlation = prefix_direction * suffix_direction
    
    # Также считаем корреляцию Пирсона между сегментами
    # (нормализуем длины для сравнения)
    min_len = min(len(prefix_segment), len(suffix_segment))
    if min_len > 2:
        prefix_norm = prefix_segment[:min_len] - np.mean(prefix_segment[:min_len])
        suffix_norm = suffix_segment[:min_len] - np.mean(suffix_segment[:min_len])
        
        # Корреляция Пирсона
        if np.std(prefix_norm) > 0 and np.std(suffix_norm) > 0:
            pearson_corr, _ = pearsonr(prefix_norm, suffix_norm)
        else:
            pearson_corr = None
    else:
        pearson_corr = None
    
    return {
        'prefix_suffix_direction_correlation': correlation,
        'prefix_suffix_pearson_correlation': pearson_corr,
    }


def calculate_segment_features(signal, time, group_prefix_suffix):
    """Calculate features for each segment and return dataset rows."""
    dataset = []

    for group, pts in group_prefix_suffix.items():
        i1 = pts['first_index']
        i2 = pts['middle_max_index']
        i3 = pts['last_index']

        # Базовые фичи сегмента
        segment = signal[i1:i3]
        duration = time[i3] - time[i1]
        mean = np.mean(segment)
        std = np.std(segment)
        slope = (segment[-1] - segment[0]) / duration if duration > 0 else 0

        # 1-2) Фичи префикса и суффикса
        prefix_features = extract_prefix_features(signal, time, i1, i2)
        suffix_features = extract_suffix_features(signal, time, i2, i3)

        # 3) Тип нажатия (будет добавлен позже на уровне файла)
        
        # 4-5) Амплитуды слова (обычные и по модулю)
        amplitude_features = calculate_word_amplitudes(signal, i1, i3)

        # 6-8) Длительности
        duration_features = calculate_word_duration(time, i1, i3)

        # 9) Анализ формы пика (производные)
        shape_features = calculate_peak_shape_analysis(signal, i1, i3)

        # 10) Интегралы (площадь под кривой)
        integral_features = calculate_signal_integrals(signal, time, i1, i2, i3)

        # 11) Направление пика префикса
        direction_features = detect_peak_direction(signal, i1, i2)

        # 12) Корреляция префикса и суффикса
        correlation_features = calculate_prefix_suffix_correlation(signal, i1, i2, i3)

        # Объединяем все фичи
        features = {
            'group': group,
            'duration': duration,
            'mean': mean,
            'std': std,
            'slope': slope,
        }
        
        features.update(prefix_features)
        features.update(suffix_features)
        features.update(amplitude_features)
        features.update(duration_features)
        features.update(shape_features)
        features.update(integral_features)
        features.update(direction_features)
        features.update(correlation_features)

        dataset.append(features)

    return dataset


def process_single_file(filename, press_type=None):
    """Process a single file and return dataset with file metadata."""
    signal, time = load_signal_data(filename)
    sorted_peak_index, sorted_index_to_height = find_all_peaks(signal)
    group_to_indices = group_peaks(sorted_peak_index)
    group_prefix_suffix = extract_key_points(group_to_indices, sorted_index_to_height)
    dataset = calculate_segment_features(signal, time, group_prefix_suffix)

    # Определяем тип нажатия из имени файла, если не указан
    if press_type is None:
        press_type = detect_press_type_from_filename(filename)

    # Add file-level metadata to each row
    for row in dataset:
        row['file'] = filename
        row['press_type'] = press_type

    return dataset


def visualize_file_processing(filename, press_type=None, save_path=None):
    """
    Визуализировать для одного выбранного файла:
    - График сигнала
    - Выделенные слова (префикс и суффикс)
    - Метки пиков
    - Надписи с длительностями, амплитудами, производными и т.д.
    """
    # Обрабатываем файл
    signal, time = load_signal_data(filename)
    sorted_peak_index, sorted_index_to_height = find_all_peaks(signal)
    group_to_indices = group_peaks(sorted_peak_index)
    group_prefix_suffix = extract_key_points(group_to_indices, sorted_index_to_height)
    dataset = calculate_segment_features(signal, time, group_prefix_suffix)
    
    if press_type is None:
        press_type = detect_press_type_from_filename(filename)
    
    # Создаём график
    fig = go.Figure()
    
    # Добавляем сигнал
    fig.add_trace(
        go.Scatter(
            x=time,
            y=signal,
            mode='lines',
            marker=dict(color='black'),
            name='Сигнал',
            line=dict(width=1.5)
        )
    )
    
    # Добавляем все пики
    peaks_plus = find_peaks(signal, height=0.03)
    peaks_minus = find_peaks(-signal, height=0.03)
    
    fig.add_trace(
        go.Scatter(
            x=time[peaks_plus[0]],
            y=signal[peaks_plus[0]],
            mode='markers',
            marker=dict(color='red', size=6, symbol='circle'),
            name='Положительные пики'
        )
    )
    
    fig.add_trace(
        go.Scatter(
            x=time[peaks_minus[0]],
            y=signal[peaks_minus[0]],
            mode='markers',
            marker=dict(color='blue', size=6, symbol='circle'),
            name='Отрицательные пики'
        )
    )
    
    # Для каждой группы добавляем выделенные зоны и аннотации
    colors_prefix = ['rgba(0, 255, 0, 0.2)', 'rgba(0, 200, 0, 0.2)', 'rgba(0, 150, 0, 0.2)']
    colors_suffix = ['rgba(255, 165, 0, 0.2)', 'rgba(255, 140, 0, 0.2)', 'rgba(255, 100, 0, 0.2)']
    
    for idx, (group, pts) in enumerate(group_prefix_suffix.items()):
        i1 = pts['first_index']
        i2 = pts['middle_max_index']
        i3 = pts['last_index']
        
        color_idx = idx % len(colors_prefix)
        
        # Выделяем префикс (зелёный)
        if i2 is not None and i1 < i2:
            fig.add_vrect(
                x0=time[i1],
                x1=time[i2],
                fillcolor=colors_prefix[color_idx],
                opacity=0.3,
                layer="below",
                line_width=1,
                line_color="green",
            )
        
        # Выделяем суффикс (оранжевый)
        if i2 is not None and i2 < i3:
            fig.add_vrect(
                x0=time[i2],
                x1=time[i3],
                fillcolor=colors_suffix[color_idx],
                opacity=0.3,
                layer="below",
                line_width=1,
                line_color="orange",
            )
        
        # Ключевые точки
        fig.add_trace(
            go.Scatter(
                x=[time[i1]],
                y=[signal[i1]],
                mode='markers+text',
                marker=dict(color='green', size=10, symbol='diamond'),
                text=[f'G{group}'],
                textposition='top center',
                textfont=dict(size=10, color='green'),
                name=f'Group {group} start',
                showlegend=False
            )
        )
        
        if i2 is not None:
            fig.add_trace(
                go.Scatter(
                    x=[time[i2]],
                    y=[signal[i2]],
                    mode='markers',
                    marker=dict(color='yellow', size=10, symbol='star'),
                    name=f'Group {group} middle',
                    showlegend=False
                )
            )
        
        fig.add_trace(
            go.Scatter(
                x=[time[i3]],
                y=[signal[i3]],
                mode='markers',
                marker=dict(color='red', size=10, symbol='x'),
                name=f'Group {group} end',
                showlegend=False
            )
        )
        
        # Аннотация с информацией о слове
        if i2 is not None:
            # Находим соответствующие данные из dataset
            group_data = next((d for d in dataset if d['group'] == group), None)
            
            if group_data:
                # Helper function to format values, handling None
                def fmt(val, decimals=3):
                    if val is None:
                        return 'N/A'
                    return f"{val:.{decimals}f}"
                
                annotation_text = (
                    f"Group {group}<br>"
                    f"Prefix: {fmt(group_data.get('prefix_duration'))}s<br>"
                    f"Suffix: {fmt(group_data.get('suffix_duration'))}s<br>"
                    f"Total: {fmt(group_data.get('word_duration'))}s<br>"
                    f"Max Amp: {fmt(group_data.get('word_max_amp'))}<br>"
                    f"Min Amp: {fmt(group_data.get('word_min_amp'))}<br>"
                    f"2nd Deriv: {fmt(group_data.get('second_deriv_mean'))}<br>"
                    f"3rd Deriv: {fmt(group_data.get('third_deriv_mean'))}<br>"
                    f"Prefix Int: {fmt(group_data.get('prefix_integral'))}<br>"
                    f"Suffix Int: {fmt(group_data.get('suffix_integral'))}<br>"
                    f"Prefix Dir: {group_data.get('prefix_direction', 'N/A')}<br>"
                    f"Correlation: {fmt(group_data.get('prefix_suffix_pearson_correlation'))}"
                )
                
                # Позиция аннотации - над максимальным значением в группе
                y_max_idx = max(range(i1, i3), key=lambda x: abs(signal[x]))
                
                fig.add_annotation(
                    x=time[y_max_idx],
                    y=signal[y_max_idx],
                    text=annotation_text,
                    showarrow=True,
                    arrowhead=2,
                    arrowsize=1,
                    arrowwidth=1,
                    arrowcolor='black',
                    bgcolor='white',
                    bordercolor='black',
                    borderwidth=1,
                    font=dict(size=9),
                    yshift=10
                )
    
    # Обновляем layout
    fig.update_layout(
        title=f"Файл: {filename}<br>Тип нажатия: {press_type}",
        xaxis_title="Время (с)",
        yaxis_title="Амплитуда (В)",
        width=1400,
        height=800,
        showlegend=True,
        legend=dict(
            yanchor="top",
            y=0.99,
            xanchor="left",
            x=0.01
        )
    )
    
    # Показываем или сохраняем
    if save_path:
        fig.write_html(save_path)
        print(f"График сохранён: {save_path}")
    else:
        fig.show()
    
    return fig