pelinker.plotting

build_fit_cluster_viz_plot_df(report, *, exclude_noise=True, hdbscan_fit_scope=True, cluster_labels=None)

Build a :func:plot_cluster_viz frame from a :class:~pelinker.reporting.ModelSelectionReport.

Source code in pelinker/plotting.py

def build_fit_cluster_viz_plot_df(
    report: ModelSelectionReport,
    *,
    exclude_noise: bool = True,
    hdbscan_fit_scope: bool = True,
    cluster_labels: dict[int, str] | None = None,
) -> tuple[pd.DataFrame | None, str]:
    """Build a :func:`plot_cluster_viz` frame from a :class:`~pelinker.reporting.ModelSelectionReport`."""
    cluster_viz = report.cluster_viz
    if cluster_viz is None or cluster_viz.size == 0 or cluster_viz.shape[1] < 1:
        return None, report.cluster_viz_method
    assign = filter_assignments_for_cluster_viz(
        report.assignments,
        exclude_noise=exclude_noise,
        hdbscan_fit_scope=hdbscan_fit_scope,
    )
    if len(assign) == 0:
        return None, report.cluster_viz_method
    n_dims = int(cluster_viz.shape[1])
    viz_cols = [f"cviz_{j:02d}" for j in range(n_dims)]
    viz_df = pd.DataFrame(
        cluster_viz, columns=viz_cols, index=report.assignments.index
    ).loc[assign.index]
    assign = assign.copy()
    if "cluster" in assign.columns:
        assign["cluster_id"] = assign["cluster"].astype(int)
    rename: dict[str, str] = {}
    if "cluster" in assign.columns:
        rename["cluster"] = "class"
    plot_assign = assign.rename(columns=rename)
    if cluster_labels and "cluster_id" in plot_assign.columns:
        plot_assign["class"] = plot_assign["cluster_id"].map(
            lambda cid: cluster_labels.get(int(cid), str(cid))
        )
    cols = [c for c in ("entity", "class") if c in plot_assign.columns]
    extra = [
        c
        for c in (
            "cluster_id",
            "pmid",
            "mention",
            "a_abs",
            "b_abs",
            "screener_score",
            "cluster_score",
        )
        if c in plot_assign.columns
    ]
    return pd.concat(
        [plot_assign[cols + extra], viz_df], axis=1
    ), report.cluster_viz_method

diagnostics_to_pairgrid_dataframe(diag)

Build a DataFrame for :func:plot_pca_quality_pairgrid from linker fit diagnostics.

class_label is negative / positive from oov_label (1 / 0).

Source code in pelinker/plotting.py

def diagnostics_to_pairgrid_dataframe(diag: LinkerFitDiagnostics) -> pd.DataFrame:
    """
    Build a DataFrame for :func:`plot_pca_quality_pairgrid` from linker fit diagnostics.

    ``class_label`` is ``negative`` / ``positive`` from ``oov_label`` (1 / 0).
    """
    ol = np.asarray(diag.oov_label, dtype=np.int64).ravel()
    return pd.DataFrame(
        {
            "pca_residual": np.asarray(diag.pca_residual, dtype=np.float64),
            "pca_mahalanobis": np.asarray(diag.pca_mahalanobis, dtype=np.float64),
            "pca_spectral_entropy": np.asarray(
                diag.pca_spectral_entropy, dtype=np.float64
            ),
            "class_label": np.where(ol == 1, "negative", "positive"),
        }
    )

enrich_fit_cluster_viz_plot_df_with_context(plot_df, pmid_text_table_path, *, words_before=5, words_after=5)

Add a context column: five words before/after each mention (from a_abs/b_abs).

Source code in pelinker/plotting.py

def enrich_fit_cluster_viz_plot_df_with_context(
    plot_df: pd.DataFrame,
    pmid_text_table_path: str | pathlib.Path,
    *,
    words_before: int = 5,
    words_after: int = 5,
) -> pd.DataFrame:
    """Add a ``context`` column: five words before/after each mention (from ``a_abs``/``b_abs``)."""
    required = {"pmid", "a_abs", "b_abs"}
    if not required.issubset(plot_df.columns):
        return plot_df

    pmid_texts = load_pmid_texts(
        pmid_text_table_path,
        {str(p) for p in plot_df["pmid"].astype(str).unique()},
    )
    contexts: list[str] = []
    for pmid, a_abs, b_abs in zip(
        plot_df["pmid"].astype(str),
        plot_df["a_abs"],
        plot_df["b_abs"],
        strict=True,
    ):
        text = pmid_texts.get(pmid)
        if text is None or pd.isna(a_abs) or pd.isna(b_abs):
            contexts.append("")
            continue
        contexts.append(
            mention_context_window(
                text,
                int(a_abs),
                int(b_abs),
                words_before=words_before,
                words_after=words_after,
            )
        )

    out = plot_df.copy()
    out["context"] = contexts
    return out

filter_assignments_for_cluster_viz(assign, *, exclude_noise=True, hdbscan_fit_scope=True)

Restrict cluster viz rows to HDBSCAN-fit + screener/OOV-pass mentions when flagged.

Source code in pelinker/plotting.py

def filter_assignments_for_cluster_viz(
    assign: pd.DataFrame,
    *,
    exclude_noise: bool = True,
    hdbscan_fit_scope: bool = True,
) -> pd.DataFrame:
    """Restrict cluster viz rows to HDBSCAN-fit + screener/OOV-pass mentions when flagged."""
    from pelinker.cluster_composition_viz import filter_emergent_assignments

    out = assign.copy()
    if exclude_noise:
        out = filter_emergent_assignments(out)
    if not hdbscan_fit_scope:
        return out
    present = [c for c in CLUSTER_VIZ_MEMBERSHIP_COLUMNS if c in out.columns]
    if not present:
        return out
    mask = np.ones(len(out), dtype=bool)
    for col in present:
        mask &= out[col].astype(bool).to_numpy()
    return out.loc[mask]

