Add unit tests for colony collection tasks (#34)

* Remove hello world

* Bump mypy (1.3.0 -> 1.11.0)

* Add docstring and unit test for calculate centrality measures task

* Add docstring and unit test for calculate degree measures task

* Add docstring and unit test for calculate distance measures task

* Add docstring and unit test for convert to network task

* Add docstring and unit test for get depth map task

* Add docstring and unit tests for get neighbors map task

* Add docstring and unit tests for make voxels array task

pyproject.toml

@@ -20,7 +20,7 @@ networkx = "^3.0"
 [tool.poetry.group.dev.dependencies]
 black = "^24.3.0"
 isort = "^5.12.0"
-mypy = "^1.3.0"
+mypy = "^1.10.0"
 pylint = "^2.16.2"
 pytest = "^7.3.0"
 pytest-cov = "^4.0.0"

src/abm_colony_collection/__main__.py

@@ -1,2 +0,0 @@
-if __name__ == "__main__":
-    print("hello world")

src/abm_colony_collection/calculate_centrality_measures.py

@@ -3,6 +3,28 @@
 
 
 def calculate_centrality_measures(network: nx.Graph) -> pd.DataFrame:
+    """
+    Calculate centrality measures for each node in network.
+
+    Measures include:
+
+    - Degree centrality = quantifies how many neighbors a node has
+    - Closeness centrality = quantifies how close a node is to all other nodes
+      in the network
+    - Betweenness centrality = quantifies the extent to which a vertex lies on
+      shortest paths between other vertices
+
+    Parameters
+    ----------
+    network
+        The network object.
+
+    Returns
+    -------
+    :
+        Centrality measures for each node in the network.
+    """
+
     # Calculate different centrality measures for network.
     degree_centralities = nx.degree_centrality(network)
     closeness_centralities = nx.closeness_centrality(network)

src/abm_colony_collection/calculate_degree_measures.py

@@ -3,6 +3,24 @@
 
 
 def calculate_degree_measures(network: nx.Graph) -> pd.DataFrame:
+    """
+    Calculate degree measures for each node in network.
+
+    Measures include:
+
+    - Degree = number of edges adjacent to the node
+
+    Parameters
+    ----------
+    network
+        The network object.
+
+    Returns
+    -------
+    :
+        Degree measures for each node in the network.
+    """
+
     # Extract degree for each node in network.
     measures = [
         {

src/abm_colony_collection/calculate_distance_measures.py

@@ -5,6 +5,24 @@
 
 
 def calculate_distance_measures(network: nx.Graph) -> pd.DataFrame:
+    """
+    Calculate distance measures for each node in network.
+
+    Measures include:
+
+    - Eccentricity = maximum distance from node to all other nodes
+
+    Parameters
+    ----------
+    network
+        The network object.
+
+    Returns
+    -------
+    :
+        Distance measures for each node in the network.
+    """
+
     measures: list[dict[str, Union[int, float]]] = []
 
     for component in nx.connected_components(network):

src/abm_colony_collection/convert_to_network.py

@@ -3,6 +3,20 @@
 
 
 def convert_to_network(neighbors: pd.DataFrame) -> nx.Graph:
+    """
+    Convert lists of neighbors to a network object.
+
+    Parameters
+    ----------
+    neighbors
+        Lists of neighbors for each node id.
+
+    Returns
+    -------
+    :
+        The network object.
+    """
+
     nodes = list(neighbors["ID"].values)
     edges = [
         (node_id, neighbor_id)

src/abm_colony_collection/get_depth_map.py

@@ -4,9 +4,33 @@
 
 
 def get_depth_map(array: np.ndarray, neighbors_map: dict) -> dict:
+    """
+    Get map of id to depth starting from depth = 1 at the edge if the region.
+
+    All ids at the edge of the array are assigned a depth of 1. Immediate
+    neighbors of those edges are assigned a depth of 2, and so on, until all ids
+    have an assigned depth.
+
+    Parameters
+    ----------
+    array
+        Segmentation array.
+    neighbors_map
+        Map of ids to lists of neighbors.
+
+    Returns
+    -------
+    :
+        Map of id to depth from edge.
+    """
+
     depth_map = {cell_id: 0 for cell_id in np.unique(array)}
     depth_map.pop(0, None)
 
+    # Return empty depth map if there are no cell ids in the array
+    if not depth_map:
+        return depth_map
+
     edge_ids = find_edge_ids(array)
     visited = set(edge_ids)
     queue = edge_ids.copy()
@@ -28,6 +52,20 @@ def get_depth_map(array: np.ndarray, neighbors_map: dict) -> dict:
 
 
 def find_edge_ids(array: np.ndarray) -> list[int]:
+    """
+    Gets ids of regions closest to the edge of the array.
+
+    Parameters
+    ----------
+    array
+        Segmentation array.
+
+    Returns
+    -------
+    :
+        List of edge arrays.
+    """
+
     slice_index = np.argmax(np.count_nonzero(array, axis=(1, 2)))
     array_slice = array[slice_index, :, :]
 

src/abm_colony_collection/get_neighbors_map.py

@@ -1,11 +1,28 @@
-from typing import Optional
+from typing import Callable, Optional
 
 import numpy as np
 from scipy import ndimage
 from skimage import measure
 
 
 def get_neighbors_map(array: np.ndarray) -> dict:
+    """
+    Creates map of region ids to lists of neighbors.
+
+    Each region id is also assigned a group number, where all regions in a given
+    group are simply connected.
+
+    Parameters
+    ----------
+    array
+        Segmentation array.
+
+    Returns
+    -------
+    :
+        Map of id to group and neighbor ids.
+    """
+
     neighbors_map: dict = {cell_id: {} for cell_id in np.unique(array)}
     neighbors_map.pop(0, None)
 
@@ -16,9 +33,6 @@ def get_neighbors_map(array: np.ndarray) -> dict:
     # Label connected groups.
     labels, groups = measure.label(mask, connectivity=2, return_num=True)
 
-    # In line function that returns a filter lambda for a given id
-    voxel_filter = lambda voxel_id: lambda v: voxel_id in v
-
     for group in range(1, groups + 1):
         group_crop = get_cropped_array(array, group, labels)
         voxel_ids = [i for i in np.unique(group_crop) if i != 0]
@@ -28,7 +42,7 @@ def get_neighbors_map(array: np.ndarray) -> dict:
             voxel_crop = get_cropped_array(group_crop, voxel_id, crop_original=True)
 
             # Apply custom filter to get border locations.
-            border_mask = ndimage.generic_filter(voxel_crop, voxel_filter(voxel_id), size=3)
+            border_mask = ndimage.generic_filter(voxel_crop, _get_voxel_id_filter(voxel_id), size=3)
 
             # Find neighbors overlapping border.
             neighbor_list = np.unique(voxel_crop[border_mask == 1])
@@ -38,8 +52,28 @@ def get_neighbors_map(array: np.ndarray) -> dict:
     return neighbors_map
 
 
+def _get_voxel_id_filter(voxel_id: int) -> Callable:
+    """Create filtering lambda for given id."""
+    return lambda v: voxel_id in v
+
+
 def get_bounding_box(array: np.ndarray) -> tuple[int, int, int, int, int, int]:
-    """Finds bounding box around binary array."""
+    """
+    Find bounding box around array.
+
+    Bounds are calculated with a one-voxel border, if possible.
+
+    Parameters
+    ----------
+    array
+        Segmentation array.
+
+    Returns
+    -------
+    :
+        The bounding box (xmin, xmax, ymin, ymax, zmin, zmax) indices
+    """
+
     x, y, z = array.shape
 
     xbounds = np.any(array, axis=(1, 2))
@@ -65,6 +99,26 @@ def get_bounding_box(array: np.ndarray) -> tuple[int, int, int, int, int, int]:
 def get_cropped_array(
     array: np.ndarray, label: int, labels: Optional[np.ndarray] = None, crop_original: bool = False
 ) -> np.ndarray:
+    """
+    Crop array around label region.
+
+    Parameters
+    ----------
+    array
+        Array to crop.
+    label
+        Region label.
+    labels
+        Array of all region labels.
+    crop_original
+        True to crop the original array keeping all labels, False otherwise.
+
+    Returns
+    -------
+    :
+        Cropped array.
+    """
+
     # Set all voxels not matching label to zero.
     array_mask = array.copy()
     array_filter = labels if labels is not None else array_mask