summaryrefslogtreecommitdiff
path: root/experiments/larkum_public_pilot.py
diff options
context:
space:
mode:
Diffstat (limited to 'experiments/larkum_public_pilot.py')
-rw-r--r--experiments/larkum_public_pilot.py460
1 files changed, 460 insertions, 0 deletions
diff --git a/experiments/larkum_public_pilot.py b/experiments/larkum_public_pilot.py
new file mode 100644
index 0000000..67773c8
--- /dev/null
+++ b/experiments/larkum_public_pilot.py
@@ -0,0 +1,460 @@
+#!/usr/bin/env python3
+"""Frozen public-data pilot for dendritic innovation and trial outcomes."""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Iterable
+
+import numpy as np
+from scipy.io import loadmat
+from sklearn.linear_model import LogisticRegression, Ridge
+from sklearn.metrics import roc_auc_score
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import StandardScaler
+
+
+PROTOCOL = "larkum_public_innovation_pilot_v1"
+CONDITIONS = ("Sal", "DCZ")
+METHODS = ("raw_branch", "task_residual", "innovation", "local_spine", "choice_only")
+N_PREDICTOR_FOLDS = 5
+RIDGE_ALPHA = 1.0
+BOOTSTRAP_SAMPLES = 5000
+BOOTSTRAP_SEED = 20270810
+
+
+@dataclass(frozen=True)
+class Trial:
+ animal: str
+ session: str
+ condition: str
+ trial_index: int
+ instruction: int
+ choice: int
+ correct: int
+ features: dict[str, np.ndarray]
+
+
+def _trial_array(value: object) -> np.ndarray:
+ array = np.asarray(value, dtype=np.float64)
+ if array.ndim != 2 or array.shape[0] != 180 or array.shape[1] < 1:
+ raise ValueError(f"expected a 180 x ROI trial array, found {array.shape}")
+ if not np.isfinite(array).all():
+ raise ValueError("non-finite neural activity")
+ return array - array[:30].mean(axis=0, keepdims=True)
+
+
+def _design(spines: np.ndarray, instruction: int) -> np.ndarray:
+ frames = spines.shape[0]
+ frame_basis = np.eye(frames, dtype=np.float64)
+ interaction = frame_basis * float(instruction)
+ return np.concatenate((spines, frame_basis, interaction), axis=1)
+
+
+def _fit_predictor(
+ spine_trials: list[np.ndarray],
+ branch_trials: list[np.ndarray],
+ instructions: np.ndarray,
+ train_indices: np.ndarray,
+) -> tuple[object, np.ndarray]:
+ x = np.concatenate([_design(spine_trials[i], int(instructions[i])) for i in train_indices])
+ y = np.concatenate([branch_trials[i] for i in train_indices])
+ model = make_pipeline(StandardScaler(), Ridge(alpha=RIDGE_ALPHA))
+ model.fit(x, y)
+ return model, y
+
+
+def _predict_trials(
+ model: object,
+ spine_trials: list[np.ndarray],
+ instructions: np.ndarray,
+ indices: Iterable[int],
+) -> dict[int, np.ndarray]:
+ return {
+ int(i): np.asarray(model.predict(_design(spine_trials[i], int(instructions[i]))))
+ for i in indices
+ }
+
+
+def _template_predict(
+ branch_trials: list[np.ndarray],
+ instructions: np.ndarray,
+ train_indices: np.ndarray,
+ test_indices: Iterable[int],
+) -> dict[int, np.ndarray]:
+ templates: dict[int, np.ndarray] = {}
+ for side in (0, 1):
+ selected = [branch_trials[i] for i in train_indices if instructions[i] == side]
+ if not selected:
+ raise ValueError(f"training fold has no instruction-{side} trials")
+ templates[side] = np.mean(selected, axis=0)
+ return {int(i): templates[int(instructions[i])] for i in test_indices}
+
+
+def _summarize(trace: np.ndarray) -> np.ndarray:
+ time_axis = np.linspace(-3.0, 3.0, 180)
+ indices = np.flatnonzero((time_axis >= 0.0) & (time_axis < 1.0))
+ bins = np.array_split(indices, 6)
+ values: list[float] = []
+ for frame_indices in bins:
+ block = trace[frame_indices]
+ values.extend((float(block.mean()), float(np.abs(block).mean()), float(np.sqrt(np.mean(block**2)))))
+ result = np.asarray(values, dtype=np.float64)
+ if result.shape != (18,) or not np.isfinite(result).all():
+ raise ValueError("invalid summary features")
+ return result
+
+
+def _session_trials(path: Path) -> list[Trial]:
+ session = path.stem
+ animal = session.split("_")[0]
+ data = loadmat(path, simplify_cells=True)["cont_data"]
+ prepared: dict[str, dict[str, object]] = {}
+
+ for condition in CONDITIONS:
+ condition_data = data[condition]
+ choices = np.asarray(condition_data["Choice"])[:, 0]
+ instructions = np.asarray(condition_data["TrialTypes"])[:, 0].astype(np.int64)
+ correct = np.asarray(condition_data["DirOut"]).astype(np.int64)
+ valid_indices: list[int] = []
+ spines: list[np.ndarray] = []
+ branches: list[np.ndarray] = []
+ for i in range(len(correct)):
+ if not np.isfinite(choices[i]):
+ continue
+ try:
+ spine = _trial_array(condition_data["spine_local"][i])
+ branch = _trial_array(condition_data["branch"][i])
+ except ValueError:
+ continue
+ valid_indices.append(i)
+ spines.append(spine)
+ branches.append(branch)
+ prepared[condition] = {
+ "source_indices": np.asarray(valid_indices, dtype=np.int64),
+ "choices": choices[valid_indices].astype(np.int64),
+ "instructions": instructions[valid_indices],
+ "correct": correct[valid_indices],
+ "spines": spines,
+ "branches": branches,
+ }
+
+ saline = prepared["Sal"]
+ sal_spines = saline["spines"]
+ sal_branches = saline["branches"]
+ sal_instructions = saline["instructions"]
+ n_sal = len(sal_spines)
+ if n_sal < N_PREDICTOR_FOLDS:
+ raise ValueError(f"{session} has only {n_sal} usable saline trials")
+
+ crossfit_prediction: dict[int, np.ndarray] = {}
+ crossfit_template: dict[int, np.ndarray] = {}
+ for test_indices in np.array_split(np.arange(n_sal), N_PREDICTOR_FOLDS):
+ train_indices = np.setdiff1d(np.arange(n_sal), test_indices)
+ model, _ = _fit_predictor(sal_spines, sal_branches, sal_instructions, train_indices)
+ crossfit_prediction.update(_predict_trials(model, sal_spines, sal_instructions, test_indices))
+ crossfit_template.update(
+ _template_predict(sal_branches, sal_instructions, train_indices, test_indices)
+ )
+
+ all_sal_indices = np.arange(n_sal)
+ dcz = prepared["DCZ"]
+ dcz_model, _ = _fit_predictor(sal_spines, sal_branches, sal_instructions, all_sal_indices)
+ dcz_prediction = _predict_trials(
+ dcz_model, dcz["spines"], dcz["instructions"], range(len(dcz["spines"]))
+ )
+ dcz_template = _template_predict(
+ sal_branches,
+ sal_instructions,
+ all_sal_indices,
+ range(len(dcz["spines"])),
+ )
+
+ output: list[Trial] = []
+ for condition in CONDITIONS:
+ entry = prepared[condition]
+ predictions = crossfit_prediction if condition == "Sal" else dcz_prediction
+ templates = crossfit_template if condition == "Sal" else dcz_template
+ for i, source_index in enumerate(entry["source_indices"]):
+ branch = entry["branches"][i]
+ spine = entry["spines"][i]
+ innovation = branch - predictions[i]
+ task_residual = branch - templates[i]
+ output.append(
+ Trial(
+ animal=animal,
+ session=session,
+ condition=condition,
+ trial_index=int(source_index),
+ instruction=int(entry["instructions"][i]),
+ choice=int(entry["choices"][i]),
+ correct=int(entry["correct"][i]),
+ features={
+ "raw_branch": _summarize(branch),
+ "task_residual": _summarize(task_residual),
+ "innovation": _summarize(innovation),
+ "local_spine": _summarize(spine),
+ "choice_only": np.asarray([float(entry["choices"][i])]),
+ },
+ )
+ )
+ return output
+
+
+def _decode(trials: list[Trial], condition: str, method: str) -> list[dict[str, object]]:
+ selected = [trial for trial in trials if trial.condition == condition]
+ animals = sorted({trial.animal for trial in selected})
+ predictions: list[dict[str, object]] = []
+ for held_animal in animals:
+ for train_instruction in (0, 1):
+ test_instruction = 1 - train_instruction
+ train = [
+ trial
+ for trial in selected
+ if trial.animal != held_animal and trial.instruction == train_instruction
+ ]
+ test = [
+ trial
+ for trial in selected
+ if trial.animal == held_animal and trial.instruction == test_instruction
+ ]
+ y_train = np.asarray([trial.correct for trial in train], dtype=np.int64)
+ if len(np.unique(y_train)) != 2 or not test:
+ raise ValueError(
+ f"invalid decoder fold: {condition=} {method=} {held_animal=} {train_instruction=}"
+ )
+ x_train = np.stack([trial.features[method] for trial in train])
+ x_test = np.stack([trial.features[method] for trial in test])
+ decoder = make_pipeline(
+ StandardScaler(),
+ LogisticRegression(C=1.0, class_weight="balanced", max_iter=5000, random_state=0),
+ )
+ decoder.fit(x_train, y_train)
+ probabilities = decoder.predict_proba(x_test)[:, 1]
+ for trial, probability in zip(test, probabilities):
+ predictions.append(
+ {
+ "animal": trial.animal,
+ "session": trial.session,
+ "condition": condition,
+ "method": method,
+ "trial_index": trial.trial_index,
+ "instruction": trial.instruction,
+ "choice": trial.choice,
+ "correct": trial.correct,
+ "probability_correct": float(probability),
+ }
+ )
+ expected = len(selected)
+ if len(predictions) != expected:
+ raise RuntimeError(f"decoder produced {len(predictions)} predictions for {expected} trials")
+ keys = {(p["session"], p["condition"], p["trial_index"]) for p in predictions}
+ if len(keys) != expected:
+ raise RuntimeError("duplicate or missing out-of-fold predictions")
+ return predictions
+
+
+def _auc(rows: list[dict[str, object]]) -> float:
+ labels = np.asarray([row["correct"] for row in rows], dtype=np.int64)
+ scores = np.asarray([row["probability_correct"] for row in rows], dtype=np.float64)
+ if len(np.unique(labels)) != 2:
+ return float("nan")
+ return float(roc_auc_score(labels, scores))
+
+
+def _bootstrap_plan(
+ rows: list[dict[str, object]], rng: np.random.Generator
+) -> dict[str, list[str]]:
+ plan: dict[str, list[str]] = {}
+ animals = sorted({str(row["animal"]) for row in rows})
+ for animal in animals:
+ sessions = sorted({str(row["session"]) for row in rows if row["animal"] == animal})
+ plan[animal] = [str(value) for value in rng.choice(sessions, size=len(sessions), replace=True)]
+ return plan
+
+
+def _bootstrap_indices(
+ rows: list[dict[str, object]], plan: dict[str, list[str]]
+) -> np.ndarray:
+ indices: list[int] = []
+ for animal, sessions in plan.items():
+ for session in sessions:
+ indices.extend(
+ i
+ for i, row in enumerate(rows)
+ if row["animal"] == animal and row["session"] == session
+ )
+ return np.asarray(indices, dtype=np.int64)
+
+
+def _bootstrap_auc(
+ rows: list[dict[str, object]], sample_plans: list[dict[str, list[str]]]
+) -> np.ndarray:
+ values: list[float] = []
+ for plan in sample_plans:
+ indices = _bootstrap_indices(rows, plan)
+ sample = [rows[i] for i in indices]
+ value = _auc(sample)
+ if np.isfinite(value):
+ values.append(value)
+ return np.asarray(values, dtype=np.float64)
+
+
+def _interval(values: np.ndarray) -> list[float]:
+ return [float(v) for v in np.quantile(values, (0.025, 0.975))]
+
+
+def _summarize_predictions(all_predictions: list[dict[str, object]]) -> dict[str, object]:
+ grouped: dict[tuple[str, str], list[dict[str, object]]] = {}
+ for row in all_predictions:
+ grouped.setdefault((str(row["condition"]), str(row["method"])), []).append(row)
+
+ rng = np.random.default_rng(BOOTSTRAP_SEED)
+ reference = grouped[("Sal", "raw_branch")]
+ sample_plans = [_bootstrap_plan(reference, rng) for _ in range(BOOTSTRAP_SAMPLES)]
+ results: dict[str, object] = {}
+ bootstrap_values: dict[tuple[str, str], np.ndarray] = {}
+ for condition in CONDITIONS:
+ results[condition] = {}
+ for method in METHODS:
+ rows = grouped[(condition, method)]
+ values = _bootstrap_auc(rows, sample_plans)
+ bootstrap_values[(condition, method)] = values
+ by_animal = {}
+ for animal in sorted({str(row["animal"]) for row in rows}):
+ animal_rows = [row for row in rows if row["animal"] == animal]
+ by_animal[animal] = {
+ "auc": _auc(animal_rows),
+ "trials": len(animal_rows),
+ "errors": int(sum(1 - int(row["correct"]) for row in animal_rows)),
+ }
+ results[condition][method] = {
+ "auc": _auc(rows),
+ "session_block_bootstrap_95ci": _interval(values),
+ "trials": len(rows),
+ "errors": int(sum(1 - int(row["correct"]) for row in rows)),
+ "by_animal": by_animal,
+ }
+
+ comparisons: dict[str, object] = {}
+ for condition in CONDITIONS:
+ difference = (
+ bootstrap_values[(condition, "innovation")]
+ - bootstrap_values[(condition, "raw_branch")]
+ )
+ comparisons[f"innovation_minus_raw_{condition.lower()}"] = {
+ "difference": float(
+ results[condition]["innovation"]["auc"]
+ - results[condition]["raw_branch"]["auc"]
+ ),
+ "session_block_bootstrap_95ci": _interval(difference),
+ "bootstrap_probability_positive": float(np.mean(difference > 0.0)),
+ }
+ interaction = (
+ bootstrap_values[("Sal", "innovation")]
+ - bootstrap_values[("DCZ", "innovation")]
+ )
+ comparisons["innovation_sal_minus_dcz"] = {
+ "difference": float(
+ results["Sal"]["innovation"]["auc"]
+ - results["DCZ"]["innovation"]["auc"]
+ ),
+ "session_block_bootstrap_95ci": _interval(interaction),
+ "bootstrap_probability_positive": float(np.mean(interaction > 0.0)),
+ }
+ return {"endpoints": results, "comparisons": comparisons}
+
+
+def _write_predictions(path: Path, rows: list[dict[str, object]]) -> None:
+ path.parent.mkdir(parents=True, exist_ok=True)
+ fields = [
+ "animal",
+ "session",
+ "condition",
+ "method",
+ "trial_index",
+ "instruction",
+ "choice",
+ "correct",
+ "probability_correct",
+ ]
+ with path.open("w", newline="", encoding="utf-8") as handle:
+ writer = csv.DictWriter(handle, fieldnames=fields)
+ writer.writeheader()
+ writer.writerows(rows)
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--data-root", type=Path, required=True, help="Extracted public archive root")
+ parser.add_argument("--output", type=Path, default=Path("results/larkum_public_pilot.json"))
+ parser.add_argument(
+ "--predictions",
+ type=Path,
+ default=Path("results/larkum_public_pilot_predictions.csv"),
+ )
+ parser.add_argument("--audit-only", action="store_true", help="Report eligible files and counts only")
+ args = parser.parse_args()
+
+ files = sorted((args.data_root / "Figure2" / "Data").glob("*_dff.mat"))
+ if len(files) != 11:
+ raise SystemExit(f"expected 11 Figure 2 session files, found {len(files)}")
+ if args.audit_only:
+ print(json.dumps({"protocol": PROTOCOL, "files": [path.name for path in files]}, indent=2))
+ return
+
+ trials: list[Trial] = []
+ for path in files:
+ trials.extend(_session_trials(path))
+
+ predictions: list[dict[str, object]] = []
+ for condition in CONDITIONS:
+ for method in METHODS:
+ predictions.extend(_decode(trials, condition, method))
+ summary = _summarize_predictions(predictions)
+
+ counts = {}
+ for condition in CONDITIONS:
+ subset = [trial for trial in trials if trial.condition == condition]
+ counts[condition] = {
+ "trials": len(subset),
+ "correct": int(sum(trial.correct for trial in subset)),
+ "errors": int(sum(1 - trial.correct for trial in subset)),
+ "animals": sorted({trial.animal for trial in subset}),
+ "sessions": len({trial.session for trial in subset}),
+ }
+ result = {
+ "protocol": PROTOCOL,
+ "source": {
+ "paper_doi": "10.1126/science.adx4358",
+ "dataset_doi": "10.12751/g-node.etlk5k",
+ "files": [path.name for path in files],
+ },
+ "fixed_configuration": {
+ "ridge_alpha": RIDGE_ALPHA,
+ "predictor_folds": N_PREDICTOR_FOLDS,
+ "decoder": "StandardScaler + balanced logistic regression, C=1.0",
+ "evaluation": "cross-instruction leave-one-animal-out",
+ "outcome_window_seconds": [0.0, 1.0],
+ "bootstrap_samples": BOOTSTRAP_SAMPLES,
+ "bootstrap_seed": BOOTSTRAP_SEED,
+ },
+ "counts": counts,
+ **summary,
+ "claim_boundary": (
+ "Exploratory biological consistency test only; the release has three animals, "
+ "no paired soma signal, and no trial-linked relearning endpoint."
+ ),
+ }
+ args.output.parent.mkdir(parents=True, exist_ok=True)
+ args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
+ _write_predictions(args.predictions, predictions)
+ print(json.dumps(result, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+ main()
generated by cgit v1.2.3 (git 2.25.1) at 2026-09-15 15:05:59 +0000