plot_topk_overlap_heatmap

Plot pairwise top-k overlap between GRADIEND models.

Parameters:

Name	Type	Description	Default
models	Dict[str, object]	Mapping of display/model ids to GRADIEND models.	required
topk	Union[int, float]	Number or fraction of top weights to compare.	1000
part	str	Model part used to select weights.	'decoder-weight'
value	str	Overlap value to display.	'intersection'
order	Union[str, List[str]]	"input", "cluster", or explicit model order.	'input'
cluster	bool	Whether to cluster rows/columns.	False
annot	Union[bool, str]	Whether/how to annotate cells.	'auto'
fmt	Optional[str]	Deprecated alias for annot_fmt.	None
annot_fmt	Optional[str]	Cell annotation format.	None
figsize	Optional[Tuple[float, float]]	Optional figure size.	None
cmap	str	Heatmap colormap.	'viridis'
vmin	Optional[float]	Optional lower value bound.	None
vmax	Optional[float]	Optional upper value bound.	None
title	Optional[Union[str, bool]]	Optional plot title, or False to omit.	False
output_path	Optional[str]	Optional output path.	None
show	bool	Whether to display the plot.	True
return_data	bool	Whether to include computed data in the return payload.	True
return_fig_ax	bool	Whether to include matplotlib figure/axis.	False
ax	Optional[object]	Optional existing matplotlib axis.	None
pretty_groups	Optional[Dict[str, List[str]]]	Optional model-id groups displayed as brackets.	None
scale	str	Color scale type.	'linear'
scale_gamma	Optional[float]	Optional gamma for power scaling.	None
annot_fontsize	Optional[Union[int, float]]	Optional annotation font size.	None
tick_label_fontsize	Optional[Union[int, float]]	Optional tick-label font size.	None
group_label_fontsize	Optional[Union[int, float]]	Optional group-label font size.	None
group_label_rotation_top	Union[int, float]	Rotation for top group labels.	0
group_label_rotation_right	Union[int, float]	Rotation for right group labels.	0
cbar_pad	Optional[float]	Optional colorbar padding.	None
cbar_y_pad	Optional[float]	Optional vertical colorbar offset as a fraction of the heatmap height; negative values move it down.	None
cbar_fontsize	Optional[Union[int, float]]	Optional colorbar font size.	None
cbar_shrink	Optional[float]	Optional colorbar shrink factor (width relative to heatmap).	None
percentages	bool	Whether overlap values are shown as percentages.	True
row_metric	Optional[Dict[str, float]]	Optional side metric by row id.	None
row_metric_label	Optional[str]	Label for the side metric.	'corr'
row_metric_cmap	str	Colormap for the side metric.	'magma'
row_metric_vmin	Optional[float]	Optional side-metric lower bound.	None
row_metric_vmax	Optional[float]	Optional side-metric upper bound.	None
row_label_mapping	Optional[Dict[str, str]]	Optional mapping for row display labels.	None
column_label_mapping	Optional[Dict[str, str]]	Optional mapping for column display labels.	None
highlight_non_convergence	bool	Whether labels mark non-converged runs.	True
seed_pairing_mode	str	"all_pairs" or positionally aligned "matched".	'matched'
converged_by_id	Optional[Dict[str, Optional[bool]]]	Optional explicit convergence status by stable model id.	None

Source code in gradiend/visualizer/topk/pairwise_heatmap.py

