bmtool 0.7.0.3__py3-none-any.whl → 0.7.0.5__py3-none-any.whl

bmtool/analysis/entrainment.py

@@ -7,58 +7,71 @@ from scipy import signal
 import numba
 from numba import cuda
 import pandas as pd
-import xarray as xr
-from .lfp import wavelet_filter,butter_bandpass_filter
-from typing import Dict, List
+from .lfp import wavelet_filter,butter_bandpass_filter,get_lfp_power, get_lfp_phase
+from typing import Dict, List, Optional
 from tqdm.notebook import tqdm
 import scipy.stats as stats
-import seaborn as sns
-import matplotlib.pyplot as plt
 
 
-def calculate_signal_signal_plv(x1: np.ndarray, x2: np.ndarray, fs: float, freq_of_interest: float = None, 
-                  method: str = 'wavelet', lowcut: float = None, highcut: float = None, 
+def calculate_signal_signal_plv(signal1: np.ndarray, signal2: np.ndarray, fs: float, freq_of_interest: float = None, 
+                  filter_method: str = 'wavelet', lowcut: float = None, highcut: float = None, 
                   bandwidth: float = 2.0) -> np.ndarray:
     """
     Calculate Phase Locking Value (PLV) between two signals using wavelet or Hilbert method.
     
-    Parameters:
-    - x1, x2: Input signals (1D arrays, same length)
-    - fs: Sampling frequency
-    - freq_of_interest: Desired frequency for wavelet PLV calculation
-    - method: 'wavelet' or 'hilbert' to choose the PLV calculation method
-    - lowcut, highcut: Cutoff frequencies for the Hilbert method
-    - bandwidth: Bandwidth parameter for the wavelet
-    
-    Returns:
-    - plv: Phase Locking Value (1D array)
+    Parameters
+    ----------
+    signal1 : np.ndarray
+        First input signal (1D array)
+    signal2 : np.ndarray
+        Second input signal (1D array, same length as signal1)
+    fs : float
+        Sampling frequency in Hz
+    freq_of_interest : float, optional
+        Desired frequency for wavelet PLV calculation, required if filter_method='wavelet'
+    filter_method : str, optional
+        Method to use for filtering, either 'wavelet' or 'butter' (default: 'wavelet')
+    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'
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 2.0)
+    
+    Returns
+    -------
+    np.ndarray
+        Phase Locking Value (1D array)
     """
-    if len(x1) != len(x2):
+    if len(signal1) != len(signal2):
         raise ValueError("Input signals must have the same length.")
     
-    if method == 'wavelet':
+    if filter_method == 'wavelet':
         if freq_of_interest is None:
             raise ValueError("freq_of_interest must be provided for the wavelet method.")
         
         # Apply CWT to both signals
-        theta1 = wavelet_filter(x=x1, freq=freq_of_interest, fs=fs, bandwidth=bandwidth)
-        theta2 = wavelet_filter(x=x2, freq=freq_of_interest, fs=fs, bandwidth=bandwidth)
+        theta1 = wavelet_filter(x=signal1, freq=freq_of_interest, fs=fs, bandwidth=bandwidth)
+        theta2 = wavelet_filter(x=signal2, freq=freq_of_interest, fs=fs, bandwidth=bandwidth)
     
-    elif method == 'hilbert':
+    elif filter_method == 'butter':
         if lowcut is None or highcut is None:
-            print("Lowcut and or highcut were not definded, signal will not be filter and just take hilbert transform for plv calc")
+            print("Lowcut and/or highcut were not defined, signal will not be filtered and will just take Hilbert transform for PLV calculation")
         
         if lowcut and highcut:
             # Bandpass filter and get the analytic signal using the Hilbert transform
-            x1 = butter_bandpass_filter(x1, lowcut, highcut, fs)
-            x2 = butter_bandpass_filter(x2, lowcut, highcut, fs)
-        
-        # Get phase using the Hilbert transform
-        theta1 = signal.hilbert(x1)
-        theta2 = signal.hilbert(x2)
+            filtered_signal1 = butter_bandpass_filter(data=signal1, lowcut=lowcut, highcut=highcut, fs=fs)
+            filtered_signal2 = butter_bandpass_filter(data=signal2, lowcut=lowcut, highcut=highcut, fs=fs)
+            # Get phase using the Hilbert transform
+            theta1 = signal.hilbert(filtered_signal1)
+            theta2 = signal.hilbert(filtered_signal2)
+        else:
+            # Get phase using the Hilbert transform without filtering
+            theta1 = signal.hilbert(signal1)
+            theta2 = signal.hilbert(signal2)
     
     else:
-        raise ValueError("Invalid method. Choose 'wavelet' or 'hilbert'.")
+        raise ValueError("Invalid method. Choose 'wavelet' or 'butter'.")
     
     # Calculate phase difference
     phase_diff = np.angle(theta1) - np.angle(theta2)
@@ -69,29 +82,43 @@ def calculate_signal_signal_plv(x1: np.ndarray, x2: np.ndarray, fs: float, freq_
     return plv
 
 
-def calculate_spike_lfp_plv(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None, spike_fs : float = None,
-                   lfp_fs: float = None, method: str = 'hilbert', freq_of_interest: float = None,
-                   lowcut: float = None, highcut: float = None,
-                   bandwidth: float = 2.0) -> tuple:
+def calculate_spike_lfp_plv(spike_times: np.ndarray = None, lfp_data: np.ndarray = None, spike_fs: float = None,
+                   lfp_fs: float = None, filter_method: str = 'butter', freq_of_interest: float = None,
+                   lowcut: float = None, highcut: float = None, bandwidth: float = 2.0,
+                   filtered_lfp_phase: np.ndarray = None) -> float:
     """
-    Calculate spike-lfp phase locking value Based on https://www.sciencedirect.com/science/article/pii/S1053811910000959
-    
-    Parameters:
-    - spike_times: Array of spike times 
-    - lfp_signal: Local field potential time series
-    - spike_fs: Sampling frequency in Hz of the spike times only needed if spikes times and lfp has different fs
-    - lfp_fs : Sampling frequency in Hz of the LFP
-    - method: 'wavelet' or 'hilbert' to choose the phase extraction method
-    - freq_of_interest: Desired frequency for wavelet phase extraction
-    - lowcut, highcut: Cutoff frequencies for bandpass filtering (Hilbert method)
-    - bandwidth: Bandwidth parameter for the wavelet
-    
-    Returns:
-    - ppc1: Phase-Phase Coupling value
-    - spike_phases: Phases at spike times
+    Calculate spike-lfp unbiased phase locking value 
+    
+    Parameters
+    ----------
+    spike_times : np.ndarray
+        Array of spike times
+    lfp_data : np.ndarray
+        Local field potential time series data. Not required if filtered_lfp_phase is provided.
+    spike_fs : float, optional
+        Sampling frequency in Hz of the spike times, only needed if spike times and LFP have different sampling rates
+    lfp_fs : float
+        Sampling frequency in Hz of the LFP data
+    filter_method : str, optional
+        Method to use for filtering, either 'wavelet' or 'butter' (default: 'butter')
+    freq_of_interest : float, optional
+        Desired frequency for wavelet phase extraction, required if filter_method='wavelet'
+    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'
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 2.0)
+    filtered_lfp_phase : np.ndarray, optional
+        Pre-computed instantaneous phase of the filtered LFP. If provided, the function will skip the filtering step.
+    
+    Returns
+    -------
+    float
+        Phase Locking Value (unbiased)
     """
     
-    if spike_fs == None:
+    if spike_fs is None:
         spike_fs = lfp_fs
     # Convert spike times to sample indices
     spike_times_seconds = spike_times / spike_fs
@@ -100,50 +127,41 @@ def calculate_spike_lfp_plv(spike_times: np.ndarray = None, lfp_signal: np.ndarr
     spike_indices = np.round(spike_times_seconds * lfp_fs).astype(int)
     
     # Filter indices to ensure they're within bounds of the LFP signal
-    valid_indices = [idx for idx in spike_indices if 0 <= idx < len(lfp_signal)]
+    if filtered_lfp_phase is not None:
+        valid_indices = [idx for idx in spike_indices if 0 <= idx < len(filtered_lfp_phase)]
+    else:
+        valid_indices = [idx for idx in spike_indices if 0 <= idx < len(lfp_data)]
+        
     if len(valid_indices) <= 1:
-        return 0, np.array([])
+        return 0
     
-    # Extract phase using the specified method
-    if method == 'wavelet':
-        if freq_of_interest is None:
-            raise ValueError("freq_of_interest must be provided for the wavelet method.")
-        
-        # Apply CWT to extract phase at the frequency of interest
-        lfp_complex = wavelet_filter(x=lfp_signal, freq=freq_of_interest, fs=lfp_fs, bandwidth=bandwidth)
-        instantaneous_phase = np.angle(lfp_complex)
-        
-    elif method == 'hilbert':
-        if lowcut is None or highcut is None:
-            print("Lowcut and/or highcut were not defined, signal will not be filtered and will just take Hilbert transform for PPC1 calculation")
-            filtered_lfp = lfp_signal
-        else:
-            # Bandpass filter the signal
-            filtered_lfp = butter_bandpass_filter(lfp_signal, lowcut, highcut, lfp_fs)
-        
-        # Get phase using the Hilbert transform
-        analytic_signal = signal.hilbert(filtered_lfp)
-        instantaneous_phase = np.angle(analytic_signal)
-        
+    # Get instantaneous phase
+    if filtered_lfp_phase is None:
+        instantaneous_phase = get_lfp_phase(lfp_data=lfp_data, filter_method=filter_method, 
+                                           freq_of_interest=freq_of_interest, lowcut=lowcut, 
+                                           highcut=highcut, bandwidth=bandwidth, fs=lfp_fs)
     else:
-        raise ValueError("Invalid method. Choose 'wavelet' or 'hilbert'.")
+        instantaneous_phase = filtered_lfp_phase
     
     # Get phases at spike times
     spike_phases = instantaneous_phase[valid_indices]
-    
-    # Calculate PPC1
-    n = len(spike_phases)
+
+    # Number of spikes
+    N = len(spike_phases)
     
     # Convert phases to unit vectors in the complex plane
     unit_vectors = np.exp(1j * spike_phases)
     
-    # Calculate the resultant vector
+    # Sum of all unit vectors (resultant vector)
     resultant_vector = np.sum(unit_vectors)
-    
-    # Plv is the squared length of the resultant vector divided by n²
-    plv = (np.abs(resultant_vector) ** 2) / (n ** 2)
-    
-    return plv
+
+    # Calculate plv^2 * N
+    plv2n = (resultant_vector * resultant_vector.conjugate()).real / N  # plv^2 * N
+    plv = (plv2n / N) ** 0.5
+    ppc = (plv2n - 1) / (N - 1)  # ppc = (plv^2 * N - 1) / (N - 1)
+    plv_unbiased = np.fmax(ppc, 0.) ** 0.5  # ensure non-negative
+
+    return plv_unbiased
 
 
 @numba.njit(parallel=True, fastmath=True)
@@ -181,27 +199,43 @@ def _ppc_gpu(spike_phases):
     return (2/(len(spike_phases)*(len(spike_phases)-1))) * total
 
 
-def calculate_ppc(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None, spike_fs: float = None,
-                  lfp_fs: float = None, method: str = 'hilbert', freq_of_interest: float = None,
-                  lowcut: float = None, highcut: float = None,
-                  bandwidth: float = 2.0,ppc_method: str = 'numpy') -> tuple:
+def calculate_ppc(spike_times: np.ndarray = None, lfp_data: np.ndarray = None, spike_fs: float = None,
+                  lfp_fs: float = None, filter_method: str = 'wavelet', freq_of_interest: float = None,
+                  lowcut: float = None, highcut: float = None, bandwidth: float = 2.0, 
+                  ppc_method: str = 'numpy', filtered_lfp_phase: np.ndarray = None) -> float:
     """
     Calculate Pairwise Phase Consistency (PPC) between spike times and LFP signal.
     Based on https://www.sciencedirect.com/science/article/pii/S1053811910000959
     
-    Parameters:
-    - spike_times: Array of spike times 
-    - lfp_signal: Local field potential time series
-    - spike_fs: Sampling frequency in Hz of the spike times only needed if spikes times and lfp has different fs
-    - lfp_fs: Sampling frequency in Hz of the LFP
-    - method: 'wavelet' or 'hilbert' to choose the phase extraction method
-    - freq_of_interest: Desired frequency for wavelet phase extraction
-    - lowcut, highcut: Cutoff frequencies for bandpass filtering (Hilbert method)
-    - bandwidth: Bandwidth parameter for the wavelet
-    - ppc_method: which algo to use for PPC calculate can be numpy, numba or gpu
-    
-    Returns:
-    - ppc: Pairwise Phase Consistency value
+    Parameters
+    ----------
+    spike_times : np.ndarray
+        Array of spike times
+    lfp_data : np.ndarray
+        Local field potential time series data. Not required if filtered_lfp_phase is provided.
+    spike_fs : float, optional
+        Sampling frequency in Hz of the spike times, only needed if spike times and LFP have different sampling rates
+    lfp_fs : float
+        Sampling frequency in Hz of the LFP data
+    filter_method : str, optional
+        Method to use for filtering, either 'wavelet' or 'butter' (default: 'wavelet')
+    freq_of_interest : float, optional
+        Desired frequency for wavelet phase extraction, required if filter_method='wavelet'
+    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'
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 2.0)
+    ppc_method : str, optional
+        Algorithm to use for PPC calculation: 'numpy', 'numba', or 'gpu' (default: 'numpy')
+    filtered_lfp_phase : np.ndarray, optional
+        Pre-computed instantaneous phase of the filtered LFP. If provided, the function will skip the filtering step.
+    
+    Returns
+    -------
+    float
+        Pairwise Phase Consistency value
     """
     if spike_fs is None:
         spike_fs = lfp_fs
@@ -212,33 +246,21 @@ def calculate_ppc(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None,
     spike_indices = np.round(spike_times_seconds * lfp_fs).astype(int)
     
     # Filter indices to ensure they're within bounds of the LFP signal
-    valid_indices = [idx for idx in spike_indices if 0 <= idx < len(lfp_signal)]
+    if filtered_lfp_phase is not None:
+        valid_indices = [idx for idx in spike_indices if 0 <= idx < len(filtered_lfp_phase)]
+    else:
+        valid_indices = [idx for idx in spike_indices if 0 <= idx < len(lfp_data)]
+        
     if len(valid_indices) <= 1:
-        return 0, np.array([])
+        return 0
     
-    # Extract phase using the specified method
-    if method == 'wavelet':
-        if freq_of_interest is None:
-            raise ValueError("freq_of_interest must be provided for the wavelet method.")
-        
-        # Apply CWT to extract phase at the frequency of interest
-        lfp_complex = wavelet_filter(x=lfp_signal, freq=freq_of_interest, fs=lfp_fs, bandwidth=bandwidth)
-        instantaneous_phase = np.angle(lfp_complex)
-        
-    elif method == 'hilbert':
-        if lowcut is None or highcut is None:
-            print("Lowcut and/or highcut were not defined, signal will not be filtered and will just take Hilbert transform for PPC calculation")
-            filtered_lfp = lfp_signal
-        else:
-            # Bandpass filter the signal
-            filtered_lfp = butter_bandpass_filter(lfp_signal, lowcut, highcut, lfp_fs)
-        
-        # Get phase using the Hilbert transform
-        analytic_signal = signal.hilbert(filtered_lfp)
-        instantaneous_phase = np.angle(analytic_signal)
-        
+    # Get instantaneous phase
+    if filtered_lfp_phase is None:
+        instantaneous_phase = get_lfp_phase(lfp_data=lfp_data, filter_method=filter_method, 
+                                           freq_of_interest=freq_of_interest, lowcut=lowcut, 
+                                           highcut=highcut, bandwidth=bandwidth, fs=lfp_fs)
     else:
-        raise ValueError("Invalid method. Choose 'wavelet' or 'hilbert'.")
+        instantaneous_phase = filtered_lfp_phase
     
     # Get phases at spike times
     spike_phases = instantaneous_phase[valid_indices]
@@ -247,28 +269,10 @@ def calculate_ppc(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None,
 
     # Calculate PPC (Pairwise Phase Consistency)
     if n_spikes <= 1:
-        return 0, spike_phases
+        return 0
     
     # Explicit calculation of pairwise phase consistency
-    sum_cos_diff = 0.0
-    
-    # # Σᵢ Σⱼ₍ᵢ₊₁₎ f(θᵢ, θⱼ)
-    # for i in range(n_spikes - 1):  # For each spike i
-    #     for j in range(i + 1, n_spikes):  # For each spike j > i
-    #         # Calculate the phase difference between spikes i and j
-    #         phase_diff = spike_phases[i] - spike_phases[j]
-            
-    #         #f(θᵢ, θⱼ) = cos(θᵢ)cos(θⱼ) + sin(θᵢ)sin(θⱼ) can become #f(θᵢ, θⱼ) = cos(θᵢ - θⱼ)
-    #         cos_diff = np.cos(phase_diff)
-            
-    #         # Add to the sum
-    #         sum_cos_diff += cos_diff
-    
-    # # Calculate PPC according to the equation
-    # # PPC = (2 / (N(N-1))) * Σᵢ Σⱼ₍ᵢ₊₁₎ f(θᵢ, θⱼ)
-    # ppc = ((2 / (n_spikes * (n_spikes - 1))) * sum_cos_diff)
-    
-    # same as above (i think) but with vectorized computation and memory fixes so it wont take forever to run.
+    # Vectorized computation for efficiency
     if ppc_method == 'numpy':
         i, j = np.triu_indices(n_spikes, k=1)
         phase_diff = spike_phases[i] - spike_phases[j]
@@ -279,14 +283,14 @@ def calculate_ppc(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None,
     elif ppc_method == 'gpu':
         ppc = _ppc_gpu(spike_phases)
     else:
-        raise ExceptionType("Please use a supported ppc method currently that is numpy, numba or gpu")
+        raise ValueError("Please use a supported ppc method currently that is numpy, numba or gpu")
     return ppc
 
     
-def calculate_ppc2(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None, spike_fs: float = None,
-                  lfp_fs: float = None, method: str = 'hilbert', freq_of_interest: float = None,
-                  lowcut: float = None, highcut: float = None,
-                  bandwidth: float = 2.0) -> tuple:
+def calculate_ppc2(spike_times: np.ndarray = None, lfp_data: np.ndarray = None, spike_fs: float = None,
+                  lfp_fs: float = None, filter_method: str = 'wavelet', freq_of_interest: float = None,
+                  lowcut: float = None, highcut: float = None, bandwidth: float = 2.0,
+                  filtered_lfp_phase: np.ndarray = None) -> float:
     """
     # -----------------------------------------------------------------------------
     # PPC2 Calculation (Vinck et al., 2010) 
@@ -297,18 +301,33 @@ def calculate_ppc2(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None
     #   PPC = (|sum(e^(i*φ_j))|^2 - n) / (n * (n - 1))
     # -----------------------------------------------------------------------------
         
-    Parameters:
-    - spike_times: Array of spike times 
-    - lfp_signal: Local field potential time series
-    - spike_fs: Sampling frequency in Hz of the spike times only needed if spikes times and lfp has different fs
-    - lfp_fs: Sampling frequency in Hz of the LFP
-    - method: 'wavelet' or 'hilbert' to choose the phase extraction method
-    - freq_of_interest: Desired frequency for wavelet phase extraction
-    - lowcut, highcut: Cutoff frequencies for bandpass filtering (Hilbert method)
-    - bandwidth: Bandwidth parameter for the wavelet
-    
-    Returns:
-    - ppc2: Pairwise Phase Consistency 2 value
+    Parameters
+    ----------
+    spike_times : np.ndarray
+        Array of spike times
+    lfp_data : np.ndarray
+        Local field potential time series data. Not required if filtered_lfp_phase is provided.
+    spike_fs : float, optional
+        Sampling frequency in Hz of the spike times, only needed if spike times and LFP have different sampling rates
+    lfp_fs : float
+        Sampling frequency in Hz of the LFP data
+    filter_method : str, optional
+        Method to use for filtering, either 'wavelet' or 'butter' (default: 'wavelet')
+    freq_of_interest : float, optional
+        Desired frequency for wavelet phase extraction, required if filter_method='wavelet'
+    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'
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 2.0)
+    filtered_lfp_phase : np.ndarray, optional
+        Pre-computed instantaneous phase of the filtered LFP. If provided, the function will skip the filtering step.
+    
+    Returns
+    -------
+    float
+        Pairwise Phase Consistency 2 (PPC2) value
     """
     
     if spike_fs is None:
@@ -320,33 +339,21 @@ def calculate_ppc2(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None
     spike_indices = np.round(spike_times_seconds * lfp_fs).astype(int)
     
     # Filter indices to ensure they're within bounds of the LFP signal
-    valid_indices = [idx for idx in spike_indices if 0 <= idx < len(lfp_signal)]
+    if filtered_lfp_phase is not None:
+        valid_indices = [idx for idx in spike_indices if 0 <= idx < len(filtered_lfp_phase)]
+    else:
+        valid_indices = [idx for idx in spike_indices if 0 <= idx < len(lfp_data)]
+        
     if len(valid_indices) <= 1:
-        return 0, np.array([])
+        return 0
     
-    # Extract phase using the specified method
-    if method == 'wavelet':
-        if freq_of_interest is None:
-            raise ValueError("freq_of_interest must be provided for the wavelet method.")
-        
-        # Apply CWT to extract phase at the frequency of interest
-        lfp_complex = wavelet_filter(x=lfp_signal, freq=freq_of_interest, fs=lfp_fs, bandwidth=bandwidth)
-        instantaneous_phase = np.angle(lfp_complex)
-        
-    elif method == 'hilbert':
-        if lowcut is None or highcut is None:
-            print("Lowcut and/or highcut were not defined, signal will not be filtered and will just take Hilbert transform for PPC2 calculation")
-            filtered_lfp = lfp_signal
-        else:
-            # Bandpass filter the signal
-            filtered_lfp = butter_bandpass_filter(lfp_signal, lowcut, highcut, lfp_fs)
-        
-        # Get phase using the Hilbert transform
-        analytic_signal = signal.hilbert(filtered_lfp)
-        instantaneous_phase = np.angle(analytic_signal)
-        
+    # Get instantaneous phase
+    if filtered_lfp_phase is None:
+        instantaneous_phase = get_lfp_phase(lfp_data=lfp_data, filter_method=filter_method, 
+                                           freq_of_interest=freq_of_interest, lowcut=lowcut, 
+                                           highcut=highcut, bandwidth=bandwidth, fs=lfp_fs)
     else:
-        raise ValueError("Invalid method. Choose 'wavelet' or 'hilbert'.")
+        instantaneous_phase = filtered_lfp_phase
     
     # Get phases at spike times
     spike_phases = instantaneous_phase[valid_indices]
@@ -355,7 +362,7 @@ def calculate_ppc2(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None
     n = len(spike_phases)
     
     if n <= 1:
-        return 0, spike_phases
+        return 0
     
     # Convert phases to unit vectors in the complex plane
     unit_vectors = np.exp(1j * spike_phases)
@@ -369,40 +376,77 @@ def calculate_ppc2(spike_times: np.ndarray = None, lfp_signal: np.ndarray = None
     return ppc2
 
 
-def calculate_ppc_per_cell(spike_df: pd.DataFrame, lfp_signal: np.ndarray,
-                            spike_fs: float, lfp_fs:float,
-                            pop_names: List[str],freqs: List[float]) -> Dict[str, Dict[int, Dict[float, float]]]:
+def calculate_entrainment_per_cell(spike_df: pd.DataFrame=None, lfp_data: np.ndarray=None, filter_method: str='wavelet', pop_names: List[str]=None,
+                            entrainment_method: str='plv', lowcut: float=None, highcut: float=None,
+                            spike_fs: float=None, lfp_fs: float=None, bandwidth: float=2,
+                            freqs: List[float]=None, ppc_method: str='numpy',) -> Dict[str, Dict[int, Dict[float, float]]]:
     """
-    Calculate pairwise phase consistency (PPC) per neuron (cell) for specified frequencies across different populations.
+    Calculate neural entrainment (PPC, PLV) per neuron (cell) for specified frequencies across different populations.
 
-    This function computes the PPC for each neuron within the specified populations based on their spike times
-    and a single-channel local field potential (LFP) signal.
+    This function computes the entrainment metrics for each neuron within the specified populations based on their spike times
+    and the provided LFP signal. It returns a nested dictionary structure containing the entrainment values
+    organized by population, node ID, and frequency.
 
-    Args:
-        spike_df (pd.DataFrame): Spike dataframe use bmtool.analysis.load_spikes_to_df
-        lfp (xr.DataArray): xarray DataArray representing the LFP use bmtool.analysis.ecp_to_lfp
-        spike_fs (float): sampling rate of spikes BMTK default is 1000
-        lfp_fs (float): sampling rate of lfp 
-        pop_names (List[str]): List of population names (as strings) to compute PPC for. pop_names should be in spike_df
-        freqs (List[float]): List of frequencies (in Hz) at which to calculate PPC.
+    Parameters
+    ----------
+    spike_df : pd.DataFrame
+        DataFrame containing spike data with columns 'pop_name', 'node_ids', and 'timestamps'
+    lfp_data : np.ndarray
+        Local field potential (LFP) time series data
+    filter_method : str, optional
+        Method to use for filtering, either 'wavelet' or 'butter' (default: 'wavelet')
+    entrainment_method : str, optional
+        Method to use for entrainment calculation, either 'plv', 'ppc', or 'ppc2' (default: 'plv')
+    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'
+    spike_fs : float
+        Sampling frequency of the spike times in Hz
+    lfp_fs : float
+        Sampling frequency of the LFP signal in Hz
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 2.0)
+    ppc_method : str, optional
+        Algorithm to use for PPC calculation: 'numpy', 'numba', or 'gpu' (default: 'numpy')
+    pop_names : List[str]
+        List of population names to analyze
+    freqs : List[float]
+        List of frequencies (in Hz) at which to calculate entrainment
 
-    Returns:
-        Dict[str, Dict[int, Dict[float, float]]]: Nested dictionary where the structure is:
-            {
-                population_name: {
-                    node_id: {
-                        frequency: PPC value
-                    }
+    Returns
+    -------
+    Dict[str, Dict[int, Dict[float, float]]]
+        Nested dictionary where the structure is:
+        {
+            population_name: {
+                node_id: {
+                    frequency: entrainment value
                 }
             }
-            PPC values are floats representing the pairwise phase consistency at each frequency.
+        }
+        Entrainment values are floats representing the metric (PPC, PLV) at each frequency
     """
-    ppc_dict = {}
+    # pre filter lfp to speed up calculate of entrainment
+    filtered_lfp_phases = {}
+    for freq in range(len(freqs)):
+        phase = get_lfp_phase(
+            lfp_data=lfp_data, 
+            freq_of_interest=freqs[freq], 
+            fs=lfp_fs, 
+            filter_method=filter_method,
+            lowcut=lowcut, 
+            highcut=highcut, 
+            bandwidth=bandwidth
+        )
+        filtered_lfp_phases[freqs[freq]] = phase
+            
+    entrainment_dict = {}
     for pop in pop_names:
         skip_count = 0
         pop_spikes = spike_df[spike_df['pop_name'] == pop]
         nodes = pop_spikes['node_ids'].unique()
-        ppc_dict[pop] = {}
+        entrainment_dict[pop] = {}
         print(f'Processing {pop} population')
         for node in tqdm(nodes):
             node_spikes = pop_spikes[pop_spikes['node_ids'] == node]
@@ -412,24 +456,58 @@ def calculate_ppc_per_cell(spike_df: pd.DataFrame, lfp_signal: np.ndarray,
                 skip_count += 1
                 continue
 
-            ppc_dict[pop][node] = {}
+            entrainment_dict[pop][node] = {}
             for freq in freqs:
-                ppc = calculate_ppc2(
-                    node_spikes['timestamps'].values,
-                    lfp_signal,
-                    spike_fs=spike_fs,
-                    lfp_fs=lfp_fs,
-                    freq_of_interest=freq,
-                    method='wavelet'
-                )
-                ppc_dict[pop][node][freq] = ppc
+                # Calculate entrainment based on the selected method using the pre-filtered phases
+                if entrainment_method == 'plv':
+                    entrainment_dict[pop][node][freq] = calculate_spike_lfp_plv(
+                        node_spikes['timestamps'].values,
+                        lfp_data,
+                        spike_fs=spike_fs,
+                        lfp_fs=lfp_fs,
+                        freq_of_interest=freq,
+                        bandwidth=bandwidth,
+                        lowcut=lowcut,
+                        highcut=highcut,
+                        filter_method=filter_method,
+                        filtered_lfp_phase=filtered_lfp_phases[freq]
+                    )
+                elif entrainment_method == 'ppc2':
+                    entrainment_dict[pop][node][freq] = calculate_ppc2(
+                        node_spikes['timestamps'].values,
+                        lfp_data,
+                        spike_fs=spike_fs,
+                        lfp_fs=lfp_fs,
+                        freq_of_interest=freq,
+                        bandwidth=bandwidth,
+                        lowcut=lowcut,
+                        highcut=highcut,
+                        filter_method=filter_method,
+                        filtered_lfp_phase=filtered_lfp_phases[freq]
+                    )
+                elif entrainment_method == 'ppc':
+                    entrainment_dict[pop][node][freq] = calculate_ppc(
+                        node_spikes['timestamps'].values,
+                        lfp_data,
+                        spike_fs=spike_fs,
+                        lfp_fs=lfp_fs,
+                        freq_of_interest=freq,
+                        bandwidth=bandwidth,
+                        lowcut=lowcut,
+                        highcut=highcut,
+                        filter_method=filter_method,
+                        ppc_method=ppc_method,
+                        filtered_lfp_phase=filtered_lfp_phases[freq]
+                    )
 
-        print(f'Calculated PPC for {pop} population with {len(nodes)-skip_count} valid cells, skipped {skip_count} cells for lack of spikes')
+        print(f'Calculated {entrainment_method.upper()} for {pop} population with {len(nodes)-skip_count} valid cells, skipped {skip_count} cells for lack of spikes')
 
-    return ppc_dict
+    return entrainment_dict
 
 
-def calculate_spike_rate_power_correlation(spike_rate, lfp, fs, pop_names, freq_range=(10, 100), freq_step=5):
+def calculate_spike_rate_power_correlation(spike_rate, lfp_data, fs, pop_names, filter_method='wavelet',
+                                          bandwidth=2.0, lowcut=None, highcut=None,
+                                          freq_range=(10, 100), freq_step=5):
     """
     Calculate correlation between population spike rates and LFP power across frequencies
     using wavelet filtering. This function assumes the fs of the spike_rate and lfp are the same.
@@ -438,16 +516,24 @@ def calculate_spike_rate_power_correlation(spike_rate, lfp, fs, pop_names, freq_
     -----------
     spike_rate : DataFrame
         Pre-calculated population spike rates at the same fs as lfp
-    lfp : np.array
+    lfp_data : np.array
         LFP data
     fs : float
         Sampling frequency
     pop_names : list
         List of population names to analyze
-    freq_range : tuple
-        Min and max frequency to analyze
-    freq_step : float
-        Step size for frequency analysis
+    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)
     
     Returns:
     --------
@@ -463,14 +549,11 @@ def calculate_spike_rate_power_correlation(spike_rate, lfp, fs, pop_names, freq_
     # Dictionary to store results
     correlation_results = {pop: {} for pop in pop_names}
     
-    # Calculate power at each frequency band using wavelet filter
+    # Calculate power at each frequency band using specified filter
     power_by_freq = {}
     for freq in frequencies:
-        # Use the wavelet_filter function from bmlfp
-        filtered_signal = wavelet_filter(lfp, freq, fs)
-        # Calculate power (magnitude squared of complex wavelet transform)
-        power = np.abs(filtered_signal)**2
-        power_by_freq[freq] = power
+        power_by_freq[freq] = get_lfp_power(lfp_data, freq, fs, filter_method, 
+                                           lowcut=lowcut, highcut=highcut, bandwidth=bandwidth)
     
     # Calculate correlation for each population
     for pop in pop_names:
@@ -481,7 +564,7 @@ def calculate_spike_rate_power_correlation(spike_rate, lfp, fs, pop_names, freq_
         for freq in frequencies:
             # Make sure the lengths match
             if len(pop_rate) != len(power_by_freq[freq]):
-                raise Exception(f"Mismatched lengths for {pop} at {freq} Hz len(pop_rate): {len(pop_rate)}, len(power_by_freq): {len(power_by_freq[freq])}")
+                raise ValueError(f"Mismatched lengths for {pop} at {freq} Hz len(pop_rate): {len(pop_rate)}, len(power_by_freq): {len(power_by_freq[freq])}")
             # use spearman for non-parametric correlation
             corr, p_val = stats.spearmanr(pop_rate, power_by_freq[freq])
             correlation_results[pop][freq] = {'correlation': corr, 'p_value': p_val}

bmtool/analysis/lfp.py

@@ -273,10 +273,83 @@ def calculate_SNR(fooof_model: FOOOF, freq_band: tuple) -> float:
     return normalized_power
 
 
-def wavelet_filter(x: np.ndarray, freq: float, fs: float, bandwidth: float = 1.0, axis: int = -1) -> np.ndarray:
+def calculate_wavelet_passband(center_freq, bandwidth, threshold=0.3):
+    """
+    Calculate the passband of a complex Morlet wavelet filter.
+    
+    Parameters
+    ----------
+    center_freq : float
+        Center frequency (Hz) of the wavelet filter
+    bandwidth : float
+        Bandwidth parameter of the wavelet filter
+    threshold : float, optional
+        Power threshold to define the passband edges (default: 0.5 = -3dB point)
+        
+    Returns
+    -------
+    tuple
+        (lower_bound, upper_bound, passband_width) of the frequency passband in Hz
+    """
+    # Create a high-resolution frequency axis around the center frequency
+    # Extend range to 3x the expected width to ensure we capture the full passband
+    expected_width = center_freq * bandwidth / 2
+    freq_min = max(0.1, center_freq - 3 * expected_width)
+    freq_max = center_freq + 3 * expected_width
+    freq_axis = np.linspace(freq_min, freq_max, 1000)
+    
+    # Calculate the theoretical frequency response of the Morlet wavelet
+    # For a complex Morlet wavelet, the frequency response approximates a Gaussian
+    # centered at the center frequency with width related to the bandwidth parameter
+    sigma_f = bandwidth * center_freq / 8  # Approximate relationship for cmor wavelet
+    response = np.exp(-((freq_axis - center_freq)**2) / (2 * sigma_f**2))
+    
+    # Find the passband edges (where response crosses the threshold)
+    above_threshold = response >= threshold
+    if not np.any(above_threshold):
+        return (center_freq, center_freq, 0)  # No passband found
+    
+    # Find the first and last indices where response is above threshold
+    indices = np.where(above_threshold)[0]
+    lower_idx = indices[0]
+    upper_idx = indices[-1]
+    
+    # Get the corresponding frequencies
+    lower_bound = freq_axis[lower_idx]
+    upper_bound = freq_axis[upper_idx]
+    passband_width = upper_bound - lower_bound
+    
+    return (lower_bound, upper_bound, passband_width)
+
+
+def wavelet_filter(x: np.ndarray, freq: float, fs: float, bandwidth: float = 1.0, axis: int = -1,show_passband: bool = False) -> np.ndarray:
     """
     Compute the Continuous Wavelet Transform (CWT) for a specified frequency using a complex Morlet wavelet.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        Input signal
+    freq : float
+        Target frequency for the wavelet filter
+    fs : float
+        Sampling frequency of the signal
+    bandwidth : float, optional
+        Bandwidth parameter of the wavelet filter (default is 1.0)
+    axis : int, optional
+        Axis along which to compute the CWT (default is -1)
+    show_passband : bool, optional
+        If True, print the passband of the wavelet filter (default is False)
+    
+    Returns
+    -------
+    np.ndarray
+        Continuous Wavelet Transform of the input signal
     """
+    if show_passband:
+        lower_bound, upper_bound, passband_width = calculate_wavelet_passband(freq, bandwidth, threshold=0.3) # kinda made up threshold gives the rough idea
+        print(f"Wavelet filter at {freq:.1f} Hz Bandwidth: {bandwidth:.1f} Hz:")
+        print(f"  Passband: {lower_bound:.1f} - {upper_bound:.1f} Hz (width: {passband_width:.1f} Hz)")
     wavelet = 'cmor' + str(2 * bandwidth ** 2) + '-1.0'
     scale = pywt.scale2frequency(wavelet, 1) * fs / freq
     x_a = pywt.cwt(x, [scale], wavelet=wavelet, axis=axis)[0][0]
@@ -292,6 +365,108 @@ def butter_bandpass_filter(data: np.ndarray, lowcut: float, highcut: float, fs:
     return x_a
 
 
+def get_lfp_power(lfp_data: np.ndarray, freq: float, fs: float, filter_method: str = 'wavelet',
+                   lowcut: float = None, highcut: float = None, bandwidth: float = 1.0) -> np.ndarray:
+    """
+    Compute the power of the raw LFP signal in a specified frequency band.
+    
+    Parameters
+    ----------
+    lfp_data : np.ndarray
+        Raw local field potential (LFP) time series data
+    freq : float
+        Center frequency (Hz) for wavelet filtering method
+    fs : float
+        Sampling frequency (Hz) of the input data
+    filter_method : str, optional
+        Filtering method to use, either 'wavelet' or 'butter' (default: 'wavelet')
+    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'
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 1.0)
+        
+    Returns
+    -------
+    np.ndarray
+        Power of the filtered signal (magnitude squared)
+        
+    Notes
+    -----
+    - The 'wavelet' method uses a complex Morlet wavelet centered at the specified frequency
+    - The 'butter' method uses a Butterworth bandpass filter with the specified cutoff frequencies
+    - When using the 'butter' method, both lowcut and highcut must be provided
+    """
+    if filter_method == 'wavelet':
+        filtered_signal = wavelet_filter(lfp_data, freq, fs, bandwidth)
+    elif filter_method == 'butter':
+        if lowcut is None or highcut is None:
+            raise ValueError("Both lowcut and highcut must be specified when using 'butter' method.")
+        filtered_signal = butter_bandpass_filter(lfp_data, lowcut, highcut, fs)
+    else:
+        raise ValueError("Invalid method. Choose 'wavelet' or 'butter'.")
+    
+    # Calculate power (magnitude squared of filtered signal)
+    power = np.abs(filtered_signal)**2
+    return power
+
+
+def get_lfp_phase(lfp_data: np.ndarray, freq_of_interest: float, fs: float, filter_method: str = 'wavelet',
+                   lowcut: float = None, highcut: float = None, bandwidth: float = 1.0) -> np.ndarray:
+    """
+    Calculate the phase of the filtered signal.
+    
+    Parameters
+    ----------
+    lfp_data : np.ndarray
+        Input LFP data
+    fs : float
+        Sampling frequency (Hz)
+    freq : float
+        Frequency of interest (Hz)
+    filter_method : str, optional
+        Method for filtering the signal ('wavelet' or 'butter')
+    bandwidth : float, optional
+        Bandwidth parameter for wavelet filter when method='wavelet' (default: 1.0)
+    lowcut : float, optional
+        Low cutoff frequency for Butterworth filter when method='butter'
+    highcut : float, optional
+        High cutoff frequency for Butterworth filter when method='butter'
+        
+    Returns
+    -------
+    np.ndarray
+        Phase of the filtered signal
+        
+    Notes
+    -----
+    - The 'wavelet' method uses a complex Morlet wavelet centered at the specified frequency
+    - The 'butter' method uses a Butterworth bandpass filter with the specified cutoff frequencies
+      followed by Hilbert transform to extract the phase
+    - When using the 'butter' method, both lowcut and highcut must be provided
+    """
+    if filter_method == 'wavelet':
+        if freq_of_interest is None:
+            raise ValueError("freq_of_interest must be provided for the wavelet method.")
+        # Wavelet filter returns complex values directly
+        filtered_signal = wavelet_filter(lfp_data, freq_of_interest, fs, bandwidth)
+        # Phase is the angle of the complex signal
+        phase = np.angle(filtered_signal)
+    elif filter_method == 'butter':
+        if lowcut is None or highcut is None:
+            raise ValueError("Both lowcut and highcut must be specified when using 'butter' method.")
+        # Butterworth filter returns real values
+        filtered_signal = butter_bandpass_filter(lfp_data, lowcut, highcut, fs)
+        # Apply Hilbert transform to get analytic signal (complex)
+        analytic_signal = signal.hilbert(filtered_signal)
+        # Phase is the angle of the analytic signal
+        phase = np.angle(analytic_signal)
+    else:
+        raise ValueError(f"Invalid method {filter_method}. Choose 'wavelet' or 'butter'.")
+    
+    return phase
+
 # windowing functions 
 def windowed_xarray(da, windows, dim='time',
                     new_coord_name='cycle', new_coord=None):

bmtool/analysis/spikes.py

@@ -11,22 +11,33 @@ from scipy.stats import mannwhitneyu
 import os
 
 
-def load_spikes_to_df(spike_file: str, network_name: str, sort: bool = True, config: str = None, groupby: str = 'pop_name') -> pd.DataFrame:
+def load_spikes_to_df(spike_file: str, network_name: str, sort: bool = True, config: str = None, groupby: Union[str, List[str]] = 'pop_name') -> pd.DataFrame:
     """
     Load spike data from an HDF5 file into a pandas DataFrame.
 
-    Args:
-        spike_file (str): Path to the HDF5 file containing spike data.
-        network_name (str): The name of the network within the HDF5 file from which to load spike data.
-        sort (bool, optional): Whether to sort the DataFrame by 'timestamps'. Defaults to True.
-        config (str, optional): Will label the cell type of each spike.
-        groupby (str or list of str, optional): The column(s) to group by. Defaults to 'pop_name'.
-
-    Returns:
-        pd.DataFrame: A pandas DataFrame containing 'node_ids' and 'timestamps' columns from the spike data.
+    Parameters
+    ----------
+    spike_file : str
+        Path to the HDF5 file containing spike data
+    network_name : str
+        The name of the network within the HDF5 file from which to load spike data
+    sort : bool, optional
+        Whether to sort the DataFrame by 'timestamps' (default: True)
+    config : str, optional
+        Path to configuration file to label the cell type of each spike (default: None)
+    groupby : Union[str, List[str]], optional
+        The column(s) to group by (default: 'pop_name')
+
+    Returns
+    -------
+    pd.DataFrame
+        A pandas DataFrame containing 'node_ids' and 'timestamps' columns from the spike data,
+        with additional columns if a config file is provided
     
-    Example:
-        df = load_spikes_to_df("spikes.h5", "cortex")
+    Examples
+    --------
+    >>> df = load_spikes_to_df("spikes.h5", "cortex")
+    >>> df = load_spikes_to_df("spikes.h5", "cortex", config="config.json", groupby=["pop_name", "model_type"])
     """
     with h5py.File(spike_file) as f:
         spikes_df = pd.DataFrame({
@@ -126,23 +137,31 @@ def compute_firing_rate_stats(df: pd.DataFrame, groupby: Union[str, List[str]] =
 
 
 def _pop_spike_rate(spike_times: Union[np.ndarray, list], time: Optional[Tuple[float, float, float]] = None, 
-                   time_points: Optional[Union[np.ndarray, list]] = None, frequeny: bool = False) -> np.ndarray:
+                   time_points: Optional[Union[np.ndarray, list]] = None, frequency: bool = False) -> np.ndarray:
     """
     Calculate the spike count or frequency histogram over specified time intervals.
 
-    Args:
-        spike_times (Union[np.ndarray, list]): Array or list of spike times in milliseconds.
-        time (Optional[Tuple[float, float, float]], optional): Tuple specifying (start, stop, step) in milliseconds. 
-            Used to create evenly spaced time points if `time_points` is not provided. Default is None.
-        time_points (Optional[Union[np.ndarray, list]], optional): Array or list of specific time points for binning. 
-            If provided, `time` is ignored. Default is None.
-        frequeny (bool, optional): If True, returns spike frequency in Hz; otherwise, returns spike count. Default is False.
-
-    Returns:
-        np.ndarray: Array of spike counts or frequencies, depending on the `frequeny` flag.
-
-    Raises:
-        ValueError: If both `time` and `time_points` are None.
+    Parameters
+    ----------
+    spike_times : Union[np.ndarray, list]
+        Array or list of spike times in milliseconds
+    time : Optional[Tuple[float, float, float]], optional
+        Tuple specifying (start, stop, step) in milliseconds. Used to create evenly spaced time points 
+        if `time_points` is not provided. Default is None.
+    time_points : Optional[Union[np.ndarray, list]], optional
+        Array or list of specific time points for binning. If provided, `time` is ignored. Default is None.
+    frequency : bool, optional
+        If True, returns spike frequency in Hz; otherwise, returns spike count. Default is False.
+
+    Returns
+    -------
+    np.ndarray
+        Array of spike counts or frequencies, depending on the `frequency` flag.
+
+    Raises
+    ------
+    ValueError
+        If both `time` and `time_points` are None.
     """
     if time_points is None:
         if time is None:
@@ -156,43 +175,57 @@ def _pop_spike_rate(spike_times: Union[np.ndarray, list], time: Optional[Tuple[f
     bins = np.append(time_points, time_points[-1] + dt)
     spike_rate, _ = np.histogram(np.asarray(spike_times), bins)
     
-    if frequeny:
+    if frequency:
         spike_rate = 1000 / dt * spike_rate
     
     return spike_rate
 
 
-def get_population_spike_rate(spikes: pd.DataFrame, fs: float = 400.0, t_start: float = 0, t_stop: Optional[float] = None, 
+def get_population_spike_rate(spike_data: pd.DataFrame, fs: float = 400.0, t_start: float = 0, t_stop: Optional[float] = None, 
                               config: Optional[str] = None, network_name: Optional[str] = None,
                               save: bool = False, save_path: Optional[str] = None,
                               normalize: bool = False) -> Dict[str, np.ndarray]:
     """
     Calculate the population spike rate for each population in the given spike data, with an option to normalize.
 
-    Args:
-        spikes (pd.DataFrame): A DataFrame containing spike data with columns 'pop_name', 'timestamps', and 'node_ids'.
-        fs (float, optional): Sampling frequency in Hz, which determines the time bin size for calculating the spike rate. Default is 400.
-        t_start (float, optional): Start time (in milliseconds) for spike rate calculation. Default is 0.
-        t_stop (Optional[float], optional): Stop time (in milliseconds) for spike rate calculation. If None, defaults to the maximum timestamp in the data.
-        config (Optional[str], optional): Path to a configuration file containing node information, used to determine the correct number of nodes per population. 
-            If None, node count is estimated from unique node spikes. Default is None.
-        network_name (Optional[str], optional): Name of the network used in the configuration file, allowing selection of nodes for that network. 
-            Required if `config` is provided. Default is None.
-        save (bool, optional): Whether to save the calculated population spike rate to a file. Default is False.
-        save_path (Optional[str], optional): Directory path where the file should be saved if `save` is True. If `save` is True and `save_path` is None, a ValueError is raised.
-        normalize (bool, optional): Whether to normalize the spike rates for each population to a range of [0, 1]. Default is False.
-
-    Returns:
-        Dict[str, np.ndarray]: A dictionary where keys are population names, and values are arrays representing the spike rate over time for each population. 
-            If `normalize` is True, each population's spike rate is scaled to [0, 1].
-
-    Raises:
-        ValueError: If `save` is True but `save_path` is not provided.
-
-    Notes:
-        - If `config` is None, the function assumes all cells in each population have fired at least once; otherwise, the node count may be inaccurate.
-        - If normalization is enabled, each population's spike rate is scaled using Min-Max normalization based on its own minimum and maximum values.
-
+    Parameters
+    ----------
+    spike_data : pd.DataFrame
+        A DataFrame containing spike data with columns 'pop_name', 'timestamps', and 'node_ids'
+    fs : float, optional
+        Sampling frequency in Hz, which determines the time bin size for calculating the spike rate (default: 400.0)
+    t_start : float, optional
+        Start time (in milliseconds) for spike rate calculation (default: 0)
+    t_stop : Optional[float], optional
+        Stop time (in milliseconds) for spike rate calculation. If None, defaults to the maximum timestamp in the data
+    config : Optional[str], optional
+        Path to a configuration file containing node information, used to determine the correct number of nodes per population.
+        If None, node count is estimated from unique node spikes (default: None)
+    network_name : Optional[str], optional
+        Name of the network used in the configuration file, allowing selection of nodes for that network.
+        Required if `config` is provided (default: None)
+    save : bool, optional
+        Whether to save the calculated population spike rate to a file (default: False)
+    save_path : Optional[str], optional
+        Directory path where the file should be saved if `save` is True (default: None)
+    normalize : bool, optional
+        Whether to normalize the spike rates for each population to a range of [0, 1] (default: False)
+
+    Returns
+    -------
+    Dict[str, np.ndarray]
+        A dictionary where keys are population names, and values are arrays representing the spike rate over time for each population.
+        If `normalize` is True, each population's spike rate is scaled to [0, 1].
+
+    Raises
+    ------
+    ValueError
+        If `save` is True but `save_path` is not provided.
+
+    Notes
+    -----
+    - If `config` is None, the function assumes all cells in each population have fired at least once; otherwise, the node count may be inaccurate.
+    - If normalization is enabled, each population's spike rate is scaled using Min-Max normalization based on its own minimum and maximum values.
     """
     pop_spikes = {}
     node_number = {}
@@ -205,8 +238,8 @@ def get_population_spike_rate(spikes: pd.DataFrame, fs: float = 400.0, t_start:
         if not network_name:
             print("Grabbing first network; specify a network name to ensure correct node population is selected.")
 
-    for pop_name in spikes['pop_name'].unique():
-        ps = spikes[spikes['pop_name'] == pop_name]
+    for pop_name in spike_data['pop_name'].unique():
+        ps = spike_data[spike_data['pop_name'] == pop_name]
         
         if config:
             nodes = load_nodes_from_config(config)
@@ -220,12 +253,12 @@ def get_population_spike_rate(spikes: pd.DataFrame, fs: float = 400.0, t_start:
             node_number[pop_name] = ps['node_ids'].nunique()
 
         if t_stop is None:
-            t_stop = spikes['timestamps'].max()
+            t_stop = spike_data['timestamps'].max()
 
-        filtered_spikes = spikes[
-            (spikes['pop_name'] == pop_name) & 
-            (spikes['timestamps'] > t_start) & 
-            (spikes['timestamps'] < t_stop)
+        filtered_spikes = spike_data[
+            (spike_data['pop_name'] == pop_name) & 
+            (spike_data['timestamps'] > t_start) & 
+            (spike_data['timestamps'] < t_stop)
         ]
         pop_spikes[pop_name] = filtered_spikes
 
@@ -254,11 +287,30 @@ def get_population_spike_rate(spikes: pd.DataFrame, fs: float = 400.0, t_start:
     return spike_rate
 
 
-def compare_firing_over_times(spike_df,group_by, time_window_1, time_window_2):
+def compare_firing_over_times(spike_df: pd.DataFrame, group_by: str, time_window_1: List[float], time_window_2: List[float]) -> None:
     """
-    Compares the firing rates of a population during two different time windows
-    time_window_1 and time_window_2 should be a list of [start, stop] in milliseconds
-    Returns firing rates and results of a Mann-Whitney U test (non-parametric)
+    Compares the firing rates of a population during two different time windows and performs
+    a statistical test to determine if there is a significant difference.
+    
+    Parameters
+    ----------
+    spike_df : pd.DataFrame
+        DataFrame containing spike data with columns for timestamps, node_ids, and grouping variable
+    group_by : str
+        Column name to group spikes by (e.g., 'pop_name')
+    time_window_1 : List[float]
+        First time window as [start, stop] in milliseconds
+    time_window_2 : List[float]
+        Second time window as [start, stop] in milliseconds
+    
+    Returns
+    -------
+    None
+        Results are printed to the console
+    
+    Notes
+    -----
+    Uses Mann-Whitney U test (non-parametric) to compare firing rates between the two windows
     """
     # Filter spikes for the population of interest
     for pop_name in spike_df[group_by].unique():

{bmtool-0.7.0.3.dist-info → bmtool-0.7.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bmtool
-Version: 0.7.0.3
+Version: 0.7.0.5
 Summary: BMTool
 Home-page: https://github.com/cyneuro/bmtool
 Download-URL: 

{bmtool-0.7.0.3.dist-info → bmtool-0.7.0.5.dist-info}/RECORD

@@ -8,10 +8,10 @@ bmtool/plot_commands.py,sha256=Tqujyf0c0u8olhiHOMwgUSJXIIE1hgjv6otb25G9cA0,12298
 bmtool/singlecell.py,sha256=imcdxIzvYVkaOLSGDxYp8WGGssGwXXBCRhzhlqVp7hA,44267
 bmtool/synapses.py,sha256=Ow2fZavA_3_5BYCjcgPjW0YsyVOetn1wvLxL7hQvbZo,64556
 bmtool/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-bmtool/analysis/entrainment.py,sha256=IMhjbLYw-rL-MfRuFT3uCkyUFObNVJhcxmYV0R9Uh-M,20007
-bmtool/analysis/lfp.py,sha256=3dZkpyVqDtmssvxAqvX4zOLafhwmFnyMehSgVvnj5lM,16754
+bmtool/analysis/entrainment.py,sha256=7lFlGMApL_2snwdvIPDLFW1KKPdyiuCnZ5ADa7ujx5o,24439
+bmtool/analysis/lfp.py,sha256=6u9cHnac-5Fzpk9ecQew7MmXBAolzKZakRsnPn3-C2U,24109
 bmtool/analysis/netcon_reports.py,sha256=7moyoUC45Cl1_6sGqwZ5aKphK_8i4AimroePXcgUnIo,3057
-bmtool/analysis/spikes.py,sha256=x24kd0RUhumJkiunfHNEE7mM6JUqdWy1gqabmkMM4cU,14129
+bmtool/analysis/spikes.py,sha256=kcJZQsvPVzQgcuiO-El_4OODW57hwNwdok_RsFMITCg,15097
 bmtool/bmplot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bmtool/bmplot/connections.py,sha256=re6QZX_NfQnIaWayGt3EhMINhCeMMSQ6rFR2sJbFeWk,51385
 bmtool/bmplot/entrainment.py,sha256=3IBD6tfW7lvkuB6DTan7rAVAeznOOzmHLr1qA2rgtCY,1671
@@ -26,9 +26,9 @@ bmtool/util/commands.py,sha256=zJF-fiLk0b8LyzHDfvewUyS7iumOxVnj33IkJDzux4M,64396
 bmtool/util/util.py,sha256=XR0qZnv_Q47jMBKQpFzCSkCuKe9u8L3YSGJAOpP2zT0,57630
 bmtool/util/neuron/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bmtool/util/neuron/celltuner.py,sha256=xSRpRN6DhPFz4q5buq_W8UmsD7BbUrkzYBEbKVloYss,87194
-bmtool-0.7.0.3.dist-info/licenses/LICENSE,sha256=qrXg2jj6kz5d0EnN11hllcQt2fcWVNumx0xNbV05nyM,1068
-bmtool-0.7.0.3.dist-info/METADATA,sha256=3jYm8B3kDgHTyLyl5C5zdvxZ8tiGoC4hJvAiVu5QLSM,2768
-bmtool-0.7.0.3.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-bmtool-0.7.0.3.dist-info/entry_points.txt,sha256=0-BHZ6nUnh0twWw9SXNTiRmKjDnb1VO2DfG_-oprhAc,45
-bmtool-0.7.0.3.dist-info/top_level.txt,sha256=gpd2Sj-L9tWbuJEd5E8C8S8XkNm5yUE76klUYcM-eWM,7
-bmtool-0.7.0.3.dist-info/RECORD,,
+bmtool-0.7.0.5.dist-info/licenses/LICENSE,sha256=qrXg2jj6kz5d0EnN11hllcQt2fcWVNumx0xNbV05nyM,1068
+bmtool-0.7.0.5.dist-info/METADATA,sha256=5A-VT9HRvmYInIJg4FvfxYSYGL70RzSgOaAaRULRXYs,2768
+bmtool-0.7.0.5.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+bmtool-0.7.0.5.dist-info/entry_points.txt,sha256=0-BHZ6nUnh0twWw9SXNTiRmKjDnb1VO2DfG_-oprhAc,45
+bmtool-0.7.0.5.dist-info/top_level.txt,sha256=gpd2Sj-L9tWbuJEd5E8C8S8XkNm5yUE76klUYcM-eWM,7
+bmtool-0.7.0.5.dist-info/RECORD,,