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

nettracer3d/proximity.py

@@ -469,68 +469,121 @@ def extract_pairwise_connections(connections):
     return output
 
 
-def average_nearest_neighbor_distances(point_centroids, root_set, compare_set, xy_scale=1.0, z_scale=1.0, num = 1):
+def average_nearest_neighbor_distances(point_centroids, root_set, compare_set, xy_scale=1.0, z_scale=1.0, num=1, do_borders=False):
     """
     Calculate the average distance between each point in root_set and its nearest neighbor in compare_set.
     
     Args:
-        point_centroids (dict): Dictionary mapping point IDs to [Z, Y, X] coordinates
-        root_set (set): Set of point IDs to find nearest neighbors for
-        compare_set (set): Set of point IDs to search for nearest neighbors in
+        point_centroids (dict): Dictionary mapping point IDs to [Z, Y, X] coordinates (when do_borders=False)
+                               OR dictionary mapping labels to border coordinates (when do_borders=True)
+        root_set (set or dict): Set of point IDs (when do_borders=False) 
+                               OR dict {label: border_coords} (when do_borders=True)
+        compare_set (set or numpy.ndarray): Set of point IDs (when do_borders=False)
+                                           OR 1D array of border coordinates (when do_borders=True)
         xy_scale (float): Scaling factor for X and Y coordinates
         z_scale (float): Scaling factor for Z coordinate
+        num (int): Number of nearest neighbors (ignored when do_borders=True, always uses 1)
+        do_borders (bool): If True, perform border-to-border distance calculation
     
     Returns:
-        float: Average distance to nearest neighbors
+        tuple: (average_distance, distances_dict)
     """
     
-    # Extract and scale coordinates for compare_set
-    compare_coords = []
-    compare_ids = list(compare_set)
-    
-    for point_id in compare_ids:
-        z, y, x = point_centroids[point_id]
-        compare_coords.append([z * z_scale, y * xy_scale, x * xy_scale])
-    
-    compare_coords = np.array(compare_coords)
-    
-    # Build KDTree for efficient nearest neighbor search
-    tree = KDTree(compare_coords)
-    
-    distances = {}
-    same_sets = root_set == compare_set
+    if do_borders:
+        # Border comparison mode
+        if not isinstance(compare_set, np.ndarray):
+            raise ValueError("When do_borders=True, compare_set must be a numpy array of coordinates")
+        
+        # Vectorized scaling for compare coordinates
+        compare_coords_scaled = compare_set.astype(float)
+        compare_coords_scaled[:, 0] *= z_scale  # Z coordinates
+        compare_coords_scaled[:, 1:] *= xy_scale  # Y and X coordinates
+        
+        distances = {}
+        
+        for label, border_coords in root_set.items():
+            if len(border_coords) == 0:
+                continue
+                
+            # Vectorized scaling for border coordinates
+            border_coords_scaled = border_coords.astype(float)
+            border_coords_scaled[:, 0] *= z_scale  # Z coordinates
+            border_coords_scaled[:, 1:] *= xy_scale  # Y and X coordinates
+            
+            # Remove overlapping coordinates to avoid distance = 0
+            # Create a set of tuples for fast membership testing
+            border_coords_set = set(map(tuple, border_coords_scaled))
+            
+            # Filter out overlapping coordinates from compare set
+            non_overlapping_mask = np.array([
+                tuple(coord) not in border_coords_set 
+                for coord in compare_coords_scaled
+            ])
+            
+            if not np.any(non_overlapping_mask):
+                # All compare coordinates overlap - skip this object or set to NaN
+                distances[label] = np.nan
+                continue
+            
+            filtered_compare_coords = compare_coords_scaled[non_overlapping_mask]
+            
+            # Build KDTree with filtered coordinates
+            tree = KDTree(filtered_compare_coords)
+            
+            # Vectorized nearest neighbor search for all border points at once
+            distances_to_all, _ = tree.query(border_coords_scaled, k=1)
+            
+            # Find minimum distance for this object
+            distances[label] = np.min(distances_to_all)
+        
+        # Calculate average excluding NaN values
+        valid_distances = [d for d in distances.values() if not np.isnan(d)]
+        avg = np.mean(valid_distances) if valid_distances else np.nan
+        return avg, distances
     
-    for root_id in root_set:
-        # Get scaled coordinates for root point
-        z, y, x = point_centroids[root_id]
-        root_coord = np.array([z * z_scale, y * xy_scale, x * xy_scale])
+    else:
+        # Original centroid comparison mode (unchanged)
+        # Extract coordinates for compare_set
+        compare_coords = np.array([point_centroids[point_id] for point_id in compare_set])
         
+        # Vectorized scaling for compare coordinates
+        compare_coords_scaled = compare_coords.astype(float)
+        compare_coords_scaled[:, 0] *= z_scale  # Z coordinates
+        compare_coords_scaled[:, 1:] *= xy_scale  # Y and X coordinates
+        
+        # Build KDTree for efficient nearest neighbor search
+        tree = KDTree(compare_coords_scaled)
+        
+        distances = {}
+        same_sets = root_set == compare_set
+        
+        # Extract and scale root coordinates all at once
+        root_coords = np.array([point_centroids[root_id] for root_id in root_set])
+        root_coords_scaled = root_coords.astype(float)
+        root_coords_scaled[:, 0] *= z_scale  # Z coordinates
+        root_coords_scaled[:, 1:] *= xy_scale  # Y and X coordinates
+        
+        # Vectorized nearest neighbor search for all root points
         if same_sets:
-            # When sets are the same, find 2 nearest neighbors and take the second one
-            # (first one would be the point itself)
-            distances_to_all, indices = tree.query(root_coord, k= (num + 1))
-
-            temp_dist = 0
-            for i in range(1, len(distances_to_all)):
-                temp_dist += distances_to_all[i]
-
-            distances[root_id] = temp_dist/(len(distances_to_all) - 1)
-
+            distances_to_all, indices = tree.query(root_coords_scaled, k=(num + 1))
+            # Remove self-matches (first column) and average the rest
+            if num == 1:
+                distances_array = distances_to_all[:, 1]  # Just take second nearest
+            else:
+                distances_array = np.mean(distances_to_all[:, 1:], axis=1)
         else:
-            # Different sets, find nearest neighbors
-            distances_to_all, _ = tree.query(root_coord, k=num)
-            temp_dist = 0
-            for val in distances_to_all:
-                temp_dist += val
-
-            distances[root_id] = temp_dist/(len(distances_to_all))
-
-    avg = np.mean(list(distances.values())) if list(distances.values()) else 0.0
-
-    
-    # Return average distance
-    return avg, distances
-
+            distances_to_all, _ = tree.query(root_coords_scaled, k=num)
+            if num == 1:
+                distances_array = distances_to_all.flatten()
+            else:
+                distances_array = np.mean(distances_to_all, axis=1)
+        
+        # Map back to root_ids
+        for i, root_id in enumerate(root_set):
+            distances[root_id] = distances_array[i]
+        
+        avg = np.mean(distances_array) if len(distances_array) > 0 else 0.0
+        return avg, distances
 
 
 #voronois: