diff --git a/src/dynamic_foraging_processing/qc/processed/plots.py b/src/dynamic_foraging_processing/qc/processed/plots.py index 4f40673..65bd0d8 100644 --- a/src/dynamic_foraging_processing/qc/processed/plots.py +++ b/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,12 +374,39 @@ 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", @@ -395,8 +414,10 @@ def _add_behavior_plot( "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) diff --git a/tests/test_qc/test_plots.py b/tests/test_qc/test_plots.py index a157f27..2ddbb02 100644 --- a/tests/test_qc/test_plots.py +++ b/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))