#!/usr/bin/env python3 """Example 01 — mEPSC Poisson Models Under Constant and Washout Magnesium. This example demonstrates: 1) Homogeneous Poisson modeling for constant Mg2+ conditions. 2) Piecewise baseline modeling under Mg2+ washout conditions. 3) Model comparison using KS plots, time-rescaling diagnostics, and estimated conditional intensity functions. Data provenance: Uses installer-downloaded nSTAT example data from ``data/mEPSCs``: ``epsc2.txt``, ``washout1.txt``, ``washout2.txt`` Expected outputs: - Figure 1: Constant Mg2+ raster + diagnostics + lambda estimate. - Figure 2: Constant vs decreasing Mg2+ raster overview. - Figure 3: Piecewise model diagnostics and lambda comparison. Paper mapping: Section 2.3.1 (mEPSC analysis); Figs. 3 and 10 (nSTAT paper, 2012). """ from __future__ import annotations import argparse import sys from pathlib import Path import matplotlib import matplotlib.pyplot as plt import numpy as np # --------------------------------------------------------------------------- # Ensure nstat is importable when running from the examples/paper directory. # --------------------------------------------------------------------------- THIS_DIR = Path(__file__).resolve().parent REPO_ROOT = THIS_DIR.parents[1] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) import nstat # noqa: E402 from nstat import ( # noqa: E402 Analysis, ConfigColl, CovColl, nspikeTrain, nstColl, Trial, TrialConfig, ) from nstat.signal import Covariate # noqa: E402 from nstat.data_manager import ensure_example_data # noqa: E402 # ========================================================================= # Helper: load mEPSC spike times from text file # ========================================================================= def _load_mepsc_times_seconds(path: Path) -> np.ndarray: """Load spike times from mEPSC text file, returning times in seconds.""" data = np.loadtxt(path, skiprows=1) # Column 2 is spike time in milliseconds at 1000 Hz times_ms = data[:, 1] if data.ndim == 2 else data return times_ms / 1000.0 def _matlab_colon(start: float, step: float, stop: float) -> np.ndarray: """Replicate MATLAB ``start:step:stop`` exactly. ``np.arange`` accumulates floating-point error over many steps and can produce off-by-one length mismatches. This function computes the element count first (like MATLAB's colon operator), then multiplies by integer indices — giving bit-exact parity. """ n = int(np.floor((stop - start) / step)) + 1 return start + np.arange(n) * step # ========================================================================= # Helper: export figure # ========================================================================= def _maybe_export(fig, export_dir: Path | None, name: str, dpi: int = 250): """Save figure to disk if export_dir is set.""" saved = [] if export_dir is not None: export_dir.mkdir(parents=True, exist_ok=True) png_path = export_dir / f"{name}.png" fig.savefig(png_path, dpi=dpi, facecolor="w", edgecolor="none") saved.append(png_path) print(f" Saved {png_path}") return saved # ========================================================================= # Main example function # ========================================================================= def run_example01(*, export_figures: bool = False, export_dir: Path | None = None, visible: bool = True): """Run Example 01: mEPSC Poisson models.""" if not visible: matplotlib.use("Agg") data_dir = ensure_example_data(download=True) mepsc_dir = data_dir / "mEPSCs" figure_files: list[Path] = [] sampleRate = 1000 # Hz # ================================================================== # Part 1: Constant magnesium concentration — Homogeneous Poisson model # ================================================================== print("=== Part 1: Constant Mg2+ — Homogeneous Poisson ===") epsc2 = _load_mepsc_times_seconds(mepsc_dir / "epsc2.txt") # Create spike train and time vector nstConst = nspikeTrain(epsc2) timeConst = _matlab_colon(0, 1.0 / sampleRate, nstConst.maxTime) # Create baseline covariate baseline = Covariate( timeConst, np.ones((len(timeConst), 1)), "Baseline", "time", "s", "", dataLabels=["\\mu"], ) covarColl = CovColl([baseline]) spikeCollConst = nstColl(nstConst) trialConst = Trial(spikeCollConst, covarColl) # Configure: single constant-rate model tcConst = TrialConfig([("Baseline", "\\mu")], sampleRate, []) tcConst.setName("Constant Baseline") configConst = ConfigColl([tcConst]) # Fit GLM resultConst = Analysis.RunAnalysisForAllNeurons(trialConst, configConst, 0) resultConst.lambda_signal.setDataLabels(["\\lambda_{const}"]) print(f" Spikes: {len(epsc2)}") print(f" AIC: {resultConst.AIC}") print(f" BIC: {resultConst.BIC}") # --- Figure 1: Constant Mg2+ diagnostics --- # MATLAB layout: 2x2 with (1,1) raster, (1,2) InvGausTrans, # (2,1) KSPlot, (2,2) lambda plot. fig1, axes1 = plt.subplots(2, 2, figsize=(14, 9)) ax = axes1[0, 0] spikeCollConst.plot(handle=ax) ax.set_title(r"Neural Raster with constant Mg$^{2+}$ Concentration", fontweight="bold", fontsize=12, fontfamily="Arial") ax.set_xlabel("time [s]", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_ylabel("mEPSCs", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_yticks([0, 1]) ax = axes1[0, 1] resultConst.plotInvGausTrans(handle=ax) ax = axes1[1, 0] resultConst.KSPlot(handle=ax) ax = axes1[1, 1] lam = resultConst.lambda_signal ax.plot(lam.time, lam.data[:, 0], "b", linewidth=2) ax.set_xlabel("time [s]", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_ylabel(f"{lam.name} [{lam.yunits}]" if lam.yunits else lam.name, fontsize=12, fontweight="bold", fontfamily="Arial") ax.legend([r"$\lambda_{const}$"], loc="upper right") fig1.tight_layout() figure_files.extend(_maybe_export(fig1, export_dir, "fig01_constant_mg_summary")) # ================================================================== # Part 2: Varying magnesium concentration — Piecewise baseline model # ================================================================== print("\n=== Part 2: Decreasing Mg2+ — Piecewise Baseline ===") washout1 = _load_mepsc_times_seconds(mepsc_dir / "washout1.txt") washout2 = _load_mepsc_times_seconds(mepsc_dir / "washout2.txt") spikeTimes1 = 260.0 + washout1 spikeTimes2 = np.sort(washout2) + 745.0 nstWashout = nspikeTrain(np.concatenate([spikeTimes1, spikeTimes2])) timeWashout = _matlab_colon(260.0, 1.0 / sampleRate, nstWashout.maxTime) # --- Figure 2: Constant vs Decreasing Mg2+ rasters --- fig2, axes2 = plt.subplots(2, 1, figsize=(14, 9)) ax = axes2[0] nstConst.plot(handle=ax) ax.set_yticks([0, 1]) ax.set_ylabel("mEPSCs", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_title(r"Neural Raster with constant Mg$^{2+}$ Concentration", fontweight="bold", fontsize=12, fontfamily="Arial") ax.xaxis.label.set(fontsize=12, fontweight="bold", fontfamily="Arial") ax = axes2[1] nstWashout.plot(handle=ax) ax.set_yticks([0, 1]) ax.set_ylabel("mEPSCs", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_title(r"Neural Raster with decreasing Mg$^{2+}$ Concentration", fontweight="bold", fontsize=12, fontfamily="Arial") ax.xaxis.label.set(fontsize=12, fontweight="bold", fontfamily="Arial") fig2.tight_layout() figure_files.extend(_maybe_export(fig2, export_dir, "fig02_washout_raster_overview")) # ================================================================== # Part 3: Piecewise baseline model and model comparison # ================================================================== print("\n=== Part 3: Piecewise Baseline Model Comparison ===") # Build piecewise indicator covariates # Matlab: timeInd1 = find(time < 495, 1, 'last') → last 1-based index < 495 # Equivalent Python: first 0-based index >= 495 (searchsorted side='left'), # so rate1[:idx] covers [260, 494.999] and rate2[idx:] starts at 495. timeInd1 = np.searchsorted(timeWashout, 495.0, side="left") timeInd2 = np.searchsorted(timeWashout, 765.0, side="left") N = len(timeWashout) constantRate = np.ones((N, 1)) rate1 = np.zeros((N, 1)) rate2 = np.zeros((N, 1)) rate3 = np.zeros((N, 1)) rate1[:timeInd1] = 1.0 rate2[timeInd1:timeInd2] = 1.0 rate3[timeInd2:] = 1.0 baselineWashout = Covariate( timeWashout, np.column_stack([constantRate, rate1, rate2, rate3]), "Baseline", "time", "s", "", dataLabels=["\\mu", "\\mu_{1}", "\\mu_{2}", "\\mu_{3}"], ) spikeCollWashout = nstColl(nstWashout) trialWashout = Trial(spikeCollWashout, CovColl([baselineWashout])) # Configure: (1) constant baseline, (2) piecewise baseline tc1 = TrialConfig([("Baseline", "\\mu")], sampleRate, []) tc1.setName("Constant Baseline") tc2 = TrialConfig([("Baseline", "\\mu_{1}", "\\mu_{2}", "\\mu_{3}")], sampleRate, []) tc2.setName("Diff Baseline") configWashout = ConfigColl([tc1, tc2]) resultWashout = Analysis.RunAnalysisForAllNeurons(trialWashout, configWashout, 0) resultWashout.lambda_signal.setDataLabels(["\\lambda_{const}", "\\lambda_{const-epoch}"]) print(f" AIC: {resultWashout.AIC}") print(f" BIC: {resultWashout.BIC}") # --- Figure 3: Piecewise model diagnostics --- # MATLAB layout: 2x2 with (1,1) raster + red epoch lines, # (1,2) InvGausTrans, (2,1) KSPlot, (2,2) dual-lambda plot. fig3, axes3 = plt.subplots(2, 2, figsize=(14, 9)) ax = axes3[0, 0] spikeCollWashout.plot(handle=ax) ax.set_title(r"Neural Raster with decreasing Mg$^{2+}$ Concentration", fontweight="bold", fontsize=12, fontfamily="Arial") ax.set_xlabel("time [s]", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_yticklabels([]) # Red epoch-transition markers (MATLAB: plot([495;495],[0,1],'r','LineWidth',4)) ax.plot([495, 495], [0, 1], "r", linewidth=4) ax.plot([765, 765], [0, 1], "r", linewidth=4) ax = axes3[0, 1] resultWashout.plotInvGausTrans(handle=ax) ax = axes3[1, 0] resultWashout.KSPlot(handle=ax) ax = axes3[1, 1] lam = resultWashout.lambda_signal ax.plot(lam.time, lam.getSubSignal(1).data[:, 0], "b", linewidth=2) ax.plot(lam.time, lam.getSubSignal(2).data[:, 0], "g", linewidth=2) ax.set_ylim(0, 5) ax.set_xlabel("time [s]", fontsize=12, fontweight="bold", fontfamily="Arial") ax.set_ylabel(f"{lam.name} [{lam.yunits}]" if lam.yunits else lam.name, fontsize=12, fontweight="bold", fontfamily="Arial") ax.legend([r"$\lambda_{const}$", r"$\lambda_{const-epoch}$"], loc="upper right") fig3.tight_layout() figure_files.extend(_maybe_export(fig3, export_dir, "fig03_piecewise_baseline_comparison")) if visible: plt.show() print(f"\nExample 01 complete. Generated {len(figure_files)} figure(s).") return figure_files # ========================================================================= # CLI entry point # ========================================================================= if __name__ == "__main__": parser = argparse.ArgumentParser(description="Example 01: mEPSC Poisson Models") parser.add_argument("--repo-root", type=Path, default=REPO_ROOT, help="Repository root (default: auto-detected).") parser.add_argument("--export-figures", action="store_true", help="Export figures to disk.") parser.add_argument("--export-dir", type=Path, default=None, help="Directory for exported figures.") parser.add_argument("--no-display", action="store_true", help="Run without displaying figures (headless).") args = parser.parse_args() export_dir = args.export_dir if args.export_figures and export_dir is None: export_dir = THIS_DIR / "figures" / "example01" run_example01( export_figures=args.export_figures, export_dir=export_dir if args.export_figures else None, visible=not args.no_display, )