nettracer3d 0.7.9__py3-none-any.whl → 0.8.1__py3-none-any.whl

nettracer3d/community_extractor.py

@@ -492,48 +492,39 @@ def generate_distinct_colors(n_colors: int) -> List[Tuple[int, int, int]]:
     return colors
 
 def assign_node_colors(node_list: List[int], labeled_array: np.ndarray) -> Tuple[np.ndarray, Dict[int, str]]:
-    """
-    Assign distinct colors to nodes and create an RGBA image.
-    
-    Args:
-        node_list: List of node IDs
-        labeled_array: 3D numpy array with labels corresponding to node IDs
+    """fast version using lookup table approach."""
     
-    Returns:
-        Tuple of (RGBA-coded numpy array (H, W, D, 4), dictionary mapping nodes to color names)
-    """
-    
-    # Sort communities by size (descending)
-    sorted_nodes= sorted(node_list, reverse=True)
+    # Sort nodes by size (descending)
+    sorted_nodes = sorted(node_list, reverse=True)
     
     # Generate distinct colors
     colors = generate_distinct_colors(len(node_list))
-    random.shuffle(colors) #Randomly sorted to make adjacent structures likely stand out
+    random.shuffle(colors)  # Randomly sorted to make adjacent structures likely stand out
     
     # Convert RGB colors to RGBA by adding alpha channel
-    colors_rgba = [(r, g, b, 255) for r, g, b in colors]  # Full opacity for colored regions
+    colors_rgba = np.array([(r, g, b, 255) for r, g, b in colors], dtype=np.uint8)
     
-    # Create mapping from community to color
+    # Create mapping from node to color
     node_to_color = {node: colors_rgba[i] for i, node in enumerate(sorted_nodes)}
     
-    # Create RGBA array (initialize with transparent background)
-    rgba_array = np.zeros((*labeled_array.shape, 4), dtype=np.uint8)
+    # Create lookup table
+    max_label = max(max(labeled_array.flat), max(node_list) if node_list else 0)
+    color_lut = np.zeros((max_label + 1, 4), dtype=np.uint8)  # Transparent by default
     
-    # Assign colors to each voxel based on its label
-    for label in np.unique(labeled_array):
-        if label in node_to_color:  # Skip background (usually label 0)
-            mask = labeled_array == label
-            for i in range(4):  # RGBA channels
-                rgba_array[mask, i] = node_to_color[label][i]
-
-    # Convert the RGB portion of community_to_color back to RGB for color naming
+    for node_id, color in node_to_color.items():
+        color_lut[node_id] = color
+    
+    # Single vectorized operation - eliminates all loops!
+    rgba_array = color_lut[labeled_array]
+    
+    # Convert colors for naming
     node_to_color_rgb = {k: tuple(v[:3]) for k, v in node_to_color.items()}
     node_to_color_names = convert_node_colors_to_names(node_to_color_rgb)
     
     return rgba_array, node_to_color_names
 
 def assign_community_colors(community_dict: Dict[int, int], labeled_array: np.ndarray) -> Tuple[np.ndarray, Dict[int, str]]:
-    """Ultra-fast version using lookup table approach."""
+    """fast version using lookup table approach."""
     
     # Same setup as before
     communities = set(community_dict.values())

nettracer3d/neighborhoods.py

