Coverage for src/cell_abm_pipeline/tasks/make_graph_figure.py: 0%

1import matplotlib.figure as mpl 

2import matplotlib.pyplot as plt 

3import pandas as pd 

4from prefect import task 

5 

6 

7@task 

8def make_graph_figure( 

9 node_data: pd.DataFrame, edge_data: pd.DataFrame, colormap: str 

10) -> mpl.Figure: 

11 fig = plt.figure(figsize=(4, 4), constrained_layout=True) 

12 

13 ax = fig.add_subplot() 

14 ax.set_box_aspect(1) 

15 ax.invert_yaxis() 

16 ax.get_xaxis().set_ticks([]) 

17 ax.get_yaxis().set_ticks([]) 

18 

19 for spine in ax.spines.values(): 

20 spine.set_edgecolor("#dddddd") 

21 

22 node_data = node_data.set_index("id") 

23 

24 for edge in edge_data.to_dict("records"): 

25 x1, y1 = node_data.loc[edge["id1"]][["x", "y"]] 

26 x2, y2 = node_data.loc[edge["id2"]][["x", "y"]] 

27 

28 ax.plot([x1, x2], [y1, y2], color="k", lw=0.5, zorder=1) 

29 

30 ax.scatter(node_data["x"], node_data["y"], c=node_data["v"], cmap=colormap, s=20) 

31 

32 return fig