plotting.py

import itertools


def add_significance_markers(
        ax, data,
        x_values,
        hue_values=None,
        x_label="treatment_size",
        hue_label="model",
        statistic_label="t"):
    """
    Add significance markers on a bar plot.
    """
    params = itertools.product(x_values, hue_values)
    sorted_patches = sorted(ax.patches, key=lambda patch: patch.get_x())
    for i, ((x_val, hue_val), patch) in enumerate(zip(params, sorted_patches)):
        x = patch.get_x() + patch.get_width() / 2
        y = patch.get_y() + patch.get_height()
        p_value = data.query(f"`{hue_label}` == '{hue_val}' and `{x_label}` == {repr(x_val)}").p.mean()
        statistic_value = data.query(
            f"`{hue_label}` == '{hue_val}' and `{x_label}` == {repr(x_val)}")[statistic_label].mean()

        print(f"{x_val} - {hue_val}. p={p_value}, {statistic_label}={statistic_value}")
        marker = "."
        color = "black"
        size = 7

        if 0.001 < p_value <= 0.01:
            marker = "o"
            color = "white"
            size = 7
        elif p_value <= 0.001:
            marker = "D"
            color = "white"
            size = 10

        ax.plot(x, y, marker=marker,
                markersize=size,
                color=color,
                markeredgecolor="black",
                markeredgewidth=2,
                zorder=1000)