doc: skeleton graph

Author: anna-grim

src/segmentation_skeleton_metrics/skeleton_graph.py

@@ -20,29 +20,30 @@
 
 class SkeletonGraph(nx.Graph):
     """
-    A subclass of the NetworkX.Graph that represents a skeleton graph with
-    additional functionality for handling labels and voxel coordinates
-    corresponding to the nodes. Note that node IDs directly index into the
-    "labels" and "voxels" attributes.
+    A subclass of the NetworkX.Graph designed for graphs built from SWC files.
+    This class extends the functionality of the standard Graph by adding
+    support for handling node labels and voxel coordinates. In this subclass,
+    node IDs serve as direct indices for accessing the labels and voxels
+    attributes.
 
     Attributes
     ----------
-    anisotropy : ArrayLike
+    anisotropy : np.ndarray
         Image to physical coordinates scaling factors to account for the
         anisotropy of the microscope.
     run_length : float
         Physical path length of the graph.
     labels : numpy.ndarray
         A 1D array that contains a label value associated with each node.
     voxels : numpy.ndarray
-        A 3D array that contains a voxel coordinate of each node.
+        A 3D array that contains a voxel coordinate for each node.
 
     """
 
     def __init__(self, anisotropy=(1.0, 1.0, 1.0)):
         """
         Initializes a SkeletonGraph, including setting the anisotropy and
-        initializing the run length attribute.
+        initializing the run length attributes.
 
         Parameters
         ----------
@@ -78,6 +79,8 @@ def init_labels(self):
         None
 
         """
+        error_msg = "Graph must have nodes to initialize labels!"
+        assert self.number_of_nodes() > 0, error_msg
         self.labels = np.zeros((self.number_of_nodes()), dtype=int)
 
     def init_voxels(self, voxels):
@@ -86,7 +89,8 @@ def init_voxels(self, voxels):
 
         Parameters
         ----------
-        None
+        voxels : ArrayLike
+            Voxel coordinates for each node in the graph.
 
         Returns
         -------
@@ -112,21 +116,20 @@ def set_filename(self, filename):
         """
         self.filename = filename
 
-    def set_nodes(self):
+    def set_nodes(self, num_nodes):
         """
-        Adds nodes to the graph. The nodes are assigned indices from 0 to the
-        total number of voxels in the image.
+        Adds nodes to the graph. The nodes are assigned indices from 0 to
+        "num_nodes".
 
         Parameters
         ----------
-        None
+        num_nodes
 
         Returns
         -------
         None
 
         """
-        num_nodes = len(self.voxels)
         self.add_nodes_from(np.arange(num_nodes))
 
     # --- Getters ---
@@ -140,9 +143,8 @@ def get_labels(self):
 
         Returns
         -------
-        numpy.ndarray
-            A 1D array of unique non-zero labels assigned to nodes in the
-            graph.
+        Set[int]
+            Unique non-zero label values assigned to nodes in the graph.
 
         """
         labels = set(np.unique(self.labels))
@@ -202,7 +204,8 @@ def physical_dist(self, i, j):
         Returns
         -------
         float
-            Euclidea distance between physical coordinates of the given nodes.
+            Euclidean distance between physical coordinates of the given
+            nodes.
 
         """
         xyz_i = self.voxels[i] * self.anisotropy
@@ -211,13 +214,13 @@ def physical_dist(self, i, j):
 
     def get_bbox(self, nodes):
         """
-        Calculates the bounding box that contains the voxel coordinates for a
-        given collection of nodes.
+        Calculates the minimal bounding box containing the voxel coordinates
+        for a given collection of nodes.
 
         Parameters
         ----------
         nodes : Container
-            A collection of node indices for which to compute the bounding box.
+            Node indices for which to compute the bounding box.
 
         Returns
         -------
@@ -281,8 +284,8 @@ def run_length_from(self, root):
 
         Parameters
         ----------
-        graph : networkx.Graph
-            Graph to be parsed.
+        root : int
+            Node contained in connected component to compute run length of.
 
         Returns
         -------
@@ -297,14 +300,14 @@ def run_length_from(self, root):
 
     def upd_labels(self, nodes, label):
         """
-        Updates the label of each node in "nodes" with "label".
+        Updates the label of the given nodes with a specified label.
 
         Parameters
         ----------
         nodes : List[int]
             Nodes to be updated.
         label : int
-            Updated label.
+            New label of nodes.
 
         Returns
         -------
@@ -316,14 +319,16 @@ def upd_labels(self, nodes, label):
 
     def to_zipped_swc(self, zip_writer, color=None):
         """
-        Writes a graph to an SWC file that is to be stored in a zip.
+        Writes the graph to an SWC file format, which is then stored in a ZIP
+        archive.
 
         Parameters
         ----------
         zip_writer : zipfile.ZipFile
-            ...
+            A ZipFile object that will store the generated SWC file.
         color : str, optional
-            ...
+            A string representing the color (e.g., "[1.0 0.0 0.0]") of the SWC
+            file. The default is None.
 
         Returns
         -------
@@ -336,24 +341,21 @@ def to_zipped_swc(self, zip_writer, color=None):
             text_buffer.write("# id, type, z, y, x, r, pid")
 
             # Write entries
-            n_entries = 0
             node_to_idx = dict()
-            r = 5 if color else 3
+            r = 6 if color else 3
             for i, j in nx.dfs_edges(self):
                 # Special Case: Root
                 x, y, z = tuple(self.voxels[i] * self.anisotropy)
                 if len(node_to_idx) == 0:
                     parent = -1
                     node_to_idx[i] = 1
                     text_buffer.write("\n" + f"1 2 {x} {y} {z} {r} {parent}")
-                    n_entries += 1
 
                 # General Case
-                node = n_entries + 1
+                node = len(node_to_idx) + 1
                 parent = node_to_idx[i]
-                node_to_idx[j] = n_entries + 1
+                node_to_idx[j] = node
                 text_buffer.write("\n" + f"{node} 2 {x} {y} {z} {r} {parent}")
-                n_entries += 1
 
             # Finish
             zip_writer.writestr(self.filename, text_buffer.getvalue())

src/segmentation_skeleton_metrics/skeleton_metric.py

@@ -537,7 +537,6 @@ def adjust_metrics(self, key):
                     # Update graph
                     self.graphs[key].remove_nodes_from(nodes)
                     visited.add(label)
-                    print("# nodes deleted:", len(nodes))
 
     def find_label_intersections(self):
         """

src/segmentation_skeleton_metrics/utils/graph_util.py

@@ -89,7 +89,7 @@ def to_graph(self, swc_dict):
         graph = SkeletonGraph(anisotropy=self.anisotropy)
         graph.init_voxels(swc_dict["voxel"])
         graph.set_filename(swc_dict["swc_id"] + ".swc")
-        graph.set_nodes()
+        graph.set_nodes(len(swc_dict["id"]))
 
         # Build graph
         id_lookup = dict()