nettracer3d 1.2.5__py3-none-any.whl → 1.3.1__py3-none-any.whl

nettracer3d/network_analysis.py

@@ -601,42 +601,35 @@ def _find_centroids_old(nodes, node_list = None, down_factor = None):
 
 
 def _find_centroids(nodes, node_list=None, down_factor=None):
-    """Internal use version to get centroids without saving"""
-    
+    """Parallel version using sum accumulation instead of storing coordinates"""
     
-    def compute_indices_in_chunk(chunk, y_offset):
-        """
-        Alternative approach using np.where for even better performance on sparse arrays.
-        """
-        indices_dict_chunk = {}
+    def compute_sums_in_chunk(chunk, y_offset):
+        """Accumulate sums and counts - much less memory than storing coords"""
+        sums_dict = {}
+        counts_dict = {}
         
-        # Get all coordinates where chunk is non-zero
         z_coords, y_coords, x_coords = np.where(chunk != 0)
         
         if len(z_coords) == 0:
-            return indices_dict_chunk
+            return sums_dict, counts_dict
         
-        # Adjust Y coordinates
         y_coords_adjusted = y_coords + y_offset
-        
-        # Get labels at these coordinates
         labels = chunk[z_coords, y_coords, x_coords]
-        
-        # Group by unique labels
         unique_labels = np.unique(labels)
         
         for label in unique_labels:
-            if label == 0:  # Skip background
+            if label == 0:
                 continue
             mask = (labels == label)
-            # Stack coordinates into the expected format [z, y, x]
-            indices_dict_chunk[label] = np.column_stack((
-                z_coords[mask], 
-                y_coords_adjusted[mask], 
-                x_coords[mask]
-            ))
+            # Just store sums and counts - O(1) memory per label
+            sums_dict[label] = np.array([
+                z_coords[mask].sum(dtype=np.float64),
+                y_coords_adjusted[mask].sum(dtype=np.float64),
+                x_coords[mask].sum(dtype=np.float64)
+            ])
+            counts_dict[label] = mask.sum()
         
-        return indices_dict_chunk
+        return sums_dict, counts_dict
     
     def chunk_3d_array(array, num_chunks):
         """Split the 3D array into smaller chunks along the y-axis."""
@@ -644,49 +637,44 @@ def _find_centroids(nodes, node_list=None, down_factor=None):
         return y_slices
     
     # Handle input processing
-    if isinstance(nodes, str):  # Open into numpy array if filepath
+    if isinstance(nodes, str):
         nodes = tifffile.imread(nodes)
-        if len(np.unique(nodes)) == 2:  # Label if binary
+        if len(np.unique(nodes)) == 2:
             structure_3d = np.ones((3, 3, 3), dtype=int)
             nodes, num_nodes = ndimage.label(nodes)
     
     if down_factor is not None:
         nodes = downsample(nodes, down_factor)
-    else:
-        down_factor = 1
     
-    indices_dict = {}
+    sums_total = {}
+    counts_total = {}
     num_cpus = mp.cpu_count()
     
-    # Chunk the 3D array along the y-axis
     node_chunks = chunk_3d_array(nodes, num_cpus)
-    
-    # Calculate Y offset for each chunk
     chunk_sizes = [chunk.shape[1] for chunk in node_chunks]
     y_offsets = np.cumsum([0] + chunk_sizes[:-1])
     
-    # Parallel computation using the optimized single-pass approach
     with ThreadPoolExecutor(max_workers=num_cpus) as executor:
-        futures = {executor.submit(compute_indices_in_chunk, chunk, y_offset): chunk_id
-                  for chunk_id, (chunk, y_offset) in enumerate(zip(node_chunks, y_offsets))}
+        futures = [executor.submit(compute_sums_in_chunk, chunk, y_offset)
+                  for chunk, y_offset in zip(node_chunks, y_offsets)]
         
         for future in as_completed(futures):
-            indices_chunk = future.result()
-            # Merge indices for each label
-            for label, indices in indices_chunk.items():
-                if label in indices_dict:
-                    indices_dict[label] = np.vstack((indices_dict[label], indices))
+            sums_chunk, counts_chunk = future.result()
+            
+            # Merge is now just addition - O(1) instead of vstack
+            for label in sums_chunk:
+                if label in sums_total:
+                    sums_total[label] += sums_chunk[label]
+                    counts_total[label] += counts_chunk[label]
                 else:
-                    indices_dict[label] = indices
-    
-    # Compute centroids from collected indices
-    centroid_dict = {}
-    for label, indices in indices_dict.items():
-        centroid = np.round(np.mean(indices, axis=0)).astype(int)
-        centroid_dict[label] = centroid
+                    sums_total[label] = sums_chunk[label]
+                    counts_total[label] = counts_chunk[label]
     
-    # Remove background label if it exists
-    centroid_dict.pop(0, None)
+    # Compute centroids from accumulated sums
+    centroid_dict = {
+        label: np.round(sums_total[label] / counts_total[label]).astype(int)
+        for label in sums_total if label != 0
+    }
     
     return centroid_dict
 
@@ -935,8 +923,15 @@ def get_distance_list(centroids, network, xy_scale, z_scale):
     return distance_list
 
 
-def prune_samenode_connections(networkfile, nodeIDs):
-    """Even faster numpy-based version for very large datasets"""
+def prune_samenode_connections(networkfile, nodeIDs, target=None):
+    """Even faster numpy-based version for very large datasets
+    
+    Args:
+        networkfile: Network file path or list of node pairs
+        nodeIDs: Node identity mapping (file path or dict)
+        target: Optional string. If provided, only prunes pairs where BOTH nodes 
+                have this specific identity. If None, prunes all same-identity pairs.
+    """
     import numpy as np
     
     # Handle nodeIDs input
@@ -965,8 +960,13 @@ def prune_samenode_connections(networkfile, nodeIDs):
     idsA = np.array([data_dict.get(node) for node in nodesA])
     idsB = np.array([data_dict.get(node) for node in nodesB])
     
-    # Create boolean mask - keep where IDs are different
-    keep_mask = idsA != idsB
+    # Create boolean mask based on target parameter
+    if target is None:
+        # Original behavior: keep where IDs are different
+        keep_mask = idsA != idsB
+    else:
+        # New behavior: only remove pairs where BOTH nodes have the target identity
+        keep_mask = ~((idsA == target) & (idsB == target))
     
     # Apply filter
     filtered_nodesA = nodesA[keep_mask].tolist()

nettracer3d/network_draw.py

@@ -264,20 +264,20 @@ def draw_network_from_centroids(nodes, network, centroids, twod_bool, directory
 
     if twod_bool:
         output_stack = output_stack[0,:,:] | output_stack[0,:,:]
-
-    if directory is None:
-        try:
-            tifffile.imwrite("drawn_network.tif", output_stack)
-        except Exception as e:
-            print("Could not save network lattice to active directory")
-            print("Network lattice saved as drawn_network.tif")
-
-    if directory is not None:
-        try:
-            tifffile.imwrite(f"{directory}/drawn_network.tif", output_stack)
-            print(f"Network lattice saved to {directory}/drawn_network.tif")
-        except Exception as e:
-            print(f"Could not save network lattice to {directory}")
+    
+    #if directory is None:
+     #   try:
+      #      tifffile.imwrite("drawn_network.tif", output_stack)
+       # except Exception as e:
+        #    print("Could not save network lattice to active directory")
+         #   print("Network lattice saved as drawn_network.tif")
+
+    #if directory is not None:
+     #   try:
+      #      tifffile.imwrite(f"{directory}/drawn_network.tif", output_stack)
+       #     print(f"Network lattice saved to {directory}/drawn_network.tif")
+        #except Exception as e:
+         #   print(f"Could not save network lattice to {directory}")
 
     return output_stack
 
@@ -340,6 +340,7 @@ def draw_network_from_centroids_GPU(nodes, network, centroids, twod_bool, direct
 
     output_stack = cp.asnumpy(output_stack)
 
+    """
     if directory is None:
         try:
             tifffile.imwrite("drawn_network.tif", output_stack)
@@ -353,7 +354,7 @@ def draw_network_from_centroids_GPU(nodes, network, centroids, twod_bool, direct
             print(f"Network lattice saved to {directory}/drawn_network.tif")
         except Exception as e:
             print(f"Could not save network lattice to {directory}")
-
+    """
 
 if __name__ == '__main__':