load_pmid_texts(table_path, pmids, *, chunk_size=10000)

Stream a PMID/text table and return rows for the requested pmids only.

Source code in pelinker/plotting.py

def load_pmid_texts(
    table_path: str | pathlib.Path,
    pmids: set[str],
    *,
    chunk_size: int = 10_000,
) -> dict[str, str]:
    """Stream a PMID/text table and return rows for the requested ``pmids`` only."""
    from pelinker.ops import load_pmid_texts_from_table

    return load_pmid_texts_from_table(table_path, pmids, chunk_size=chunk_size)

mention_context_window(text, a_abs, b_abs, *, words_before=5, words_after=5)

Return a short phrase around [a_abs, b_abs) with the mention marked «…».

Source code in pelinker/plotting.py

def mention_context_window(
    text: str,
    a_abs: int,
    b_abs: int,
    *,
    words_before: int = 5,
    words_after: int = 5,
) -> str:
    """Return a short phrase around ``[a_abs, b_abs)`` with the mention marked «…»."""
    spans = _word_spans(text)
    if not spans:
        return ""

    start_word: int | None = None
    end_word: int | None = None
    for i, (wa, wb) in enumerate(spans):
        if wb <= a_abs:
            continue
        if wa >= b_abs:
            break
        if start_word is None:
            start_word = i
        end_word = i

    if start_word is None:
        for i, (wa, _wb) in enumerate(spans):
            if wa >= a_abs:
                start_word = i
                end_word = i
                break
        if start_word is None:
            return ""

    assert end_word is not None
    win_start = max(0, start_word - words_before)
    win_end = min(len(spans) - 1, end_word + words_after)
    parts: list[str] = []

    for i in range(win_start, win_end + 1):
        wa, wb = spans[i]
        word = text[wa:wb]

        if i == start_word:
            word = f"<b>{word}"

        if i == end_word:
            word = f"{word}</b>"

        parts.append(word)
    return " ".join(parts)

plot_cluster_entity_sankey(flow_df, *, save_dir, basename='fit_cluster_entity_sankey', max_clusters=None, max_entities=None, min_within_cluster_fraction=0.0, cluster_labels=None, cluster_label_include_id=False, inches_per_label=_SANKEY_DEFAULT_INCHES_PER_LABEL, min_fig_height=_SANKEY_DEFAULT_MIN_FIG_HEIGHT, min_band_height=_SANKEY_DEFAULT_MIN_BAND_HEIGHT)

Bipartite KB-in entity→cluster Sankey from a long flow table (cluster, entity, count).

Expects unbundled entity mass (no pie/bar Other (...) rows). Caps by dropping the long tail. pySankey sizes bands by weight only; inches_per_label sets figure height; min_band_height uniformly scales weights for readability.

Source code in pelinker/plotting.py

def plot_cluster_entity_sankey(
    flow_df: pd.DataFrame,
    *,
    save_dir: pathlib.Path,
    basename: str = "fit_cluster_entity_sankey",
    max_clusters: int | None = None,
    max_entities: int | None = None,
    min_within_cluster_fraction: float = 0.0,
    cluster_labels: dict[int, str] | None = None,
    cluster_label_include_id: bool = False,
    inches_per_label: float = _SANKEY_DEFAULT_INCHES_PER_LABEL,
    min_fig_height: float = _SANKEY_DEFAULT_MIN_FIG_HEIGHT,
    min_band_height: float = _SANKEY_DEFAULT_MIN_BAND_HEIGHT,
) -> list[pathlib.Path]:
    """
    Bipartite KB-in entity→cluster Sankey from a long flow table (cluster, entity, count).

    Expects unbundled entity mass (no pie/bar ``Other (...)`` rows). Caps by dropping
    the long tail. pySankey sizes bands by weight only; ``inches_per_label`` sets
    figure height; ``min_band_height`` uniformly scales weights for readability.
    """
    if flow_df.empty:
        return []
    from pelinker.cluster_composition_viz import limit_entity_flow_for_plots

    work = limit_entity_flow_for_plots(
        flow_df,
        max_clusters=max_clusters,
        max_entities=max_entities,
        min_within_cluster_fraction=min_within_cluster_fraction,
    )
    if work.empty:
        return []
    work = work.copy()
    work["entity"] = work["entity"].astype(str)
    if cluster_labels:
        work["cluster"] = (
            work["cluster"]
            .astype(int)
            .map(
                lambda cid: _format_cluster_sankey_label(
                    int(cid),
                    cluster_labels,
                    include_id=cluster_label_include_id,
                )
            )
        )
    else:
        work["cluster"] = work["cluster"].astype(str)
    left = work["entity"].to_numpy()
    right = work["cluster"].to_numpy()
    weights = work["count"].astype(float).to_numpy()
    weights = _scale_sankey_weights_for_min_band(
        work,
        weights,
        left_column="entity",
        min_band_height=min_band_height,
    )

    n_labels = max(work["entity"].nunique(), work["cluster"].nunique())
    fig_height = max(min_fig_height, n_labels * inches_per_label)

    # pySankey always calls plt.figure() internally; figure_name is a path string, not a Figure.
    sankey(
        left,
        right,
        leftWeight=weights,
        aspect=12,
        fontsize=10,
        closePlot=False,
    )
    fig = plt.gcf()
    fig.set_size_inches(14, fig_height)
    written: list[pathlib.Path] = []
    for path in _save_figure_multi_format(fig, save_dir / basename):
        written.append(path)
    plt.close(fig)
    return written

plot_dbcv_vs_ari_from_grid(df_grid, output_path, *, optimization_config=None, grid_cluster_count_reward=None, grid_n_entities=None, grid_objective=None, optimization_method=None)

