nettracer3d 0.8.9__py3-none-any.whl → 0.9.1__py3-none-any.whl

nettracer3d/nettracer.py

@@ -500,7 +500,7 @@ def _upsample_3d_array(data, factor, original_shape):
     else:
         trimmed_rows = trimmed_planes[:, sub_before[1]:-sub_after[1], :]
     
-    # Remove columns from the beginning and end
+    # Remove    columns from the beginning and end
     if sub_dims[2] == 0:
         trimmed_array = trimmed_rows
     else:
@@ -508,6 +508,101 @@ def _upsample_3d_array(data, factor, original_shape):
     
     return trimmed_array
 
+
+def remove_branches_new(skeleton, length):
+    """Used to compensate for overly-branched skeletons resulting from the scipy 3d skeletonization algorithm"""
+    def find_coordinate_difference(arr):
+        try:
+            arr[1,1,1] = 0
+            # Find the indices of non-zero elements
+            indices = np.array(np.nonzero(arr)).T
+            
+            # Calculate the difference
+            diff = np.array([1,1,1]) - indices[0]
+            
+            return diff
+        except:
+            return None
+    
+    skeleton = np.pad(skeleton, pad_width=1, mode='constant', constant_values=0) #Add black planes over the 3d space to avoid index errors
+    image_copy = np.copy(skeleton)
+    
+    # Find all endpoints ONCE at the beginning
+    nonzero_coords = np.transpose(np.nonzero(image_copy))
+    endpoints = []
+    nubs = []
+    
+    for x, y, z in nonzero_coords:
+        mini = image_copy[x-1:x+2, y-1:y+2, z-1:z+2]
+        nearby_sum = np.sum(mini)
+        threshold = 2 * image_copy[x, y, z]
+        
+        if nearby_sum <= threshold:
+            endpoints.append((x, y, z))
+    
+    x, y, z = endpoints[0]
+    original_val = image_copy[x, y, z]
+
+    # Process each endpoint individually for nub assessment
+    for start_x, start_y, start_z in endpoints:
+            
+        # Trace the branch from this endpoint, removing points as we go
+        branch_coords = []
+        current_coord = (start_x, start_y, start_z)
+        nub_reached = False
+        
+        for step in range(length):
+            x, y, z = current_coord
+            
+            # Store original value and coordinates
+            branch_coords.append((x, y, z))
+            
+            # Remove this point temporarily
+            image_copy[x, y, z] = 0
+            
+            # If we've reached the maximum length without hitting a nub, break
+            if step == length - 1:
+                break
+            
+            # Find next coordinate in the branch
+            mini = image_copy[x-1:x+2, y-1:y+2, z-1:z+2]
+            dif = find_coordinate_difference(mini.copy())
+            if dif is None:
+                break
+                
+            next_coord = (x - dif[0], y - dif[1], z - dif[2])
+            
+            # Check if next coordinate is valid and exists
+            nx, ny, nz = next_coord
+            
+            # Check if next point is a nub (has more neighbors than expected)
+            next_mini = image_copy[nx-1:nx+2, ny-1:ny+2, nz-1:nz+2]
+            next_nearby_sum = np.sum(next_mini)
+            next_threshold = 2 * image_copy[nx, ny, nz]
+            
+            if next_nearby_sum > next_threshold:
+                nub_reached = True
+                nubs.append(next_coord)
+                nubs.append(current_coord) # Note, if we don't add the current coord here (and restore it below), the behavior of this method can be changed to trim branches beneath previous branches, which could be neat but its somewhat unpredictable so I opted out of it.
+                image_copy[x, y, z] = original_val
+                #image_copy[nx, ny, nz] = 0
+                break
+                
+            current_coord = next_coord
+        
+        # If no nub was reached, restore all the points we removed
+        if not nub_reached:
+            for i, (bx, by, bz) in enumerate(branch_coords):
+                image_copy[bx, by, bz] = original_val
+        # If nub was reached, points stay removed (branch is eliminated)
+
+    for item in nubs: #The nubs are endpoints of length = 1. They appear a bit different in the array so we just note when one is created and remove them all at the end in a batch.
+        image_copy[item[0], item[1], item[2]] = 0 # Removing the nub itself leaves a hole in the skeleton but for branchpoint detection that doesn't matter, which is why it behaves this way. To fill the hole, one option is to dilate once then erode/skeletonize again, but we want to avoid making anything that looks like local branching so I didn't bother.
+
+    # Remove padding and return
+    image_copy = (image_copy[1:-1, 1:-1, 1:-1]).astype(np.uint8)
+    return image_copy
+
 def remove_branches(skeleton, length):
     """Used to compensate for overly-branched skeletons resulting from the scipy 3d skeletonization algorithm"""
 