@@ -202,10 +202,12 @@ def visualize_cluster_composition_umap(cluster_data: Dict[int, np.ndarray],
     
     return embedding
 
-def create_community_heatmap(community_intensity, node_community, node_centroids, is_3d=True, 
-                           figsize=(12, 8), point_size=50, alpha=0.7, colorbar_label="Community Intensity"):
+def create_community_heatmap(community_intensity, node_community, node_centroids, shape=None, is_3d=True, 
+                           labeled_array=None, figsize=(12, 8), point_size=50, alpha=0.7, 
+                           colorbar_label="Community Intensity", title="Community Intensity Heatmap"):
     """
     Create a 2D or 3D heatmap showing nodes colored by their community intensities.
+    Can return either matplotlib plot or numpy RGB array for overlay purposes.
     
     Parameters:
     -----------
@@ -220,25 +222,39 @@ def create_community_heatmap(community_intensity, node_community, node_centroids
         Dictionary mapping node IDs to centroids
         Centroids should be [Z, Y, X] for 3D or [1, Y, X] for pseudo-3D
         
+    shape : tuple, optional
+        Shape of the output array in [Z, Y, X] format
+        If None, will be inferred from node_centroids
+        
     is_3d : bool, default=True
-        If True, create 3D plot. If False, create 2D plot.
+        If True, create 3D plot/array. If False, create 2D plot/array.
+        
+    labeled_array : np.ndarray, optional
+        If provided, returns numpy RGB array overlay using this labeled array template
+        instead of matplotlib plot. Uses lookup table approach for efficiency.
         
     figsize : tuple, default=(12, 8)
-        Figure size (width, height)
+        Figure size (width, height) - only used for matplotlib
         
     point_size : int, default=50
-        Size of scatter plot points
+        Size of scatter plot points - only used for matplotlib
         
     alpha : float, default=0.7
-        Transparency of points (0-1)
+        Transparency of points (0-1) - only used for matplotlib
         
     colorbar_label : str, default="Community Intensity"
-        Label for the colorbar
+        Label for the colorbar - only used for matplotlib
+        
+    title : str, default="Community Intensity Heatmap"
+        Title for the plot
         
     Returns:
     --------
-    fig, ax : matplotlib figure and axis objects
+    If labeled_array is None: fig, ax (matplotlib figure and axis objects)
+    If labeled_array is provided: np.ndarray (RGB heatmap array with community intensity colors)
     """
+    import numpy as np
+    import matplotlib.pyplot as plt
     
     # Convert numpy int64 keys to regular ints for consistency
     community_intensity_clean = {}
@@ -254,6 +270,10 @@ def create_community_heatmap(community_intensity, node_community, node_centroids
     
     for node_id, centroid in node_centroids.items():
         try:
+            # Convert node_id to regular int if it's numpy
+            if hasattr(node_id, 'item'):
+                node_id = node_id.item()
+                
             # Get community for this node
             community_id = node_community[node_id]
             
@@ -266,74 +286,392 @@ def create_community_heatmap(community_intensity, node_community, node_centroids
             
             node_positions.append(centroid)
             node_intensities.append(intensity)
-        except:
+        except KeyError:
+            # Skip nodes that don't have community assignments or community intensities
             pass
     
     # Convert to numpy arrays
     positions = np.array(node_positions)
     intensities = np.array(node_intensities)
     
-    # Determine min and max intensities for color scaling
-    min_intensity = np.min(intensities)
-    max_intensity = np.max(intensities)
+    # Determine shape if not provided
+    if shape is None:
+        if len(positions) > 0:
+            max_coords = np.max(positions, axis=0).astype(int)
+            shape = tuple(max_coords + 1)
+        else:
+            shape = (100, 100, 100) if is_3d else (1, 100, 100)
     
-    # Create figure
-    fig = plt.figure(figsize=figsize)
+    # Determine min and max intensities for scaling
+    if len(intensities) > 0:
+        min_intensity = np.min(intensities)
+        max_intensity = np.max(intensities)
+    else:
+        min_intensity, max_intensity = 0, 1
     
-    if is_3d:
-        # 3D plot
-        ax = fig.add_subplot(111, projection='3d')
+    if labeled_array is not None:
+        # Create numpy RGB array output using labeled array and lookup table approach
+        
+        # Create mapping from node ID to community intensity value
+        node_to_community_intensity = {}
+        for node_id, centroid in node_centroids.items():
+            # Convert node_id to regular int if it's numpy
+            if hasattr(node_id, 'item'):
+                node_id = node_id.item()
+            
+            try:
+                # Get community for this node
+                community_id = node_community[node_id]
+                
+                # Convert community_id to regular int if it's numpy
+                if hasattr(community_id, 'item'):
+                    community_id = community_id.item()
+                    
+                # Get intensity for this community
+                if community_id in community_intensity_clean:
+                    node_to_community_intensity[node_id] = community_intensity_clean[community_id]
+            except KeyError:
+                # Skip nodes that don't have community assignments
+                pass
+        
+        # Create colormap function (RdBu_r - red for high, blue for low, yellow/white for middle)
+        def intensity_to_rgb(intensity, min_val, max_val):
+            """Convert intensity value to RGB using RdBu_r colormap logic"""
+            if max_val == min_val:
+                # All same value, use neutral color
+                return np.array([255, 255, 255], dtype=np.uint8)  # White
+            
+            # Normalize to -1 to 1 range (like RdBu_r colormap)
+            normalized = 2 * (intensity - min_val) / (max_val - min_val) - 1
+            normalized = np.clip(normalized, -1, 1)
+            
+            if normalized > 0:
+                # Positive values: white to red
+                r = 255
+                g = int(255 * (1 - normalized))
+                b = int(255 * (1 - normalized))
+            else:
+                # Negative values: white to blue
+                r = int(255 * (1 + normalized))
+                g = int(255 * (1 + normalized))
+                b = 255
+            
+            return np.array([r, g, b], dtype=np.uint8)
         
-        # Extract coordinates (assuming [Z, Y, X] format)
-        z_coords = positions[:, 0]
-        y_coords = positions[:, 1]
-        x_coords = positions[:, 2]
+        # Create lookup table for RGB colors
+        max_label = max(max(labeled_array.flat), max(node_to_community_intensity.keys()) if node_to_community_intensity else 0)
+        color_lut = np.zeros((max_label + 1, 3), dtype=np.uint8)  # Default to black (0,0,0)
         
-        # Create scatter plot
-        scatter = ax.scatter(x_coords, y_coords, z_coords, 
-                           c=intensities, s=point_size, alpha=alpha,
-                           cmap='RdBu_r', vmin=min_intensity, vmax=max_intensity)
+        # Fill lookup table with RGB colors based on community intensity
+        for node_id, intensity in node_to_community_intensity.items():
+            rgb_color = intensity_to_rgb(intensity, min_intensity, max_intensity)
+            color_lut[int(node_id)] = rgb_color
         
-        ax.set_xlabel('X')
-        ax.set_ylabel('Y')
-        ax.set_zlabel('Z')
-        ax.set_title('3D Community Intensity Heatmap')
+        # Apply lookup table to labeled array - single vectorized operation
+        if is_3d:
+            # Return full 3D RGB array [Z, Y, X, 3]
+            heatmap_array = color_lut[labeled_array]
+        else:
+            # Return 2D RGB array
+            if labeled_array.ndim == 3:
+                # Take middle slice for 2D representation
+                middle_slice = labeled_array.shape[0] // 2
+                heatmap_array = color_lut[labeled_array[middle_slice]]
+            else:
+                # Already 2D
+                heatmap_array = color_lut[labeled_array]
         
+        return heatmap_array
+    
     else:
-        # 2D plot (using Y, X coordinates, ignoring Z/first dimension)
-        ax = fig.add_subplot(111)
+        # Create matplotlib plot
+        fig = plt.figure(figsize=figsize)
         
-        # Extract Y, X coordinates
-        y_coords = positions[:, 1]
-        x_coords = positions[:, 2]
+        if is_3d:
+            # 3D plot
+            ax = fig.add_subplot(111, projection='3d')
+            
+            # Extract coordinates (assuming [Z, Y, X] format)
+            z_coords = positions[:, 0]
+            y_coords = positions[:, 1]
+            x_coords = positions[:, 2]
+            
+            # Create scatter plot
+            scatter = ax.scatter(x_coords, y_coords, z_coords, 
+                               c=intensities, s=point_size, alpha=alpha,
+                               cmap='RdBu_r', vmin=min_intensity, vmax=max_intensity)
+            
+            ax.set_xlabel('X')
+            ax.set_ylabel('Y')
+            ax.set_zlabel('Z')
+            ax.set_title(f'{title}')
+            
+            # Set axis limits based on shape
+            ax.set_xlim(0, shape[2])
+            ax.set_ylim(0, shape[1])
+            ax.set_zlim(0, shape[0])
+            
+        else:
+            # 2D plot (using Y, X coordinates, ignoring Z/first dimension)
+            ax = fig.add_subplot(111)
+            
+            # Extract Y, X coordinates
+            y_coords = positions[:, 1]
+            x_coords = positions[:, 2]
+            
+            # Create scatter plot
+            scatter = ax.scatter(x_coords, y_coords, 
+                               c=intensities, s=point_size, alpha=alpha,
+                               cmap='RdBu_r', vmin=min_intensity, vmax=max_intensity)
+            
+            ax.set_xlabel('X')
+            ax.set_ylabel('Y')
+            ax.set_title(f'{title}')
+            ax.grid(True, alpha=0.3)
+            
+            # Set axis limits based on shape
+            ax.set_xlim(0, shape[2])
+            ax.set_ylim(0, shape[1])
+            
+            # Set origin to top-left (invert Y-axis)
+            ax.invert_yaxis()
         
-        # Create scatter plot
-        scatter = ax.scatter(x_coords, y_coords, 
-                           c=intensities, s=point_size, alpha=alpha,
-                           cmap='RdBu_r', vmin=min_intensity, vmax=max_intensity)
+        # Add colorbar
+        cbar = plt.colorbar(scatter, ax=ax, shrink=0.8)
+        cbar.set_label(colorbar_label)
         
-        ax.set_xlabel('X')
-        ax.set_ylabel('Y')
-        ax.set_title('2D Community Intensity Heatmap')
-        ax.grid(True, alpha=0.3)
+        # Add text annotations for min/max values
+        cbar.ax.text(1.05, 0, f'Min: {min_intensity:.3f}\n(Blue)', 
+                    transform=cbar.ax.transAxes, va='bottom')
+        cbar.ax.text(1.05, 1, f'Max: {max_intensity:.3f}\n(Red)', 
+                    transform=cbar.ax.transAxes, va='top')
         
-        # Set origin to top-left (invert Y-axis)
-        ax.invert_yaxis()
+        plt.tight_layout()
+        plt.show()
+
+
+def create_node_heatmap(node_intensity, node_centroids, shape=None, is_3d=True, 
+                        labeled_array=None, figsize=(12, 8), point_size=50, alpha=0.7, 
+                        colorbar_label="Node Intensity", title="Node Clustering Intensity Heatmap"):
+    """
+    Create a 2D or 3D heatmap showing nodes colored by their individual intensities.
+    Can return either matplotlib plot or numpy array for overlay purposes.
     
-    # Add colorbar
-    cbar = plt.colorbar(scatter, ax=ax, shrink=0.8)
-    cbar.set_label(colorbar_label)
+    Parameters:
+    -----------
+    node_intensity : dict
+        Dictionary mapping node IDs to intensity values
+        Keys can be np.int64 or regular ints
+        
+    node_centroids : dict
+        Dictionary mapping node IDs to centroids
+        Centroids should be [Z, Y, X] for 3D or [1, Y, X] for pseudo-3D
+        
+    shape : tuple, optional
+        Shape of the output array in [Z, Y, X] format
+        If None, will be inferred from node_centroids
+        
+    is_3d : bool, default=True
+        If True, create 3D plot/array. If False, create 2D plot/array.
+        
+    labeled_array : np.ndarray, optional
+        If provided, returns numpy array overlay using this labeled array template
+        instead of matplotlib plot. Uses lookup table approach for efficiency.
+        
+    figsize : tuple, default=(12, 8)
+        Figure size (width, height) - only used for matplotlib
+        
+    point_size : int, default=50
+        Size of scatter plot points - only used for matplotlib
+        
+    alpha : float, default=0.7
+        Transparency of points (0-1) - only used for matplotlib
+        
+    colorbar_label : str, default="Node Intensity"
+        Label for the colorbar - only used for matplotlib
+        
+    Returns:
+    --------
+    If labeled_array is None: fig, ax (matplotlib figure and axis objects)
+    If labeled_array is provided: np.ndarray (heatmap array with intensity values)
+    """
+    import numpy as np
+    import matplotlib.pyplot as plt
+    
+    # Convert numpy int64 keys to regular ints for consistency
+    node_intensity_clean = {}
+    for k, v in node_intensity.items():
+        if hasattr(k, 'item'):  # numpy scalar
+            node_intensity_clean[k.item()] = v
+        else:
+            node_intensity_clean[k] = v
     
-    # Add text annotations for min/max values
-    cbar.ax.text(1.05, 0, f'Min: {min_intensity:.3f}\n(Blue)', 
-                transform=cbar.ax.transAxes, va='bottom')
-    cbar.ax.text(1.05, 1, f'Max: {max_intensity:.3f}\n(Red)', 
-                transform=cbar.ax.transAxes, va='top')
+    # Prepare data for plotting/array creation
+    node_positions = []
+    node_intensities = []
     
-    plt.tight_layout()
-    plt.show()
-
-
+    for node_id, centroid in node_centroids.items():
+        try:
+            # Convert node_id to regular int if it's numpy
+            if hasattr(node_id, 'item'):
+                node_id = node_id.item()
+                
+            # Get intensity for this node
+            intensity = node_intensity_clean[node_id]
+            
+            node_positions.append(centroid)
+            node_intensities.append(intensity)
+        except KeyError:
+            # Skip nodes that don't have intensity values
+            pass
+    
+    # Convert to numpy arrays
+    positions = np.array(node_positions)
+    intensities = np.array(node_intensities)
+    
+    # Determine shape if not provided
+    if shape is None:
+        if len(positions) > 0:
+            max_coords = np.max(positions, axis=0).astype(int)
+            shape = tuple(max_coords + 1)
+        else:
+            shape = (100, 100, 100) if is_3d else (1, 100, 100)
+    
+    # Determine min and max intensities for scaling
+    if len(intensities) > 0:
+        min_intensity = np.min(intensities)
+        max_intensity = np.max(intensities)
+    else:
+        min_intensity, max_intensity = 0, 1
+    
+    if labeled_array is not None:
+        # Create numpy RGB array output using labeled array and lookup table approach
+        
+        # Create mapping from node ID to intensity value (keep original float values)
+        node_to_intensity = {}
+        for node_id, centroid in node_centroids.items():
+            # Convert node_id to regular int if it's numpy
+            if hasattr(node_id, 'item'):
+                node_id = node_id.item()
+            
+            # Only include nodes that have intensity values
+            if node_id in node_intensity_clean:
+                node_to_intensity[node_id] = node_intensity_clean[node_id]
+        
+        # Create colormap function (RdBu_r - red for high, blue for low, yellow/white for middle)
+        def intensity_to_rgb(intensity, min_val, max_val):
+            """Convert intensity value to RGB using RdBu_r colormap logic"""
+            if max_val == min_val:
+                # All same value, use neutral color
+                return np.array([255, 255, 255], dtype=np.uint8)  # White
+            
+            # Normalize to -1 to 1 range (like RdBu_r colormap)
+            normalized = 2 * (intensity - min_val) / (max_val - min_val) - 1
+            normalized = np.clip(normalized, -1, 1)
+            
+            if normalized > 0:
+                # Positive values: white to red
+                r = 255
+                g = int(255 * (1 - normalized))
+                b = int(255 * (1 - normalized))
+            else:
+                # Negative values: white to blue
+                r = int(255 * (1 + normalized))
+                g = int(255 * (1 + normalized))
+                b = 255
+            
+            return np.array([r, g, b], dtype=np.uint8)
+        
+        # Create lookup table for RGB colors
+        max_label = max(max(labeled_array.flat), max(node_to_intensity.keys()) if node_to_intensity else 0)
+        color_lut = np.zeros((max_label + 1, 3), dtype=np.uint8)  # Default to black (0,0,0)
+        
+        # Fill lookup table with RGB colors based on intensity
+        for node_id, intensity in node_to_intensity.items():
+            rgb_color = intensity_to_rgb(intensity, min_intensity, max_intensity)
+            color_lut[int(node_id)] = rgb_color
+        
+        # Apply lookup table to labeled array - single vectorized operation
+        if is_3d:
+            # Return full 3D RGB array [Z, Y, X, 3]
+            heatmap_array = color_lut[labeled_array]
+        else:
+            # Return 2D RGB array
+            if labeled_array.ndim == 3:
+                # Take middle slice for 2D representation
+                middle_slice = labeled_array.shape[0] // 2
+                heatmap_array = color_lut[labeled_array[middle_slice]]
+            else:
+                # Already 2D
+                heatmap_array = color_lut[labeled_array]
+        
+        return heatmap_array
+    
+    else:
+        # Create matplotlib plot
+        fig = plt.figure(figsize=figsize)
+        
+        if is_3d:
+            # 3D plot
+            ax = fig.add_subplot(111, projection='3d')
+            
+            # Extract coordinates (assuming [Z, Y, X] format)
+            z_coords = positions[:, 0]
+            y_coords = positions[:, 1]
+            x_coords = positions[:, 2]
+            
+            # Create scatter plot
+            scatter = ax.scatter(x_coords, y_coords, z_coords, 
+                               c=intensities, s=point_size, alpha=alpha,
+                               cmap='RdBu_r', vmin=min_intensity, vmax=max_intensity)
+            
+            ax.set_xlabel('X')
+            ax.set_ylabel('Y')
+            ax.set_zlabel('Z')
+            ax.set_title(f'{title}')
+            
+            # Set axis limits based on shape
+            ax.set_xlim(0, shape[2])
+            ax.set_ylim(0, shape[1])
+            ax.set_zlim(0, shape[0])
+            
+        else:
+            # 2D plot (using Y, X coordinates, ignoring Z/first dimension)
+            ax = fig.add_subplot(111)
+            
+            # Extract Y, X coordinates
+            y_coords = positions[:, 1]
+            x_coords = positions[:, 2]
+            
+            # Create scatter plot
+            scatter = ax.scatter(x_coords, y_coords, 
+                               c=intensities, s=point_size, alpha=alpha,
+                               cmap='RdBu_r', vmin=min_intensity, vmax=max_intensity)
+            
+            ax.set_xlabel('X')
+            ax.set_ylabel('Y')
+            ax.set_title(f'{title}')
+            ax.grid(True, alpha=0.3)
+            
+            # Set axis limits based on shape
+            ax.set_xlim(0, shape[2])
+            ax.set_ylim(0, shape[1])
+            
+            # Set origin to top-left (invert Y-axis)
+            ax.invert_yaxis()
+        
+        # Add colorbar
+        cbar = plt.colorbar(scatter, ax=ax, shrink=0.8)
+        cbar.set_label(colorbar_label)
+        
+        # Add text annotations for min/max values
+        cbar.ax.text(1.05, 0, f'Min: {min_intensity:.3f}\n(Blue)', 
+                    transform=cbar.ax.transAxes, va='bottom')
+        cbar.ax.text(1.05, 1, f'Max: {max_intensity:.3f}\n(Red)', 
+                    transform=cbar.ax.transAxes, va='top')
+        
+        plt.tight_layout()
+        plt.show()
 
 # Example usage:
 if __name__ == "__main__":
@@ -350,5 +688,4 @@ if __name__ == "__main__":
     fig, ax = plot_dict_heatmap(sample_dict, sample_id_set, 
                                title="Sample Heatmap Visualization")
     
-    plt.show()
-
+    plt.show()