Coverage for src/arcade_collection/input/convert_to_locations_file.py: 100%
29 statements
« prev ^ index » next coverage.py v7.1.0, created at 2024-12-09 19:07 +0000
1from __future__ import annotations
3from typing import TYPE_CHECKING
5if TYPE_CHECKING:
6 import pandas as pd
9def convert_to_locations_file(samples: pd.DataFrame) -> list[dict]:
10 """
11 Convert all samples to location objects.
13 Parameters
14 ----------
15 samples
16 Sample cell ids and coordinates.
18 Returns
19 -------
20 :
21 List of location objects formatted for ARCADE.
22 """
24 locations: list[dict] = []
25 samples_by_id = samples.groupby("id")
27 for i, (_, group) in enumerate(samples_by_id):
28 locations.append(convert_to_location(i + 1, group))
30 return locations
33def convert_to_location(cell_id: int, samples: pd.DataFrame) -> dict:
34 """
35 Convert samples to location object.
37 Parameters
38 ----------
39 cell_id
40 Unique cell id.
41 samples
42 Sample coordinates for a single object.
44 Returns
45 -------
46 :
47 Location object formatted for ARCADE.
48 """
50 center = get_center_voxel(samples)
52 if "region" in samples.columns and not samples["region"].isna().all():
53 voxels = [
54 {"region": region, "voxels": get_location_voxels(samples, region)}
55 for region in samples["region"].unique()
56 ]
57 else:
58 voxels = [{"region": "UNDEFINED", "voxels": get_location_voxels(samples)}]
60 return {
61 "id": cell_id,
62 "center": center,
63 "location": voxels,
64 }
67def get_center_voxel(samples: pd.DataFrame) -> tuple[int, int, int]:
68 """
69 Get coordinates of center voxel of samples.
71 Parameters
72 ----------
73 samples
74 Sample cell ids and coordinates.
76 Returns
77 -------
78 :
79 Center voxel.
80 """
82 center_x = int(samples["x"].mean())
83 center_y = int(samples["y"].mean())
84 center_z = int(samples["z"].mean())
85 return (center_x, center_y, center_z)
88def get_location_voxels(
89 samples: pd.DataFrame, region: str | None = None
90) -> list[tuple[int, int, int]]:
91 """
92 Get list of voxel coordinates from samples dataframe.
94 Parameters
95 ----------
96 samples
97 Sample cell ids and coordinates.
98 region
99 Region key.
101 Returns
102 -------
103 :
104 List of voxel coordinates.
105 """
107 if region is not None:
108 region_samples = samples[samples["region"] == region]
109 voxels_x = region_samples["x"]
110 voxels_y = region_samples["y"]
111 voxels_z = region_samples["z"]
112 else:
113 voxels_x = samples["x"]
114 voxels_y = samples["y"]
115 voxels_z = samples["z"]
117 return list(zip(voxels_x, voxels_y, voxels_z))