Promote timing_source into PsychophysicalKernel primary key

docs/MIGRATION.md

@@ -151,8 +151,8 @@ Within a NIDAQ-timed session, mapped flashes and port edges must agree with the
 Bpod trial alignment within 100 ms. Trials with missing or misaligned hardware
 events are excluded rather than replaced with Bpod timestamps.
 
-`PsychophysicalKernel.timing_source` records `nidaq`, `bpod`, or `mixed` for the
-pooled result. Detailed session provenance remains derivable from
+`PsychophysicalKernel.timing_source` is part of the primary key and records
+`nidq` or `bpod` for the fit. Detailed session provenance remains derivable from
 `BehaviorAnalysisSet.TrialSet` and the ephys `EventMapping`; no duplicate
 per-session provenance table is needed.
 

labdata_plugin/analysisschema.py

@@ -193,17 +193,17 @@ class PsychophysicalKernelFitConfig(dj.Lookup):
 
 @rojasbowe_schema
 class PsychophysicalKernel(dj.Computed):
-    """One pooled kernel per analysis set, subject, condition, and config."""
+    """One pooled kernel per analysis set, subject, condition, config, and clock."""
 
     definition = """
     -> BehaviorAnalysisSet
     -> Subject
     trialset_description                 : varchar(54)
     -> PsychophysicalKernelFitConfig
+    timing_source                        : enum('nidq', 'bpod')
     ---
     fit_status                           : enum('fit', 'skipped')
     fit_message = NULL                   : varchar(256)
-    timing_source = 'bpod'               : enum('nidaq', 'bpod', 'mixed')
     n_trials_fit                         : int
     n_bins_fit = NULL                    : int
     n_observed_per_bin = NULL            : longblob
@@ -227,7 +227,29 @@ def key_source(self):
             .aggr(BehaviorAnalysisSet.TrialSet(), n_trialsets="count(*)")
             .proj()
         )
-        return subject_conditions * PsychophysicalKernelFitConfig()
+        base = subject_conditions * PsychophysicalKernelFitConfig()
+        key_fields = (
+            "analysis_set_id",
+            "subject_name",
+            "trialset_description",
+            "kernel_fit_config_id",
+        )
+        nidq_keys = []
+        for row in base.fetch(as_dict=True):
+            trialset_keys = _selected_trialset_keys(row)
+            from behavior_analyses.kernel_timing import available_timing_sources
+
+            if "nidq" in available_timing_sources(trialset_keys):
+                nidq_keys.append({field: row[field] for field in key_fields})
+
+        key_relation = dj.U(*key_fields, "timing_source")
+        bpod = key_relation & base.proj(*key_fields, timing_source="'bpod'")
+        if not nidq_keys:
+            return bpod
+        nidq = key_relation & (base & nidq_keys).proj(
+            *key_fields, timing_source="'nidq'"
+        )
+        return bpod + nidq
 
     def make(self, key):
         config = (PsychophysicalKernelFitConfig() & key).fetch1()
@@ -274,6 +296,7 @@ def _kernel_payload(key, config, trialset_keys):
         trialset_keys,
         key["trialset_description"],
         observation_window=str(config["observation_window"]),
+        timing_source=str(key["timing_source"]),
     )
     residual, choices, n_observed, bin_centers, expected_counts = (
         build_residual_rate_matrix(
@@ -297,7 +320,6 @@ def _kernel_payload(key, config, trialset_keys):
         expected_counts=expected_counts,
     )
     base = {
-        "timing_source": inputs["timing_source"],
         "n_trials_fit": int(result["n_trials_fit"]),
         "n_bins_fit": int(result["n_bins_fit"]),
         "n_observed_per_bin": n_observed,

src/behavior_analyses/kernel_timing.py

@@ -5,6 +5,7 @@
 import numpy as np
 
 TRIALSET_DATASET_KEY_FIELDS = ("subject_name", "session_name", "dataset_name")
+TIMING_SOURCES = ("nidq", "bpod")
 REQUIRED_NIDAQ_EVENTS = (
     "visual_stim",
     "trial_start",
@@ -15,18 +16,45 @@
 MAX_NIDAQ_ALIGNMENT_ERROR_S = 0.1
 
 
+def available_timing_sources(trialset_keys: list[dict[str, Any]]) -> list[str]:
+    """Return timing sources the selected trial sets can supply."""
+    if not trialset_keys:
+        return []
+    session_keys = {(key["subject_name"], key["session_name"]) for key in trialset_keys}
+    sources = ["bpod"]
+    if all(
+        has_nidq_timing(
+            _fetch_event_mapping_rows(
+                {"subject_name": subject, "session_name": session}
+            )
+        )
+        for subject, session in session_keys
+    ):
+        sources.insert(0, "nidq")
+    return sources
+
+
 def fetch_pooled_kernel_inputs(
     trialset_keys: list[dict[str, Any]],
     trialset_description: str,
     *,
     observation_window: str,
+    timing_source: str,
 ) -> dict[str, Any]:
-    """Fetch pooled trial inputs, preferring validated NIDAQ timing per session."""
+    """Fetch pooled trial inputs from one requested timing source."""
     if observation_window not in {"center_exit", "response"}:
         raise ValueError(
             "observation_window must be 'center_exit' or 'response', "
             f"got {observation_window!r}"
         )
+    if timing_source not in TIMING_SOURCES:
+        raise ValueError(
+            f"timing_source must be one of {TIMING_SOURCES}, got {timing_source!r}"
+        )
+    if timing_source not in available_timing_sources(trialset_keys):
+        raise ValueError(
+            f"Selected trial sets cannot supply timing_source={timing_source!r}"
+        )
 
     session_inputs = []
     seen_datasets = set()
@@ -44,28 +72,32 @@ def fetch_pooled_kernel_inputs(
             field: dataset_key[field] for field in ("subject_name", "session_name")
         }
         mapping_rows = _fetch_event_mapping_rows(session_key)
-        if has_nidaq_visual_timing(mapping_rows):
+        if timing_source == "nidq":
+            if not has_nidq_timing(mapping_rows):
+                raise ValueError(
+                    f"Session {session_key['subject_name']} "
+                    f"{session_key['session_name']} cannot supply nidq timing: "
+                    "incomplete EventMapping"
+                )
             event_rows = _fetch_mapped_digital_event_rows(session_key, mapping_rows)
-            aligned_events = resolve_nidaq_event_arrays(
+            aligned_events = resolve_nidq_event_arrays(
                 event_rows,
                 mapping_rows,
                 session_key["subject_name"],
                 session_key["session_name"],
             )
-            inputs = extract_nidaq_kernel_inputs(
+            inputs = extract_nidq_kernel_inputs(
                 aligned_events,
                 trial_rows,
                 trialset_description,
                 observation_window=observation_window,
             )
-            inputs["timing_source"] = "nidaq"
         else:
             inputs = extract_bpod_kernel_inputs(
                 trial_rows,
                 trialset_description,
                 observation_window=observation_window,
             )
-            inputs["timing_source"] = "bpod"
         session_inputs.append(inputs)
 
     if not session_inputs:
@@ -109,7 +141,7 @@ def extract_bpod_kernel_inputs(
     return result
 
 
-def extract_nidaq_kernel_inputs(
+def extract_nidq_kernel_inputs(
     aligned_events: dict[str, np.ndarray],
     trial_rows: list[dict[str, Any]],
     trialset_description: str,
@@ -134,12 +166,12 @@ def extract_nidaq_kernel_inputs(
             "Insufficient finite Bpod/NIDAQ sync points to interpolate trial timing"
         )
     bpod_sync = np.asarray([row["t_sync"] for row in sync_rows], dtype=float)
-    nidaq_sync = np.asarray(
+    nidq_sync = np.asarray(
         [trial_starts[int(row["trial_num"])] for row in sync_rows], dtype=float
     )
     order = np.argsort(bpod_sync)
     bpod_sync = bpod_sync[order]
-    nidaq_sync = nidaq_sync[order]
+    nidq_sync = nidq_sync[order]
 
     stims = np.asarray(aligned_events["visual_stim"], dtype=float)
     center_exits = np.asarray(aligned_events["center_port_exit"], dtype=float)
@@ -174,11 +206,11 @@ def extract_nidaq_kernel_inputs(
             np.interp(
                 float(row["t_sync"]) + float(bpod_stims[0]),
                 bpod_sync,
-                nidaq_sync,
+                nidq_sync,
             )
         )
         interpolated_exit = float(
-            np.interp(float(row["t_react"]), bpod_sync, nidaq_sync)
+            np.interp(float(row["t_react"]), bpod_sync, nidq_sync)
         )
         trial_exits = center_exits[
             (center_exits > trial_start) & (center_exits < trial_end)
@@ -193,7 +225,7 @@ def extract_nidaq_kernel_inputs(
             if response_time is None or not np.isfinite(response_time):
                 continue
             interpolated_response = float(
-                np.interp(float(response_time), bpod_sync, nidaq_sync)
+                np.interp(float(response_time), bpod_sync, nidq_sync)
             )
             response_entries = (
                 right_entries if int(row["response"]) == 1 else left_entries
@@ -222,12 +254,15 @@ def extract_nidaq_kernel_inputs(
     return result
 
 
-def has_nidaq_visual_timing(mapping_rows: list[dict[str, Any]]) -> bool:
-    """Return whether a session declares a mapped NIDAQ/OneBox visual stream."""
-    return any(row.get("event_name") == "visual_stim" for row in mapping_rows)
+def has_nidq_timing(mapping_rows: list[dict[str, Any]]) -> bool:
+    """Return whether a session has one mapping for every required NIDAQ event."""
+    mapped_names = [row.get("event_name") for row in mapping_rows]
+    return len(mapped_names) == len(set(mapped_names)) and all(
+        name in mapped_names for name in REQUIRED_NIDAQ_EVENTS
+    )
 
 
-def resolve_nidaq_event_arrays(
+def resolve_nidq_event_arrays(
     event_rows: list[dict[str, Any]],
     mapping_rows: list[dict[str, Any]],
     subject: str,
@@ -293,9 +328,7 @@ def resolve_nidaq_event_arrays(
 
 
 def combine_kernel_inputs(session_inputs: list[dict[str, Any]]) -> dict[str, Any]:
-    """Combine per-session inputs and summarize their timing provenance."""
-    sources = {inputs["timing_source"] for inputs in session_inputs}
-    timing_source = next(iter(sources)) if len(sources) == 1 else "mixed"
+    """Combine per-session inputs from one timing source."""
     return {
         "stim_times_per_trial": [
             stims
@@ -310,7 +343,6 @@ def combine_kernel_inputs(session_inputs: list[dict[str, Any]]) -> dict[str, Any
         ),
         "response_values": _concatenate(session_inputs, "response_values", dtype=int),
         "trial_rate_hz": _concatenate(session_inputs, "trial_rate_hz", dtype=float),
-        "timing_source": timing_source,
     }