Coverage for src/cell_abm_pipeline/tasks/physicell/convert_physicell_to_simularium.py: 0%

1import json 

2import tempfile 

3from tarfile import TarFile 

4 

5import numpy as np 

6from simulariumio import DISPLAY_TYPE, CameraData, DisplayData, JsonWriter, MetaData, UnitData 

7from simulariumio.filters import TranslateFilter 

8from simulariumio.physicell import PhysicellConverter, PhysicellData 

9 

10 

11def convert_physicell_to_simularium( 

12 tar_file: TarFile, 

13 box_size: tuple[float, float, float], 

14 timestep: float, 

15 frame_spec: tuple[float, float, float], 

16 substrate_color: str, 

17 cell_colors: list[str], 

18 max_owner_cells: int = 10000, 

19) -> str: 

20 with tempfile.TemporaryDirectory() as temp_directory: 

21 # Extract specified contents of archive into temporary directory. 

22 tar_file.extract("initial_mesh0.mat", temp_directory) 

23 for timepoint in np.arange(*frame_spec): 

24 tar_file.extract(f"output{timepoint:08d}.xml", temp_directory) 

25 tar_file.extract(f"output{timepoint:08d}_cells_physicell.mat", temp_directory) 

26 

27 # Set sizing. 

28 box_size_array = np.array(box_size) 

29 scale = 0.1 

30 offset = scale * (-0.5 * box_size_array) 

31 

32 # Create data object. 

33 data = PhysicellData( 

34 timestep=timestep, 

35 path_to_output_dir=temp_directory, 

36 meta_data=MetaData( 

37 box_size=box_size_array, 

38 scale_factor=scale, 

39 camera_defaults=CameraData(position=np.array([30.0, -10.0, 30.0])), 

40 ), 

41 display_data={ 

42 0: DisplayData( 

43 name="Substrate", 

44 display_type=DISPLAY_TYPE.SPHERE, 

45 ) 

46 }, 

47 time_units=UnitData("min"), 

48 ) 

49 

50 # Create converter object. 

51 converter = PhysicellConverter(data) 

52 converter._data.spatial_units = UnitData(name="micron", magnitude=10.0) 

53 

54 # Update agent names. 

55 type_names = [] 

56 for time_index, _ in enumerate(converter._data.agent_data.types): 

57 for agent_index, _ in enumerate(converter._data.agent_data.types[time_index]): 

58 tn = converter._data.agent_data.types[time_index][agent_index] 

59 new_name = "" 

60 

61 if "Substrate" in tn: 

62 new_name = "Substrate" 

63 

64 if "cell" in tn: 

65 cell_id = int(tn[4 : tn.index("#")]) 

66 

67 while cell_id > max_owner_cells: 

68 cell_id -= max_owner_cells 

69 

70 new_name = "Stem cell#" + str(cell_id) 

71 

72 if new_name not in type_names: 

73 type_names.append(new_name) 

74 

75 if not new_name: 

76 continue 

77 

78 converter._data.agent_data.types[time_index][agent_index] = new_name 

79 

80 # Set substrate agent color. 

81 converter._data.agent_data.display_data = { 

82 "Substrate": DisplayData( 

83 name="Substrate", 

84 display_type=DISPLAY_TYPE.SPHERE, 

85 color=substrate_color, 

86 ) 

87 } 

88 

89 # Set cell agent colors. 

90 color_map = { 

91 type_name: cell_colors[type_index % len(cell_colors)] 

92 for type_index, type_name in enumerate(type_names) 

93 } 

94 for type_name in type_names: 

95 converter._data.agent_data.display_data[type_name] = DisplayData( 

96 name=type_name, 

97 display_type=DISPLAY_TYPE.SPHERE, 

98 color=color_map[type_name], 

99 ) 

100 

101 # Remove the huge cell artifact. 

102 converter._data.agent_data.n_agents[0] -= 1 

103 

104 # Center cells in box. 

105 filtered_data = converter.filter_data([TranslateFilter(default_translation=offset)]) 

106 

107 json_data = JsonWriter.format_trajectory_data(filtered_data) 

108 return json.dumps(json_data)