Function plot

benchmarking/plot_bsz_sweep_runtime.py:109–197

(
    df: pd.DataFrame,
    out_path: Path,
    fmt: str,
    y_col: str,
    y_label: str,
    y_cap: float | None = None,
    log_y: bool = False,
    style: str = "line",
)

def plot(
    df: pd.DataFrame,
    out_path: Path,
    fmt: str,
    y_col: str,
    y_label: str,
    y_cap: float | None = None,
    log_y: bool = False,
    style: str = "line",
) -> plt.Figure:
    methods = list(df["method"].unique())

    palette = {m: PROVIDER_COLORS[m] for m in methods}
    fig, ax = plt.subplots()

    if style == "bar":
        sns.barplot(
            df,
            x="bsz",
            y=y_col,
            hue="method",
            hue_order=methods,
            palette=palette,
            ax=ax,
        )
        hatches = _apply_hatches(ax, methods)
        if y_cap is not None:
            ax.set_ylim(0, y_cap)
            for container in ax.containers:
                for bar in container:
                    if bar.get_height() > y_cap:
                        true_val = bar.get_height()
                        bar.set_height(y_cap)
                        bx = bar.get_x() + bar.get_width() / 2
                        ax.text(
                            bx,
                            y_cap * 1.01,
                            f"{true_val:.0f}\u00b5s",
                            ha="center",
                            va="bottom",
                            fontsize=9,
                            fontweight="bold",
                            clip_on=False,
                        )
        sns.move_legend(ax, "upper center", title="Method", bbox_to_anchor=(0.5, 1.35), ncol=1)
        for handle, hatch in zip(ax.get_legend().legend_handles, hatches):
            handle.set_hatch(hatch)
        ax.set_xticks(ax.get_xticks(), labels=df["bsz"].unique().astype(int))
    else:
        for method in methods:
            mdf = df.query("method == @method")
            ax.plot(
                mdf["bsz"],
                mdf[y_col],
                color=PROVIDER_COLORS[method],
                marker=PROVIDER_MARKERS.get(method, "o"),
                label=method,
                linewidth=2,