fix: manual water annotation qc plot

src/dynamic_foraging_processing/qc/processed/plots.py

@@ -27,6 +27,11 @@
 #: this far below, so the pair reads as arrows straddling y=0.
 _MOVE_MARKER_OFFSET = 0.06
 
+#: Left edge of every trial-indexed x-axis. Padded below trial 0 so events on the
+#: first trial (e.g. manual water given at the start of a session) are drawn clear
+#: of the y-axis spine instead of being hidden behind it.
+_TRIAL_AXIS_LEFT = -1
+
 
 def plot_lick_intervals(
     left_lick_times: np.ndarray, right_lick_times: np.ndarray, results_folder: str
@@ -179,7 +184,7 @@ def _add_bias_plot(
     trials = np.arange(len(bias))
     ax.plot(trials, bias, "k", linewidth=2)
     if len(bias):
-        ax.set_xlim([0, len(bias)])
+        ax.set_xlim([_TRIAL_AXIS_LEFT, len(bias)])
 
     plotted = False
     if anti_bias_right_water is not None:
@@ -333,7 +338,12 @@ def _add_behavior_plot(
     manual_right_times: t.Optional[np.ndarray],
     go_cue_times: t.Optional[np.ndarray],
 ) -> None:
-    """Draw the per-trial behavior raster (choices, rewards, water)."""
+    """Draw the per-trial behavior raster (choices, rewards, water).
+
+    Each water type gets its own row per side: manual water is not autowater, so
+    sharing a band with it (as an earlier version did) made the two
+    indistinguishable and left manual deliveries reading as mislabeled autowater.
+    """
     choices = np.asarray(animal_response)
     ax.vlines(np.where(choices == 1)[0], 0.8, 1, linewidth=1, color="gray", label="Choice")
     ax.vlines(np.where(choices == 0)[0], 0, 0.2, linewidth=1, color="gray")
@@ -346,24 +356,6 @@ def _add_behavior_plot(
         right_rewards = np.where(np.asarray(rewarded_right))[0]
         ax.vlines(right_rewards, 1, 1.2, linewidth=1, color="black")
 
-    if manual_right_times is not None and go_cue_times is not None:
-        ax.vlines(
-            _time_to_trial_index(go_cue_times, manual_right_times),
-            1.2,
-            1.4,
-            linewidth=1,
-            color="blue",
-            label="Manual Water",
-        )
-    if manual_left_times is not None and go_cue_times is not None:
-        ax.vlines(
-            _time_to_trial_index(go_cue_times, manual_left_times),
-            -0.4,
-            -0.2,
-            linewidth=1,
-            color="blue",
-        )
-
     if autowater_right is not None:
         ax.vlines(
             np.where(np.asarray(autowater_right) == 1)[0],
@@ -382,21 +374,50 @@ def _add_behavior_plot(
             color="cyan",
         )
 
-    ax.set_ylim([-0.4, 1.4])
-    ax.set_xlim([0, len(choices)])
+    # Manual water sits outside the autowater rows and is dashed, so it reads as
+    # distinct from the solid earned and auto ticks even where colour alone is
+    # hard to judge. It is labelled only when the session actually has
+    # deliveries, so the legend never claims manual water for a session that had
+    # none.
+    manual_label: t.Optional[str] = "Manual Water"
+    for times, bottom, top in (
+        (manual_right_times, 1.4, 1.6),
+        (manual_left_times, -0.6, -0.4),
+    ):
+        if times is None or go_cue_times is None:
+            continue
+        trial_indices = _time_to_trial_index(go_cue_times, times)
+        if not trial_indices:
+            continue
+        ax.vlines(
+            trial_indices,
+            bottom,
+            top,
+            linewidth=1,
+            color="blue",
+            linestyles="dashed",
+            label=manual_label,
+        )
+        manual_label = None
+
+    ax.set_ylim([-0.6, 1.6])
+    ax.set_xlim([_TRIAL_AXIS_LEFT, len(choices)])
     ax.set_xlabel("Trial #")
     ax.set_yticks(
-        [-0.3, -0.1, 0.1, 0.5, 0.9, 1.1, 1.3],
+        [-0.5, -0.3, -0.1, 0.1, 0.5, 0.9, 1.1, 1.3, 1.5],
         labels=[
+            "L Manual Water",
             "L Auto Water",
             "L Reward",
             "L Choice",
             "Ignore",
             "R Choice",
             "R Reward",
             "R Auto Water",
+            "R Manual Water",
         ],
     )
+    _legend_outside(ax)
 
 
 def _add_reward_probabilities(
@@ -511,11 +532,11 @@ def plot_side_bias(
 
     # Align the x-axis across every panel so trials line up vertically. The
     # panels are all indexed by trial, but some auto-scale (adding margins) while
-    # others set [0, N]; pin them all to a common [0, n_trials].
+    # others set their own limits; pin them all to a common range.
     n_trials = max(len(np.asarray(side_bias)), len(np.asarray(animal_response)))
     if n_trials:
         for axis in ax:
-            axis.set_xlim([0, n_trials])
+            axis.set_xlim([_TRIAL_AXIS_LEFT, n_trials])
 
     fig.savefig(Path(results_folder) / SIDE_BIAS_PLOT, dpi=300, bbox_inches="tight")
     plt.close(fig)

tests/test_qc/test_plots.py

@@ -132,6 +132,50 @@ def test_add_lickspout_position_plot_splits_automatic_and_manual_moves():
     plt.close(fig)
 
 
+def test_add_behavior_plot_gives_manual_water_its_own_rows():
+    """Manual water sits on its own rows, outside the autowater bands."""
+    fig, ax = plt.subplots()
+    _plots._add_behavior_plot(
+        ax,
+        np.array([0, 1, 0, 1]),
+        rewarded_left=None,
+        rewarded_right=None,
+        autowater_left=np.array([1, 0, 0, 0]),
+        autowater_right=np.array([0, 0, 0, 1]),
+        manual_left_times=np.array([1.6]),
+        manual_right_times=np.array([2.6]),
+        go_cue_times=np.array([0.5, 1.5, 2.5, 3.5]),
+    )
+    labels = [text.get_text() for text in ax.get_legend().get_texts()]
+    # A single "Manual Water" entry covers both sides.
+    assert labels.count("Manual Water") == 1
+    assert "Auto Water" in labels
+    tick_labels = [text.get_text() for text in ax.get_yticklabels()]
+    assert "L Manual Water" in tick_labels and "R Manual Water" in tick_labels
+    # The manual rows are outside the autowater bands, so the two never overlap.
+    assert ax.get_ylim() == (-0.6, 1.6)
+    plt.close(fig)
+
+
+def test_add_behavior_plot_omits_manual_water_legend_when_absent():
+    """A session with no manual deliveries gets no manual-water legend entry."""
+    fig, ax = plt.subplots()
+    _plots._add_behavior_plot(
+        ax,
+        np.array([0, 1]),
+        rewarded_left=None,
+        rewarded_right=None,
+        autowater_left=None,
+        autowater_right=None,
+        manual_left_times=np.array([]),
+        manual_right_times=np.array([]),
+        go_cue_times=np.array([0.5, 1.5]),
+    )
+    labels = [text.get_text() for text in ax.get_legend().get_texts()]
+    assert "Manual Water" not in labels
+    plt.close(fig)
+
+
 def test_plot_side_bias_minimal_inputs(tmp_path):
     """With only choices supplied, the optional panels are skipped cleanly."""
     name = _plots.plot_side_bias(np.array([]), np.array([]), str(tmp_path))