bmtool 0.7.1.7__py3-none-any.whl → 0.7.2.1__py3-none-any.whl

bmtool/bmplot/entrainment.py

@@ -1,56 +1,299 @@
+from typing import List, Tuple, Union
+
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import seaborn as sns
+import xarray as xr
 from matplotlib.gridspec import GridSpec
 from scipy import stats
 
-
-def plot_spike_power_correlation(correlation_results, frequencies, pop_names):
+from bmtool.analysis import entrainment as bmentr
+from bmtool.analysis import spikes as bmspikes
+from bmtool.analysis.lfp import get_lfp_power
+
+
+def plot_spike_power_correlation(
+    spike_df: pd.DataFrame,
+    lfp_data: xr.DataArray,
+    firing_quantile: float,
+    fs: float,
+    pop_names: list,
+    filter_method: str = "wavelet",
+    bandwidth: float = 2.0,
+    lowcut: float = None,
+    highcut: float = None,
+    freq_range: tuple = (10, 100),
+    freq_step: float = 5,
+    type_name: str = "raw",
+    time_windows: list = None,
+    error_type: str = "ci",  # New parameter: "ci" for confidence interval, "sem" for standard error, "std" for standard deviation
+):
     """
-    Plot the correlation between population spike rates and LFP power.
-
-    Parameters:
-    -----------
-    correlation_results : dict
-        Dictionary with correlation results for calculate_spike_rate_power_correlation
-    frequencies : array
-        Array of frequencies analyzed
+    Calculate and plot correlation between population spike rates and LFP power across frequencies.
+    Supports both single-signal and trial-based analysis with error bars.
+
+    Parameters
+    ----------
+    spike_rate : xr.DataArray
+        Population spike rates with dimensions (time, population[, type])
+    lfp_data : xr.DataArray
+        LFP data
+    fs : float
+        Sampling frequency
     pop_names : list
-        List of population names
+        List of population names to analyze
+    filter_method : str, optional
+        Filtering method to use, either 'wavelet' or 'butter' (default: 'wavelet')
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 2.0)
+    lowcut : float, optional
+        Lower frequency bound (Hz) for butterworth bandpass filter, required if filter_method='butter'
+    highcut : float, optional
+        Upper frequency bound (Hz) for butterworth bandpass filter, required if filter_method='butter'
+    freq_range : tuple, optional
+        Min and max frequency to analyze (default: (10, 100))
+    freq_step : float, optional
+        Step size for frequency analysis (default: 5)
+    type_name : str, optional
+        Which type of spike rate to use if 'type' dimension exists (default: 'raw')
+    time_windows : list, optional
+        List of (start, end) time tuples for trial-based analysis. If None, analyze entire signal
+    error_type : str, optional
+        Type of error bars to plot: "ci" for 95% confidence interval, "sem" for standard error, "std" for standard deviation
     """
-    sns.set_style("whitegrid")
-    plt.figure(figsize=(10, 6))
 
