Fix missing cue onset and stimulus for legacy Go/NoGo trials

CHANGELOG.md

@@ -7,6 +7,30 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- **Stimulus and cue onset of Go/NoGo trials in pre-0.5 Visiomode sessions.** In
+  sessions without a per-trial stimulus, misses (outcome `no_response`) were not
+  matched to the Go stimulus, so they had an empty `cue_onset` and `stim_*` fields.
+  In sessions whose trials record `{"target": {}, "distractor": {}}`, no trial had
+  a `cue_onset` or `stim_*` fields. In sessions that record per-trial target
+  parameters (e.g. the contrast of `variablecontrastmovinggrating`), precued
+  trials got a `cue_onset` after the trial had ended, NoGo trials got the previous
+  grating's contrast, and there was no `stim_id`. For all three, the shown stimulus
+  is now worked out from the outcome and response: precued trials have no
+  stimulus, and the per-trial contrast appears as `stim_trial_contrast` on Go trials
+  only (it was `target_trial_contrast` on every trial). Regressors now give misses
+  their Go stimulus and hold regressors, and sessions of the last two kinds no longer
+  fail with `KeyError: 'stim_id'`. The overall reaction time plot in session reports,
+  which only counts trials with a cue onset, was empty for the second kind and
+  included precued pushes for the third. Session and subject summaries were not
+  affected. Re-run `session` (and `regressors`) on these sessions to update their
+  trials CSVs, reports and regressors.
+- **Go/NoGo regressors for sessions with non-default stimuli.** `regressors` assumed
+  the Go stimulus was `movinggrating` and the NoGo stimulus `isoluminantgray`. It now
+  uses the session's target and distractor, so sessions with other stimuli (e.g.
+  `variablecontrastmovinggrating`) get stimulus regressors.
+
 ## [0.4.1] - 2026-10-02
 
 ### Fixed

src/visiomode_analysis/session/__init__.py

