Skip to content

refactor: update reward size references #51

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
arjunsridhar12345 merged 6 commits into from
Jul 27, 2026
Merged

refactor: update reward size references #51

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
97e9e24
build: update lock file
arjunsridhar12345 Jul 24, 2026
f1cf383
refactor: use updated reward size reference and minor updates
arjunsridhar12345 Jul 25, 2026
0d05b35
test: update tests
arjunsridhar12345 Jul 25, 2026
851be30
docs: update trial table docs
arjunsridhar12345 Jul 25, 2026
0164c10
refactor: remove unnecessary None guards for required Trial model
arjunsridhar12345 Jul 27, 2026
491805b
test: update tests
arjunsridhar12345 Jul 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 4 additions & 16 deletions docs/trials_table_mapping.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,21 +42,6 @@ Columns are grouped by the raw source they map from.
| `block_beta`, `block_duration`, `block_min`, `block_max` | `block_length` |
| `delay_beta`, `delay_duration`, `delay_min`, `delay_max` | `quiescent_duration_key` (scalar distribution, so no beta/min/max) |

### From `task_logic_input` (under `task_parameters`)

| Trials column | Source field |
| --- | --- |
| `reward_size_left` | `task_parameters.reward_size.left_value_volume` — the reward volume (uL) at the left port. |
| `reward_size_right` | `task_parameters.reward_size.right_value_volume` — the reward volume (uL) at the right port. |

> **Note:** `reward_size` is read from the task parameters, not the trial
> generator, so it is populated even when no summarising generator is resolved.
> The acquisition system can in principle vary reward size per trial, but the
> current data format only exposes a single session-level value, so these
> columns are constant across trials. They are **required** (non-nullable): a
> missing `TaskLogic` stream raises rather than silently producing null reward
> sizes when there are trials to build.

### From `TrialMetrics.json` (`SoftwareEvents` stream)

| Trials column | Mapping |
Expand All @@ -81,14 +66,16 @@ Columns are grouped by the raw source they map from.
| `response_duration` | `response_deadline_duration`. |
| `reward_consumption_duration` | `Trial -> reward_consumption_duration`. |
| `reward_probabilityL` / `reward_probabilityR` | The **block** probability from `Trial -> metadata -> p_reward_left` / `p_reward_right`. The top-level `trial.p_reward_left` / `p_reward_right` is the per-trial probability, not the block probability, so it is not used here. `None` when the trial or its metadata is missing. |
| `reward_size_left` | `Trial -> reward_size.left` — the reward volume (uL) at the left port. Defaults to `2.0` when not set on the trial. `None` when the trial is missing. |
| `reward_size_right` | `Trial -> reward_size.right` — the reward volume (uL) at the right port. Defaults to `2.0` when not set on the trial. `None` when the trial is missing. |
| `rewarded_historyL` / `rewarded_historyR` | Filter `is_rewarded == True`, then on `is_right_choice`. |

### From `TrialGeneratorSpec.json` (`SoftwareEvents` stream)

| Trials column | Mapping |
| --- | --- |
| `base_reward_probability_sum` | If `type == "CoupledTrialGenerator"`, look at `reward_probability_parameters`. |
| `min_reward_each_block` | Present when `type == "CoupledTrialGenerator"`; otherwise `None`. |
| `min_reward_each_block` | Present when `type == "CoupledWarmupTrialGenerator"` (has `min_block_reward`); otherwise `None`. |

### From `QuiescentPeriod.json` (`SoftwareEvents` stream)

Expand Down Expand Up @@ -155,3 +142,4 @@ These were mapped during exploration but are no longer in scope:
| 2026-06-17 | `auto_waterL` / `auto_waterR` now encode no auto-response (`is_auto_reward_right` is `None`) and missing trials as `0` instead of `NULL`. The columns are non-nullable (`int`, default `0`). |
| 2026-06-20 | Added `reward_size_left` / `reward_size_right` (reward volume in uL) from `task_parameters.reward_size`, and `side_bias` from the per-trial `TrialMetrics` event (`bias` field). |
| 2026-06-20 | `reward_probabilityL` / `reward_probabilityR` now read the block probability from `trial.metadata.p_reward_left` / `p_reward_right` instead of the top-level per-trial `trial.p_reward_left` / `p_reward_right`. |
| 2026-07-24 | `reward_size_left` / `reward_size_right` moved from session-level `task_parameters.reward_size` to per-trial `Trial.reward_size` (fields `.left` / `.right`). The columns are now nullable — `None` when the trial is missing. A missing `TaskLogic` stream no longer raises; session distribution columns are simply null. `min_reward_each_block` moved from `CoupledTrialGenerator` to `CoupledWarmupTrialGenerator`. |
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ readme = "README.md"
version = "0.1.0"