Scatter of mean DBCV vs mean ARI per (model, layer); shape = arity (△/□/○), fill colors = base encoder model(s); text = layer spec only (e.g. fusion 2+3). 95% covariance ellipses when n_sample ≥ 2.

Uses (dbcv, ari) at chosen_min_cluster_size for each bootstrap sample_idx (same hyperparameter as the vertical line on per-combination error-bar plots). Both axes are fixed to [0, _AXIS_MAX].

When any solver argument is set (optimization_config or grid override kwargs), chosen_min_cluster_size is re-computed per (model, layer) from the grid metrics instead of using values stored in df_grid.

Returns:

Type	Description
bool	True if a figure was written, False if required data were absent.

Source code in pelinker/plotting.py

def plot_dbcv_vs_ari_from_grid(
    df_grid: pd.DataFrame,
    output_path: pathlib.Path,
    *,
    optimization_config: ClusteringOptimizationConfig | None = None,
    grid_cluster_count_reward: float | None = None,
    grid_n_entities: int | None = None,
    grid_objective: GridObjectiveSpec | None = None,
    optimization_method: str | None = None,
) -> bool:
    """
    Scatter of mean DBCV vs mean ARI per (model, layer); shape = arity (△/□/○),
    fill colors = base encoder model(s); text = layer spec only (e.g. fusion ``2+3``).
    95% covariance ellipses when ``n_sample`` ≥ 2.

    Uses ``(dbcv, ari)`` at ``chosen_min_cluster_size`` for each bootstrap ``sample_idx``
    (same hyperparameter as the vertical line on per-combination error-bar plots).
    Both axes are fixed to ``[0, _AXIS_MAX]``.

    When any solver argument is set (``optimization_config`` or grid override kwargs),
    ``chosen_min_cluster_size`` is re-computed per (model, layer) from the grid metrics
    instead of using values stored in ``df_grid``.

    Returns:
        True if a figure was written, False if required data were absent.
    """
    if not has_grid_points_for_dbcv_ari_scatter(df_grid):
        return False

    df_plot = df_grid
    if grid_solver_overrides_active(
        optimization_config=optimization_config,
        grid_cluster_count_reward=grid_cluster_count_reward,
        grid_n_entities=grid_n_entities,
        grid_objective=grid_objective,
        optimization_method=optimization_method,
    ):
        chosen_by_combo = resolve_chosen_min_cluster_size_by_combo_from_grid(
            df_grid,
            optimization_config,
            grid_cluster_count_reward=grid_cluster_count_reward,
            grid_n_entities=grid_n_entities,
            grid_objective=grid_objective,
            optimization_method=optimization_method,
        )
        if chosen_by_combo:
            df_plot = apply_chosen_min_cluster_size_to_grid(df_grid, chosen_by_combo)

    try:
        df = select_grid_points_at_chosen_min_cluster_size(df_plot)
    except ValueError:
        return False
    if df.empty:
        return False

    all_models: list[str] = []
    for m, lyr in df[["model", "layer"]].drop_duplicates().itertuples(index=False):
        all_models.extend(_base_models_in_row(str(m), str(lyr)))
    color_by_model = _model_color_map(all_models)

    fig, ax = plt.subplots(figsize=(8.5, 8.5))
    arities_present: set[str] = set()

    for (model, layer), g in df.groupby(["model", "layer"], sort=False):
        xy = np.column_stack(
            [
                g["dbcv"].to_numpy(dtype=np.float64),
                g["ari"].to_numpy(dtype=np.float64),
            ]
        )
        mean = xy.mean(axis=0)
        n = xy.shape[0]
        arity = _arity_from_model(str(model))
        arities_present.add(arity)
        models_row = _base_models_in_row(str(model), str(layer))

        # Halo drawn first: light filled ellipse + clear dashed rim, slightly inflated so
        # it remains visible through transparent markers.
        if n >= 2:
            cov = np.cov(xy, rowvar=False, ddof=1)
            ell = _covariance_ellipse_95(cov)
            if ell is not None:
                w, h, ang = ell
                wi, hi = w * _ELLIPSE_INFLATE, h * _ELLIPSE_INFLATE
                patch = Ellipse(
                    xy=(float(mean[0]), float(mean[1])),
                    width=wi,
                    height=hi,
                    angle=ang,
                    facecolor=(0.45, 0.48, 0.52, _ELLIPSE_FILL_ALPHA),
                    edgecolor=(0.12, 0.14, 0.18, _ELLIPSE_EDGE_ALPHA),
                    linewidth=1.45,
                    linestyle=(0, (4.5, 3.0)),
                    zorder=2,
                )
                ax.add_patch(patch)

        _draw_arity_marker(
            ax,
            float(mean[0]),
            float(mean[1]),
            arity=arity,
            models=models_row,
            color_by_model=color_by_model,
            zorder=5,
        )
        layer_code = _layer_spec_code(str(model), str(layer))
        if len(layer_code) > 22:
            layer_code = layer_code[:19] + "…"
        ax.annotate(
            layer_code,
            (mean[0], mean[1]),
            textcoords="offset points",
            # xytext=(8, 8),
            xytext=(-3 - 1.5 * (len(layer_code) - 1), -3),
            fontsize=8,
            alpha=0.88,
            zorder=6,
        )

    ax.set_xlim(0.0, _AXIS_MAX)
    ax.set_ylim(0.0, _AXIS_MAX)
    ax.set_aspect("equal")
    ax.set_xlabel("DBCV at chosen min_cluster_size (mean over samples)")
    ax.set_ylabel("ARI at chosen min_cluster_size (mean over samples)")
    ax.set_title(
        "DBCV vs ARI at consensus hyperparameter; dashed ellipse ≈95% (n_sample >= 2)"
    )

    order_a = ["singleton", "fusion2", "fusion3"]
    arity_labels = {
        "singleton": "singleton",
        "fusion2": "pair fusion",
        "fusion3": "triple fusion",
    }
    edge_legend = (0.0, 0.0, 0.0, _MARKER_OUTLINE_ALPHA)
    legend_shapes = [
        Line2D(
            [0],
            [0],
            marker=ARITY_MARKER_SCATTER[a],
            color="none",
            label=arity_labels[a],
            markerfacecolor=_rgba("#bbbbbb", _MARKER_FACE_ALPHA),
            markeredgecolor=edge_legend,
            markersize=10,
        )
        for a in order_a
        if a in arities_present
    ]
    legend_colors = [
        Patch(
            facecolor=_rgba(color_by_model[m], _MARKER_FACE_ALPHA),
            edgecolor=edge_legend,
            linewidth=0.5,
            label=m,
        )
        for m in sorted(color_by_model.keys())
    ]
    if legend_shapes or legend_colors:
        leg1 = ax.legend(
            handles=legend_shapes,
            title="Arity",
            loc="upper left",
            bbox_to_anchor=(1.02, 1.0),
            borderaxespad=0.0,
            frameon=True,
        )
        ax.add_artist(leg1)
        ax.legend(
            handles=legend_colors,
            title="Base model",
            loc="lower left",
            bbox_to_anchor=(1.02, 0.0),
            borderaxespad=0.0,
            frameon=True,
        )

    ax.grid(True, alpha=0.28, linestyle="--", zorder=0)
    plt.tight_layout()
    _save_figure_multi_format(fig, output_path)
    plt.close(fig)
    return True

