plot flow unit.py

"""
plot_flow_unit.py
从 NWB 文件加载数据，用 flow_library 里的函数分析，
复现 "Tuning Curve + PSTH Heatmap" 图（hierarchy: speed → dirStd → dir）。

用法：
    python plot_flow_unit.py
"""

import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from pathlib import Path
from pynwb import NWBHDF5IO

# ── 把 library 所在目录加进 path（按实际路径修改）──────────────────────
sys.path.insert(0, str(Path("/Users/jiahuan/Desktop/test_data")))
from SC_SF_TF_flow_library import (
    extract_trials_spike_times,
    calculate_conditions_firing_rate,
    calculate_conditions_psth,
)

# ══════════════════════════════════════════════════════════════════
#  配置
# ══════════════════════════════════════════════════════════════════
NWB_PATH   = Path("/Users/jiahuan/Desktop/test_data/SC data/NWB/NP-SC05-d1-2x192.nwb")
TRIAL_NAME = "flow_trials"          # NWB 里的 TimeIntervals 名称
UNIT_ID    = 165                    # 要画的 unit
PRE_STIM   = 0.5                    # 刺激前窗口 (s)
POST_STIM  = 0.0                    # 刺激后窗口 (s)
STIM_DUR   = 1.0                    # 刺激时长 (s)
OUT_PATH   = Path("unit165_flow.png")

# ══════════════════════════════════════════════════════════════════
#  1. 从 NWB 加载 spike times 和 trials
# ══════════════════════════════════════════════════════════════════
def load_from_nwb(nwb_path, trial_name):
    with NWBHDF5IO(str(nwb_path), "r") as io:
        nwb = io.read()

        # ── spike times ──────────────────────────────────────────
        units_df = nwb.units.to_dataframe()
        spike_times = {}
        for uid, row in units_df.iterrows():
            spike_times[uid] = np.array(row["spike_times"])

        # ── trials ───────────────────────────────────────────────
        trials_raw = nwb.intervals[trial_name].to_dataframe()

    print(f"Units: {len(spike_times)},  Trials: {len(trials_raw)}")
    print(f"Trial columns: {list(trials_raw.columns)}")
    return spike_times, trials_raw

# ══════════════════════════════════════════════════════════════════
#  2. 构建 condition_params 和 trials（library 格式）
# ══════════════════════════════════════════════════════════════════
def build_conditions(trials_raw):
    """
    给每个 (speed, dirStd, direction) 组合分配 condition_num。
    返回：
      trials_df   — 含 stim_on, stim_off, condition_num
      cond_params — 含 condition_num, speed, dirStd, dir
    """
    # 唯一条件组合
    keys = ["speed", "dirStd", "direction"]
    combos = (trials_raw[keys]
              .drop_duplicates()
              .sort_values(keys)
              .reset_index(drop=True))
    combos["condition_num"] = combos.index + 1  # 1-based

    # merge 回 trials
    trials_df = trials_raw.merge(combos, on=keys, how="left")
    trials_df = trials_df.rename(columns={
        "start_time": "stim_on",
        "stop_time":  "stim_off",
    })

    cond_params = combos.rename(columns={"direction": "dir"})

    n_conds = len(cond_params)
    speeds   = sorted(cond_params["speed"].unique())
    dirstds  = sorted(cond_params["dirStd"].unique())
    dirs     = sorted(cond_params["dir"].unique())
    print(f"Conditions: {n_conds}  "
          f"speeds={speeds}  dirStds={dirstds}  dirs={dirs}")
    return trials_df, cond_params

# ══════════════════════════════════════════════════════════════════
#  3. 计算 FR 和 PSTH（library 函数）
# ══════════════════════════════════════════════════════════════════
def compute_responses(spike_times, trials_df, unit_id):
    # 只保留单个 unit，加快速度
    st = {unit_id: spike_times[unit_id]}

    _, conditions_spike_times = extract_trials_spike_times(
        st, trials_df, pre_stim=PRE_STIM, post_stim=POST_STIM)

    conditions_FR = calculate_conditions_firing_rate(
        conditions_spike_times, stim_duration=STIM_DUR)

    conditions_PSTH = calculate_conditions_psth(
        conditions_spike_times,
        bin_size=0.02, sigma=0.05,
        pre_stim=PRE_STIM, post_stim=POST_STIM,
        stim_duration=STIM_DUR)

    return conditions_FR, conditions_PSTH