+    if not (0 <= firing_quantile < 1):
+        raise ValueError("firing_quantile must be between 0 and 1")
+
+    if error_type not in ["ci", "sem", "std"]:
+        raise ValueError(
+            "error_type must be 'ci' for confidence interval, 'sem' for standard error, or 'std' for standard deviation"
+        )
+
+    # Setup
+    is_trial_based = time_windows is not None
+
+    # Convert spike_df to spike rate with trial-based filtering of high firing cells
+    if is_trial_based:
+        # Initialize storage for trial-based spike rates
+        trial_rates = []
+
+        for start_time, end_time in time_windows:
+            # Get spikes for this trial
+            trial_spikes = spike_df[
+                (spike_df["timestamps"] >= start_time) & (spike_df["timestamps"] <= end_time)
+            ].copy()
+
+            # Filter for high firing cells within this trial
+            trial_spikes = bmspikes.find_highest_firing_cells(
+                trial_spikes, upper_quantile=firing_quantile
+            )
+            # Calculate rate for this trial's filtered spikes
+            trial_rate = bmspikes.get_population_spike_rate(
+                trial_spikes, fs=fs, t_start=start_time, t_stop=end_time
+            )
+            trial_rates.append(trial_rate)
+
+        # Combine all trial rates
+        spike_rate = xr.concat(trial_rates, dim="trial")
+    else:
+        # For non-trial analysis, proceed as before
+        spike_df = bmspikes.find_highest_firing_cells(spike_df, upper_quantile=firing_quantile)
+        spike_rate = bmspikes.get_population_spike_rate(spike_df)
+
+    # Setup frequencies for analysis
+    frequencies = np.arange(freq_range[0], freq_range[1] + 1, freq_step)
+
+    # Pre-calculate LFP power for all frequencies
+    power_by_freq = {}
+    for freq in frequencies:
+        power_by_freq[freq] = get_lfp_power(
+            lfp_data, freq, fs, filter_method, lowcut=lowcut, highcut=highcut, bandwidth=bandwidth
+        )
+
+    # Calculate correlations
+    results = {}
     for pop in pop_names:
-        # Extract correlation values for each frequency
-        corr_values = []
-        valid_freqs = []
+        pop_spike_rate = spike_rate.sel(population=pop, type=type_name)
+        results[pop] = {}
 
         for freq in frequencies:
-            if freq in correlation_results[pop]:
-                corr_values.append(correlation_results[pop][freq]["correlation"])
-                valid_freqs.append(freq)
+            lfp_power = power_by_freq[freq]
 
-        # Plot correlation line
-        plt.plot(valid_freqs, corr_values, marker="o", label=pop, linewidth=2, markersize=6)
+            if not is_trial_based:
+                # Single signal analysis
+                if len(pop_spike_rate) != len(lfp_power):
+                    print(f"Warning: Length mismatch for {pop} at {freq} Hz")
+                    continue
+
+                corr, p_val = stats.spearmanr(pop_spike_rate, lfp_power)
+                results[pop][freq] = {
+                    "correlation": corr,
+                    "p_value": p_val,
+                }
+            else:
+                # Trial-based analysis using pre-filtered trial rates
+                trial_correlations = []
+
+                for trial_idx in range(len(time_windows)):
+                    # Get time window first
+                    start_time, end_time = time_windows[trial_idx]
+
+                    # Get the pre-filtered spike rate for this trial
+                    trial_spike_rate = pop_spike_rate.sel(trial=trial_idx)
+
+                    # Get corresponding LFP power for this trial window
+                    trial_lfp_power = lfp_power.sel(time=slice(start_time, end_time))
+
+                    # Ensure both signals have same time points
+                    common_times = np.intersect1d(trial_spike_rate.time, trial_lfp_power.time)
+
+                    if len(common_times) > 0:
+                        trial_sr = trial_spike_rate.sel(time=common_times).values
+                        trial_lfp = trial_lfp_power.sel(time=common_times).values
+
+                        if (
+                            len(trial_sr) > 1 and len(trial_lfp) > 1
+                        ):  # Need at least 2 points for correlation
+                            corr, _ = stats.spearmanr(trial_sr, trial_lfp)
+                            if not np.isnan(corr):
+                                trial_correlations.append(corr)
+
+                # Calculate trial statistics
+                if len(trial_correlations) > 0:
+                    trial_correlations = np.array(trial_correlations)
+                    mean_corr = np.mean(trial_correlations)
+
+                    if len(trial_correlations) > 1:
+                        if error_type == "ci":
+                            # Calculate 95% confidence interval using t-distribution
+                            df = len(trial_correlations) - 1
+                            sem = stats.sem(trial_correlations)
+                            t_critical = stats.t.ppf(0.975, df)  # 95% CI, two-tailed
+                            error_val = t_critical * sem
+                            error_lower = mean_corr - error_val
+                            error_upper = mean_corr + error_val
+                        elif error_type == "sem":
+                            # Calculate standard error of the mean
+                            sem = stats.sem(trial_correlations)
+                            error_lower = mean_corr - sem
+                            error_upper = mean_corr + sem
+                        elif error_type == "std":
+                            # Calculate standard deviation
+                            std = np.std(trial_correlations, ddof=1)
+                            error_lower = mean_corr - std
+                            error_upper = mean_corr + std
+                    else:
+                        error_lower = error_upper = mean_corr
+
+                    results[pop][freq] = {
+                        "correlation": mean_corr,
+                        "error_lower": error_lower,
+                        "error_upper": error_upper,
+                        "n_trials": len(trial_correlations),
+                        "trial_correlations": trial_correlations,
+                    }
+                else:
+                    # No valid trials
+                    results[pop][freq] = {
+                        "correlation": np.nan,
+                        "error_lower": np.nan,
+                        "error_upper": np.nan,
+                        "n_trials": 0,
+                        "trial_correlations": np.array([]),
+                    }
+
+    # Plotting
+    sns.set_style("whitegrid")
+    plt.figure(figsize=(12, 8))
 