@@ -532,6 +627,7 @@ def remove_branches(skeleton, length):
     x, y, z = nonzero_coords[0]
     threshold = 2 * skeleton[x, y, z]
     nubs = []
+    
 
     for b in range(length):
 
@@ -628,6 +724,8 @@ def break_and_label_skeleton(skeleton, peaks = 1, branch_removal = 0, comp_dil =
     else:
         broken_skele = None
 
+    #old_skeleton = copy.deepcopy(skeleton) # The skeleton might get modified in label_vertices so we can make a preserved copy of it to use later
+
     if nodes is None:
 
         verts = label_vertices(skeleton, peaks = peaks, branch_removal = branch_removal, comp_dil = comp_dil, max_vol = max_vol, return_skele = return_skele)
@@ -637,6 +735,8 @@ def break_and_label_skeleton(skeleton, peaks = 1, branch_removal = 0, comp_dil =
 
     verts = invert_array(verts)
 
+    #skeleton = old_skeleton
+
     image_copy = skeleton * verts
 
  
@@ -1031,10 +1131,93 @@ def remove_trunk(edges, num_iterations=1):
     
     return edges
 
-def hash_inners(search_region, inner_edges, GPU = True):
+def get_all_label_coords(labeled_array, background=0):
+    """
+    Get coordinates for all labels using single pass method.
+    
+    Parameters:
+    -----------
+    labeled_array : numpy.ndarray
+        Labeled array with integer labels
+    background : int, optional
+        Background label to exclude (default: 0)
+    
+    Returns:
+    --------
+    dict : {label: coordinates_array}
+        Dictionary mapping each label to its coordinate array
+    """
+    coords_dict = {}
+    
+    # Get all non-background coordinates at once
+    all_coords = np.argwhere(labeled_array != background)
+    
+    if len(all_coords) == 0:
+        return coords_dict
+    
+    # Get the label values at those coordinates
+    labels_at_coords = labeled_array[tuple(all_coords.T)]
+    
+    # Group by label
+    unique_labels = np.unique(labels_at_coords)
+    for label in unique_labels:
+        mask = labels_at_coords == label
+        coords_dict[label] = all_coords[mask]
+    
+    return coords_dict
+
+def approx_boundaries(array, iden_set = None, node_identities = None, keep_labels = False):
+
+    """Hollows out an array, can do it for only a set number of identities. Returns coords as dict if labeled or as 1d numpy array if binary is desired"""
+
+    if node_identities is not None:
+
+        nodes = []
+
+        for node in node_identities:
+
+            if node_identities[node] in iden_set: #Filter out only idens we need
+                nodes.append(node)
+
+        mask = np.isin(array, nodes)
+
+        if keep_labels:
+
+            array = array * mask
+        else:
+            array = mask
+        del mask
+
+    from skimage.segmentation import find_boundaries
+
+    borders = find_boundaries(array, mode='thick')
+    array = array * borders
+    del borders
+    if not keep_labels:
+        return np.argwhere(array != 0)
+    else:
+        return get_all_label_coords(array)
+
+
+
+def hash_inners(search_region, inner_edges, GPU = False):
     """Internal method used to help sort out inner edge connections. The inner edges of the array will not differentiate between what nodes they contact if those nodes themselves directly touch each other.
     This method allows these elements to be efficiently seperated from each other"""
 
+    from skimage.segmentation import find_boundaries
+
+    borders = find_boundaries(search_region, mode='thick')
+
+    inner_edges = inner_edges * borders #And as a result, we can mask out only 'inner edges' that themselves exist within borders
+
+    inner_edges = dilate_3D_old(inner_edges, 3, 3, 3) #Not sure if dilating is necessary. Want to ensure that the inner edge pieces still overlap with the proper nodes after the masking.
+
+    return inner_edges
+
+def hash_inners_old(search_region, inner_edges, GPU = True):
+    """Internal method used to help sort out inner edge connections. The inner edges of the array will not differentiate between what nodes they contact if those nodes themselves directly touch each other.
+    This method allows these elements to be efficiently seperated from each other. Originally this was implemented using the gaussian blur because i didn't yet realize skimage could do the same more efficiently."""
+
     print("Performing gaussian blur to hash inner edges.")
 
     blurred_search = smart_dilate.gaussian(search_region, GPU = GPU) 
@@ -2170,6 +2353,9 @@ def label_vertices(array, peaks = 0, branch_removal = 0, comp_dil = 0, max_vol =
 
     array = skeletonize(array)
 
+    if return_skele:
+        old_skeleton = copy.deepcopy(array) # The skeleton might get modified in label_vertices so we can make a preserved copy of it to use later
+
     if branch_removal > 0:
         array = remove_branches(array, branch_removal)
 
@@ -2235,7 +2421,7 @@ def label_vertices(array, peaks = 0, branch_removal = 0, comp_dil = 0, max_vol =
 
     if return_skele:
 
-        return labeled_image, (array[1:-1, 1:-1, 1:-1]).astype(np.uint8)
+        return labeled_image, old_skeleton
 
     else:
 
@@ -5306,7 +5492,7 @@ class Network_3D:
         neighborhoods.visualize_cluster_composition_umap(self.node_centroids, None, id_dictionary = self.node_identities, graph_label = "Node ID", title = 'UMAP Visualization of Node Centroids') 
 
 
-    def community_id_info_per_com(self, umap = False, label = False, limit = 0, proportional = False):
+    def community_id_info_per_com(self, umap = False, label = 0, limit = 0, proportional = False, neighbors = None):
 
         community_dict = invert_dict(self.communities)
         summation = 0
@@ -5395,7 +5581,19 @@ class Network_3D:
         if umap:
             from . import neighborhoods
 
-            neighborhoods.visualize_cluster_composition_umap(umap_dict, id_set, label = label) 
+
+            if self.communities is not None and label == 2:
+                neighbor_group = {}
+                for node, com in self.communities.items():
+                    neighbor_group[com] = neighbors[node]
+                neighborhoods.visualize_cluster_composition_umap(umap_dict, id_set, neighborhoods = neighbor_group)
+            elif label == 1:
+                neighborhoods.visualize_cluster_composition_umap(umap_dict, id_set, label = True) 
+            else:
+                neighborhoods.visualize_cluster_composition_umap(umap_dict, id_set, label = False) 
+
+
+            #neighborhoods.visualize_cluster_composition_umap(umap_dict, id_set, label = label) 
 
         return output, id_set
 
@@ -5668,75 +5866,141 @@ class Network_3D:
                     pass
 
 
-    def nearest_neighbors_avg(self, root, targ, xy_scale = 1, z_scale = 1, num = 1, heatmap = False, threed = True, numpy = False, quant = False):
+    def nearest_neighbors_avg(self, root, targ, xy_scale = 1, z_scale = 1, num = 1, heatmap = False, threed = True, numpy = False, quant = False, centroids = True):
 
-        def get_theoretical_nearest_neighbor_distance(compare_set, num_neighbors, volume, is_2d=False):
-            """
-            Calculate theoretical expected distance to k-th nearest neighbor 
-            assuming random uniform distribution in 2D or 3D space.
-            """
-            import math
+        def distribute_points_uniformly(n, shape, z_scale, xy_scale, num, is_2d=False):
+
+            from scipy.spatial import KDTree
+
+            if n <= 1:
+                return 0
+                        
+            # Calculate total positions and sampling step
+            total_positions = np.prod(shape)
+            if n >= total_positions:
+                # If we want more points than positions, just return scaled unit distance
+                return xy_scale if is_2d else min(z_scale, xy_scale)
             
-            if len(compare_set) == 0 or volume <= 0:
-                raise ValueError("Invalid input: empty set or non-positive volume")
+            # Create uniformly spaced indices
+            indices = np.linspace(0, total_positions - 1, n, dtype=int)
             
-            density = len(compare_set) / volume
-            k = num_neighbors
+            # Convert flat indices to coordinates
+            coords = []
+            for idx in indices:
+                coord = np.unravel_index(idx, shape)
+                if len(shape) == 3:
+                    # Apply scaling: [z, y, x] with respective scales
+                    scaled_coord = [coord[0] * z_scale, coord[1] * xy_scale, coord[2] * xy_scale]
+                elif len(shape) == 2:
+                    # Apply scaling: [y, x] with xy_scale
+                    scaled_coord = [coord[0] * xy_scale, coord[1] * xy_scale]
+                coords.append(scaled_coord)
             
-            if is_2d:
-                # Expected distance to k-th nearest neighbor in 2D
-                # μ1' = Γ(k + 1/2) / (Γ(k) × √(m × π))
-                expected_distance = math.gamma(k + 0.5) / (math.gamma(k) * math.sqrt(density * math.pi))
-            else:
-                # Expected distance to k-th nearest neighbor in 3D
-                # μ1' = Γ(k + 1/3) / (Γ(k) × (m × Φ)^(1/3))
-                # where Φ = π^(3/2) / Γ(3/2 + 1) = π^(3/2) / Γ(5/2) = 4π/3
-                phi_3d = 4 * math.pi / 3  # Volume of unit sphere in 3D
-                expected_distance = math.gamma(k + 1/3) / (math.gamma(k) * (density * phi_3d)**(1/3))
+            coords = np.array(coords)
+            
+            # Build KDTree
+            tree = KDTree(coords)
+            
+            # Pick a point near the middle of the array
+            middle_idx = len(coords) // 2
+            query_point = coords[middle_idx]
+            
+            # Find the num+1 nearest neighbors (including the point itself)
+            distances, indices = tree.query(query_point, k=num+1)
+            
+            # Exclude the point itself (distance 0) and get the actual neighbors
+            neighbor_distances = distances[1:num+1]
+
+            if num == n:
+                neighbor_distances[-1] = neighbor_distances[-2]
+            
+            avg_distance = np.mean(neighbor_distances)
             
-            return expected_distance
+            
+            return avg_distance
 
-        root_set = []
+        do_borders = not centroids
 
-        compare_set = []
+        if centroids:
+            root_set = []
 
-        if root is None:
+            compare_set = []
 
-            root_set = list(self.node_centroids.keys())
-            compare_set = root_set
-            title = "Nearest Neighbors Between Nodes Heatmap"
+            if root is None:
 
-        else:
+                root_set = list(self.node_centroids.keys())
+                compare_set = root_set
+                title = "Nearest Neighbors Between Nodes Heatmap"
+
+            else:
 
-            title = f"Nearest Neighbors of ID {targ} from ID {root} Heatmap"
+                title = f"Nearest Neighbors of ID {targ} from ID {root} Heatmap"
 
-            for node, iden in self.node_identities.items():
+                for node, iden in self.node_identities.items():
 
-                if iden == root:
+                    if iden == root:
 
-                    root_set.append(node)
+                        root_set.append(node)
 
-                elif (iden == targ) or (targ == 'All Others (Excluding Self)'):
+                    elif (iden == targ) or (targ == 'All Others (Excluding Self)'):
 
-                    compare_set.append(node)
+                        compare_set.append(node)
 
-        if root == targ:
+            if root == targ:
+
+                compare_set = root_set
+                if len(compare_set) - 1 < num:
+
+                    num = len(compare_set) - 1
+
+                    print(f"Error: Not enough neighbor nodes for requested number of neighbors. Using max available neighbors: {num}")
+                    
 
-            compare_set = root_set
-            if len(compare_set) - 1 < num:
+            if len(compare_set) < num:
 
-                num = len(compare_set) - 1
+                num = len(compare_set)
 
                 print(f"Error: Not enough neighbor nodes for requested number of neighbors. Using max available neighbors: {num}")
-                
 
-        if len(compare_set) < num:
+            avg, output = proximity.average_nearest_neighbor_distances(self.node_centroids, root_set, compare_set, xy_scale=self.xy_scale, z_scale=self.z_scale, num = num, do_borders = do_borders)
 
-            num = len(compare_set)
+        else:
+            if heatmap:
+                root_set = []
+                compare_set = []
+                if root is None:
 
-            print(f"Error: Not enough neighbor nodes for requested number of neighbors. Using max available neighbors: {num}")
-            
-        avg, output = proximity.average_nearest_neighbor_distances(self.node_centroids, root_set, compare_set, xy_scale=self.xy_scale, z_scale=self.z_scale, num = num)
+                    root_set = list(self.node_centroids.keys())
+                    compare_set = root_set
+                else:
+                    for node, iden in self.node_identities.items():
+
+                        if iden == root:
+
+                            root_set.append(node)
+
+                        elif (iden == targ) or (targ == 'All Others (Excluding Self)'):
+
+                            compare_set.append(node)
+
+            if root is None:
+                title = "Nearest Neighbors Between Nodes Heatmap"
+                root_set_neigh = approx_boundaries(self.nodes, keep_labels = True)
+                compare_set_neigh = approx_boundaries(self.nodes, keep_labels = False)
+            else:
+                title = f"Nearest Neighbors of ID {targ} from ID {root} Heatmap"
+
+                root_set_neigh = approx_boundaries(self.nodes, [root], self.node_identities, keep_labels = True)
+
+                if targ == 'All Others (Excluding Self)':
+                    compare_set_neigh = set(self.node_identities.values())
+                    compare_set_neigh.remove(root)
+                    targ = compare_set_neigh
+                else:
+                    targ = [targ]
+
+                compare_set_neigh = approx_boundaries(self.nodes, targ, self.node_identities, keep_labels = False)
+                avg, output = proximity.average_nearest_neighbor_distances(self.node_centroids, root_set_neigh, compare_set_neigh, xy_scale=self.xy_scale, z_scale=self.z_scale, num = num, do_borders = do_borders)
 
         if quant:
             try:
@@ -5778,8 +6042,7 @@ class Network_3D:
             else:
                 is_2d = False
 
-            pred = get_theoretical_nearest_neighbor_distance(compare_set, num, volume, is_2d = is_2d)
-            #pred = avg
+            pred = distribute_points_uniformly(len(compare_set), bounds, self.z_scale, self.xy_scale, num = num, is_2d = is_2d)
 
             node_intensity = {}
             import math
@@ -5787,18 +6050,22 @@ class Network_3D:
 
             for node in root_set:
                 node_intensity[node] = math.log(pred/output[node])
+                #print(output[node])
                 node_centroids[node] = self.node_centroids[node]
 
             if numpy:
 
                 overlay = neighborhoods.create_node_heatmap(node_intensity, node_centroids, shape = shape, is_3d=threed, labeled_array = self.nodes, colorbar_label="Clustering Intensity", title = title)
 
-                return avg, output, overlay, quant_overlay
+                return avg, output, overlay, quant_overlay, pred
 
             else:
                 neighborhoods.create_node_heatmap(node_intensity, node_centroids, shape = shape, is_3d=threed, labeled_array = None, colorbar_label="Clustering Intensity", title = title)
 
-        return avg, output, quant_overlay
+        else:
+            pred = None
+
+        return avg, output, quant_overlay, pred