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

1""" 

2Workflow for making simulation movies. 

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 ├── movies 

13 │ ├── movies.CENTROIDS 

14 │ │ ├── (name)_(key)_(seed).CENTROIDS.gif 

15 │ │ └── (name)_(key)_(seed) 

16 │ │ └── (frame).CENTROIDS.png 

17 │ └── movies.SCAN 

18 │ ├── (name)_(key)_(seed)_(view)_(frame).SCAN.gif 

19 │ └── (name)_(key)_(seed)_(view)_(frame) 

20 │ └── (index).SCAN.png 

21 └── results 

22 └── (name)_(key)_(seed).csv 

23 

24Different formats use inputs from **results** and **data.LOCATIONS**. Movies are 

25saved to **movies**. 

26""" 

27 

28from dataclasses import dataclass, field 

29 

30import numpy as np 

31from arcade_collection.output import get_voxel_contours 

32from io_collection.keys import check_key, make_key 

33from io_collection.load import load_dataframe, load_tar 

34from io_collection.save import save_figure, save_gif 

35from prefect import flow, get_run_logger 

36 

37from cell_abm_pipeline.flows.plot_cell_shapes import REGION_COLORS 

38from cell_abm_pipeline.tasks import make_centroids_figure, make_contour_figure 

39 

40FORMATS: list[str] = [ 

41 "centroids", 

42 "scan", 

43] 

44 

45 

46@dataclass 

47class ParametersConfigScan: 

48 """Parameter configuration for make simulation movies flow - scan.""" 

49 

50 seeds: list[int] = field(default_factory=lambda: [0]) 

51 """Simulation seeds to use for creating scan movies.""" 

52 

53 frame_spec: tuple[int, int, int] = (0, 1153, 1152) 

54 """Specification for simulation ticks to use for creating scan movies.""" 

55 

56 index_spec: tuple[int, int, int] = (0, 1, 1) 

57 """Specification for contour indices to use for creating scan movies.""" 

58 

59 regions: list[str] = field(default_factory=lambda: ["DEFAULT"]) 

60 """List of subcellular regions.""" 

61 

62 box: tuple[int, int, int] = field(default_factory=lambda: (1, 1, 1)) 

63 """Size of bounding box.""" 

64 

65 view: str = "top" 

66 """Projection view.""" 

67 

68 x_bounds: tuple[int, int] = field(default_factory=lambda: (0, 1)) 

69 """Size of x bounds.""" 

70 

71 y_bounds: tuple[int, int] = field(default_factory=lambda: (0, 1)) 

72 """Size of y bounds.""" 

73 

74 region_colors: dict[str, str] = field(default_factory=lambda: REGION_COLORS) 

75 """Colors for each cell region.""" 

76 

77 

78@dataclass 

79class ParametersConfigCentroids: 

80 """Parameter configuration for make simulation movies flow - centroids.""" 

81 

82 seeds: list[int] = field(default_factory=lambda: [0]) 

83 """Simulation seeds to use for creating centroid movies.""" 

84 

85 frame_spec: tuple[int, int, int] = (0, 1, 1) 

86 """Specification for simulation ticks to use for creating centroid movies.""" 

87 

88 x_bounds: tuple[int, int] = field(default_factory=lambda: (0, 1)) 

89 """Size of x bounds.""" 

90 

91 y_bounds: tuple[int, int] = field(default_factory=lambda: (0, 1)) 

92 """Size of y bounds.""" 

93 

94 dt: float = 1.0 

95 """Temporal scaling in hours/tick.""" 

96 

97 window: int = 0 

98 """Window size for centroid tail.""" 

99 

100 

101@dataclass 

102class ParametersConfig: 

103 """Parameter configuration for make simulation movies flow.""" 

104 

105 formats: list[str] = field(default_factory=lambda: FORMATS) 

106 """List of movie formats.""" 

107 

108 scan: ParametersConfigScan = ParametersConfigScan() 

109 """Parameters for scan movie subflow.""" 

110 

111 centroids: ParametersConfigCentroids = ParametersConfigCentroids() 

112 """Parameters for centroids movie subflow.""" 

113 

114 

115@dataclass 

116class ContextConfig: 

117 """Context configuration for make simulation movies flow.""" 

118 

119 working_location: str 

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

121 

122 

123@dataclass 

124class SeriesConfig: 