# ══════════════════════════════════════════════════════════════════
#  4. 画图（复现 Image2 的 hierarchy: speed → dirStd → dir）
# ══════════════════════════════════════════════════════════════════
def plot_tuning_and_psth(unit_id, conditions_FR, conditions_PSTH,
                          cond_params, out_path):

    speeds  = sorted(cond_params["speed"].unique())
    dirstds = sorted(cond_params["dirStd"].unique())
    dirs    = sorted(cond_params["dir"].unique())
    n_speed  = len(speeds)
    n_dirstd = len(dirstds)
    n_dir    = len(dirs)
    n_cols   = n_speed * n_dirstd          # 每个 (speed,dirStd) 一列

    # ── colormap：每个 (speed,dirStd) 一个颜色 ──────────────────
    cmap_base = plt.cm.get_cmap("tab20", n_cols)
    col_colors = {(s, ds): cmap_base(i)
                  for i, (s, ds) in enumerate(
                      [(s, ds) for s in speeds for ds in dirstds])}

    fig = plt.figure(figsize=(max(24, n_cols * 1.2), 10))
    gs  = fig.add_gridspec(2, 1, height_ratios=[1, 1.4], hspace=0.4)

    # ── 上图：tuning curve ───────────────────────────────────────
    ax_top = fig.add_subplot(gs[0])

    # 构建 x 轴：每个 (speed,dirStd) 占 n_dir 个点，组间加间隔
    gap = 2
    x_all, x_ticks_dir, x_ticks_group = [], [], []
    x_labels_dir, x_labels_group = [], []
    cursor = 0
    group_x_map = {}   # (speed,dirStd) → list of x positions
    for s in speeds:
        for ds in dirstds:
            xs = list(range(cursor, cursor + n_dir))
            group_x_map[(s, ds)] = xs
            x_all.extend(xs)
            x_ticks_dir.extend(xs)
            x_labels_dir.extend([str(int(d)) for d in dirs])
            x_ticks_group.append(cursor + n_dir / 2 - 0.5)
            x_labels_group.append(f"sp={s}\nds={ds}")
            cursor += n_dir + gap

    # 画每个 (speed,dirStd) 的 tuning curve
    baseline_fr = None
    for s in speeds:
        for ds in dirstds:
            color = col_colors[(s, ds)]
            fr_vals, fr_sems = [], []
            for d in dirs:
                row = cond_params[
                    (cond_params["speed"] == s) &
                    (cond_params["dirStd"] == ds) &
                    (cond_params["dir"] == d)]
                if row.empty:
                    fr_vals.append(np.nan); fr_sems.append(0)
                    continue
                cid = int(row["condition_num"].values[0])
                stats = conditions_FR.get(cid, {}).get(unit_id, {})
                fr_vals.append(stats.get("mean", np.nan))
                fr_sems.append(stats.get("sem", 0))

            xs  = group_x_map[(s, ds)]
            frs = np.array(fr_vals)
            sms = np.array(fr_sems)
            ax_top.plot(xs, frs, "o-", color=color, lw=1.2, ms=3,
                        label=f"speed={s}, dirStd={ds}")
            ax_top.fill_between(xs, frs - sms, frs + sms,
                                 color=color, alpha=0.15)

    # baseline (mean across all conditions)
    all_means = [conditions_FR[c][unit_id]["mean"]
                 for c in conditions_FR if unit_id in conditions_FR[c]]
    baseline_fr = np.nanmean(all_means)
    ax_top.axhline(baseline_fr, color="gray", ls="--", lw=1,
                   label=f"Baseline FR ({baseline_fr:.1f} Hz)")

    ax_top.set_xticks(x_ticks_dir, x_labels_dir, fontsize=5, rotation=45)
    ax_top.set_ylabel("Firing rate (Hz)")
    ax_top.set_title(f"Unit {unit_id} Tuning Curve and PSTH Heatmaps\n"
                     f"Hierarchy: speed → dirStd → dir")
    ax_top.legend(loc="upper right", fontsize=5,
                  ncol=max(1, n_cols // 6), framealpha=0.5)

    # group labels (speed + dirStd) on secondary x
    ax2 = ax_top.twiny()
    ax2.set_xlim(ax_top.get_xlim())
    ax2.set_xticks(x_ticks_group)
    ax2.set_xticklabels(x_labels_group, fontsize=6)
    ax2.tick_params(axis="x", length=0)

    # ── 下图：PSTH heatmaps（每列一个 (speed,dirStd)）────────────
    gs_bot = gs[1].subgridspec(1, n_cols, wspace=0.05)
    axes_bot = [fig.add_subplot(gs_bot[0, c]) for c in range(n_cols)]

    # 统一颜色范围
    all_psth = []
    for s in speeds:
        for ds in dirstds:
            for d in dirs:
                row = cond_params[
                    (cond_params["speed"] == s) &
                    (cond_params["dirStd"] == ds) &
                    (cond_params["dir"] == d)]
                if row.empty: continue
                cid = int(row["condition_num"].values[0])
                p = conditions_PSTH.get(cid, {}).get(unit_id, {})
                if "psth" in p:
                    all_psth.append(p["psth"])
    vmin = np.nanmin(np.concatenate(all_psth)) if all_psth else 0
    vmax = np.nanmax(np.concatenate(all_psth)) if all_psth else 1

    time_axis = None
    col_idx = 0
    for s in speeds:
        for ds in dirstds:
            ax = axes_bot[col_idx]
            heat = []
            for d in dirs:
                row = cond_params[
                    (cond_params["speed"] == s) &
                    (cond_params["dirStd"] == ds) &
                    (cond_params["dir"] == d)]
                if row.empty:
                    heat.append(np.zeros(50))
                    continue
                cid = int(row["condition_num"].values[0])
                p = conditions_PSTH.get(cid, {}).get(unit_id, {})
                if "psth" in p:
                    heat.append(p["psth"])
                    time_axis = p["time"]
                else:
                    heat.append(np.zeros(len(time_axis) if time_axis is not None else 50))

            heat = np.array(heat)  # (n_dir, n_time)
            if time_axis is not None:
                im = ax.imshow(heat, aspect="auto", origin="lower",
                               cmap="viridis", vmin=vmin, vmax=vmax,
                               extent=[time_axis[0], time_axis[-1], 0, n_dir])
                ax.axvline(0, color="white", lw=0.8, ls="--")
                ax.axvline(STIM_DUR, color="white", lw=0.8, ls="--")
            ax.set_title(f"speed={s}\ndirStd={ds}", fontsize=6, pad=2)
            ax.set_yticks(np.arange(n_dir) + 0.5)
            if col_idx == 0:
                ax.set_yticklabels([str(int(d)) for d in dirs], fontsize=5)
                ax.set_ylabel("Direction (deg)", fontsize=7)
            else:
                ax.set_yticklabels([])
            ax.set_xlabel("Time (s)", fontsize=6)
            ax.tick_params(axis="x", labelsize=5)
            col_idx += 1

    # colorbar
    if time_axis is not None:
        cbar_ax = fig.add_axes([0.92, 0.05, 0.008, 0.25])
        fig.colorbar(im, cax=cbar_ax, label="Firing rate (Hz)")

    plt.savefig(out_path, dpi=150, bbox_inches="tight")
    print(f"Saved: {out_path}")
    plt.close()


# ══════════════════════════════════════════════════════════════════
#  Main
# ══════════════════════════════════════════════════════════════════
if __name__ == "__main__":
    print("Loading NWB...")
    spike_times, trials_raw = load_from_nwb(NWB_PATH, TRIAL_NAME)

    print("Building conditions...")
    trials_df, cond_params = build_conditions(trials_raw)

    print(f"Computing responses for unit {UNIT_ID}...")
    conditions_FR, conditions_PSTH = compute_responses(
        spike_times, trials_df, UNIT_ID)

    print("Plotting...")
    plot_tuning_and_psth(UNIT_ID, conditions_FR, conditions_PSTH,
                          cond_params, OUT_PATH)