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

1from typing import Optional 

2 

3import matplotlib.figure as mpl 

4import matplotlib.pyplot as plt 

5import pandas as pd 

6from prefect import task 

7 

8 

9@task 

10def make_scatter_figure(data: pd.DataFrame, colormap: Optional[dict] = None) -> 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 

16 if colormap is not None: 

17 ax.scatter(data["x"], data["y"], c=data["v"].map(colormap), s=10) 

18 else: 

19 ax.scatter(data["x"], data["y"], c=data["v"], cmap="magma_r", s=10) 

20 

21 return fig