125 """Series configuration for make simulation movies flow.""" 

126 

127 name: str 

128 """Name of the simulation series.""" 

129 

130 conditions: list[dict] 

131 """List of series condition dictionaries (must include unique condition "key").""" 

132 

133 

134@flow(name="make-simulation-movies") 

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

136 """ 

137 Main make simulation movies flow. 

138 

139 Calls the following subflows, if the format is specified. 

140 

141 - :py:func:`run_flow_make_centroids_movie` 

142 - :py:func:`run_flow_make_scan_movie` 

143 """ 

144 

145 if "centroids" in parameters.formats: 

146 run_flow_make_centroids_movie(context, series, parameters.centroids) 

147 

148 if "scan" in parameters.formats: 

149 run_flow_make_scan_movie(context, series, parameters.scan) 

150 

151 

152@flow(name="make-simulation-movies_make-centroids-movie") 

153def run_flow_make_centroids_movie( 

154 context: ContextConfig, series: SeriesConfig, parameters: ParametersConfigCentroids 

155) -> None: 

156 """Make simulation movies subflow for centroids.""" 

157 

158 movie_key = make_key(series.name, "movies", "movies.CENTROIDS") 

159 keys = [condition["key"] for condition in series.conditions] 

160 

161 for key in keys: 

162 for seed in parameters.seeds: 

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

164 

165 results_key = make_key(series.name, "results", f"{series_key}.csv") 

166 results = load_dataframe(context.working_location, results_key) 

167 

168 frame_keys = [] 

169 

170 for frame in np.arange(*parameters.frame_spec): 

171 frame_key = make_key(movie_key, f"{series_key}", f"{frame:06d}.CENTROIDS.png") 

172 frame_keys.append(frame_key) 

173 

174 if check_key(context.working_location, frame_key): 

175 continue 

176 

177 save_figure( 

178 context.working_location, 

179 frame_key, 

180 make_centroids_figure( 

181 results, 

182 frame, 

183 parameters.x_bounds, 

184 parameters.y_bounds, 

185 parameters.dt, 

186 parameters.window, 

187 ), 

188 ) 

189 

190 output_key = make_key(movie_key, f"{series_key}.CENTROIDS.gif") 

191 save_gif(context.working_location, output_key, frame_keys) 

192 

193 

194@flow(name="make-simulation-movies_make-scan-movie") 

195def run_flow_make_scan_movie( 

196 context: ContextConfig, series: SeriesConfig, parameters: ParametersConfigScan 

197) -> None: 

198 """Make simulation movies subflow for scan.""" 

199 

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

201 movie_key = make_key(series.name, "movies", "movies.SCAN") 

202 keys = [condition["key"] for condition in series.conditions] 

203 

204 if parameters.view not in ("top", "side"): 

205 logger = get_run_logger() 

206 logger.error("View [ %s ] not valid for scan movie.", parameters.view) 

207 return 

208 

209 indices = list(np.arange(*parameters.index_spec)) 

210 view = "top" if parameters.view == "top" else "side1" 

211 

212 for key in keys: 

213 for seed in parameters.seeds: 

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

215 

216 tar_key = make_key(data_key, f"{series_key}.LOCATIONS.tar.xz") 

217 tar = load_tar(context.working_location, tar_key) 

218 

219 for frame in np.arange(*parameters.frame_spec): 

220 frame_key = f"{series_key}_{parameters.view}_{frame:06d}" 

221 contours = get_voxel_contours( 

222 series_key, 

223 tar, 

224 frame, 

225 parameters.regions, 

226 parameters.box, 

227 {view: indices}, 

228 ) 

229 

230 index_keys = [] 

231 

232 for index in indices: 

233 index_key = make_key(movie_key, frame_key, f"{index:03d}.SCAN.png") 

234 index_keys.append(index_key) 

235 

236 if check_key(context.working_location, index_key): 

237 continue 

238 

239 save_figure( 

240 context.working_location, 

241 index_key, 

242 make_contour_figure( 

243 contours, 

244 index, 

245 view, 

246 parameters.regions, 

247 parameters.x_bounds, 

248 parameters.y_bounds, 

249 parameters.region_colors, 

250 ), 

251 ) 

252 

253 output_key = make_key(movie_key, f"{frame_key}.SCAN.gif") 

254 save_gif(context.working_location, output_key, index_keys)