plot_heatmap(df_results, output_path, metric='best_score', metric_label=None, *, secondary_metric='best_size')

Create a heatmap with model (rows) and layer (columns). Color represents the specified metric; text overlays secondary_metric and metric value.

Parameters:

Name	Type	Description	Default
df_results	DataFrame	DataFrame with columns: model, layer, …	required
output_path	Path	Path to save the heatmap figure	required
metric	str	Column name for the metric to display as color (default: "best_score")	'best_score'
metric_label	str | None	Label for the metric (default: uses metric column name)	None
secondary_metric	str | None	Column for text annotation besides the metric cell (default: "best_size"); use None to annotate metric value only.	'best_size'

Source code in pelinker/plotting.py

def plot_heatmap(
    df_results: pd.DataFrame,
    output_path: pathlib.Path,
    metric: str = "best_score",
    metric_label: str | None = None,
    *,
    secondary_metric: str | None = "best_size",
):
    """
    Create a heatmap with model (rows) and layer (columns).
    Color represents the specified metric; text overlays ``secondary_metric`` and metric value.

    Args:
        df_results: DataFrame with columns: model, layer, …
        output_path: Path to save the heatmap figure
        metric: Column name for the metric to display as color (default: "best_score")
        metric_label: Label for the metric (default: uses metric column name)
        secondary_metric: Column for text annotation besides the metric cell (default:
            ``"best_size"``); use ``None`` to annotate metric value only.
    """
    if metric_label is None:
        metric_label = metric.replace("_", " ").title()

    # Create pivot tables
    score_pivot = df_results.pivot(index="model", columns="layer", values=metric)
    size_pivot: pd.DataFrame | None = None
    if secondary_metric is not None and secondary_metric in df_results.columns:
        size_pivot = df_results.pivot(
            index="model", columns="layer", values=secondary_metric
        )

    # Create figure
    fig, ax = plt.subplots(
        figsize=(
            max(8, len(score_pivot.columns) * 0.8),
            max(6, len(score_pivot.index) * 0.6),
        )
    )

    # Create heatmap with metric as color
    # Use RdBu_r (Red-Blue reversed) for clear visual distinction: red=high, blue=low
    sns.heatmap(
        score_pivot,
        annot=False,  # We'll add custom annotations
        fmt=".3f",
        cmap="RdBu_r",
        center=None,  # Center colormap at the median for better contrast
        cbar_kws={"label": metric_label, "shrink": 0.8},
        ax=ax,
        linewidths=0.5,
        linecolor="white",
        square=False,
    )

    # Add best_size and metric name as text annotations
    # Calculate mean score for text color threshold
    valid_scores = score_pivot.values[~pd.isna(score_pivot.values)]
    mean_score = valid_scores.mean() if len(valid_scores) > 0 else 0

    for i in range(len(score_pivot.index)):
        for j in range(len(score_pivot.columns)):
            score_val = score_pivot.iloc[i, j]
            if pd.isna(score_val):
                continue
            secondary_ok = True
            size_val: float | None = None
            if size_pivot is not None:
                sv = size_pivot.iloc[i, j]
                if pd.isna(sv):
                    secondary_ok = False
                else:
                    size_val = float(sv)

            if not secondary_ok and size_pivot is not None:
                continue

            text_color = "white" if score_val < mean_score else "black"
            if abs(score_val) < 0.01:
                metric_str = f"{score_val:.2e}"
            elif abs(score_val) < 1:
                metric_str = f"{score_val:.3f}"
            else:
                metric_str = f"{score_val:.2f}"

            if size_val is None:
                anno = metric_str
            else:
                size_display = (
                    str(int(round(size_val)))
                    if secondary_metric == "best_size"
                    else f"{size_val:.2f}"
                )
                anno = f"{size_display}\n{metric_str}"

            ax.text(
                j + 0.5,
                i + 0.5,
                anno,
                ha="center",
                va="center",
                color=text_color,
                fontweight="bold",
                fontsize=8,
                linespacing=1.2,
            )

    sub_label = secondary_metric.replace("_", " ").title() if secondary_metric else None
    if sub_label:
        ax.set_title(f"Results: {metric_label} (color) and {sub_label} + value (text)")
    else:
        ax.set_title(f"Results: {metric_label} (color)")
    ax.set_xlabel("Layer")
    ax.set_ylabel("Model")

    plt.tight_layout()
    _save_figure_multi_format(fig, output_path)
    plt.close()