def plot_topk_overlap_heatmap(
    models: Dict[str, object],
    topk: Union[int, float] = 1000,
    part: str = "decoder-weight",
    value: str = "intersection",
    order: Union[str, List[str]] = "input",
    cluster: bool = False,
    annot: Union[bool, str] = "auto",
    fmt: Optional[str] = None,
    annot_fmt: Optional[str] = None,
    figsize: Optional[Tuple[float, float]] = None,
    cmap: str = "viridis",
    vmin: Optional[float] = None,
    vmax: Optional[float] = None,
    title: Optional[Union[str, bool]] = False,
    output_path: Optional[str] = None,
    show: bool = True,
    return_data: bool = True,
    return_fig_ax: bool = False,
    ax: Optional[object] = None,
    pretty_groups: Optional[Dict[str, List[str]]] = None,
    scale: str = "linear",
    scale_gamma: Optional[float] = None,
    annot_fontsize: Optional[Union[int, float]] = None,
    tick_label_fontsize: Optional[Union[int, float]] = None,
    axis_label_fontsize: Optional[Union[int, float]] = None,
    group_label_fontsize: Optional[Union[int, float]] = None,
    group_label_rotation_top: Union[int, float] = 0,
    group_label_rotation_right: Union[int, float] = 0,
    cbar_pad: Optional[float] = None,
    cbar_y_pad: Optional[float] = None,
    cbar_fontsize: Optional[Union[int, float]] = None,
    cbar_shrink: Optional[float] = None,
    percentages: bool = True,
    row_metric: Optional[Dict[str, float]] = None,
    row_metric_label: Optional[str] = "corr",
    row_metric_cmap: str = "magma",
    row_metric_vmin: Optional[float] = None,
    row_metric_vmax: Optional[float] = None,
    row_label_mapping: Optional[Dict[str, str]] = None,
    column_label_mapping: Optional[Dict[str, str]] = None,
    highlight_non_convergence: bool = True,
    seed_aggregate: str = "mean",
    dispersion: str = "none",
    seed_pairing_mode: str = "matched",
    dispersion_display: str = "none",
    converged_by_id: Optional[Dict[str, Optional[bool]]] = None,
):
    """
    Plot pairwise top-k overlap between GRADIEND models.

    Args:
        models: Mapping of display/model ids to GRADIEND models.
        topk: Number or fraction of top weights to compare.
        part: Model part used to select weights.
        value: Overlap value to display.
        order: ``"input"``, ``"cluster"``, or explicit model order.
        cluster: Whether to cluster rows/columns.
        annot: Whether/how to annotate cells.
        fmt: Deprecated alias for ``annot_fmt``.
        annot_fmt: Cell annotation format.
        figsize: Optional figure size.
        cmap: Heatmap colormap.
        vmin: Optional lower value bound.
        vmax: Optional upper value bound.
        title: Optional plot title, or False to omit.
        output_path: Optional output path.
        show: Whether to display the plot.
        return_data: Whether to include computed data in the return payload.
        return_fig_ax: Whether to include matplotlib figure/axis.
        ax: Optional existing matplotlib axis.
        pretty_groups: Optional model-id groups displayed as brackets.
        scale: Color scale type.
        scale_gamma: Optional gamma for power scaling.
        annot_fontsize: Optional annotation font size.
        tick_label_fontsize: Optional tick-label font size.
        group_label_fontsize: Optional group-label font size.
        group_label_rotation_top: Rotation for top group labels.
        group_label_rotation_right: Rotation for right group labels.
        cbar_pad: Optional colorbar padding.
        cbar_y_pad: Optional vertical colorbar offset as a fraction of the
            heatmap height; negative values move it down.
        cbar_fontsize: Optional colorbar font size.
        cbar_shrink: Optional colorbar shrink factor (width relative to heatmap).
        percentages: Whether overlap values are shown as percentages.
        row_metric: Optional side metric by row id.
        row_metric_label: Label for the side metric.
        row_metric_cmap: Colormap for the side metric.
        row_metric_vmin: Optional side-metric lower bound.
        row_metric_vmax: Optional side-metric upper bound.
        row_label_mapping: Optional mapping for row display labels.
        column_label_mapping: Optional mapping for column display labels.
        highlight_non_convergence: Whether labels mark non-converged runs.
        seed_pairing_mode: ``"all_pairs"`` or positionally aligned ``"matched"``.
        converged_by_id: Optional explicit convergence status by stable model id.
    """
    return plot_similarity_heatmap(
        models,
        measure="topk_overlap",
        part=part,
        topk=topk,
        value=value,
        order=order,
        cluster=cluster,
        annot=annot,
        fmt=fmt,
        annot_fmt=annot_fmt,
        figsize=figsize,
        cmap=cmap,
        vmin=vmin,
        vmax=vmax,
        title=title,
        output_path=output_path,
        show=show,
        return_data=return_data,
        return_fig_ax=return_fig_ax,
        ax=ax,
        pretty_groups=pretty_groups,
        scale=scale,
        scale_gamma=scale_gamma,
        annot_fontsize=annot_fontsize,
        tick_label_fontsize=tick_label_fontsize,
        axis_label_fontsize=axis_label_fontsize,
        group_label_fontsize=group_label_fontsize,
        group_label_rotation_top=group_label_rotation_top,
        group_label_rotation_right=group_label_rotation_right,
        cbar_pad=cbar_pad,
        cbar_y_pad=cbar_y_pad,
        cbar_fontsize=cbar_fontsize,
        cbar_shrink=cbar_shrink,
        percentages=percentages,
        row_metric=row_metric,
        row_metric_label=row_metric_label,
        row_metric_cmap=row_metric_cmap,
        row_metric_vmin=row_metric_vmin,
        row_metric_vmax=row_metric_vmax,
        row_label_mapping=row_label_mapping,
        column_label_mapping=column_label_mapping,
        highlight_non_convergence=highlight_non_convergence,
        seed_aggregate=seed_aggregate,
        dispersion=dispersion,
        seed_pairing_mode=seed_pairing_mode,
        dispersion_display=dispersion_display,
        converged_by_id=converged_by_id,
    )