bmtool 0.7.0.6.4__py3-none-any.whl → 0.7.1.1__py3-none-any.whl

bmtool/analysis/entrainment.py

@@ -2,23 +2,62 @@
 Module for entrainment analysis
 """
 
-import numpy as np
-from scipy import signal
+from typing import Dict, List
+
 import numba
-from numba import cuda
+import numpy as np
 import pandas as pd
-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 xarray as xr
+from numba import cuda
+from scipy import signal
+from tqdm.notebook import tqdm
+
+from .lfp import butter_bandpass_filter, get_lfp_phase, get_lfp_power, wavelet_filter
+
+
+def align_spike_times_with_lfp(lfp: xr.DataArray, timestamps: np.ndarray) -> np.ndarray:
+    """the lfp xarray should have a time axis. use that to align the spike times since the lfp can start at a
+    non-zero time after sliced. Both need to be on same fs for this to be correct.
 
+    Parameters
+    ----------
+    lfp : xarray.DataArray
+        LFP data with time coordinates
+    timestamps : np.ndarray
+        Array of spike timestamps
 
-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:
+    Returns
+    -------
+    np.ndarray
+        Copy of timestamps with adjusted timestamps to align with lfp.
+    """
+    # print("Pairing LFP and Spike Times")
+    # print(lfp.time.values)
+    # print(f"LFP starts at {lfp.time.values[0]}ms")
+    # need to make sure lfp and spikes have the same time axis
+    # align spikes with lfp
+    timestamps = timestamps[
+        (timestamps >= lfp.time.values[0]) & (timestamps <= lfp.time.values[-1])
+    ].copy()
+    # set the time axis of the spikes to match the lfp
+    timestamps = timestamps - lfp.time.values[0]
+    return timestamps
+
+
+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
     ----------
     signal1 : np.ndarray
@@ -37,7 +76,7 @@ def calculate_signal_signal_plv(signal1: np.ndarray, signal2: np.ndarray, fs: fl
         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
@@ -45,23 +84,29 @@ def calculate_signal_signal_plv(signal1: np.ndarray, signal2: np.ndarray, fs: fl
     """
     if len(signal1) != len(signal2):
         raise ValueError("Input signals must have the same length.")
-    
-    if filter_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=signal1, freq=freq_of_interest, fs=fs, bandwidth=bandwidth)
         theta2 = wavelet_filter(x=signal2, freq=freq_of_interest, fs=fs, bandwidth=bandwidth)
-    
-    elif filter_method == 'butter':
+
+    elif filter_method == "butter":
         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 PLV calculation")