plot_metrics_with_error_bars(metrics_list, output_path, *, chosen_min_cluster_size=None, grid_solve=None, optimization_config=None, grid_cluster_count_reward=None, grid_n_entities=None, grid_objective=None, optimization_method=None)

Plot metrics across multiple runs with error bars using seaborn lineplot.

Parameters:

Name	Type	Description	Default
metrics_list	list[DataFrame]	List of DataFrames, each with columns: min_cluster_size, icm, n_clusters, dbcv, ari	required
output_path	Path	Path to save the figure	required
chosen_min_cluster_size	float | None	Optional vertical marker for the selected grid value (e.g. from smoother / argmax).	None
grid_solve	SmoothedGridOptimumResult | None	Precomputed pooled grid diagnostics (avoids a second solve; drives objective panel).	None
optimization_config	ClusteringOptimizationConfig | None	When set (or when any grid override kwarg is set and chosen_min_cluster_size is omitted), re-run the pooled grid solver for the vertical marker.	None
grid_cluster_count_reward	float | None	Override :attr:~pelinker.config.ClusteringOptimizationConfig.grid_cluster_count_reward.	None
grid_n_entities	int | None	Override :attr:~pelinker.config.ClusteringOptimizationConfig.grid_n_entities.	None
grid_objective	GridObjectiveSpec | None	Override :attr:~pelinker.config.ClusteringOptimizationConfig.grid_objective.	None
optimization_method	str | None	Override :attr:~pelinker.config.ClusteringOptimizationConfig.optimization_method.	None

Source code in pelinker/plotting.py

def plot_metrics_with_error_bars(
    metrics_list: list[pd.DataFrame],
    output_path: pathlib.Path,
    *,
    chosen_min_cluster_size: float | None = None,
    grid_solve: SmoothedGridOptimumResult | None = None,
    optimization_config: ClusteringOptimizationConfig | None = None,
    grid_cluster_count_reward: float | None = None,
    grid_n_entities: int | None = None,
    grid_objective: GridObjectiveSpec | None = None,
    optimization_method: str | None = None,
):
    """
    Plot metrics across multiple runs with error bars using seaborn lineplot.

    Args:
        metrics_list: List of DataFrames, each with columns: min_cluster_size, icm, n_clusters, dbcv, ari
        output_path: Path to save the figure
        chosen_min_cluster_size: Optional vertical marker for the selected grid value (e.g. from smoother / argmax).
        grid_solve: Precomputed pooled grid diagnostics (avoids a second solve; drives objective panel).
        optimization_config: When set (or when any grid override kwarg is set and ``chosen_min_cluster_size``
            is omitted), re-run the pooled grid solver for the vertical marker.
        grid_cluster_count_reward: Override :attr:`~pelinker.config.ClusteringOptimizationConfig.grid_cluster_count_reward`.
        grid_n_entities: Override :attr:`~pelinker.config.ClusteringOptimizationConfig.grid_n_entities`.
        grid_objective: Override :attr:`~pelinker.config.ClusteringOptimizationConfig.grid_objective`.
        optimization_method: Override :attr:`~pelinker.config.ClusteringOptimizationConfig.optimization_method`.
    """
    if grid_solve is None and should_resolve_chosen_min_cluster_size(
        chosen_min_cluster_size=chosen_min_cluster_size,
        optimization_config=optimization_config,
        grid_cluster_count_reward=grid_cluster_count_reward,
        grid_n_entities=grid_n_entities,
        grid_objective=grid_objective,
        optimization_method=optimization_method,
    ):
        grid_solve = solve_pooled_grid_from_metrics_list(
            metrics_list,
            optimization_config,
            grid_cluster_count_reward=grid_cluster_count_reward,
            grid_n_entities=grid_n_entities,
            grid_objective=grid_objective,
            optimization_method=optimization_method,
        )

    if grid_solve is not None and chosen_min_cluster_size is None:
        chosen_min_cluster_size = float(grid_solve.chosen_min_cluster_size)

    # Combine all metrics DataFrames, adding a run_id column
    combined_metrics = []
    for run_id, df in enumerate(metrics_list):
        df_copy = df.copy()
        df_copy["run_id"] = run_id
        combined_metrics.append(df_copy)

    df_combined = pd.concat(combined_metrics, ignore_index=True)

    # Filter out trivial points where n_clusters <= 1
    df_combined = df_combined[df_combined["n_clusters"] > 1].copy()

    if len(df_combined) == 0:
        print(
            f"Warning: No valid data points after filtering (n_clusters > 1) for {output_path}"
        )
        return

    has_ari = "ari" in df_combined.columns and bool(df_combined["ari"].notna().any())
    colors = ["#2E86AB", "#A23B72", "#C44E52", "#F18F01"]  # Blue, Purple, Red, Orange

    if has_ari:
        fig, axes = plt.subplots(1, 4, figsize=(24, 5))
        ax_dbcv, ax_ari, ax_k, ax_obj = axes[0], axes[1], axes[2], axes[3]
    else:
        fig, axes = plt.subplots(1, 3, figsize=(18, 5))
        ax_dbcv, ax_k, ax_obj = axes[0], axes[1], axes[2]
        ax_ari = None

    def _maybe_vline(ax: plt.Axes) -> None:
        if chosen_min_cluster_size is None:
            return
        ax.axvline(
            chosen_min_cluster_size,
            color="0.35",
            linestyle="--",
            linewidth=1.5,
            alpha=0.9,
            zorder=0,
        )

    sns.lineplot(
        data=df_combined,
        x="min_cluster_size",
        y="dbcv",
        ax=ax_dbcv,
        errorbar="sd",
        marker="o",
        color=colors[0],
        linewidth=2,
        markersize=8,
        err_kws={"alpha": 0.3, "linewidth": 1.5},
    )
    _maybe_vline(ax_dbcv)
    _style_metrics_axis(
        ax_dbcv,
        ylabel="DBCV Score",
        title="DBCV vs. min_cluster_size",
        color=colors[0],
    )

    if ax_ari is not None:
        sns.lineplot(
            data=df_combined,
            x="min_cluster_size",
            y="ari",
            ax=ax_ari,
            errorbar="sd",
            marker="D",
            color=colors[1],
            linewidth=2,
            markersize=7,
            err_kws={"alpha": 0.3, "linewidth": 1.5},
        )
        _maybe_vline(ax_ari)
        _style_metrics_axis(
            ax_ari,
            ylabel="ARI",
            title="ARI vs. min_cluster_size",
            color=colors[1],
        )

    sns.lineplot(
        data=df_combined,
        x="min_cluster_size",
        y="n_clusters",
        ax=ax_k,
        errorbar="sd",
        marker="^",
        color=colors[2],
        linewidth=2,
        markersize=8,
        err_kws={"alpha": 0.3, "linewidth": 1.5},
    )
    _maybe_vline(ax_k)
    _style_metrics_axis(
        ax_k,
        ylabel="n clusters",
        title="Number of Clusters vs. min_cluster_size",
        color=colors[2],
    )

    obj_color = colors[3]
    if grid_solve is not None and len(grid_solve.x) > 0:
        x_obj = np.array(grid_solve.x, dtype=np.float64)
        y_obj = np.array(grid_solve.y_objective, dtype=np.float64)
        y_smooth = np.array(grid_solve.y_smooth, dtype=np.float64)
        ax_obj.plot(
            x_obj,
            y_obj,
            marker="s",
            color=obj_color,
            linewidth=2,
            markersize=7,
            label="objective",
            zorder=2,
        )
        if np.any(np.isfinite(y_smooth)):
            ax_obj.plot(
                x_obj,
                y_smooth,
                linestyle="--",
                color=obj_color,
                linewidth=1.5,
                alpha=0.65,
                label="smoothed",
                zorder=1,
            )
        has_cluster_term = any(
            abs(v) > 1e-12 for v in grid_solve.y_cluster_term if np.isfinite(v)
        )
        obj_title = (
            "Grid objective (pooled + cluster penalty)"
            if has_cluster_term
            else "Grid objective (pooled)"
        )
    else:
        obj_title = "Grid objective (unavailable)"
    _maybe_vline(ax_obj)
    _style_metrics_axis(
        ax_obj, ylabel="Grid objective", title=obj_title, color=obj_color
    )

    plt.tight_layout()
    _save_figure_multi_format(fig, output_path)
    plt.close(fig)

