Add docstring and unit tests for make voxels array task

src/abm_colony_collection/get_depth_map.py

@@ -27,6 +27,10 @@ def get_depth_map(array: np.ndarray, neighbors_map: dict) -> dict:
     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()

src/abm_colony_collection/make_voxels_array.py

@@ -2,6 +2,20 @@
 
 
 def make_voxels_array(locations: list) -> np.ndarray:
+    """
+    Convert list of locations to a segmentated image array.
+
+    Parameters
+    ----------
+    locations
+        List of location dictionaries containing id and voxels.
+
+    Returns
+    -------
+    :
+        Segmentation array.
+    """
+
     # Extract all voxel positions with id.
     all_ids: list[int] = []
     all_xyz: list[tuple[int, int, int]] = []
@@ -11,16 +25,16 @@ def make_voxels_array(locations: list) -> np.ndarray:
         all_xyz = all_xyz + xyz
         all_ids = all_ids + [cell_id] * len(xyz)
 
+    # Return if no voxels.
+    if len(all_ids) == 0:
+        return np.zeros((1, 1, 1))
+
     # Create empty array.
     mins = np.min(all_xyz, axis=0)
     maxs = np.max(all_xyz, axis=0)
     length, width, height = np.subtract(maxs, mins) + 3
     array = np.zeros((height, width, length), dtype=np.uint16)
 
-    # Return if no voxels.
-    if len(all_ids) == 0:
-        return array
-
     # Fill voxel array.
     all_xyz_offset = [(z - mins[2] + 1, y - mins[1] + 1, x - mins[0] + 1) for x, y, z in all_xyz]
     array[tuple(np.transpose(all_xyz_offset))] = all_ids

tests/abm_colony_collection/test_get_depth_map.py

@@ -52,6 +52,11 @@ def test_get_depth_map(self):
         depth_map = get_depth_map(self.array, self.neighbors_map)
         self.assertDictEqual(expected_depth_map, depth_map)
 
+    def test_get_depth_map_empty_array(self):
+        array = np.zeros((1, 1, 1))
+        depth_map = get_depth_map(array, {})
+        self.assertDictEqual({}, depth_map)
+
     def test_find_edge_ids(self):
         expected_edge_ids = [5, 6, 7, 8]
         edges_ids = find_edge_ids(self.array)

tests/abm_colony_collection/test_get_neighbors_map.py

@@ -36,7 +36,6 @@ def setUp(self) -> None:
         self.array = array
 
     def test_get_neighbors_map(self):
-
         expected_neighbors_map = {
             1: {"group": 1, "neighbors": [2, 3, 9]},
             2: {"group": 1, "neighbors": [1, 4, 9]},
@@ -54,6 +53,11 @@ def test_get_neighbors_map(self):
 
         self.assertDictEqual(expected_neighbors_map, neighbors_map)
 
+    def test_get_neighbors_map_empty_array(self):
+        array = np.zeros((1, 1, 1))
+        neighbors_map = get_neighbors_map(array)
+        self.assertDictEqual({}, neighbors_map)
+
     def test_get_cropped_array_no_labels_no_crop_original(self):
         label = 1
 

tests/abm_colony_collection/test_make_voxels_array.py

@@ -0,0 +1,84 @@
+import unittest
+
+import numpy as np
+
+from abm_colony_collection.make_voxels_array import make_voxels_array
+
+
+class TestMakeVoxelsArray(unittest.TestCase):
+    def test_make_voxels_array_no_locations(self):
+        locations = [
+            {
+                "id": 1,
+                "location": [
+                    {
+                        "region": "A",
+                        "voxels": [],
+                    },
+                ],
+            },
+            {
+                "id": 2,
+                "location": [
+                    {
+                        "region": "A",
+                        "voxels": [],
+                    },
+                ],
+            },
+        ]
+
+        expected_array = np.zeros((1, 1, 1))
+
+        array = make_voxels_array(locations)
+
+        self.assertTrue((expected_array == array).all())
+
+    def test_make_voxels_array(self):
+        locations = [
+            {
+                "id": 1,
+                "location": [
+                    {
+                        "region": "A",
+                        "voxels": [
+                            [1, 1, 1],
+                            [1, 2, 1],
+                            [1, 1, 2],
+                        ],
+                    },
+                    {
+                        "region": "B",
+                        "voxels": [
+                            [1, 2, 2],
+                        ],
+                    },
+                ],
+            },
+            {
+                "id": 2,
+                "location": [
+                    {
+                        "region": "A",
+                        "voxels": [
+                            [2, 5, 2],
+                            [2, 5, 3],
+                            [2, 6, 2],
+                            [2, 6, 3],
+                        ],
+                    },
+                ],
+            },
+        ]
+
+        expected_array = np.zeros((5, 8, 4))
+        expected_array[1:3, 1:3, 1] = 1
+        expected_array[2:4, 5:7, 2] = 2
+
+        array = make_voxels_array(locations)
+
+        self.assertTrue((expected_array == array).all())
+
+
+if __name__ == "__main__":
+    unittest.main()