-        
+            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
-            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)
+            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)
@@ -69,26 +114,34 @@ def calculate_signal_signal_plv(signal1: np.ndarray, signal2: np.ndarray, fs: fl
             # Get phase using the Hilbert transform without filtering
             theta1 = signal.hilbert(signal1)
             theta2 = signal.hilbert(signal2)
-    
+
     else:
         raise ValueError("Invalid method. Choose 'wavelet' or 'butter'.")
-    
+
     # Calculate phase difference
     phase_diff = np.angle(theta1) - np.angle(theta2)
-    
+
     # Calculate PLV from standard equation from Measuring phase synchrony in brain signals(1999)
     plv = np.abs(np.mean(np.exp(1j * phase_diff), axis=-1))
-    
+
     return plv
 
 
-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:
+def calculate_spike_lfp_plv(
+    spike_times: np.ndarray = None,
+    lfp_data=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 unbiased phase locking value 
-    
+    Calculate spike-lfp unbiased phase locking value
+
     Parameters
     ----------
     spike_times : np.ndarray
@@ -111,13 +164,13 @@ def calculate_spike_lfp_plv(spike_times: np.ndarray = None, lfp_data: np.ndarray
         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 is None:
         spike_fs = lfp_fs
     # Convert spike times to sample indices
@@ -125,33 +178,50 @@ def calculate_spike_lfp_plv(spike_times: np.ndarray = None, lfp_data: np.ndarray
 
     # Then convert from seconds to samples at the new sampling rate
     spike_indices = np.round(spike_times_seconds * lfp_fs).astype(int)
-    
+
     # Filter indices to ensure they're within bounds of the 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 isinstance(lfp_data, xr.DataArray):
+        if filtered_lfp_phase is not None:
+            valid_indices = align_spike_times_with_lfp(
+                lfp=filtered_lfp_phase, timestamps=spike_indices
+            )
+        else:
+            valid_indices = align_spike_times_with_lfp(lfp=lfp_data, timestamps=spike_indices)
+    elif isinstance(lfp_data, np.ndarray):
+        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
-    
+
     # 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)
+        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:
         instantaneous_phase = filtered_lfp_phase
-    
+
     # Get phases at spike times
-    spike_phases = instantaneous_phase[valid_indices]
+    if isinstance(instantaneous_phase, xr.DataArray):
+        spike_phases = instantaneous_phase.sel(time=valid_indices).values
+    else:
+        spike_phases = instantaneous_phase[valid_indices]
 
     # Number of spikes
     N = len(spike_phases)
-    
+
     # Convert phases to unit vectors in the complex plane
     unit_vectors = np.exp(1j * spike_phases)
-    
+
     # Sum of all unit vectors (resultant vector)
     resultant_vector = np.sum(unit_vectors)
 
@@ -159,7 +229,7 @@ def calculate_spike_lfp_plv(spike_times: np.ndarray = None, lfp_data: np.ndarray
     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
+    plv_unbiased = np.fmax(ppc, 0.0) ** 0.5  # ensure non-negative
 
     return plv_unbiased
 
@@ -181,7 +251,7 @@ def _ppc_cuda_kernel(spike_phases, out):
     i = cuda.grid(1)
     if i < len(spike_phases):
         local_sum = 0.0
-        for j in range(i+1, len(spike_phases)):
+        for j in range(i + 1, len(spike_phases)):
             local_sum += np.cos(spike_phases[i] - spike_phases[j])
         out[i] = local_sum
 
@@ -190,23 +260,32 @@ def _ppc_gpu(spike_phases):
     """GPU-accelerated implementation"""
     d_phases = cuda.to_device(spike_phases)
     d_out = cuda.device_array(len(spike_phases), dtype=np.float64)
-    
+
     threads = 256
     blocks = (len(spike_phases) + threads - 1) // threads
-    
+
     _ppc_cuda_kernel[blocks, threads](d_phases, d_out)
     total = d_out.copy_to_host().sum()
-    return (2/(len(spike_phases)*(len(spike_phases)-1))) * total
-
-
-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:
+    return (2 / (len(spike_phases) * (len(spike_phases) - 1))) * total
+
+
+def calculate_ppc(
+    spike_times: np.ndarray = None,
+    lfp_data=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 : np.ndarray
@@ -231,7 +310,7 @@ def calculate_ppc(spike_times: np.ndarray = None, lfp_data: np.ndarray = None, s
         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
@@ -244,63 +323,88 @@ def calculate_ppc(spike_times: np.ndarray = None, lfp_data: np.ndarray = None, s
 
     # Then convert from seconds to samples at the new sampling rate
     spike_indices = np.round(spike_times_seconds * lfp_fs).astype(int)
-    
+
     # Filter indices to ensure they're within bounds of the 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 isinstance(lfp_data, xr.DataArray):
+        if filtered_lfp_phase is not None:
+            valid_indices = align_spike_times_with_lfp(
+                lfp=filtered_lfp_phase, timestamps=spike_indices
+            )
+        else:
+            valid_indices = align_spike_times_with_lfp(lfp=lfp_data, timestamps=spike_indices)
+    elif isinstance(lfp_data, np.ndarray):
+        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
-    
+
     # 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)
+        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:
         instantaneous_phase = filtered_lfp_phase
-    
+
     # Get phases at spike times
-    spike_phases = instantaneous_phase[valid_indices]
-    
+    if isinstance(instantaneous_phase, xr.DataArray):
+        spike_phases = instantaneous_phase.sel(time=valid_indices).values
+    else:
+        spike_phases = instantaneous_phase[valid_indices]
+
     n_spikes = len(spike_phases)
 
     # Calculate PPC (Pairwise Phase Consistency)
     if n_spikes <= 1:
         return 0
-    
+
     # Explicit calculation of pairwise phase consistency
     # Vectorized computation for efficiency
-    if ppc_method == 'numpy':
+    if ppc_method == "numpy":
         i, j = np.triu_indices(n_spikes, k=1)
         phase_diff = spike_phases[i] - spike_phases[j]
         sum_cos_diff = np.sum(np.cos(phase_diff))
-        ppc = ((2 / (n_spikes * (n_spikes - 1))) * sum_cos_diff)
-    elif ppc_method == 'numba':
+        ppc = (2 / (n_spikes * (n_spikes - 1))) * sum_cos_diff
+    elif ppc_method == "numba":
         ppc = _ppc_parallel_numba(spike_phases)
-    elif ppc_method == 'gpu':
+    elif ppc_method == "gpu":
         ppc = _ppc_gpu(spike_phases)
     else:
         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_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:
+
+def calculate_ppc2(
+    spike_times: np.ndarray = None,
+    lfp_data=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) 
+    # PPC2 Calculation (Vinck et al., 2010)
     # -----------------------------------------------------------------------------
     # Equation(Original):
     #   PPC = (2 / (n * (n - 1))) * sum(cos(φ_i - φ_j) for all i < j)
     # Optimized Formula (Algebraically Equivalent):
     #   PPC = (|sum(e^(i*φ_j))|^2 - n) / (n * (n - 1))
     # -----------------------------------------------------------------------------
-        
+
     Parameters
     ----------
     spike_times : np.ndarray
@@ -323,13 +427,13 @@ def calculate_ppc2(spike_times: np.ndarray = None, lfp_data: np.ndarray = None,
         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:
         spike_fs = lfp_fs
     # Convert spike times to sample indices
@@ -337,49 +441,75 @@ def calculate_ppc2(spike_times: np.ndarray = None, lfp_data: np.ndarray = None,
 
     # Then convert from seconds to samples at the new sampling rate
     spike_indices = np.round(spike_times_seconds * lfp_fs).astype(int)
-    
+
     # Filter indices to ensure they're within bounds of the 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 isinstance(lfp_data, xr.DataArray):
+        if filtered_lfp_phase is not None:
+            valid_indices = align_spike_times_with_lfp(
+                lfp=filtered_lfp_phase, timestamps=spike_indices
+            )
+        else:
+            valid_indices = align_spike_times_with_lfp(lfp=lfp_data, timestamps=spike_indices)
+    elif isinstance(lfp_data, np.ndarray):
+        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
-    
+
     # 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)
+        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:
         instantaneous_phase = filtered_lfp_phase
-    
+
     # Get phases at spike times
-    spike_phases = instantaneous_phase[valid_indices]
-    
+    if isinstance(instantaneous_phase, xr.DataArray):
+        spike_phases = instantaneous_phase.sel(time=valid_indices).values
+    else:
+        spike_phases = instantaneous_phase[valid_indices]
     # Calculate PPC2 according to Vinck et al. (2010), Equation 6
     n = len(spike_phases)
-    
+
     if n <= 1:
         return 0
-    
+
     # Convert phases to unit vectors in the complex plane
     unit_vectors = np.exp(1j * spike_phases)
-    
+
     # Calculate the resultant vector
     resultant_vector = np.sum(unit_vectors)
-    
+
     # PPC2 = (|∑(e^(i*φ_j))|² - n) / (n * (n - 1))
-    ppc2 = (np.abs(resultant_vector)**2 - n) / (n * (n - 1))
-    
+    ppc2 = (np.abs(resultant_vector) ** 2 - n) / (n * (n - 1))
+
     return ppc2
 
 
-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]]]:
+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 neural entrainment (PPC, PLV) per neuron (cell) for specified frequencies across different populations.
 
@@ -431,26 +561,26 @@ def calculate_entrainment_per_cell(spike_df: pd.DataFrame=None, lfp_data: np.nda
     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, 
+            lfp_data=lfp_data,
+            freq_of_interest=freqs[freq],
+            fs=lfp_fs,
             filter_method=filter_method,
-            lowcut=lowcut, 
-            highcut=highcut, 
-            bandwidth=bandwidth
+            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()
+        pop_spikes = spike_df[spike_df["pop_name"] == pop]
+        nodes = pop_spikes["node_ids"].unique()
         entrainment_dict[pop] = {}
-        print(f'Processing {pop} population')
+        print(f"Processing {pop} population")
         for node in tqdm(nodes):
-            node_spikes = pop_spikes[pop_spikes['node_ids'] == node]
-            
+            node_spikes = pop_spikes[pop_spikes["node_ids"] == node]
+
             # Skip nodes with less than or equal to 1 spike
             if len(node_spikes) <= 1:
                 skip_count += 1
@@ -459,9 +589,9 @@ def calculate_entrainment_per_cell(spike_df: pd.DataFrame=None, lfp_data: np.nda
             entrainment_dict[pop][node] = {}
             for freq in freqs:
                 # Calculate entrainment based on the selected method using the pre-filtered phases
-                if entrainment_method == 'plv':
+                if entrainment_method == "plv":
                     entrainment_dict[pop][node][freq] = calculate_spike_lfp_plv(
-                        node_spikes['timestamps'].values,
+                        node_spikes["timestamps"].values,
                         lfp_data,
                         spike_fs=spike_fs,
                         lfp_fs=lfp_fs,
@@ -470,11 +600,11 @@ def calculate_entrainment_per_cell(spike_df: pd.DataFrame=None, lfp_data: np.nda
                         lowcut=lowcut,
                         highcut=highcut,
                         filter_method=filter_method,
-                        filtered_lfp_phase=filtered_lfp_phases[freq]
+                        filtered_lfp_phase=filtered_lfp_phases[freq],
                     )
-                elif entrainment_method == 'ppc2':
+                elif entrainment_method == "ppc2":
                     entrainment_dict[pop][node][freq] = calculate_ppc2(
-                        node_spikes['timestamps'].values,
+                        node_spikes["timestamps"].values,
                         lfp_data,
                         spike_fs=spike_fs,
                         lfp_fs=lfp_fs,
@@ -483,11 +613,11 @@ def calculate_entrainment_per_cell(spike_df: pd.DataFrame=None, lfp_data: np.nda
                         lowcut=lowcut,
                         highcut=highcut,
                         filter_method=filter_method,
-                        filtered_lfp_phase=filtered_lfp_phases[freq]
+                        filtered_lfp_phase=filtered_lfp_phases[freq],
                     )
-                elif entrainment_method == 'ppc':
+                elif entrainment_method == "ppc":
                     entrainment_dict[pop][node][freq] = calculate_ppc(
-                        node_spikes['timestamps'].values,
+                        node_spikes["timestamps"].values,
                         lfp_data,
                         spike_fs=spike_fs,
                         lfp_fs=lfp_fs,
@@ -497,21 +627,32 @@ def calculate_entrainment_per_cell(spike_df: pd.DataFrame=None, lfp_data: np.nda
                         highcut=highcut,
                         filter_method=filter_method,
                         ppc_method=ppc_method,
-                        filtered_lfp_phase=filtered_lfp_phases[freq]
+                        filtered_lfp_phase=filtered_lfp_phases[freq],
                     )
 
-        print(f'Calculated {entrainment_method.upper()} 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 entrainment_dict
 
 
-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):
+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.
-    
+
     Parameters:
     -----------
     spike_rate : DataFrame
@@ -534,7 +675,7 @@ def calculate_spike_rate_power_correlation(spike_rate, lfp_data, fs, pop_names,
         Min and max frequency to analyze (default: (10, 100))
     freq_step : float, optional
         Step size for frequency analysis (default: 5)
-    
+
     Returns:
     --------
     correlation_results : dict
@@ -542,32 +683,34 @@ def calculate_spike_rate_power_correlation(spike_rate, lfp_data, fs, pop_names,
     frequencies : array
         Array of frequencies analyzed
     """
-    
+
     # Define frequency bands to analyze
     frequencies = np.arange(freq_range[0], freq_range[1] + 1, freq_step)
-    
+
     # Dictionary to store results
     correlation_results = {pop: {} for pop in pop_names}
-    
+
     # Calculate power at each frequency band using specified filter
     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)
-    
+        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:
         # Extract spike rate for this population
         pop_rate = spike_rate[pop]
-        
+
         # Calculate correlation with power at each frequency
         for freq in frequencies:
             # Make sure the lengths match
             if len(pop_rate) != 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])}")
+                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}
-    
-    return correlation_results, frequencies
+            correlation_results[pop][freq] = {"correlation": corr, "p_value": p_val}
 
+    return correlation_results, frequencies