+    for i, pop in enumerate(pop_names):
+        # Extract data for plotting
+        plot_freqs = []
+        plot_corrs = []
+        plot_ci_lower = []
+        plot_ci_upper = []
+
+        for freq in frequencies:
+            if freq in results[pop] and not np.isnan(results[pop][freq]["correlation"]):
+                plot_freqs.append(freq)
+                plot_corrs.append(results[pop][freq]["correlation"])
+
+                if is_trial_based:
+                    plot_ci_lower.append(results[pop][freq]["error_lower"])
+                    plot_ci_upper.append(results[pop][freq]["error_upper"])
+
+        if len(plot_freqs) == 0:
+            continue
+
+        # Convert to arrays
+        plot_freqs = np.array(plot_freqs)
+        plot_corrs = np.array(plot_corrs)
+
+        # Get color for this population
+        colors = plt.get_cmap("tab10")
+        color = colors(i)
+
+        # Plot main line
+        plt.plot(
+            plot_freqs, plot_corrs, marker="o", label=pop, linewidth=2, markersize=6, color=color
+        )
+
+        # Plot error bands for trial-based analysis
+        if is_trial_based and len(plot_ci_lower) > 0:
+            plot_ci_lower = np.array(plot_ci_lower)
+            plot_ci_upper = np.array(plot_ci_upper)
+            plt.fill_between(plot_freqs, plot_ci_lower, plot_ci_upper, alpha=0.2, color=color)
+
+    # Formatting
     plt.xlabel("Frequency (Hz)", fontsize=12)
     plt.ylabel("Spike Rate-Power Correlation", fontsize=12)
-    plt.title("Spike rate LFP power correlation during stimulus", fontsize=14)
-    plt.grid(True, alpha=0.3)
-    plt.legend(fontsize=12)
-    plt.xticks(frequencies[::2])  # Display every other frequency on x-axis
 
-    # Add horizontal line at zero for reference
+    # Calculate percentage for title
+    firing_percentage = round(float((1 - firing_quantile) * 100), 1)
+    if is_trial_based:
+        title = f"Trial-averaged Spike Rate-LFP Power Correlation\nTop {firing_percentage}% Firing Cells (95% CI)"
+    else:
+        title = f"Spike Rate-LFP Power Correlation\nTop {firing_percentage}% Firing Cells"
+
+    plt.title(title, fontsize=14)
+    plt.grid(True, alpha=0.3)
     plt.axhline(y=0, color="gray", linestyle="-", alpha=0.5)
 
