bmtool 0.7.1.6__py3-none-any.whl → 0.7.2__py3-none-any.whl

bmtool/analysis/entrainment.py

@@ -7,13 +7,12 @@ from typing import Dict, List, Optional, Union
 import numba
 import numpy as np
 import pandas as pd
-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
+from .lfp import butter_bandpass_filter, get_lfp_phase, wavelet_filter
 
 
 def align_spike_times_with_lfp(lfp: xr.DataArray, timestamps: np.ndarray) -> np.ndarray:
@@ -635,85 +634,6 @@ def calculate_entrainment_per_cell(
     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,
-):
-    """
-    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
-        Pre-calculated population spike rates at the same fs as lfp
-    lfp_data : np.array
-        LFP data
-    fs : float
-        Sampling frequency
-    pop_names : list
-        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)
-
-    Returns:
-    --------
-    correlation_results : dict
-        Dictionary with correlation results for each population and frequency
-    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
-        )
-
-    # 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])}"
-                )
-            # 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
-
-
 def get_spikes_in_cycle(
     spike_df,
     lfp_data,

bmtool/analysis/spikes.py

@@ -281,10 +281,13 @@ def get_population_spike_rate(
     node_number = {}
 
     if config is None:
-        print(
-            "Note: Node number is obtained by counting unique node spikes in the network.\nIf the network did not run for a sufficient duration, and not all cells fired, this count might be incorrect."
-        )
-        print("You can provide a config to calculate the correct amount of nodes!")
+        pass
+        # print(
+        #     "Note: Node number is obtained by counting unique node spikes in the network.\nIf the network did not run for a sufficient duration, or not all cells fired,\nthen this count will not include all nodes so the firing rate will not be of the whole population!"
+        # )
+        # print(
+        #     "You can provide a config to calculate the correct amount of nodes! for a true population rate."
+        # )
 
     if config:
         if not network_name:
@@ -600,3 +603,51 @@ def find_bursting_cells(
         )
 
     return burst_cells
+
+
+def find_highest_firing_cells(
+    df: pd.DataFrame, upper_quantile: float, groupby: str = "pop_name"
+) -> pd.DataFrame:
+    """
+    Identifies and returns spikes from cells with firing rates above a specified upper quantile,
+    grouped by a population label.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        DataFrame containing spike data with at least the following columns:
+        - 'timestamps': Time of each spike event
+        - 'node_ids': Identifier for each neuron
+        - groupby (e.g., 'pop_name'): Population labels or grouping identifiers for neurons
+
+    upper_quantile : float
+        The upper quantile threshold (between 0 and 1).
+        Cells with firing rates in the top (1 - upper_quantile) fraction are selected.
+        For example, upper_quantile=0.8 selects the top 20% of high-firing cells.
+
+    groupby : str, optional
+        The column name used to group neurons by population. Default is 'pop_name'.
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame containing only the spikes from the high-firing cells across all groupbys.
+    """
+    if upper_quantile == 0:
+        return df
+    df_list = []
+    for pop in df[groupby].unique():
+        pop_df = df[df[groupby] == pop]
+        _, pop_fr = compute_firing_rate_stats(pop_df, groupby=groupby)
+
+        # Identify high firing cells
+        threshold = pop_fr["firing_rate"].quantile(upper_quantile)
+        high_firing_cells = pop_fr[pop_fr["firing_rate"] >= threshold]["node_ids"]
+
+        # Filter spikes for high firing cells
+        pop_spikes = pop_df[pop_df["node_ids"].isin(high_firing_cells)]
+        df_list.append(pop_spikes)
+
+    # Combine all high firing spikes into one DataFrame
+    result_df = pd.concat(df_list, ignore_index=True)
+    return result_df

bmtool/bmplot/connections.py

@@ -1088,7 +1088,6 @@ def plot_connection_info(
     Function to plot connection information as a heatmap, including handling missing source and target values.
     If there is no source or target, set the value to 0.
     """
-
     # Ensure text dimensions match num dimensions
     num_source = len(source_labels)
     num_target = len(target_labels)