plot_pca_quality_pairgrid(df, output_path, *, class_col='class_label', subtitle=None, max_scatter_points=4000, max_kde_points=20000)

PairGrid of PCA quality features (pca_residual, pca_spectral_entropy, pca_mahalanobis) with hue for class (negative / positive).

Upper triangle: balanced class subsample (equal cap per class) with distinct markers. Lower triangle & diagonal: KDE per hue with common_norm=False so each class is normalized on its own scale (not pooled against the majority class).

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with feature columns and class_col.	required
output_path	Path	PNG path (PDF sibling also written).	required
class_col	str	Column used for hue.	'class_label'
subtitle	str | None	Combo label for the figure suptitle (e.g. model/layer/sample).	None
max_scatter_points	int	Cap for upper-triangle scatter layer (split evenly by class).	4000
max_kde_points	int	Cap for lower-triangle KDE and diagonal KDE layers.	20000

Source code in pelinker/plotting.py

def plot_pca_quality_pairgrid(
    df: pd.DataFrame,
    output_path: pathlib.Path,
    *,
    class_col: str = "class_label",
    subtitle: str | None = None,
    max_scatter_points: int = 4_000,
    max_kde_points: int = 20_000,
) -> bool:
    """
    PairGrid of PCA quality features (pca_residual, pca_spectral_entropy, pca_mahalanobis)
    with hue for class (negative / positive).

    Upper triangle: balanced class subsample (equal cap per class) with distinct markers.
    Lower triangle & diagonal: KDE per hue with ``common_norm=False`` so each class is
    normalized on its own scale (not pooled against the majority class).

    Args:
        df: DataFrame with feature columns and class_col.
        output_path: PNG path (PDF sibling also written).
        class_col: Column used for hue.
        subtitle: Combo label for the figure suptitle (e.g. model/layer/sample).
        max_scatter_points: Cap for upper-triangle scatter layer (split evenly by class).
        max_kde_points: Cap for lower-triangle KDE and diagonal KDE layers.
    """
    feature_cols = ["pca_residual", "pca_spectral_entropy", "pca_mahalanobis"]
    needed = set(feature_cols + [class_col])
    if not needed.issubset(df.columns):
        return False

    plot_df = df.loc[:, feature_cols + [class_col]].dropna().copy()
    if plot_df.empty:
        return False

    scatter_df = _balanced_subsample_by_class(plot_df, max_scatter_points, class_col)
    kde_df = _balanced_subsample_by_class(plot_df, max_kde_points, class_col)

    classes = sorted(plot_df[class_col].unique())
    palette = sns.color_palette("Set2", n_colors=len(classes))
    color_map = dict(zip(classes, palette, strict=True))
    marker_map = {
        cls: _PCA_CLASS_MARKERS.get(str(cls), ("o", "^")[i % 2])
        for i, cls in enumerate(classes)
    }
    class_sizes = plot_df[class_col].value_counts()
    # Draw majority first, minority on top in scatter panels.
    classes_by_size = sorted(classes, key=lambda c: int(class_sizes.get(c, 0)))

    g = sns.PairGrid(
        kde_df, vars=feature_cols, hue=class_col, palette=color_map, diag_sharey=False
    )
    g.map_lower(
        sns.kdeplot,
        fill=True,
        alpha=0.4,
        thresh=0.05,
        common_norm=False,
        legend=False,
    )
    g.map_diag(sns.kdeplot, lw=2, common_norm=False, fill=False, legend=False)

    n_vars = len(feature_cols)
    for row in range(n_vars):
        for col in range(n_vars):
            if col <= row:
                continue
            ax = g.axes[row, col]
            x_col = feature_cols[col]
            y_col = feature_cols[row]
            for z, cls in enumerate(classes_by_size):
                sub = scatter_df.loc[scatter_df[class_col] == cls]
                if sub.empty:
                    continue
                ax.scatter(
                    sub[x_col].to_numpy(),
                    sub[y_col].to_numpy(),
                    color=color_map[cls],
                    marker=marker_map[cls],
                    alpha=0.55,
                    s=16,
                    linewidths=0.35,
                    edgecolors="white",
                    rasterized=True,
                    zorder=z + 1,
                )

    legend_handles = [
        Line2D(
            [0],
            [0],
            marker=marker_map[cls],
            color="w",
            markerfacecolor=color_map[cls],
            markeredgecolor="white",
            markeredgewidth=0.35,
            markersize=7,
            label=str(cls),
            linestyle="",
        )
        for cls in classes
    ]
    g.axes[0, 0].legend(
        handles=legend_handles, title="class", loc="best", framealpha=0.9
    )

    g.figure.suptitle(
        _pca_pairgrid_title(
            subtitle=subtitle,
            plot_df=plot_df,
            class_col=class_col,
            scatter_df=scatter_df,
            kde_df=kde_df,
        ),
        y=1.02,
        fontsize=10,
    )
    g.figure.tight_layout()
    _save_figure_multi_format(g.figure, output_path)
    plt.close(g.figure)
    return True

