ephys-mcp 0.1.0__tar.gz

ephys_mcp-0.1.0/PKG-INFO

@@ -0,0 +1,123 @@
+Metadata-Version: 2.4
+Name: ephys-mcp
+Version: 0.1.0
+Summary: MCP server for read-only analysis of intracortical (spike-level) brain-computer-interface recordings
+Keywords: mcp,bci,electrophysiology,neuroscience,spikes,nwb,dandi
+Author: happyc0der
+Author-email: happyc0der <Keshav.Rajput4@gmail.com>
+License-Expression: CC0-1.0
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: CC0 1.0 Universal (CC0 1.0) Public Domain Dedication
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering
+Requires-Dist: h5py>=3.16.0
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: matplotlib>=3.11.2
+Requires-Dist: mcp[cli]>=2.2.0
+Requires-Dist: numpy>=2.5.3
+Requires-Dist: pynwb>=4.2.0
+Requires-Dist: remfile>=0.1.14
+Requires-Dist: scipy>=1.18.1
+Requires-Python: >=3.12
+Project-URL: Homepage, https://github.com/happyc0der/ephys-mcp
+Project-URL: Repository, https://github.com/happyc0der/ephys-mcp
+Project-URL: Issues, https://github.com/happyc0der/ephys-mcp/issues
+Description-Content-Type: text/markdown
+
+# ephys-mcp
+
+An [MCP](https://modelcontextprotocol.io) server that lets an LLM analyse **intracortical (spike-level) brain-computer-interface recordings**: signal quality, spike detection, firing rates, and cursor-velocity decoding.
+
+Existing BCI MCP servers target scalp EEG. This one targets the kind of data a high-channel-count implant produces, and defines a read-only adapter contract so a live device backend can be added when a vendor publishes an API.
+
+> Research and education software. **Not a medical device. Not for clinical use.**
+> Not affiliated with or endorsed by Neuralink Corp. or any other implant manufacturer.
+
+## Status
+
+v0.1, early. Working today: local NWB files, local broadband WAV recordings, streaming from the DANDI Archive, a synthetic motor-cortex source with ground truth, spike detection, quality metrics, ridge and Kalman decoders, trial-aligned PSTHs, and figures. Planned: PyPI release and registry listings.
+
+## Install and run
+
+Needs [uv](https://docs.astral.sh/uv/). No install step: `uvx ephys-mcp` fetches the package and starts the server on stdio.
+
+Claude Code:
+
+```bash
+claude mcp add ephys -- uvx ephys-mcp
+```
+
+Claude Desktop (`claude_desktop_config.json`):
+
+```json
+{ "mcpServers": { "ephys": { "command": "uvx", "args": ["ephys-mcp"] } } }
+```
+
+From a checkout, use `uv run ephys-mcp` instead, or `uv --directory /path/to/ephys-mcp run ephys-mcp` in the configs above.
+
+Then ask, for real data: *"Find a small motor cortex dataset on DANDI, open it, and tell me how well hand velocity can be decoded."*
+Or offline: *"Open a synthetic session, check signal quality, fit a Kalman decoder and show me a decoded window."*
+
+## Data sources
+
+| Source | What it opens |
+| --- | --- |
+| `synthetic` | Simulated units tuned to cursor velocity, with broadband signal and ground truth |
+| `nwb` | A local `.nwb` file (`params.path`) |
+| `wav_dir` | Local broadband WAV (`params.path`): a folder of mono clips, one channel each, or one multi-channel file |
+| `dandi` | An NWB file streamed from the [DANDI Archive](https://dandiarchive.org) by HTTP range requests; nothing is mirrored |
+| `n1_stub` | Not implemented. Documents the contract for a live implant adapter |
+
+Dataset licence and citation come from the archive and are returned by `open_session`, so the model can attribute the data. Many datasets record only during trials; the server tracks those spans (`recorded_fraction`) and leaves the gaps out of rates and decoding instead of reading them as silence.
+
+WAV samples carry no physical unit, so amplitudes are reported as ADC counts unless you pass `uv_per_count`; every amplitude result names its unit. Clips in a folder are separate recordings, so the server says that timing across those channels is not meaningful. Spike times from WAV are threshold crossings, not sorted units.
+
+Reference result on MC_Maze_Small (DANDI 000140, 142 units, last 20% held out, 50 ms bins): ridge R² 0.50, Kalman R² 0.34 for hand velocity. These are simple causal linear baselines, not state of the art.
+
+## Tools
+
+| Tool | Purpose |
+| --- | --- |
+| `list_sources` | Source types and their parameters |
+| `search_datasets` | Search DANDI, or list curated intracortical datasets |
+| `list_dataset_files` | Licence, citation and NWB files of a DANDI dataset |
+| `open_session` / `close_session` | Session lifecycle |
+| `get_session_info` | Channels, rates, behaviour signals, licence, citation |
+| `get_signal_quality` | Noise, SNR, dead/noisy channels |
+| `detect_spikes` | Threshold crossings; precision/recall when truth exists |
+| `get_firing_rates` | Population rate summary |
+| `fit_decoder` | Ridge or Kalman, scored on held-out data; hyperparameters chosen inside the training split |
+| `decode_window` | Decoded-vs-true preview for a window |
+| `get_psth` | Firing aligned to a trial event, optionally grouped by a trial column or limited to some units |
+| `plot_psth` | Figure: PSTH per group with SEM, above a unit-by-time heatmap of change from baseline |
+| `plot_raster` | Figure: spike raster, unrecorded spans shaded |
+| `plot_decoding` | Figure: decoded against actual behaviour, one panel per dimension |
+
+Resource: `ephys://sessions`. Prompt: `analyze_session`.
+
+Tools return summaries, never raw arrays, so results fit in a model's context.
+
+Plot tools return the PNG inline, so a vision-capable model can read the figure, and also save it under `~/.cache/ephys-mcp/plots` (override with `EPHYS_MCP_OUTPUT_DIR`). Figures use a categorical palette checked for colour-blind separation, with direct labels so identity never rests on colour alone.
+
+## Design rules
+
+- **Read-only.** The `NeuralSource` contract has no write, stimulate or configure method. None will be added without a separate safety design.
+- **Local by default.** stdio transport, no telemetry. Neural data is sensitive.
+- **No bundled third-party data.** See [DATA_LICENSES.md](DATA_LICENSES.md).
+
+## Writing a source adapter
+
+Subclass `ephys_mcp.sources.base.NeuralSource` (`info`, `read_raw`, `spike_times`, `behavior`) and register it in `ephys_mcp/sources/__init__.py`. `sources/n1_stub.py` documents what a live implant adapter would need.
+
+## Development
+
+```bash
+uv run pytest              # offline
+uv run pytest -m network   # also streams a real file from DANDI
+uv run ruff check .
+```
+
+## Licence
+
+[CC0 1.0 Universal](LICENSE). The authors waive all copyright and related rights to the extent the law allows. Use it for anything, no attribution required. CC0 does not grant patent or trademark rights.

ephys_mcp-0.1.0/README.md

@@ -0,0 +1,96 @@
+# ephys-mcp
+
+An [MCP](https://modelcontextprotocol.io) server that lets an LLM analyse **intracortical (spike-level) brain-computer-interface recordings**: signal quality, spike detection, firing rates, and cursor-velocity decoding.
+
+Existing BCI MCP servers target scalp EEG. This one targets the kind of data a high-channel-count implant produces, and defines a read-only adapter contract so a live device backend can be added when a vendor publishes an API.
+
+> Research and education software. **Not a medical device. Not for clinical use.**
+> Not affiliated with or endorsed by Neuralink Corp. or any other implant manufacturer.
+
+## Status
+
+v0.1, early. Working today: local NWB files, local broadband WAV recordings, streaming from the DANDI Archive, a synthetic motor-cortex source with ground truth, spike detection, quality metrics, ridge and Kalman decoders, trial-aligned PSTHs, and figures. Planned: PyPI release and registry listings.
+
+## Install and run
+
+Needs [uv](https://docs.astral.sh/uv/). No install step: `uvx ephys-mcp` fetches the package and starts the server on stdio.
+
+Claude Code:
+
+```bash
+claude mcp add ephys -- uvx ephys-mcp
+```
+
+Claude Desktop (`claude_desktop_config.json`):
+
+```json
+{ "mcpServers": { "ephys": { "command": "uvx", "args": ["ephys-mcp"] } } }
+```
+
+From a checkout, use `uv run ephys-mcp` instead, or `uv --directory /path/to/ephys-mcp run ephys-mcp` in the configs above.
+
+Then ask, for real data: *"Find a small motor cortex dataset on DANDI, open it, and tell me how well hand velocity can be decoded."*
+Or offline: *"Open a synthetic session, check signal quality, fit a Kalman decoder and show me a decoded window."*
+
+## Data sources
+
+| Source | What it opens |
+| --- | --- |
+| `synthetic` | Simulated units tuned to cursor velocity, with broadband signal and ground truth |
+| `nwb` | A local `.nwb` file (`params.path`) |
+| `wav_dir` | Local broadband WAV (`params.path`): a folder of mono clips, one channel each, or one multi-channel file |
+| `dandi` | An NWB file streamed from the [DANDI Archive](https://dandiarchive.org) by HTTP range requests; nothing is mirrored |
+| `n1_stub` | Not implemented. Documents the contract for a live implant adapter |
+
+Dataset licence and citation come from the archive and are returned by `open_session`, so the model can attribute the data. Many datasets record only during trials; the server tracks those spans (`recorded_fraction`) and leaves the gaps out of rates and decoding instead of reading them as silence.
+
+WAV samples carry no physical unit, so amplitudes are reported as ADC counts unless you pass `uv_per_count`; every amplitude result names its unit. Clips in a folder are separate recordings, so the server says that timing across those channels is not meaningful. Spike times from WAV are threshold crossings, not sorted units.
+
+Reference result on MC_Maze_Small (DANDI 000140, 142 units, last 20% held out, 50 ms bins): ridge R² 0.50, Kalman R² 0.34 for hand velocity. These are simple causal linear baselines, not state of the art.
+
+## Tools
+
+| Tool | Purpose |
+| --- | --- |
+| `list_sources` | Source types and their parameters |
+| `search_datasets` | Search DANDI, or list curated intracortical datasets |
+| `list_dataset_files` | Licence, citation and NWB files of a DANDI dataset |
+| `open_session` / `close_session` | Session lifecycle |
+| `get_session_info` | Channels, rates, behaviour signals, licence, citation |
+| `get_signal_quality` | Noise, SNR, dead/noisy channels |
+| `detect_spikes` | Threshold crossings; precision/recall when truth exists |
+| `get_firing_rates` | Population rate summary |
+| `fit_decoder` | Ridge or Kalman, scored on held-out data; hyperparameters chosen inside the training split |
+| `decode_window` | Decoded-vs-true preview for a window |
+| `get_psth` | Firing aligned to a trial event, optionally grouped by a trial column or limited to some units |
+| `plot_psth` | Figure: PSTH per group with SEM, above a unit-by-time heatmap of change from baseline |
+| `plot_raster` | Figure: spike raster, unrecorded spans shaded |
+| `plot_decoding` | Figure: decoded against actual behaviour, one panel per dimension |
+
+Resource: `ephys://sessions`. Prompt: `analyze_session`.
+
+Tools return summaries, never raw arrays, so results fit in a model's context.
+
+Plot tools return the PNG inline, so a vision-capable model can read the figure, and also save it under `~/.cache/ephys-mcp/plots` (override with `EPHYS_MCP_OUTPUT_DIR`). Figures use a categorical palette checked for colour-blind separation, with direct labels so identity never rests on colour alone.
+
+## Design rules
+
+- **Read-only.** The `NeuralSource` contract has no write, stimulate or configure method. None will be added without a separate safety design.
+- **Local by default.** stdio transport, no telemetry. Neural data is sensitive.
+- **No bundled third-party data.** See [DATA_LICENSES.md](DATA_LICENSES.md).
+
+## Writing a source adapter
+
+Subclass `ephys_mcp.sources.base.NeuralSource` (`info`, `read_raw`, `spike_times`, `behavior`) and register it in `ephys_mcp/sources/__init__.py`. `sources/n1_stub.py` documents what a live implant adapter would need.
+
+## Development
+
+```bash
+uv run pytest              # offline
+uv run pytest -m network   # also streams a real file from DANDI
+uv run ruff check .
+```
+
+## Licence
+
+[CC0 1.0 Universal](LICENSE). The authors waive all copyright and related rights to the extent the law allows. Use it for anything, no attribution required. CC0 does not grant patent or trademark rights.

ephys_mcp-0.1.0/pyproject.toml

@@ -0,0 +1,68 @@
+[project]
+name = "ephys-mcp"
+version = "0.1.0"
+description = "MCP server for read-only analysis of intracortical (spike-level) brain-computer-interface recordings"
+license = "CC0-1.0"
+keywords = [
+    "mcp",
+    "bci",
+    "electrophysiology",
+    "neuroscience",
+    "spikes",
+    "nwb",
+    "dandi",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: CC0 1.0 Universal (CC0 1.0) Public Domain Dedication",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+]
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "h5py>=3.16.0",
+    "httpx>=0.28.1",
+    "matplotlib>=3.11.2",
+    "mcp[cli]>=2.2.0",
+    "numpy>=2.5.3",
+    "pynwb>=4.2.0",
+    "remfile>=0.1.14",
+    "scipy>=1.18.1",
+]
+
+[[project.authors]]
+name = "happyc0der"
+email = "Keshav.Rajput4@gmail.com"
+
+[project.urls]
+Homepage = "https://github.com/happyc0der/ephys-mcp"
+Repository = "https://github.com/happyc0der/ephys-mcp"
+Issues = "https://github.com/happyc0der/ephys-mcp/issues"
+
+[project.scripts]
+ephys-mcp = "ephys_mcp.server:main"
+
+[build-system]
+requires = ["uv_build>=0.12.10,<0.13.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
+    "ruff>=0.16.8",
+]
+
+[tool.ruff]
+line-length = 120
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+markers = ["network: needs internet access to the DANDI Archive"]
+addopts = "-m 'not network'"
+filterwarnings = [
+    "ignore::UserWarning:hdmf",
+    "ignore::UserWarning:pynwb",
+]

ephys_mcp-0.1.0/pyproject.toml.orig

@@ -0,0 +1,56 @@
+[project]
+name = "ephys-mcp"
+version = "0.1.0"
+description = "MCP server for read-only analysis of intracortical (spike-level) brain-computer-interface recordings"
+license = "CC0-1.0"
+keywords = ["mcp", "bci", "electrophysiology", "neuroscience", "spikes", "nwb", "dandi"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "License :: CC0 1.0 Universal (CC0 1.0) Public Domain Dedication",
+    "Programming Language :: Python :: 3",
+    "Topic :: Scientific/Engineering",
+]
+readme = "README.md"
+authors = [
+    { name = "happyc0der", email = "Keshav.Rajput4@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "h5py>=3.16.0",
+    "httpx>=0.28.1",
+    "matplotlib>=3.11.2",
+    "mcp[cli]>=2.2.0",
+    "numpy>=2.5.3",
+    "pynwb>=4.2.0",
+    "remfile>=0.1.14",
+    "scipy>=1.18.1",
+]
+
+[project.urls]
+Homepage = "https://github.com/happyc0der/ephys-mcp"
+Repository = "https://github.com/happyc0der/ephys-mcp"
+Issues = "https://github.com/happyc0der/ephys-mcp/issues"
+
+[project.scripts]
+ephys-mcp = "ephys_mcp.server:main"
+
+[build-system]
+requires = ["uv_build>=0.12.10,<0.13.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=9.1.1",
+    "pytest-asyncio>=1.4.0",
+    "ruff>=0.16.8",
+]
+
+[tool.ruff]
+line-length = 120
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+markers = ["network: needs internet access to the DANDI Archive"]
+addopts = "-m 'not network'"
+filterwarnings = ["ignore::UserWarning:hdmf", "ignore::UserWarning:pynwb"]

ephys_mcp-0.1.0/src/ephys_mcp/__init__.py

@@ -0,0 +1,2 @@
+def hello() -> str:
+    return "Hello from ephys-mcp!"

ephys_mcp-0.1.0/src/ephys_mcp/plots.py

@@ -0,0 +1,188 @@
+"""Static figures, saved as PNG. Colour is assigned by job: categorical hues in a
+fixed order for groups, a blue-gray-red diverging map for signed change, ink for
+single series. The categorical set is validated for colour-blind separation."""
+
+from __future__ import annotations
+
+import os
+import re
+import time
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import LinearSegmentedColormap
+
+CATEGORICAL = ["#2a78d6", "#eb6834", "#1baf7a", "#eda100", "#e87ba4", "#008300", "#4a3aa7", "#e34948"]
+DIVERGING = LinearSegmentedColormap.from_list("blue_gray_red", ["#184f95", "#86b6ef", "#f0efec", "#ec835a", "#a82a2a"])
+SURFACE, INK, INK_2, MUTED, GRID, AXIS = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9", "#c3c2b7"
+UNRECORDED = "#d5d3c8"
+MAX_DIRECT_LABELS = 4
+
+plt.rcParams.update(
+    {
+        "figure.facecolor": SURFACE,
+        "axes.facecolor": SURFACE,
+        "savefig.facecolor": SURFACE,
+        "axes.edgecolor": AXIS,
+        "axes.labelcolor": INK_2,
+        "axes.titlecolor": INK,
+        "axes.titlesize": 11,
+        "axes.titleweight": "bold",
+        "axes.titlelocation": "left",
+        "axes.labelsize": 9,
+        "axes.spines.top": False,
+        "axes.spines.right": False,
+        "xtick.color": MUTED,
+        "ytick.color": MUTED,
+        "xtick.labelsize": 8,
+        "ytick.labelsize": 8,
+        "grid.color": GRID,
+        "grid.linewidth": 0.6,
+        "legend.frameon": False,
+        "legend.fontsize": 8,
+        "legend.labelcolor": INK_2,
+        "lines.linewidth": 2.0,
+        "font.family": "sans-serif",
+    }
+)
+
+
+def output_dir() -> Path:
+    d = Path(os.environ.get("EPHYS_MCP_OUTPUT_DIR", Path.home() / ".cache" / "ephys-mcp" / "plots")).expanduser()
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def _save(fig, stem: str) -> Path:
+    safe = re.sub(r"[^A-Za-z0-9_-]+", "-", stem)
+    path = output_dir() / f"{safe}-{int(time.time() * 1000) % 10**9}.png"
+    fig.savefig(path, dpi=150, bbox_inches="tight")
+    plt.close(fig)
+    return path
+
+
+def _event_line(ax):
+    ax.axvline(0, color=INK_2, linewidth=1.0, linestyle=(0, (3, 3)), zorder=1)
+
+
+def _title_and_legend(ax, title: str, legend_cols: int):
+    """Title on top, legend in its own band beneath it, so neither ever covers data."""
+    if legend_cols:
+        legend = ax.legend(
+            loc="lower left",
+            bbox_to_anchor=(0, 1.0),
+            ncols=legend_cols,
+            handlelength=1.4,
+            columnspacing=1.2,
+            borderpad=0,
+        )
+        rows = -(-len(legend.get_texts()) // legend_cols)
+        ax.set_title(title, pad=10 + 15 * rows)
+    else:
+        ax.set_title(title)
+
+
+def plot_psth(
+    stem: str, title: str, event: str, t, groups: dict[str, tuple], unit_z: np.ndarray, rate_label: str
+) -> Path:
+    """groups: label -> (mean, sem, n) of the population rate. unit_z: (n_units, n_bins) baseline z-scores."""
+    heatmap = len(unit_z) > 1  # a one-row heatmap says nothing the curve above does not
+    if heatmap:
+        fig, (top, bottom) = plt.subplots(
+            2, 1, figsize=(7.2, 6.4), sharex=True, gridspec_kw={"height_ratios": [1, 1.25], "hspace": 0.16}
+        )
+    else:
+        fig, top = plt.subplots(figsize=(7.2, 3.6))
+    for (label, (mean, sem, n)), colour in zip(groups.items(), CATEGORICAL):
+        top.fill_between(t, mean - sem, mean + sem, color=colour, alpha=0.16, linewidth=0)
+        top.plot(t, mean, color=colour, label=f"{label} (n={n})", solid_capstyle="round")
+    if 1 < len(groups) <= MAX_DIRECT_LABELS:
+        ends = sorted(groups.items(), key=lambda kv: kv[1][0][-1])
+        gap = 0.075 * (top.get_ylim()[1] - top.get_ylim()[0])
+        y_prev = -np.inf
+        for label, (mean, _, _) in ends:  # nudge labels apart so they never collide
+            y = max(mean[-1], y_prev + gap)
+            top.annotate(
+                label, (t[-1], y), xytext=(5, 0), textcoords="offset points", va="center", fontsize=8, color=INK_2
+            )
+            y_prev = y
+    _event_line(top)
+    top.grid(axis="y")
+    top.set_ylabel(rate_label)
+    _title_and_legend(top, title, min(len(groups), 4) if len(groups) > 1 else 0)
+    top.margins(x=0)
+    if not heatmap:
+        top.set_xlabel(f"Time from {event} (s)")
+        return _save(fig, stem)
+
+    order = np.argsort(np.abs(unit_z).argmax(axis=1))
+    lim = max(2.0, float(np.percentile(np.abs(unit_z), 99)))
+    step = t[1] - t[0]
+    im = bottom.imshow(
+        unit_z[order],
+        aspect="auto",
+        cmap=DIVERGING,
+        vmin=-lim,
+        vmax=lim,
+        extent=(t[0] - step / 2, t[-1] + step / 2, len(order), 0),
+        interpolation="nearest",
+    )
+    _event_line(bottom)
+    bottom.set_ylabel("Units, sorted by time of peak change")
+    bottom.set_xlabel(f"Time from {event} (s)")
+    bar = fig.colorbar(im, ax=[top, bottom], fraction=0.035, pad=0.09, shrink=0.5, anchor=(0, 0.1))
+    bar.outline.set_visible(False)
+    bar.set_label("Change from baseline (SD)", color=INK_2, fontsize=8)
+    return _save(fig, stem)
+
+
+def plot_raster(stem: str, title: str, spikes: list[np.ndarray], t0: float, t1: float, valid: np.ndarray) -> Path:
+    fig, ax = plt.subplots(figsize=(7.2, 4.6))
+    edges = np.concatenate([[t0], np.clip(valid, t0, t1).ravel(), [t1]])
+    for a, b in edges.reshape(-1, 2):  # spans between recorded intervals
+        if b > a:
+            ax.axvspan(a, b, color=UNRECORDED, linewidth=0, zorder=0)
+    ax.eventplot(spikes, colors=INK, linewidths=0.6, linelengths=0.8, lineoffsets=np.arange(len(spikes)) + 0.5)
+    ax.set_xlim(t0, t1)
+    ax.set_ylim(len(spikes), 0)
+    ax.set_xlabel("Time (s)")
+    ax.set_ylabel("Unit")
+    ax.set_title(title)
+    if (np.diff(edges.reshape(-1, 2), axis=1) > 0).any():
+        ax.annotate(
+            "shaded = not recorded",
+            (1, 1.01),
+            xycoords="axes fraction",
+            ha="right",
+            va="bottom",
+            fontsize=8,
+            color=MUTED,
+        )
+    return _save(fig, stem)
+
+
+def plot_decode(stem: str, title: str, t, true: np.ndarray, decoded: np.ndarray, unit: str, r2: np.ndarray) -> Path:
+    dims = true.shape[1]
+    fig, axes = plt.subplots(dims, 1, figsize=(7.2, 1.6 + 1.9 * dims), sharex=True, squeeze=False)
+    names = ["x", "y", "z"] if dims <= 3 else [str(i) for i in range(dims)]
+    gaps = np.flatnonzero(np.diff(t) > 1.5 * np.median(np.diff(t))) + 1  # break lines across unrecorded spans
+    for i, ax in enumerate(axes[:, 0]):
+        for series, colour, label in (
+            (true[:, i], CATEGORICAL[0], "Actual"),
+            (decoded[:, i], CATEGORICAL[1], "Decoded"),
+        ):
+            ax.plot(np.insert(t, gaps, np.nan), np.insert(series, gaps, np.nan), color=colour, label=label)
+        ax.axhline(0, color=AXIS, linewidth=0.8, zorder=0)
+        ax.grid(axis="y")
+        ax.margins(x=0)
+        ax.set_ylabel(f"{names[i]} ({unit})" if unit else names[i])
+        ax.annotate(
+            f"R² {r2[i]:.2f}", (1, 1.03), xycoords="axes fraction", ha="right", va="bottom", fontsize=8, color=INK_2
+        )
+    _title_and_legend(axes[0, 0], title, 2)
+    axes[-1, 0].set_xlabel("Time (s)")
+    return _save(fig, stem)

ephys_mcp-0.1.0/src/ephys_mcp/processing/decode.py

@@ -0,0 +1,121 @@
+"""Linear decoders from binned spike counts to a continuous behavioural signal.
+
+Both decoders are causal: the estimate at bin t uses counts up to bin t only.
+Hyperparameters left as None are chosen on the tail of the training data, so
+the caller's test split is never touched.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+VALIDATION_FRACTION = 0.2
+ALPHAS = (1.0, 10.0, 100.0, 1000.0)
+LEADS = (0, 1, 2, 3)  # bins by which neural activity leads behaviour
+
+
+def r2_score(y: np.ndarray, yhat: np.ndarray) -> np.ndarray:
+    ss_res = ((y - yhat) ** 2).sum(axis=0)
+    ss_tot = ((y - y.mean(axis=0)) ** 2).sum(axis=0)
+    return 1.0 - ss_res / np.where(ss_tot == 0, 1.0, ss_tot)
+
+
+def _select(make, candidates, counts: np.ndarray, y: np.ndarray):
+    """Best candidate by R² on the tail of the training data."""
+    cut = int(len(y) * (1 - VALIDATION_FRACTION))
+    scores = [r2_score(y[cut:], make(c)._fit(counts[:cut], y[:cut]).predict(counts[cut:])).mean() for c in candidates]
+    return candidates[int(np.argmax(scores))]
+
+
+class RidgeDecoder:
+    """Ridge regression on a short history of spike counts."""
+
+    kind = "ridge"
+
+    def __init__(self, n_lags: int = 10, alpha: float | None = None):
+        self.n_lags, self.alpha = n_lags, alpha
+
+    @property
+    def params(self) -> dict:
+        return {"history_bins": self.n_lags, "alpha": self.alpha}
+
+    def _design(self, counts: np.ndarray) -> np.ndarray:
+        lagged = [np.roll(counts, k, axis=0) for k in range(self.n_lags)]
+        X = np.concatenate(lagged, axis=1)
+        X[: self.n_lags] = 0.0
+        return np.concatenate([X, np.ones((X.shape[0], 1))], axis=1)
+
+    def _fit(self, counts: np.ndarray, y: np.ndarray) -> RidgeDecoder:
+        X = self._design(counts)
+        reg = self.alpha * np.eye(X.shape[1])
+        reg[-1, -1] = 0.0
+        self.W = np.linalg.solve(X.T @ X + reg, X.T @ y)
+        return self
+
+    def fit(self, counts: np.ndarray, y: np.ndarray) -> RidgeDecoder:
+        if self.alpha is None:
+            self.alpha = _select(lambda a: RidgeDecoder(self.n_lags, a), ALPHAS, counts, y)
+        return self._fit(counts, y)
+
+    def predict(self, counts: np.ndarray) -> np.ndarray:
+        return self._design(counts) @ self.W
+
+
+class KalmanDecoder:
+    """Kalman filter (Wu et al. 2006): x_t = A x_{t-1} + w, z_t = H x_t + q.
+
+    The state holds the target at the current and previous bin, which gives the
+    smooth second-order dynamics real movements have. `lead_bins` models motor
+    cortex firing ahead of the movement it drives.
+    """
+
+    kind = "kalman"
+
+    def __init__(self, lead_bins: int | None = None):
+        self.lead_bins = lead_bins
+
+    @property
+    def params(self) -> dict:
+        return {"lead_bins": self.lead_bins, "state": "target at t and t-1"}
+
+    def _fit(self, counts: np.ndarray, y: np.ndarray) -> KalmanDecoder:
+        k = self.lead_bins
+        self.d = y.shape[1]
+        if k:
+            counts, y = counts[:-k], y[k:]
+        counts, y = counts[1:], np.hstack([y[1:], y[:-1]])
+        self.z_mean, self.x_mean = counts.mean(axis=0), y.mean(axis=0)
+        Z, X = counts - self.z_mean, y - self.x_mean
+        X0, X1 = X[:-1], X[1:]
+        self.A = np.linalg.lstsq(X0, X1, rcond=None)[0].T
+        self.W = np.cov((X1 - X0 @ self.A.T).T) + 1e-9 * np.eye(X.shape[1])
+        self.H = np.linalg.lstsq(X, Z, rcond=None)[0].T
+        self.Q = np.cov((Z - X @ self.H.T).T) + 1e-6 * np.eye(Z.shape[1])
+        return self
+
+    def fit(self, counts: np.ndarray, y: np.ndarray) -> KalmanDecoder:
+        if self.lead_bins is None:
+            self.lead_bins = _select(KalmanDecoder, LEADS, counts, y)
+        return self._fit(counts, y)
+
+    def predict(self, counts: np.ndarray) -> np.ndarray:
+        Z = counts - self.z_mean
+        n = self.A.shape[0]
+        x, P = np.zeros(n), np.eye(n)
+        HtQinv = self.H.T @ np.linalg.inv(self.Q)
+        HtQinvH = HtQinv @ self.H
+        out = np.empty((Z.shape[0], n))
+        for t in range(Z.shape[0]):
+            x = self.A @ x
+            P = self.A @ P @ self.A.T + self.W
+            P = np.linalg.inv(np.linalg.inv(P) + HtQinvH)  # information form: small inverse only
+            x = x + P @ HtQinv @ (Z[t] - self.H @ x)
+            out[t] = x
+        est = (out + self.x_mean)[:, : self.d]
+        k = self.lead_bins
+        if k:  # counts at bin t describe behaviour at t + k; realign to the input bins
+            est = np.vstack([np.repeat(est[:1], k, axis=0), est[:-k]])
+        return est
+
+
+DECODERS = {"ridge": RidgeDecoder, "kalman": KalmanDecoder}