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

1""" 

2Workflow for plotting colony dynamics. 

3 

4Working location structure: 

5 

6.. code-block:: bash 

7 

8 (name) 

9 ├── groups 

10 │ └── groups.COLONIES 

11 │ ├── (name).feature_distributions.(feature).json 

12 │ ├── (name).feature_temporal.(key).(feature).json 

13 │ ├── (name).neighbor_positions.(key).(seed).(tick).csv 

14 │ └── (name).neighbor_positions.(key).(seed).(tick).(feature).csv 

15 └── plots 

16 └── plots.COLONIES 

17 ├── (name).feature_distributions.(feature).png 

18 ├── (name).feature_temporal.(key).(feature).json 

19 └── (name).neighbor_positions.(key).(seed).(tick).(feature).png 

20 

21Plots use grouped data from **groups.COLONIES**. Plots are saved to 

22**plots.COLONIES**. 

23""" 

24 

25from dataclasses import dataclass, field 

26 

27from io_collection.keys import make_key 

28from io_collection.load import load_dataframe, load_json 

29from io_collection.save import save_figure 

30from prefect import flow 

31 

32from cell_abm_pipeline.flows.group_colony_dynamics import ( 

33 DISTRIBUTION_FEATURES, 

34 POSITION_FEATURES, 

35 TEMPORAL_FEATURES, 

36) 

37from cell_abm_pipeline.tasks import make_graph_figure, make_histogram_figure, make_range_figure 

38 

39PLOTS: list[str] = [ 

40 "feature_distributions", 

41 "feature_temporal", 

42 "neighbor_positions", 

43] 

44 

45 

46FEATURE_COLORMAPS = {"depth": "magma_r", "group": "tab10"} 

47 

48 

49@dataclass 

50class ParametersConfigFeatureDistributions: 

51 """Parameter configuration for plot colony dynamics subflow - feature distributions.""" 

52 

53 features: list[str] = field(default_factory=lambda: DISTRIBUTION_FEATURES) 

54 """List of colony features.""" 

55 

56 

57@dataclass 

58class ParametersConfigFeatureTemporal: 

59 """Parameter configuration for plot colony dynamics subflow - feature temporal.""" 

60 

61 features: list[str] = field(default_factory=lambda: TEMPORAL_FEATURES) 

62 """List of temporal features.""" 

63 

64 

65@dataclass 

66class ParametersConfigNeighborPositions: 

67 """Parameter configuration for plot colony dynamics subflow - neighbor positions.""" 

68 

69 features: list[str] = field(default_factory=lambda: POSITION_FEATURES) 

70 """List of position features.""" 

71 

72 seed: int = 0 

73 """Simulation seed to use for plotting neighbor positions.""" 

74 

75 ticks: list[int] = field(default_factory=lambda: [0]) 

76 """Simulation ticks to use for plotting neighbor positions.""" 

77 

78 colormaps: dict[str, str] = field(default_factory=lambda: FEATURE_COLORMAPS) 

79 """Colormaps for each feature.""" 

80 

81 

82@dataclass 

83class ParametersConfig: 

84 """Parameter configuration for plot colony dynamics flow.""" 

85 

86 plots: list[str] = field(default_factory=lambda: PLOTS) 

87 """List of colony dynamics plots.""" 

88 

89 feature_distributions: ParametersConfigFeatureDistributions = ( 

90 ParametersConfigFeatureDistributions() 

91 ) 

92 """Parameters for plot feature distributions subflow.""" 

93 

94 feature_temporal: ParametersConfigFeatureTemporal = ParametersConfigFeatureTemporal() 

95 """Parameters for plot feature temporal subflow.""" 

96 

97 neighbor_positions: ParametersConfigNeighborPositions = ParametersConfigNeighborPositions() 

98 """Parameters for plot neighbor positions subflow.""" 

99 

100 

101@dataclass 

102class ContextConfig: 

103 """Context configuration for plot colony dynamics flow.""" 

104 

105 working_location: str 

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

107 

108 

109@dataclass 

110class SeriesConfig: 

111 """Series configuration for plot colony dynamics flow.""" 

112 

113 name: str 

114 """Name of the simulation series.""" 

115 

116 conditions: list[dict] 

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

118 

119 

120@flow(name="plot-colony-dynamics") 

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

