tcMultiTestNPZ.py

import os
import glob
import sys
import argparse
import numpy as np
import pandas as pd
from natsort import natsorted
import matplotlib.pyplot as plt

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import pyBindMat as pb


def build_rows(cell_ids, results, mouse, date, analysis_type):
    rows = []
    for cid, r in zip(cell_ids, results):
        rows.append({
            'cell_id':          int(cid),
            'mouseName':        mouse,
            'date':             date,
            'analysisType':     analysis_type,
            'meanScore':        float(r['meanScore']),
            'baseScore':        float(r['baseScore']),
            'percentileScore':  float(r['percentileScore']),
            'sigMean':          bool(r['sigMean']),
            'sigBootstrap':     bool(r['sigBootstrap']),
            'fracTrialsFired':  float(r['fracTrialsFired']),
            'meanPkIdx':        int(r['meanPkIdx']),
            'numPeaks':         len(r['allPkIndices']),
            'allPkIndices':     ';'.join(str(x) for x in r['allPkIndices']),
            'allPkScores':      ';'.join(f'{x:.6f}' for x in r['allPkScores']),
            'allPkPvalues':     ';'.join(f'{x:.6f}' for x in r['allPkPvalues']),
        })
    return rows


def _heatmap_ax(ax, traces, pk_indices, time_axis, method, us_time):
    if not traces:
        ax.text(0.5, 0.5, f'No significant\n{method} cells',
                ha='center', va='center', transform=ax.transAxes, fontsize=11)
        ax.set_title(method)
        return

    order = np.argsort(pk_indices)
    mat   = np.array([traces[i] for i in order], dtype=float)
    rmin  = mat.min(axis=1, keepdims=True)
    rmax  = mat.max(axis=1, keepdims=True)
    denom = rmax - rmin
    denom[denom == 0] = 1.0
    mat   = (mat - rmin) / denom

    ncells = mat.shape[0]
    dt = time_axis[1] - time_axis[0]
    im = ax.imshow(mat, aspect='auto', origin='upper', cmap='hot', vmin=0, vmax=1,
                   extent=[time_axis[0] - dt/2, time_axis[-1] + dt/2, ncells, 0])
    ax.axvline(0.0,     color='cyan', lw=1.0, ls='--', label='CS onset')
    ax.axvline(us_time, color='lime', lw=1.0, ls='--', label='US onset')
    ax.legend(loc='upper right', fontsize=7)
    ax.set_xlabel('Time re CS onset (s)')
    ax.set_ylabel(f'Cell (sorted by 1st peak)  n={ncells}')
    ax.set_title(method)
    plt.colorbar(im, ax=ax, label='Norm. ΔF/F', fraction=0.03, pad=0.02)


def plot_heatmaps(r2b_traces, r2b_pk, r2b_time,
                  ti_traces,  ti_pk,  ti_time,
                  r2b_traces_mp, r2b_pk_mp,
                  ti_traces_mp,  ti_pk_mp,
                  us_time, title):
    n_all = max(len(r2b_traces),    len(ti_traces),    1)
    n_mp  = max(len(r2b_traces_mp), len(ti_traces_mp), 1)
    fig, axes = plt.subplots(
        2, 2, figsize=(14, n_all * 0.12 + n_mp * 0.12 + 3.0),
        gridspec_kw={'height_ratios': [n_all, n_mp]})
    fig.suptitle(title)
    _heatmap_ax(axes[0, 0], r2b_traces,    r2b_pk,    r2b_time, 'R2B — all significant', us_time)
    _heatmap_ax(axes[0, 1], ti_traces,     ti_pk,     ti_time,  'TI — all significant',  us_time)
    _heatmap_ax(axes[1, 0], r2b_traces_mp, r2b_pk_mp, r2b_time, 'R2B — 2+ peaks',        us_time)
    _heatmap_ax(axes[1, 1], ti_traces_mp,  ti_pk_mp,  ti_time,  'TI — 2+ peaks',         us_time)
    plt.tight_layout()
    plt.show()


