Coverage for src/cell_abm_pipeline/flows/calculate_neighbors.py: 0%

1""" 

2Workflow for calculating neighbor connections. 

3 

4Working location structure: 

5 

6.. code-block:: bash 

7 

8 (name) 

9 ├── data 

10 │ └── data.LOCATIONS 

11 │ └── (name)_(key)_(seed).LOCATIONS.tar.xz 

12 └── calculations 

13 └── calculations.NEIGHBORS 

14 └── (name)_(key)_(seed)_(tick).NEIGHBORS.csv 

15 

16Data from **data.LOCATIONS** are used to calculate neighbors, which are saved to 

17**calculations.NEIGHBORS**. 

18""" 

19 

20from dataclasses import dataclass 

21 

22import pandas as pd 

23from abm_colony_collection import get_depth_map, get_neighbors_map, make_voxels_array 

24from arcade_collection.output import extract_tick_json 

25from io_collection.keys import make_key 

26from io_collection.load import load_tar 

27from io_collection.save import save_dataframe 

28from prefect import flow 

29 

30 

31@dataclass 

32class ParametersConfig: 

33 """Parameter configuration for calculate neighbors flow.""" 

34 

35 key: str 

36 """Simulation key to calculate.""" 

37 

38 seed: int 

39 """Simulation random seed to calculate.""" 

40 

41 tick: int 

42 """Simulation tick to calculate.""" 

43 

44 

45@dataclass 

46class ContextConfig: 

47 """Context configuration for calculate neighbors flow.""" 

48 

49 working_location: str 

50 """Location for input and output files (local path or S3 bucket).""" 

51 

52 

53@dataclass 

54class SeriesConfig: 

55 """Series configuration for calculate neighbors flow.""" 

56 

57 name: str 

58 """Name of the simulation series.""" 

59 

60 

61@flow(name="calculate-neighbors") 

62def run_flow(context: ContextConfig, series: SeriesConfig, parameters: ParametersConfig) -> None: 

63 """Main calculate neighbors flow.""" 

64 

65 data_key = make_key(series.name, "data", "data.LOCATIONS") 

66 calc_key = make_key(series.name, "calculations", "calculations.NEIGHBORS") 

67 series_key = f"{series.name}_{parameters.key}_{parameters.seed:04d}" 

68 

69 locations_key = make_key(data_key, f"{series_key}.LOCATIONS.tar.xz") 

70 locations_tar = load_tar(context.working_location, locations_key) 

71 locations_json = extract_tick_json(locations_tar, series_key, parameters.tick, "LOCATIONS") 

72 

73 array = make_voxels_array(locations_json) 

74 

75 neighbors_map = get_neighbors_map(array) 

76 depth_map = get_depth_map(array, neighbors_map) 

77 center_map = {location["id"]: location["center"] for location in locations_json} 

78 

79 attributes = {"KEY": parameters.key, "SEED": parameters.seed, "TICK": parameters.tick} 

80 all_neighbors = [] 

81 

82 for voxel_id, voxel_neighbors in neighbors_map.items(): 

83 voxel_neighbors = { 

84 "ID": voxel_id, 

85 "GROUP": voxel_neighbors["group"], 

86 "NEIGHBORS": voxel_neighbors["neighbors"], 

87 "CX": center_map[voxel_id][0], 

88 "CY": center_map[voxel_id][1], 

89 "CZ": center_map[voxel_id][2], 

90 "DEPTH": depth_map[voxel_id], 

91 } 

92 voxel_neighbors.update(attributes) 

93 all_neighbors.append(voxel_neighbors) 

94 

95 neighbors_dataframe = pd.DataFrame(all_neighbors) 

96 neighbors_key = make_key(calc_key, f"{series_key}_{parameters.tick:06d}.NEIGHBORS.csv") 

97 save_dataframe(context.working_location, neighbors_key, neighbors_dataframe, index=False)