def plot_value_counts()

in src/beanmachine/tutorials/utils/hearts.py [0:0]

• 63 lines of code
• 4 McCabe index (conditional complexity)

def plot_value_counts(df: pd.DataFrame) -> List[List[Figure]]:
    """
    Plot the pre-drug and post-drug value counts as a bar plot.

    :param df: Pandas dataframe object of the model data.
    :type df: pd.DataFrame
    :return: A list of figures to display in the notebook.
    :rtype: List[List[Figure]]
    """
    predrug_data = df["predrug"].value_counts().sort_index()
    predrug_index = predrug_data.index.values
    postdrug_data = df["postdrug"].value_counts().sort_index()
    postdrug_index = postdrug_data.index.values
    PADDING = 2
    x = np.arange(0, max(predrug_data.index.max(), postdrug_data.index.max()) + PADDING)
    top_predrug = np.zeros(len(x))
    top_predrug[predrug_index] = predrug_data
    top_postdrug = np.zeros(len(x))
    top_postdrug[postdrug_index] = postdrug_data
    OFFSET = 0.5
    left = x - OFFSET
    right = x + OFFSET
    bottom = np.zeros(len(x))

    figs = []
    for i, column in enumerate(["predrug", "postdrug"]):
        # Create the figure.
        p = figure(
            plot_width=700,
            plot_height=300,
            y_axis_label="Counts",
            x_axis_label="PVC events",
            title=f'PVC events "{column}"',
            y_range=[0, 8],
            x_range=[-2, 52],
            outline_line_color="black",
        )

        # Prepare data for the figure.
        source_data = {
            "x": x,
            "left": left,
            "top": None,
            "right": right,
            "bottom": bottom,
        }
        if i == 0:
            source_data["top"] = top_predrug
        else:
            source_data["top"] = top_postdrug
        source = ColumnDataSource(source_data)

        # Add data to the figure.
        glyph = p.quad(
            left="left",
            top="top",
            right="right",
            bottom="bottom",
            source=source,
            fill_color="steelblue",
            line_color="white",
            fill_alpha=0.7,
            hover_fill_color="orange",
            hover_line_color="black",
            hover_alpha=1,
        )

        # Add tooltips to the figure.
        tips = HoverTool(
            renderers=[glyph],
            tooltips=[("Count", "@top"), ("PVC events", "@x")],
        )
        p.add_tools(tips)

        # Style the figure
        p.grid.grid_line_alpha = 0.2
        p.grid.grid_line_color = "gray"
        p.grid.grid_line_width = 0.3
        p.yaxis.minor_tick_line_color = None
        figs.append(p)

    return [[figs[0]], [figs[1]]]