plot_roc_comparison(scores_df, output_path, *, combo_keys)

ROC curves for screener_best_score, oov_score, and combined_score pooled over samples per combo_key.

Source code in pelinker/plotting.py

def plot_roc_comparison(
    scores_df: pd.DataFrame,
    output_path: pathlib.Path,
    *,
    combo_keys: list[str],
) -> bool:
    """
    ROC curves for ``screener_best_score``, ``oov_score``, and ``combined_score``
    pooled over samples per ``combo_key``.
    """
    if scores_df.empty or not combo_keys:
        return False
    from sklearn.metrics import auc as sk_auc, roc_curve

    avail = scores_df.loc[scores_df["combo_key"].isin(combo_keys)]
    usable = avail["combo_key"].unique().tolist()
    if not usable:
        return False

    cols = {"y_true", "screener_best_score", "oov_score", "combined_score"}
    if not cols.issubset(scores_df.columns):
        return False

    n_p = len(usable)
    ncols = min(3, n_p)
    nrows = int(np.ceil(n_p / ncols)) if n_p else 1
    # ncells = max(1, nrows * ncols)
    fig, axes = plt.subplots(nrows, ncols, figsize=(5 * ncols, 4.5 * nrows))
    axes_flat = np.atleast_1d(axes).ravel()

    for ix, ck in enumerate(usable):
        if ix >= axes_flat.shape[0]:
            break
        ax = axes_flat[ix]
        sub = scores_df.loc[scores_df["combo_key"] == ck].copy()
        y = np.asarray(sub["y_true"], dtype=np.int64)
        if np.unique(y).size < 2:
            ax.set_visible(False)
            continue
        for name, serie, ls in (
            ("Screener", sub["screener_best_score"], "-"),
            ("OOV", sub["oov_score"], "--"),
            ("Combined", sub["combined_score"], "-."),
        ):
            s = np.asarray(serie, dtype=np.float64)
            try:
                fpr, tpr, _ = roc_curve(y, s)
                a = float(sk_auc(fpr, tpr))
            except ValueError:
                continue
            ax.plot(fpr, tpr, ls=ls, lw=2, label=f"{name} (AUC={a:.3f})")

        ax.plot([0, 1], [0, 1], "k:", lw=1, alpha=0.35)
        ax.set_xlim(0.0, 1.0)
        ax.set_ylim(0.0, 1.05)
        ax.set_xlabel("FPR")
        ax.set_ylabel("TPR")
        ax.legend(fontsize=7, loc="lower right")
        ax.set_title(str(ck)[:60], fontsize=9)

    for k in range(n_p, len(axes_flat)):
        axes_flat[k].set_visible(False)

    plt.tight_layout()
    _save_figure_multi_format(fig, output_path)
    plt.close()
    return True

plot_scale_curve(curve, output_path, *, formats=('png', 'pdf'))

Log-log plot of chosen min_cluster_size vs realized N, with the fitted law.

Rungs whose chosen value sat on a grid bound are drawn hollow and red: on those points the grid, not the sample size, decided the answer, and the slope through them is not evidence.

Source code in pelinker/plotting.py

