Adding grid/zoom options and add fast forward functionality; improved app.py

Player.py

@@ -21,15 +21,24 @@ def __init__(
         pos=(0.125, 0.92),
         times=None,
         t_units="",
+        grid_state=False,
+        zoom_state=False,
+        on_toggle_grid=None,
+        on_toggle_zoom=None,
         **kwargs,
     ):
         self.i = 0
         self.min = mini
         self.max = maxi
         self.runs = True
         self.forwards = True
+        self.step = 1
         self.fig = fig
         self.func = func
+        self.grid_state = grid_state
+        self.zoom_state = zoom_state
+        self.on_toggle_grid = on_toggle_grid
+        self.on_toggle_zoom = on_toggle_zoom
         self.setup(pos)
         FuncAnimation.__init__(
             self,
@@ -46,10 +55,11 @@ def __init__(
 
     def play(self):
         while self.runs:
-            self.i = self.i + self.forwards - (not self.forwards)
-            if self.i > self.min and self.i < self.max:
+            self.i = self.i + self.step * (1 if self.forwards else -1)
+            if self.min < self.i < self.max:
                 yield self.i
             else:
+                self.i = self.max if self.forwards else self.min
                 self.stop()
                 yield self.i
 
@@ -63,10 +73,17 @@ def stop(self, event=None):
 
     def forward(self, event=None):
         self.forwards = True
+        self.step = 1
         self.start()
 
     def backward(self, event=None):
         self.forwards = False
+        self.step = 1
+        self.start()
+
+    def fastforward(self, event=None):
+        self.forwards = True
+        self.step = 10
         self.start()
 
     def oneforward(self, event=None):
@@ -95,17 +112,22 @@ def setup(self, pos):
         bax = divider.append_axes("right", size="80%", pad=0.05)
         sax = divider.append_axes("right", size="80%", pad=0.05)
         fax = divider.append_axes("right", size="80%", pad=0.05)
+        ffax = divider.append_axes("right", size="80%", pad=0.05)
         ofax = divider.append_axes("right", size="100%", pad=0.05)
         sliderax = divider.append_axes("right", size="500%", pad=0.07)
         self.button_oneback = matplotlib.widgets.Button(playerax, label="$\u29cf$")
         self.button_back = matplotlib.widgets.Button(bax, label="$\u25c0$")
         self.button_stop = matplotlib.widgets.Button(sax, label="$\u25a0$")
         self.button_forward = matplotlib.widgets.Button(fax, label="$\u25b6$")
+        self.button_fastforward = matplotlib.widgets.Button(
+            ffax, label="$\u25b6\u25b6$"
+        )
         self.button_oneforward = matplotlib.widgets.Button(ofax, label="$\u29d0$")
         self.button_oneback.on_clicked(self.onebackward)
         self.button_back.on_clicked(self.backward)
         self.button_stop.on_clicked(self.stop)
         self.button_forward.on_clicked(self.forward)
+        self.button_fastforward.on_clicked(self.fastforward)
         self.button_oneforward.on_clicked(self.oneforward)
         self.slider = matplotlib.widgets.Slider(
             sliderax, "", self.min, self.max, valinit=self.i
@@ -120,6 +142,27 @@ def setup(self, pos):
 
         # self.text_box = matplotlib.widgets.TextBox(textax, label="", initial="0 ms")
 
+        checkax = self.fig.add_axes([0.02, 0.01, 0.13, 0.08])
+        self.check_labels = ["Grid", "Zoom"]
+        self.check_buttons = matplotlib.widgets.CheckButtons(
+            checkax, self.check_labels, [self.grid_state, self.zoom_state]
+        )
+        self.check_buttons.on_clicked(self.on_check_clicked)
+
+    def on_check_clicked(self, label):
+        status = dict(zip(self.check_labels, self.check_buttons.get_status()))
+        if label == "Grid":
+            self.grid_state = status["Grid"]
+            if self.on_toggle_grid is not None:
+                self.on_toggle_grid(self.grid_state)
+        elif label == "Zoom":
+            self.zoom_state = status["Zoom"]
+            if self.on_toggle_zoom is not None:
+                self.on_toggle_zoom(self.zoom_state)
+
+        self.func(self.i)
+        self.fig.canvas.draw_idle()
+
     def set_pos(self, i):
         self.i = int(self.slider.val)
         self.func(self.i)

WormView.py

@@ -1,4 +1,5 @@
 from matplotlib import pyplot as plt
+from matplotlib.ticker import MultipleLocator, FormatStrFormatter
 import numpy as np
 import math
 import os
@@ -8,6 +9,17 @@
 from SimpleWCON import SimpleWCON
 
 
+class HalfIntegerLocator(MultipleLocator):
+    """Locator placing ticks at half-integer values (..., -0.5, 0.5, 1.5, ...),
+    regardless of the current view limits (so it stays fixed while panning/zooming)."""
+
+    def __init__(self):
+        super().__init__(base=1.0)
+
+    def tick_values(self, vmin, vmax):
+        return super().tick_values(vmin - 0.5, vmax - 0.5) + 0.5
+
+
 def validate_file(file_path):
     if not os.path.exists(file_path):
         raise argparse.ArgumentTypeError(f"The file {file_path} does not exist.")
@@ -22,8 +34,17 @@ class WormView:
     head_plot = None
     times = None
 
-    def __init__(self, show_head=False):
+    zoom_to_worm = False
+    show_grid = False
+    zoom_side = 1.5  # in mm
+
+    def __init__(
+        self, show_head=False, zoom_to_worm=False, show_grid=False, zoom_side=1.5
+    ):
         self.show_head = show_head
+        self.zoom_to_worm = zoom_to_worm
+        self.show_grid = show_grid
+        self.zoom_side = zoom_side
         print(" - Initializing WormView")
 
     def get_perimeter(self, x, y, r):
@@ -78,6 +99,50 @@ def reset(self):
         self.head_plot = None
         plt.close("all")
 
+    def set_full_view_limits(self):
+        factor = 0.05
+        if abs(self.wcon.xmax - self.wcon.xmin) > abs(self.wcon.ymax - self.wcon.ymin):
+            side_x = abs(self.wcon.xmax - self.wcon.xmin)
+            self.ax.set_xlim(
+                [self.wcon.xmin - side_x * factor, self.wcon.xmax + side_x * factor]
+            )
+            mid = (self.wcon.ymax + self.wcon.ymin) / 2
+            self.ax.set_ylim(
+                [mid - side_x * (0.5 + factor), mid + side_x * (0.5 + factor)]
+            )
+        else:
+            side_y = abs(self.wcon.ymax - self.wcon.ymin)
+            self.ax.set_ylim(
+                [self.wcon.ymin - side_y * factor, self.wcon.ymax + side_y * factor]
+            )
+            mid = (self.wcon.xmax + self.wcon.xmin) / 2
+            self.ax.set_xlim(
+                [mid - side_y * (0.5 + factor), mid + side_y * (0.5 + factor)]
+            )
+
+    def set_zoom_to_worm(self, value):
+        print(" - Setting zoom_to_worm: %s" % value)
+        self.zoom_to_worm = value
+        if not value:
+            self.set_full_view_limits()
+
+    def set_show_grid(self, value):
+        print(" - Setting show_grid: %s" % value)
+        self.show_grid = value
+        if value:
+            self.ax.xaxis.set_major_locator(MultipleLocator(1))
+            self.ax.yaxis.set_major_locator(MultipleLocator(1))
+            self.ax.xaxis.set_minor_locator(HalfIntegerLocator())
+            self.ax.yaxis.set_minor_locator(HalfIntegerLocator())
+            self.ax.xaxis.set_minor_formatter(FormatStrFormatter("%.1f"))
+            self.ax.yaxis.set_minor_formatter(FormatStrFormatter("%.1f"))
+            self.ax.tick_params(axis="both", which="minor", labelsize=8)
+            self.ax.grid(True, which="major", linestyle="-", linewidth=0.5, alpha=0.7)
+            self.ax.grid(True, which="minor", linestyle=":", linewidth=0.5, alpha=0.7)
+        else:
+            self.ax.grid(False, which="major")
+            self.ax.grid(False, which="minor")
+
     def get_plot(self, args):
         # global times, t_units, x, y, px, py, ax
 
@@ -91,21 +156,10 @@ def get_plot(self, args):
         self.ax.set_xlabel("x (%s)" % self.wcon.x_units_used)
         self.ax.set_ylabel("y (%s)" % self.wcon.y_units_used)
 
-        factor = 0.05
-        if abs(self.wcon.xmax - self.wcon.xmin) > abs(self.wcon.ymax - self.wcon.ymin):
-            side = abs(self.wcon.xmax - self.wcon.xmin)
-            self.ax.set_xlim(
-                [self.wcon.xmin - side * factor, self.wcon.xmax + side * factor]
-            )
-            mid = (self.wcon.ymax + self.wcon.ymin) / 2
-            self.ax.set_ylim([mid - side * (0.5 + factor), mid + side * (0.5 + factor)])
-        else:
-            side = abs(self.wcon.ymax - self.wcon.ymin)
-            self.ax.set_ylim(
-                [self.wcon.ymin - side * factor, self.wcon.ymax + side * factor]
-            )
-            mid = (self.wcon.xmax + self.wcon.xmin) / 2
-            self.ax.set_xlim([mid - side * (0.5 + factor), mid + side * (0.5 + factor)])
+        if not self.zoom_to_worm:
+            self.set_full_view_limits()
+
+        self.set_show_grid(self.show_grid)
 
         if "@CelegansNeuromechanicalGaitModulation" in self.wcon.extras:
             center_x_arr = self.wcon.extras["@CelegansNeuromechanicalGaitModulation"][
@@ -196,6 +250,21 @@ def update(self, ti):
             else:
                 self.head_plot.set_data([self.wcon.x[0, ti]], [self.wcon.y[0, ti]])
 
+        if self.zoom_to_worm:
+            xi = self.wcon.x[:, ti]
+            yi = self.wcon.y[:, ti]
+            if self.px is not None and self.py is not None:
+                xi = np.concatenate([xi, self.px[:, ti]])
+                yi = np.concatenate([yi, self.py[:, ti]])
+
+            cx = (np.nanmax(xi) + np.nanmin(xi)) / 2
+            cy = (np.nanmax(yi) + np.nanmin(yi)) / 2
+
+            half = self.zoom_side / 2
+
+            self.ax.set_xlim([cx - half, cx + half])
+            self.ax.set_ylim([cy - half, cy + half])
+
         return self.midline_plot, self.perimeter_plot, self.head_plot
 
 
@@ -239,6 +308,18 @@ def parse_args():
         action="store_true",
         help="Show the head of the worm.",
     )
+    parser.add_argument(
+        "-zoom",
+        "--zoom_to_worm",
+        action="store_true",
+        help="Zoom the view to the extent of the worm.",
+    )
+    parser.add_argument(
+        "-grid",
+        "--show_grid",
+        action="store_true",
+        help="Show the grid in the view.",
+    )
 
     args = parser.parse_args()
 
@@ -252,7 +333,11 @@ def main():
 
     print(" - Arguments parsed: %s" % args)
 
-    wv = WormView(show_head=args.show_head)
+    wv = WormView(
+        show_head=args.show_head,
+        zoom_to_worm=args.zoom_to_worm,
+        show_grid=args.show_grid,
+    )
 
     fig, ax = wv.get_plot(args)
 
@@ -266,8 +351,17 @@ def update(ti):
         maxi=len(wv.wcon.times) - 1,
         times=[t for t in wv.wcon.times],
         t_units=wv.wcon.t_units,
+        grid_state=wv.show_grid,
+        zoom_state=wv.zoom_to_worm,
+        on_toggle_grid=wv.set_show_grid,
+        on_toggle_zoom=wv.set_zoom_to_worm,
     )
 
+    if args.nogui:
+        # Checkboxes are only useful for interactive control; don't bake them
+        # into the frames used for the saved movie.
+        anim.check_buttons.ax.set_visible(False)
+
     if not args.nogui:
         plt.show()
     else:

app.py

@@ -20,11 +20,11 @@ def make_args(wcon_file):
     )
 
 
-def load_view(wcon_file, show_head):
+def load_view(wcon_file, show_head, show_grid=False, zoom_to_worm=False):
     """Load and parse the WCON file once and build the figure. The result is
     kept in session_state so the (potentially large) file is parsed only when
     the file or options change, not on every frame."""
-    wv = WormView(show_head=show_head)
+    wv = WormView(show_head=show_head, show_grid=show_grid, zoom_to_worm=zoom_to_worm)
     wv.get_plot(make_args(wcon_file))
     return wv
 
@@ -33,13 +33,17 @@ def load_view(wcon_file, show_head):
 
 wcon_file = st.text_input("WCON file", value=DEFAULT_WCON)
 show_head = st.checkbox("Show head", value=False)
+show_grid = st.checkbox("Show grid", value=False)
+zoom_to_worm = st.checkbox("Zoom to worm", value=False)
 
 # Load only when the file or options change; otherwise reuse the cached view.
-key = (wcon_file, show_head)
+key = (wcon_file, show_head, show_grid, zoom_to_worm)
 if st.session_state.get("loaded_key") != key:
     try:
         with st.spinner(f"Loading {wcon_file}..."):
-            st.session_state.wv = load_view(wcon_file, show_head)
+            st.session_state.wv = load_view(
+                wcon_file, show_head, show_grid, zoom_to_worm
+            )
     except Exception as e:
         st.error(f"Could not load '{wcon_file}': {e}")
         st.stop()
@@ -54,15 +58,41 @@ def load_view(wcon_file, show_head):
 st.session_state.setdefault("running", False)
 
 # --- Controls (rendered before any rerun so they are always clickable) ---
-c1, c2, c3, c4 = st.columns(4)
-if c1.button("Play", use_container_width=True):
+c1, c2, c3, c4, c5, c6, c7 = st.columns(7)
+
+# one back
+if c1.button("$\u29cf$", use_container_width=True):
+    st.session_state.running = False
+    st.session_state.t -= 1
+
+# play backwards
+if c2.button("$\u25c0$", use_container_width=True):
     st.session_state.running = True
-if c2.button("Pause", use_container_width=True):
+    st.session_state.t -= 1
+    st.session_state.step = -1
+
+# stop
+if c3.button("$\u25a0$", use_container_width=True):
     st.session_state.running = False
-if c3.button("Step", use_container_width=True):
+
+# play
+if c4.button("$\u25b6$", use_container_width=True):
+    st.session_state.running = True
+    st.session_state.t += 1
+    st.session_state.step = 1
+
+# fast forward
+if c5.button("$\u25b6\u25b6$", use_container_width=True):
+    st.session_state.running = True
+    st.session_state.t += 10
+    st.session_state.step = 10
+
+# step forward
+if c6.button("$\u29d0$", use_container_width=True):
     st.session_state.running = False
     st.session_state.t += 1
-if c4.button("Reset", use_container_width=True):
+
+if c7.button("Reset", use_container_width=True):
     st.session_state.running = False
     st.session_state.t = 0
 
@@ -86,6 +116,6 @@ def load_view(wcon_file, show_head):
     if ti >= n_frames - 1:
         st.session_state.running = False  # reached the end
     else:
-        st.session_state.t += 1
+        st.session_state.t += st.session_state.step
         time.sleep(FRAME_DELAY)
         st.rerun()