fix: reward annotations still incorrect

docs/trials_table_mapping.md

@@ -48,23 +48,19 @@ from "every time the valve opened":
 reward-consumption and ITI periods, so a delivery can sit nearer the *previous*
 trial's outcome and inherit its `is_auto_reward_right`. The valve opens within
 milliseconds of the response, so the response anchors each delivery to its own
-trial. Verified on `864253_2026-08-11_12-52-08`, `864253_2026-08-04_12-54-55`,
-and `864253_2026-08-10_13-09-30`: nearest-`Response` agrees with the trial whose
-`[quiescent_start_time, ITI_start_time)` window contains the delivery on all 859
-valve openings, and every `earned` delivery follows a lick on that same port
-(median 2.8 ms, max 32 ms).
-
-**The series is reward-keyed**: a delivery whose matched trial reports
-`is_rewarded=False` is dropped rather than annotated, so the retained count
-equals the metadata mapper's `sum(is_rewarded)` (248 / 298 / 288 on those three
-sessions, against 261 / 298 / 300 valve openings). Manual water is
-experimenter-driven, unrelated to the trial outcome, and is never dropped. Free
-water is triggered immediately at the go cue and the trial then continues
-normally, so `is_rewarded` reports the outcome of the animal's *own choice* — a
-separate event from the free water. The consequence is that free water delivered
-on a trial that did not pay out is real water the animal received and is **not**
-in this series; the trials table's ungated `auto_water*` /`anti_bias_*` columns
-still record it.
+trial. Two independent checks back this: the nearest-`Response` trial agrees
+with the trial whose `[quiescent_start_time, ITI_start_time)` window contains the
+delivery, and every `earned` delivery follows a lick on that same port within a
+few milliseconds.
+
+**The series records every valve opening**: nothing is filtered out, so a
+delivery on a trial reporting `is_rewarded=False` is still annotated. Free water
+is triggered immediately at the go cue and the trial then continues normally, so
+`is_rewarded` reports the outcome of the animal's *own choice* — a separate
+event from the free water. A single trial can therefore contribute an `auto`
+delivery at the port the task watered and an `earned` delivery at the port the
+animal chose. The series is a water record rather than a reward-trial count, so
+its length is not expected to equal `sum(is_rewarded)` over the trials.
 
 ## Trials Table
 
@@ -98,7 +94,7 @@ Columns are grouped by the raw source they map from.
 | Trials column | Mapping |
 | --- | --- |
 | `auto_waterL` / `auto_waterR` | **Scheduled autowater only**: `1` when `trial.metadata.extra.is_autowater` is `True` **and** `is_auto_reward_right` points to that side. `0` otherwise, including when the trial's free water came from the anti-bias algorithm — that is reported by `anti_bias_left_water` / `anti_bias_right_water`. `is_auto_reward_right` is only the delivery *channel* (free water fired, and to which side); the mechanism comes from the metadata, so the two columns are mutually exclusive. Not gated on `is_rewarded`: the column records what the task did, and free water fires at the go cue regardless of how the animal's own choice resolves. Note this is narrower than the legacy `dynamic-foraging-task` column of the same name, which was the ungated channel ("Autowater given at Left", straight from `B_AutoWaterTrial`) and predates anti-bias water. |
