diff --git a/README.md b/README.md index aa4a7ac..94fc6de 100644 --- a/README.md +++ b/README.md @@ -44,8 +44,10 @@ Chipmunk data. Archived `djchurchland` notebooks are under ```bash uv run python scripts/analyses/seed_behavior_analysis_set.py --help uv run python scripts/analyses/migrate_behavior_analysis_schema.py --help +uv run python scripts/analyses/migrate_kernel_timing_source.py --help uv run python scripts/analyses/populate_behavior_tables.py --help uv run python scripts/analyses/plot_psychometrics.py --help ``` -Use `--dry-run` on seed/populate before any database writes. +Seed and populate support `--dry-run`. Schema migrations print their planned +changes unless explicitly run with `--apply`. diff --git a/docs/MIGRATION.md b/docs/MIGRATION.md index 8fcbb48..473df79 100644 --- a/docs/MIGRATION.md +++ b/docs/MIGRATION.md @@ -47,8 +47,8 @@ and **Plan Analysis** v0.5. | `PsychometricFitConfig` | Lookup | One versioned psychometric eligibility configuration | none | minimum choices/stimulus values, analysis version | | `PsychometricSessionFit` | Computed | One fit for one upstream TrialSet and config | `DecisionTask.TrialSet` + config | status, fit sample size, curve/parameters/diagnostics | | `PsychometricSubjectFit` | Computed | One pooled fit for one analysis set, subject, condition, and config | analysis set + `Subject` + config | status, fit sample size, curve/parameters/diagnostics | -| `PsychophysicalKernelFitConfig` | Lookup | One versioned pooled-kernel configuration | none | bins, CV folds, seed, calibration rate, regularization, version | -| `PsychophysicalKernel` | Computed | One pooled kernel for one analysis set, subject, condition, and config | analysis set + `Subject` + kernel config | status, fit sample size, weights, held-out scores, bias | +| `PsychophysicalKernelFitConfig` | Lookup | One versioned pooled-kernel configuration | none | binning method, observation window, evidence encoding, CV, regularization, version | +| `PsychophysicalKernel` | Computed | One pooled kernel for one analysis set, subject, condition, and config | analysis set + `Subject` + kernel config | status, timing source, fit sample size, weights, held-out scores, bias | Keep: @@ -130,3 +130,33 @@ sizes, frameless legends, and vector output. Raw figures, commands, configuration IDs, and observation-first notes are on the LAB-TASKS-479 [Results](https://www.notion.so/3b1ecf086b7c814a959aebd29420b453) subpage. + +## Kernel timing-source extension + +The canonical pooled kernel now supports fixed 100 ms windows in addition to +the migrated variable-width `v1_10bin_10fold` fit. For fixed-window configs, +each selected session: + +1. uses NIDAQ/OneBox visual-flash and port-event times when a `visual_stim` + `EventMapping` exists and all required mappings validate; +2. otherwise uses the equivalent Bpod `stim_events`, `t_react`, and + `t_response` timestamps; +3. raises on incomplete or invalid NIDAQ mappings rather than silently falling + back to Bpod. + +`PsychophysicalKernel.timing_source` records `nidaq`, `bpod`, or `mixed` for the +pooled result. Detailed session provenance remains derivable from +`BehaviorAnalysisSet.TrialSet` and the ephys `EventMapping`; no duplicate +per-session provenance table is needed. + +Before importing this version against the live schema, run: + +```bash +uv run python scripts/analyses/migrate_kernel_timing_source.py +uv run python scripts/analyses/migrate_kernel_timing_source.py --apply +``` + +The first command is read-only and prints the missing `ALTER TABLE` statements. +The second adds the new nullable/defaulted columns and inserts the two v2 +fixed-window configs while preserving all existing v1 rows as Bpod-timed +legacy fits. diff --git a/labdata_plugin/analysisschema.py b/labdata_plugin/analysisschema.py index 8710664..e913984 100644 --- a/labdata_plugin/analysisschema.py +++ b/labdata_plugin/analysisschema.py @@ -158,9 +158,57 @@ class PsychophysicalKernelFitConfig(dj.Lookup): random_state : int max_rate_hz : float # calibration rate regularization_c : float + kernel_method = 'legacy_variable' : enum('legacy_variable', 'fixed_window') + bin_width_s = NULL : float + evidence_encoding = NULL : enum('max_rate', 'trial_rate') + min_trials_per_bin = NULL : int + observation_window = NULL : enum('center_exit', 'response') analysis_version : varchar(32) """ - contents = [("v1_10bin_10fold", 10, 10, 0, 20.0, 1.0, "v1")] # noqa: RUF012 + contents = [ # noqa: RUF012 + ( + "v1_10bin_10fold", + 10, + 10, + 0, + 20.0, + 1.0, + "legacy_variable", + None, + None, + None, + None, + "v1", + ), + ( + "v2_100ms_10bin_center_rate", + 10, + 10, + 0, + 20.0, + 1.0, + "fixed_window", + 0.1, + "trial_rate", + 50, + "center_exit", + "v2", + ), + ( + "v2_100ms_10bin_response_rate", + 10, + 10, + 0, + 20.0, + 1.0, + "fixed_window", + 0.1, + "trial_rate", + 50, + "response", + "v2", + ), + ] @rojasbowe_schema @@ -175,14 +223,21 @@ class PsychophysicalKernel(dj.Computed): --- 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 + bin_centers_s = NULL : longblob # time from first flash weights = NULL : longblob # cv fold x stimulus time bin weights_mean = NULL : longblob weights_error = NULL : longblob scores = NULL : longblob # held-out accuracy by fold score_mean = NULL : float + majority_accuracy = NULL : float + score_above_majority = NULL : float bias = NULL : longblob # intercept by fold bias_mean = NULL : float + interpretation = NULL : varchar(32) """ @property @@ -194,51 +249,16 @@ def key_source(self): return subject_conditions * PsychophysicalKernelFitConfig() def make(self, key): - from behavior_analyses.io import get_chipmunk_table - from behavior_analyses.kernels import fit_psychophysical_kernel - config = (PsychophysicalKernelFitConfig() & key).fetch1() trialset_keys = _selected_trialset_keys(key) - Chipmunk = get_chipmunk_table() - relation = ( - Chipmunk() * Chipmunk.Trial() * Chipmunk.TrialParameters() - & trialset_keys - & {"rewarded_modality": key["trialset_description"]} - ) - stim_events, response_values = relation.fetch("stim_events", "response") - result = fit_psychophysical_kernel( - stim_events, - response_values, - timebins=int(config["timebins"]), - cv_splits=int(config["cv_splits"]), - random_state=int(config["random_state"]), - max_rate_hz=float(config["max_rate_hz"]), - regularization_c=float(config["regularization_c"]), - ) - n_trials_fit = int(result["choice_right"].size) - if result["weights"].size == 0: - self.insert1( - { - **key, - "fit_status": "skipped", - "fit_message": "insufficient trials or response classes for CV", - "n_trials_fit": n_trials_fit, - } - ) - return - + if config["kernel_method"] == "legacy_variable": + payload = _legacy_kernel_payload(key, config, trialset_keys) + else: + payload = _fixed_kernel_payload(key, config, trialset_keys) self.insert1( { **key, - "fit_status": "fit", - "n_trials_fit": n_trials_fit, - "weights": result["weights"], - "weights_mean": np.mean(result["weights"], axis=0), - "weights_error": np.mean(result["error"], axis=0), - "scores": result["scores"], - "score_mean": float(np.mean(result["scores"])), - "bias": result["bias"], - "bias_mean": float(np.mean(result["bias"])), + **payload, } ) @@ -264,6 +284,142 @@ def _fetch_trialset_rows_for_subject(key): return list((DecisionTask.TrialSet() & trialset_keys).fetch(as_dict=True)) +def _legacy_kernel_payload(key, config, trialset_keys): + from behavior_analyses.io import get_chipmunk_table + from behavior_analyses.kernels import ( + fit_psychophysical_kernel, + interpret_kernel_weights, + ) + + Chipmunk = get_chipmunk_table() + relation = ( + Chipmunk() * Chipmunk.Trial() * Chipmunk.TrialParameters() + & trialset_keys + & {"rewarded_modality": key["trialset_description"]} + ) + stim_events, response_values = relation.fetch("stim_events", "response") + result = fit_psychophysical_kernel( + stim_events, + response_values, + timebins=int(config["timebins"]), + cv_splits=int(config["cv_splits"]), + random_state=int(config["random_state"]), + max_rate_hz=float(config["max_rate_hz"]), + regularization_c=float(config["regularization_c"]), + ) + n_trials_fit = int(result["choice_right"].size) + if result["weights"].size == 0: + return { + "fit_status": "skipped", + "fit_message": "insufficient trials or response classes for CV", + "timing_source": "bpod", + "n_trials_fit": n_trials_fit, + "n_bins_fit": 0, + } + + n_observed = np.full(int(config["timebins"]), n_trials_fit, dtype=int) + score_mean = float(np.mean(result["scores"])) + majority_accuracy = float( + max(np.mean(result["choice_right"]), 1.0 - np.mean(result["choice_right"])) + ) + weights_mean = np.mean(result["weights"], axis=0) + return { + "fit_status": "fit", + "timing_source": "bpod", + "n_trials_fit": n_trials_fit, + "n_bins_fit": int(config["timebins"]), + "n_observed_per_bin": n_observed, + "weights": result["weights"], + "weights_mean": weights_mean, + "weights_error": np.mean(result["error"], axis=0), + "scores": result["scores"], + "score_mean": score_mean, + "majority_accuracy": majority_accuracy, + "score_above_majority": score_mean - majority_accuracy, + "bias": result["bias"], + "bias_mean": float(np.mean(result["bias"])), + "interpretation": interpret_kernel_weights( + weights_mean, + n_observed, + min_trials_per_bin=int(config["cv_splits"]), + ), + } + + +def _fixed_kernel_payload(key, config, trialset_keys): + from behavior_analyses.kernel_timing import fetch_pooled_kernel_inputs + from behavior_analyses.kernels import ( + build_fixed_residual_rate_matrix, + fit_fixed_psychophysical_kernel, + interpret_kernel_weights, + ) + + inputs = fetch_pooled_kernel_inputs( + trialset_keys, + key["trialset_description"], + observation_window=str(config["observation_window"]), + ) + trial_rate_hz = ( + inputs["trial_rate_hz"] if config["evidence_encoding"] == "trial_rate" else None + ) + residual, choices, n_observed, bin_centers, expected_counts = ( + build_fixed_residual_rate_matrix( + inputs["stim_times_per_trial"], + inputs["first_stim_times"], + inputs["observation_end_times"], + inputs["response_values"], + timebins=int(config["timebins"]), + bin_width_s=float(config["bin_width_s"]), + max_rate_hz=float(config["max_rate_hz"]), + trial_rate_hz=trial_rate_hz, + ) + ) + result = fit_fixed_psychophysical_kernel( + residual, + choices, + n_observed_per_bin=n_observed, + cv_splits=int(config["cv_splits"]), + random_state=int(config["random_state"]), + min_trials_per_bin=int(config["min_trials_per_bin"]), + regularization_c=float(config["regularization_c"]), + expected_counts=( + expected_counts if config["evidence_encoding"] == "trial_rate" else None + ), + ) + 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, + "bin_centers_s": bin_centers, + } + if not result["fit_converged"]: + return { + **base, + "fit_status": "skipped", + "fit_message": "insufficient trials or response classes for CV", + } + + return { + **base, + "fit_status": "fit", + "weights": result["weights"], + "weights_mean": result["weights_mean"], + "weights_error": result["weights_error"], + "scores": result["scores"], + "score_mean": result["score_mean"], + "majority_accuracy": result["majority_accuracy"], + "score_above_majority": result["score_above_majority"], + "bias": result["bias"], + "bias_mean": result["bias_mean"], + "interpretation": interpret_kernel_weights( + result["weights_mean"], + n_observed, + min_trials_per_bin=int(config["min_trials_per_bin"]), + ), + } + + def _psychometric_fit_payload(row, config): from behavior_analyses.psychometrics import fit_psychometric_labdata diff --git a/scripts/analyses/migrate_kernel_timing_source.py b/scripts/analyses/migrate_kernel_timing_source.py new file mode 100644 index 0000000..8e9da67 --- /dev/null +++ b/scripts/analyses/migrate_kernel_timing_source.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import argparse + +import _bootstrap # noqa: F401 + + +CONFIG_TABLE = "#psychophysical_kernel_fit_config" +KERNEL_TABLE = "__psychophysical_kernel" +CONFIG_COLUMNS = ( + ( + "kernel_method", + "enum('legacy_variable','fixed_window') NOT NULL DEFAULT 'legacy_variable'", + "regularization_c", + ), + ("bin_width_s", "float DEFAULT NULL", "kernel_method"), + ( + "evidence_encoding", + "enum('max_rate','trial_rate') DEFAULT NULL", + "bin_width_s", + ), + ("min_trials_per_bin", "int DEFAULT NULL", "evidence_encoding"), + ( + "observation_window", + "enum('center_exit','response') DEFAULT NULL", + "min_trials_per_bin", + ), +) +KERNEL_COLUMNS = ( + ( + "timing_source", + "enum('nidaq','bpod','mixed') NOT NULL DEFAULT 'bpod'", + "fit_message", + ), + ("n_bins_fit", "int DEFAULT NULL", "n_trials_fit"), + ("n_observed_per_bin", "longblob DEFAULT NULL", "n_bins_fit"), + ("bin_centers_s", "longblob DEFAULT NULL", "n_observed_per_bin"), + ("majority_accuracy", "float DEFAULT NULL", "score_mean"), + ("score_above_majority", "float DEFAULT NULL", "majority_accuracy"), + ("interpretation", "varchar(32) DEFAULT NULL", "bias_mean"), +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument( + "--apply", + action="store_true", + help="Alter the live kernel tables and insert fixed-window configs.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + import datajoint as dj + import labdata.schema as labdata_schema + + connection = dj.conn() + database = f"{labdata_schema.dbase_name}_user" + table_columns = { + CONFIG_TABLE: CONFIG_COLUMNS, + KERNEL_TABLE: KERNEL_COLUMNS, + } + existing = { + table: _existing_columns(connection, database, table) for table in table_columns + } + statements = [ + statement + for table, columns in table_columns.items() + for statement in _missing_column_statements( + database, table, columns, existing[table] + ) + ] + + if not args.apply: + for statement in statements: + print(statement) + if not statements: + print("Kernel timing-source columns are already present.") + print("Dry run only. Re-run with --apply after exact live-write approval.") + return + + for statement in statements: + connection.query(statement) + + from labdata_plugin.analysisschema import PsychophysicalKernelFitConfig + + PsychophysicalKernelFitConfig.insert( + PsychophysicalKernelFitConfig.contents, + skip_duplicates=True, + ) + _validate_migration(connection, database) + print("Kernel timing-source schema migration complete.") + + +def _existing_columns(connection, database: str, table: str) -> set[str]: + return { + row[0] + for row in connection.query( + "SELECT column_name FROM information_schema.columns " + f"WHERE table_schema={database!r} AND table_name={table!r}" + ).fetchall() + } + + +def _missing_column_statements( + database: str, + table: str, + columns: tuple[tuple[str, str, str], ...], + existing: set[str], +) -> list[str]: + return [ + f"ALTER TABLE `{database}`.`{table}` ADD COLUMN `{name}` " + f"{definition} AFTER `{after}`" + for name, definition, after in columns + if name not in existing + ] + + +def _validate_migration(connection, database: str) -> None: + expected = { + CONFIG_TABLE: {name for name, _, _ in CONFIG_COLUMNS}, + KERNEL_TABLE: {name for name, _, _ in KERNEL_COLUMNS}, + } + missing = { + table: sorted(columns - _existing_columns(connection, database, table)) + for table, columns in expected.items() + } + missing = {table: columns for table, columns in missing.items() if columns} + if missing: + raise RuntimeError(f"Kernel timing-source migration incomplete: {missing}") + + invalid_sources = connection.query( + f"SELECT COUNT(*) FROM `{database}`.`{KERNEL_TABLE}` " + "WHERE timing_source NOT IN ('nidaq','bpod','mixed')" + ).fetchone()[0] + if invalid_sources: + raise RuntimeError( + f"Kernel timing-source migration found {invalid_sources} invalid rows" + ) + + +if __name__ == "__main__": + main() diff --git a/src/behavior_analyses/io.py b/src/behavior_analyses/io.py index d49af1f..68a7fb8 100644 --- a/src/behavior_analyses/io.py +++ b/src/behavior_analyses/io.py @@ -47,16 +47,28 @@ def get_chipmunk_table() -> Any: return Chipmunk except ModuleNotFoundError: + registered = _registered_chipmunk_table() + if registered is not None: + return registered plugin_path = _configured_chipmunk_plugin_path() if plugin_path is None: raise ModuleNotFoundError( - "Chipmunk plugin not importable as `chipmunk`, and no " - "CHIPMUNK_PLUGIN_PATH / tool.behavior_analyses.chipmunk_plugin_path " - "is configured." + "Chipmunk plugin is not importable or registered with labdata, " + "and no CHIPMUNK_PLUGIN_PATH / " + "tool.behavior_analyses.chipmunk_plugin_path is configured." ) from None return _load_local_chipmunk_plugin(plugin_path).Chipmunk +def _registered_chipmunk_table() -> Any | None: + import labdata + + try: + return labdata.plugins["chipmunk"].Chipmunk + except KeyError: + return None + + def _load_local_chipmunk_plugin(plugin_root: Path) -> Any: module_name = "_behavior_analyses_chipmunk_plugin" if module_name in sys.modules: diff --git a/src/behavior_analyses/kernel_timing.py b/src/behavior_analyses/kernel_timing.py new file mode 100644 index 0000000..7146d31 --- /dev/null +++ b/src/behavior_analyses/kernel_timing.py @@ -0,0 +1,402 @@ +from __future__ import annotations + +from importlib import import_module +from typing import Any + +import numpy as np + +TRIALSET_DATASET_KEY_FIELDS = ("subject_name", "session_name", "dataset_name") +REQUIRED_NIDAQ_EVENTS = ( + "visual_stim", + "trial_start", + "left_port", + "center_port", + "right_port", +) + + +def fetch_pooled_kernel_inputs( + trialset_keys: list[dict[str, Any]], + trialset_description: str, + *, + observation_window: str, +) -> dict[str, Any]: + """Fetch pooled trial inputs, preferring validated NIDAQ timing per session.""" + if observation_window not in {"center_exit", "response"}: + raise ValueError( + "observation_window must be 'center_exit' or 'response', " + f"got {observation_window!r}" + ) + + session_inputs = [] + seen_datasets = set() + for trialset_key in trialset_keys: + dataset_key = { + field: trialset_key[field] for field in TRIALSET_DATASET_KEY_FIELDS + } + dataset_identity = tuple(dataset_key.values()) + if dataset_identity in seen_datasets: + continue + seen_datasets.add(dataset_identity) + + trial_rows = _fetch_chipmunk_trial_rows(dataset_key) + session_key = { + 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): + event_rows = _fetch_mapped_digital_event_rows(session_key, mapping_rows) + aligned_events = resolve_nidaq_event_arrays( + event_rows, + mapping_rows, + session_key["subject_name"], + session_key["session_name"], + ) + inputs = extract_nidaq_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: + raise ValueError("No selected Chipmunk trial sets were available") + return combine_kernel_inputs(session_inputs) + + +def extract_bpod_kernel_inputs( + trial_rows: list[dict[str, Any]], + trialset_description: str, + *, + observation_window: str, +) -> dict[str, Any]: + """Extract fixed-window inputs using native Bpod trial timestamps.""" + result = _new_kernel_inputs() + end_field = "t_react" if observation_window == "center_exit" else "t_response" + for row in sorted(trial_rows, key=lambda item: int(item["trial_num"])): + if row["rewarded_modality"] != trialset_description: + continue + response = row["response"] + observation_end = row.get(end_field) + stims = np.asarray(row.get("stim_events", []), dtype=float) + stims = stims[np.isfinite(stims)] + if ( + response not in (-1, 1) + or observation_end is None + or not np.isfinite(observation_end) + or stims.size == 0 + ): + continue + stims = stims[stims < float(observation_end)] + if stims.size == 0 or observation_end <= stims[0]: + continue + _append_kernel_trial(result, stims, observation_end, row) + return result + + +def extract_nidaq_kernel_inputs( + aligned_events: dict[str, np.ndarray], + trial_rows: list[dict[str, Any]], + trialset_description: str, + *, + observation_window: str, +) -> dict[str, Any]: + """Extract fixed-window inputs from NIDAQ events aligned to Bpod trials.""" + rows = sorted(trial_rows, key=lambda item: int(item["trial_num"])) + trial_starts = np.asarray(aligned_events["trial_start"], dtype=float) + if trial_starts.size == 0: + raise ValueError("NIDAQ trial_start contains no rising edges") + + sync_rows = [ + row + for row in rows + if int(row["trial_num"]) < trial_starts.size + and row.get("t_sync") is not None + and np.isfinite(row["t_sync"]) + ] + if len(sync_rows) < 2: + raise ValueError( + "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( + [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] + + stims = np.asarray(aligned_events["visual_stim"], dtype=float) + center_entries = np.asarray(aligned_events["center_port"], dtype=float) + center_exits = np.asarray(aligned_events["center_port_exit"], dtype=float) + left_entries = np.asarray(aligned_events["left_port"], dtype=float) + right_entries = np.asarray(aligned_events["right_port"], dtype=float) + result = _new_kernel_inputs() + + for row in rows: + if ( + row["rewarded_modality"] != trialset_description + or row["response"] not in (-1, 1) + or row.get("t_react") is None + or not np.isfinite(row["t_react"]) + ): + continue + trial_number = int(row["trial_num"]) + if trial_number >= trial_starts.size: + raise ValueError(f"NIDAQ trial_start is missing trial_num={trial_number}") + trial_start = trial_starts[trial_number] + trial_end = ( + trial_starts[trial_number + 1] + if trial_number + 1 < trial_starts.size + else np.inf + ) + interpolated_exit = float( + np.interp(float(row["t_react"]), bpod_sync, nidaq_sync) + ) + trial_exits = center_exits[ + (center_exits > trial_start) & (center_exits < trial_end) + ] + if trial_exits.size: + center_exit = float( + trial_exits[np.argmin(np.abs(trial_exits - interpolated_exit))] + ) + else: + center_exit = interpolated_exit + + center_mask = ( + (center_entries > trial_start) + & (center_entries < center_exit) + & (center_entries < trial_end) + ) + if not center_mask.any(): + raise ValueError(f"No NIDAQ center-port entry for trial_num={trial_number}") + center_entry = float(center_entries[center_mask][-1]) + + response_entries = right_entries if int(row["response"]) == 1 else left_entries + response_mask = (response_entries > center_exit) & ( + response_entries < trial_end + ) + if not response_mask.any(): + raise ValueError( + f"No NIDAQ response-port entry for trial_num={trial_number}" + ) + response_entry = float(response_entries[response_mask][0]) + observation_end = ( + center_exit if observation_window == "center_exit" else response_entry + ) + trial_stims = stims[(stims >= center_entry) & (stims < observation_end)] + if trial_stims.size == 0: + continue + _append_kernel_trial(result, trial_stims, observation_end, row) + 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 resolve_nidaq_event_arrays( + event_rows: list[dict[str, Any]], + mapping_rows: list[dict[str, Any]], + subject: str, + session: str, +) -> dict[str, np.ndarray]: + """Resolve and validate logical NIDAQ event arrays for one session.""" + mapped_names = [row["event_name"] for row in mapping_rows] + missing = [name for name in REQUIRED_NIDAQ_EVENTS if name not in mapped_names] + if missing: + raise ValueError( + f"Missing EventMapping rows for {subject} {session}: {missing}" + ) + duplicates = sorted({name for name in mapped_names if mapped_names.count(name) > 1}) + if duplicates: + raise ValueError( + f"Duplicate EventMapping rows for {subject} {session}: {duplicates}" + ) + + resolved = {} + for logical_name in REQUIRED_NIDAQ_EVENTS: + mapping = next(row for row in mapping_rows if row["event_name"] == logical_name) + matches = [ + row + for row in event_rows + if row["dataset_name"] == mapping["source_dataset_name"] + and row["stream_name"] == mapping["source_stream_name"] + and row["event_name"] == mapping["source_event_name"] + ] + if len(matches) != 1: + raise ValueError( + f"Expected one mapped NIDAQ row for {subject} {session} " + f"{logical_name}; found {len(matches)}" + ) + timestamps = np.asarray(matches[0]["event_timestamps"], dtype=float) + if timestamps.size == 0 or np.any(~np.isfinite(timestamps)): + raise ValueError( + f"Mapped NIDAQ row is empty or nonfinite for {subject} {session} " + f"{logical_name}" + ) + values = matches[0].get("event_values") + if values is not None: + values = np.asarray(values) + if values.shape != timestamps.shape: + raise ValueError( + f"NIDAQ event values do not match timestamps for " + f"{subject} {session} {logical_name}" + ) + resolved[logical_name] = (timestamps, values) + + visual_stim = _merge_visual_stim_edges(resolved["visual_stim"][0]) + if visual_stim.size == 0: + raise ValueError(f"No visual flashes found for {subject} {session}") + return { + "visual_stim": visual_stim, + "trial_start": _digital_onsets(*resolved["trial_start"]), + "left_port": _port_entries(*resolved["left_port"]), + "left_port_exit": _port_exits(*resolved["left_port"]), + "center_port": _port_entries(*resolved["center_port"]), + "center_port_exit": _port_exits(*resolved["center_port"]), + "right_port": _port_entries(*resolved["right_port"]), + "right_port_exit": _port_exits(*resolved["right_port"]), + } + + +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" + return { + "stim_times_per_trial": [ + stims + for inputs in session_inputs + for stims in inputs["stim_times_per_trial"] + ], + "first_stim_times": _concatenate( + session_inputs, "first_stim_times", dtype=float + ), + "observation_end_times": _concatenate( + session_inputs, "observation_end_times", dtype=float + ), + "response_values": _concatenate(session_inputs, "response_values", dtype=int), + "trial_rate_hz": _concatenate(session_inputs, "trial_rate_hz", dtype=float), + "timing_source": timing_source, + } + + +def _fetch_chipmunk_trial_rows(dataset_key: dict[str, Any]) -> list[dict[str, Any]]: + from behavior_analyses.io import get_chipmunk_table + + Chipmunk = get_chipmunk_table() + relation = Chipmunk() * Chipmunk.Trial() * Chipmunk.TrialParameters() & dataset_key + return list( + relation.fetch( + "trial_num", + "rewarded_modality", + "stim_events", + "stim_rate_vision", + "response", + "t_sync", + "t_react", + "t_response", + as_dict=True, + order_by="trial_num", + ) + ) + + +def _fetch_event_mapping_rows( + session_key: dict[str, Any], +) -> list[dict[str, Any]]: + try: + import labdata + + labdata.plugins["gephys"].__file__ + module = import_module("gephys.analysisschema") + except KeyError: + return [] + except ModuleNotFoundError as error: + if error.name in {"gephys", "gephys.analysisschema"}: + return [] + raise + return list((module.EventMapping() & session_key).fetch(as_dict=True)) + + +def _fetch_mapped_digital_event_rows( + session_key: dict[str, Any], + mapping_rows: list[dict[str, Any]], +) -> list[dict[str, Any]]: + from labdata.schema import DatasetEvents + + source_keys = [ + { + **session_key, + "dataset_name": row["source_dataset_name"], + "stream_name": row["source_stream_name"], + "event_name": row["source_event_name"], + } + for row in mapping_rows + if row["event_name"] in REQUIRED_NIDAQ_EVENTS + ] + return list((DatasetEvents.Digital() & source_keys).fetch_synced()) + + +def _new_kernel_inputs() -> dict[str, Any]: + return { + "stim_times_per_trial": [], + "first_stim_times": [], + "observation_end_times": [], + "response_values": [], + "trial_rate_hz": [], + } + + +def _append_kernel_trial( + result: dict[str, Any], + stims: np.ndarray, + observation_end: float, + row: dict[str, Any], +) -> None: + rate = row.get("stim_rate_vision") + if rate is None or not np.isfinite(rate): + raise ValueError(f"Trial {row['trial_num']} has no finite visual stimulus rate") + result["stim_times_per_trial"].append(np.asarray(stims, dtype=float)) + result["first_stim_times"].append(float(stims[0])) + result["observation_end_times"].append(float(observation_end)) + result["response_values"].append(int(row["response"])) + result["trial_rate_hz"].append(float(rate)) + + +def _digital_onsets(timestamps: np.ndarray, values: np.ndarray | None) -> np.ndarray: + return timestamps[::2] if values is None else timestamps[values == 1] + + +def _port_entries(timestamps: np.ndarray, values: np.ndarray | None) -> np.ndarray: + return timestamps if values is None else timestamps[values == 1] + + +def _port_exits(timestamps: np.ndarray, values: np.ndarray | None) -> np.ndarray: + return np.array([], dtype=float) if values is None else timestamps[values == 0] + + +def _merge_visual_stim_edges(timestamps: np.ndarray) -> np.ndarray: + timestamps = np.sort(np.asarray(timestamps, dtype=float)) + if timestamps.size == 0: + return timestamps + split_indices = np.where(np.diff(timestamps) > 0.020)[0] + 1 + return np.asarray([burst[0] for burst in np.split(timestamps, split_indices)]) + + +def _concatenate( + session_inputs: list[dict[str, Any]], field: str, *, dtype: type +) -> np.ndarray: + values = [np.asarray(inputs[field], dtype=dtype) for inputs in session_inputs] + return np.concatenate(values) if values else np.empty((0,), dtype=dtype) diff --git a/src/behavior_analyses/kernels.py b/src/behavior_analyses/kernels.py index 038ce97..979fbd1 100644 --- a/src/behavior_analyses/kernels.py +++ b/src/behavior_analyses/kernels.py @@ -1,5 +1,8 @@ from __future__ import annotations +from collections.abc import Sequence +from typing import Any + import numpy as np from sklearn.linear_model import LogisticRegression from sklearn.model_selection import StratifiedKFold @@ -86,3 +89,280 @@ def fit_psychophysical_kernel( "bias": np.asarray(biases, dtype=float), "error": np.asarray(errors, dtype=float), } + + +def build_fixed_residual_rate_matrix( + stim_times_per_trial: Sequence[np.ndarray], + first_stim_times: Sequence[float], + observation_end_times: Sequence[float], + response_values: Sequence[Any], + *, + timebins: int = 10, + bin_width_s: float = 0.1, + max_rate_hz: float = 20.0, + trial_rate_hz: Sequence[float] | None = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Build fixed-width residual flash counts for a timing-aware kernel.""" + if timebins < 1: + raise ValueError("timebins must be >= 1") + if bin_width_s <= 0: + raise ValueError("bin_width_s must be > 0") + + responses = np.asarray(response_values) + choice_mask = np.isin(responses, (-1, 1)) + choices = (responses[choice_mask] == 1).astype(int) + stim_list = [ + np.asarray(times, dtype=float) + for times, keep in zip(stim_times_per_trial, choice_mask) + if keep + ] + first_stim = np.asarray(first_stim_times, dtype=float)[choice_mask] + observation_end = np.asarray(observation_end_times, dtype=float)[choice_mask] + + if trial_rate_hz is None: + expected_rates = np.full(choice_mask.shape, max_rate_hz, dtype=float) + else: + expected_rates = np.asarray(trial_rate_hz, dtype=float) + if expected_rates.shape != choice_mask.shape: + raise ValueError("trial_rate_hz must match response_values") + if np.any(~np.isfinite(expected_rates)) or np.any(expected_rates < 0): + raise ValueError("trial_rate_hz must contain finite nonnegative rates") + expected_rates = expected_rates[choice_mask] + + bin_edges = np.arange(timebins + 1, dtype=float) * bin_width_s + bin_centers_s = 0.5 * (bin_edges[:-1] + bin_edges[1:]) + residual = np.full((len(stim_list), timebins), np.nan, dtype=float) + expected_counts = np.full_like(residual, np.nan) + + for trial_index, (stims, first, observation_end, expected_rate) in enumerate( + zip(stim_list, first_stim, observation_end, expected_rates) + ): + if ( + not np.isfinite(first) + or not np.isfinite(observation_end) + or observation_end <= first + ): + continue + + duration = float(observation_end - first) + relative_stims = stims[np.isfinite(stims)] - first + relative_stims = relative_stims[ + (relative_stims >= 0.0) & (relative_stims < duration) + ] + for bin_index in range(timebins): + bin_start = bin_edges[bin_index] + if bin_start >= duration: + continue + observed_end = min(bin_edges[bin_index + 1], duration) + observed_duration = observed_end - bin_start + count = np.sum( + (relative_stims >= bin_start) & (relative_stims < observed_end) + ) + expected = expected_rate * observed_duration + residual[trial_index, bin_index] = float(count - expected) + expected_counts[trial_index, bin_index] = expected + + n_observed_per_bin = np.sum(np.isfinite(residual), axis=0).astype(int) + return residual, choices, n_observed_per_bin, bin_centers_s, expected_counts + + +def fit_fixed_psychophysical_kernel( + residual: np.ndarray, + choice_right: np.ndarray, + *, + n_observed_per_bin: np.ndarray | None = None, + cv_splits: int = 10, + random_state: int = 0, + min_trials_per_bin: int = 50, + regularization_c: float = 1.0, + expected_counts: np.ndarray | None = None, +) -> dict[str, Any]: + """Fit an L2 logistic kernel using the longest complete-case bin prefix.""" + residual = np.asarray(residual, dtype=float) + choice_right = np.asarray(choice_right, dtype=int) + n_bins = residual.shape[1] if residual.ndim == 2 else 0 + if n_observed_per_bin is None: + n_observed_per_bin = np.sum(np.isfinite(residual), axis=0).astype(int) + else: + n_observed_per_bin = np.asarray(n_observed_per_bin, dtype=int) + + empty = { + "choice_right": choice_right, + "weights": np.empty((0, n_bins)), + "weights_mean": np.full(n_bins, np.nan), + "weights_error": np.full(n_bins, np.nan), + "scores": np.empty((0,)), + "score_mean": np.nan, + "majority_accuracy": np.nan, + "score_above_majority": np.nan, + "bias": np.empty((0,)), + "bias_mean": np.nan, + "n_observed_per_bin": n_observed_per_bin, + "n_trials_fit": 0, + "n_bins_fit": 0, + "fit_converged": False, + } + if residual.size == 0 or choice_right.size == 0 or np.unique(choice_right).size < 2: + return empty + + n_bins_fit = _complete_case_prefix( + residual, + n_observed_per_bin, + min_trials_per_bin=min_trials_per_bin, + ) + if n_bins_fit < 1: + return empty + + complete = np.all(np.isfinite(residual[:, :n_bins_fit]), axis=1) + x = residual[complete, :n_bins_fit] + y = choice_right[complete] + if x.shape[0] < max(cv_splits, min_trials_per_bin) or np.unique(y).size < 2: + return empty + n_splits = int(min(cv_splits, np.min(np.bincount(y)))) + if n_splits < 2: + return empty + + design, coefficient_to_weights = _fixed_kernel_design( + residual, + x, + complete, + n_bins_fit, + expected_counts, + ) + splitter = StratifiedKFold( + n_splits=n_splits, shuffle=True, random_state=random_state + ) + weights = [] + errors = [] + scores = [] + biases = [] + for train_index, test_index in splitter.split(design, y): + x_train, x_test = design[train_index], design[test_index] + y_train, y_test = y[train_index], y[test_index] + model = LogisticRegression( + solver="liblinear", + C=regularization_c, + fit_intercept=True, + ).fit(x_train, y_train) + predict_prob = model.predict_proba(x_train) + variance = np.prod(predict_prob, axis=1) + covariance = np.linalg.pinv(np.dot(x_train.T * variance, x_train)) + + weight_full = np.full(n_bins, np.nan, dtype=float) + error_full = np.full(n_bins, np.nan, dtype=float) + weight_full[:n_bins_fit] = coefficient_to_weights @ model.coef_[0] + weight_covariance = ( + coefficient_to_weights @ covariance @ coefficient_to_weights.T + ) + error_full[:n_bins_fit] = np.sqrt(np.diag(weight_covariance)) + weights.append(weight_full) + errors.append(error_full) + scores.append(model.score(x_test, y_test)) + biases.append(float(model.intercept_[0])) + + weights = np.asarray(weights, dtype=float) + errors = np.asarray(errors, dtype=float) + scores = np.asarray(scores, dtype=float) + biases = np.asarray(biases, dtype=float) + weights_mean = np.full(n_bins, np.nan) + weights_error = np.full(n_bins, np.nan) + weights_mean[:n_bins_fit] = np.mean(weights[:, :n_bins_fit], axis=0) + weights_error[:n_bins_fit] = np.mean(errors[:, :n_bins_fit], axis=0) + score_mean = float(np.mean(scores)) + majority_accuracy = float(max(np.mean(y), 1.0 - np.mean(y))) + return { + "choice_right": choice_right, + "weights": weights, + "weights_mean": weights_mean, + "weights_error": weights_error, + "scores": scores, + "score_mean": score_mean, + "majority_accuracy": majority_accuracy, + "score_above_majority": score_mean - majority_accuracy, + "bias": biases, + "bias_mean": float(np.mean(biases)), + "n_observed_per_bin": n_observed_per_bin, + "n_trials_fit": int(x.shape[0]), + "n_bins_fit": int(n_bins_fit), + "fit_converged": True, + } + + +def interpret_kernel_weights( + weights_mean: np.ndarray, + n_observed_per_bin: np.ndarray, + *, + min_trials_per_bin: int = 50, + ratio_threshold: float = 1.5, +) -> str: + """Return a descriptive early, late, flat, or failed kernel label.""" + weights_mean = np.asarray(weights_mean, dtype=float) + n_observed_per_bin = np.asarray(n_observed_per_bin, dtype=int) + usable = np.isfinite(weights_mean) & (n_observed_per_bin >= min_trials_per_bin) + if int(np.sum(usable)) < 3: + return "failed_fit" + + values = np.abs(weights_mean[usable]) + midpoint = max(1, values.size // 2) + early = float(np.mean(values[:midpoint])) + late = float(np.mean(values[midpoint:])) + if early == 0: + ratio = np.inf if late > 0 else 1.0 + else: + ratio = late / early + if ratio > ratio_threshold: + return "late_integrator" + if ratio < 1.0 / ratio_threshold: + return "early_integrator" + return "flat_indeterminate" + + +def _complete_case_prefix( + residual: np.ndarray, + n_observed_per_bin: np.ndarray, + *, + min_trials_per_bin: int, +) -> int: + prefix = 0 + for bin_index in range(residual.shape[1]): + complete = np.all(np.isfinite(residual[:, : bin_index + 1]), axis=1) + if ( + int(n_observed_per_bin[bin_index]) < min_trials_per_bin + or int(np.sum(complete)) < min_trials_per_bin + ): + break + prefix = bin_index + 1 + return prefix + + +def _fixed_kernel_design( + residual: np.ndarray, + complete_residual: np.ndarray, + complete: np.ndarray, + n_bins_fit: int, + expected_counts: np.ndarray | None, +) -> tuple[np.ndarray, np.ndarray]: + if expected_counts is None: + return complete_residual, np.eye(n_bins_fit) + + expected_counts = np.asarray(expected_counts, dtype=float) + if expected_counts.shape != residual.shape: + raise ValueError("expected_counts must match residual") + expected_fit = expected_counts[complete, :n_bins_fit] + if np.any(~np.isfinite(expected_fit)): + raise ValueError("expected_counts must be finite for fitted bins") + + mean_rate = np.sum(expected_fit, axis=1) + if n_bins_fit == 1: + return mean_rate[:, None], np.ones((1, 1)) + + basis_source = np.column_stack( + [ + np.eye(n_bins_fit)[:, index] - np.eye(n_bins_fit)[:, -1] + for index in range(n_bins_fit - 1) + ] + ) + basis = np.linalg.qr(basis_source, mode="reduced")[0] + design = np.column_stack((complete_residual @ basis, mean_rate)) + coefficient_to_weights = np.column_stack((basis, np.ones(n_bins_fit))) + return design, coefficient_to_weights diff --git a/tests/test_analysis_functions.py b/tests/test_analysis_functions.py index ec10f73..c1bd951 100644 --- a/tests/test_analysis_functions.py +++ b/tests/test_analysis_functions.py @@ -102,6 +102,58 @@ def test_kernel_design_matrix_skips_no_choice_and_short_stim_trials(self): np.testing.assert_allclose(x[0], [1 - 20 / 3, 2 - 20 / 3]) np.testing.assert_array_equal(y, [1]) + def test_fixed_kernel_keeps_unobserved_late_bins_as_nan(self): + from behavior_analyses.kernels import build_fixed_residual_rate_matrix + + residual, choices, n_observed, centers, expected = ( + build_fixed_residual_rate_matrix( + [np.array([1.0, 1.05, 1.12])], + [1.0], + [1.15], + [1], + timebins=3, + bin_width_s=0.1, + trial_rate_hz=[20.0], + ) + ) + + np.testing.assert_array_equal(choices, [1]) + np.testing.assert_array_equal(n_observed, [1, 1, 0]) + np.testing.assert_allclose(centers, [0.05, 0.15, 0.25]) + self.assertTrue(np.isnan(residual[0, 2])) + self.assertTrue(np.isnan(expected[0, 2])) + + def test_fixed_kernel_fit_is_deterministic(self): + from behavior_analyses.kernels import fit_fixed_psychophysical_kernel + + rng = np.random.default_rng(4) + residual = rng.normal(size=(120, 4)) + choices = (residual[:, 0] > 0).astype(int) + residual[60:, 3] = np.nan + n_observed = np.sum(np.isfinite(residual), axis=0) + + first = fit_fixed_psychophysical_kernel( + residual, + choices, + n_observed_per_bin=n_observed, + cv_splits=5, + random_state=7, + min_trials_per_bin=50, + ) + second = fit_fixed_psychophysical_kernel( + residual, + choices, + n_observed_per_bin=n_observed, + cv_splits=5, + random_state=7, + min_trials_per_bin=50, + ) + + self.assertTrue(first["fit_converged"]) + self.assertEqual(first["n_bins_fit"], 4) + np.testing.assert_allclose(first["weights_mean"], second["weights_mean"]) + np.testing.assert_allclose(first["scores"], second["scores"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_cli_and_migration_contracts.py b/tests/test_cli_and_migration_contracts.py index bdf9eff..c34178c 100644 --- a/tests/test_cli_and_migration_contracts.py +++ b/tests/test_cli_and_migration_contracts.py @@ -51,6 +51,20 @@ def populate(*_args, **_kwargs): class CliContractTests(unittest.TestCase): + def test_kernel_timing_migration_only_adds_missing_columns(self): + module = runpy.run_path(str(SCRIPTS / "migrate_kernel_timing_source.py")) + + statements = module["_missing_column_statements"]( + "labdata_user", + module["KERNEL_TABLE"], + module["KERNEL_COLUMNS"], + {"timing_source", "n_trials_fit", "fit_message"}, + ) + + self.assertEqual(len(statements), len(module["KERNEL_COLUMNS"]) - 1) + self.assertTrue(all("timing_source" not in sql for sql in statements)) + self.assertTrue(all(sql.startswith("ALTER TABLE") for sql in statements)) + def test_schema_migration_archive_names_leave_room_for_foreign_keys(self): module = runpy.run_path(str(SCRIPTS / "migrate_behavior_analysis_schema.py")) @@ -206,6 +220,17 @@ def test_configured_path_reads_env(self): path = io_mod._configured_chipmunk_plugin_path() self.assertEqual(path, Path("/tmp/chipmunk-plugin")) + def test_registered_chipmunk_plugin_is_used(self): + from behavior_analyses import io as io_mod + + table = object() + fake_labdata = types.ModuleType("labdata") + fake_labdata.plugins = {"chipmunk": types.SimpleNamespace(Chipmunk=table)} + with patch.dict(sys.modules, {"labdata": fake_labdata}): + registered = io_mod._registered_chipmunk_table() + + self.assertIs(registered, table) + class PsychometricPlotHelperTests(unittest.TestCase): def test_fetch_uses_labdata_rates_and_session_names(self): diff --git a/tests/test_kernel_timing.py b/tests/test_kernel_timing.py new file mode 100644 index 0000000..bf32979 --- /dev/null +++ b/tests/test_kernel_timing.py @@ -0,0 +1,161 @@ +from __future__ import annotations + +from pathlib import Path +import sys +import unittest + +import numpy as np + + +REPO_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(REPO_ROOT / "src")) + + +class KernelTimingTests(unittest.TestCase): + def setUp(self): + self.trial_rows = [ + { + "trial_num": 0, + "rewarded_modality": "visual", + "stim_events": np.array([0.2, 0.35, 0.7]), + "stim_rate_vision": 12.0, + "response": 1, + "t_sync": 0.0, + "t_react": 0.5, + "t_response": 0.9, + }, + { + "trial_num": 1, + "rewarded_modality": "visual", + "stim_events": np.array([2.2]), + "stim_rate_vision": 8.0, + "response": 0, + "t_sync": 2.0, + "t_react": None, + "t_response": None, + }, + ] + + def test_bpod_window_uses_reaction_or_response_time(self): + from behavior_analyses.kernel_timing import extract_bpod_kernel_inputs + + center = extract_bpod_kernel_inputs( + self.trial_rows, "visual", observation_window="center_exit" + ) + response = extract_bpod_kernel_inputs( + self.trial_rows, "visual", observation_window="response" + ) + + np.testing.assert_allclose(center["stim_times_per_trial"][0], [0.2, 0.35]) + np.testing.assert_allclose( + response["stim_times_per_trial"][0], [0.2, 0.35, 0.7] + ) + np.testing.assert_allclose(center["observation_end_times"], [0.5]) + np.testing.assert_allclose(response["observation_end_times"], [0.9]) + + def test_nidaq_visual_mapping_requires_complete_valid_events(self): + from behavior_analyses.kernel_timing import resolve_nidaq_event_arrays + + with self.assertRaisesRegex(ValueError, "Missing EventMapping"): + resolve_nidaq_event_arrays( + [], + [ + { + "event_name": "visual_stim", + "source_dataset_name": "ephys", + "source_stream_name": "nidq", + "source_event_name": "ai0", + } + ], + "GRB006", + "session", + ) + + def test_nidaq_window_uses_mapped_flash_and_port_times(self): + from behavior_analyses.kernel_timing import ( + extract_nidaq_kernel_inputs, + resolve_nidaq_event_arrays, + ) + + mapping_rows, event_rows = self._mapped_event_fixture() + aligned = resolve_nidaq_event_arrays( + event_rows, mapping_rows, "GRB006", "session" + ) + center = extract_nidaq_kernel_inputs( + aligned, + self.trial_rows, + "visual", + observation_window="center_exit", + ) + response = extract_nidaq_kernel_inputs( + aligned, + self.trial_rows, + "visual", + observation_window="response", + ) + + np.testing.assert_allclose(center["stim_times_per_trial"][0], [0.2, 0.35]) + np.testing.assert_allclose( + response["stim_times_per_trial"][0], [0.2, 0.35, 0.7] + ) + np.testing.assert_allclose(center["observation_end_times"], [0.5]) + np.testing.assert_allclose(response["observation_end_times"], [0.9]) + + def test_combined_provenance_is_mixed(self): + from behavior_analyses.kernel_timing import ( + combine_kernel_inputs, + extract_bpod_kernel_inputs, + ) + + first = extract_bpod_kernel_inputs( + self.trial_rows, "visual", observation_window="center_exit" + ) + second = extract_bpod_kernel_inputs( + self.trial_rows, "visual", observation_window="center_exit" + ) + first["timing_source"] = "bpod" + second["timing_source"] = "nidaq" + + combined = combine_kernel_inputs([first, second]) + + self.assertEqual(combined["timing_source"], "mixed") + self.assertEqual(len(combined["stim_times_per_trial"]), 2) + + @staticmethod + def _mapped_event_fixture(): + sources = { + "visual_stim": ("ai0", [0.2, 0.35, 0.7, 2.2], None), + "trial_start": ("line0", [0.0, 2.0], [1, 1]), + "left_port": ("line1", [2.3, 2.4], [1, 0]), + "center_port": ( + "line2", + [0.1, 0.5, 2.1, 2.5], + [1, 0, 1, 0], + ), + "right_port": ("line3", [0.9, 1.0], [1, 0]), + } + mapping_rows = [] + event_rows = [] + for logical_name, (source_name, timestamps, values) in sources.items(): + mapping_rows.append( + { + "event_name": logical_name, + "source_dataset_name": "ephys", + "source_stream_name": "nidq", + "source_event_name": source_name, + } + ) + event_row = { + "dataset_name": "ephys", + "stream_name": "nidq", + "event_name": source_name, + "event_timestamps": np.asarray(timestamps, dtype=float), + } + if values is not None: + event_row["event_values"] = np.asarray(values) + event_rows.append(event_row) + return mapping_rows, event_rows + + +if __name__ == "__main__": + unittest.main()