397 lines
14 KiB
Python
397 lines
14 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import time
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from typing import Dict, Optional, Tuple
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
from PyQt5 import QtCore, QtWidgets
|
||
|
|
|
||
|
|
try:
|
||
|
|
import pyqtgraph as pg
|
||
|
|
except Exception as exc:
|
||
|
|
pg = None
|
||
|
|
PYQTGRAPH_IMPORT_ERROR = exc
|
||
|
|
else:
|
||
|
|
PYQTGRAPH_IMPORT_ERROR = None
|
||
|
|
|
||
|
|
|
||
|
|
if pg is None:
|
||
|
|
InteractivePlotViewBox = None
|
||
|
|
else:
|
||
|
|
class InteractivePlotViewBox(pg.ViewBox):
|
||
|
|
def wheelEvent(self, ev, axis=None):
|
||
|
|
modifiers = QtWidgets.QApplication.keyboardModifiers()
|
||
|
|
target_axis = 0 if modifiers & QtCore.Qt.ControlModifier else 1
|
||
|
|
super().wheelEvent(ev, axis=target_axis)
|
||
|
|
|
||
|
|
|
||
|
|
MAX_RENDER_POINTS = 600
|
||
|
|
REFRESH_INTERVAL_MS = 80
|
||
|
|
LEGEND_LIMIT = 8
|
||
|
|
POSITION_Y_RANGE = (-3.141592653589793, 3.141592653589793)
|
||
|
|
VELOCITY_Y_RANGE = (-5.0, 5.0)
|
||
|
|
TIME_WINDOW_OPTIONS = (
|
||
|
|
("1 min", 60.0),
|
||
|
|
("5 min", 300.0),
|
||
|
|
("All", None),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class MultiResolutionSeries:
|
||
|
|
levels: list[list[tuple[float, float]]] = field(default_factory=lambda: [[]])
|
||
|
|
carry: list[list[tuple[float, float]]] = field(default_factory=lambda: [[]])
|
||
|
|
|
||
|
|
def clear(self) -> None:
|
||
|
|
self.levels = [[]]
|
||
|
|
self.carry = [[]]
|
||
|
|
|
||
|
|
def append(self, t_value: float, y_value: float) -> None:
|
||
|
|
self._append_level(0, (t_value, y_value))
|
||
|
|
|
||
|
|
def sample(self, max_points: int, min_time: Optional[float] = None) -> tuple[np.ndarray, np.ndarray]:
|
||
|
|
raw_points = self.levels[0]
|
||
|
|
if not raw_points:
|
||
|
|
return np.empty(0, dtype=float), np.empty(0, dtype=float)
|
||
|
|
|
||
|
|
points = self._slice_points(raw_points, min_time)
|
||
|
|
if not points:
|
||
|
|
return np.empty(0, dtype=float), np.empty(0, dtype=float)
|
||
|
|
|
||
|
|
for level_points in self.levels:
|
||
|
|
sliced_points = self._slice_points(level_points, min_time)
|
||
|
|
if not sliced_points:
|
||
|
|
continue
|
||
|
|
points = sliced_points
|
||
|
|
if len(sliced_points) <= max_points:
|
||
|
|
break
|
||
|
|
|
||
|
|
if len(points) > max_points:
|
||
|
|
step = max(1, len(points) // max_points)
|
||
|
|
points = points[::step]
|
||
|
|
|
||
|
|
sampled = list(points)
|
||
|
|
first_raw = self._slice_points(raw_points, min_time)[0]
|
||
|
|
last_raw = raw_points[-1]
|
||
|
|
if sampled[0][0] > first_raw[0]:
|
||
|
|
sampled.insert(0, first_raw)
|
||
|
|
if sampled[-1][0] < last_raw[0]:
|
||
|
|
sampled.append(last_raw)
|
||
|
|
|
||
|
|
x_values = np.fromiter((point[0] for point in sampled), dtype=float, count=len(sampled))
|
||
|
|
y_values = np.fromiter((point[1] for point in sampled), dtype=float, count=len(sampled))
|
||
|
|
return x_values, y_values
|
||
|
|
|
||
|
|
def _append_level(self, level: int, point: tuple[float, float]) -> None:
|
||
|
|
while len(self.levels) <= level:
|
||
|
|
self.levels.append([])
|
||
|
|
self.carry.append([])
|
||
|
|
|
||
|
|
self.levels[level].append(point)
|
||
|
|
bucket = self.carry[level]
|
||
|
|
bucket.append(point)
|
||
|
|
if len(bucket) < 2:
|
||
|
|
return
|
||
|
|
|
||
|
|
point_a, point_b = bucket
|
||
|
|
bucket.clear()
|
||
|
|
merged = (
|
||
|
|
(point_a[0] + point_b[0]) * 0.5,
|
||
|
|
(point_a[1] + point_b[1]) * 0.5,
|
||
|
|
)
|
||
|
|
self._append_level(level + 1, merged)
|
||
|
|
|
||
|
|
def _slice_points(
|
||
|
|
self,
|
||
|
|
points: list[tuple[float, float]],
|
||
|
|
min_time: Optional[float],
|
||
|
|
) -> list[tuple[float, float]]:
|
||
|
|
if not points:
|
||
|
|
return []
|
||
|
|
if min_time is None or min_time <= points[0][0]:
|
||
|
|
return points
|
||
|
|
|
||
|
|
left = 0
|
||
|
|
right = len(points)
|
||
|
|
while left < right:
|
||
|
|
mid = (left + right) // 2
|
||
|
|
if points[mid][0] < min_time:
|
||
|
|
left = mid + 1
|
||
|
|
else:
|
||
|
|
right = mid
|
||
|
|
|
||
|
|
start = max(0, left - 1)
|
||
|
|
return points[start:]
|
||
|
|
|
||
|
|
|
||
|
|
class PlotPage(QtWidgets.QWidget):
|
||
|
|
def __init__(self, joint_names: list[str], parent=None) -> None:
|
||
|
|
super().__init__(parent)
|
||
|
|
self.joint_names = joint_names
|
||
|
|
self.pos_lines: Dict[str, object] = {}
|
||
|
|
self.vel_lines: Dict[str, object] = {}
|
||
|
|
self.pos_hist: Dict[str, MultiResolutionSeries] = {
|
||
|
|
name: MultiResolutionSeries() for name in joint_names
|
||
|
|
}
|
||
|
|
self.vel_hist: Dict[str, MultiResolutionSeries] = {
|
||
|
|
name: MultiResolutionSeries() for name in joint_names
|
||
|
|
}
|
||
|
|
self._plot_start = time.time()
|
||
|
|
self._elapsed_time = 0.0
|
||
|
|
self._last_plot_time = None
|
||
|
|
self._pending_selected: set[str] = set()
|
||
|
|
self._dirty = False
|
||
|
|
self._last_legend_key: tuple[str, ...] = ()
|
||
|
|
self._color_map = {}
|
||
|
|
self._auto_follow_x = True
|
||
|
|
self._applying_x_range = False
|
||
|
|
self._build_ui()
|
||
|
|
self._refresh_timer = QtCore.QTimer(self)
|
||
|
|
self._refresh_timer.setSingleShot(True)
|
||
|
|
self._refresh_timer.setInterval(REFRESH_INTERVAL_MS)
|
||
|
|
self._refresh_timer.timeout.connect(self._flush_refresh)
|
||
|
|
|
||
|
|
def _build_ui(self) -> None:
|
||
|
|
layout = QtWidgets.QVBoxLayout(self)
|
||
|
|
layout.setContentsMargins(0, 0, 0, 0)
|
||
|
|
layout.setSpacing(4)
|
||
|
|
|
||
|
|
controls = QtWidgets.QHBoxLayout()
|
||
|
|
controls.addWidget(QtWidgets.QLabel("Time Window:"))
|
||
|
|
self.window_combo = QtWidgets.QComboBox()
|
||
|
|
for label, seconds in TIME_WINDOW_OPTIONS:
|
||
|
|
self.window_combo.addItem(label, seconds)
|
||
|
|
self.window_combo.setCurrentIndex(0)
|
||
|
|
self.window_combo.currentIndexChanged.connect(self._on_window_changed)
|
||
|
|
controls.addWidget(self.window_combo)
|
||
|
|
controls.addStretch(1)
|
||
|
|
self.clear_btn = QtWidgets.QPushButton("Clear Plot")
|
||
|
|
self.clear_btn.clicked.connect(self.reset_data)
|
||
|
|
controls.addWidget(self.clear_btn)
|
||
|
|
self.reset_view_btn = QtWidgets.QPushButton("Reset View")
|
||
|
|
self.reset_view_btn.clicked.connect(self._reset_view)
|
||
|
|
controls.addWidget(self.reset_view_btn)
|
||
|
|
layout.addLayout(controls)
|
||
|
|
|
||
|
|
if pg is None:
|
||
|
|
message = QtWidgets.QLabel(f"pyqtgraph unavailable: {PYQTGRAPH_IMPORT_ERROR}")
|
||
|
|
message.setStyleSheet("color: red; padding: 12px;")
|
||
|
|
layout.addWidget(message)
|
||
|
|
return
|
||
|
|
|
||
|
|
pg.setConfigOptions(antialias=False)
|
||
|
|
self.plot_splitter = QtWidgets.QSplitter(QtCore.Qt.Vertical)
|
||
|
|
self.plot_splitter.setChildrenCollapsible(False)
|
||
|
|
self.plot_splitter.setHandleWidth(8)
|
||
|
|
layout.addWidget(self.plot_splitter, 1)
|
||
|
|
|
||
|
|
self.pos_plot_widget = pg.PlotWidget(viewBox=InteractivePlotViewBox())
|
||
|
|
self.pos_plot = self.pos_plot_widget.getPlotItem()
|
||
|
|
self.pos_plot.setTitle("Position (rad)")
|
||
|
|
self.plot_splitter.addWidget(self.pos_plot_widget)
|
||
|
|
|
||
|
|
self.vel_plot_widget = pg.PlotWidget(viewBox=InteractivePlotViewBox())
|
||
|
|
self.vel_plot = self.vel_plot_widget.getPlotItem()
|
||
|
|
self.vel_plot.setTitle("Velocity (rad/s)")
|
||
|
|
self.plot_splitter.addWidget(self.vel_plot_widget)
|
||
|
|
self.plot_splitter.setStretchFactor(0, 1)
|
||
|
|
self.plot_splitter.setStretchFactor(1, 1)
|
||
|
|
|
||
|
|
self.vel_plot.setXLink(self.pos_plot)
|
||
|
|
self.pos_plot.getViewBox().sigXRangeChanged.connect(self._on_plot_x_range_changed)
|
||
|
|
self.vel_plot.getViewBox().sigXRangeChanged.connect(self._on_plot_x_range_changed)
|
||
|
|
|
||
|
|
self._configure_plot(self.pos_plot, "Position", "rad", POSITION_Y_RANGE, show_bottom=False)
|
||
|
|
self._configure_plot(self.vel_plot, "Velocity", "rad/s", VELOCITY_Y_RANGE, show_bottom=True)
|
||
|
|
|
||
|
|
self.pos_legend = self.pos_plot.addLegend(offset=(-10, 10))
|
||
|
|
self.vel_legend = self.vel_plot.addLegend(offset=(-10, 10))
|
||
|
|
self.pos_legend.setVisible(False)
|
||
|
|
self.vel_legend.setVisible(False)
|
||
|
|
self._apply_time_window()
|
||
|
|
|
||
|
|
def _configure_plot(self, plot, label: str, units: str, y_range: tuple[float, float], show_bottom: bool) -> None:
|
||
|
|
plot.showGrid(x=True, y=True, alpha=0.25)
|
||
|
|
plot.setLabel("left", label, units=units)
|
||
|
|
plot.setMouseEnabled(x=True, y=True)
|
||
|
|
plot.setMenuEnabled(False)
|
||
|
|
plot.setYRange(y_range[0], y_range[1], padding=0.0)
|
||
|
|
plot.getAxis("bottom").setStyle(showValues=show_bottom)
|
||
|
|
if show_bottom:
|
||
|
|
plot.setLabel("bottom", "Time", units="s")
|
||
|
|
|
||
|
|
def update_data(
|
||
|
|
self,
|
||
|
|
data: Dict[str, Tuple[float, float]],
|
||
|
|
selected: set[str],
|
||
|
|
refresh: bool = True,
|
||
|
|
) -> None:
|
||
|
|
sample_time = time.time() - self._plot_start
|
||
|
|
if self._last_plot_time is not None and sample_time <= self._last_plot_time:
|
||
|
|
sample_time = self._last_plot_time + 1e-6
|
||
|
|
self._last_plot_time = sample_time
|
||
|
|
self._elapsed_time = max(self._elapsed_time, sample_time)
|
||
|
|
|
||
|
|
for name in selected:
|
||
|
|
pos, vel = data.get(name, (0.0, 0.0))
|
||
|
|
self.pos_hist[name].append(sample_time, pos)
|
||
|
|
self.vel_hist[name].append(sample_time, vel)
|
||
|
|
|
||
|
|
if refresh:
|
||
|
|
self.refresh(selected)
|
||
|
|
|
||
|
|
def refresh(self, selected: set[str]) -> None:
|
||
|
|
self._pending_selected = set(selected)
|
||
|
|
self._dirty = True
|
||
|
|
if not self._refresh_timer.isActive():
|
||
|
|
self._refresh_timer.start()
|
||
|
|
|
||
|
|
def reset_data(self) -> None:
|
||
|
|
for name in self.joint_names:
|
||
|
|
self.pos_hist[name].clear()
|
||
|
|
self.vel_hist[name].clear()
|
||
|
|
self._plot_start = time.time()
|
||
|
|
self._elapsed_time = 0.0
|
||
|
|
self._last_plot_time = None
|
||
|
|
self._pending_selected.clear()
|
||
|
|
self._dirty = False
|
||
|
|
self._last_legend_key = ()
|
||
|
|
self._auto_follow_x = True
|
||
|
|
|
||
|
|
if pg is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
for line in self.pos_lines.values():
|
||
|
|
line.setData([], [])
|
||
|
|
for line in self.vel_lines.values():
|
||
|
|
line.setData([], [])
|
||
|
|
|
||
|
|
self._apply_time_window()
|
||
|
|
self.pos_plot.setYRange(POSITION_Y_RANGE[0], POSITION_Y_RANGE[1], padding=0.0)
|
||
|
|
self.vel_plot.setYRange(VELOCITY_Y_RANGE[0], VELOCITY_Y_RANGE[1], padding=0.0)
|
||
|
|
self._sync_legends(set())
|
||
|
|
|
||
|
|
def _flush_refresh(self) -> None:
|
||
|
|
if not self._dirty:
|
||
|
|
return
|
||
|
|
self._dirty = False
|
||
|
|
self._refresh_plot(self._pending_selected)
|
||
|
|
|
||
|
|
def _refresh_plot(self, selected: set[str]) -> None:
|
||
|
|
if pg is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
min_time = self._window_min_time()
|
||
|
|
|
||
|
|
for name in list(self.pos_lines.keys()):
|
||
|
|
if name not in selected:
|
||
|
|
self.pos_plot.removeItem(self.pos_lines[name])
|
||
|
|
self.vel_plot.removeItem(self.vel_lines[name])
|
||
|
|
del self.pos_lines[name]
|
||
|
|
del self.vel_lines[name]
|
||
|
|
|
||
|
|
for name in selected:
|
||
|
|
if name not in self.pos_lines:
|
||
|
|
color = self._color_for_joint(name)
|
||
|
|
self.pos_lines[name] = self._create_curve(self.pos_plot, color)
|
||
|
|
self.vel_lines[name] = self._create_curve(self.vel_plot, color)
|
||
|
|
|
||
|
|
pos_x, pos_y = self.pos_hist[name].sample(MAX_RENDER_POINTS, min_time=min_time)
|
||
|
|
vel_x, vel_y = self.vel_hist[name].sample(MAX_RENDER_POINTS, min_time=min_time)
|
||
|
|
self.pos_lines[name].setData(pos_x, pos_y)
|
||
|
|
self.vel_lines[name].setData(vel_x, vel_y)
|
||
|
|
|
||
|
|
self._apply_time_window()
|
||
|
|
|
||
|
|
self._sync_legends(selected)
|
||
|
|
|
||
|
|
def _create_curve(self, plot, color):
|
||
|
|
curve = plot.plot([], [], pen=pg.mkPen(color=color, width=1.6), antialias=False)
|
||
|
|
curve.setClipToView(True)
|
||
|
|
curve.setDownsampling(auto=True, method="peak")
|
||
|
|
return curve
|
||
|
|
|
||
|
|
def _color_for_joint(self, name: str):
|
||
|
|
if name not in self._color_map:
|
||
|
|
index = len(self._color_map)
|
||
|
|
self._color_map[name] = pg.intColor(index, hues=max(8, len(self.joint_names)))
|
||
|
|
return self._color_map[name]
|
||
|
|
|
||
|
|
def _sync_legends(self, selected: set[str]) -> None:
|
||
|
|
if pg is None:
|
||
|
|
return
|
||
|
|
|
||
|
|
selection_key = tuple(sorted(selected))
|
||
|
|
if selection_key == self._last_legend_key:
|
||
|
|
return
|
||
|
|
|
||
|
|
self.pos_legend.clear()
|
||
|
|
self.vel_legend.clear()
|
||
|
|
|
||
|
|
show_legend = bool(selected) and len(selected) <= LEGEND_LIMIT
|
||
|
|
self.pos_legend.setVisible(show_legend)
|
||
|
|
self.vel_legend.setVisible(show_legend)
|
||
|
|
if show_legend:
|
||
|
|
for name in selection_key:
|
||
|
|
pos_line = self.pos_lines.get(name)
|
||
|
|
vel_line = self.vel_lines.get(name)
|
||
|
|
if pos_line is not None:
|
||
|
|
self.pos_legend.addItem(pos_line, name)
|
||
|
|
if vel_line is not None:
|
||
|
|
self.vel_legend.addItem(vel_line, name)
|
||
|
|
self._last_legend_key = selection_key
|
||
|
|
|
||
|
|
def _window_min_time(self) -> Optional[float]:
|
||
|
|
window_seconds = self.window_combo.currentData() if hasattr(self, "window_combo") else None
|
||
|
|
if window_seconds is None:
|
||
|
|
return None
|
||
|
|
return max(0.0, self._elapsed_time - float(window_seconds))
|
||
|
|
|
||
|
|
def _apply_time_window(self) -> None:
|
||
|
|
if pg is None or not hasattr(self, "window_combo"):
|
||
|
|
return
|
||
|
|
if not self._auto_follow_x:
|
||
|
|
return
|
||
|
|
|
||
|
|
window_seconds = self.window_combo.currentData()
|
||
|
|
if window_seconds is None:
|
||
|
|
min_time = 0.0
|
||
|
|
max_time = max(1.0, self._elapsed_time)
|
||
|
|
else:
|
||
|
|
duration = float(window_seconds)
|
||
|
|
min_time = max(0.0, self._elapsed_time - duration)
|
||
|
|
max_time = duration if self._elapsed_time < duration else self._elapsed_time
|
||
|
|
max_time = max(min_time + 1.0, max_time)
|
||
|
|
|
||
|
|
self._set_x_range(min_time, max_time)
|
||
|
|
|
||
|
|
def _set_x_range(self, min_time: float, max_time: float) -> None:
|
||
|
|
self._applying_x_range = True
|
||
|
|
try:
|
||
|
|
self.pos_plot.setXRange(min_time, max_time, padding=0.0)
|
||
|
|
self.vel_plot.setXRange(min_time, max_time, padding=0.0)
|
||
|
|
finally:
|
||
|
|
self._applying_x_range = False
|
||
|
|
|
||
|
|
def _on_plot_x_range_changed(self, *_args) -> None:
|
||
|
|
if self._applying_x_range:
|
||
|
|
return
|
||
|
|
self._auto_follow_x = False
|
||
|
|
|
||
|
|
def _reset_view(self) -> None:
|
||
|
|
if pg is None:
|
||
|
|
return
|
||
|
|
self._auto_follow_x = True
|
||
|
|
self._apply_time_window()
|
||
|
|
self.pos_plot.setYRange(POSITION_Y_RANGE[0], POSITION_Y_RANGE[1], padding=0.0)
|
||
|
|
self.vel_plot.setYRange(VELOCITY_Y_RANGE[0], VELOCITY_Y_RANGE[1], padding=0.0)
|
||
|
|
|
||
|
|
def _on_window_changed(self, _index: int) -> None:
|
||
|
|
self._auto_follow_x = True
|
||
|
|
active = self._pending_selected or set(self.pos_lines.keys())
|
||
|
|
if active:
|
||
|
|
self.refresh(active)
|
||
|
|
else:
|
||
|
|
self._apply_time_window()
|