Coverage for src/arcade_collection/convert/convert_to_simularium_objects.py: 100%
27 statements
« prev ^ index » next coverage.py v7.1.0, created at 2024-12-09 19:07 +0000
1import numpy as np
2import pandas as pd
4from arcade_collection.convert.convert_to_simularium import convert_to_simularium
7def convert_to_simularium_objects(
8 series_key: str,
9 simulation_type: str,
10 categories: pd.DataFrame,
11 regions: list[str],
12 frame_spec: tuple[int, int, int],
13 box: tuple[int, int, int],
14 ds: tuple[float, float, float],
15 dt: float,
16 colors: dict[str, str],
17 group_size: int,
18 url: str,
19 jitter: float = 1.0,
20) -> str:
21 """
22 Convert data to Simularium trajectory using mesh objects.
24 Parameters
25 ----------
26 series_key
27 Simulation series key.
28 simulation_type : {'potts'}
29 Simulation type.
30 categories
31 Simulation data containing ID, FRAME, and CATEGORY.
32 regions
33 List of regions.
34 frame_spec
35 Specification for simulation ticks.
36 box
37 Size of bounding box.
38 ds
39 Spatial scaling in um/voxel.
40 dt
41 Temporal scaling in hours/tick.
42 colors
43 Map of category to colors.
44 group_size
45 Number of objects in each mesh group.
46 url
47 URL for mesh object files.
48 jitter
49 Relative jitter applied to colors (set to 0 for exact colors).
51 Returns
52 -------
53 :
54 Simularium trajectory.
55 """
57 if simulation_type == "potts":
58 frames = list(map(int, np.arange(*frame_spec)))
59 length, width, height = box
60 data = format_potts_for_objects(
61 categories, frames, group_size, regions, length, width, height
62 )
63 else:
64 message = f"invalid simulation type {simulation_type}"
65 raise ValueError(message)
67 return convert_to_simularium(
68 series_key, simulation_type, data, length, width, height, ds, dt, colors, url, jitter
69 )
72def format_potts_for_objects(
73 categories: pd.DataFrame,
74 frames: list[int],
75 group_size: int,
76 regions: list[str],
77 length: int,
78 width: int,
79 height: int,
80) -> pd.DataFrame:
81 """
82 Format ``potts`` simulation data for object-based Simularium trajectory.
84 Parameters
85 ----------
86 categories
87 Simulation data containing ID, FRAME, and CATEGORY.
88 frames
89 List of frames.
90 group_size
91 Number of objects in each mesh group.
92 regions
93 List of regions.
94 length
95 Length of bounding box.
96 width
97 Width of bounding box.
98 height
99 Height of bounding box.
101 Returns
102 -------
103 :
104 Data formatted for trajectory.
105 """
107 data: list[list[object]] = []
108 center = [length / 2, width / 2, height / 2]
110 for frame in frames:
111 frame_categories = categories[categories["FRAME"] == frame]
112 index_offset = 0
114 for category, category_group in frame_categories.groupby("CATEGORY"):
115 ids = list(category_group["ID"].values)
116 group_ids = [ids[i : i + group_size] for i in range(0, len(ids), group_size)]
118 for i, _ in enumerate(group_ids):
119 index = i + index_offset
121 for region in regions:
122 name = f"{region}#{category}#{index}#{frame}"
123 data = [*data, [name, int(frame), 1, *center, [], "OBJ"]]
125 index_offset = index_offset + len(group_ids)
127 return pd.DataFrame(
128 data, columns=["name", "frame", "radius", "x", "y", "z", "points", "display"]
129 )