From ac370942ed9e876d7ef25294616d5e27df693221 Mon Sep 17 00:00:00 2001
From: Hanjin Liu <hanjin.liu@bs.s.u-tokyo.ac.jp>
Date: Thu, 21 Sep 2023 21:50:22 +0900
Subject: [PATCH] improve plot UI

---
 magicclass/widgets/_mpl_canvas.py | 51 ++++++++++++++++++-------------
 1 file changed, 29 insertions(+), 22 deletions(-)

diff --git a/magicclass/widgets/_mpl_canvas.py b/magicclass/widgets/_mpl_canvas.py
index ddac24dd..062493b6 100644
--- a/magicclass/widgets/_mpl_canvas.py
+++ b/magicclass/widgets/_mpl_canvas.py
@@ -1,7 +1,6 @@
 from __future__ import annotations
 from typing import TYPE_CHECKING
 import numpy as np
-from matplotlib.backends.backend_qt5agg import FigureCanvas
 from matplotlib.backend_bases import MouseEvent, MouseButton
 from qtpy import QtWidgets as QtW, QtGui
 
@@ -9,6 +8,12 @@
     from matplotlib.figure import Figure
     from matplotlib.axes import Axes
 
+    class FigureCanvas(QtW.QWidget):
+        ...  # fmt: skip
+
+else:
+    from matplotlib.backends.backend_qt5agg import FigureCanvas
+
 
 class InteractiveFigureCanvas(FigureCanvas):
     """A figure canvas implemented with mouse callbacks."""
@@ -60,28 +65,25 @@ def mouseMoveEvent(self, event):
         cursor exists will be translated.
         """
         ax = self.last_axis
-        if (
-            self.pressed not in (MouseButton.LEFT, MouseButton.RIGHT)
-            or self.lastx_pressed is None
-            or not self._interactive
-            or not ax
-        ):
+        if self.lastx_pressed is None or not self._interactive or not ax:
             return
-
         event = self.get_mouse_event(event)
-        x, y = event.xdata, event.ydata
+        xdata, ydata = _get_xy_data(ax, event)
 
-        if x is None or y is None:
+        if xdata is None or ydata is None:
             return None
 
         if self.pressed == MouseButton.LEFT:
-            _translate_x(ax, self.lastx_pressed, x)
-            _translate_y(ax, self.lasty_pressed, y)
+            _translate_x(ax, self.lastx_pressed, xdata)
+            _translate_y(ax, self.lasty_pressed, ydata)
         elif self.pressed == MouseButton.RIGHT:
-            _zoom_x(ax, self.lastx, x)
-            _zoom_y(ax, self.lasty, y)
+            _zoom_x(ax, self.lastx, xdata)
+            _zoom_y(ax, self.lasty, ydata)
+            xdata, ydata = _get_xy_data(ax, event)  # ticks changed!
+        else:
+            return None
 
-        self.lastx, self.lasty = x, y
+        self.lastx, self.lasty = xdata, ydata
         self.figure.canvas.draw()
         return None
 
@@ -89,7 +91,7 @@ def mouseReleaseEvent(self, event):
         """Stop dragging state."""
         pos = event.pos()
         event = self.get_mouse_event(event)
-        if self.lastx == event.xdata and self.lasty == event.ydata:
+        if self.lastx_pressed == event.xdata and self.lasty_pressed == event.ydata:
             if self.pressed == MouseButton.LEFT:
                 for callbacks in self._mouse_click_callbacks:
                     callbacks(event)
@@ -212,10 +214,8 @@ def _zoom_x(ax: Axes, xstart: float, xstop: float):
     xscale = ax.get_xscale()
     x0, x1 = ax.get_xlim()
     if xscale == "linear":
-        _u = x1 + x0
-        _v = x1 - x0
         dx = xstop - xstart
-        ax.set_xlim([_u / 2 - _v / 2 + dx, _u / 2 + _v / 2 - dx])
+        ax.set_xlim([x0 + dx, x1 - dx])
     elif xscale == "log":
         if xstart <= 0 or xstop <= 0:
             ax.autoscale(axis="x")
@@ -227,10 +227,8 @@ def _zoom_y(ax: Axes, ystart: float, ystop: float):
     yscale = ax.get_yscale()
     y0, y1 = ax.get_ylim()
     if yscale == "linear":
-        _u = y1 + y0
-        _v = y1 - y0
         dy = ystop - ystart
-        ax.set_ylim([_u / 2 - _v / 2 + dy, _u / 2 + _v / 2 - dy])
+        ax.set_ylim([y0 + dy, y1 - dy])
     elif yscale == "log":
         if ystart <= 0 or ystop <= 0:
             ax.autoscale(axis="y")
@@ -271,3 +269,12 @@ def _zoom_y_wheel(ax: Axes, factor: float):
             y0_t = (y0 / yc) ** factor
             y1_t = (y1 / yc) ** factor
             ax.set_ylim([y0_t * yc, y1_t * yc])
+
+
+def _get_xy_data(ax: Axes, event: MouseEvent) -> tuple[float, float]:
+    try:
+        trans = ax.transData.inverted()
+        xdata, ydata = trans.transform((event.x, event.y))
+    except Exception:
+        xdata, ydata = event.xdata, event.ydata
+    return xdata, ydata