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

1import matplotlib.figure as mpl 

2import matplotlib.pyplot as plt 

3import numpy as np 

4from prefect import task 

5 

6 

7@task 

8def make_heatmap_figure(rows: list[str], cols: list[str], values: list[list]) -> mpl.Figure: 

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

10 

11 ax = fig.add_subplot() 

12 

13 ax.imshow(values, cmap="magma_r") 

14 

15 ax.set_xticks(np.arange(len(cols)), labels=cols) 

16 ax.set_yticks(np.arange(len(rows)), labels=rows) 

17 

18 for i in range(len(rows)): 

19 for j in range(len(cols)): 

20 ax.text(j, i, f"{values[i][j]:.2f}", ha="center", va="center", color="#999") 

21 

22 return fig