update unit & probe widget to be on the same widget

Author: JaerongA

element_array_ephys/plotting/widget.py

@@ -4,12 +4,12 @@
 import matplotlib.pyplot as plt
 import pathlib
 import typing as T
-from IPython.display import display
+from IPython.display import display, clear_output
 from .. import ephys_no_curation as ephys
 from skimage import io
 
 
-# Widgets
+# Build selection widgets
 probe_dropdown_wg = widgets.Dropdown(
     options=ephys.CuratedClustering & ephys_report.ProbeLevelReport.fetch("KEY"),
     description="Select Probe Insertion : ",
@@ -30,6 +30,13 @@
     style={"description_width": "100px"},
 )
 
+shank_toggle_wg = widgets.ToggleButtons(
+    options=(ephys_report.ProbeLevelReport & probe_dropdown_wg.value).fetch("shank"),
+    description="Select Shank : ",
+    layout=widgets.Layout(width="auto"),
+    style={"button_width": "50px"},
+)
+
 unit_dropdown_wg = widgets.Dropdown(
     options=(
         (ephys_report.UnitLevelReport & probe_dropdown_wg.value)
@@ -43,48 +50,98 @@
     style={"description_width": "100px"},
 )
 
-# Update shank & unit dropdown according to probe change
-def update_shank_options(change):
-    probe_key = change.new
-    shank_dropdown_wg.options = (ephys_report.ProbeLevelReport & probe_key.value).fetch(
-        "shank"
-    )
 
+def probe_widget():
+    def plot_probe_figure(probe_key, shank):
 
-def update_unit_dropdown(change):
-    probe_key = change.new
-    unit_dropdown_wg.options = (
-        (
-            (ephys_report.UnitLevelReport & probe_key.value)
-            & "cluster_quality_label='good'"
-        ).fetch("unit"),
-    )
+        fig_name = (
+            ephys_report.ProbeLevelReport & probe_key & f"shank={shank}"
+        ).fetch1("drift_map_plot")
 
+        # Constants
+        img_width = 2000
+        img_height = 1000
+        scale_factor = 0.5
 
-probe_dropdown_wg.observe(update_shank_options, "value")
-probe_dropdown_wg.observe(update_unit_dropdown, "value")
+        img = io.imread(fig_name)
+        probe_fig = px.imshow(img)
 
+        # Configure other layout
+        probe_fig.update_layout(
+            width=img_width * scale_factor,
+            height=img_height * scale_factor,
+            margin={"l": 50, "r": 0, "t": 0, "b": 0},
+            hovermode=False,
+            xaxis_visible=False,
+            yaxis_visible=False,
+        )
 
-def plot_probe_level_report(probe_key: T.Dict[str, T.Any]) -> None:
-    fig_name = (ephys_report.ProbeLevelReport & probe_key).fetch1("drift_map_plot")
-    fig, ax = plt.subplots(1, 1, figsize=(10, 5))
-    img = mpimg.imread(fig_name)
-    ax.imshow(img)
-    ax.axis("off")
-    plt.show()
-    pathlib.Path(fig_name).unlink()
+        pathlib.Path(fig_name).unlink()
+        return go.FigureWidget(probe_fig)
 
+    def probe_dropdown_evt(change):
+        probe_key = change.new
+        shank_dropdown_wg.options = (
+            ephys_report.ProbeLevelReport & probe_key.value
+        ).fetch("shank")
+        unit_dropdown_wg.options = (
+            (
+                (ephys_report.UnitLevelReport & probe_key.value)
+                & "cluster_quality_label='good'"
+            ).fetch("unit"),
+        )
+        clear_output()
+        display(
+            widgets.VBox(
+                [
+                    probe_dropdown_wg,
+                    shank_dropdown_wg,
+                    plot_probe_figure(probe_key, shank_dropdown_wg.value),
+                ]
+            )
+        )
 
-def plot_unit_level_report(unit_key: T.Dict[str, T.Any]) -> None:
-    waveform_fig, autocorrelogram_fig, depth_waveform_fig = (
-        ephys_report.UnitLevelReport & unit_key
-    ).fetch1("waveform_plotly", "autocorrelogram_plotly", "depth_waveform_plotly")
-    waveform_fig = go.FigureWidget(waveform_fig)
-    autocorrelogram_fig = go.FigureWidget(autocorrelogram_fig)
-    depth_waveform_fig = go.FigureWidget(depth_waveform_fig)
+    probe_dropdown_wg.observe(probe_dropdown_evt, "value")
 
-    display(
-        widgets.HBox(
-            [widgets.VBox([waveform_fig, autocorrelogram_fig]), depth_waveform_fig]
-        )
+    return widgets.VBox(
+        [
+            probe_dropdown_wg,
+            shank_dropdown_wg,
+            plot_probe_figure(probe_dropdown_wg.value, shank_dropdown_wg.value),
+        ]
     )
+
+
+def unit_widget():
+    def plot_unit_figure(unit):
+        # Build a unit widgets
+        waveform_fig, autocorrelogram_fig, depth_waveform_fig = (
+            ephys_report.UnitLevelReport & probe_dropdown_wg.value & f"unit={unit}"
+        ).fetch1("waveform_plotly", "autocorrelogram_plotly", "depth_waveform_plotly")
+        waveform_fig = go.FigureWidget(waveform_fig).update_layout(
+            width=300, height=300
+        )
+        autocorrelogram_fig = go.FigureWidget(autocorrelogram_fig).update_layout(
+            width=300, height=300
+        )
+        depth_waveform_fig = go.FigureWidget(depth_waveform_fig)
+        depth_waveform_fig.update_layout(
+            width=300,
+            height=600,
+            autosize=False,
+            margin={"l": 0, "r": 0, "t": 100, "b": 100},
+        )
+
+        unit_fig_wg = widgets.HBox(
+            [widgets.VBox([waveform_fig, autocorrelogram_fig]), depth_waveform_fig],
+            layout=widgets.Layout(margin="0 0 0 100px"),
+        )
+        return unit_fig_wg
+
+    def unit_dropdown_evt(change):
+        unit = change.new
+        clear_output()
+        display(unit_dropdown_wg, plot_unit_figure(unit))
+
+    unit_dropdown_wg.observe(unit_dropdown_evt, "value")
+    return widgets.VBox([unit_dropdown_wg, plot_unit_figure(unit_dropdown_wg.value)])