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

1from math import sqrt 

2 

3import matplotlib.figure as mpl 

4import matplotlib.pyplot as plt 

5import numpy as np 

6import pandas as pd 

7from matplotlib import colormaps 

8from matplotlib.patches import RegularPolygon 

9from prefect import task 

10 

11 

12@task 

13def make_density_figure(data: pd.DataFrame, scale: float) -> mpl.Figure: 

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

15 

16 ax = fig.add_subplot() 

17 ax.set_box_aspect(1) 

18 

19 cmap = colormaps["magma_r"] 

20 data["vn"] = (data["v"] - data["v"].min()) / (data["v"].max() - data["v"].min()) 

21 

22 for _, row in data.iterrows(): 

23 hexagon = RegularPolygon( 

24 (row["x"], row["y"]), 

25 numVertices=6, 

26 radius=(scale / sqrt(3)), 

27 orientation=np.radians(30), 

28 facecolor=cmap(row["vn"]), 

29 ) 

30 ax.add_patch(hexagon) 

31 

32 ax.set_xbound(data["x"].min() - scale, data["x"].max() + scale) 

33 ax.set_ybound(data["y"].min() - scale, data["y"].max() + scale) 

34 

35 return fig