dependencies = [
"aind-behavior-dynamic-foraging[data] @ git+https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git@baab12133b22f599c1ba0583260eca9eca216cc0",
"aind-behavior-dynamic-foraging[data] @ git+https://github.com/AllenNeuralDynamics/Aind.Behavior.DynamicForaging.git@ac5ddbf909c9375b9e8875d6a5f90796cfa98653",
"ipykernel",
]

Expand Down
88 changes: 34 additions & 54 deletions src/dynamic_foraging_processing/processing/_trial_table.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,7 @@ def _animal_response(payload: t.Any) -> int:
return 1 if bool(choice) else 0

@staticmethod
def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]:
def _parse_outcome(payload: t.Any) -> TrialOutcome:
"""Parse a ``TrialOutcome`` software-event payload into its domain model.

Parameters
Expand All @@ -247,11 +247,16 @@ def _parse_outcome(payload: t.Any) -> t.Optional[TrialOutcome]:

Returns
-------
TrialOutcome or None
The parsed model, or ``None`` if ``payload`` is empty.
TrialOutcome
The parsed model.

Raises
------
ValueError
If ``payload`` is ``None``.
"""
if payload is None:
return None
raise ValueError("TrialOutcome payload is required but received None.")
if isinstance(payload, TrialOutcome):
return payload
if isinstance(payload, str):
Expand Down Expand Up @@ -320,7 +325,7 @@ def _rewarded_history(
return is_rewarded and (is_right_choice is is_right)

@staticmethod
def _is_baited(trial: t.Optional[Trial], *, is_right: bool) -> bool:
def _is_baited(trial: Trial, *, is_right: bool) -> bool:
"""Return whether the requested lickport is baited on this trial.

A port is "baited" when reward is guaranteed there (its reward
Expand All @@ -343,16 +348,15 @@ def _is_baited(trial: t.Optional[Trial], *, is_right: bool) -> bool:

Parameters
----------
trial : Trial or None
The per-trial task-logic model, or ``None`` when unavailable.
trial : Trial
The per-trial task-logic model.
is_right : bool
``True`` for the right port, ``False`` for the left port.

Returns
-------
bool
Whether the requested side is baited. A missing ``trial`` is treated
as not baited (``False``).
Whether the requested side is baited.

Examples
--------
Expand All @@ -374,8 +378,6 @@ def _is_baited(trial: t.Optional[Trial], *, is_right: bool) -> bool:
>>> TrialTableBuilder._is_baited(trial, is_right=False)
False
"""
if trial is None:
return False
auto = trial.is_auto_reward_right
if is_right:
# Right stays baited unless the animal was auto-responded right.
Expand All @@ -384,20 +386,19 @@ def _is_baited(trial: t.Optional[Trial], *, is_right: bool) -> bool:
return trial.p_reward_left == 1 and auto in (None, True)

@staticmethod
def _auto_water(trial: t.Optional[Trial], *, is_right: bool) -> int:
def _auto_water(trial: Trial, *, is_right: bool) -> int:
"""Encode autowater for a side from ``is_auto_reward_right``.

Returns ``1`` if the auto response was to the requested side, else ``0``.
A missing trial or no auto-response (``is_auto_reward_right`` is
``None``) counts as no autowater (``0``). ``is_right`` is ``True`` for
right.
No auto-response (``is_auto_reward_right`` is ``None``) counts as no
autowater (``0``). ``is_right`` is ``True`` for right.
"""
if trial is None or trial.is_auto_reward_right is None:
if trial.is_auto_reward_right is None:
return 0
return int(trial.is_auto_reward_right is is_right)

@staticmethod
def _block_reward_probability(trial: t.Optional[Trial], *, is_right: bool) -> t.Optional[float]:
def _block_reward_probability(trial: Trial, *, is_right: bool) -> t.Optional[float]:
"""Return the block reward probability for a side from the trial metadata.

The top-level ``trial.p_reward_left/right`` is the *per-trial* probability;
Expand All @@ -406,18 +407,17 @@ def _block_reward_probability(trial: t.Optional[Trial], *, is_right: bool) -> t.

Parameters
----------
trial : Trial or None
The per-trial task-logic model, or ``None`` when unavailable.
trial : Trial
The per-trial task-logic model.
is_right : bool
``True`` for the right port, ``False`` for the left port.

Returns
-------
float or None
The block reward probability, or ``None`` when the trial or its
metadata is unavailable.
The block reward probability, or ``None`` when metadata is unavailable.
"""
if trial is None or trial.metadata is None:
if trial.metadata is None:
return None
return trial.metadata.p_reward_right if is_right else trial.metadata.p_reward_left

Expand Down Expand Up @@ -483,14 +483,6 @@ def _session_columns(self, task_logic: AindDynamicForagingTaskLogic) -> t.Dict[s
if task_logic is None:
return columns

# Reward volumes live on the task parameters, not the trial generator, so
# populate them before the generator resolution (which may bail out).
# Known limitation: the acquisition system can vary reward size per trial,
# but the current data format only exposes a single session-level value.
reward_size = task_logic.task_parameters.reward_size
columns["reward_size_left"] = reward_size.left_value_volume
columns["reward_size_right"] = reward_size.right_value_volume

generator = self._summary_generator(task_logic.task_parameters.trial_generator)
if generator is None:
return columns
Expand Down Expand Up @@ -519,7 +511,7 @@ def _session_columns(self, task_logic: AindDynamicForagingTaskLogic) -> t.Dict[s
delay_max=delay_max,
base_reward_probability_sum=base_reward_sum,
)
# ``min_block_reward`` is coupled-only; uncoupled generators omit it.
# ``min_block_reward`` is warmup-generator-only; main coupled generators omit it.
if hasattr(generator, "min_block_reward"):
columns["min_reward_each_block"] = generator.min_block_reward
return columns
Expand Down Expand Up @@ -566,7 +558,7 @@ def _lickspout_columns(
def _build_row(
self,
*,
outcome: t.Optional[TrialOutcome],
outcome: TrialOutcome,
start: float,
stop: float,
response: t.Any,
Expand All @@ -578,9 +570,9 @@ def _build_row(
lickspout: t.Dict[str, t.Optional[float]],
) -> TrialConfig:
"""Assemble a single ``TrialConfig`` from aligned per-trial inputs."""
trial = outcome.trial if outcome is not None else None
is_right_choice = outcome.is_right_choice if outcome is not None else None
is_rewarded = bool(outcome.is_rewarded) if outcome is not None else False
trial = outcome.trial
is_right_choice = outcome.is_right_choice
is_rewarded = bool(outcome.is_rewarded)

return TrialConfig(
start_time=start,
Expand All @@ -596,13 +588,13 @@ def _build_row(
bait_right=self._is_baited(trial, is_right=True),
reward_probabilityL=self._block_reward_probability(trial, is_right=False),
reward_probabilityR=self._block_reward_probability(trial, is_right=True),
reward_size_left=trial.reward_size.left,
reward_size_right=trial.reward_size.right,
side_bias=side_bias,
response_duration=trial.response_deadline_duration if trial is not None else None,
reward_consumption_duration=(
trial.reward_consumption_duration if trial is not None else None
),
ITI_duration=trial.inter_trial_interval_duration if trial is not None else None,
delay_duration=trial.quiescence_period_duration if trial is not None else None,
response_duration=trial.response_deadline_duration,
reward_consumption_duration=trial.reward_consumption_duration,
ITI_duration=trial.inter_trial_interval_duration,
delay_duration=trial.quiescence_period_duration,
auto_waterL=self._auto_water(trial, is_right=False),
auto_waterR=self._auto_water(trial, is_right=True),
**session,
Expand Down Expand Up @@ -669,9 +661,7 @@ def build(self) -> pd.DataFrame:
Raises
------
ValueError
If the ``TaskLogic`` stream is missing while there are trials to
build (the required reward-size columns are sourced from it), or if
``raise_on_error`` is ``True`` and a per-trial stream length
If ``raise_on_error`` is ``True`` and a per-trial stream length
disagrees with the ``TrialOutcome`` trial count.
"""
outcomes = self._load("Behavior", "SoftwareEvents", "TrialOutcome")
Expand All @@ -696,16 +686,6 @@ def build(self) -> pd.DataFrame:
# Guard the positional alignment before we pair streams by index.
n_trials = len(outcome_payloads)

# Reward size is sourced from the task logic and is a required column, so
# a missing TaskLogic stream cannot yield a valid table when there are
# trials to build. Surface that clearly rather than failing later with a
# cryptic per-row validation error.
if n_trials > 0 and task_logic is None:
raise ValueError(
"TaskLogic stream is required to build the trials table "
f"(reward sizes are sourced from it) but it failed to load for {n_trials} trials."
)

warnings = self._check_aligned(
n_trials,
{
Expand Down
Loading