@@ -1096,103 +1095,149 @@ def plot_connection_info(
     # Set color map
     matplotlib.rc("image", cmap="viridis")
 
-    # Create figure and axis for the plot
-    fig1, ax1 = plt.subplots(figsize=(num_source, num_target))
-    num = np.nan_to_num(num, nan=0)  # replace NaN with 0
-    im1 = ax1.imshow(num)
+    # Calculate square cell size to ensure proper aspect ratio
+    base_cell_size = 0.6  # Base size per cell
+
+    # Calculate figure dimensions with proper aspect ratio
+    # Make sure width and height are proportional to the matrix dimensions
+    fig_width = max(8, num_target * base_cell_size + 4)  # Width based on columns
+    fig_height = max(6, num_source * base_cell_size + 3)  # Height based on rows
+
+    # Ensure minimum readable size
+    min_fig_size = 8
+    if fig_width < min_fig_size or fig_height < min_fig_size:
+        scale_factor = min_fig_size / min(fig_width, fig_height)
+        fig_width *= scale_factor
+        fig_height *= scale_factor
+
+    # Create figure and axis
+    fig1, ax1 = plt.subplots(figsize=(fig_width, fig_height))
 
-    # Set ticks and labels for source and target
+    # Replace NaN with 0 and create heatmap
+    num_clean = np.nan_to_num(num, nan=0)
+    # if string is nan\nnan make it 0
+
+    # Use 'auto' aspect ratio to let matplotlib handle it properly
+    # This prevents the stretching issue
+    im1 = ax1.imshow(num_clean, aspect="auto", interpolation="nearest")
+
+    # Set ticks and labels
     ax1.set_xticks(list(np.arange(len(target_labels))))
     ax1.set_yticks(list(np.arange(len(source_labels))))
     ax1.set_xticklabels(target_labels)
-    ax1.set_yticklabels(source_labels, size=12, weight="semibold")
+    ax1.set_yticklabels(source_labels)
 
-    # Rotate the tick labels for better visibility
+    # Improved font sizing based on matrix size
+    label_font_size = max(8, min(14, 120 / max(num_source, num_target)))
+
+    # Style the tick labels
+    ax1.tick_params(axis="y", labelsize=label_font_size, pad=5)
     plt.setp(
         ax1.get_xticklabels(),
         rotation=45,
         ha="right",
         rotation_mode="anchor",
-        size=12,
-        weight="semibold",
+        fontsize=label_font_size,
     )
 
     # Dictionary to store connection information
     graph_dict = {}
 
+    # Improved text size calculation - more readable for larger matrices
+    text_size = max(6, min(12, 80 / max(num_source, num_target)))
+
     # Loop over data dimensions and create text annotations
     for i in range(num_source):
         for j in range(num_target):
-            # Get the edge info, or set it to '0' if it's missing
-            edge_info = text[i, j] if text[i, j] is not None else 0
+            edge_info = text[i, j] if text[i, j] is not None else "0\n0"
 
-            # Initialize the dictionary for the source node if not already done
             if source_labels[i] not in graph_dict:
                 graph_dict[source_labels[i]] = {}
-
-            # Add edge info for the target node
             graph_dict[source_labels[i]][target_labels[j]] = edge_info
 
-            # Set text annotations based on syn_info type
-            if syn_info == "2" or syn_info == "3":
-                if num_source > 8 and num_source < 20:
-                    fig_text = ax1.text(
-                        j,
-                        i,
-                        edge_info,
-                        ha="center",
-                        va="center",
-                        color="w",
-                        rotation=37.5,
-                        size=8,
-                        weight="semibold",
-                    )
-                elif num_source > 20:
-                    fig_text = ax1.text(
-                        j,
-                        i,
-                        edge_info,
-                        ha="center",
-                        va="center",
-                        color="w",
-                        rotation=37.5,
-                        size=7,
-                        weight="semibold",
-                    )
+            # Skip displaying text for NaN values to reduce clutter
+            if edge_info == "nan\nnan":
+                edge_info = "0\n±0"
+
+            # Format the text display
+            if isinstance(edge_info, str) and "\n" in edge_info:
+                # For mean/std format (e.g. "15.5\n4.0")
+                parts = edge_info.split("\n")
+                if len(parts) == 2:
+                    try:
+                        mean_val = float(parts[0])
+                        std_val = float(parts[1])
+                        display_text = f"{mean_val:.1f}\n±{std_val:.1f}"
+                    except ValueError:
+                        display_text = edge_info
                 else:
-                    fig_text = ax1.text(
-                        j,
-                        i,
-                        edge_info,
-                        ha="center",
-                        va="center",
-                        color="w",
-                        rotation=37.5,
-                        size=11,
-                        weight="semibold",
-                    )
+                    display_text = edge_info
+            else:
+                display_text = str(edge_info)
+
+            # Add text to plot with better contrast
+            text_color = "white" if num_clean[i, j] < (np.nanmax(num_clean) * 0.9) else "black"
+
+            if syn_info == "2" or syn_info == "3":
+                ax1.text(
+                    j,
+                    i,
+                    display_text,
+                    ha="center",
+                    va="center",
+                    color=text_color,
+                    rotation=37.5,
+                    fontsize=text_size,
+                    weight="bold",
+                )
             else:
-                fig_text = ax1.text(
-                    j, i, edge_info, ha="center", va="center", color="w", size=11, weight="semibold"
+                ax1.text(
+                    j,
+                    i,
+                    display_text,
+                    ha="center",
+                    va="center",
+                    color=text_color,
+                    fontsize=text_size,
+                    weight="bold",
                 )
 
-    # Set labels and title for the plot
-    ax1.set_ylabel("Source", size=11, weight="semibold")
-    ax1.set_xlabel("Target", size=11, weight="semibold")
-    ax1.set_title(title, size=20, weight="semibold")
+    # Set labels and title
+    title_font_size = max(12, min(18, label_font_size + 4))
+    ax1.set_ylabel("Source", fontsize=title_font_size, weight="bold", labelpad=10)
+    ax1.set_xlabel("Target", fontsize=title_font_size, weight="bold", labelpad=10)
+    ax1.set_title(title, fontsize=title_font_size + 2, weight="bold", pad=20)
+
+    # Add colorbar
+    cbar = plt.colorbar(im1, shrink=0.8)
+    cbar.ax.tick_params(labelsize=label_font_size)
+
+    # Adjust layout to minimize whitespace and prevent stretching
+    plt.tight_layout(pad=1.5)
+
+    # Force square cells by setting equal axis limits if needed
+    ax1.set_xlim(-0.5, num_target - 0.5)
+    ax1.set_ylim(num_source - 0.5, -0.5)  # Inverted for proper matrix orientation
+
+    # Display or save the plot
+    try:
+        # Check if running in notebook
+        from IPython import get_ipython
+
+        notebook = get_ipython() is not None
+    except ImportError:
+        notebook = False
 
-    # Display the plot or save it based on the environment and arguments
-    notebook = is_notebook()  # Check if running in a Jupyter notebook
     if not notebook:
-        fig1.show()
+        plt.show()
 
     if save_file:
-        plt.savefig(save_file)
+        plt.savefig(save_file, dpi=300, bbox_inches="tight", pad_inches=0.1)
 
     if return_dict:
         return graph_dict
     else:
-        return
+        return fig1, ax1
 
 
 def connector_percent_matrix(
@@ -1467,19 +1512,23 @@ def plot_3d_positions(config=None, sources=None, sid=None, title=None, save_file
     plt.title(title)
     plt.legend(handles=handles)
 
+    # Add axis labels
+    ax.set_xlabel("X Position (μm)")
+    ax.set_ylabel("Y Position (μm)")
+    ax.set_zlabel("Z Position (μm)")
+
     # Draw the plot
     plt.draw()
+    plt.tight_layout()
 
     # Save the plot if save_file is provided
     if save_file:
         plt.savefig(save_file)
 
-    # Show the plot if running outside of a notebook
-    if not is_notebook:
+    # Show if running in notebook
+    if is_notebook:
         plt.show()
 
-    return ax
-
 
 def plot_3d_cell_rotation(
     config=None,