-| `anti_bias_left_water` / `anti_bias_right_water` | Boolean. `True` when the anti-bias algorithm delivered a water intervention to that side — i.e. `trial.metadata.extra.is_bias_water_intervention` is `True` **and** `is_auto_reward_right` points to that side (`False` → left, `True` → right). The anti-bias water uses the same auto-response channel as scheduled autowater, so the `is_bias_water_intervention` flag is what distinguishes it and the two columns are mutually exclusive. `False` otherwise. Like `auto_water*`, **not** gated on `is_rewarded`: these columns record what the algorithm did, and the intervention fires at the go cue regardless of how the animal's own choice resolves. The reward-delivery series *is* reward-keyed, so this column can exceed the series' `auto` count. |
+| `anti_bias_left_water` / `anti_bias_right_water` | Boolean. `True` when the anti-bias algorithm delivered a water intervention to that side — i.e. `trial.metadata.extra.is_bias_water_intervention` is `True` **and** `is_auto_reward_right` points to that side (`False` → left, `True` → right). The anti-bias water uses the same auto-response channel as scheduled autowater, so the `is_bias_water_intervention` flag is what distinguishes it and the two columns are mutually exclusive. `False` otherwise. Like `auto_water*`, **not** gated on `is_rewarded`: these columns record what the algorithm did, and the intervention fires at the go cue regardless of how the animal's own choice resolves. The reward-delivery series is also ungated, so an intervention on a trial that did not pay out appears in both. |
 | `anti_bias_lickspout_movement` | Signed horizontal displacement (mm, positive is rightward) the anti-bias algorithm moved the lickspouts on this trial: `trial.lickspout_offset_delta` when `trial.metadata.extra.is_bias_stage_intervention` is `True`, else `0.0`. |
 | `bait_left` / `bait_right` | Boolean, read straight from `trial.metadata.extra.is_left_baited` / `is_right_baited` — the bait state the acquisition software reports for each port. `False` when the trial carries no extra metadata. |
 | `response_duration` | `response_deadline_duration`. |
@@ -227,8 +223,8 @@ These were mapped during exploration but are no longer in scope:
 | 2026-08-06 | Confirmed and documented that the legacy `delay_*` columns describe the acquisition software's **quiescence period**: `delay_start_time` is the `QuiescentPeriod` timestamp (always equal to the new `quiescent_start_time`) and `delay_duration` / `delay_beta` / `delay_min` / `delay_max` summarize `quiescence_period_duration`. `delay_duration` is the *configured* duration — each lick restarts the quiescent period, so the realized `quiescent_stop_time - quiescent_start_time` can be longer. Column descriptions updated accordingly. |
 | 2026-08-12 | `rewarded_historyL` / `rewarded_historyR` now record **earned** reward only: an auto-reward trial (`is_auto_reward_right` set to either side) is `False` on *both* sides, since `TrialOutcome.is_rewarded` is `True` for autowater too and that water is already reported by `auto_waterL` / `auto_waterR`. This matches the `earned` / `automatic` split used for the NWB reward-delivery annotations. |
 | 2026-08-12 | `min_reward_each_block` is now `0` rather than `NULL` when the trial generator exposes no `min_block_reward` — no per-block minimum is a floor of zero, not an unknown. The column is non-nullable (`float`, default `0`). |
