bmtool 0.7.7__py3-none-any.whl → 0.7.8.1__py3-none-any.whl

bmtool/bmplot/spikes.py

@@ -1,6 +1,6 @@
 """Plotting functions for neural spikes and firing rates."""
 
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -15,12 +15,13 @@ def raster(
     spikes_df: Optional[pd.DataFrame] = None,
     config: Optional[str] = None,
     network_name: Optional[str] = None,
-    groupby: Optional[str] = "pop_name",
+    groupby: str = "pop_name",
+    sortby: Optional[str] = None,
     ax: Optional[Axes] = None,
     tstart: Optional[float] = None,
     tstop: Optional[float] = None,
     color_map: Optional[Dict[str, str]] = None,
-    dot_size: Optional[float] = 0.3,
+    dot_size: float = 0.3,
 ) -> Axes:
     """
     Plots a raster plot of neural spikes, with different colors for each population.
@@ -33,6 +34,10 @@ def raster(
         Path to the configuration file used to load node data.
     network_name : str, optional
         Specific network name to select from the configuration; if not provided, uses the first network.
+    groupby : str, optional
+        Column name to group spikes by for coloring. Default is 'pop_name'.
+    sortby : str, optional
+        Column name to sort node_ids within each group. If provided, nodes within each population will be sorted by this column.
     ax : matplotlib.axes.Axes, optional
         Axes on which to plot the raster; if None, a new figure and axes are created.
     tstart : float, optional
@@ -107,11 +112,32 @@ def raster(
 
     # Plot each population with its specified or generated color
     legend_handles = []
+    y_offset = 0  # Track y-position offset for stacking populations
+    
     for pop_name, group in spikes_df.groupby(groupby):
-        ax.scatter(group["timestamps"], group["node_ids"], color=color_map[pop_name], s=dot_size)
+        if sortby:
+            # Sort by the specified column, putting NaN values at the end
+            group_sorted = group.sort_values(by=sortby, na_position='last')
+            # Create a mapping from node_ids to consecutive y-positions based on sorted order
+            # Use the sorted order to maintain the same sequence for all spikes from same node
+            unique_nodes_sorted = group_sorted['node_ids'].drop_duplicates()
+            node_to_y = {node_id: y_offset + i for i, node_id in enumerate(unique_nodes_sorted)}
+            # Map node_ids to new y-positions for ALL spikes (not just the sorted group)
+            y_positions = group['node_ids'].map(node_to_y)
+            # Verify no data was lost
+            assert len(y_positions) == len(group), f"Data loss detected in population {pop_name}"
+            assert y_positions.isna().sum() == 0, f"Unmapped node_ids found in population {pop_name}"
+        else:
+            y_positions = group['node_ids']
+            
+        ax.scatter(group["timestamps"], y_positions, color=color_map[pop_name], s=dot_size)
         # Dummy scatter for consistent legend appearance
         handle = ax.scatter([], [], color=color_map[pop_name], label=pop_name, s=20)
         legend_handles.append(handle)
+        
+        # Update y_offset for next population if sortby is used
+        if sortby:
+            y_offset += len(unique_nodes_sorted)
 
     # Label axes
     ax.set_xlabel("Time")
@@ -211,11 +237,12 @@ def plot_firing_rate_pop_stats(
 # uses df from bmtool.analysis.spikes compute_firing_rate_stats
 def plot_firing_rate_distribution(
     individual_stats: pd.DataFrame,
-    groupby: Union[str, list],
+    groupby: Union[str, List[str]],
     ax: Optional[Axes] = None,
     color_map: Optional[Dict[str, str]] = None,
-    plot_type: Union[str, list] = "box",
+    plot_type: Union[str, List[str]] = "box",
     swarm_alpha: float = 0.6,
+    logscale: bool = False,
 ) -> Axes:
     """
     Plots a distribution of individual firing rates using one or more plot types
