Line data    Source code

       1             : # SPDX-FileCopyrightText: 2025 PairInteraction Developers
       2             : # SPDX-License-Identifier: LGPL-3.0-or-later
       3           1 : from __future__ import annotations
       4             : 
       5           1 : from typing import TYPE_CHECKING
       6             : 
       7           1 : import matplotlib.pyplot as plt
       8           1 : import numpy as np
       9           1 : from matplotlib.backends.backend_qtagg import FigureCanvasQTAgg
      10           1 : from PySide6.QtCore import Qt, QTimer
      11             : 
      12             : if TYPE_CHECKING:
      13             :     from PySide6.QtGui import QWheelEvent
      14             :     from PySide6.QtWidgets import QWidget
      15             : 
      16             : 
      17           1 : class MatplotlibCanvas(FigureCanvasQTAgg):
      18             :     """Canvas for matplotlib figures."""
      19             : 
      20           1 :     def __init__(self, parent: QWidget | None = None) -> None:
      21             :         """Initialize the canvas with a figure."""
      22           1 :         self.fig, self.ax = plt.subplots()
      23           1 :         super().__init__(self.fig)
      24             : 
      25           1 :         self.setup_zoom()
      26             : 
      27           1 :     def setup_zoom(self) -> None:
      28             :         """Set up mouse wheel zoom functionality."""
      29             :         # Wheel event accumulation variables
      30           1 :         self.wheel_accumulation: float = 0
      31           1 :         self.last_wheel_pos: list[list[float]] = []
      32           1 :         self.wheel_timer = QTimer(self)
      33           1 :         self.wheel_timer.setSingleShot(True)
      34           1 :         self.wheel_timer.timeout.connect(self.apply_accumulated_zoom)
      35           1 :         self.setFocusPolicy(Qt.FocusPolicy.StrongFocus)
      36           1 :         self.setMouseTracking(True)
      37             : 
      38           1 :     def wheelEvent(self, event: QWheelEvent) -> None:
      39             :         """Handle mouse wheel events for zooming."""
      40           0 :         self.wheel_accumulation += event.angleDelta().y() / 120
      41           0 :         self.last_wheel_pos.append([event.position().x(), event.position().y()])
      42           0 :         self.wheel_timer.start(100)  # Apply zoom after 100ms of inactivity
      43             : 
      44           1 :     def apply_accumulated_zoom(self) -> None:
      45             :         """Apply the accumulated zoom from wheel events."""
      46           0 :         if not self.wheel_accumulation:
      47           0 :             return
      48             : 
      49           0 :         x_min, x_max = self.ax.get_xlim()
      50           0 :         y_min, y_max = self.ax.get_ylim()
      51             : 
      52           0 :         scale_factor = 1 - 0.1 * self.wheel_accumulation
      53             : 
      54             :         # Get the mouse position in data coordinates
      55           0 :         wheel_pos_mean = np.mean(self.last_wheel_pos, axis=0)
      56           0 :         x_data, y_data = self.ax.transData.inverted().transform(wheel_pos_mean)
      57           0 :         y_data = -(y_data - (y_max + y_min) / 2) + (y_max + y_min) / 2  # y_data is mirrored bottom / top
      58             : 
      59           0 :         self.wheel_accumulation = 0
      60           0 :         self.last_wheel_pos = []
      61             : 
      62           0 :         if x_data > x_max or y_data > y_max or (x_data < x_min and y_data < y_min):
      63           0 :             return
      64             : 
      65           0 :         if x_min <= x_data <= x_max:
      66           0 :             x_min_new = x_data - (x_data - x_min) * scale_factor
      67           0 :             x_max_new = x_data + (x_max - x_data) * scale_factor
      68           0 :             self.ax.set_xbound(x_min_new, x_max_new)
      69           0 :             self.ax.set_autoscalex_on(False)
      70             : 
      71           0 :         if y_min <= y_data <= y_max:
      72           0 :             y_min_new = y_data - (y_data - y_min) * scale_factor
      73           0 :             y_max_new = y_data + (y_max - y_data) * scale_factor
      74           0 :             self.ax.set_ybound(y_min_new, y_max_new)
      75           0 :             self.ax.set_autoscaley_on(False)
      76             : 
      77             :         # Redraw the canvas
      78           0 :         self.draw()