-| 2026-08-17 | Reward-delivery annotations now match each delivery to its trial by the `Response` software-event timestamp rather than the `TrialOutcome` timestamp. `TrialOutcome` fires at the *end* of a trial (after the reward-consumption and ITI periods), so a delivery could land nearer the *previous* trial's outcome and inherit its `is_auto_reward_right`, flipping `earned` and `auto`. The valve opens within milliseconds of the response, so the response anchors the delivery to its own trial. Verified 0-mismatch against trial-window containment on 859 valve openings across three sessions. |
-| 2026-08-17 | The reward-delivery series is now reward-keyed: a delivery whose matched trial reports `is_rewarded=False` is dropped rather than annotated (manual water exempt), so the retained count equals the metadata mapper's `sum(is_rewarded)`. Free water fires at the go cue and the trial then continues normally, so `is_rewarded` describes the animal's own choice, not the free water — meaning free water on an unrewarded trial is real water the animal received and is excluded from this series. The trials table's `auto_water*` / `anti_bias_*` columns remain ungated and still record it. |
+| 2026-08-17 | Reward-delivery annotations now match each delivery to its trial by the `Response` software-event timestamp rather than the `TrialOutcome` timestamp. `TrialOutcome` fires at the *end* of a trial (after the reward-consumption and ITI periods), so a delivery could land nearer the *previous* trial's outcome and inherit its `is_auto_reward_right`, flipping `earned` and `auto`. The valve opens within milliseconds of the response, so the response anchors the delivery to its own trial. Verified 0-mismatch against trial-window containment. |
+| 2026-08-17 | The reward-delivery series annotates every valve opening; deliveries are no longer filtered on `is_rewarded`. Free water fires at the go cue and the trial then continues normally, so `is_rewarded` describes the outcome of the animal's own choice, not the water — a trial where the task waters one port and the animal earns reward at the other is valid and contributes one `auto` and one `earned` delivery. The series is therefore a complete water record rather than a reward-trial count, and its total is not expected to equal the metadata mapper's `sum(is_rewarded)`. |
 | 2026-08-17 | `auto_waterL` / `auto_waterR` now read `trial.metadata.extra.is_autowater` rather than the `is_auto_reward_right` channel, making them **scheduled autowater only** and mutually exclusive with `anti_bias_left_water` / `anti_bias_right_water`. `is_auto_reward_right` says free water fired and on which side but not what kind; the mechanism is in the metadata. Neither column is gated on `is_rewarded`, since both record what the task did. This is narrower than the legacy `dynamic-foraging-task` column of the same name, which was the ungated channel and predates anti-bias water. |
 | 2026-08-17 | The reward-delivery labels stay `earned` / `auto` / `manual`: free water is `auto` whatever mechanism produced it, so the series does not split scheduled autowater from anti-bias water. That split lives in the trials table. Consequence: the series' `auto` count tracks the channel while `auto_waterL` / `auto_waterR` track `is_autowater`, so the two are not expected to be equal. |
 | 2026-08-20 | `block_max` is now one below `block_length`'s configured maximum, which accounts for the floor applied upstream: a block is a whole number of trials, so the configured bound is never itself reachable. `block_min`, `block_beta`, and the `ITI_*` / `delay_*` bounds are unchanged — those durations are continuous and take no such adjustment. |

src/dynamic_foraging_processing/nwb/acquisition/acquisition_builder.py

