qt_plot.py

#!/usr/bin/env python3
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                               QHBoxLayout, QListWidget, QPushButton, QSplitter,
                               QTableWidget, QTableWidgetItem, QHeaderView)
from PySide6.QtCore import Qt
from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg, NavigationToolbar2QT
from matplotlib.figure import Figure
import matplotlib.ticker as ticker
import numpy as np
import sys


class PlotViewer(QMainWindow):
    def __init__(self, all_data):
        super().__init__()
        
        self.all_data = all_data
        self.current_entry = None
        self.acceleration_visible = False
        self.markers_visible = False
        self.y_range_when_acceleration_hidden = (-5, 5)
        
        self.setWindowTitle("Data Plot Viewer")
        self.setGeometry(100, 100, 1200, 800)
        
        # Create central widget and main layout
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QHBoxLayout(central_widget)
        
        # Create splitter for resizable panels
        splitter = QSplitter(Qt.Horizontal)
        
        # Left panel - List widget
        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)
        left_layout.setContentsMargins(5, 5, 5, 5)
        
        # List widget
        self.list_widget = QListWidget()
        self.list_widget.currentRowChanged.connect(self.on_select)
        
        # Populate list
        for entry in all_data:
            self.list_widget.addItem(entry['name'])
        
        left_layout.addWidget(self.list_widget)
        
        # Create annotation table
        self.annotation_table = QTableWidget(6, 2)
        self.annotation_table.setHorizontalHeaderLabels(['Metric', 'Value'])
        self.annotation_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.annotation_table.verticalHeader().setVisible(False)
        self.annotation_table.setMaximumHeight(210)
        self.annotation_table.setEditTriggers(QTableWidget.NoEditTriggers)
        left_layout.addWidget(self.annotation_table)
        
        # Toggle acceleration button
        self.toggle_acceleration_btn = QPushButton("Toggle Acceleration")
        self.toggle_acceleration_btn.clicked.connect(self.toggle_acceleration)
        left_layout.addWidget(self.toggle_acceleration_btn)

        # Toggle markers button
        self.toggle_markers_btn = QPushButton("Toggle Markers")
        self.toggle_markers_btn.clicked.connect(self.toggle_markers)
        left_layout.addWidget(self.toggle_markers_btn)
        
        # Right panel - Plot widget
        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)
        right_layout.setContentsMargins(5, 5, 5, 5)
        
        # Create matplotlib figure and canvas
        self.fig = Figure(figsize=(10, 8))
        self.ax = self.fig.add_subplot(111)
        self.canvas = FigureCanvasQTAgg(self.fig)
        
        # Add navigation toolbar
        toolbar = NavigationToolbar2QT(self.canvas, right_widget)
        
        right_layout.addWidget(toolbar)
        right_layout.addWidget(self.canvas)
        
        # Add widgets to splitter
        splitter.addWidget(left_widget)
        splitter.addWidget(right_widget)
        splitter.setStretchFactor(0, 1)  # Left panel
        splitter.setStretchFactor(1, 4)  # Right panel (larger)
        
        main_layout.addWidget(splitter)
        
        # Store plot lines
        self.lines = {}
        
        # Cursor tracking
        self.cursor_line = None
        self.cursor_text = None
        self.canvas.mpl_connect('motion_notify_event', self.on_mouse_move)
        
        # Select first item by default
        if all_data:
            self.list_widget.setCurrentRow(0)
    
    def on_select(self, index):
        if index < 0 or index >= len(self.all_data):
            return
        
        self.current_entry = self.all_data[index]
        self.plot_entry(self.current_entry)
    
    def plot_entry(self, entry):
        name = entry["name"]
        data = entry["data"]
        
        # Clear previous plot
        self.ax.clear()
        self.lines = {}
        
        # Plot data
        line1, = self.ax.plot(data['time'], data['displacement'], 
                              label='Displacement (m)', marker='x', 
                              markersize=0, color='C0', linewidth=0.8)
        line3, = self.ax.plot(data['time'], data['velocity'], 
                      label=r'Velocity [Btw 4$^{\mathrm{th}}$] (m/s²)', marker='x', 
                      markersize=0, color='C1', linewidth=0.8, alpha=1.0)
        line2, = self.ax.plot(data['time'], data['velocity_raw'], 
                              label='Velocity [raw] (m/s)', marker='x', 
                              markersize=0, color='C1', linewidth=0.8, alpha=0.5)
        line4, = self.ax.plot(data['time'], data['acceleration'], 
                              label='Acceleration (m/s²)', marker='x', 
                              markersize=0, color='C2', linewidth=0.8)
        
        self.lines['displacement'] = line1

        self.lines['velocity_raw'] = line2
        self.lines['velocity'] = line3

        self.lines['acceleration'] = line4
        
        # Set visibility based on current state
        line4.set_visible(self.acceleration_visible)
        
        self.ax.set_title(name, fontsize=12)
        self.ax.set_xlabel('Time (s)')
        
        # Plot vertical lines for start and end indices
        start = entry["start"]
        end = entry["end"]
        self.ax.axvline(x=data['time'][start], color='gray', 
                       linestyle='--', linewidth=0.8, label='Start')
        self.ax.axvline(x=data['time'][end], color='black', 
                       linestyle='--', linewidth=0.8, label='End')
        
        # Plot max displacement marker
        if 'max_displacement_index' in entry:
            max_idx = entry['max_displacement_index']
            max_time = data['time'][max_idx]
            max_disp = entry['max_displacement']
            self.ax.plot(max_time, max_disp, 'rx', markersize=8, 
                        label=f'Max Disp ({max_disp:.3f}m)', zorder=5)
        
        # Plot free fall marker
        if 'free_fall_index' in entry:
            ff_idx = entry['free_fall_index']
            ff_time = data['time'][ff_idx]
            ff_disp = entry['displacement_at_free_fall']
            self.ax.plot(ff_time, ff_disp, 'gx', markersize=8, 
                        label=f'Free Fall Disp ({ff_disp:.3f}m)', zorder=5)
        
        # Update annotation table
        start_time = data['time'][start]
        end_time = data['time'][end]
        
        self.annotation_table.setItem(0, 0, QTableWidgetItem('Start Time'))
        self.annotation_table.setItem(0, 1, QTableWidgetItem(f'{start_time:.3f}s'))
        self.annotation_table.setItem(1, 0, QTableWidgetItem('End Time'))
        self.annotation_table.setItem(1, 1, QTableWidgetItem(f'{end_time:.3f}s'))
        
        if 'max_displacement_index' in entry:
            max_idx = entry['max_displacement_index']
            max_time = data['time'][max_idx]
            max_disp = entry['max_displacement']
            self.annotation_table.setItem(2, 0, QTableWidgetItem('Max Disp'))
            self.annotation_table.setItem(2, 1, QTableWidgetItem(f'{max_disp:.3f}m'))
            self.annotation_table.setItem(3, 0, QTableWidgetItem('Max Disp Time'))
            self.annotation_table.setItem(3, 1, QTableWidgetItem(f'{max_time:.3f}s'))
        
        if 'free_fall_index' in entry:
            ff_idx = entry['free_fall_index']
            ff_time = data['time'][ff_idx]
            ff_disp = entry['displacement_at_free_fall']
            self.annotation_table.setItem(4, 0, QTableWidgetItem('Free Fall Disp'))
            self.annotation_table.setItem(4, 1, QTableWidgetItem(f'{ff_disp:.3f}m'))
            self.annotation_table.setItem(5, 0, QTableWidgetItem('Free Fall Time'))
            self.annotation_table.setItem(5, 1, QTableWidgetItem(f'{ff_time:.3f}s'))
        
        # Scale the x-axis to focus on the analysis window
        context = 1  # seconds before / after
        self.ax.set_xlim(data['time'][start] - context, data['time'][end] + context)
        
        # Set major ticks every second and minor ticks every 200ms
        self.ax.xaxis.set_major_locator(ticker.MultipleLocator(1.0))
        self.ax.xaxis.set_minor_locator(ticker.MultipleLocator(0.2))
        
        # Add grid lines
        self.ax.grid(True, which='major', alpha=0.5, linewidth=0.8)
        self.ax.grid(True, which='minor', alpha=0.3, linewidth=0.4)
        
        # Set y-axis limits when acceleration is disabled
        if not self.acceleration_visible:
            self.ax.set_ylim(*self.y_range_when_acceleration_hidden)
        
        # Update legend
        self.ax.legend()
        
        # Initialize cursor elements
        self.cursor_line = self.ax.axvline(x=0, color='black', linestyle='-', 
                                          linewidth=0.8, visible=False)
        self.cursor_text = self.ax.text(0.02, 0.98, '', transform=self.ax.transAxes,
                                       verticalalignment='top', 
                                       bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8),
                                       visible=False)
        
        self.canvas.draw()
    
    def toggle_acceleration(self):
        if not self.current_entry:
            return
        
        self.acceleration_visible = not self.acceleration_visible
        
        # Toggle line visibility
        self.lines['acceleration'].set_visible(self.acceleration_visible)

        # Adjust y-axis limits
        if not self.acceleration_visible:
            self.ax.set_ylim(*self.y_range_when_acceleration_hidden)
        else:
            data = self.current_entry['data']
            y_min = min(data['displacement'].min(), data['velocity'].min(), 
                       data['acceleration'].min())
            y_max = max(data['displacement'].max(), data['velocity'].max(), 
                       data['acceleration'].max())
            y_range = y_max - y_min
            padding = y_range * 0.1
            self.ax.set_ylim(y_min - padding, y_max + padding)

        self.ax.legend()
        
        self.canvas.draw()

    def toggle_markers(self):
        if not self.current_entry:
            return
        
        self.markers_visible = not self.markers_visible
        
        marker_size = 5 if self.markers_visible else 0
        
        for line in self.lines.values():
            line.set_markersize(marker_size)
        
        self.canvas.draw()
    
    def on_mouse_move(self, event):
        if event.inaxes != self.ax or not self.current_entry:
            if self.cursor_line:
                self.cursor_line.set_visible(False)
            if self.cursor_text:
                self.cursor_text.set_visible(False)
            self.canvas.draw_idle()
            return
        
        data = self.current_entry['data']
        cursor_time = event.xdata
        
        if cursor_time is not None:
            # Find closest data point
            time_values = data['time']
            closest_idx = np.argmin(np.abs(time_values - cursor_time))
            closest_time = time_values[closest_idx]
            
            # Get values
            displacement_val = data['displacement'][closest_idx]
            velocity_val = data['velocity'][closest_idx]
            velocity_raw_val = data['velocity_raw'][closest_idx]
            acceleration_val = data['acceleration'][closest_idx]
            
            # Update cursor line
            self.cursor_line.set_xdata([closest_time, closest_time])
            self.cursor_line.set_visible(True)
            
            # Update text
            if self.acceleration_visible:
                text_content = f'Time: {closest_time:.3f}s\nDisp: {displacement_val:.3f}m\nVel: {velocity_val:.3f}m/s\nVel (raw): {velocity_raw_val:.3f}m/s\nAcc: {acceleration_val:.3f}m/s²'
            else:
                text_content = f'Time: {closest_time:.3f}s\nDisp: {displacement_val:.3f}m\nVel: {velocity_val:.3f}m/s\nVel (raw): {velocity_raw_val:.3f}m/s'
            
            self.cursor_text.set_text(text_content)
            self.cursor_text.set_visible(True)
            
            self.canvas.draw_idle()


def plot_all_data(all_data):
    """Launch Qt viewer with list selection"""
    app = QApplication.instance()
    if app is None:
        app = QApplication(sys.argv)
    
    viewer = PlotViewer(all_data)
    viewer.show()
    app.exec()