-    # Set y-axis limits to make zero visible
+    # Legend
+    # Create legend elements for each population
+    from matplotlib.lines import Line2D
+
+    colors = plt.get_cmap("tab10")
+    legend_elements = [
+        Line2D([0], [0], color=colors(i), marker="o", linestyle="-", label=pop)
+        for i, pop in enumerate(pop_names)
+    ]
+
+    # Add error band legend element for trial-based analysis
+    if is_trial_based:
+        # Map error type to legend label
+        error_labels = {"ci": "95% CI", "sem": "±SEM", "std": "±1 SD"}
+        error_label = error_labels[error_type]
+
+        legend_elements.append(
+            Line2D([0], [0], color="gray", alpha=0.3, linewidth=10, label=error_label)
+        )
+
+    plt.legend(handles=legend_elements, fontsize=10, loc="best")
+
+    # Axis formatting
+    if len(frequencies) > 10:
+        plt.xticks(frequencies[::2])
+    else:
+        plt.xticks(frequencies)
+    plt.xlim(frequencies[0], frequencies[-1])
+
     y_min, y_max = plt.ylim()
     plt.ylim(min(y_min, -0.1), max(y_max, 0.1))
 
     plt.tight_layout()
-
     plt.show()
 
 
@@ -366,3 +609,257 @@ def plot_entrainment_swarm_plot(ppc_dict, pop_names, freq, save_path=None, title
         plt.savefig(f"{save_path}/ppc_change_swarm_plot_{freq}Hz.png", dpi=300, bbox_inches="tight")
 
     plt.show()
+
+
+def plot_trial_avg_entrainment(
+    spike_df: pd.DataFrame,
+    lfp: np.ndarray,
+    time_windows: List[Tuple[float, float]],
+    entrainment_method: str,
+    pop_names: List[str],
+    freqs: Union[List[float], np.ndarray],
+    firing_quantile: float,
+    spike_fs: float = 1000,
+    error_type: str = "ci",  # New parameter: "ci" for confidence interval, "sem" for standard error, "std" for standard deviation
+) -> None:
+    """
+    Plot trial-averaged entrainment for specified population names. Only supports wavelet filter current, could easily add other support
+
+    Parameters:
+    -----------
+    spike_df : pd.DataFrame
+        Spike data containing timestamps, node_ids, and pop_name columns
+    spike_fs : float
+        fs for spike data. Default is 1000
+    lfp : xarray
+        Xarray for a channel of the lfp data
+    time_windows : List[Tuple[float, float]]
+        List of windows to analysis with start and stp time [(start_time, end_time), ...] for each trial
+    entrainment_method : str
+        Method for entrainment calculation ('ppc', 'ppc2' or 'plv')
+    pop_names : List[str]
+        List of population names to process (e.g., ['FSI', 'LTS'])
+    freqs : Union[List[float], np.ndarray]
+        Array of frequencies to analyze (Hz)
+    firing_quantile : float
+        Upper quantile threshold for selecting high-firing cells (e.g., 0.8 for top 20%)
+    error_type : str
+        Type of error bars to plot: "ci" for 95% confidence interval, "sem" for standard error, "std" for standard deviation
+
+    Raises:
+    -------
+    ValueError
+        If entrainment_method is not 'ppc', 'ppc2' or 'plv'
+        If error_type is not 'ci', 'sem', or 'std'
+        If no spikes found for a population in a trial
+
+    Returns:
+    --------
+    None
+        Displays plot and prints summary statistics
+    """
+    sns.set_style("whitegrid")
+    # Validate inputs
+    if entrainment_method not in ["ppc", "plv", "ppc2"]:
+        raise ValueError("entrainment_method must be 'ppc', ppc2 or 'plv'")
+
+    if error_type not in ["ci", "sem", "std"]:
+        raise ValueError(
+            "error_type must be 'ci' for confidence interval, 'sem' for standard error, or 'std' for standard deviation"
+        )
+
+    if not (0 <= firing_quantile < 1):
+        raise ValueError("firing_quantile must be between 0 and 1")
+
+    # Convert freqs to numpy array for easier indexing
+    freqs = np.array(freqs)
+
+    # Collect all PPC/PLV values across trials for each population
+    all_plv_data = {}  # Dictionary to store results for each population
+
+    # Initialize storage for each population
+    for pop_name in pop_names:
+        all_plv_data[pop_name] = []  # Will be shape (n_trials, n_freqs)
+
+    # Loop through all pulse groups to collect data
+    for trial_idx in range(len(time_windows)):
+        plv_lists = {}  # Store PLV lists for this trial
+
+        # Initialize PLV lists for each population
+        for pop_name in pop_names:
+            plv_lists[pop_name] = []
+
+        # Filter spikes for this trial
+        network_spikes = spike_df[
+            (spike_df["timestamps"] >= time_windows[trial_idx][0])
+            & (spike_df["timestamps"] <= time_windows[trial_idx][1])
+        ].copy()
+
+        # Process each population
+        pop_spike_data = {}
+        for pop_name in pop_names:
+            # Get spikes for this population
+            pop_spikes = network_spikes[network_spikes["pop_name"] == pop_name]
+
+            if len(pop_spikes) == 0:
+                print(f"Warning: No spikes found for population {pop_name} in trial {trial_idx}")
+                # Add NaN values for this trial/population
+                plv_lists[pop_name] = [np.nan] * len(freqs)
+                continue
+
+            # Filter to get the top firing cells
+            # firing_quantile of 0.8 gets the top 20% of firing cells to use
+            pop_spikes = bmspikes.find_highest_firing_cells(
+                pop_spikes, upper_quantile=firing_quantile
+            )
+
+            if len(pop_spikes) == 0:
+                print(
+                    f"Warning: No high-firing spikes found for population {pop_name} in trial {trial_idx}"
+                )
+                plv_lists[pop_name] = [np.nan] * len(freqs)
+                continue
+
+            pop_spike_data[pop_name] = pop_spikes
+
+        # Calculate PPC/PLV for each frequency and each population
+        for freq_idx, freq in enumerate(freqs):
+            for pop_name in pop_names:
+                if pop_name not in pop_spike_data:
+                    continue  # Skip if no data for this population
+
+                pop_spikes = pop_spike_data[pop_name]
+
+                try:
+                    if entrainment_method == "ppc":
+                        result = bmentr.calculate_ppc(
+                            pop_spikes["timestamps"].values,
+                            lfp,
+                            spike_fs=spike_fs,
+                            lfp_fs=lfp.fs,
+                            freq_of_interest=freq,
+                            filter_method="wavelet",
+                            ppc_method="gpu",
+                        )
+                    elif entrainment_method == "plv":
+                        result = bmentr.calculate_spike_lfp_plv(
+                            pop_spikes["timestamps"].values,
+                            lfp,
+                            spike_fs=spike_fs,
+                            lfp_fs=lfp.fs,
+                            freq_of_interest=freq,
+                            filter_method="wavelet",
+                        )
+                    elif entrainment_method == "ppc2":
+                        result = bmentr.calculate_ppc2(
+                            pop_spikes["timestamps"].values,
+                            lfp,
+                            spike_fs=spike_fs,
+                            lfp_fs=lfp.fs,
+                            freq_of_interest=freq,
+                            filter_method="wavelet",
+                        )
+
+                    plv_lists[pop_name].append(result)
+
+                except Exception as e:
+                    print(
+                        f"Warning: Error calculating {entrainment_method} for {pop_name} at {freq}Hz in trial {trial_idx}: {e}"
+                    )
+                    plv_lists[pop_name].append(np.nan)
+
+        # Store this trial's results for each population
+        for pop_name in pop_names:
+            if pop_name in plv_lists and len(plv_lists[pop_name]) == len(freqs):
+                all_plv_data[pop_name].append(plv_lists[pop_name])
+            else:
+                # Fill with NaNs if data is missing
+                all_plv_data[pop_name].append([np.nan] * len(freqs))
+
+    # Convert to numpy arrays and calculate statistics
+    mean_plv = {}
+    error_plv = {}
+
+    for pop_name in pop_names:
+        all_plv_data[pop_name] = np.array(all_plv_data[pop_name])  # Shape: (n_trials, n_freqs)
+
+        # Calculate statistics across trials, ignoring NaN values
+        with np.errstate(invalid="ignore"):  # Suppress warnings for all-NaN slices
+            mean_plv[pop_name] = np.nanmean(all_plv_data[pop_name], axis=0)
+
+            if error_type == "ci":
+                # Calculate 95% confidence intervals using SEM
+                valid_counts = np.sum(~np.isnan(all_plv_data[pop_name]), axis=0)
+                sem_plv = np.nanstd(all_plv_data[pop_name], axis=0, ddof=1) / np.sqrt(valid_counts)
+
+                # For 95% CI, multiply SEM by appropriate t-value
+                # Use minimum valid count across frequencies for conservative t-value
+                min_valid_trials = np.min(valid_counts[valid_counts > 1])  # Avoid division by zero
+                if min_valid_trials > 1:
+                    t_value = stats.t.ppf(0.975, min_valid_trials - 1)  # 95% CI, two-tailed
+                    error_plv[pop_name] = t_value * sem_plv
+                else:
+                    error_plv[pop_name] = np.full_like(sem_plv, np.nan)
+
+            elif error_type == "sem":
+                # Calculate standard error of the mean
+                valid_counts = np.sum(~np.isnan(all_plv_data[pop_name]), axis=0)
+                error_plv[pop_name] = np.nanstd(all_plv_data[pop_name], axis=0, ddof=1) / np.sqrt(
+                    valid_counts
+                )
+
+            elif error_type == "std":
+                # Calculate standard deviation
+                error_plv[pop_name] = np.nanstd(all_plv_data[pop_name], axis=0, ddof=1)
+
+    # Create the combined plot
+    plt.figure(figsize=(12, 8))
+
+    # Define markers and colors for different populations
+    markers = ["o-", "s-", "^-", "D-", "v-", "<-", ">-", "p-"]
+    colors = sns.color_palette(n_colors=len(pop_names))
+
+    # Plot each population
+    for i, pop_name in enumerate(pop_names):
+        marker = markers[i % len(markers)]  # Cycle through markers if more populations than markers
+        color = colors[i]
+
+        # Only plot if we have valid data
+        valid_mask = ~np.isnan(mean_plv[pop_name])
+        if np.any(valid_mask):
+            plt.plot(
+                freqs[valid_mask],
+                mean_plv[pop_name][valid_mask],
+                marker,
+                linewidth=2,
+                label=pop_name,
+                color=color,
+                markersize=6,
+            )
+
+            # Add error bars/shading if available
+            if not np.all(np.isnan(error_plv[pop_name])):
+                plt.fill_between(
+                    freqs[valid_mask],
+                    (mean_plv[pop_name] - error_plv[pop_name])[valid_mask],
+                    (mean_plv[pop_name] + error_plv[pop_name])[valid_mask],
+                    alpha=0.3,
+                    color=color,
+                )
+
+    plt.xlabel("Frequency (Hz)", fontsize=12)
+    plt.ylabel(f"{entrainment_method.upper()}", fontsize=12)
+
+    # Calculate percentage for title and update title based on error type
+    firing_percentage = round(float((1 - firing_quantile) * 100), 1)
+    error_labels = {"ci": "95% CI", "sem": "±SEM", "std": "±1 SD"}
+    error_label = error_labels[error_type]
+    plt.title(
+        f"{entrainment_method.upper()} Across Trials for Top {firing_percentage}% Firing Cells ({error_label})",
+        fontsize=14,
+    )
+
+    plt.legend(fontsize=10)
+    plt.grid(True, alpha=0.3)
+    plt.tight_layout()
+    plt.show()