def main():
    parser = argparse.ArgumentParser(description="Multi-peak time cell analysis on NPZ files.")
    parser.add_argument('--path',       type=str,   default='/home/bhalla/homework/2p_analysis/NIDHI_TEST_DAT_MAY2026/')
    parser.add_argument('--out_file',   type=str,   default='multi_peak_results.csv')
    parser.add_argument('--mouse',      type=int,   default=54)
    parser.add_argument('--csOnset',    type=int,   default=240)
    parser.add_argument('--usOnset',    type=int,   default=249)
    parser.add_argument('--circPad',    type=int,   default=60)
    parser.add_argument('--binFrames',  type=int,   default=1)
    parser.add_argument('--numShuffle', type=int,   default=1000)
    parser.add_argument('--r2bThresh',              type=float, default=3.0)
    parser.add_argument('--r2bPercentile',          type=float, default=99.5)
    parser.add_argument('--transientThresh',        type=float, default=2.0)
    parser.add_argument('--tiPercentile',           type=float, default=99.0)
    parser.add_argument('--fracTrialsFiredThresh',  type=float, default=0.25)
    parser.add_argument('--frameDt',                type=float, default=1.0/30.0)
    parser.add_argument('--minPeakSep',             type=float, default=0.4)
    args = parser.parse_args()

    us_time = (args.usOnset - args.csOnset) * args.frameDt

    npz_files = natsorted(glob.glob(os.path.join(args.path, '*.npz')))
    if not npz_files:
        print(f"No .npz files found in {args.path}")
        return

    all_rows  = []
    total_r2b = {0: 0, 1: 0, 2: 0, 3: 0}
    total_ti  = {0: 0, 1: 0, 2: 0, 3: 0}
    r2b_traces,    r2b_pk    = [], []
    ti_traces,     ti_pk     = [], []
    r2b_traces_mp, r2b_pk_mp = [], []   # 2+ peaks only
    ti_traces_mp,  ti_pk_mp  = [], []
    r2b_time, ti_time = None, None   # established from first file

    for fpath in npz_files:
        fname    = os.path.basename(fpath)
        date     = os.path.splitext(fname)[0].split('_')[0]
        data     = np.load(fpath)
        dfbf     = data['dfbf_clean']   # [cell, trial, frame]
        cell_ids = data['cell_ids']
        numFrames = dfbf.shape[2]
        circShuffleFrames = numFrames - (args.csOnset - args.circPad)
        print(f"\n{fname}: {dfbf.shape[0]} cells, {dfbf.shape[1]} trials, {numFrames} frames"
              f"  circShuff={circShuffleFrames}")

        this_r2b_time = np.array([(ff - args.circPad) * args.frameDt
                                  for ff in range(circShuffleFrames)])
        numBins = circShuffleFrames // args.binFrames
        this_ti_time  = np.array([(bb * args.binFrames - args.circPad) * args.frameDt
                                  for bb in range(numBins)])
        if r2b_time is None:
            r2b_time, ti_time = this_r2b_time, this_ti_time

        r2b_res = pb.runR2Banalysis(
            dfbf, args.csOnset, args.usOnset, args.circPad,
            circShuffleFrames, args.binFrames, args.numShuffle,
            args.r2bThresh, args.r2bPercentile,
            frameDt=args.frameDt, minPeakSep=args.minPeakSep)

        ti_res = pb.runTIanalysis(
            dfbf, args.csOnset, args.usOnset, args.circPad,
            circShuffleFrames, args.binFrames, args.numShuffle,
            args.transientThresh, args.tiPercentile,
            args.fracTrialsFiredThresh, args.frameDt,
            minPeakSep=args.minPeakSep)

        r2b_counts = {0: 0, 1: 0, 2: 0, 3: 0}
        ti_counts  = {0: 0, 1: 0, 2: 0, 3: 0}
        for r in r2b_res:
            r2b_counts[min(len(r['allPkIndices']), 3)] += 1
            total_r2b[min(len(r['allPkIndices']), 3)]  += 1
        for r in ti_res:
            ti_counts[min(len(r['allPkIndices']), 3)]  += 1
            total_ti[min(len(r['allPkIndices']), 3)]   += 1

        print(f"  R2B: 0 peaks={r2b_counts[0]}  1={r2b_counts[1]}  2={r2b_counts[2]}  3+={r2b_counts[3]}")
        print(f"  TI:  0 peaks={ti_counts[0]}   1={ti_counts[1]}   2={ti_counts[2]}   3+={ti_counts[3]}")

        all_rows.extend(build_rows(cell_ids, r2b_res, args.mouse, date, 'R2B'))
        all_rows.extend(build_rows(cell_ids, ti_res,  args.mouse, date, 'TI'))

        for r in r2b_res:
            n = len(r['allPkIndices'])
            if n >= 1:
                t = np.array(r['meanTrace'])
                if len(t) == len(r2b_time):
                    r2b_traces.append(t)
                    r2b_pk.append(r['allPkIndices'][0])
                    if n >= 2:
                        r2b_traces_mp.append(t)
                        r2b_pk_mp.append(r['allPkIndices'][0])
        for r in ti_res:
            n = len(r['allPkIndices'])
            if n >= 1:
                t = np.array(r['meanTrace'])
                if len(t) == len(ti_time):
                    ti_traces.append(t)
                    ti_pk.append(r['allPkIndices'][0])
                    if n >= 2:
                        ti_traces_mp.append(t)
                        ti_pk_mp.append(r['allPkIndices'][0])

    print(f"\n=== Totals ===")
    print(f"  R2B: 0 peaks={total_r2b[0]}  1={total_r2b[1]}  2={total_r2b[2]}  3+={total_r2b[3]}")
    print(f"  TI:  0 peaks={total_ti[0]}   1={total_ti[1]}   2={total_ti[2]}   3+={total_ti[3]}")

    df = pd.DataFrame(all_rows)
    df.to_csv(args.out_file, index=False)
    print(f"Saved {len(df)} rows to {args.out_file}")

    # ── 2+ peak table ────────────────────────────────────────────────────────
    multi = df[df['numPeaks'] >= 2].copy()
    if multi.empty:
        print("\nNo cells with 2+ significant peaks.")
    else:
        print(f"\n=== Cells with 2+ significant peaks ===")
        hdr = f"{'source':<20}\t{'cell_id':>7}\t{'method':<5}\t{'pk#':>3}\t{'time(s)':>8}\t{'score':>10}\t{'p-val':>8}"
        print(hdr)
        print("-" * len(hdr.expandtabs(8)))
        for _, row in multi.sort_values(['date', 'cell_id', 'analysisType']).iterrows():
            indices = [int(x)   for x in row['allPkIndices'].split(';')]
            scores  = [float(x) for x in row['allPkScores'].split(';')]
            pvals   = [float(x) for x in row['allPkPvalues'].split(';')]
            bf = args.binFrames if row['analysisType'] == 'TI' else 1
            for pk_num, (idx, score, pval) in enumerate(zip(indices, scores, pvals), 1):
                t = (idx * bf - args.circPad) * args.frameDt
                print(f"{row['date']:<20}\t{row['cell_id']:>7}\t{row['analysisType']:<5}\t"
                      f"{pk_num:>3}\t{t:>8.3f}\t{score:>10.4f}\t{pval:>8.4f}")

    # ── Combined heatmap ─────────────────────────────────────────────────────
    plot_heatmaps(r2b_traces,    r2b_pk,    r2b_time,
                  ti_traces,     ti_pk,     ti_time,
                  r2b_traces_mp, r2b_pk_mp,
                  ti_traces_mp,  ti_pk_mp,
                  us_time, "All sessions")


if __name__ == '__main__':
    main()