diff --git a/README.md b/README.md index 48b16ab..4b87c04 100644 --- a/README.md +++ b/README.md @@ -65,8 +65,13 @@ The core functionality of our workflow is separated into three files... There are also example scripts that use the main function and demonstrate example usage of the utils scripts. Example calls are provided at the top of each docstring for the example files. +1) example_make_correlation_matrix.py +2) example_make_raster.py +3) example_run_permutation.py +4) example_sequence_matching.py + ### What's in the __'\data' Folder__? -If you look in the '\data' folder, you will see two sub-folders labelled 'test_data' and 'full_data'. The first of these sub-folders holds short, sample versions of our data that we can use for quick testing. +After running the download_data.sh script in "scripts/", you will have two sub-folders labelled 'test_data' and 'full_data'. test_data holds short, sample versions of our data that we can use for quick testing, but make no biological sense. full_data includes subfolders for two prairie voles, each with two types of social interaction recordings. One Partner introduction with a novel opposite-sex vole and one Same-sex (SS) introduction. The snakemake will use all these files to make basic raster plots and correlation plots of co-activity. They are also randomly selected in the different "example usage" portion of each example_..._.py file. Below are more details on each of the types of files in the folders. There are five primary types of data contained here... diff --git a/Snakefile b/Snakefile index ce5ce1f..109430d 100644 --- a/Snakefile +++ b/Snakefile @@ -37,7 +37,7 @@ def get_swr_csv_name(dataset): # handle special case for sleepyvole if dataset == "7742/SSIntro": - return f"{animal_id}_{session}_cropped_SWRs_ca2.csv" + return f"{animal_id}_{session}_sleepyvole_SWRs_ca2.csv" # all other cases return f"{animal_id}_{session}_SWRs_ca2.csv" @@ -154,4 +154,4 @@ rule generate_raster_plot: window=params.window, ripple_index=params.ripple_index, save_path=output.raster - ) \ No newline at end of file + ) diff --git a/scripts/download_data.sh b/scripts/download_data.sh index ad5c079..685bbf1 100755 --- a/scripts/download_data.sh +++ b/scripts/download_data.sh @@ -6,7 +6,7 @@ set -euo pipefail OUT_DIR="data" # Google Drive folder URL (replace or add more as needed) -FOLDER_URL="https://drive.google.com/drive/folders/1RcxsdghOMMvaF-mQlVCOtkmXXj2TR7fg?usp=drive_link" +FOLDER_URL="https://drive.google.com/drive/folders/1MCQce7FXHNKEQg97zs5YrXFSZbxB_iDU?usp=drive_link" # Create output directory if needed if [ -d "${OUT_DIR}" ]; then diff --git a/src/analysis_utils.py b/src/analysis_utils.py index f932b8b..644ee32 100644 --- a/src/analysis_utils.py +++ b/src/analysis_utils.py @@ -1,9 +1,9 @@ +import os +import sys import numpy as np import pandas as pd import matplotlib.pyplot as plt from tqdm import tqdm -import os -import sys def make_correlation_dictionary( @@ -12,59 +12,44 @@ def make_correlation_dictionary( normalize: bool = True, ) -> dict: """ - Generates a dictionary that indicates how many - times each neuron cluster fired when any other - cluster had also fired. The output is a nested - dictionary where the first key indicates the - cluster of interest, and the second key indicates - whether or not another given neuron fired. So, - for example... + Generate a nested dictionary describing co-firing frequencies + between neuron clusters. + + corr_dict[a][b] returns how many times cluster b fired in SWRs + where cluster a also fired. - corr_dict[7][16] + Parameters + ---------- + df : pd.DataFrame + DataFrame where the ID column contains a list of cluster IDs + firing in each SWR event. - ...would tell you how many times cluster 16 fired - out of all times that cluster 7 fired. + id_column_name : str + Name of the column that stores cluster ID lists. - Parameters: - ----------- - df :: pd.DataFrame - A dataframe that contains information about - which neuron clusters fired for each detected - SWR. This function assumes that the 'ID' column - has a list for every row of the dataframe. - id_column_name :: str - Name of the column that cluster ID is stored in. - normalize :: bool - Controls whether recorded counts are normalized. - If 'True', each second key (i.e., 16) is divided - by the total number of counts for the first key - (i.e., 7). - - Returns: - -------- - corr_dict :: dict - A nested dictionary showing how many times each - cluster fired with another neuron cluster. - """ + normalize : bool + If True, divide counts by the total SWR count for that baseline + cluster. - # Using 'np.sum()' appends all lists together, and 'flatten()' makes it 1D + Returns + ------- + dict + Nested dictionary mapping cluster → cluster → count/probability. + """ all_ids = np.array(np.sum(df[id_column_name])).flatten() unique_ids = np.unique(all_ids) - # Initialized with 0 in case a cluster never fires with the reference one - corr_dict = {key: {key: 0 for key in unique_ids} for key in unique_ids} + corr_dict = { + key: {key: 0 for key in unique_ids} for key in unique_ids + } for baseline_id in tqdm(unique_ids, desc="Building correlation dict."): for compared_id in unique_ids: for row in df[id_column_name]: - - # Only true if both clusters fires during the same SWR if (baseline_id in row) and (compared_id in row): corr_dict[baseline_id][compared_id] += 1 if normalize: - - # The original dictionary value is modified, so we have to store it baseline_value = corr_dict[baseline_id][baseline_id] for compared_id in unique_ids: try: @@ -80,46 +65,34 @@ def visualize_correlation_dictionary( save_directory: str = None, ): """ - Plots a 'correlation matrix' plot for a given - correlation dictionary. The x-axis shows the baseline - neuron cluster id, and the y-axis shows the compared cluster - id. The color is controlled by the value found for that - baseline-referenced pair in the dictionary. - - Parameters: - ----------- - corr_matrix :: dict - A nested dictionary containing information about the - co-firing of two neuron clusters. - save_directory :: str - The directory (path + filename) to save the correlation - matrix image to. If no argument is provided, the image - will not be saved and 'plt.show()' will be called - instead. - """ + Visualize a correlation dictionary as a matrix heatmap. - # Initializes grid with zeros in case no information is found for an id + Parameters + ---------- + corr_matrix : dict + Nested dictionary of cluster → cluster → value. + + save_directory : str, optional + Path (directory + filename) to save figure. If None, figure + is shown instead. + """ shape = (len(corr_matrix), len(corr_matrix)) grid = np.zeros(shape) keys = list(corr_matrix.keys()) - # Populates matrix with values found in the dictionary for i in range(len(grid)): for j in range(len(grid)): grid[j][i] = corr_matrix[keys[i]][keys[j]] - # Plots correlation matrix and adds a colorbar plt.rcParams["figure.figsize"] = (9, 7) plt.imshow(grid, origin="lower", interpolation="none", cmap="inferno") plt.colorbar(label="Probability of Co-Activity") - # Adds labels to plot plt.title("Correlation Matrix for Neuron Clusters") plt.xlabel("Cluster ID (baseline)") plt.ylabel("Cluster ID (compared)") - # Replacing default xtick labels with cluster ids plt.xticks(range(len(keys)), labels=keys, rotation=90, fontsize=7) plt.yticks(range(len(keys)), labels=keys, fontsize=7) @@ -131,26 +104,100 @@ def visualize_correlation_dictionary( plt.show() -def match_up( - df: pd.DataFrame, - swr_dir: str = None, - swr_df: pd.DataFrame = None, - only_keep_good: bool = True, - progress: bool = True, -) -> pd.DataFrame: +def make_spike_df( + time_dir: str, + cluster_dir: str, + label_dir: str, +): + """ + Load spike times, cluster IDs, and KS labels into a DataFrame. + + Parameters + ---------- + time_dir :: str + Path to spike_times.npy file. + cluster_dir :: str + Path to spike_clusters.npy file. + label_dir :: str + Path to cluster_KSLabel.tsv file. + + Returns + ------- + spike_df :: pd.DataFrame + DataFrame with columns: Time, Cluster ID, KSLabel. + filtered for "good" clusters """ - Assigns each SWR its list of spike times and cluster IDs. + try: + spike_times = np.load(time_dir).flatten() / 30000 # convert to seconds + except FileNotFoundError: + print(f"File not found: {time_dir}") + sys.exit(1) + + try: + spike_clusters = np.load(cluster_dir).flatten() + except FileNotFoundError: + print(f"File not found: {cluster_dir}") + sys.exit(1) + + try: + cluster_labels = pd.read_csv(label_dir, sep="\t") + except FileNotFoundError: + print(f"File not found: {label_dir}") + sys.exit(1) + + # Standardize column names + cluster_labels = cluster_labels.rename( + columns={"cluster_id": "Cluster ID"} + ) + + # Merge to add KSLabel to every spike event + spike_df = pd.DataFrame( + {"Time": spike_times, "Cluster ID": spike_clusters} + ) + spike_df = spike_df.merge(cluster_labels, on="Cluster ID", how="left") + + # Filter for good clusters only + spike_df = spike_df[spike_df["KSLabel"] == "good"] - You can pass either: - - swr_dir: path to SWR CSV - - swr_df: already-loaded SWR DataFrame (faster for permutations) + return spike_df + + +def match_up(df, + swr_dir=None, + swr_df=None, + only_keep_good=True, + progress=True): + """ + Assign spike times and cluster IDs to each SWR interval. + + One of: + • swr_dir — path to SWR CSV + • swr_df — already-loaded DataFrame + + Parameters + ---------- + df :: pd.DataFrame + DataFrame of spike times and cluster IDs. + swr_dir :: str + Path to SWR CSV file. + swr_df :: pd.DataFrame + DataFrame of SWR events. + only_keep_good :: bool + If True, only consider spikes from clusters labeled "good". + progress :: bool + If True, show progress bar. + + Returns + ------- + swr_df :: pd.DataFrame + DataFrame of SWR events with assigned spike times and + cluster IDs. """ if swr_df is None and swr_dir is None: raise ValueError("Provide either swr_dir or swr_df") spike_df = df.copy() - # Load SWR data if path is given if swr_df is None: try: swr_df = pd.read_csv(swr_dir) @@ -173,10 +220,13 @@ def match_up( cluster_ids = np.array(spike_df["Cluster ID"]) for idx in tqdm( - range(len(swr_df)), desc="Checking SWR Data", disable=not progress + range(len(swr_df)), + desc="Checking SWR Data", + disable=not progress, ): - mask = (swr_df["Start"][idx] <= spike_times) & ( - spike_times <= swr_df["Stop"][idx] + mask = ( + (swr_df["Start"][idx] <= spike_times) & + (spike_times <= swr_df["Stop"][idx]) ) swr_df.at[idx, "Spike Times (s)"] = list(spike_times[mask]) swr_df.at[idx, "Cluster IDs"] = list(cluster_ids[mask]) @@ -184,28 +234,22 @@ def match_up( return swr_df -def count_spikes( - swr_df: pd.DataFrame, - mode: str = "all", -) -> dict: +def count_spikes(swr_df: pd.DataFrame, mode="all"): """ Count spikes per cluster across SWRs. - Parameters ---------- - swr_df : pd.DataFrame - Must contain columns: - - "Spike Times (s)" : list/array of spike times - - "Cluster IDs" : list/array of cluster IDs - mode : str - "all" count all spikes of each cluster within each SWR - "first" count only the first spike cluster per SWR + swr_df :: pd.DataFrame + DataFrame of SWR events with assigned spikes. + mode :: str + mode="first": count only the first spike per SWR. + mode="all": count all spikes. Returns ------- - dict : {cluster_id: count} + counts :: dict + Dictionary mapping cluster ID to spike count. """ - if mode not in ("first", "all"): raise ValueError("mode must be 'first' or 'all'") @@ -215,7 +259,6 @@ def count_spikes( times = row["Spike Times (s)"] clusters = row["Cluster IDs"] - # Skip empty SWRs if times is None or len(times) == 0: continue @@ -223,12 +266,10 @@ def count_spikes( clusters = np.array(clusters, dtype=float) if mode == "first": - # find first spike first_cluster = clusters[np.argmin(times)] counts[first_cluster] = counts.get(first_cluster, 0) + 1 - elif mode == "all": - # count every spike event + else: unique, freqs = np.unique(clusters, return_counts=True) for cid, f in zip(unique, freqs): counts[cid] = counts.get(cid, 0) + f @@ -236,370 +277,266 @@ def count_spikes( return counts -def compute_mean_firing_rates( - spike_df: pd.DataFrame, - total_duration: float = None -): - """Compute mean firing rate (Hz) for each cluster.""" - if total_duration is None: - total_duration = spike_df["Time"].max() - counts = spike_df["Cluster ID"].value_counts() - firing_rates = counts / total_duration # spikes per second - - return firing_rates.to_dict() - - -def compute_true_counts( - swr_df: pd.DataFrame, - spike_df: pd.DataFrame, - total_duration: float, - mode: str, -) -> tuple[dict, dict]: - """ - Compute firing rates and observed spike counts during SWRs. - - Parameters - ---------- - swr_df : pd.DataFrame - DataFrame containing SWR windows and their associated spike times. - spike_df : pd.DataFrame - Full spike train dataset containing at least 'Time' and 'ClusterID'. - total_duration : float - Total recording duration in seconds. - mode : {"first", "all"} - Spike counting mode: - - "first": count only the first spike per SWR. - - "all": count all spikes occurring in each SWR. - - Returns - ------- - firing_rates : dict - Mapping of cluster_id -> firing rate (Hz). - true_counts : dict - Mapping of cluster_id -> observed spike count in true SWRs. - """ - firing_rates = compute_mean_firing_rates(spike_df, total_duration) - true_counts = count_spikes(swr_df, mode=mode) - return firing_rates, true_counts - - -def generate_shifts( - n_permutations: int, - shift_range_seconds: float, -) -> np.ndarray: +def compute_mean_firing_rates(spike_df: pd.DataFrame, + total_duration: float = None): """ - Generate random circular shift values for permutation testing. + Compute mean firing rate (Hz) for each cluster. Parameters ---------- - n_permutations : int - Number of permutations to generate. - shift_range_seconds : float - Range of uniform shifts. Values will be sampled from: - [-shift_range_seconds, +shift_range_seconds]. + spike_df :: pd.DataFrame + DataFrame of spike times and cluster IDs. + total_duration :: float + Total recording duration (seconds). If None, inferred from + spike_df. Returns ------- - np.ndarray - Array of random shift values of shape (n_permutations,). - """ - return np.random.uniform( - -shift_range_seconds, shift_range_seconds, size=n_permutations - ) - - -def apply_circular_shift( - swr_df: pd.DataFrame, - shift: float, - total_duration: float, -): + firing_rates :: dict + Dictionary mapping cluster ID to mean firing rate (Hz). """ - Apply a circular temporal shift to SWR start/stop timestamps. + if total_duration is None: + total_duration = spike_df["Time"].max() - Parameters - ---------- - swr_df : pd.DataFrame - Original SWR dataframe with 'Start' and 'Stop' columns. - shift : float - Amount of time (seconds) to shift SWR windows. - total_duration : float - Total duration of the recording. Used for modulo wrap-around. + counts = spike_df["Cluster ID"].value_counts() + firing_rates = counts / total_duration - Returns - ------- - pd.DataFrame - New SWR dataframe with shifted (wrapped) start and stop times. - """ - shifted = swr_df.copy() - shifted["Start"] = (shifted["Start"] + shift) % total_duration - shifted["Stop"] = (shifted["Stop"] + shift) % total_duration - return shifted + return firing_rates.to_dict() -def compute_permutation_counts( - shifts: list | np.ndarray, +def circular_permutation_test_with_firing_rate( swr_df: pd.DataFrame, spike_df: pd.DataFrame, - total_duration: float, - mode: str, + tail: str = "two", + total_duration: float = None, + n_permutations: int = 1000, + shift_range_seconds: float = 3.0, progress: bool = True, + save_path: str = None, + mode: str = "all", ): """ - Generate spike-count dictionaries for each permutation. - + Perform circular permutation tests of SWR windows with firing-rate + normalization. Parameters ---------- - shifts : array-like - List/array of shift values to apply for each permutation. - swr_df : pd.DataFrame - Original SWR dataframe. - spike_df : pd.DataFrame - Full spike dataset, used for matching spikes to shifted SWRs. - total_duration : float - Recording duration for circular wrap-around. - mode : {"first", "all"} - Spike counting mode for SWR windows. - progress : bool, default True - Whether to show tqdm progress bar. - - Yields - ------ - dict - Mapping cluster_id -> spike counts for that permutation. - """ - for shift in tqdm(shifts, disable=not progress): - shifted = apply_circular_shift(swr_df, shift, total_duration) - shifted = match_up( - df=spike_df, swr_df=shifted, only_keep_good=False, progress=False - ) - yield count_spikes(shifted, mode=mode) - - -def aggregate_permuted_counts( - true_counts: dict, - perm_generator: dict, -) -> dict: - """ - Collect permutation spike counts across all permutations. - - Parameters - ---------- - true_counts : dict - Observed true spike counts per cluster. Needed to initialize keys. - perm_generator : iterable of dict - Yields dictionaries of permutation spike counts. - + swr_df :: pd.DataFrame + DataFrame of SWR events with assigned spikes. + spike_df :: pd.DataFrame + DataFrame of spike times and cluster IDs. + tail :: str + "one" for one-tailed test, "two" for two-tailed test. + total_duration :: float + Total recording duration (seconds). If None, inferred from + spike_df. + n_permutations :: int + Number of permutations to run. + shift_range_seconds :: float + Maximum absolute time shift (seconds) for circular permutation. + progress :: bool + If True, show progress bar. + save_path :: str + If provided, save results DataFrame to this path. + mode :: str + "first" to count only first spike per SWR, "all" to count all + spikes. Returns ------- - dict - Mapping cluster_id -> list of spike counts across permutations. + pd.DataFrame: + results DataFrame, with columns: + Cluster ID + Firing Rate (Hz) + True Count + Mean Permuted + Normalized True + Normalized Mean + Z-Score + p-value (one- or two-tailed) + true counts :: dict + true observed counts + all_permuted_counts :: dict + permuted counts """ - out = {cid: [] for cid in true_counts.keys()} - - for perm_counts in perm_generator: - for cid in out.keys(): - out[cid].append(perm_counts.get(cid, 0)) + if total_duration is None: + total_duration = spike_df["Time"].max() - return out + firing_rates = compute_mean_firing_rates(spike_df, total_duration) + true_counts = count_spikes(swr_df, mode) + all_permuted_counts = {cid: [] for cid in true_counts.keys()} -def compute_stats_for_cluster( - cid: int | float, - true_val: int, - perm_vals: list | np.ndarray, - rate: float, - tail: str = "two", -): - """ - Compute statistical metrics for one neuron. + lower_bound = -shift_range_seconds + upper_bound = shift_range_seconds - Parameters - ---------- - cid : int or float - Cluster ID. - true_val : int - Observed spike count in the true SWR dataset. - perm_vals : array-like - Null distribution of spike counts for this neuron. - rate : float - Neuron firing rate in Hz. - tail : {"one", "two"} - Type of p-value test: - - "one" one-sided is greater - - "two" two-sided is greater or less (default) + shifts = np.random.uniform( + low=lower_bound, + high=upper_bound, + size=n_permutations, + ) - Returns - ------- - dict - Dictionary containing statistical metrics including: - z-score, p-value, normalized counts, etc. - """ - perm_vals = np.array(perm_vals) - mean_perm = perm_vals.mean() - std_perm = perm_vals.std(ddof=1) + 1e-6 + for shift in tqdm( + shifts, + desc="Running permutations", + disable=not progress, + ): + shifted_swr = swr_df.copy() + shifted_swr["Start"] = ( + shifted_swr["Start"] + shift + ) % total_duration + shifted_swr["Stop"] = ( + shifted_swr["Stop"] + shift + ) % total_duration + + shifted_swr = match_up( + spike_df, + swr_df=shifted_swr, + only_keep_good=False, + progress=False, + ) + perm_counts = count_spikes(shifted_swr) - # --- p-value --- - if tail == "one": - p_value = (np.sum(perm_vals >= true_val) + 1) / (len(perm_vals) + 1) - else: - p_value = ( - np.sum( - np.abs(perm_vals - mean_perm) >= np.abs(true_val - mean_perm) + for cid in all_permuted_counts.keys(): + all_permuted_counts[cid].append( + perm_counts.get(cid, 0) ) - + 1 - ) / (len(perm_vals) + 1) - - # --- normalized counts --- - normalized_true = true_val / rate if rate > 0 else np.nan - normalized_mean = mean_perm / rate if rate > 0 else np.nan - # --- z-score --- - z_score = (true_val - mean_perm) / std_perm + results = [] - return { - "Cluster ID": cid, - "Firing Rate (Hz)": rate, - "True Count": true_val, - "Mean Permuted": mean_perm, - "Normalized True": normalized_true, - "Normalized Mean": normalized_mean, - "Z-Score": z_score, - f"p-value ({tail}-tailed)": p_value, - } + for cid, true_val in true_counts.items(): + perm_vals = np.array(all_permuted_counts[cid]) + n_perm = len(perm_vals) + mean_perm = perm_vals.mean() + std_perm = perm_vals.std(ddof=1) + 1e-6 + rate = firing_rates.get(cid, np.nan) -def build_results_table( - true_counts: dict, - all_permuted_counts: dict, - firing_rates: dict, - tail: str = "two", -): - """ - Build a results DataFrame from permutation outcomes. + if tail == "one": + p_value = ( + np.sum(perm_vals >= true_val) + 1 + ) / (n_perm + 1) + + elif tail == "two": + p_value = ( + np.sum( + np.abs(perm_vals - mean_perm) >= + np.abs(true_val - mean_perm) + ) + 1 + ) / (n_perm + 1) + + else: + raise ValueError( + "Parameter 'tail' must be 'one' or 'two'." + ) - Parameters - ---------- - true_counts : dict - Observed spike counts per cluster. - all_permuted_counts : dict - Permutation null distribution for each cluster. - firing_rates : dict - Firing rates per cluster in Hz. - tail : {"one", "two"} - Tailedness of p-value calculation. + normalized_true = ( + true_val / rate if rate > 0 else np.nan + ) + normalized_mean = ( + mean_perm / rate if rate > 0 else np.nan + ) + z_score = ( + (true_val - mean_perm) / std_perm + if len(perm_vals) > 1 else np.nan + ) - Returns - ------- - pd.DataFrame - Table of permutation statistics for all neurons. - """ - rows = [] - for cid, true_val in true_counts.items(): - perm_vals = all_permuted_counts[cid] - rate = firing_rates.get(cid, np.nan) - rows.append( - compute_stats_for_cluster(cid, true_val, perm_vals, rate, tail) + results.append( + { + "Cluster ID": cid, + "Firing Rate (Hz)": rate, + "True Count": true_val, + "Mean Permuted": mean_perm, + "Normalized True": normalized_true, + "Normalized Mean": normalized_mean, + "Z-Score": z_score, + f"p-value ({tail}-tailed)": p_value, + } ) - return pd.DataFrame(rows) + result_df = pd.DataFrame(results) -def save_results( - result_df: pd.DataFrame, - save_path: str, -): - """ - Save results as a CSV file if a save path is provided. + if save_path is not None: + os.makedirs( + os.path.dirname(save_path), exist_ok=True + ) + result_df.to_csv(save_path, index=False) - Parameters - ---------- - result_df : pd.DataFrame - DataFrame of permutation metrics. - save_path : str or None - Path to save the CSV file. If None, no file is saved. - """ - if save_path is None: - return + print(f"\n Results saved to: {save_path}") - os.makedirs(os.path.dirname(save_path), exist_ok=True) - result_df.to_csv(save_path, index=False) + return result_df, true_counts, all_permuted_counts -# Full wrapper function to create circular permutation test -# and histograms for individual neuron spike data in SWRs -def circular_permutation_test_with_firing_rate( - swr_df: pd.DataFrame, - spike_df: pd.DataFrame, - tail: str = "two", - total_duration: float = None, - n_permutations: int = 1000, - shift_range_seconds: float = 3.0, - progress: bool = True, - save_path: str = None, - mode: str = "all", -) -> tuple[pd.DataFrame, dict, dict]: +def plot_permutation_histogram( + cid: int, + all_permuted_counts, + true_counts, + binsize=20, + two_tailed=True, +): """ - Run a full circular-permutation test on SWR spike data. - - This wrapper function coordinates the modular components: - - firing rate computation - - true spike counts - - shift generation - - permutation counting - - null distribution aggregation - - statistical testing - - optional saving - + Plot permutation distribution and overlay the observed spike count. Parameters ---------- - swr_df : pd.DataFrame - SWR dataframe with spike timings per event. - spike_df : pd.DataFrame - Full spike dataset (Time, ClusterID). - tail : {"one", "two"}, default "two" - Tailedness of p-value test. - total_duration : float, optional - Recording duration in seconds. If None, uses max spike time. - n_permutations : int, default 1000 - Number of circular permutations. - shift_range_seconds : float, default 3.0 - Uniform shift range (±shift_range_seconds). - progress : bool, default True - Whether to show tqdm progress bar. - save_path : str or None - Output path for saving results CSV. - mode : {"first", "all"}, default "all" - Spike counting strategy for SWRs. - first: count only first spike per SWR. - all: count all spikes in each SWR. + cid :: int + Cluster ID to plot. + all_permuted_counts :: dict + Dictionary mapping cluster ID to list of permuted counts. + true_counts :: dict + Dictionary mapping cluster ID to true observed count. + binsize :: int + Number of bins in the histogram. + two_tailed :: bool + If True, compute two-tailed p-value; else one-tailed. Returns ------- - result_df : pd.DataFrame - Final results table per cluster. - true_counts : dict - Observed spike counts per cluster. - all_permuted_counts : dict - Permutation null distributions per cluster. + Plot histogram and print p-value and z-score. """ - if total_duration is None: - total_duration = spike_df["Time"].max() + perm_vals = np.array(all_permuted_counts[cid]) + true_val = true_counts[cid] + mean_perm = np.mean(perm_vals) - firing_rates, true_counts = compute_true_counts( - swr_df, spike_df, total_duration, mode - ) + if two_tailed: + diff_true = abs(true_val - mean_perm) + diff_perm = abs(perm_vals - mean_perm) + p_value = ( + np.sum(diff_perm >= diff_true) + 1 + ) / (len(perm_vals) + 1) + else: + p_value = ( + np.sum(perm_vals >= true_val) + 1 + ) / (len(perm_vals) + 1) - shifts = generate_shifts(n_permutations, shift_range_seconds) + std_perm = perm_vals.std(ddof=1) + 1e-6 + z_score = (true_val - mean_perm) / std_perm - perm_gen = compute_permutation_counts( - shifts, swr_df, spike_df, total_duration, mode, progress + plt.figure(figsize=(6, 4)) + plt.hist( + perm_vals, + bins=binsize, + color="lightblue", + edgecolor="black", ) - - all_permuted_counts = aggregate_permuted_counts(true_counts, perm_gen) - - result_df = build_results_table( - true_counts, all_permuted_counts, firing_rates, tail + plt.axvline( + true_val, + color="red", + linestyle="--", + linewidth=2, + label=f"True = {true_val}", + ) + plt.title( + f"Cluster {cid} | " + f"{'Two' if two_tailed else 'One'}-tailed " + f"p = {p_value:.3f}, z = {z_score:.2f}" + ) + plt.xlabel("Spike count (permutation)") + plt.ylabel("Frequency") + plt.legend() + plt.tight_layout() + plt.show() + + print( + f"Cluster {cid}: true = {true_val}, " + f"mean_perm = {mean_perm:.2f}, " + f"p = {p_value:.3f}, " + f"z = {z_score:.3f}" ) - - save_results(result_df, save_path) - - return result_df, true_counts, all_permuted_counts diff --git a/src/example_make_raster.py b/src/example_make_raster.py index e968e16..25cfaa3 100644 --- a/src/example_make_raster.py +++ b/src/example_make_raster.py @@ -49,7 +49,6 @@ """ import argparse -import pandas as pd from loading_utils import load_spike_data, match_times from raster_plot_utils import (