@@ -203,9 +203,8 @@ def _reward_delivery_series(
 
         Only valve-open events (``port_column`` is truthy) are reward
         deliveries; the ``data`` field annotates each as earned, manual, or auto
-        via :func:`get_reward_deliveries`, which also drops deliveries on trials
-        that did not pay out, so the series reports reward rather than every
-        valve opening.
+        via :func:`get_reward_deliveries`. Every valve opening is reported, so
+        the series is a complete record of the water delivered at this port.
 
         Parameters
         ----------
@@ -237,7 +236,7 @@ def _reward_delivery_series(
         open_writes = writes[writes[port_column].fillna(False).astype(bool)]
         delivery_times = open_writes.index.to_numpy()
         manual_water_times = manual_water.index[manual_water["data"] == is_right].to_numpy()
-        delivery_times, annotations = get_reward_deliveries(
+        annotations = get_reward_deliveries(
             delivery_times,
             trial_outcomes,
             manual_water_times,

src/dynamic_foraging_processing/utils/rewards.py

@@ -61,8 +61,8 @@ def get_reward_deliveries(
     trial_outcome_df: pd.DataFrame,
     manual_water_times: np.ndarray,
     response_times: np.ndarray,
-) -> t.Tuple[np.ndarray, np.ndarray]:
-    """Get one lick port's reward deliveries, classified by how the water was given.
+) -> np.ndarray:
+    """Classify one lick port's reward deliveries by how the water was given.
 
     Annotates the deliveries of a single lick port. Each delivery is classified
     as follows, with ``manual`` taking precedence because manual water is not
@@ -80,17 +80,14 @@ def get_reward_deliveries(
       ``anti_bias_left_water``/``anti_bias_right_water``.
     - ``earned`` -- otherwise: water the animal worked for.
 
-    Deliveries on a trial reporting ``is_rewarded=False`` are dropped rather than
-    annotated, so the series reports only water that counted as reward. In
-    practice these are all free water: it is triggered immediately at the go cue
-    and the trial then continues normally, so a trial whose own choice did not
-    pay out still carries the delivery. Manual water is experimenter-driven and
-    is never dropped. The surviving timestamps are returned alongside their
-    annotations so the two stay aligned.
-
-    Note this makes the series reward-keyed rather than a complete record of the
-    hardware's valve openings: free water delivered on an unrewarded trial is
-    real water the animal received, and it is excluded here.
+    Every valve opening is annotated and none is filtered out, so the series is a
+    complete record of the water this port delivered. In particular a trial
+    reporting ``is_rewarded=False`` keeps its delivery: free water is triggered
+    immediately at the go cue and the trial then continues normally, so
+    ``is_rewarded`` reports the outcome of the animal's *own choice* -- a
+    separate event from the water being classified here. A trial can therefore
+    contribute an ``auto`` delivery on the side the task watered and an
+    ``earned`` delivery on the side the animal chose.
 
     Deliveries are matched to trials by the ``Response`` software-event
     timestamp: each delivery takes the annotation of the trial whose response is
@@ -121,11 +118,8 @@ def get_reward_deliveries(
     Returns
     -------
     numpy.ndarray
-        The retained reward-delivery timestamps: ``reward_delivery_times`` less
-        the deliveries on unrewarded trials.
-    numpy.ndarray
-        The matching annotations, one per retained timestamp, each ``"earned"``,
-        ``"auto"``, or ``"manual"``.
+        Array of the same shape as ``reward_delivery_times`` whose entries are
+        ``"earned"``, ``"auto"``, or ``"manual"``.
 
     Raises
     ------
@@ -142,18 +136,16 @@ def get_reward_deliveries(
 
     reward_times = np.asarray(reward_delivery_times)
     if reward_times.size == 0:
-        return reward_times, np.array([], dtype=object)
+        return np.array([], dtype=object)
 
     # Annotate each delivery from its originating trial: query with reward_times so we
     # get one trial position per reward delivery.
     trial_indices_in_reward_times = find_closest_timestamps(reward_times, response_times)
 
     annotated_rewards = []
-    is_unrewarded = []
     for trial_index in trial_indices_in_reward_times:
         outcome = _parse_outcome(trial_outcome_df.iloc[trial_index]["data"])
         annotated_rewards.append(_free_water_label(outcome.trial if outcome is not None else None))
-        is_unrewarded.append(outcome is not None and not outcome.is_rewarded)
 
     # Object dtype, not the inferred fixed-width string dtype: a run of only "auto"
     # and "earned" entries would be too narrow to hold "manual" and would truncate it.
@@ -164,14 +156,8 @@ def get_reward_deliveries(
     # manual-water software event to its closest reward delivery; the returned
     # positions index into reward_times, i.e. the deliveries that are manual.
     manual_water_times = np.asarray(manual_water_times)
-    manual_mask = np.zeros(reward_times.size, dtype=bool)
     if manual_water_times.size:
         manual_indices_in_reward_times = find_closest_timestamps(manual_water_times, reward_times)
-        manual_mask[manual_indices_in_reward_times] = True
-        annotated_rewards[manual_mask] = "manual"
-
-    # Downstream analysis is keyed on reward, so a delivery whose trial did not pay out
-    # is excluded. Manual water is experimenter-driven, unrelated to the trial's
-    # outcome, and keeps its delivery.
-    keep = ~(np.array(is_unrewarded, dtype=bool) & ~manual_mask)
-    return reward_times[keep], annotated_rewards[keep]
+        annotated_rewards[manual_indices_in_reward_times] = "manual"
+
+    return annotated_rewards