def plot_scale_curve(
    curve: ScaleCurve,
    output_path: pathlib.Path,
    *,
    formats: tuple[str, ...] = ("png", "pdf"),
) -> tuple[pathlib.Path, ...]:
    """Log-log plot of chosen ``min_cluster_size`` vs realized N, with the fitted law.

    Rungs whose chosen value sat on a grid bound are drawn hollow and red: on those
    points the grid, not the sample size, decided the answer, and the slope through
    them is not evidence.
    """
    ns = np.array([r.n_rows_realized for r in curve.rungs], dtype=np.float64)
    mcs = np.array([r.chosen_min_cluster_size for r in curve.rungs], dtype=np.float64)
    pinned = np.array([r.is_pinned for r in curve.rungs], dtype=bool)

    fig, ax = plt.subplots(figsize=(7.0, 5.0))

    # Fitted law across the measured span, extended by 2x to show the extrapolation.
    x_line = np.logspace(
        np.log10(ns.min()), np.log10(ns.max() * 2.0), num=100, dtype=np.float64
    )
    y_line = np.exp(curve.log_intercept + curve.log_slope * np.log(x_line))
    ax.plot(
        x_line,
        y_line,
        color="#3b6ea5",
        linewidth=1.8,
        label=(
            f"fit: log(MCS) = {curve.log_intercept:.3f} + "
            f"{curve.log_slope:.3f}·log(N)   (R²={curve.r_squared:.3f})"
        ),
    )
    ax.axvspan(
        ns.max(),
        x_line.max(),
        color="#999999",
        alpha=0.12,
        label="extrapolation (beyond measured N)",
    )

    if np.any(~pinned):
        ax.scatter(
            ns[~pinned],
            mcs[~pinned],
            s=70,
            color="#3b6ea5",
            edgecolor="black",
            zorder=3,
            label="measured rung",
        )
    if np.any(pinned):
        ax.scatter(
            ns[pinned],
            mcs[pinned],
            s=90,
            facecolor="none",
            edgecolor="#c0392b",
            linewidth=2.0,
            zorder=4,
            label="pinned on grid bound (not evidence)",
        )

    lo = curve.rungs[0].grid_min_scale
    hi = curve.rungs[0].grid_max_scale
    if lo is not None:
        ax.axhline(lo, color="#c0392b", linestyle=":", linewidth=1.0, alpha=0.7)
    if hi is not None:
        ax.axhline(hi, color="#c0392b", linestyle=":", linewidth=1.0, alpha=0.7)

    ax.set_xscale("log")
    ax.set_yscale("log")
    ax.set_xlabel("mention rows clustered (realized N)")
    ax.set_ylabel("chosen min_cluster_size")
    ax.set_title("min_cluster_size scaling with sample size")
    ax.grid(True, which="both", alpha=0.25, linewidth=0.5)
    ax.legend(loc="best", fontsize=8, framealpha=0.9)

    written = _save_figure_multi_format(fig, output_path, formats=formats)
    plt.close(fig)
    return written

plot_screener_oov_bar(summary_df, output_path)

Grouped bar chart per (model, layer): mean AUC for screener / OOV / combined.

Requires screener_auc_mean, oov_auc_mean, combined_auc_mean. Single-embedding rows only (excludes fusion* model labels).

Source code in pelinker/plotting.py

def plot_screener_oov_bar(
    summary_df: pd.DataFrame,
    output_path: pathlib.Path,
) -> bool:
    """
    Grouped bar chart per (model, layer): mean AUC for screener / OOV / combined.

    Requires ``screener_auc_mean``, ``oov_auc_mean``, ``combined_auc_mean``.
    Single-embedding rows only (excludes fusion* model labels).
    """
    required = ("screener_auc_mean", "oov_auc_mean", "combined_auc_mean")
    if not all(c in summary_df.columns for c in required):
        return False
    df = (
        summary_df[~summary_df["model"].astype(str).str.startswith("fusion")]
        .dropna(subset=list(required))
        .copy()
    )
    if df.empty:
        return False
    df["_label"] = (
        df["model"].astype(str) + " / " + df["layer"].astype(str).str.slice(0, 24)
    )
    df = df.sort_values("combined_auc_mean", ascending=False)
    df = df.reset_index(drop=True)
    labels = df["_label"].tolist()
    n = len(labels)
    x = np.arange(n)
    w = 0.25

    ys = df["screener_auc_mean"].to_numpy()
    yo = df["oov_auc_mean"].to_numpy()
    yc = df["combined_auc_mean"].to_numpy()

    errs = []
    for col_std in ("screener_auc_std", "oov_auc_std", "combined_auc_std"):
        if col_std in df.columns:
            errs.append(df[col_std].fillna(0.0).to_numpy())
        else:
            errs.append(np.zeros(n))

    colors = ["#4C72B0", "#55A868", "#C44E52"]

    all_vals = np.concatenate([ys, yo, yc])
    all_errs = np.concatenate([errs[0], errs[1], errs[2]])
    data_min = max(0.0, (all_vals - all_errs).min())
    data_max = min(1.0, (all_vals + all_errs).max())
    margin = max(0.02, (data_max - data_min) * 0.15)
    y_lo = max(0.0, data_min - margin)
    y_hi = min(1.0, data_max + margin)

    fig, ax = plt.subplots(figsize=(max(10, n * 0.55), 6))
    ax.bar(
        x - w,
        ys,
        width=w,
        yerr=errs[0],
        capsize=3,
        label="Screener AUC",
        color=colors[0],
        edgecolor="white",
        linewidth=0.5,
        error_kw={"elinewidth": 1.0, "ecolor": "dimgray", "capthick": 1.0},
    )
    ax.bar(
        x,
        yo,
        width=w,
        yerr=errs[1],
        capsize=3,
        label="OOV AUC",
        color=colors[1],
        edgecolor="white",
        linewidth=0.5,
        error_kw={"elinewidth": 1.0, "ecolor": "dimgray", "capthick": 1.0},
    )
    ax.bar(
        x + w,
        yc,
        width=w,
        yerr=errs[2],
        capsize=3,
        label="Combined AUC",
        color=colors[2],
        edgecolor="white",
        linewidth=0.5,
        error_kw={"elinewidth": 1.0, "ecolor": "dimgray", "capthick": 1.0},
    )
    ax.set_xticks(x)
    ax.set_xticklabels(labels, rotation=65, ha="right", fontsize=7)
    ax.set_ylim(y_lo, y_hi)
    ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda v, _: f"{v:.3f}"))
    ax.grid(axis="y", linestyle="--", linewidth=0.5, alpha=0.6, color="gray")
    ax.set_axisbelow(True)
    ax.legend(loc="upper right", framealpha=0.9)
    ax.set_ylabel("AUC")
    ax.set_title("Mean AUC: screener / OOV / combined (± std)")
    plt.tight_layout()
    _save_figure_multi_format(fig, output_path)
    plt.close()
    return True