@@ -235,6 +262,8 @@ def plot_firing_rate_distribution(
         List of plot types to generate. Options: "box", "violin", "swarm". Default is "box".
     swarm_alpha : float, optional
         Transparency of swarm plot points. Default is 0.6.
+    logscale : bool, optional
+        If True, use logarithmic scale for the y-axis (default is False).
 
     Returns:
     -------
@@ -316,40 +345,46 @@ def plot_firing_rate_distribution(
     ax.set_title("Firing Rate Distribution for individual cells")
     ax.grid(axis="y", linestyle="--", alpha=0.7)
 
+    if logscale:
+        ax.set_yscale('log')
+
     return ax
 
 
 def plot_firing_rate_vs_node_attribute(
-    individual_stats: Optional[pd.DataFrame] = None,
+    individual_stats: pd.DataFrame,
+    groupby: str,
+    attribute: str,
     config: Optional[str] = None,
     nodes: Optional[pd.DataFrame] = None,
-    groupby: Optional[str] = None,
     network_name: Optional[str] = None,
-    attribute: Optional[str] = None,
-    figsize=(12, 8),
+    figsize: Tuple[float, float] = (12, 8),
     dot_size: float = 3,
+    color_map: Optional[Dict[str, str]] = None,
 ) -> plt.Figure:
     """
     Plot firing rate vs node attribute for each group in separate subplots.
 
     Parameters
     ----------
-    individual_stats : pd.DataFrame, optional
+    individual_stats : pd.DataFrame
         DataFrame containing individual cell firing rates from compute_firing_rate_stats
+    groupby : str
+        Column name in individual_stats to group plots by
+    attribute : str
+        Node attribute column name to plot against firing rate
     config : str, optional
         Path to configuration file for loading node data
     nodes : pd.DataFrame, optional
         Pre-loaded node data as alternative to loading from config
-    groupby : str, optional
-        Column name in individual_stats to group plots by
     network_name : str, optional
         Name of network to load from config file
-    attribute : str, optional
-        Node attribute column name to plot against firing rate
-    figsize : tuple[int, int], optional
+    figsize : Tuple[float, float], optional
         Figure dimensions (width, height) in inches
     dot_size : float, optional
         Size of scatter plot points
+    color_map : dict, optional
+        Dictionary specifying colors for each group. Keys should be group names, and values should be color values.
 
     Returns
     -------
@@ -407,12 +442,26 @@ def plot_firing_rate_vs_node_attribute(
         axes = np.array([axes])
     axes = axes.flatten()
 
+    # Generate colors if no color_map is provided
+    if color_map is None:
+        cmap = plt.get_cmap("tab10")
+        color_map = {group: cmap(i / len(unique_groups)) for i, group in enumerate(unique_groups)}
+    else:
+        # Ensure color_map contains all groups
+        missing_colors = [group for group in unique_groups if group not in color_map]
+        if missing_colors:
+            raise ValueError(f"color_map is missing colors for groups: {missing_colors}")
+
     # Plot each group
     for i, group in enumerate(unique_groups):
         group_df = merged_df[merged_df[groupby] == group]
-        axes[i].scatter(group_df["firing_rate"], group_df[attribute], s=dot_size)
+        axes[i].scatter(group_df["firing_rate"], group_df[attribute], s=dot_size, color=color_map[group])
         axes[i].set_xlabel("Firing Rate (Hz)")
         axes[i].set_ylabel(attribute)
+        
+        # Calculate and display mean firing rate in legend
+        mean_fr = group_df["firing_rate"].mean()
+        axes[i].legend([f"Mean FR: {mean_fr:.2f} Hz"], loc="upper right")
         axes[i].set_title(f"{groupby}: {group}")
 
     # Hide unused subplots
@@ -420,4 +469,112 @@ def plot_firing_rate_vs_node_attribute(
         axes[j].set_visible(False)
 
     plt.tight_layout()
-    plt.show()
+    return fig
+
+
+def plot_firing_rate_histogram(
+    individual_stats: pd.DataFrame,
+    groupby: str = "pop_name",
+    ax: Optional[Axes] = None,
+    color_map: Optional[Dict[str, str]] = None,
+    bins: int = 30,
+    alpha: float = 0.7,
+    figsize: Tuple[float, float] = (12, 8),
+    stacked: bool = False,
+    logscale: bool = False,
+    min_fr: Optional[float] = None,
+) -> plt.Figure:
+    """
+    Plot histograms of firing rates for each population group.
+
+    Parameters:
+    ----------
+    individual_stats : pd.DataFrame
+        DataFrame containing individual firing rates with group labels.
+    groupby : str, optional
+        Column name to group by (default is "pop_name").
+    ax : matplotlib.axes.Axes, optional
+        Axes on which to plot; if None, a new figure is created.
+    color_map : dict, optional
+        Dictionary specifying colors for each group. Keys should be group names, and values should be color values.
+    bins : int, optional
+        Number of bins for the histogram (default is 30).
+    alpha : float, optional
+        Transparency level for the histograms (default is 0.7).
+    figsize : Tuple[float, float], optional
+        Figure size if creating a new figure (default is (12, 8)).
+    stacked : bool, optional
+        If True, plot all histograms on a single axes stacked (default is False).
+    logscale : bool, optional
+        If True, use logarithmic scale for the x-axis (default is False).
+    min_fr : float, optional
+        Minimum firing rate for log scale bins (default is None).
+
+    Returns:
+    -------
+    matplotlib.figure.Figure
+        Figure containing the histogram subplots.
+    """
+    sns.set_style("whitegrid")
+
+    # Get unique groups
+    unique_groups = individual_stats[groupby].unique()
+
+    # Generate colors if no color_map is provided
+    if color_map is None:
+        cmap = plt.get_cmap("tab10")
+        color_map = {group: cmap(i / len(unique_groups)) for i, group in enumerate(unique_groups)}
+    else:
+        # Ensure color_map contains all groups
+        missing_colors = [group for group in unique_groups if group not in color_map]
+        if missing_colors:
+            raise ValueError(f"color_map is missing colors for groups: {missing_colors}")
+
+    # Group data by population
+    pop_fr = {}
+    for group in unique_groups:
+        pop_fr[group] = individual_stats[individual_stats[groupby] == group]["firing_rate"].values
+
+    if logscale and min_fr is not None:
+        pop_fr = {p: np.fmax(fr, min_fr) for p, fr in pop_fr.items()}
+    fr = np.concatenate(list(pop_fr.values()))
+    if logscale:
+        fr = fr[fr > 0]
+        bins_array = np.geomspace(fr.min(), fr.max(), bins + 1)
+    else:
+        bins_array = np.linspace(fr.min(), fr.max(), bins + 1)
+
+    # Setup subplot layout or single plot
+    n_groups = len(unique_groups)
+    if stacked or not stacked:  # Always use single ax for now, since stacked means overlaid
+        fig, ax = plt.subplots(figsize=figsize)
+    else:
+        # If not stacked, but since overlaid is default, perhaps keep as is
+        fig, ax = plt.subplots(figsize=figsize)
+
+    if stacked:
+        ax.hist(pop_fr.values(), bins=bins_array, label=list(pop_fr.keys()),
+                color=[color_map[p] for p in pop_fr.keys()], stacked=True)
+    else:
+        for p, fr_vals in pop_fr.items():
+            ax.hist(fr_vals, bins=bins_array, label=p, color=color_map[p], alpha=alpha)
+
+    if logscale:
+        ax.set_xscale('log')
+        plt.draw()
+        xt = ax.get_xticks()
+        xtl = [f'{x:g}' for x in xt]
+        if min_fr is not None:
+            xt = np.append(xt, min_fr)
+            xtl.append('0')
+        ax.set_xticks(xt)
+        ax.set_xticklabels(xtl)
+
+    ax.set_xlim(bins_array[0], bins_array[-1])
+    ax.legend(loc='upper right')
+    ax.set_title('Firing Rate Histogram')
+    ax.set_xlabel('Frequency (Hz)')
+    ax.set_ylabel('Count')
+    return fig
+
+

bmtool/singlecell.py

@@ -1042,6 +1042,7 @@ class Profiler:
         self.mechanism_dir = None
         self.templates = None  # Initialize templates attribute
         self.config = config  # Store config path
+        self.last_figure = None  # Store reference to last generated figure
 
         # If a BMTK config is provided, load mechanisms/templates from it
         if config is not None:
@@ -1202,6 +1203,7 @@ class Profiler:
             plt.title("Passive Cell Current Injection")
             plt.xlabel("Time (ms)")
             plt.ylabel("Membrane Potential (mV)")
+            self.last_figure = plt.gcf()
             plt.show()
 
         return time, amp
@@ -1230,6 +1232,8 @@ class Profiler:
             plt.title("Current Injection")
             plt.xlabel("Time (ms)")
             plt.ylabel("Membrane Potential (mV)")
+            plt.xlim(ccl.inj_delay - 10, ccl.inj_delay + ccl.inj_dur + 10)
+            self.last_figure = plt.gcf()
             plt.show()
 
         return time, amp
@@ -1278,6 +1282,7 @@ class Profiler:
             plt.title("FI Curve")
             plt.xlabel("Injection (pA)")
             plt.ylabel("# Spikes")
+            self.last_figure = plt.gcf()
             plt.show()
 
         return amp, nspk
@@ -1317,18 +1322,22 @@ class Profiler:
             plt.title("ZAP Response")
             plt.xlabel("Time (ms)")
             plt.ylabel("Membrane Potential (mV)")
+            self.last_figure = plt.gcf()
 
             plt.figure()
             plt.plot(time, zap.zap_vec)
             plt.title("ZAP Current")
             plt.xlabel("Time (ms)")
             plt.ylabel("Current Injection (nA)")
+            # Note: This will overwrite last_figure with the current plot
+            self.last_figure = plt.gcf()
 
             plt.figure()
             plt.plot(*zap.get_impedance(smooth=smooth))
             plt.title("Impedance Amplitude Profile")
             plt.xlabel("Frequency (Hz)")
             plt.ylabel("Impedance (MOhms)")
+            self.last_figure = plt.gcf()
             plt.show()
 
         return time, amp
@@ -1461,6 +1470,15 @@ class Profiler:
             layout=widgets.Layout(width='150px')
         )
         
+        save_path_text = widgets.Text(value='', description='Save Path:', placeholder='e.g., plot.png', style={'description_width': 'initial'}, layout=widgets.Layout(width='300px'))
+        
+        save_button = widgets.Button(
+            description='Save Plot', 
+            button_style='success', 
+            icon='save',
+            layout=widgets.Layout(width='120px')
+        )
+        
         output_area = widgets.Output(
             layout=widgets.Layout(border='1px solid #ccc', padding='10px', margin='10px 0 0 0')
         )
@@ -1475,7 +1493,8 @@ class Profiler:
         # Button row - main controls
         button_row = widgets.HBox([
             run_button,
-            reset_button
+            reset_button,
+            save_button
         ], layout=widgets.Layout(margin='0 0 10px 0'))
         
         # Section row - recording and injection sections  
@@ -1485,6 +1504,11 @@ class Profiler:
             post_init_text
         ], layout=widgets.Layout(margin='0 0 10px 0'))
         
+        # Save row
+        save_row = widgets.HBox([
+            save_path_text
+        ], layout=widgets.Layout(margin='0 0 10px 0'))
+        
         # Parameter columns - organized in columns like synapse tuner
         injection_params_col1 = widgets.VBox([
             inj_amp_slider,
@@ -1681,15 +1705,36 @@ class Profiler:
                 update_slider_values(method)
                 print(f"Reset all parameters to defaults for {method}")
         
+        # Save function
+        def save_plot(b):
+            path = save_path_text.value
+            if not path:
+                with output_area:
+                    print("Please enter a save path")
+                return
+            if self.last_figure is None:
+                with output_area:
+                    print("No plot to save. Run an analysis first.")
+                return
+            try:
+                self.last_figure.savefig(path)
+                with output_area:
+                    print(f"Plot saved to {path}")
+            except Exception as e:
+                with output_area:
+                    print(f"Error saving plot: {e}")
+        
         run_button.on_click(run_analysis)
         reset_button.on_click(reset_to_defaults)
+        save_button.on_click(save_plot)
         
         # Create main UI layout - matching synapse tuner structure
         ui = widgets.VBox([
             selection_row,
             button_row, 
             section_row,
-            param_columns
+            param_columns,
+            save_row
         ], layout=widgets.Layout(padding='10px'))
         
         # Display the interface - UI on top, output below (like synapse tuner)