diff --git a/examples/lsl_viewer.py b/examples/lsl_viewer.py index ca7232e..e78c925 100644 --- a/examples/lsl_viewer.py +++ b/examples/lsl_viewer.py @@ -9,6 +9,7 @@ python lsl_viewer.py --name MyStream # resolve stream by name python lsl_viewer.py --type EEG # resolve stream by type python lsl_viewer.py --visible 32 # show 32 channels at a time + python lsl_viewer.py --scatter # scatter/heatmap (needs channel locations) Requires: pylsl, phosphor pip install pylsl phosphor @@ -23,7 +24,7 @@ from PySide6.QtCore import QTimer from PySide6.QtWidgets import QApplication -from phosphor import SweepConfig, SweepWidget +from phosphor import ScatterConfig, ScatterWidget, SweepConfig, SweepWidget # Map pylsl channel formats to numpy dtypes _LSL_DTYPES = { @@ -57,12 +58,69 @@ def resolve_stream(name: str | None, type_: str | None) -> pylsl.StreamInfo: return results[0] +def parse_channel_info( + info: pylsl.StreamInfo, +) -> tuple[list[str] | None, np.ndarray | None]: + """Parse channel labels and locations from LSL stream info. + + Returns ``(labels, positions)`` where *positions* is ``(n_channels, 3)`` + float32 or ``None`` if no locations are present. + """ + n_ch = info.channel_count() + chans = info.desc().child("channels") + if chans.empty(): + return None, None + + labels: list[str] = [] + positions: list[list[float]] = [] + has_locations = False + + ch_elem = chans.first_child() + while not ch_elem.empty(): + # Label + label_val = ch_elem.child("label").child_value() + labels.append(label_val if label_val else "") + + # Location — valvalval + loc_elem = ch_elem.child("location") + if not loc_elem.empty(): + x_val = loc_elem.child("X").child_value() + y_val = loc_elem.child("Y").child_value() + z_val = loc_elem.child("Z").child_value() + x = float(x_val) if x_val else 0.0 + y = float(y_val) if y_val else 0.0 + z = float(z_val) if z_val else 0.0 + positions.append([x, y, z]) + if x != 0.0 or y != 0.0 or z != 0.0: + has_locations = True + else: + positions.append([0.0, 0.0, 0.0]) + + ch_elem = ch_elem.next_sibling() + + if len(labels) != n_ch: + return None, None + + # Fill empty labels with channel indices + for i, lbl in enumerate(labels): + if not lbl: + labels[i] = f"Ch {i}" + + pos_array = np.array(positions, dtype=np.float32) if has_locations else None + return labels, pos_array + + def main(): parser = argparse.ArgumentParser(description="LSL Viewer (phosphor)") parser.add_argument("--name", type=str, default=None, help="Resolve stream by name") parser.add_argument("--type", type=str, default=None, help="Resolve stream by type") parser.add_argument("--dur", type=float, default=2.0, help="Display duration in seconds") parser.add_argument("--visible", type=int, default=None, help="Visible channels (default: all)") + parser.add_argument( + "--scatter", + action="store_true", + help="Use scatter/heatmap view (requires channel locations in stream metadata)", + ) args = parser.parse_args() app = QApplication(sys.argv) @@ -79,21 +137,46 @@ def main(): max_buflen=int(max(args.dur * 2, 1)), processing_flags=pylsl.proc_clocksync | pylsl.proc_dejitter, ) + inlet.open_stream() + + # Retrieve full stream info (includes description XML with channel metadata) + full_info = inlet.info() + channel_labels, channel_positions = parse_channel_info(full_info) + + if channel_labels: + print(f"Channel labels: {channel_labels[0]} … {channel_labels[-1]}") + if channel_positions is not None: + print(f"Channel locations: found for {n_channels} channels") # Pre-allocate pull buffer matching the stream's native dtype max_samples = max(1024, math.ceil(srate / 30) * 2) stream_dtype = _LSL_DTYPES.get(info.channel_format(), np.float32) pull_buffer = np.empty((max_samples, n_channels), dtype=stream_dtype, order="C") - config = SweepConfig( - n_channels=n_channels, - srate=srate, - display_dur=args.dur, - n_visible=n_visible, - ) - widget = SweepWidget(config) + use_scatter = args.scatter + if use_scatter and channel_positions is None: + print("Warning: --scatter requested but no channel locations found; falling back to sweep view") + use_scatter = False + + if use_scatter: + config = ScatterConfig( + positions=channel_positions, + channel_labels=channel_labels, + ) + widget = ScatterWidget(config) + widget.resize(800, 800) + else: + config = SweepConfig( + n_channels=n_channels, + srate=srate, + display_dur=args.dur, + n_visible=n_visible, + channel_labels=channel_labels, + ) + widget = SweepWidget(config) + widget.resize(1200, 800) + widget.setWindowTitle(f"LSL: {info.name()} ({n_channels}ch @ {srate}Hz)") - widget.resize(1200, 800) widget.show() def pull_and_push(): diff --git a/examples/scatter_demo.py b/examples/scatter_demo.py new file mode 100644 index 0000000..77ddfd8 --- /dev/null +++ b/examples/scatter_demo.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python +""" +Scatter Demo — simulated scalp-map heatmap with the ScatterWidget. + +Generates electrode positions on concentric rings (mimicking a 10-20 EEG +montage layout) and drives them with a rotating spatial wave so color and +size ripple outward from a wandering hotspot. + +Usage: + python scatter_demo.py # 64 electrodes, color+size + python scatter_demo.py --channels 128 # more electrodes + python scatter_demo.py --no-size # color only + python scatter_demo.py --cmap plasma # different colormap + python scatter_demo.py --fixed-range 0 1 # fixed vmin/vmax + +Requires: phosphor + pip install phosphor +""" + +import argparse +import sys + +import numpy as np +from PySide6.QtCore import QTimer +from PySide6.QtWidgets import QApplication + +from phosphor import ScatterConfig, ScatterWidget + + +def make_scalp_positions(n: int) -> np.ndarray: + """Generate electrode positions on concentric rings like a scalp map. + + Returns an (n, 2) array with positions roughly within a unit circle. + """ + positions = [] + + # Place one electrode at the center (Cz-like) + positions.append([0.0, 0.0]) + remaining = n - 1 + + ring = 1 + while remaining > 0: + radius = ring * 0.25 + # More electrodes on outer rings + count = min(remaining, 6 * ring) + angles = np.linspace(0, 2 * np.pi, count, endpoint=False) + # Offset alternate rings for nicer packing + angles += (ring % 2) * np.pi / count + for a in angles: + positions.append([radius * np.cos(a), radius * np.sin(a)]) + remaining -= count + ring += 1 + + return np.array(positions[:n], dtype=np.float32) + + +def make_electrode_labels(n: int) -> list[str]: + """Generate electrode labels (E0, E1, …).""" + return [f"E{i}" for i in range(n)] + + +def main(): + parser = argparse.ArgumentParser(description="Phosphor scatter/heatmap demo") + parser.add_argument("--channels", type=int, default=64, help="Number of electrodes") + parser.add_argument("--cmap", type=str, default="viridis", help="Colormap name") + parser.add_argument("--no-size", action="store_true", help="Disable size modulation") + parser.add_argument( + "--fixed-range", + type=float, + nargs=2, + metavar=("VMIN", "VMAX"), + default=None, + help="Fixed value range (default: autoscale)", + ) + parser.add_argument("--marker-size", type=float, default=14.0, help="Base marker size") + parser.add_argument("--fps", type=float, default=60.0, help="Data push rate (Hz)") + args = parser.parse_args() + + app = QApplication(sys.argv) + + n_ch = args.channels + positions = make_scalp_positions(n_ch) + labels = make_electrode_labels(n_ch) + + vmin = args.fixed_range[0] if args.fixed_range else None + vmax = args.fixed_range[1] if args.fixed_range else None + + config = ScatterConfig( + positions=positions, + cmap=args.cmap, + modulate_color=True, + modulate_size=not args.no_size, + marker_size=args.marker_size, + size_range=(4.0, 24.0), + vmin=vmin, + vmax=vmax, + channel_labels=labels, + ) + + widget = ScatterWidget(config) + widget.setWindowTitle(f"Phosphor Scatter Demo — {n_ch} electrodes") + widget.resize(800, 800) + widget.show() + + # --- Simulation state --- + t = 0.0 + dt = 1.0 / args.fps + + def push_data(): + nonlocal t + + # Hotspot wanders in a figure-8 (Lissajous) pattern + hx = 0.5 * np.sin(t * 0.7) + hy = 0.5 * np.sin(t * 1.1) + + # Distance from each electrode to the hotspot + dx = positions[:, 0] - hx + dy = positions[:, 1] - hy + dist = np.sqrt(dx**2 + dy**2) + + # Gaussian blob centered on the hotspot, plus a small amount of noise + values = np.exp(-(dist**2) / 0.15) + np.random.randn(n_ch).astype(np.float32) * 0.03 + + widget.push_data(values.astype(np.float32)) + t += dt + + timer = QTimer() + timer.timeout.connect(push_data) + timer.start(max(1, int(1000 / args.fps))) + + app.aboutToQuit.connect(timer.stop) + sys.exit(app.exec()) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 0c7adcc..d6313f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,10 +8,9 @@ readme = "README.md" requires-python = ">=3.11" dynamic = ["version"] dependencies = [ + "fastplotlib>=0.6.1", "numpy>=2.4.2", "pyside6>=6.10.2", - "rendercanvas>=2.6.1", - "wgpu>=0.29.0", ] [dependency-groups] diff --git a/src/phosphor/__init__.py b/src/phosphor/__init__.py index 0cdea95..9c77df6 100644 --- a/src/phosphor/__init__.py +++ b/src/phosphor/__init__.py @@ -2,11 +2,14 @@ from .__version__ import __version__ as __version__ from .channel_plot import ChannelPlotWidget +from .scatter_widget import ScatterConfig, ScatterWidget from .spectrum_widget import SpectrumConfig, SpectrumWidget from .sweep_widget import SweepConfig, SweepWidget __all__ = [ "ChannelPlotWidget", + "ScatterConfig", + "ScatterWidget", "SpectrumConfig", "SpectrumWidget", "SweepConfig", diff --git a/src/phosphor/__main__.py b/src/phosphor/__main__.py index caa1073..67adf3a 100644 --- a/src/phosphor/__main__.py +++ b/src/phosphor/__main__.py @@ -7,12 +7,19 @@ from PySide6.QtCore import QTimer from PySide6.QtWidgets import QApplication -from phosphor import SpectrumConfig, SpectrumWidget, SweepConfig, SweepWidget +from phosphor import ( + ScatterConfig, + ScatterWidget, + SpectrumConfig, + SpectrumWidget, + SweepConfig, + SweepWidget, +) def main(): parser = argparse.ArgumentParser(description="Phosphor renderer demo") - parser.add_argument("--mode", choices=["sweep", "spectrum"], default="sweep", help="Render mode") + parser.add_argument("--mode", choices=["sweep", "spectrum", "scatter"], default="sweep", help="Render mode") parser.add_argument("--channels", type=int, default=128, help="Total channels") parser.add_argument("--srate", type=float, default=30000.0, help="Sample rate (Hz)") parser.add_argument("--dur", type=float, default=2.0, help="Display duration (s)") @@ -27,7 +34,31 @@ def main(): chunk_size = max(1, int(args.srate / 60)) # ~500 samples per timer tick sample_counter = 0 - if args.mode == "spectrum": + if args.mode == "scatter": + n_ch = args.channels + # Generate positions on a unit circle + angles = np.linspace(0, 2 * np.pi, n_ch, endpoint=False) + positions = np.column_stack([np.cos(angles), np.sin(angles)]).astype(np.float32) + labels = [f"E{i}" for i in range(n_ch)] + config = ScatterConfig( + positions=positions, + cmap="viridis", + modulate_color=True, + modulate_size=True, + channel_labels=labels, + ) + widget = ScatterWidget(config) + widget.setWindowTitle("Phosphor Demo — Scatter") + phase = 0.0 + + def push_chunk(): + nonlocal phase + # Rotating wave around the circle + values = (np.sin(angles + phase) * 0.5 + 0.5).astype(np.float32) + widget.push_data(values) + phase += 0.05 + + elif args.mode == "spectrum": fft_size = int(args.srate) n_bins = fft_size // 2 config = SpectrumConfig( diff --git a/src/phosphor/channel_plot.py b/src/phosphor/channel_plot.py index 7d642b0..4e08c28 100644 --- a/src/phosphor/channel_plot.py +++ b/src/phosphor/channel_plot.py @@ -1,10 +1,10 @@ -"""Base class for multichannel plot widgets with shared canvas, scrolling, and key handling.""" +"""Base class for multichannel plot widgets backed by fastplotlib.""" from __future__ import annotations -from PySide6.QtCore import QEvent, Qt +import fastplotlib as fpl +from PySide6.QtCore import Qt from PySide6.QtWidgets import QLabel, QToolTip, QVBoxLayout, QWidget -from rendercanvas.qt import QRenderWidget from .constants import CHANNEL_COLORS @@ -12,20 +12,21 @@ class ChannelPlotWidget(QWidget): - """Base widget for GPU-rendered multichannel plots. + """Base widget for fastplotlib-rendered multichannel plots. - Provides canvas setup, channel scrolling (↑↓ PgUp/PgDn [ ]), - y-scale keys (- = A), tooltip on hover, and range label overlay. + Provides canvas setup, channel scrolling (scroll / ↑↓ / PgUp/PgDn / [ ]), + amplitude zoom (Shift+scroll / - = A), and range label overlay. Subclass contract: - - Set ``self._buffer`` to a buffer object before calling ``_init_rendering()``. - Expected interface: ``.n_visible``, ``.n_channels``, ``.channel_offset``, - ``.set_channel_offset(int)``, ``.set_n_visible(int)``, - ``.adjust_y_scale(float)``, ``.toggle_autoscale()``. + + - Set ``self._buffer`` before calling ``_init_rendering()``. + Expected interface: ``.n_visible``, ``.n_channels``, ``.channel_offset``, + ``.set_channel_offset(int)``, ``.set_n_visible(int)``. - Add axis widgets to ``self.layout()`` after ``super().__init__``. - - Implement ``_draw_frame(self)`` (called by rendercanvas scheduler). - - Override ``_handle_key(key)`` to intercept subclass keys, calling - ``super()._handle_key(key)`` for common keys. + - Implement ``_update_graphics()`` (called every frame via animation callback). + - Override ``_on_ctrl_scroll(delta)`` for time/freq zoom. + - Override ``_on_key_down(key)`` for subclass-specific keys, calling + ``super()._on_key_down(key)`` for common keys. """ def __init__( @@ -39,22 +40,22 @@ def __init__( super().__init__(parent) self._n_channels = n_channels self._channel_labels = channel_labels + self._autoscale_enabled = True # Layout layout = QVBoxLayout(self) layout.setContentsMargins(0, 0, 0, 0) - # Render canvas - self.canvas = QRenderWidget(parent=self) - layout.addWidget(self.canvas) + # fastplotlib figure with a single subplot + self._figure = fpl.Figure() + self._subplot = self._figure[0, 0] - # Keyboard focus & mouse tracking: intercept events on the canvas - self.canvas.setFocusPolicy(Qt.FocusPolicy.StrongFocus) - self.canvas.setMouseTracking(True) - self.canvas.installEventFilter(self) + # Get the Qt widget from fastplotlib and embed it + self._fpl_widget = self._figure.show() + layout.addWidget(self._fpl_widget) - # Channel range overlay label (parented to canvas so it floats on top) - self._range_label = QLabel(self.canvas) + # Channel range overlay label (parented to fpl widget so it floats on top) + self._range_label = QLabel(self._fpl_widget) self._range_label.setStyleSheet( "background: rgba(25,25,30,200); color: #b4b4b4;" " padding: 2px 6px; font-size: 9pt;" @@ -67,90 +68,145 @@ def __init__( self._buffer = None # ------------------------------------------------------------------ - # Subclass hook: call after buffer + renderer are ready + # Subclass hook: call after buffer is ready # ------------------------------------------------------------------ def _init_rendering(self) -> None: - """Start continuous rendering and initialize the range label. + """Start rendering and register event handlers. - Call from subclass ``__init__`` after ``self._buffer`` and any - renderer are fully constructed. + Call from subclass ``__init__`` after ``self._buffer`` is set and + initial graphics are created via ``_setup_graphics()``. """ - from .constants import DEFAULT_MAX_FPS + # Disable built-in pan/zoom controller, axes, and default title + self._subplot.controller = None + self._subplot.axes.visible = False + self._subplot.title.visible = False + + # Register fpl event handlers on the subplot's pygfx renderer + renderer = self._subplot.renderer + renderer.add_event_handler(self._on_key_down_event, "key_down") + renderer.add_event_handler(self._on_wheel_event, "wheel") + renderer.add_event_handler(self._on_pointer_move_event, "pointer_move") + + # Register animation callback (wrap in lambda to avoid getfullargspec + # issue with bound methods under `from __future__ import annotations`) + self._figure.add_animations(lambda: self._animation_callback()) - self.canvas.set_update_mode("continuous", max_fps=DEFAULT_MAX_FPS) - self.canvas.request_draw(self._draw_frame) self._update_range_label() # ------------------------------------------------------------------ # Subclass must implement # ------------------------------------------------------------------ - def _draw_frame(self) -> None: + def _update_graphics(self) -> None: + """Update LineStack data each frame. Called from animation callback.""" raise NotImplementedError + def _on_ctrl_scroll(self, delta: float) -> None: + """Handle Ctrl+scroll for time/freq zoom. Override in subclass.""" + # ------------------------------------------------------------------ - # Event filter (keyboard + mouse) + # Animation callback # ------------------------------------------------------------------ - def eventFilter(self, obj, event): - if obj is self.canvas: - if event.type() == QEvent.Type.KeyPress: - self._handle_key(event.key()) - return True - if event.type() == QEvent.Type.MouseMove: - self._handle_mouse_move(event) - return False # don't swallow — let canvas process too - return super().eventFilter(obj, event) - - def resizeEvent(self, event) -> None: - super().resizeEvent(event) - self._update_range_label() + def _animation_callback(self) -> None: + self._update_graphics() + if self._autoscale_enabled: + self._subplot.auto_scale(maintain_aspect=False, zoom=1.0) # ------------------------------------------------------------------ - # Keyboard controls (shared) + # Event handlers (fpl native events) # ------------------------------------------------------------------ - def _handle_key(self, key: int) -> None: + def _on_key_down_event(self, event) -> None: + self._on_key_down(event.key) + + def _on_key_down(self, key: str) -> None: + """Handle keyboard events. Override in subclass, call super for common keys.""" buf = self._buffer - if key == Qt.Key.Key_Up: + if key == "ArrowUp": buf.set_channel_offset(buf.channel_offset - 1) - elif key == Qt.Key.Key_Down: + elif key == "ArrowDown": buf.set_channel_offset(buf.channel_offset + 1) - elif key == Qt.Key.Key_PageUp: + elif key == "PageUp": buf.set_channel_offset(buf.channel_offset - buf.n_visible) - elif key == Qt.Key.Key_PageDown: + elif key == "PageDown": buf.set_channel_offset(buf.channel_offset + buf.n_visible) - elif key == Qt.Key.Key_BracketLeft: + elif key == "[": buf.set_n_visible(max(1, buf.n_visible // 2)) - elif key == Qt.Key.Key_BracketRight: + elif key == "]": buf.set_n_visible(min(buf.n_channels, buf.n_visible * 2)) - elif key == Qt.Key.Key_Minus: - buf.adjust_y_scale(0.8) # zoom out - elif key == Qt.Key.Key_Equal: - buf.adjust_y_scale(1.25) # zoom in + elif key == "-": + self._zoom_amplitude(0.8) + elif key == "=": + self._zoom_amplitude(1.25) - elif key == Qt.Key.Key_A: - buf.toggle_autoscale() + elif key in ("a", "A"): + self._autoscale_enabled = not self._autoscale_enabled + if self._autoscale_enabled: + self._subplot.auto_scale(maintain_aspect=False, zoom=1.0) self._update_range_label() + def _on_wheel_event(self, event) -> None: + delta = event.dy + + if "Control" in getattr(event, "modifiers", ()): + # Ctrl+scroll → time/freq zoom (subclass hook) + self._on_ctrl_scroll(delta) + elif "Shift" in getattr(event, "modifiers", ()): + # Shift+scroll → amplitude zoom + factor = 1.1 if delta > 0 else 0.9 + self._zoom_amplitude(factor) + else: + # Unmodified scroll → channel scroll + buf = self._buffer + step = 1 if delta < 0 else -1 + buf.set_channel_offset(buf.channel_offset + step) + self._update_range_label() + + def _on_pointer_move_event(self, event) -> None: + self._handle_mouse_move(event) + + # ------------------------------------------------------------------ + # Amplitude zoom helper + # ------------------------------------------------------------------ + + def _zoom_amplitude(self, factor: float) -> None: + """Adjust camera y-scale and disable autoscale.""" + camera = self._subplot.camera + camera.world.scale_y *= factor + self._autoscale_enabled = False + # ------------------------------------------------------------------ # Mouse hover tooltip # ------------------------------------------------------------------ def _handle_mouse_move(self, event) -> None: - h = self.canvas.height() - if h < 1: + if self._line_stack is None: return buf = self._buffer - mouse_y = event.position().y() - ndc_y = 1.0 - 2.0 * (mouse_y / h) - ch_index = int((1.0 - ndc_y) * buf.n_visible / 2.0) - ch_index = max(0, min(ch_index, buf.n_visible - 1)) + + # Convert screen position to world coordinates + world = self._subplot.map_screen_to_world(event) + if world is None: + return + wy = float(world[1]) + + # Find the nearest line by comparing world Y positions + best_idx = 0 + best_dist = float("inf") + for i in range(buf.n_visible): + line_y = float(self._line_stack[i].world_object.world.position[1]) + dist = abs(wy - line_y) + if dist < best_dist: + best_dist = dist + best_idx = i + + ch_index = best_idx abs_ch = buf.channel_offset + ch_index labels = self._channel_labels @@ -159,7 +215,9 @@ def _handle_mouse_move(self, event) -> None: rgba = CHANNEL_COLORS[ch_index % len(CHANNEL_COLORS)] hex_color = f"#{int(rgba[0]*255):02x}{int(rgba[1]*255):02x}{int(rgba[2]*255):02x}" html = f'\u25a0 {label}' - QToolTip.showText(event.globalPosition().toPoint(), html, self.canvas) + from PySide6.QtCore import QPoint + + QToolTip.showText(self._fpl_widget.mapToGlobal(QPoint(int(event.x), int(event.y))), html, self._fpl_widget) # ------------------------------------------------------------------ # Channel range overlay @@ -174,7 +232,11 @@ def _update_range_label(self) -> None: total = buf.n_channels self._range_label.setText(f"Ch {first}\u2013{last} / {total}") self._range_label.adjustSize() - # Position at bottom-left of canvas + # Position at bottom-left of fpl widget margin = 4 - y = self.canvas.height() - self._range_label.height() - margin + y = self._fpl_widget.height() - self._range_label.height() - margin self._range_label.move(margin, max(0, y)) + + def resizeEvent(self, event) -> None: + super().resizeEvent(event) + self._update_range_label() diff --git a/src/phosphor/constants.py b/src/phosphor/constants.py index ade0ed6..857d326 100644 --- a/src/phosphor/constants.py +++ b/src/phosphor/constants.py @@ -6,9 +6,6 @@ DEFAULT_MAX_FPS = 60 CURSOR_GAP_COLUMNS = 5 -AUTOSCALE_TIME_CONSTANT = 2.0 -AUTOSCALE_N_SIGMA = 3.0 - BG_COLOR = (0.10, 0.10, 0.12, 1.0) CURSOR_COLOR = (0.25, 0.25, 0.28, 0.85) diff --git a/src/phosphor/gpu_renderer.py b/src/phosphor/gpu_renderer.py deleted file mode 100644 index b8fb966..0000000 --- a/src/phosphor/gpu_renderer.py +++ /dev/null @@ -1,312 +0,0 @@ -"""wgpu device, pipeline, buffer management, and draw calls.""" - -import struct - -import wgpu - -from .constants import BG_COLOR, CURSOR_COLOR, CURSOR_GAP_COLUMNS -from .shader import CURSOR_SHADER, SWEEP_SHADER - - -class GPURenderer: - def __init__(self, canvas): - self.canvas = canvas - self.adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance") - self.device = self.adapter.request_device_sync() - self.context = canvas.get_context("wgpu") - self.texture_format = self.context.get_preferred_format(self.adapter) - self.context.configure(device=self.device, format=self.texture_format) - - self._n_visible = 0 - self._n_columns = 0 - self._buf_version = -1 - self._initialized = False - - # GPU resources (created in setup) - self.data_buffer = None - self.channel_params_buffer = None - self.uniforms_buffer = None - self.cursor_uniforms_buffer = None - self.sweep_pipeline = None - self.cursor_pipeline = None - self.sweep_bind_group = None - self.cursor_bind_group = None - - def needs_setup(self, n_visible: int, n_columns: int, buf_version: int) -> bool: - return ( - not self._initialized - or n_visible != self._n_visible - or n_columns != self._n_columns - or buf_version != self._buf_version - ) - - def setup(self, n_visible: int, n_columns: int, buf_version: int) -> None: - """Create or recreate all GPU resources for the given dimensions.""" - self._n_visible = n_visible - self._n_columns = n_columns - self._buf_version = buf_version - - device = self.device - - # --- Buffers --- - data_size = n_columns * n_visible * 2 * 4 # float32 - self.data_buffer = device.create_buffer( - size=max(data_size, 4), - usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST, - ) - - params_size = n_visible * 8 * 4 # 8 floats per channel - self.channel_params_buffer = device.create_buffer( - size=max(params_size, 4), - usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST, - ) - - self.uniforms_buffer = device.create_buffer( - size=32, - usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST, - ) - - self.cursor_uniforms_buffer = device.create_buffer( - size=32, - usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST, - ) - - # --- Sweep pipeline --- - sweep_shader = device.create_shader_module(code=SWEEP_SHADER) - - sweep_bgl = device.create_bind_group_layout( - entries=[ - { - "binding": 0, - "visibility": wgpu.ShaderStage.VERTEX, - "buffer": { - "type": wgpu.BufferBindingType.read_only_storage, - }, - }, - { - "binding": 1, - "visibility": wgpu.ShaderStage.VERTEX, - "buffer": { - "type": wgpu.BufferBindingType.read_only_storage, - }, - }, - { - "binding": 2, - "visibility": wgpu.ShaderStage.VERTEX, - "buffer": { - "type": wgpu.BufferBindingType.uniform, - }, - }, - ] - ) - - sweep_layout = device.create_pipeline_layout(bind_group_layouts=[sweep_bgl]) - - self.sweep_pipeline = device.create_render_pipeline( - layout=sweep_layout, - vertex={ - "module": sweep_shader, - "entry_point": "vs_main", - "buffers": [], - }, - primitive={ - "topology": wgpu.PrimitiveTopology.line_strip, - }, - fragment={ - "module": sweep_shader, - "entry_point": "fs_main", - "targets": [ - { - "format": self.texture_format, - "blend": { - "color": { - "src_factor": wgpu.BlendFactor.src_alpha, - "dst_factor": wgpu.BlendFactor.one_minus_src_alpha, - "operation": wgpu.BlendOperation.add, - }, - "alpha": { - "src_factor": wgpu.BlendFactor.one, - "dst_factor": wgpu.BlendFactor.one_minus_src_alpha, - "operation": wgpu.BlendOperation.add, - }, - }, - } - ], - }, - ) - - self.sweep_bind_group = device.create_bind_group( - layout=sweep_bgl, - entries=[ - { - "binding": 0, - "resource": { - "buffer": self.data_buffer, - "offset": 0, - "size": self.data_buffer.size, - }, - }, - { - "binding": 1, - "resource": { - "buffer": self.channel_params_buffer, - "offset": 0, - "size": self.channel_params_buffer.size, - }, - }, - { - "binding": 2, - "resource": { - "buffer": self.uniforms_buffer, - "offset": 0, - "size": self.uniforms_buffer.size, - }, - }, - ], - ) - - # --- Cursor pipeline --- - cursor_shader = device.create_shader_module(code=CURSOR_SHADER) - - cursor_bgl = device.create_bind_group_layout( - entries=[ - { - "binding": 0, - "visibility": wgpu.ShaderStage.VERTEX | wgpu.ShaderStage.FRAGMENT, - "buffer": { - "type": wgpu.BufferBindingType.uniform, - }, - }, - ] - ) - - cursor_layout = device.create_pipeline_layout(bind_group_layouts=[cursor_bgl]) - - self.cursor_pipeline = device.create_render_pipeline( - layout=cursor_layout, - vertex={ - "module": cursor_shader, - "entry_point": "vs_cursor", - "buffers": [], - }, - primitive={ - "topology": wgpu.PrimitiveTopology.triangle_list, - }, - fragment={ - "module": cursor_shader, - "entry_point": "fs_cursor", - "targets": [ - { - "format": self.texture_format, - "blend": { - "color": { - "src_factor": wgpu.BlendFactor.src_alpha, - "dst_factor": wgpu.BlendFactor.one_minus_src_alpha, - "operation": wgpu.BlendOperation.add, - }, - "alpha": { - "src_factor": wgpu.BlendFactor.one, - "dst_factor": wgpu.BlendFactor.one_minus_src_alpha, - "operation": wgpu.BlendOperation.add, - }, - }, - } - ], - }, - ) - - self.cursor_bind_group = device.create_bind_group( - layout=cursor_bgl, - entries=[ - { - "binding": 0, - "resource": { - "buffer": self.cursor_uniforms_buffer, - "offset": 0, - "size": self.cursor_uniforms_buffer.size, - }, - }, - ], - ) - - self._initialized = True - - def update_and_draw(self, buf) -> None: - """Upload data from buffer to GPU and render one frame.""" - if self.needs_setup(buf.n_visible, buf.n_columns, buf.version): - self.setup(buf.n_visible, buf.n_columns, buf.version) - # Full upload after setup - data = buf.get_gpu_data() - self.device.queue.write_buffer(self.data_buffer, 0, data.tobytes()) - else: - # Incremental upload - result = buf.get_dirty_gpu_data() - if result is not None: - data, col_start, n_cols = result - byte_offset = col_start * buf.n_visible * 2 * 4 - self.device.queue.write_buffer(self.data_buffer, byte_offset, data.tobytes()) - - # Upload channel params - params = buf.get_channel_params() - self.device.queue.write_buffer(self.channel_params_buffer, 0, params.tobytes()) - - # Upload sweep uniforms - y_scale = buf.y_scale - uniforms = struct.pack( - " float: + if self._fixed_vmin is not None: + return self._fixed_vmin + return self._ew_min if self._ew_min is not None else 0.0 + + @property + def vmax(self) -> float: + if self._fixed_vmax is not None: + return self._fixed_vmax + return self._ew_max if self._ew_max is not None else 1.0 + + def push_data(self, data: np.ndarray) -> None: + """Push scalar values. Shape ``(n_channels,)`` or ``(n_samples, n_channels)``.""" + if data.size == 0: + return + with self._lock: + if data.ndim == 1: + if data.shape[0] != self.n_channels: + return + finite = np.isfinite(data) + if finite.any(): + vals = np.where(finite, data.astype(np.float64), 0.0) + self._accum += vals + self._per_ch_count += finite.astype(np.int64) + self._dirty = True + elif data.ndim == 2: + if data.shape[1] != self.n_channels: + return + finite = np.isfinite(data) + if finite.any(): + vals = np.where(finite, data.astype(np.float64), 0.0) + self._accum += vals.sum(axis=0) + self._per_ch_count += finite.astype(np.int64).sum(axis=0) + self._dirty = True + + def consume(self) -> np.ndarray | None: + """Return mean of accumulated data since last consume, or None if clean. + + NaN channels (no finite samples since last consume) are preserved as + NaN in the output. Updates EWMA vmin/vmax from finite values only. + """ + with self._lock: + if not self._dirty: + return None + has_data = self._per_ch_count > 0 + if not has_data.any(): + self._accum[:] = 0.0 + self._per_ch_count[:] = 0 + self._dirty = False + return None + + mean = np.full(self.n_channels, np.nan, dtype=np.float32) + mean[has_data] = (self._accum[has_data] / self._per_ch_count[has_data]).astype(np.float32) + self._accum[:] = 0.0 + self._per_ch_count[:] = 0 + self._dirty = False + self._current[:] = np.nan_to_num(mean, nan=0.0) + + # Update EWMA for auto-ranging (finite values only) + finite_vals = mean[np.isfinite(mean)] + if finite_vals.size > 0: + frame_min = float(finite_vals.min()) + frame_max = float(finite_vals.max()) + a = self._alpha + if self._ew_min is None: + self._ew_min = frame_min + self._ew_max = frame_max + else: + self._ew_min = self._ew_min * (1 - a) + frame_min * a + self._ew_max = self._ew_max * (1 - a) + frame_max * a + + return mean diff --git a/src/phosphor/scatter_widget.py b/src/phosphor/scatter_widget.py new file mode 100644 index 0000000..05eb0ff --- /dev/null +++ b/src/phosphor/scatter_widget.py @@ -0,0 +1,172 @@ +"""GPU-accelerated real-time scatter/heatmap widget.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import fastplotlib as fpl +import numpy as np +from PySide6.QtCore import QPoint, Qt +from PySide6.QtWidgets import QLabel, QToolTip, QVBoxLayout, QWidget + +from .scatter_buffer import ScatterBuffer + +__all__ = ["ScatterConfig", "ScatterWidget"] + + +@dataclass +class ScatterConfig: + positions: np.ndarray # (n_channels, 2) or (n_channels, 3) + cmap: str = "viridis" + modulate_color: bool = True + modulate_size: bool = False + marker_size: float = 10.0 # base marker size (screen pixels) + size_range: tuple[float, float] = (4.0, 20.0) + vmin: float | None = None + vmax: float | None = None + channel_labels: list[str] | None = None + + +class ScatterWidget(QWidget): + """Embeddable QWidget that renders a GPU-accelerated scatter heatmap. + + Each channel has a fixed 2D position; incoming scalar data modulates the + color and/or size of the marker at that position. + + Usage:: + + widget = ScatterWidget(ScatterConfig( + positions=electrode_positions, # (n_ch, 2) + cmap="viridis", + modulate_color=True, + )) + widget.show() + widget.push_data(values) # (n_channels,) + """ + + def __init__(self, config: ScatterConfig, parent: QWidget | None = None): + super().__init__(parent) + self._config = config + n_channels = config.positions.shape[0] + + # Use only xy for rendering + self._positions_2d = config.positions[:, :2].astype(np.float32) + + # Buffer + self._buffer = ScatterBuffer( + n_channels=n_channels, + vmin=config.vmin, + vmax=config.vmax, + ) + + # Layout + layout = QVBoxLayout(self) + layout.setContentsMargins(0, 0, 0, 0) + + # fastplotlib figure + self._figure = fpl.Figure() + self._subplot = self._figure[0, 0] + self._fpl_widget = self._figure.show() + layout.addWidget(self._fpl_widget) + + # Value label overlay + self._value_label = QLabel(self._fpl_widget) + self._value_label.setStyleSheet( + "background: rgba(25,25,30,200); color: #b4b4b4;" + " padding: 2px 6px; font-size: 9pt;" + " font-family: 'Menlo', 'Consolas', 'DejaVu Sans Mono', monospace;" + ) + self._value_label.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents) + self._value_label.hide() + + # Graphics + self._scatter = None + self._setup_graphics() + + # Disable pan/zoom, axes, and default title + self._subplot.controller = None + self._subplot.axes.visible = False + self._subplot.title.visible = False + + # Event handlers + renderer = self._subplot.renderer + renderer.add_event_handler(self._on_pointer_move_event, "pointer_move") + + # Animation callback + self._figure.add_animations(lambda: self._animation_callback()) + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def push_data(self, data: np.ndarray) -> None: + """Push scalar values. Shape ``(n_channels,)`` or ``(n_samples, n_channels)``.""" + self._buffer.push_data(data) + + # ------------------------------------------------------------------ + # Graphics setup + # ------------------------------------------------------------------ + + def _setup_graphics(self) -> None: + positions = self._positions_2d + n = positions.shape[0] + + # Build (n, 3) positions for fpl (z=0) + pts = np.zeros((n, 3), dtype=np.float32) + pts[:, :2] = positions + + self._scatter = self._subplot.add_scatter( + pts, + cmap=self._config.cmap, + cmap_transform=np.zeros(n, dtype=np.float32), + sizes=self._config.marker_size, + size_space="screen", + ) + + # ------------------------------------------------------------------ + # Animation callback + # ------------------------------------------------------------------ + + def _animation_callback(self) -> None: + values = self._buffer.consume() + if values is None: + return + + vmin, vmax = self._buffer.vmin, self._buffer.vmax + span = max(vmax - vmin, 1e-12) + norm = np.nan_to_num(np.clip((values - vmin) / span, 0.0, 1.0), nan=0.0).astype(np.float32) + + if self._config.modulate_color: + self._scatter.cmap.transform = norm + if self._config.modulate_size: + lo, hi = self._config.size_range + self._scatter.sizes = lo + norm * (hi - lo) + + self._subplot.auto_scale(maintain_aspect=True, zoom=0.9) + + # ------------------------------------------------------------------ + # Tooltip on hover + # ------------------------------------------------------------------ + + def _on_pointer_move_event(self, event) -> None: + world = self._subplot.map_screen_to_world(event) + if world is None: + self._value_label.hide() + return + + wx, wy = float(world[0]), float(world[1]) + dists = (self._positions_2d[:, 0] - wx) ** 2 + (self._positions_2d[:, 1] - wy) ** 2 + idx = int(np.argmin(dists)) + + labels = self._config.channel_labels + label = labels[idx] if labels and idx < len(labels) else f"Ch {idx}" + value = self._buffer._current[idx] + self._value_label.setText(f"{label}: {value:.4g}") + self._value_label.adjustSize() + self._value_label.show() + + QToolTip.showText( + self._fpl_widget.mapToGlobal(QPoint(int(event.x), int(event.y))), + f"{label}: {value:.4g}", + self._fpl_widget, + ) diff --git a/src/phosphor/shader.py b/src/phosphor/shader.py deleted file mode 100644 index c1ca4ae..0000000 --- a/src/phosphor/shader.py +++ /dev/null @@ -1,90 +0,0 @@ -"""WGSL shader source strings for sweep lines and cursor overlay.""" - -SWEEP_SHADER = """ -struct Uniforms { - y_scale: f32, - n_display_points: u32, - n_columns: u32, - sweep_col: u32, - cursor_gap: u32, - n_visible: u32, - _pad0: u32, - _pad1: u32, -} - -@group(0) @binding(0) var data: array; -@group(0) @binding(1) var channel_params: array; -@group(0) @binding(2) var uniforms: Uniforms; - -struct VertexOutput { - @builtin(position) position: vec4, - @location(0) color: vec4, -} - -@vertex -fn vs_main( - @builtin(vertex_index) vertex_index: u32, - @builtin(instance_index) instance_index: u32, -) -> VertexOutput { - var out: VertexOutput; - - let column = vertex_index / 2u; - let sub = vertex_index % 2u; - - // Read display data value (column-major interleaved min/max) - let data_index = (column * uniforms.n_visible + instance_index) * 2u + sub; - let value = data[data_index]; - - // Read per-channel params (8 floats per channel) - let param_base = instance_index * 8u; - let y_offset = channel_params[param_base + 0u]; - let color = vec4( - channel_params[param_base + 4u], - channel_params[param_base + 5u], - channel_params[param_base + 6u], - channel_params[param_base + 7u], - ); - - let x = (f32(column) + 0.5) / f32(uniforms.n_columns) * 2.0 - 1.0; - let y = value * uniforms.y_scale + y_offset; - - out.position = vec4(x, y, 0.0, 1.0); - out.color = color; - - return out; -} - -@fragment -fn fs_main(in: VertexOutput) -> @location(0) vec4 { - return in.color; -} -""" - -CURSOR_SHADER = """ -struct CursorUniforms { - x_left: f32, - x_right: f32, - _pad0: f32, - _pad1: f32, - color: vec4, -} - -@group(0) @binding(0) var cursor: CursorUniforms; - -@vertex -fn vs_cursor(@builtin(vertex_index) vi: u32) -> @builtin(position) vec4 { - // Full-height quad from two triangles - var px: array = array(0.0, 1.0, 0.0, 1.0, 1.0, 0.0); - var py: array = array(-1.0, -1.0, 1.0, -1.0, 1.0, 1.0); - - let x = mix(cursor.x_left, cursor.x_right, px[vi]); - let y = py[vi]; - - return vec4(x, y, 0.0, 1.0); -} - -@fragment -fn fs_cursor() -> @location(0) vec4 { - return cursor.color; -} -""" diff --git a/src/phosphor/spectrum_buffer.py b/src/phosphor/spectrum_buffer.py index fc11e3d..18cd7cb 100644 --- a/src/phosphor/spectrum_buffer.py +++ b/src/phosphor/spectrum_buffer.py @@ -4,19 +4,12 @@ import numpy as np -from .constants import ( - AUTOSCALE_N_SIGMA, - AUTOSCALE_TIME_CONSTANT, - CHANNEL_COLORS, -) - class SpectrumBuffer: """Thin storage buffer for pre-computed magnitude spectra. - Duck-types the interface expected by :class:`GPURenderer` and - :class:`ChannelPlotWidget` so both sweep and spectrum can share the - same rendering pipeline. + Duck-types the interface expected by :class:`ChannelPlotWidget` so both + sweep and spectrum widgets can share the same base class. """ def __init__(self, n_channels: int, n_bins: int, n_visible: int = 64): @@ -25,16 +18,13 @@ def __init__(self, n_channels: int, n_bins: int, n_visible: int = 64): self.n_visible = min(n_visible, n_channels) self.channel_offset = 0 - self.autoscale_enabled = True - self._manual_y_scale: float | None = None - self._lock = threading.Lock() self._version = 0 self._allocate() # ------------------------------------------------------------------ - # Alias so GPURenderer sees n_columns + # Alias so widgets can use n_columns generically # ------------------------------------------------------------------ @property @@ -49,24 +39,11 @@ def _allocate(self) -> None: self.display_mins = np.zeros((self.n_bins, self.n_visible), dtype=np.float32) self.display_maxs = np.zeros((self.n_bins, self.n_visible), dtype=np.float32) - self.ew_mean: np.ndarray | None = None - self.ew_sq_mean: np.ndarray | None = None - self._dirty_start: int | None = None self._dirty_end: int | None = None self._version += 1 - # ------------------------------------------------------------------ - # Sweep-cursor stub (pushed offscreen) - # ------------------------------------------------------------------ - - @property - def sweep_col(self) -> int: - """Return a value outside the valid column range so the cursor quad - is clipped offscreen.""" - return self.n_bins * 2 - # ------------------------------------------------------------------ # Public mutators # ------------------------------------------------------------------ @@ -102,8 +79,6 @@ def push_data(self, magnitudes: np.ndarray) -> None: self.display_mins[:] = vis self.display_maxs[:] = vis - self._update_autoscale(vis) - # Mark entire buffer dirty self._dirty_start = 0 self._dirty_end = self.n_bins - 1 @@ -139,100 +114,84 @@ def set_n_bins(self, n_bins: int) -> None: self.n_bins = n_bins self._allocate() - def adjust_y_scale(self, factor: float) -> None: - with self._lock: - self._manual_y_scale = self.y_scale * factor - self.autoscale_enabled = False - - def toggle_autoscale(self) -> None: - with self._lock: - self.autoscale_enabled = not self.autoscale_enabled - # ------------------------------------------------------------------ - # Properties and GPU data + # Properties and LineStack data # ------------------------------------------------------------------ - @property - def y_scale(self) -> float: - if not self.autoscale_enabled and self._manual_y_scale is not None: - return self._manual_y_scale - return self._compute_y_scale() - @property def version(self) -> int: return self._version - def get_gpu_data(self) -> np.ndarray: + def _compute_y_scale(self) -> float: + """Compute normalization scale. Must be called while holding ``_lock``.""" + max_abs = max(float(np.abs(self.display_mins).max()), float(np.abs(self.display_maxs).max())) + return 0.5 / max(max_abs, 1e-12) + + def _build_linestack_array(self, mins, maxs, bin_indices, freq_max, scale) -> np.ndarray: + """Build a ``[n_visible, 2*n_bins, 3]`` array. Must hold ``_lock``.""" + n_bins = mins.shape[0] + out = np.zeros((self.n_visible, 2 * n_bins, 3), dtype=np.float32) + bin_x = bin_indices.astype(np.float32) / max(self.n_bins - 1, 1) * freq_max + out[:, 0::2, 0] = bin_x[np.newaxis, :] + out[:, 1::2, 0] = bin_x[np.newaxis, :] + out[:, 0::2, 1] = mins.T * scale + out[:, 1::2, 1] = maxs.T * scale + return out + + def get_linestack_data(self, freq_max: float) -> np.ndarray: + """Full data shaped ``[n_visible, 2*n_bins, 3]`` for fastplotlib LineStack. + + Y-coordinates are normalized so the max absolute value maps to ±0.5. + """ with self._lock: - gpu = np.empty((self.n_bins, self.n_visible, 2), dtype=np.float32) - gpu[:, :, 0] = self.display_mins - gpu[:, :, 1] = self.display_maxs + self._y_scale = self._compute_y_scale() + out = self._build_linestack_array( + self.display_mins, self.display_maxs, np.arange(self.n_bins), freq_max, self._y_scale + ) self._dirty_start = None self._dirty_end = None - return gpu.reshape(-1) + return out - def get_dirty_gpu_data(self) -> tuple[np.ndarray, int, int] | None: + def get_dirty_linestack_range(self, freq_max: float) -> tuple[np.ndarray, int, int] | None: + """Incremental update for dirty bin range. + + Returns ``(data_slice, bin_start, n_bins)`` or ``None`` if clean. + If the scale changed significantly, returns a full-buffer update. + """ with self._lock: if self._dirty_start is None: return None + new_scale = self._compute_y_scale() + old_scale = getattr(self, "_y_scale", new_scale) + + if old_scale > 0 and abs(new_scale - old_scale) / old_scale > 0.2: + self._y_scale = new_scale + out = self._build_linestack_array( + self.display_mins, self.display_maxs, np.arange(self.n_bins), freq_max, self._y_scale + ) + self._dirty_start = None + self._dirty_end = None + return out, 0, self.n_bins + start = self._dirty_start end = self._dirty_end self._dirty_start = None self._dirty_end = None if end >= start: - n_cols = end - start + 1 - gpu = np.empty((n_cols, self.n_visible, 2), dtype=np.float32) - gpu[:, :, 0] = self.display_mins[start : end + 1] - gpu[:, :, 1] = self.display_maxs[start : end + 1] - return gpu.reshape(-1), start, n_cols + n_bins = end - start + 1 + out = self._build_linestack_array( + self.display_mins[start : end + 1], + self.display_maxs[start : end + 1], + np.arange(start, end + 1), + freq_max, + old_scale, + ) + return out, start, n_bins else: - gpu = np.empty((self.n_bins, self.n_visible, 2), dtype=np.float32) - gpu[:, :, 0] = self.display_mins - gpu[:, :, 1] = self.display_maxs - return gpu.reshape(-1), 0, self.n_bins - - def get_channel_params(self) -> np.ndarray: - params = np.zeros((self.n_visible, 8), dtype=np.float32) - ys = self.y_scale - for i in range(self.n_visible): - y_center = 1.0 - (2.0 * (i + 0.5)) / self.n_visible - mean_off = float(self.ew_mean[i] * ys) if self.ew_mean is not None else 0.0 - params[i, 0] = y_center - mean_off - c = CHANNEL_COLORS[i % len(CHANNEL_COLORS)] - params[i, 4] = c[0] - params[i, 5] = c[1] - params[i, 6] = c[2] - params[i, 7] = c[3] - return params.reshape(-1) - - # ------------------------------------------------------------------ - # Internal helpers - # ------------------------------------------------------------------ - - def _update_autoscale(self, vis_data: np.ndarray) -> None: - clean = np.nan_to_num(vis_data, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) - batch_mean = np.mean(clean, axis=0) - batch_sq_mean = np.mean(clean**2, axis=0) - - # Use a fixed time constant equivalent — treat each push as 1 update - alpha = 1.0 - np.exp(-1.0 / AUTOSCALE_TIME_CONSTANT) - - if self.ew_mean is None: - self.ew_mean = batch_mean.copy() - self.ew_sq_mean = batch_sq_mean.copy() - else: - self.ew_mean += alpha * (batch_mean - self.ew_mean) - self.ew_sq_mean += alpha * (batch_sq_mean - self.ew_sq_mean) - - def _compute_y_scale(self) -> float: - if self.ew_mean is None: - return 1.0 / max(self.n_visible, 1) - var = self.ew_sq_mean - self.ew_mean**2 - var = np.maximum(var, 0.0) - sigma = np.sqrt(var) - mean_sigma = float(np.mean(sigma)) - if mean_sigma < 1e-12: - return 1.0 / max(self.n_visible, 1) - return 1.0 / (AUTOSCALE_N_SIGMA * mean_sigma * self.n_visible) + self._y_scale = new_scale + out = self._build_linestack_array( + self.display_mins, self.display_maxs, np.arange(self.n_bins), freq_max, self._y_scale + ) + return out, 0, self.n_bins diff --git a/src/phosphor/spectrum_widget.py b/src/phosphor/spectrum_widget.py index 8577074..737b0cc 100644 --- a/src/phosphor/spectrum_widget.py +++ b/src/phosphor/spectrum_widget.py @@ -5,12 +5,10 @@ from dataclasses import dataclass import numpy as np -from PySide6.QtCore import Qt from PySide6.QtWidgets import QWidget from .channel_plot import ChannelPlotWidget -from .constants import DEFAULT_N_VISIBLE -from .gpu_renderer import GPURenderer +from .constants import CHANNEL_COLORS, DEFAULT_N_VISIBLE from .spectrum_buffer import SpectrumBuffer from .x_axis import XAxisWidget @@ -62,8 +60,10 @@ def __init__(self, config: SpectrumConfig, parent: QWidget | None = None): ) self._buffer = self.spectrum_buffer - # GPU renderer - self.gpu_renderer = GPURenderer(self.canvas) + # Create initial graphics + self._cached_version = -1 + self._line_stack = None + self._setup_graphics() # Start rendering self._init_rendering() @@ -103,39 +103,83 @@ def update_config(self, config: SpectrumConfig) -> None: self._sync_display() # ------------------------------------------------------------------ - # Rendering + # Graphics setup # ------------------------------------------------------------------ - def _draw_frame(self) -> None: - self.gpu_renderer.update_and_draw(self.spectrum_buffer) + def _setup_graphics(self) -> None: + """Create or recreate LineStack.""" + subplot = self._subplot + + if self._line_stack is not None: + subplot.delete_graphic(self._line_stack) + self._line_stack = None + + buf = self.spectrum_buffer + data = buf.get_linestack_data(self._display_freq_max) + + n_vis = buf.n_visible + colors = [CHANNEL_COLORS[i % len(CHANNEL_COLORS)][:3] for i in range(n_vis)] + + self._line_stack = subplot.add_line_stack( + data, + colors=colors, + separation=1.0, + separation_axis="y", + ) + + self._cached_version = buf.version + + # ------------------------------------------------------------------ + # Rendering (animation callback) + # ------------------------------------------------------------------ + + def _update_graphics(self) -> None: + buf = self.spectrum_buffer + + if buf.version != self._cached_version: + self._setup_graphics() + return + + # Incremental update + result = buf.get_dirty_linestack_range(self._display_freq_max) + if result is not None: + data_slice, bin_start, n_bins = result + idx_start = bin_start * 2 + idx_end = (bin_start + n_bins) * 2 + for ch in range(buf.n_visible): + self._line_stack[ch].data[idx_start:idx_end] = data_slice[ch] # ------------------------------------------------------------------ # Keyboard controls # ------------------------------------------------------------------ - def _handle_key(self, key: int) -> None: - if key == Qt.Key.Key_Comma: - # Halve frequency range - new_max = max(self._display_freq_max / 2.0, self._freq_min * 4) - if new_max != self._display_freq_max: - self._display_freq_max = new_max - self._sync_display() - elif key == Qt.Key.Key_Period: - # Double frequency range - new_max = min(self._display_freq_max * 2.0, self._nyquist) - if new_max != self._display_freq_max: - self._display_freq_max = new_max - self._sync_display() - elif key == Qt.Key.Key_L: + def _on_key_down(self, key: str) -> None: + if key == ",": + self._freq_zoom(0.5) + elif key == ".": + self._freq_zoom(2.0) + elif key in ("l", "L"): self._log_x = not self._log_x self._sync_display() else: - super()._handle_key(key) + super()._on_key_down(key) + + def _on_ctrl_scroll(self, delta: float) -> None: + factor = 0.5 if delta > 0 else 2.0 + self._freq_zoom(factor) # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ + def _freq_zoom(self, factor: float) -> None: + new_max = self._display_freq_max * factor + new_max = max(new_max, self._freq_min * 4) + new_max = min(new_max, self._nyquist) + if new_max != self._display_freq_max: + self._display_freq_max = new_max + self._sync_display() + def _linear_display_bins(self) -> int: """Number of linear bins covering 0 .. _display_freq_max.""" frac = self._display_freq_max / self._nyquist diff --git a/src/phosphor/sweep_buffer.py b/src/phosphor/sweep_buffer.py index 4f8384b..2f8bedb 100644 --- a/src/phosphor/sweep_buffer.py +++ b/src/phosphor/sweep_buffer.py @@ -1,16 +1,10 @@ -"""CPU-side circular buffer with incremental min/max downsampling and autoscale.""" +"""CPU-side circular buffer with incremental min/max downsampling.""" import threading import warnings import numpy as np -from .constants import ( - AUTOSCALE_N_SIGMA, - AUTOSCALE_TIME_CONSTANT, - CHANNEL_COLORS, -) - class SweepBuffer: def __init__( @@ -29,9 +23,6 @@ def __init__( self.n_visible = min(n_visible, n_channels) self.channel_offset = 0 - self.autoscale_enabled = True - self._manual_y_scale: float | None = None - self._lock = threading.Lock() self._version = 0 @@ -54,9 +45,6 @@ def _allocate(self): self.write_pos = 0 self.sweep_col = 0 - self.ew_mean: np.ndarray | None = None - self.ew_sq_mean: np.ndarray | None = None - self._dirty_start: int | None = None self._dirty_end: int | None = None @@ -94,9 +82,6 @@ def push_data(self, data: np.ndarray) -> None: vis_data = vis_data[-self.total_raw_samples :] n_samples = self.total_raw_samples - # Update autoscale statistics - self._update_autoscale(vis_data) - # Track column range before writing first_col = self._col_for_pos(self.write_pos) @@ -164,49 +149,82 @@ def set_srate(self, srate: float) -> None: self.srate = srate self._allocate() - def adjust_y_scale(self, factor: float) -> None: - """Multiply y_scale by factor and disable autoscale.""" - with self._lock: - self._manual_y_scale = self.y_scale * factor - self.autoscale_enabled = False - - def toggle_autoscale(self) -> None: - with self._lock: - self.autoscale_enabled = not self.autoscale_enabled - # ------------------------------------------------------------------ - # Properties and GPU data (called from render/UI thread) + # Properties and LineStack data (called from render/UI thread) # ------------------------------------------------------------------ - @property - def y_scale(self) -> float: - if not self.autoscale_enabled and self._manual_y_scale is not None: - return self._manual_y_scale - return self._compute_y_scale() - @property def version(self) -> int: return self._version - def get_gpu_data(self) -> np.ndarray: - """Full upload: interleave min/max into column-major flat array.""" + def _compute_y_scale(self) -> float: + """Compute normalization scale from current buffer data. + + Uses the max absolute value across all display columns/channels so + that normalized data fits within ±0.5, matching LineStack separation. + Must be called while holding ``_lock``. + """ + max_abs = max(float(np.abs(self.display_mins).max()), float(np.abs(self.display_maxs).max())) + return 0.5 / max(max_abs, 1e-12) + + def _build_linestack_array(self, mins, maxs, col_indices, scale) -> np.ndarray: + """Build a ``[n_visible, 2*n_cols, 3]`` array from min/max slices. + + Must be called while holding ``_lock``. + """ + n_cols = mins.shape[0] + out = np.zeros((self.n_visible, 2 * n_cols, 3), dtype=np.float32) + col_x = col_indices.astype(np.float32) / max(self.n_columns - 1, 1) * self.display_dur + out[:, 0::2, 0] = col_x[np.newaxis, :] + out[:, 1::2, 0] = col_x[np.newaxis, :] + out[:, 0::2, 1] = mins.T * scale + out[:, 1::2, 1] = maxs.T * scale + return out + + def get_linestack_data(self) -> np.ndarray: + """Full data shaped ``[n_visible, 2*n_columns, 3]`` for fastplotlib LineStack. + + Y-coordinates are normalized so the max absolute value maps to ±0.5, + matching LineStack separation=1.0 and preventing channel overlap. + """ with self._lock: - gpu = np.empty((self.n_columns, self.n_visible, 2), dtype=np.float32) - gpu[:, :, 0] = self.display_mins - gpu[:, :, 1] = self.display_maxs + self._y_scale = self._compute_y_scale() + out = self._build_linestack_array( + self.display_mins, + self.display_maxs, + np.arange(self.n_columns), + self._y_scale, + ) self._dirty_start = None self._dirty_end = None - return gpu.reshape(-1) + return out - def get_dirty_gpu_data(self) -> tuple[np.ndarray, int, int] | None: - """Partial upload for dirty column range. + def get_dirty_linestack_range(self) -> tuple[np.ndarray, int, int] | None: + """Incremental update for dirty column range. - Returns (flat_data, col_start, n_cols) or None if nothing dirty. + Returns ``(data_slice, col_start, n_cols)`` or ``None`` if clean. + If the scale changed significantly, returns a full-buffer update. """ with self._lock: if self._dirty_start is None: return None + # Check if scale drifted significantly + new_scale = self._compute_y_scale() + old_scale = getattr(self, "_y_scale", new_scale) + if old_scale > 0 and abs(new_scale - old_scale) / old_scale > 0.2: + # Full update with new scale + self._y_scale = new_scale + out = self._build_linestack_array( + self.display_mins, + self.display_maxs, + np.arange(self.n_columns), + self._y_scale, + ) + self._dirty_start = None + self._dirty_end = None + return out, 0, self.n_columns + start = self._dirty_start end = self._dirty_end self._dirty_start = None @@ -214,30 +232,23 @@ def get_dirty_gpu_data(self) -> tuple[np.ndarray, int, int] | None: if end >= start: n_cols = end - start + 1 - gpu = np.empty((n_cols, self.n_visible, 2), dtype=np.float32) - gpu[:, :, 0] = self.display_mins[start : end + 1] - gpu[:, :, 1] = self.display_maxs[start : end + 1] - return gpu.reshape(-1), start, n_cols + out = self._build_linestack_array( + self.display_mins[start : end + 1], + self.display_maxs[start : end + 1], + np.arange(start, end + 1), + old_scale, + ) + return out, start, n_cols else: - # Wrapped — upload full buffer - gpu = np.empty((self.n_columns, self.n_visible, 2), dtype=np.float32) - gpu[:, :, 0] = self.display_mins - gpu[:, :, 1] = self.display_maxs - return gpu.reshape(-1), 0, self.n_columns - - def get_channel_params(self) -> np.ndarray: - """Per-channel y_offset and RGBA color. Shape (n_visible * 8,) float32.""" - params = np.zeros((self.n_visible, 8), dtype=np.float32) - # ys = self.y_scale - for i in range(self.n_visible): - y_center = 1.0 - (2.0 * (i + 0.5)) / self.n_visible - params[i, 0] = y_center - c = CHANNEL_COLORS[i % len(CHANNEL_COLORS)] - params[i, 4] = c[0] - params[i, 5] = c[1] - params[i, 6] = c[2] - params[i, 7] = c[3] - return params.reshape(-1) + # Wrapped — full update + self._y_scale = new_scale + out = self._build_linestack_array( + self.display_mins, + self.display_maxs, + np.arange(self.n_columns), + self._y_scale, + ) + return out, 0, self.n_columns # ------------------------------------------------------------------ # Internal helpers @@ -263,32 +274,6 @@ def _recompute_columns(self, first_col: int, n_cols: int) -> None: self.display_mins[col] = np.nan_to_num(mins, nan=0.0) self.display_maxs[col] = np.nan_to_num(maxs, nan=0.0) - def _update_autoscale(self, new_data: np.ndarray) -> None: - clean = np.nan_to_num(new_data, nan=0.0, posinf=0.0, neginf=0.0).astype(np.float64) - batch_mean = np.mean(clean, axis=0) - batch_sq_mean = np.mean(clean**2, axis=0) - - dt = new_data.shape[0] / max(self.srate, 1.0) - alpha = 1.0 - np.exp(-dt / AUTOSCALE_TIME_CONSTANT) - - if self.ew_mean is None: - self.ew_mean = batch_mean.copy() - self.ew_sq_mean = batch_sq_mean.copy() - else: - self.ew_mean += alpha * (batch_mean - self.ew_mean) - self.ew_sq_mean += alpha * (batch_sq_mean - self.ew_sq_mean) - - def _compute_y_scale(self) -> float: - if self.ew_mean is None: - return 1.0 / max(self.n_visible, 1) - var = self.ew_sq_mean - self.ew_mean**2 - var = np.maximum(var, 0.0) - sigma = np.sqrt(var) - mean_sigma = float(np.mean(sigma)) - if mean_sigma < 1e-12: - return 1.0 / max(self.n_visible, 1) - return 1.0 / (AUTOSCALE_N_SIGMA * mean_sigma * self.n_visible) - def _mark_dirty(self, first_col: int, last_col: int) -> None: if first_col <= last_col: if self._dirty_start is None: diff --git a/src/phosphor/sweep_widget.py b/src/phosphor/sweep_widget.py index 9275dd4..d4b321d 100644 --- a/src/phosphor/sweep_widget.py +++ b/src/phosphor/sweep_widget.py @@ -5,12 +5,17 @@ from dataclasses import dataclass import numpy as np -from PySide6.QtCore import Qt from PySide6.QtWidgets import QWidget from .channel_plot import ChannelPlotWidget -from .constants import DEFAULT_DISPLAY_DUR, DEFAULT_N_COLUMNS, DEFAULT_N_VISIBLE -from .gpu_renderer import GPURenderer +from .constants import ( + CHANNEL_COLORS, + CURSOR_COLOR, + CURSOR_GAP_COLUMNS, + DEFAULT_DISPLAY_DUR, + DEFAULT_N_COLUMNS, + DEFAULT_N_VISIBLE, +) from .sweep_buffer import SweepBuffer from .x_axis import XAxisWidget @@ -60,10 +65,11 @@ def __init__(self, config: SweepConfig, parent: QWidget | None = None): ) self._buffer = self.sweep_buffer - # Create GPU renderer eagerly so the wgpu context exists - # before the first draw (rendercanvas's _draw_and_present - # cancels the draw if context is None). - self.gpu_renderer = GPURenderer(self.canvas) + # Create initial graphics + self._cached_version = -1 + self._line_stack = None + self._cursor_line = None + self._setup_graphics() # Start rendering self._init_rendering() @@ -94,26 +100,101 @@ def update_config(self, config: SweepConfig) -> None: self._update_range_label() # ------------------------------------------------------------------ - # Rendering + # Graphics setup # ------------------------------------------------------------------ - def _draw_frame(self) -> None: - self.gpu_renderer.update_and_draw(self.sweep_buffer) + def _setup_graphics(self) -> None: + """Create or recreate LineStack and cursor line.""" + subplot = self._subplot + + # Delete old graphics + if self._line_stack is not None: + subplot.delete_graphic(self._line_stack) + self._line_stack = None + if self._cursor_line is not None: + subplot.delete_graphic(self._cursor_line) + self._cursor_line = None + + buf = self.sweep_buffer + data = buf.get_linestack_data() + + # Build per-channel colors (cycling through CHANNEL_COLORS) + n_vis = buf.n_visible + colors = [CHANNEL_COLORS[i % len(CHANNEL_COLORS)][:3] for i in range(n_vis)] + + self._line_stack = subplot.add_line_stack( + data, + colors=colors, + separation=1.0, + separation_axis="y", + ) + + # Cursor: vertical line at the sweep position. + # Span the LineStack's y-extent with small margin. + sweep_x = buf.sweep_col / max(buf.n_columns - 1, 1) * buf.display_dur + cursor_color = CURSOR_COLOR[:3] + gap_w = CURSOR_GAP_COLUMNS / max(buf.n_columns - 1, 1) * buf.display_dur + y_bottom = self._line_stack[0].world_object.world.position[1] + y_top = self._line_stack[-1].world_object.world.position[1] + margin = max((y_top - y_bottom) * 0.05, 0.5) + self._cursor_y_min = y_bottom - margin + self._cursor_y_max = y_top + margin + self._cursor_line = subplot.add_line( + np.array( + [[sweep_x, self._cursor_y_min, 0], [sweep_x, self._cursor_y_max, 0]], + dtype=np.float32, + ), + colors=cursor_color, + thickness=max(1.0, gap_w * 2), + ) + + self._cached_version = buf.version # ------------------------------------------------------------------ - # Keyboard controls (sweep-specific keys) + # Rendering (animation callback) # ------------------------------------------------------------------ - def _handle_key(self, key: int) -> None: + def _update_graphics(self) -> None: buf = self.sweep_buffer - if key == Qt.Key.Key_Comma: - buf.set_display_dur(buf.display_dur / 2.0) # halve duration - self._time_axis.set_range(buf.display_dur) - self._update_range_label() - elif key == Qt.Key.Key_Period: - buf.set_display_dur(buf.display_dur * 2.0) # double duration - self._time_axis.set_range(buf.display_dur) - self._update_range_label() + if buf.version != self._cached_version: + # Version changed (scroll, resize, display_dur change) → full rebuild + self._setup_graphics() + return + + # Incremental update from dirty columns + result = buf.get_dirty_linestack_range() + if result is not None: + data_slice, col_start, n_cols = result + idx_start = col_start * 2 + idx_end = (col_start + n_cols) * 2 + for ch in range(buf.n_visible): + self._line_stack[ch].data[idx_start:idx_end] = data_slice[ch] + + # Update cursor x-position; y spans the LineStack extent (not camera, + # which would create a feedback loop with auto_scale). + sweep_x = buf.sweep_col / max(buf.n_columns - 1, 1) * buf.display_dur + self._cursor_line.data[0] = [sweep_x, self._cursor_y_min, 0] + self._cursor_line.data[1] = [sweep_x, self._cursor_y_max, 0] + + # ------------------------------------------------------------------ + # Keyboard controls (sweep-specific keys) + # ------------------------------------------------------------------ + + def _on_key_down(self, key: str) -> None: + if key == ",": + self._time_zoom(0.5) + elif key == ".": + self._time_zoom(2.0) else: - super()._handle_key(key) + super()._on_key_down(key) + + def _on_ctrl_scroll(self, delta: float) -> None: + factor = 0.5 if delta > 0 else 2.0 + self._time_zoom(factor) + + def _time_zoom(self, factor: float) -> None: + buf = self.sweep_buffer + buf.set_display_dur(buf.display_dur * factor) + self._time_axis.set_range(buf.display_dur) + self._update_range_label()