#!/usr/bin/env python3 """Example 05 — Stimulus Decoding With PPAF and PPHF. This example demonstrates neural decoding using point-process adaptive filters (PPAF) and point-process hybrid filters (PPHF) from the nSTAT toolbox. The example has three parts: Part A — Univariate Sinusoidal Stimulus (Figures 1–2): 1. Define 20-cell population with logistic (binomial) tuning to a 1-D sinusoidal stimulus. 2. Simulate spike observations from the binomial CIF. 3. Decode the stimulus using ``PPDecodeFilterLinear`` (PPAF). Part B — 4-State Arm Reach with PPAF (Figures 3–4): 4. Simulate reaching trajectories (position + velocity, 4-D state) using minimum-jerk dynamics (cosine acceleration toward target). 5. Encode with 20-cell velocity-tuned population (binomial CIF). 6. Decode with PPAF (free) and PPAF + Goal; compare across 20 simulations. Part C — Hybrid Filter (Figures 5–6): 7. Load fixture trajectory with 6-D state (pos + vel + accel) and 2 discrete movement modes (not-moving / moving) from ``paperHybridFilterExample.mat``. 8. Simulate 40-cell population with velocity-tuned binomial CIF. 9. Decode joint discrete + continuous state via ``PPHybridFilterLinear`` (both goal-directed and free), averaged over 20 simulations. Paper mapping: Sections 2.3.6–2.3.7 (decoding); Figs. 8, 9, 14 plus hybrid extension. Expected outputs: - Figure 1: CIF tuning curves and simulated spike raster. - Figure 2: Decoded stimulus vs true (with 95% confidence band). - Figure 3: Reach trajectory, position/velocity traces, neural raster, CIF. - Figure 4: PPAF decoding overlaid trajectories (free=green, goal=blue). - Figure 5: Hybrid fixture setup (reach path, traces, raster, discrete state). - Figure 6: Hybrid decoding summary (state est, probabilities, decoded path). """ from __future__ import annotations import argparse import json import sys from pathlib import Path import matplotlib.pyplot as plt import numpy as np 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)) from nstat import DecodingAlgorithms # noqa: E402 # ────────────────────────────────────────────────────────────────────────────── # Helper: simulate binomial spikes from linear-logistic CIF # ────────────────────────────────────────────────────────────────────────────── def _simulate_binomial_spikes_from_lambda(lambdaRate, delta, rng): """Simulate spikes from precomputed lambda rates via thinning. Parameters ---------- lambdaRate : (C, T) array — firing rates [spikes/sec] per cell per time delta : float — bin width in seconds rng : numpy Generator Returns ------- dN : (C, T) array — binary spike indicators """ prob = lambdaRate * delta # convert rate to probability per bin prob = np.clip(prob, 0.0, 1.0) return (rng.random(prob.shape) < prob).astype(float) def _simulate_binomial_spikes(xState, muCoeffs, beta, rng, delta=0.001): """Simulate binomial spikes from state, encoding coefficients, and tuning. Computes CIF via logistic link from state and beta, then draws spikes. Parameters ---------- xState : (ns, T) array — kinematic state (position + velocity) muCoeffs : (C,) array — baseline log-rate per cell beta : (ns, C) array — tuning coefficients per state dimension per cell rng : numpy Generator delta : float — bin width in seconds (default 0.001) Returns ------- dN : (C, T) array — binary spike indicators """ T = xState.shape[1] C = len(muCoeffs) # Build design matrix: (T, 1+ns) = [1, x1, x2, ..., xns] dataMat = np.column_stack([np.ones(T), xState.T]) # Build coefficient matrix: (C, 1+ns) = [mu, beta_1, ..., beta_ns] coeffs = np.column_stack([muCoeffs, beta.T]) lambdaRate = _logistic_cif(dataMat, coeffs, delta) return _simulate_binomial_spikes_from_lambda(lambdaRate, delta, rng) def _logistic_cif(dataMat, coeffs, delta): """Compute binomial CIF rates matching MATLAB's logistic link. Parameters ---------- dataMat : (T, p) — design matrix [1, covariates] coeffs : (C, p) — per-cell coefficients [mu, betas] delta : float — bin width Returns ------- lambdaRate : (C, T) — firing rates in spikes/sec """ C = coeffs.shape[0] T = dataMat.shape[0] lambdaRate = np.zeros((C, T)) for c in range(C): eta = dataMat @ coeffs[c, :] expEta = np.exp(np.clip(eta, -20.0, 20.0)) p = expEta / (1.0 + expEta) lambdaRate[c, :] = p / delta return lambdaRate # ────────────────────────────────────────────────────────────────────────────── # Part A — Univariate sinusoidal stimulus # ────────────────────────────────────────────────────────────────────────────── def _run_part_a(seed=0, n_cells=20): """Encode/decode a 1-D sinusoidal stimulus with 20-cell binomial CIF. Matches MATLAB: rng(0,'twister'), delta=0.001, f=2Hz, b0~N(log(10*delta),1), b1~N(0,1), logistic CIF, PPDecodeFilterLinear with A=1, Q=std(diff(stim)). """ rng = np.random.default_rng(seed) delta = 0.001 tmax = 1.0 time = np.arange(0.0, tmax + delta, delta) T = len(time) f = 2.0 # Encoding model — matches MATLAB exactly b1 = rng.standard_normal(n_cells) b0 = np.log(10.0 * delta) + rng.standard_normal(n_cells) stimSignal = np.sin(2.0 * np.pi * f * time) # Compute CIF and simulate spikes per cell dN = np.zeros((n_cells, T)) lambdaAll = np.zeros((n_cells, T)) for c in range(n_cells): eta = b1[c] * stimSignal + b0[c] expEta = np.exp(np.clip(eta, -20.0, 20.0)) p = expEta / (1.0 + expEta) lambdaAll[c, :] = p / delta dN[c, :] = (rng.random(T) < p).astype(float) # State-space model: x(t+1) = A*x(t) + w A = np.array([[1.0]]) Q_val = float(np.std(np.diff(stimSignal))) Q = np.array([[Q_val]]) x0 = np.array([0.0]) Pi0 = 0.5 * np.eye(1) # Decode beta = b1.reshape(1, -1) # (1, C) x_p, W_p, x_u, W_u, _, _, _, _ = DecodingAlgorithms.PPDecodeFilterLinear( A, Q, dN, b0, beta, "binomial", delta, None, None, x0, Pi0 ) x_decoded = x_u[0, :] sigma = np.sqrt(np.maximum(W_u[0, 0, :], 0.0)) z_val = 1.96 ci_low = np.minimum(x_decoded - z_val * sigma, x_decoded + z_val * sigma) ci_high = np.maximum(x_decoded - z_val * sigma, x_decoded + z_val * sigma) rmse = float(np.sqrt(np.mean((x_decoded - stimSignal) ** 2))) return { "time": time, "stimSignal": stimSignal, "x_decoded": x_decoded, "ci_low": ci_low, "ci_high": ci_high, "dN": dN, "lambdaAll": lambdaAll, "b0": b0, "b1": b1, "rmse": rmse, "n_cells": n_cells, "delta": delta, } # ────────────────────────────────────────────────────────────────────────────── # Part B — 4-state arm reach with PPAF # ────────────────────────────────────────────────────────────────────────────── def _simulate_reach_minjerk(delta, T_total): """Simulate a 2-D minimum-jerk reach from x0 to xT. Uses MATLAB's cosine-acceleration dynamics: xState(:,k) = A * xState(:,k-1) + (delta/2)*(pi/T)^2 * cos(pi*t/T) * [0; 0; xT(1)-x0(1); xT(2)-x0(2)] Returns time, xState (4×T), A (4×4). """ x0 = np.array([0.0, 0.0, 0.0, 0.0]) xT_target = np.array([-0.35, 0.2, 0.0, 0.0]) time = np.arange(0.0, T_total + delta, delta) T = len(time) A = np.array([ [1, 0, delta, 0], [0, 1, 0, delta], [0, 0, 1, 0], [0, 0, 0, 1], ], dtype=float) xState = np.zeros((4, T), dtype=float) xState[:, 0] = x0 accel_dir = np.array([0.0, 0.0, xT_target[0] - x0[0], xT_target[1] - x0[1]]) for k in range(1, T): accel = (delta / 2.0) * (np.pi / T_total) ** 2 * np.cos(np.pi * time[k] / T_total) xState[:, k] = A @ xState[:, k - 1] + accel * accel_dir return time, xState, A def _run_part_b(seed=0, n_cells=20, n_sims=20): """Compare PPAF free vs goal-directed decoding for arm reach. Matches MATLAB: single trajectory, 20 re-randomized encoding simulations. """ rng = np.random.default_rng(seed) delta = 0.001 # 1 ms bins (matches MATLAB) T_total = 2.0 # 2-second reach # ── Generate minimum-jerk reach trajectory ── time, xState, A = _simulate_reach_minjerk(delta, T_total) T = xState.shape[1] ns = 4 # Target = final state yT = xState[:, -1].copy() # Q from trajectory variance (MATLAB: diag(var(diff(xState,[],2),[],2))*100) Q = np.diag(np.var(np.diff(xState, axis=1), axis=1)) * 100.0 # Initial/target covariances (MATLAB: very tight) r, p = 1e-6, 1e-6 pi0 = np.diag([r, r, p, p]) piT = np.diag([r, r, p, p]) # ── Run 20 repeated simulations ── # Same trajectory, re-randomized encoding + spikes each time all_runs_goal = [] all_runs_free = [] example_run = None for sim_idx in range(n_sims): # MATLAB: bCoeffs = 10*(rand(numCells,2)-0.5); Uniform[-5,5] bCoeffs = 10.0 * (rng.random((n_cells, 2)) - 0.5) # MATLAB: muCoeffs = log(10*delta) + randn(numCells,1) muCoeffs = np.log(10.0 * delta) + rng.standard_normal(n_cells) # beta: 4×C with zeros for position, bCoeffs for velocity beta = np.zeros((ns, n_cells), dtype=float) beta[2, :] = bCoeffs[:, 0] # vx tuning beta[3, :] = bCoeffs[:, 1] # vy tuning # Simulate spikes dN = _simulate_binomial_spikes(xState, muCoeffs, beta, rng) # Initial state x0 = np.array([0.0, 0.0, 0.0, 0.0]) # --- Goal-directed decode --- _, _, x_u_goal, _, _, _, _, _ = DecodingAlgorithms.PPDecodeFilterLinear( A, Q, dN, muCoeffs, beta, "binomial", delta, None, None, x0, pi0, yT, piT, 0 ) # --- Free decode (no goal) --- _, _, x_u_free, _, _, _, _, _ = DecodingAlgorithms.PPDecodeFilterLinear( A, Q, dN, muCoeffs, beta, "binomial", delta, None, None, x0, ) all_runs_goal.append(x_u_goal) all_runs_free.append(x_u_free) if sim_idx == 0: example_run = { "time": time, "xState": xState, "dN": dN, "muCoeffs": muCoeffs, "bCoeffs": bCoeffs, "beta": beta, } return { "all_runs_goal": all_runs_goal, "all_runs_free": all_runs_free, "example": example_run, "n_cells": n_cells, "n_sims": n_sims, "xState": xState, "time": time, "delta": delta, } # ────────────────────────────────────────────────────────────────────────────── # Part C — Hybrid filter # ────────────────────────────────────────────────────────────────────────────── def _load_hybrid_fixture(): """Load the MATLAB hybrid filter fixture (paperHybridFilterExample.mat). Returns a dict with: time, delta, X (6×T), mstate (T,), A (list of 2), Q (list of 2), p_ij (2×2), Px0 (list of 2), ind. """ import scipy.io as sio # Search for fixture in multiple locations candidates = [ REPO_ROOT / "nstat" / "data" / "paperHybridFilterExample.mat", REPO_ROOT / "helpfiles" / "paperHybridFilterExample.mat", ] mat_path = None for p in candidates: if p.exists() and p.stat().st_size > 200: # skip LFS pointers mat_path = p break if mat_path is None: raise FileNotFoundError( "Cannot find paperHybridFilterExample.mat fixture. " "Ensure it is in nstat/data/ or helpfiles/." ) f = sio.loadmat(str(mat_path)) time = f["time"].ravel().astype(float) delta = float(f["delta"].ravel()[0]) X = f["X"].astype(float) # (6, T) mstate = f["mstate"].ravel().astype(int) # (T,), values 1 or 2 p_ij = f["p_ij"].astype(float) # (2, 2) # Cell arrays → Python lists A_cell = f["A"] Q_cell = f["Q"] Px0_cell = f["Px0"] ind_cell = f["ind"] A = [A_cell[0, i].astype(float) for i in range(A_cell.shape[1])] Q = [Q_cell[0, i].astype(float) for i in range(Q_cell.shape[1])] Px0 = [Px0_cell[0, i].astype(float) for i in range(Px0_cell.shape[1])] # ind: convert from MATLAB 1-indexed to Python 0-indexed ind = [ind_cell[0, i].ravel().astype(int) - 1 for i in range(ind_cell.shape[1])] return { "time": time, "delta": delta, "X": X, "mstate": mstate, "A": A, "Q": Q, "p_ij": p_ij, "Px0": Px0, "ind": ind, } def _run_part_c(seed=0, n_cells=40, n_sims=20): """PPHybridFilterLinear: joint discrete/continuous state decoding. Loads fixture trajectory, runs 20 simulations with re-randomized encoding, comparing goal-directed vs free hybrid decoding. """ rng = np.random.default_rng(seed) # ── Load fixture ── fix = _load_hybrid_fixture() time = fix["time"] delta = fix["delta"] X = fix["X"] # (6, T) mstate = fix["mstate"] # (T,) p_ij = fix["p_ij"] A_models = fix["A"] # [A_hold (2×2), A_reach (6×6)] Q_models_orig = fix["Q"] # [Q_hold (2×2), Q_reach (6×6)] Px0 = fix["Px0"] # [Px0_hold (2×2), Px0_reach (6×6)] ind = fix["ind"] # [[0,1], [0,1,2,3,4,5]] T = X.shape[1] # ── Recompute Q from trajectory variance (matching MATLAB) ── nonMovingInd = np.where((X[4, :] == 0) & (X[5, :] == 0))[0] movingInd = np.setdiff1d(np.arange(T), nonMovingInd) Q_reach = np.diag(np.var(np.diff(X[:, movingInd], axis=1), axis=1)) Q_reach[:4, :4] = 0.0 # Zero out pos/vel noise; only accel has noise varNV = np.diag(np.var(np.diff(X[:, nonMovingInd], axis=1), axis=1)) Q_hold = varNV[:2, :2] Q_models = [Q_hold, Q_reach] # State dimensions dim_hold = A_models[0].shape[0] # 2 dim_reach = A_models[1].shape[0] # 6 # ── Run 20 repeated simulations ── X_estAll = np.zeros((dim_reach, T, n_sims), dtype=float) X_estNTAll = np.zeros((dim_reach, T, n_sims), dtype=float) S_estAll = np.zeros((n_sims, T), dtype=int) S_estNTAll = np.zeros((n_sims, T), dtype=int) MU_estAll = np.zeros((2, T, n_sims), dtype=float) MU_estNTAll = np.zeros((2, T, n_sims), dtype=float) example_dN = None for n in range(n_sims): # MATLAB: muCoeffs = log(10*delta) + randn(numCells,1) muCoeffs = np.log(10.0 * delta) + rng.standard_normal(n_cells) # MATLAB: coeffs = [muCoeffs, zeros(C,2), 10*(rand(C,2)-0.5), zeros(C,2)] # = [mu, 0, 0, b_vx, b_vy, 0, 0] — tuned to velocities (states 3-4) bCoeffs_vx = 10.0 * (rng.random(n_cells) - 0.5) bCoeffs_vy = 10.0 * (rng.random(n_cells) - 0.5) # Full beta: 6×C matrix beta_full = np.zeros((6, n_cells), dtype=float) beta_full[2, :] = bCoeffs_vx # vx tuning beta_full[3, :] = bCoeffs_vy # vy tuning # Simulate spikes from full state trajectory dN = _simulate_binomial_spikes(X, muCoeffs, beta_full, rng) if n == 0: example_dN = dN.copy() # ── Initial conditions per mode ── x0_list = [X[ind[0], 0], X[ind[1], 0]] Pi0_list = Px0 # ── Target conditions per mode ── yT_list = [X[ind[0], -1], X[ind[1], -1]] piT_list = [1e-9 * np.eye(dim_hold), 1e-9 * np.eye(dim_reach)] # beta per mode: hold uses 2-dim subset, reach uses full 6-dim beta_hold = beta_full[ind[0], :] # (2, C) — will be zeros beta_reach = beta_full[ind[1], :] # (6, C) — has vx/vy tuning mu0 = np.array([0.5, 0.5]) # --- Goal-directed hybrid decode --- S_est, X_est, _, MU_est, _, _, _ = DecodingAlgorithms.PPHybridFilterLinear( A_models, Q_models, p_ij, mu0, dN, [muCoeffs, muCoeffs], [beta_hold, beta_reach], "binomial", delta, None, None, x0_list, Pi0_list, yT_list, piT_list, ) # --- Free hybrid decode (no target) --- S_estNT, X_estNT, _, MU_estNT, _, _, _ = DecodingAlgorithms.PPHybridFilterLinear( A_models, Q_models, p_ij, mu0, dN, [muCoeffs, muCoeffs], [beta_hold, beta_reach], "binomial", delta, None, None, x0_list, Pi0_list, ) X_estAll[:, :, n] = X_est X_estNTAll[:, :, n] = X_estNT S_estAll[n, :] = S_est S_estNTAll[n, :] = S_estNT MU_estAll[:, :, n] = MU_est MU_estNTAll[:, :, n] = MU_estNT print(f" Hybrid sim {n + 1}/{n_sims} done") return { "time": time, "X": X, "mstate": mstate, "dN": example_dN, "X_estAll": X_estAll, "X_estNTAll": X_estNTAll, "S_estAll": S_estAll, "S_estNTAll": S_estNTAll, "MU_estAll": MU_estAll, "MU_estNTAll": MU_estNTAll, "n_cells": n_cells, "n_sims": n_sims, } # ────────────────────────────────────────────────────────────────────────────── # Plotting # ────────────────────────────────────────────────────────────────────────────── def _plot_part_a(result): """Figure 1: stimulus + CIF + raster (3×1). Figure 2: decoded vs true.""" time = result["time"] stimSignal = result["stimSignal"] dN = result["dN"] n_cells = result["n_cells"] # ── Figure 1: stimulus, CIF, spike raster (3 panels, matching MATLAB) ── fig1, axes1 = plt.subplots(3, 1, figsize=(14, 9), sharex=True) # (3,1,1): Driving stimulus axes1[0].plot(time, stimSignal, "k", linewidth=1.5) axes1[0].set_ylabel("Stimulus") axes1[0].set_title("Driving Stimulus", fontweight="bold", fontsize=14, fontfamily="Arial") axes1[0].tick_params(labelbottom=False) # (3,1,2): CIFs overlaid in black lambdaAll = result["lambdaAll"] for c in range(n_cells): axes1[1].plot(time, lambdaAll[c, :], "k", linewidth=0.5) axes1[1].set_ylabel("Firing Rate [spikes/sec]") axes1[1].set_title("Conditional Intensity Functions", fontweight="bold", fontsize=14, fontfamily="Arial") axes1[1].tick_params(labelbottom=False) # (3,1,3): Spike raster for c in range(n_cells): spike_times = time[dN[c, :] > 0] axes1[2].plot(spike_times, np.full_like(spike_times, c + 1), "|", color="k", markersize=4) axes1[2].set_ylabel("Cell Number") axes1[2].set_xlabel("time [s]") axes1[2].set_ylim(0.5, n_cells + 0.5) axes1[2].set_yticks(np.arange(0, n_cells + 1, 10)) axes1[2].set_title("Point Process Sample Paths", fontweight="bold", fontsize=14, fontfamily="Arial") fig1.tight_layout() # ── Figure 2: Decoding results (MATLAB: black=decoded, blue=actual) ── fig2, ax2 = plt.subplots(1, 1, figsize=(14, 9)) ax2.fill_between( time, result["ci_low"], result["ci_high"], color="0.75", alpha=0.4, label="95% CI" ) ax2.plot(time, result["x_decoded"], "k-", linewidth=4, label="Decoded") ax2.plot(time, stimSignal, "b-", linewidth=4, label="Actual") ax2.set_xlabel("time [s]") ax2.set_ylabel("Stimulus") # MATLAB: ylabel('Stimulus','Interpreter','none') ax2.set_title(f"Decoded Stimulus $\\pm$ 95% CIs with {result['n_cells']} cells", fontweight="bold", fontsize=18, fontfamily="Arial") ax2.legend(["Decoded", "Actual"], loc="upper right") fig2.tight_layout() return fig1, fig2 def _plot_part_b(result): """Figure 3: Reach setup (4×2 layout). Figure 4: Overlaid decoded trajectories.""" ex = result["example"] time = ex["time"] xState = ex["xState"] dN = ex["dN"] delta = result["delta"] n_cells = result["n_cells"] # ── Figure 3: Reach trajectory and population setup (4×2 layout) ── fig3 = plt.figure(figsize=(14, 9)) # Top-left [1,3]: 2D reach path (in cm) ax_path = fig3.add_subplot(4, 2, (1, 3)) ax_path.plot(100 * xState[0, :], 100 * xState[1, :], "k", linewidth=2) ax_path.plot(100 * xState[0, 0], 100 * xState[1, 0], "bo", markersize=14, markerfacecolor="none", markeredgewidth=2) ax_path.plot(100 * xState[0, -1], 100 * xState[1, -1], "ro", markersize=14, markerfacecolor="none", markeredgewidth=2) ax_path.legend(["Path", "Start", "Finish"], loc="upper right") ax_path.set_xlabel("X Position [cm]") ax_path.set_ylabel("Y Position [cm]") ax_path.set_title("Reach Path", fontweight="bold", fontsize=14) # Middle-left [5]: position vs time (in cm) ax_pos = fig3.add_subplot(4, 2, 5) h1, = ax_pos.plot(time, 100 * xState[0, :], "k", linewidth=2) h2, = ax_pos.plot(time, 100 * xState[1, :], "k-.", linewidth=2) ax_pos.legend([h1, h2], ["x", "y"], loc="upper right") ax_pos.set_xlabel("time [s]") ax_pos.set_ylabel("Position [cm]") # Lower-left [7]: velocity vs time (in cm/s) ax_vel = fig3.add_subplot(4, 2, 7) h1, = ax_vel.plot(time, 100 * xState[2, :], "k", linewidth=2) h2, = ax_vel.plot(time, 100 * xState[3, :], "k-.", linewidth=2) ax_vel.legend([h1, h2], ["$v_x$", "$v_y$"], loc="upper right") ax_vel.set_xlabel("time [s]") ax_vel.set_ylabel("Velocity [cm/s]") # Top-right [2,4]: neural raster ax_raster = fig3.add_subplot(4, 2, (2, 4)) for c in range(n_cells): spike_t = time[dN[c, :] > 0] ax_raster.plot(spike_t, np.full_like(spike_t, c + 1), "|", color="k", markersize=4) ax_raster.set_ylabel("Cell Number") ax_raster.set_xticks([]) ax_raster.set_xticklabels([]) ax_raster.set_title("Neural Raster", fontweight="bold", fontsize=14) # Bottom-right [6,8]: CIF curves ax_cif = fig3.add_subplot(4, 2, (6, 8)) muCoeffs = ex["muCoeffs"] beta = ex["beta"] for c in range(n_cells): eta = muCoeffs[c] + beta[:, c] @ xState exp_eta = np.exp(np.clip(eta, -20, 20)) lam = (exp_eta / (1.0 + exp_eta)) / delta ax_cif.plot(time, lam, "k", linewidth=0.5) ax_cif.set_title("Neural Conditional Intensity Functions", fontweight="bold", fontsize=14) ax_cif.set_xlabel("time [s]") ax_cif.set_ylabel("Firing Rate [spikes/sec]") fig3.tight_layout() # ── Figure 4: Overlaid decoded trajectories (4×2 layout, 20 runs) ── fig4 = plt.figure(figsize=(14, 9)) # Top [1:4]: 2D estimated vs actual reach paths ax_2d = fig4.add_subplot(4, 2, (1, 4)) ax_2d.plot(100 * xState[0, :], 100 * xState[1, :], "k", linewidth=3) ax_2d.set_title("Estimated vs. Actual Reach Paths", fontweight="bold", fontsize=12) for sim_idx in range(result["n_sims"]): x_u_goal = result["all_runs_goal"][sim_idx] x_u_free = result["all_runs_free"][sim_idx] ax_2d.plot(100 * x_u_goal[0, :], 100 * x_u_goal[1, :], "b", linewidth=0.5) ax_2d.plot(100 * x_u_free[0, :], 100 * x_u_free[1, :], "g", linewidth=0.5) ax_2d.set_xlabel("x [cm]") ax_2d.set_ylabel("y [cm]") # Bottom panels: per-state traces state_labels = ["x(t) [cm]", "y(t) [cm]", "$v_x$(t) [cm/s]", "$v_y$(t) [cm/s]"] subplot_indices = [5, 6, 7, 8] scale = 100.0 # meters → cm for d, (sp_idx, ylabel) in enumerate(zip(subplot_indices, state_labels)): ax = fig4.add_subplot(4, 2, sp_idx) ax.plot(time, scale * xState[d, :], "k", linewidth=3) for sim_idx in range(result["n_sims"]): x_u_goal = result["all_runs_goal"][sim_idx] x_u_free = result["all_runs_free"][sim_idx] hB, = ax.plot(time, scale * x_u_goal[d, :], "b", linewidth=0.5) hC, = ax.plot(time, scale * x_u_free[d, :], "g", linewidth=0.5) ax.set_ylabel(ylabel) if d >= 2: ax.set_xlabel("time [s]") else: ax.set_xticklabels([]) # Add legend on y(t) panel (subplot 6), matching MATLAB if d == 1: hA, = ax.plot([], [], "k", linewidth=3) ax.legend([hA, hB, hC], ["Actual", "PPAF+Goal", "PPAF"], loc="lower right", fontsize=8) fig4.tight_layout() return fig3, fig4 def _plot_part_c(result): """Figure 5: Hybrid setup (4×2 layout). Figure 6: Hybrid decoding summary (4×3).""" time = result["time"] X = result["X"] # (6, T) mstate = result["mstate"] dN = result["dN"] n_cells = result["n_cells"] # ── Figure 5: Setup — reach path, traces, raster, discrete state (4×2) ── fig5 = plt.figure(figsize=(14, 9)) # Top-left [1,3]: 2D reach path ax_path = fig5.add_subplot(4, 2, (1, 3)) ax_path.plot(100 * X[0, :], 100 * X[1, :], "k", linewidth=2) ax_path.plot(100 * X[0, 0], 100 * X[1, 0], "bo", markersize=16, markerfacecolor="none", markeredgewidth=2) ax_path.plot(100 * X[0, -1], 100 * X[1, -1], "ro", markersize=16, markerfacecolor="none", markeredgewidth=2) ax_path.set_xlabel("X [cm]") ax_path.set_ylabel("Y [cm]") ax_path.set_title("Reach Path", fontweight="bold", fontsize=14) # Middle-left [5]: position vs time ax_pos = fig5.add_subplot(4, 2, 5) h1, = ax_pos.plot(time, 100 * X[0, :], "k", linewidth=2) h2, = ax_pos.plot(time, 100 * X[1, :], "k-.", linewidth=2) ax_pos.legend([h1, h2], ["x", "y"], loc="upper right") ax_pos.set_xlabel("time [s]") ax_pos.set_ylabel("Position [cm]") # Lower-left [7]: velocity vs time ax_vel = fig5.add_subplot(4, 2, 7) h1, = ax_vel.plot(time, 100 * X[2, :], "k", linewidth=2) h2, = ax_vel.plot(time, 100 * X[3, :], "k-.", linewidth=2) ax_vel.legend([h1, h2], ["$v_x$", "$v_y$"], loc="upper right") ax_vel.set_xlabel("time [s]") ax_vel.set_ylabel("Velocity [cm/s]") # Top-right [2,4]: neural raster (show ALL cells, matching MATLAB) ax_raster = fig5.add_subplot(4, 2, (2, 4)) for c in range(dN.shape[0]): spike_t = time[dN[c, :] > 0] ax_raster.plot(spike_t, np.full_like(spike_t, c + 1), "|", color="k", markersize=4) ax_raster.set_ylabel("Cell Number") ax_raster.set_yticklabels([]) ax_raster.set_xticks([]) ax_raster.set_xticklabels([]) ax_raster.set_title("Neural Raster", fontweight="bold", fontsize=14) # Bottom-right [6,8]: discrete movement state ax_state = fig5.add_subplot(4, 2, (6, 8)) ax_state.plot(time, mstate, "k", linewidth=2) ax_state.set_ylim(0, 3) ax_state.set_yticks([1, 2]) ax_state.set_yticklabels(["N", "M"]) ax_state.set_xlabel("time [s]") ax_state.set_ylabel("state") ax_state.set_title("Discrete Movement State", fontweight="bold", fontsize=14) fig5.tight_layout() # ── Figure 6: Hybrid decoding results (4×3 layout, averaged over 20 sims) ── fig6 = plt.figure(figsize=(14, 9)) # Mean across simulations mS_est = np.mean(result["S_estAll"], axis=0) mS_estNT = np.mean(result["S_estNTAll"], axis=0) mMU_est = np.mean(result["MU_estAll"][1, :, :], axis=1) # P(M|data) for goal mMU_estNT = np.mean(result["MU_estNTAll"][1, :, :], axis=1) # P(M|data) for free mX_est = np.mean(100 * result["X_estAll"], axis=2) mX_estNT = np.mean(100 * result["X_estNTAll"], axis=2) # Left column: state estimation + probability # [1,4]: Estimated vs actual state ax_s = fig6.add_subplot(4, 3, (1, 4)) ax_s.plot(time, mstate, "k", linewidth=3) ax_s.plot(time, mS_est, "b", linewidth=3) ax_s.plot(time, mS_estNT, "g", linewidth=3) ax_s.set_yticks([1, 2.1]) ax_s.set_yticklabels(["N", "M"]) ax_s.set_xticklabels([]) ax_s.set_ylabel("state") ax_s.set_title("Estimated vs. Actual State", fontweight="bold", fontsize=12) # [7,10]: P(s(t)=M | data) ax_prob = fig6.add_subplot(4, 3, (7, 10)) ax_prob.plot(time, mMU_est, "b", linewidth=3) ax_prob.plot(time, mMU_estNT, "g", linewidth=3) ax_prob.set_xlim(time[0], time[-1]) ax_prob.set_ylim(0, 1.1) ax_prob.set_xlabel("time [s]") ax_prob.set_ylabel("P(s(t)=M | data)") ax_prob.set_title("Probability of State", fontweight="bold", fontsize=12) # Right top [2,3,5,6]: 2D estimated vs actual reach path ax_2d = fig6.add_subplot(4, 3, (2, 6)) ax_2d.plot(100 * X[0, :], 100 * X[1, :], "k", linewidth=1) ax_2d.plot(mX_est[0, :], mX_est[1, :], "b", linewidth=3) ax_2d.plot(mX_estNT[0, :], mX_estNT[1, :], "g", linewidth=3) ax_2d.plot(100 * X[0, 0], 100 * X[1, 0], "bo", markersize=14, markerfacecolor="none", markeredgewidth=2) ax_2d.plot(100 * X[0, -1], 100 * X[1, -1], "ro", markersize=14, markerfacecolor="none", markeredgewidth=2) ax_2d.set_xlabel("x [cm]") ax_2d.set_ylabel("y [cm]") ax_2d.set_title("Estimated vs. Actual Reach Path", fontweight="bold", fontsize=12) # Bottom panels: per-state traces # [8]: x(t) ax_x = fig6.add_subplot(4, 3, 8) ax_x.plot(time, 100 * X[0, :], "k", linewidth=3) ax_x.plot(time, mX_est[0, :], "b", linewidth=3) ax_x.plot(time, mX_estNT[0, :], "g", linewidth=3) ax_x.set_ylabel("x(t) [cm]") ax_x.set_xticklabels([]) ax_x.set_title("X Position", fontweight="bold", fontsize=12) # [9]: y(t) with legend ax_y = fig6.add_subplot(4, 3, 9) h1, = ax_y.plot(time, 100 * X[1, :], "k", linewidth=3) h2, = ax_y.plot(time, mX_est[1, :], "b", linewidth=3) h3, = ax_y.plot(time, mX_estNT[1, :], "g", linewidth=3) ax_y.legend([h1, h2, h3], ["Actual", "PPAF+Goal", "PPAF"], loc="lower right", fontsize=8) ax_y.set_ylabel("y(t) [cm]") ax_y.set_xticklabels([]) ax_y.set_title("Y Position", fontweight="bold", fontsize=12) # [11]: vx(t) ax_vx = fig6.add_subplot(4, 3, 11) ax_vx.plot(time, 100 * X[2, :], "k", linewidth=3) ax_vx.plot(time, mX_est[2, :], "b", linewidth=3) ax_vx.plot(time, mX_estNT[2, :], "g", linewidth=3) ax_vx.set_ylabel("$v_x$(t) [cm/s]") ax_vx.set_xlabel("time [s]") ax_vx.set_title("X Velocity", fontweight="bold", fontsize=12) # [12]: vy(t) ax_vy = fig6.add_subplot(4, 3, 12) ax_vy.plot(time, 100 * X[3, :], "k", linewidth=3) ax_vy.plot(time, mX_est[3, :], "b", linewidth=3) ax_vy.plot(time, mX_estNT[3, :], "g", linewidth=3) ax_vy.set_ylabel("$v_y$(t) [cm/s]") ax_vy.set_xlabel("time [s]") ax_vy.set_title("Y Velocity", fontweight="bold", fontsize=12) fig6.tight_layout() return fig5, fig6 # ────────────────────────────────────────────────────────────────────────────── # Main entry point # ────────────────────────────────────────────────────────────────────────────── def run_example05(*, export_figures=False, export_dir=None, show=False): """Run Example 05: PPAF and PPHF decoding. Mirrors MATLAB ``example05_decoding_ppaf_pphf.m``: Part A — Univariate stimulus decoding (Figs 1–2): 1. 20-cell sinusoidal-tuned population, binomial CIF. 2. PPDecodeFilterLinear decoding with 95% CIs. Part B — Arm-reach PPAF: 4. Simulate 4-state minimum-jerk reaching movement. 5. Encode with 20-cell velocity-tuned population. 6. Decode with PPAF (free) and PPAF+Goal; 20 overlaid simulations. Part C — Hybrid filter: 7. Load 6-state fixture trajectory with 2 discrete modes. 8. Simulate 40-cell population with velocity tuning. 9. Decode joint discrete/continuous state via PPHybridFilterLinear (both goal-directed and free), averaged over 20 simulations. """ print("=" * 70) print("Example 05: Stimulus Decoding with PPAF and PPHF") print("=" * 70) # --- Part A --- print("\n--- Part A: Univariate Sinusoidal Stimulus ---") result_a = _run_part_a() print(f" {result_a['n_cells']} cells, decode RMSE = {result_a['rmse']:.4f}") # --- Part B --- print("\n--- Part B: Arm Reach PPAF (20 simulations) ---") result_b = _run_part_b() print(f" {result_b['n_sims']} simulations, {result_b['n_cells']} cells") # --- Part C: Hybrid filter --- print("\n--- Part C: Hybrid Filter (20 simulations) ---") result_c = _run_part_c() print(f" {result_c['n_cells']} cells, {result_c['n_sims']} simulations") # Summary summary = { "experiment5": { "num_cells": float(result_a["n_cells"]), "decode_rmse": result_a["rmse"], }, "experiment5b": { "num_cells": float(result_b["n_cells"]), "n_sims": float(result_b["n_sims"]), }, "experiment6": { "num_cells": float(result_c["n_cells"]), "n_sims": float(result_c["n_sims"]), }, } print("\n" + json.dumps(summary, indent=2)) # --- Figures --- fig1, fig2 = _plot_part_a(result_a) fig3, fig4 = _plot_part_b(result_b) fig5, fig6 = _plot_part_c(result_c) figures = [fig1, fig2, fig3, fig4, fig5, fig6] if export_figures: if export_dir is None: export_dir = THIS_DIR / "figures" / "example05" export_dir = Path(export_dir) export_dir.mkdir(parents=True, exist_ok=True) fig_names = [ "fig01_univariate_setup", "fig02_univariate_decoding", "fig03_reach_and_population_setup", "fig04_ppaf_goal_vs_free", "fig05_hybrid_setup", "fig06_hybrid_decoding_summary", ] for i, fig in enumerate(figures): path = export_dir / f"{fig_names[i]}.png" fig.savefig(path, dpi=250, facecolor="w", edgecolor="none") print(f" Saved: {path}") if show: plt.show() else: plt.close("all") return summary if __name__ == "__main__": parser = argparse.ArgumentParser( description="Example 05: Stimulus Decoding With PPAF and PPHF" ) parser.add_argument("--repo-root", type=Path, default=REPO_ROOT) parser.add_argument("--export-figures", action="store_true") parser.add_argument("--export-dir", type=Path, default=None) parser.add_argument("--output-json", type=Path, default=None) parser.add_argument("--show", action="store_true", help="Display figures interactively") args = parser.parse_args() result = run_example05( export_figures=args.export_figures, export_dir=args.export_dir, show=args.show, ) if args.output_json: args.output_json.write_text(json.dumps(result, indent=2), encoding="utf-8")