From c245654e0a5d6f28e40ea3cb02bc52972cfe487f Mon Sep 17 00:00:00 2001
From: Alex Rockhill <aprockhill@mailbox.org>
Date: Sun, 12 Nov 2023 04:36:54 -0800
Subject: [PATCH] [ENH] More natural zoom (#19)

---
 mne_gui_addons/_core.py        | 23 ++++++++++++++++++-----
 mne_gui_addons/_ieeg_locate.py |  2 +-
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/mne_gui_addons/_core.py b/mne_gui_addons/_core.py
index ead1cf3..a5c6848 100644
--- a/mne_gui_addons/_core.py
+++ b/mne_gui_addons/_core.py
@@ -44,6 +44,7 @@
 
 _IMG_LABELS = [["I", "P"], ["I", "L"], ["P", "L"]]
 _ZOOM_STEP_SIZE = 5
+_ZOOM_BORDER = 1 / 5
 
 
 @verbose
@@ -422,11 +423,22 @@ def _zoom(self, sign=1, draw=False):
             ymin, ymax = fig.axes[0].get_ylim()
             xmid = (xmin + xmax) / 2
             ymid = (ymin + ymax) / 2
-            if sign == 1:  # may need to shift if zooming in
-                if abs(xmid - xcur) > delta / 2 * rx:
-                    xmid += delta * np.sign(xcur - xmid) * rx
-                if abs(ymid - ycur) > delta / 2 * ry:
-                    ymid += delta * np.sign(ycur - ymid) * ry
+            if sign >= 0:  # may need to shift if zooming in or clicking
+                xedge = min([xmax - xcur, xcur - xmin])
+                if xedge < (xmax - xmin) * _ZOOM_BORDER:
+                    xmid += np.sign(xcur - xmid) * (
+                        (xmax - xmin) * _ZOOM_BORDER - xedge
+                    )
+                if xcur < xmin or xcur > xmax:  # out of view, reset
+                    xmid = xcur
+                yedge = min([ymax - ycur, ycur - ymin])
+                if yedge < (ymax - ymin) * _ZOOM_BORDER:
+                    ymid += np.sign(ycur - ymid) * (
+                        (ymax - ymin) * _ZOOM_BORDER - yedge
+                    )
+                if ycur < ymin or ycur > ymax:  # out of view, reset
+                    ymid = ycur
+
             xwidth = (xmax - xmin) / 2 - delta * rx
             ywidth = (ymax - ymin) / 2 - delta * ry
             if xwidth <= 0 or ywidth <= 0:
@@ -620,6 +632,7 @@ def _on_click(self, event, axis):
             logger.debug(f"Using voxel  {list(xyz)}")
             ras = apply_trans(self._vox_ras_t, xyz)
             self._set_ras(ras)
+            self._zoom(sign=0, draw=True)
 
     def _update_moved(self):
         """Update when cursor position changes."""
diff --git a/mne_gui_addons/_ieeg_locate.py b/mne_gui_addons/_ieeg_locate.py
index d9ec82a..afd947f 100644
--- a/mne_gui_addons/_ieeg_locate.py
+++ b/mne_gui_addons/_ieeg_locate.py
@@ -553,8 +553,8 @@ def _update_ch_selection(self):
         self._update_group()
         if not np.isnan(self._chs[name]).any():
             self._set_ras(self._chs[name])
+            self._zoom(sign=0, draw=True)
             self._update_camera(render=True)
-            self._draw()
 
     def _go_to_ch(self, index):
         """Change current channel to the item selected."""