Add docstrings and unit tests for convert to locations file task

src/arcade_collection/input/convert_to_locations_file.py

@@ -4,6 +4,20 @@
 
 
 def convert_to_locations_file(samples: pd.DataFrame) -> list[dict]:
+    """
+    Convert all samples to location objects.
+
+    Parameters
+    ----------
+    samples
+        Sample cell ids and coordinates.
+
+    Returns
+    -------
+    :
+        List of location objects formatted for ARCADE.
+    """
+
     locations: list[dict] = []
     samples_by_id = samples.groupby("id")
 
@@ -15,23 +29,24 @@ def convert_to_locations_file(samples: pd.DataFrame) -> list[dict]:
 
 def convert_to_location(cell_id: int, samples: pd.DataFrame) -> dict:
     """
-    Convert samples to ARCADE .LOCATIONS json format.
+    Convert samples to location object.
 
     Parameters
     ----------
     cell_id
         Unique cell id.
     samples
-        Sample cell ids and coordinates.
+        Sample coordinates for a single object.
 
     Returns
     -------
     :
-        Dictionary in ARCADE .LOCATIONS json format.
+        Location object formatted for ARCADE.
     """
+
     center = get_center_voxel(samples)
 
-    if "region" in samples.columns:
+    if "region" in samples.columns and not samples["region"].isnull().all():
         voxels = [
             {"region": region, "voxels": get_location_voxels(samples, region)}
             for region in samples["region"].unique()
@@ -44,6 +59,7 @@ def convert_to_location(cell_id: int, samples: pd.DataFrame) -> dict:
         "center": center,
         "location": voxels,
     }
+
     return location
 
 
@@ -61,6 +77,7 @@ def get_center_voxel(samples: pd.DataFrame) -> tuple[int, int, int]:
     :
         Center voxel.
     """
+
     center_x = int(samples["x"].mean())
     center_y = int(samples["y"].mean())
     center_z = int(samples["z"].mean())
@@ -86,6 +103,7 @@ def get_location_voxels(
     :
         List of voxel coordinates.
     """
+
     if region is not None:
         region_samples = samples[samples["region"] == region]
         voxels_x = region_samples["x"]

tests/arcade_collection/input/test_convert_to_locations_file.py

@@ -0,0 +1,195 @@
+import unittest
+
+import pandas as pd
+
+from arcade_collection.input.convert_to_locations_file import (
+    convert_to_location,
+    convert_to_locations_file,
+    get_center_voxel,
+    get_location_voxels,
+)
+
+
+class TestConvertToLocationsFile(unittest.TestCase):
+    def test_convert_to_locations_file(self):
+        samples = pd.DataFrame(
+            {
+                "id": [10, 10, 10, 10, 10, 11, 11, 11, 11],
+                "x": [0, 1, 1, 2, 2, 30, 31, 31, 32],
+                "y": [3, 3, 4, 5, 5, 40, 42, 42, 44],
+                "z": [6, 6, 7, 7, 8, 50, 51, 52, 52],
+                "region": [None, None, None, None, None, "A", "B", "A", "B"],
+            }
+        )
+
+        expected_locations = [
+            {
+                "id": 1,
+                "center": (1, 4, 6),
+                "location": [
+                    {
+                        "region": "UNDEFINED",
+                        "voxels": [
+                            (0, 3, 6),
+                            (1, 3, 6),
+                            (1, 4, 7),
+                            (2, 5, 7),
+                            (2, 5, 8),
+                        ],
+                    }
+                ],
+            },
+            {
+                "id": 2,
+                "center": (31, 42, 51),
+                "location": [
+                    {
+                        "region": "A",
+                        "voxels": [
+                            (30, 40, 50),
+                            (31, 42, 52),
+                        ],
+                    },
+                    {
+                        "region": "B",
+                        "voxels": [
+                            (31, 42, 51),
+                            (32, 44, 52),
+                        ],
+                    },
+                ],
+            },
+        ]
+
+        locations = convert_to_locations_file(samples)
+
+        self.assertCountEqual(expected_locations, locations)
+
+    def test_convert_to_location_no_region(self):
+        cell_id = 2
+        samples = pd.DataFrame(
+            {
+                "x": [0, 1, 1, 2, 2],
+                "y": [3, 3, 4, 5, 5],
+                "z": [6, 6, 7, 7, 8],
+            }
+        )
+        center = (1, 4, 6)
+
+        expected_location = {
+            "id": cell_id,
+            "center": center,
+            "location": [
+                {
+                    "region": "UNDEFINED",
+                    "voxels": [
+                        (0, 3, 6),
+                        (1, 3, 6),
+                        (1, 4, 7),
+                        (2, 5, 7),
+                        (2, 5, 8),
+                    ],
+                }
+            ],
+        }
+
+        location = convert_to_location(cell_id, samples)
+
+        self.assertDictEqual(expected_location, location)
+
+    def test_convert_to_location_with_region(self):
+        cell_id = 2
+        samples = pd.DataFrame(
+            {
+                "x": [0, 1, 1, 2, 2],
+                "y": [3, 3, 4, 5, 5],
+                "z": [6, 6, 7, 7, 8],
+                "region": ["A", "B", "A", "B", "A"],
+            }
+        )
+        center = (1, 4, 6)
+
+        expected_location = {
+            "id": cell_id,
+            "center": center,
+            "location": [
+                {
+                    "region": "A",
+                    "voxels": [
+                        (0, 3, 6),
+                        (1, 4, 7),
+                        (2, 5, 8),
+                    ],
+                },
+                {
+                    "region": "B",
+                    "voxels": [
+                        (1, 3, 6),
+                        (2, 5, 7),
+                    ],
+                },
+            ],
+        }
+
+        location = convert_to_location(cell_id, samples)
+
+        self.assertDictEqual(expected_location, location)
+
+    def test_get_center_voxel(self):
+        parameters = [
+            ([10, 12], [3, 5], [2, 4], (11, 4, 3)),  # exact
+            ([10, 11], [3, 4], [2, 3], (10, 3, 2)),  # rounded
+        ]
+
+        for x, y, z, expected_center in parameters:
+            with self.subTest(x=x, y=y, z=z):
+                samples = pd.DataFrame({"x": x, "y": y, "z": z})
+                center = get_center_voxel(samples)
+                self.assertTupleEqual(expected_center, center)
+
+    def test_get_location_voxels_no_region(self):
+        region = None
+        samples = pd.DataFrame(
+            {
+                "x": [0, 1, 1, 2, 2],
+                "y": [3, 3, 4, 5, 5],
+                "z": [6, 6, 7, 7, 8],
+            }
+        )
+
+        expected_voxels = [
+            (0, 3, 6),
+            (1, 3, 6),
+            (1, 4, 7),
+            (2, 5, 7),
+            (2, 5, 8),
+        ]
+
+        voxels = get_location_voxels(samples, region)
+
+        self.assertCountEqual(expected_voxels, voxels)
+
+    def test_get_location_voxels_with_region(self):
+        region = "A"
+        samples = pd.DataFrame(
+            {
+                "x": [0, 1, 1, 2, 2],
+                "y": [3, 3, 4, 5, 5],
+                "z": [6, 6, 7, 7, 8],
+                "region": ["A", "B", "A", "B", "A"],
+            }
+        )
+
+        expected_voxels = [
+            (0, 3, 6),
+            (1, 4, 7),
+            (2, 5, 8),
+        ]
+
+        voxels = get_location_voxels(samples, region)
+
+        self.assertCountEqual(expected_voxels, voxels)
+
+
+if __name__ == "__main__":
+    unittest.main()