122 """ 

123 Main plot colony dynamics flow. 

124 

125 Calls the following subflows, if the plot is specified: 

126 

127 - :py:func:`run_flow_plot_feature_distributions` 

128 - :py:func:`run_flow_plot_feature_temporal` 

129 - :py:func:`run_flow_plot_neighbor_positions` 

130 """ 

131 

132 if "feature_distributions" in parameters.plots: 

133 run_flow_plot_feature_distributions(context, series, parameters.feature_distributions) 

134 

135 if "feature_temporal" in parameters.plots: 

136 run_flow_plot_feature_temporal(context, series, parameters.feature_temporal) 

137 

138 if "neighbor_positions" in parameters.plots: 

139 run_flow_plot_neighbor_positions(context, series, parameters.neighbor_positions) 

140 

141 

142@flow(name="plot-colony-dynamics_plot-feature-distributions") 

143def run_flow_plot_feature_distributions( 

144 context: ContextConfig, series: SeriesConfig, parameters: ParametersConfigFeatureDistributions 

145) -> None: 

146 """Plot colony dynamics subflow for feature distributions.""" 

147 

148 group_key = make_key(series.name, "groups", "groups.COLONIES") 

149 plot_key = make_key(series.name, "plots", "plots.COLONIES") 

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

151 

152 for feature in parameters.features: 

153 feature_key = feature.upper() 

154 

155 group = load_json( 

156 context.working_location, 

157 make_key(group_key, f"{series.name}.feature_distributions.{feature_key}.json"), 

158 ) 

159 

160 assert isinstance(group, dict) 

161 

162 save_figure( 

163 context.working_location, 

164 make_key(plot_key, f"{series.name}.feature_distributions.{feature_key}.png"), 

165 make_histogram_figure(keys, group), 

166 ) 

167 

168 

169@flow(name="plot-colony-dynamics_plot-feature-temporal") 

170def run_flow_plot_feature_temporal( 

171 context: ContextConfig, series: SeriesConfig, parameters: ParametersConfigFeatureTemporal 

172) -> None: 

173 """Plot colony dynamics subflow for temporal features.""" 

174 

175 group_key = make_key(series.name, "groups", "groups.COLONIES") 

176 plot_key = make_key(series.name, "plots", "plots.COLONIES") 

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

178 

179 for key in keys: 

180 for feature in parameters.features: 

181 feature_key = f"{key}.{feature.upper()}" 

182 

183 group = load_json( 

184 context.working_location, 

185 make_key(group_key, f"{series.name}.feature_temporal.{feature_key}.json"), 

186 ) 

187 

188 assert isinstance(group, dict) 

189 

190 save_figure( 

191 context.working_location, 

192 make_key(plot_key, f"{series.name}.feature_temporal.{feature_key}.png"), 

193 make_range_figure(group), 

194 ) 

195 

196 

197@flow(name="plot-colony-dynamics_plot-neighbor-positions") 

198def run_flow_plot_neighbor_positions( 

199 context: ContextConfig, series: SeriesConfig, parameters: ParametersConfigNeighborPositions 

200) -> None: 

201 """Plot colony dynamics subflow for neighbor positions.""" 

202 

203 group_key = make_key(series.name, "groups", "groups.COLONIES") 

204 plot_key = make_key(series.name, "plots", "plots.COLONIES") 

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

206 

207 for key in keys: 

208 for tick in parameters.ticks: 

209 for feature in parameters.features: 

210 edge_key = f"{key}.{parameters.seed:04d}.{tick:06d}" 

211 node_key = f"{key}.{parameters.seed:04d}.{tick:06d}.{feature.upper()}" 

212 

213 edge_group = load_dataframe( 

214 context.working_location, 

215 make_key(group_key, f"{series.name}.neighbor_positions.{edge_key}.csv"), 

216 ) 

217 

218 node_group = load_dataframe( 

219 context.working_location, 

220 make_key(group_key, f"{series.name}.neighbor_positions.{node_key}.csv"), 

221 ) 

222 

223 save_figure( 

224 context.working_location, 

225 make_key(plot_key, f"{series.name}.neighbor_positions.{node_key}.png"), 

226 make_graph_figure(node_group, edge_group, parameters.colormaps[feature]), 

227 )