@@ -620,7 +620,8 @@ def generate_regressors(
     Args:
         trials_df (pd.DataFrame): A DataFrame containing trial data with columns for trial type,
             start time, and stop time.
-        metadata (dict): A dictionary containing session metadata.
+        metadata (dict): A dictionary containing session metadata. For Go/NoGo sessions, the target and distractor
+            in its `stimuli` are the Go and NoGo stimuli.
         regressor_timestamps_path (str): Path to a CSV or TXT file of ISO timestamps, or a mesoscopy H5 file with a
             `/timestamps_aligned` dataset. Typically corresponds to the frame times of an imaging session or other
             continuous recording.
@@ -661,7 +662,19 @@ def generate_regressors(
     outpath = f"{output_dir}{os.sep}sub-{metadata.get('animal_id')}_exp-{metadata.get('experiment')}_ses-{str(metadata.get('session_date')).replace('-', '')}_behaviour-{metadata.get('protocol')}_regressors.npz"
 
     if metadata.get("protocol") == "gonogo":
-        regressors, labels, trial_idx = rgr.generate_gonogo_regressors(trials_df, timestamps, trial_epoch_only=True)
+        # The Go and NoGo stimuli are the session's target and distractor; fall back to the generator's defaults for
+        # sessions whose spec doesn't name them.
+        stimulus_ids = {
+            arg: stim_id
+            for arg, stim_id in (
+                ("go_stim_id", metadata.get("stimuli", {}).get("target_id")),
+                ("nogo_stim_id", metadata.get("stimuli", {}).get("distractor_id")),
+            )
+            if stim_id
+        }
+        regressors, labels, trial_idx = rgr.generate_gonogo_regressors(
+            trials_df, timestamps, trial_epoch_only=True, **stimulus_ids
+        )
         np.savez(
             outpath,
             regressors=regressors,
@@ -869,6 +882,35 @@ def _normalise_no_response(trial: Any) -> Any:
     return trial
 
 
+def _reconstruct_gonogo_stimulus(trial: Any, metadata: dict) -> dict:
+    """Return the `stim_*` fields of the stimulus a Go/NoGo trial showed, for sessions whose trials don't record it.
+
+    The shown stimulus is inferred from the outcome and whether there was a response: the target for hits and misses
+    (`no_response`, or `incorrect` without a response), the distractor for false alarms and correct rejections, and
+    none for precued trials. Its fields come from the session-level `stimuli` metadata, plus, for sessions that record
+    per-trial parameters for both stimuli (`{"target": {...}, "distractor": {...}}`, e.g. the contrast of a variable
+    contrast grating), those of the shown stimulus.
+    """
+    responded = bool(trial.get("response"))
+    outcome = trial.get("outcome")
+    if (responded and outcome == "correct") or (not responded and outcome in ("incorrect", "no_response")):
+        shown = "target"
+    elif (responded and outcome == "incorrect") or (not responded and outcome == "correct"):
+        shown = "distractor"
+    else:
+        return {}
+
+    per_trial_params = trial.get("stimulus").get(shown) or {} if isinstance(trial.get("stimulus"), dict) else {}
+    return {
+        **{
+            f"stim_{key.replace(f'{shown}_', '')}": value
+            for key, value in metadata.get("stimuli", {}).items()
+            if key.startswith(f"{shown}_")
+        },
+        **{f"stim_{key}": value for key, value in per_trial_params.items()},
+    }
+
+
 def _flatten_trials(session: dict, metadata: dict) -> Iterator[dict]:
     session_start_time = datetime.datetime.fromisoformat(metadata.get("session_start_time", ""))
 
@@ -902,7 +944,15 @@ def _flatten_trials(session: dict, metadata: dict) -> Iterator[dict]:
         dist_y = trial.get("response").get("dist_y", 0) if trial.get("response") else None
 
         stimulus: dict = {}
-        if trial.get("stimulus"):
+        if (
+            metadata.get("protocol") == "gonogo"
+            and isinstance(trial.get("stimulus"), dict)
+            and "target" in trial.get("stimulus")
+        ):
+            # Some Visiomode versions before 0.5 record per-trial parameters for both Go/NoGo stimuli rather than
+            # which one was shown, so it's reconstructed as for sessions without a per-trial stimulus.
+            stimulus = _reconstruct_gonogo_stimulus(trial, metadata)
+        elif trial.get("stimulus"):
             if trial.get("stimulus") == "None":
                 stimulus = {}
             elif trial.get("stimulus").get("common_name"):
@@ -917,26 +967,7 @@ def _flatten_trials(session: dict, metadata: dict) -> Iterator[dict]:
                 stimulus = {**target_stim, **distractor_stim}
         else:  # handle older versions of visiomode
             if metadata.get("protocol") == "gonogo":
-                if (trial.get("response") and trial.get("outcome") == "correct") or (
-                    not trial.get("response") and trial.get("outcome") == "incorrect"
-                ):
-                    stimulus = {
-                        **{
-                            f"stim_{key.replace('target_', '')}": value
-                            for key, value in metadata.get("stimuli", {}).items()
-                            if key.startswith("target_")
-                        },
-                    }
-                elif (trial.get("response") and trial.get("outcome") == "incorrect") or (
-                    not trial.get("response") and trial.get("outcome") == "correct"
-                ):
-                    stimulus = {
-                        **{
-                            f"stim_{key.replace('distractor_', '')}": value
-                            for key, value in metadata.get("stimuli", {}).items()
-                            if key.startswith("distractor_")
-                        },
-                    }
+                stimulus = _reconstruct_gonogo_stimulus(trial, metadata)
             elif is_targetonly(metadata.get("protocol")):
                 if (trial.get("response") and (trial.get("outcome") == "correct")) or (
                     trial.get("outcome") == "no_response"

tests/test_flatten_trials.py

@@ -1,8 +1,9 @@
 """Tests for `session._flatten_trials`'s handling of different Visiomode JSON formats: sessions
 recorded before an explicit `stimulus`/`sdt_type` field existed on each trial, where the presented
 stimulus and signal-detection classification instead have to be reconstructed from the trial's
-`outcome`/`response` and the session-level `stimuli` metadata, and newer (0.5+) sessions that record
-a trial without a response as a response named "none".
+`outcome`/`response` and the session-level `stimuli` metadata (as they also are for Go/NoGo sessions
+whose per-trial stimulus holds both stimuli's parameters but not which was shown), and newer (0.5+)
+sessions that record a trial without a response as a response named "none".
 """
 
 import numpy as np
@@ -68,6 +69,94 @@ def test_legacy_gonogo_miss_reconstructs_target_stimulus_and_classifies_as_miss(
     assert trial["stim_id"] == "movinggrating"
 
 
+def test_legacy_gonogo_no_response_miss_reconstructs_target_stimulus_and_cue_onset(flatten_trial):
+    # Recorded Go/NoGo sessions label misses "no_response", not "incorrect".
+    trial = flatten_trial({**BASE_TRIAL, "response": None, "response_time": -1, "outcome": "no_response"})
+
+    assert trial["sdt_type"] == "miss"
+    assert trial["stim_id"] == "movinggrating"
+    assert trial["cue_onset"] == pytest.approx(1.0 + BASE_TRIAL["iti"])
+
+
+def test_legacy_gonogo_precued_has_no_stimulus(flatten_trial):
+    trial = flatten_trial(
+        {**BASE_TRIAL, "response": {"timestamp": "2022-01-01T00:00:03"}, "outcome": "precued"},
+    )
+
+    assert "stim_id" not in trial
+    assert np.isnan(trial["cue_onset"])
+
+
+# -- Go/NoGo sessions from Visiomode versions before 0.5 that record per-trial parameters for both stimuli, but not
+# which was shown, as `{"target": {...}, "distractor": {...}}`. --
+
+PUSH = {"timestamp": "2022-01-01T00:00:07", "pos_x": 400.0, "pos_y": 240.0}
+
+
+@pytest.mark.parametrize(
+    "outcome, response, sdt_type, stim_id",
+    [
+        ("correct", PUSH, "hit", "movinggrating"),
+        ("no_response", None, "miss", "movinggrating"),
+        ("incorrect", PUSH, "false_alarm", "isoluminantgray"),
+        ("correct", None, "correct_rejection", "isoluminantgray"),
+    ],
+)
+def test_gonogo_empty_stimulus_pair_reconstructs_shown_stimulus(flatten_trial, outcome, response, sdt_type, stim_id):
+    trial = flatten_trial(
+        {**BASE_TRIAL, "outcome": outcome, "response": response, "stimulus": {"target": {}, "distractor": {}}}
+    )
+
+    assert trial["sdt_type"] == sdt_type
+    assert trial["stim_id"] == stim_id
+    assert trial["cue_onset"] == pytest.approx(1.0 + BASE_TRIAL["iti"])
+
+
+def test_gonogo_stimulus_pair_adds_per_trial_params_to_target_trials(flatten_trial):
+    trial = flatten_trial(
+        {
+            **BASE_TRIAL,
+            "outcome": "correct",
+            "response": PUSH,
+            "stimulus": {"target": {"trial_contrast": 0.25}, "distractor": {}},
+        }
+    )
+
+    assert trial["stim_id"] == "movinggrating"
+    assert trial["stim_trial_contrast"] == 0.25
+    assert not any(key.startswith(("target_", "distractor_")) for key in trial)
+
+
+def test_gonogo_stimulus_pair_leaves_target_params_off_distractor_trials(flatten_trial):
+    # Each trial records the target's parameters, whichever stimulus was shown.
+    trial = flatten_trial(
+        {
+            **BASE_TRIAL,
+            "outcome": "correct",
+            "response": None,
+            "stimulus": {"target": {"trial_contrast": 0.25}, "distractor": {}},
+        }
+    )
+
+    assert trial["stim_id"] == "isoluminantgray"
+    assert "stim_trial_contrast" not in trial
+
+
+def test_gonogo_stimulus_pair_precued_has_no_stimulus(flatten_trial):
+    trial = flatten_trial(
+        {
+            **BASE_TRIAL,
+            "outcome": "precued",
+            "response": {"timestamp": "2022-01-01T00:00:03"},
+            "stimulus": {"target": {"trial_contrast": 0}, "distractor": {}},
+        }
+    )
+
+    assert not any(key.startswith(("stim_", "target_", "distractor_")) for key in trial)
+    # The push ended the trial before the stimulus was due.
+    assert np.isnan(trial["cue_onset"])
+
+
 def test_stimulus_literal_none_string_is_treated_as_no_stimulus(flatten_trial):
     trial = flatten_trial({**BASE_TRIAL, "response": None, "stimulus": "None"})
 
@@ -98,7 +187,8 @@ def test_stimulus_with_target_and_distractor_keys_uses_2afc_format(flatten_trial
             **BASE_TRIAL,
             "response": None,
             "stimulus": {"target": {"id": "left_shape"}, "distractor": {"id": "right_shape"}},
-        }
+        },
+        metadata_overrides={"protocol": "afc2"},
     )
 
     assert trial["target_id"] == "left_shape"

tests/test_session.py

@@ -246,6 +246,53 @@ def test_generate_regressors_rejects_unsupported_timestamps_file_extension(gonog
         session.generate_regressors(trials_df, metadata, str(bad_path), output_dir=str(tmp_path))
 
 
+def test_generate_regressors_uses_session_stimuli_as_go_and_nogo(tmp_path):
+    # A Go stimulus other than the generator's default ("movinggrating") still gets a stim_go regressor.
+    trials_df = pd.DataFrame(
+        [
+            dict(
+                start_time=0.0,
+                stop_time=5.0,
+                cue_onset=1.0,
+                stim_id="variablecontrastmovinggrating",
+                outcome="correct",
+                response="leverpush",
+                sdt_type="hit",
+                response_time=0.5,
+            ),
+            dict(
+                start_time=5.0,
+                stop_time=10.0,
+                cue_onset=6.0,
+                stim_id="checkerboard",
+                outcome="incorrect",
+                response="leverpush",
+                sdt_type="false_alarm",
+                response_time=0.3,
+            ),
+        ]
+    )
+    metadata = {
+        "protocol": "gonogo",
+        "animal_id": "MM228",
+        "experiment": "exp",
+        "session_date": "20220516",
+        "session_start_time": "2022-05-16T00:00:00",
+        "stimuli": {"target_id": "variablecontrastmovinggrating", "distractor_id": "checkerboard"},
+    }
+    timestamps_path = tmp_path / "timestamps.txt"
+    # During the Go stimulus of trial 0, and the NoGo stimulus of trial 1.
+    timestamps_path.write_text("2022-05-16T00:00:01.500000\n2022-05-16T00:00:06.200000\n")
+
+    out_path = session.generate_regressors(trials_df, metadata, str(timestamps_path), output_dir=str(tmp_path))
+
+    with np.load(out_path) as npz:
+        labels = list(npz["labels"])
+        regressors = npz["regressors"]
+    np.testing.assert_array_equal(regressors[:, labels.index("stim_go")], [1, 0])
+    np.testing.assert_array_equal(regressors[:, labels.index("stim_nogo")], [0, 1])
+
+
 def test_generate_regressors_rejects_unsupported_protocol(tmp_path):
     trials_df = pd.DataFrame(
         [dict(start_time=0.0, stop_time=1.0, cue_onset=0.5, stim_id="x", outcome="correct", response=None)]