diff --git a/docs/trials_table_mapping.md b/docs/trials_table_mapping.md index d53a059..2e22f0b 100644 --- a/docs/trials_table_mapping.md +++ b/docs/trials_table_mapping.md @@ -123,11 +123,11 @@ seconds, far too coarse for the ~tens-of-ms valve pulse. | --- | --- | | `goCue_start_time` | `PlaySoundOrFrequency` `WRITE` message. | -### From `InitialManipulatorPosition` (software event) +### From `HarpManipulator` `AccumulatedSteps` (+ `InputSchemas.Rig`) | Trials column | Mapping | | --- | --- | -| `lickspout_positions` | `data` field. | +| `lickspout_position_x` / `y1` / `y2` / `z` | Per-motor cumulative microstep count from the `AccumulatedSteps` stream, converted to millimetres via the rig manipulator calibration (`full_step_to_mm / microstep_resolution`) and re-referenced to the session-start position (displacement **relative to session start**, mm). The manipulator is a continuously-sampled hardware value, so — like the go cue — each trial takes the sample within its `[start_time, stop_time)` window nearest the start. `Motor{i}` drives `Axis(i + 1)` (X, Y1, Y2, Z). `None` when no sample falls in the trial window. The rig and `AccumulatedSteps` streams are required inputs (`build` raises if either is missing with trials present). | ### From `trainer_state.json` and `acquisition.json` (autoTrain — can be disregarded) @@ -155,3 +155,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-22 | `lickspout_position_x` / `y1` / `y2` / `z` now derive from the `HarpManipulator` `AccumulatedSteps` stream (microsteps → mm via the `InputSchemas.Rig` manipulator calibration, `full_step_to_mm / microstep_resolution`), sampled per trial via the closest sample in the `[start_time, stop_time)` window and re-referenced to the session-start position (displacement relative to session start, mm), replacing the static `InitialManipulatorPosition` software event. `Motor{i}` maps to `Axis(i + 1)` (X, Y1, Y2, Z). The rig and `AccumulatedSteps` streams are required when there are trials (`build` raises if either is missing). Column descriptions corrected from `um` to `mm`. | diff --git a/src/dynamic_foraging_processing/processing/_trial_table.py b/src/dynamic_foraging_processing/processing/_trial_table.py index b83fb63..acd6baf 100644 --- a/src/dynamic_foraging_processing/processing/_trial_table.py +++ b/src/dynamic_foraging_processing/processing/_trial_table.py @@ -10,6 +10,7 @@ import numpy as np import pandas as pd +from aind_behavior_dynamic_foraging.rig import AindDynamicForagingRig from aind_behavior_dynamic_foraging.task_logic import AindDynamicForagingTaskLogic from aind_behavior_dynamic_foraging.task_logic.trial_generators import TrialGeneratorSpec from aind_behavior_dynamic_foraging.task_logic.trial_models import ( @@ -24,6 +25,12 @@ logger = logging.getLogger(__name__) +# The four manipulator axes, ordered so that index ``i`` is the axis driven +# by ``Motor{i}`` in the ``AccumulatedSteps`` stream: ``Motor{i}`` maps to +# ``Axis(i + 1)`` (Axis.X=1, Y1=2, Y2=3, Z=4), whose names match the +# ``lickspout_position_*`` suffixes and the ``full_step_to_mm`` attributes. +_MANIPULATOR_AXES = ("x", "y1", "y2", "z") + class TrialTableBuilder: """Builds the NWB ``trials`` table from a dynamic foraging ``Dataset``. @@ -524,41 +531,140 @@ def _session_columns(self, task_logic: AindDynamicForagingTaskLogic) -> t.Dict[s columns["min_reward_each_block"] = generator.min_block_reward return columns - def _lickspout_columns( - self, manipulator: t.Optional[pd.DataFrame] + def _manipulator_mm_per_step(self, rig: AindDynamicForagingRig) -> t.Dict[str, float]: + """Return millimetres travelled per accumulated step, keyed by axis. + + ``AccumulatedSteps`` counts *microsteps*. The physical distance per + microstep is the full-step-to-mm calibration divided by the microstep + resolution (e.g. ``MICROSTEP8`` → 8 microsteps per full step), both read + from the rig's manipulator calibration. + + Parameters + ---------- + rig : AindDynamicForagingRig + The ``InputSchemas.Rig`` stream data. The rig is a required input + schema, so ``build`` guarantees it is present before this is called. + + Returns + ------- + dict of str to float + ``{axis: mm_per_step}`` for ``x`` / ``y1`` / ``y2`` / ``z``. + + Raises + ------ + ValueError + If the manipulator calibration is missing one of the four axes. + """ + calibration = rig.manipulator.calibration + resolution_by_axis = { + config.axis.name.lower(): config.microstep_resolution + for config in calibration.axis_configuration + } + mm_per_step: t.Dict[str, float] = {} + for axis in _MANIPULATOR_AXES: + resolution = resolution_by_axis.get(axis) + if resolution is None: + raise ValueError( + f"Manipulator calibration is missing axis {axis!r}; " + "cannot derive lickspout position." + ) + # Resolution enum names encode the divisor (MICROSTEP8 -> 8); the + # enum *value* (0..3) does not, so parse from the name. + microsteps = int(resolution.name.replace("MICROSTEP", "")) + mm_per_step[axis] = getattr(calibration.full_step_to_mm, axis) / microsteps + return mm_per_step + + def _manipulator_positions( + self, + accumulated_steps: pd.DataFrame, + rig: AindDynamicForagingRig, + ) -> pd.DataFrame: + """Return a time-indexed frame of lickspout position (mm) over the session. + + Built from the ``HarpManipulator`` ``AccumulatedSteps`` stream: each + ``EVENT`` row's per-motor microstep count is converted to millimetres via + the rig calibration. ``AccumulatedSteps`` is absolute (zero-referenced at + the homing position), so the count maps directly to position with no + offset. The frame is then re-referenced to the session start (the first + sample subtracted from every row), so the stored values are displacement + *relative to session start* — the units the QC plot expects. + + Parameters + ---------- + accumulated_steps : pandas.DataFrame + The ``AccumulatedSteps`` stream data (``Motor0``..``Motor3`` columns). + Required; guaranteed present by ``build``. + rig : AindDynamicForagingRig + The ``InputSchemas.Rig`` stream data (required; guaranteed present by + ``build``). + + Returns + ------- + pandas.DataFrame + Columns ``lickspout_position_x`` / ``y1`` / ``y2`` / ``z`` indexed by + time (sorted ascending), relative to the session-start position. + + Raises + ------ + ValueError + If the stream has no ``EVENT`` rows, is missing a motor column, or the + rig calibration is missing an axis. + """ + mm_per_step = self._manipulator_mm_per_step(rig) + events = accumulated_steps + if "MessageType" in events.columns: + events = events[events["MessageType"] == "EVENT"] + if len(events) == 0: + raise ValueError( + "AccumulatedSteps stream has no EVENT rows; cannot derive lickspout position." + ) + events = events.sort_index() + columns: t.Dict[str, np.ndarray] = {} + for index, axis in enumerate(_MANIPULATOR_AXES): + motor = f"Motor{index}" + if motor not in events.columns: + raise ValueError( + f"AccumulatedSteps stream is missing column {motor}; " + "cannot derive lickspout position." + ) + columns[f"lickspout_position_{axis}"] = ( + events[motor].to_numpy(dtype=float) * mm_per_step[axis] + ) + positions = pd.DataFrame(columns, index=events.index) + # Re-reference to session start so the stored values are displacement + # from the first sample (previously done in the plot as ``values[0]``). + return positions - positions.iloc[0] + + def _sample_lickspout( + self, positions: pd.DataFrame, start: float, stop: float ) -> t.Dict[str, t.Optional[float]]: - """Return lickspout position columns from ``InitialManipulatorPosition``. + """Return the lickspout position sampled within a trial's window. - The event's ``data`` payload carries the x / y1 / y2 / z positions. These - are the *initial* manipulator coordinates recorded once at experiment - start, so the columns are currently constant across trials. + The manipulator position is a continuously-sampled hardware value, so — + like the go cue — each trial takes the sample within its ``[start, stop)`` + window nearest the start (see :meth:`_closest_time_in_window`). - Known limitation: the manipulator can move within a session (e.g. the - anti-bias intervention shifts the lickspouts horizontally), which this - static initial position does not capture. + Parameters + ---------- + positions : pandas.DataFrame + The time-indexed position frame from :meth:`_manipulator_positions`. + start, stop : float + The trial window bounds (seconds); ``start`` inclusive, ``stop`` + exclusive. May be ``NaN`` when unaligned. - TODO: think about how to represent - this if it becomes relevant. The ideal solution would be to use the harp files - to track the manipulator position over time + Returns + ------- + dict of str to (float or None) + The four ``lickspout_position_*`` values, all ``None`` for a trial + whose window contains no manipulator sample. """ - keys = ( - "lickspout_position_x", - "lickspout_position_y1", - "lickspout_position_y2", - "lickspout_position_z", - ) - empty = {key: None for key in keys} - payloads = self._event_payloads(manipulator) - if not payloads: - return empty - payload = payloads[0] - components = ("x", "y1", "y2", "z") - if isinstance(payload, dict): - return {key: payload.get(comp) for key, comp in zip(keys, components)} - if isinstance(payload, (list, tuple)) and len(payload) == len(keys): - return {key: payload[idx] for idx, key in enumerate(keys)} - logger.warning("Unrecognized InitialManipulatorPosition payload; leaving lickspout null.") - return empty + keys = tuple(f"lickspout_position_{axis}" for axis in _MANIPULATOR_AXES) + times = positions.index.to_numpy(dtype=float) + sample_time = self._closest_time_in_window(times, start, stop) + if sample_time is None: + return {key: None for key in keys} + row = positions.loc[sample_time] + return {key: (None if pd.isna(row[key]) else float(row[key])) for key in keys} # ------------------------------------------------------------------ # # Per-trial assembly @@ -683,7 +789,8 @@ def build(self) -> pd.DataFrame: pulse_supply_left = self._load("Behavior", "HarpBehavior", "PulseSupplyPort0") pulse_supply_right = self._load("Behavior", "HarpBehavior", "PulseSupplyPort1") go_cue = self._load("Behavior", "HarpSoundCard", "PlaySoundOrFrequency") - manipulator = self._load("Behavior", "SoftwareEvents", "InitialManipulatorPosition") + accumulated_steps = self._load("Behavior", "HarpManipulator", "AccumulatedSteps") + rig = self._load("Behavior", "InputSchemas", "Rig") task_logic = self._load("Behavior", "InputSchemas", "TaskLogic") # Per-trial streams: one payload/timestamp per trial, aligned by index. @@ -706,6 +813,20 @@ def build(self) -> pd.DataFrame: f"(reward sizes are sourced from it) but it failed to load for {n_trials} trials." ) + # The rig and AccumulatedSteps streams are the required sources for the + # lickspout positions, so a missing one cannot yield valid position + # columns when there are trials to build. + if n_trials > 0 and rig is None: + raise ValueError( + "Rig stream is required to build the trials table " + f"(lickspout positions are sourced from it) but it failed to load for {n_trials} trials." + ) + if n_trials > 0 and accumulated_steps is None: + raise ValueError( + "AccumulatedSteps stream is required to build the trials table " + f"(lickspout positions are sourced from it) but it failed to load for {n_trials} trials." + ) + warnings = self._check_aligned( n_trials, { @@ -732,7 +853,11 @@ def build(self) -> pd.DataFrame: go_cue_times = self._write_times(go_cue) session = self._session_columns(task_logic) - lickspout = self._lickspout_columns(manipulator) + # When there are trials, the guards above guarantee both sources are + # present; with no trials the frame goes unused (the loop does not run). + lickspout_positions = ( + self._manipulator_positions(accumulated_steps, rig) if n_trials else None + ) rows: t.List[TrialConfig] = [] for i, outcome_payload in enumerate(outcome_payloads): @@ -743,6 +868,7 @@ def build(self) -> pd.DataFrame: stop = float(stop_times[i]) if i < stop_times.size else np.nan response = response_payloads[i] if i < len(response_payloads) else None side_bias = self._side_bias(metric_payloads[i] if i < len(metric_payloads) else None) + lickspout = self._sample_lickspout(lickspout_positions, start, stop) rows.append( self._build_row( diff --git a/src/dynamic_foraging_processing/processing/models/trial_config.py b/src/dynamic_foraging_processing/processing/models/trial_config.py index 8c20bcc..ed3425b 100644 --- a/src/dynamic_foraging_processing/processing/models/trial_config.py +++ b/src/dynamic_foraging_processing/processing/models/trial_config.py @@ -142,18 +142,20 @@ class TrialConfig(BaseModel): # --- lickspout_position (mapping's `lickspout_positions` -> these four components) --- lickspout_position_x: Optional[float] = Field( - default=None, description="x position (um) of the lickspout position (left-right)" + default=None, + description="x lickspout position (mm), relative to session start (left-right)", ) lickspout_position_y1: Optional[float] = Field( default=None, - description="y1 position (um) of the left lickspout position (forward-backward)", + description="y1 left lickspout position (mm), relative to session start (forward-backward)", ) lickspout_position_y2: Optional[float] = Field( default=None, - description="y2 position (um) of the right lickspout position (forward-backward)", + description="y2 right lickspout position (mm), relative to session start (forward-backward)", ) lickspout_position_z: Optional[float] = Field( - default=None, description="z position (um) of the lickspout position (up-down)" + default=None, + description="z lickspout position (mm), relative to session start (up-down)", ) @classmethod diff --git a/src/dynamic_foraging_processing/qc/processed/plots.py b/src/dynamic_foraging_processing/qc/processed/plots.py index d04c941..e46e335 100644 --- a/src/dynamic_foraging_processing/qc/processed/plots.py +++ b/src/dynamic_foraging_processing/qc/processed/plots.py @@ -122,7 +122,9 @@ def _add_lickspout_position_plot( if position is None or len(position) == 0: continue values = np.asarray(position, dtype=float) - ax.plot(values - values[0], color, label=label) + # Values already come in relative to session start (normalized by the + # trial-table builder), so plot them directly. + ax.plot(values, color, label=label) plotted = True if plotted: ax.legend() @@ -303,6 +305,14 @@ def plot_side_bias( ) _add_reward_probabilities(ax[3], reward_probability_left, reward_probability_right) + # 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]. + 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]) + fig.savefig(Path(results_folder) / SIDE_BIAS_PLOT, dpi=300, bbox_inches="tight") plt.close(fig) return SIDE_BIAS_PLOT diff --git a/src/dynamic_foraging_processing/qc/processed/results.py b/src/dynamic_foraging_processing/qc/processed/results.py index 762890f..a263f1a 100644 --- a/src/dynamic_foraging_processing/qc/processed/results.py +++ b/src/dynamic_foraging_processing/qc/processed/results.py @@ -22,15 +22,15 @@ from dynamic_foraging_processing.qc.processed.plots import plot_lick_intervals, plot_side_bias # Logical input -> trials-table column name. Centralized so the mapping is easy -# to correct against the trial-table builder; ``side_bias`` and the -# ``lickspout_*`` arrays are not yet pinned down in trials_table_mapping.md. +# to correct against the trial-table builder. The lickspout columns are the +# ``lickspout_position_*`` fields documented in trials_table_mapping.md. _COLUMNS = { "animal_response": "animal_response", "side_bias": "side_bias", - "lickspout_x": "lickspout_x", - "lickspout_y1": "lickspout_y1", - "lickspout_y2": "lickspout_y2", - "lickspout_z": "lickspout_z", + "lickspout_x": "lickspout_position_x", + "lickspout_y1": "lickspout_position_y1", + "lickspout_y2": "lickspout_position_y2", + "lickspout_z": "lickspout_position_z", "rewarded_left": "rewarded_historyL", "rewarded_right": "rewarded_historyR", "reward_probability_left": "reward_probabilityL", diff --git a/tests/test_processing/test_trial_table.py b/tests/test_processing/test_trial_table.py index e4b91a6..e5debe0 100644 --- a/tests/test_processing/test_trial_table.py +++ b/tests/test_processing/test_trial_table.py @@ -1,5 +1,7 @@ """Tests for ``dynamic_foraging_processing.processing._trial_table``.""" +from types import SimpleNamespace + import numpy as np import pandas as pd import pytest @@ -15,6 +17,7 @@ UncoupledTrialGeneratorSpec, ) from aind_behavior_dynamic_foraging.task_logic.trial_models import TrialOutcome +from aind_behavior_services.rig.aind_manipulator import Axis, MicrostepResolution from aind_behavior_services.task.distributions import ( ExponentialDistribution, ExponentialDistributionParameters, @@ -157,6 +160,35 @@ def _pulse_supply(column, value_ms): ) +def _rig(step_mm=0.01, resolution=MicrostepResolution.MICROSTEP8): + """Build a rig stub exposing the manipulator calibration the builder reads. + + A ``SimpleNamespace`` mirrors the ``AindDynamicForagingRig`` attribute path + (``manipulator.calibration.full_step_to_mm`` / ``axis_configuration``) used by + the builder, with real ``Axis`` / ``MicrostepResolution`` enums. + """ + full_step_to_mm = SimpleNamespace(x=step_mm, y1=step_mm, y2=step_mm, z=step_mm) + axis_configuration = [ + SimpleNamespace(axis=axis, microstep_resolution=resolution) + for axis in (Axis.X, Axis.Y1, Axis.Y2, Axis.Z) + ] + calibration = SimpleNamespace( + full_step_to_mm=full_step_to_mm, axis_configuration=axis_configuration + ) + return SimpleNamespace(manipulator=SimpleNamespace(calibration=calibration)) + + +def _accumulated_steps(times, motor_steps): + """Build an ``AccumulatedSteps`` frame (Motor0..Motor3 microstep counts). + + ``motor_steps`` is a sequence of ``(m0, m1, m2, m3)`` per timestamp; all rows + are ``EVENT`` messages, mirroring the HarpManipulator stream. + """ + columns = {f"Motor{i}": [row[i] for row in motor_steps] for i in range(4)} + columns["MessageType"] = ["EVENT"] * len(times) + return pd.DataFrame(columns, index=pd.Index(times, name="Time")) + + def _full_dataset(): """Assemble a two-trial fake dataset covering the common path.""" software_events = _Node( @@ -184,9 +216,6 @@ def _full_dataset(): [10.5, 20.5], [{"Item1": 10.5, "Item2": True}, {"Item1": 20.5, "Item2": None}] ) ), - "InitialManipulatorPosition": _Stream( - _events([5.0], [{"x": 1.0, "y1": 2.0, "y2": 3.0, "z": 4.0}]) - ), "TrialMetrics": _Stream(_events([10.2, 20.2], [{"bias": 0.3}, {"bias": None}])), } ) @@ -209,7 +238,20 @@ def _full_dataset(): ) } ), - "InputSchemas": _Node({"TaskLogic": _Stream(_task_logic())}), + "HarpManipulator": _Node( + { + # 0.00125 mm/microstep (0.01 / MICROSTEP8). One sample falls + # inside each trial window ([10, 20) and [20, 30)); x moves + # 10.0 -> 15.0 between the two trials. + "AccumulatedSteps": _Stream( + _accumulated_steps( + [10.5, 20.5], + [(8000, 1600, 2400, 4000), (12000, 1600, 2400, 4000)], + ) + ) + } + ), + "InputSchemas": _Node({"TaskLogic": _Stream(_task_logic()), "Rig": _Stream(_rig())}), } ) return _Node({"Behavior": behavior}) @@ -278,9 +320,14 @@ def test_build_full_dataset(): assert first["min_reward_each_block"] == 2 assert first["base_reward_probability_sum"] == pytest.approx(0.8) - # Lickspout positions from InitialManipulatorPosition. - assert first["lickspout_position_x"] == 1.0 - assert first["lickspout_position_z"] == 4.0 + # Lickspout positions from AccumulatedSteps (microsteps * 0.00125 mm), + # sampled at each trial start and re-referenced to session start. The first + # sample is the baseline (all zero); x then moves +5.0 mm for the second trial. + assert first["lickspout_position_x"] == 0.0 + assert first["lickspout_position_y1"] == 0.0 + assert first["lickspout_position_y2"] == 0.0 + assert first["lickspout_position_z"] == 0.0 + assert second["lickspout_position_x"] == 5.0 # Session-level reward volumes (uL) from task_parameters.reward_size. assert first["reward_size_left"] == 2.0 @@ -301,6 +348,24 @@ def test_build_raises_when_task_logic_missing_with_trials(): TrialTableBuilder(dataset).build() +def test_build_raises_when_rig_missing_with_trials(): + """A missing Rig stream is an error when there are trials (lickspout position needs it).""" + dataset = _full_dataset() + input_schemas = dataset.children["Behavior"].children["InputSchemas"] + input_schemas.children["Rig"] = _FailedStream() + with pytest.raises(ValueError, match="Rig stream is required"): + TrialTableBuilder(dataset).build() + + +def test_build_raises_when_accumulated_steps_missing_with_trials(): + """A missing AccumulatedSteps stream is an error when there are trials.""" + dataset = _full_dataset() + harp_manipulator = dataset.children["Behavior"].children["HarpManipulator"] + harp_manipulator.children["AccumulatedSteps"] = _FailedStream() + with pytest.raises(ValueError, match="AccumulatedSteps stream is required"): + TrialTableBuilder(dataset).build() + + def test_build_empty_dataset_returns_empty_frame(): """Missing streams yield an empty table with the full column set.""" table = TrialTableBuilder(_Node({}), raise_on_error=False).build() @@ -589,23 +654,109 @@ def test_parse_outcome_accepts_json_string(): assert isinstance(TrialTableBuilder._parse_outcome(payload), TrialOutcome) -def test_lickspout_columns_from_list_payload(): - """A list payload is mapped positionally to x/y1/y2/z.""" +def test_manipulator_mm_per_step_from_calibration(): + """mm-per-microstep is full_step_to_mm / microstep resolution, per axis.""" builder = TrialTableBuilder(_Node({})) - columns = builder._lickspout_columns(_events([1.0], [[10.0, 20.0, 30.0, 40.0]])) - assert columns["lickspout_position_x"] == 10.0 - assert columns["lickspout_position_y2"] == 30.0 + mm_per_step = builder._manipulator_mm_per_step(_rig()) + assert mm_per_step == {"x": 0.00125, "y1": 0.00125, "y2": 0.00125, "z": 0.00125} -def test_lickspout_columns_unrecognized_payload_is_null(): - """An unrecognized payload leaves all lickspout positions ``None``.""" +def test_manipulator_mm_per_step_respects_resolution(): + """A finer microstep resolution shrinks the distance per step.""" builder = TrialTableBuilder(_Node({})) - columns = builder._lickspout_columns(_events([1.0], ["unexpected"])) - assert all(value is None for value in columns.values()) + mm_per_step = builder._manipulator_mm_per_step( + _rig(step_mm=0.01, resolution=MicrostepResolution.MICROSTEP16) + ) + assert mm_per_step["x"] == pytest.approx(0.01 / 16) + + +def test_manipulator_mm_per_step_missing_axis_raises(): + """A calibration lacking one of the four axes raises ``ValueError``.""" + builder = TrialTableBuilder(_Node({})) + rig = _rig() + # Drop the Z axis configuration. + rig.manipulator.calibration.axis_configuration = [ + config + for config in rig.manipulator.calibration.axis_configuration + if config.axis is not Axis.Z + ] + with pytest.raises(ValueError, match="missing axis 'z'"): + builder._manipulator_mm_per_step(rig) + + +def test_manipulator_positions_converts_steps_to_mm_relative_to_start(): + """EVENT rows convert to mm and are re-referenced to the session-start sample.""" + builder = TrialTableBuilder(_Node({})) + steps = _accumulated_steps([8.0, 15.0], [(8000, 1600, 2400, 4000), (12000, 1600, 2400, 4000)]) + positions = builder._manipulator_positions(steps, _rig()) + assert list(positions.index) == [8.0, 15.0] + # First sample is the baseline (zeroed); x then moves +5.0 mm (4000 steps). + assert positions.loc[8.0, "lickspout_position_x"] == 0.0 + assert positions.loc[15.0, "lickspout_position_x"] == 5.0 + assert positions.loc[8.0, "lickspout_position_z"] == 0.0 + + +def test_manipulator_positions_filters_non_event_rows(): + """Only ``EVENT`` rows contribute to the position frame.""" + builder = TrialTableBuilder(_Node({})) + steps = _accumulated_steps([8.0, 15.0], [(8000, 1600, 2400, 4000), (12000, 1600, 2400, 4000)]) + steps.loc[15.0, "MessageType"] = "WRITE" + positions = builder._manipulator_positions(steps, _rig()) + assert list(positions.index) == [8.0] + + +def test_manipulator_positions_all_non_event_raises(): + """A stream with no ``EVENT`` rows raises ``ValueError``.""" + builder = TrialTableBuilder(_Node({})) + steps = _accumulated_steps([8.0], [(8000, 1600, 2400, 4000)]) + steps["MessageType"] = "WRITE" + with pytest.raises(ValueError, match="no EVENT rows"): + builder._manipulator_positions(steps, _rig()) + + +def test_manipulator_positions_missing_motor_column_raises(): + """A stream missing a Motor column raises ``ValueError``.""" + builder = TrialTableBuilder(_Node({})) + steps = _accumulated_steps([8.0], [(8000, 1600, 2400, 4000)]).drop(columns=["Motor3"]) + with pytest.raises(ValueError, match="missing column Motor3"): + builder._manipulator_positions(steps, _rig()) + + +def test_manipulator_positions_incomplete_calibration_raises(): + """Steps present but a calibration axis missing raises ``ValueError``.""" + builder = TrialTableBuilder(_Node({})) + steps = _accumulated_steps([8.0], [(8000, 1600, 2400, 4000)]) + rig = _rig() + rig.manipulator.calibration.axis_configuration = [ + config + for config in rig.manipulator.calibration.axis_configuration + if config.axis is not Axis.Z + ] + with pytest.raises(ValueError, match="missing axis 'z'"): + builder._manipulator_positions(steps, rig) + + +def _two_sample_positions(builder): + """Build a position frame relative to session start: t=8.0 (x=0.0), t=15.0 (x=5.0).""" + return builder._manipulator_positions( + _accumulated_steps([8.0, 15.0], [(8000, 1600, 2400, 4000), (12000, 1600, 2400, 4000)]), + _rig(), + ) + + +def test_sample_lickspout_picks_sample_nearest_start_in_window(): + """A trial takes the in-window sample nearest its start.""" + builder = TrialTableBuilder(_Node({})) + positions = _two_sample_positions(builder) + # Window [7.0, 12.0) contains only the t=8.0 baseline sample (0.0). + assert builder._sample_lickspout(positions, 7.0, 12.0)["lickspout_position_x"] == 0.0 + # Window [14.0, 20.0) contains only the t=15.0 sample (+5.0 mm). + assert builder._sample_lickspout(positions, 14.0, 20.0)["lickspout_position_x"] == 5.0 -def test_lickspout_columns_missing_stream_is_null(): - """A missing manipulator stream leaves all lickspout positions ``None``.""" +def test_sample_lickspout_no_sample_in_window_is_null(): + """A trial whose window contains no sample yields ``None`` columns.""" builder = TrialTableBuilder(_Node({})) - columns = builder._lickspout_columns(None) + positions = _two_sample_positions(builder) + columns = builder._sample_lickspout(positions, 20.0, 30.0) assert all(value is None for value in columns.values()) diff --git a/tests/test_qc/test_builder.py b/tests/test_qc/test_builder.py index 69507e8..bbf97e5 100644 --- a/tests/test_qc/test_builder.py +++ b/tests/test_qc/test_builder.py @@ -29,6 +29,22 @@ def test_behavior_qc_results_without_plots(): assert results[0].name == "average side bias" +def test_lickspout_columns_map_to_trial_table_names(): + """The lickspout mapping resolves to the ``lickspout_position_*`` trial columns.""" + trials = pd.DataFrame( + { + "lickspout_position_x": [16.0, 16.0], + "lickspout_position_y1": [9.6, 9.6], + "lickspout_position_y2": [9.7, 9.7], + "lickspout_position_z": [17.25, 17.25], + } + ) + for key in ("lickspout_x", "lickspout_y1", "lickspout_y2", "lickspout_z"): + column = _results._column(trials, key) + assert column is not None, key + assert column[0] == trials[_results._COLUMNS[key]][0] + + def test_behavior_qc_results_writes_plots(tmp_path): """Supplying a results folder writes both behavior plots.""" trials = pd.DataFrame(