plot_cross_encoding_heatmap

Plot a cross-encoding heatmap for pairwise GRADIEND trainers.

Pass feature_classes for oriented symmetric cross-encoding (anchor sign alignment and aggregation; see compute_anchor_aligned_encoding_matrix). Omit feature_classes for directed positive-pair trainer×trainer encoding via compute_trainer_pair_encoding_matrix (default metric positive_mean).

Parameters:

Name	Type	Description	Default
trainers	Dict[str, object]	Mapping of ids to trainers.	required
feature_classes	Optional[Sequence[str]]	Feature classes used as oriented row anchors. When omitted, the directed positive-pair matrix is plotted instead.	None
alignment	str	Column alignment for oriented mode. "auto" (default) follows trained GRADIEND source: factual/diff-source trainers use factual alignment; alternative-source trainers use counterfactual alignment. Explicit values "factual", "counterfactual", and "transition" are diagnostic overrides.	'auto'
column_ids	Optional[Sequence[str]]	Optional explicit oriented-matrix columns.	None
encoder_summary	Optional[Dict[str, Any]]	Optional precomputed encoder summary for oriented mode.	None
split	str	Encoder split used when evaluation is needed.	'test'
max_size	Optional[int]	Optional evaluation row cap.	None
use_cache	bool	Whether to use cached encoder analysis.	True
full_eval	Optional[bool]	Directed mode only; whether encoder evaluation includes all transitions. Oriented matrices use a shared cross-task test pool when encoder_summary is not supplied.	None
aggregate	str	Oriented mode aggregate across trainers per anchor.	'mean'
metric	str	Directed mode cross-encoding metric.	'positive_mean'
encoder_eval	str	Directed mode encoder evaluation policy: "auto", "cached", or "recompute".	'auto'
allow_incomplete	bool	Directed mode; tolerate missing cells.	False
seed_selection	Optional[str]	Directed mode seed selection.	None
seed_aggregate	str	Directed mode seed aggregate.	'mean'
dispersion	Optional[str]	Directed mode dispersion statistic.	None
normalize	bool	If True, divide each row by its diagonal so self-encoding is 1.0. Requires a square matrix with matching row/column ids (oriented factual/counterfactual mode, or directed trainer×trainer mode).	False
order	Any	Heatmap row/column order.	'input'
cluster	bool	Whether to cluster rows/columns.	False
highlight_non_convergence	bool	Whether labels mark non-converged runs.	True
xlabel	Optional[str]	Optional x-axis label. Oriented mode defaults to a probe-feature label.	None
ylabel	Optional[str]	Optional y-axis label. Oriented mode defaults to Orienting feature.	None
**plot_kwargs	Any	Additional options forwarded to plot_comparison_heatmap.	{}

Source code in gradiend/visualizer/heatmaps/encoding.py

def plot_cross_encoding_heatmap(
    trainers: Dict[str, object],
    feature_classes: Optional[Sequence[str]] = None,
    *,
    alignment: str = "auto",
    column_ids: Optional[Sequence[str]] = None,
    encoder_summary: Optional[Dict[str, Any]] = None,
    split: str = "test",
    max_size: Optional[int] = None,
    use_cache: bool = True,
    full_eval: Optional[bool] = None,
    aggregate: str = "mean",
    metric: str = "positive_mean",
    encoder_eval: str = "auto",
    allow_incomplete: bool = False,
    seed_selection: Optional[str] = None,
    seed_aggregate: str = "mean",
    dispersion: Optional[str] = None,
    normalize: bool = False,
    order: Any = "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[Any] = 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_fontsize: Optional[Union[int, float]] = None,
    cbar_shrink: Optional[float] = None,
    cbar_label: Optional[str] = None,
    percentages: bool = False,
    row_label_mapping: Optional[Dict[str, str]] = None,
    column_label_mapping: Optional[Dict[str, str]] = None,
    xlabel: Optional[str] = None,
    ylabel: Optional[str] = None,
    dispersion_display: str = "none",
    seed_annotation: Union[bool, Dict[str, Any]] = False,
    highlight_non_convergence: bool = True,
    **plot_kwargs: Any,
) -> Dict[str, Any]:
    """
    Plot a cross-encoding heatmap for pairwise GRADIEND trainers.

    Pass ``feature_classes`` for oriented symmetric cross-encoding (anchor sign
    alignment and aggregation; see ``compute_anchor_aligned_encoding_matrix``).
    Omit ``feature_classes`` for directed positive-pair trainer×trainer encoding via
    ``compute_trainer_pair_encoding_matrix`` (default metric ``positive_mean``).

    Args:
        trainers: Mapping of ids to trainers.
        feature_classes: Feature classes used as oriented row anchors. When
            omitted, the directed positive-pair matrix is plotted instead.
        alignment: Column alignment for oriented mode. ``"auto"`` (default)
            follows trained GRADIEND source: factual/diff-source trainers use
            factual alignment; alternative-source trainers use counterfactual
            alignment. Explicit values ``"factual"``, ``"counterfactual"``,
            and ``"transition"`` are diagnostic overrides.
        column_ids: Optional explicit oriented-matrix columns.
        encoder_summary: Optional precomputed encoder summary for oriented mode.
        split: Encoder split used when evaluation is needed.
        max_size: Optional evaluation row cap.
        use_cache: Whether to use cached encoder analysis.
        full_eval: Directed mode only; whether encoder evaluation includes all
            transitions. Oriented matrices use a shared cross-task test pool
            when ``encoder_summary`` is not supplied.
        aggregate: Oriented mode aggregate across trainers per anchor.
        metric: Directed mode cross-encoding metric.
        encoder_eval: Directed mode encoder evaluation policy:
            ``"auto"``, ``"cached"``, or ``"recompute"``.
        allow_incomplete: Directed mode; tolerate missing cells.
        seed_selection: Directed mode seed selection.
        seed_aggregate: Directed mode seed aggregate.
        dispersion: Directed mode dispersion statistic.
        normalize: If True, divide each row by its diagonal so self-encoding is
            1.0. Requires a square matrix with matching row/column ids (oriented
            factual/counterfactual mode, or directed trainer×trainer mode).
        order: Heatmap row/column order.
        cluster: Whether to cluster rows/columns.
        highlight_non_convergence: Whether labels mark non-converged runs.
        xlabel: Optional x-axis label. Oriented mode defaults to a probe-feature label.
        ylabel: Optional y-axis label. Oriented mode defaults to ``Orienting feature``.
        **plot_kwargs: Additional options forwarded to ``plot_comparison_heatmap``.
    """
    if cbar_label is None:
        cbar_label = (
            CROSS_ENCODING_ROW_NORMALIZED_CBAR_LABEL
            if normalize
            else CROSS_ENCODING_CBAR_LABEL
        )
    seed_selection = resolve_seed_selection_for_trainers(trainers, seed_selection)
    if dispersion is None:
        dispersion = resolve_dispersion_for_trainers(trainers, None)
    if feature_classes is not None:
        resolved_alignment, alignment_was_auto = resolve_oriented_cross_encoding_alignment(
            trainers,
            alignment,
        )
        if encoder_summary is None:
            comparison_data = compute_dense_anchor_aligned_encoding_matrix(
                trainers,
                feature_classes,
                alignment=resolved_alignment,
                column_ids=column_ids,
                split=split,
                max_size=max_size,
                aggregate=aggregate,
                seed_selection=seed_selection,
                seed_aggregate=seed_aggregate,
                dispersion=dispersion,
            )
        else:
            comparison_data = compute_anchor_aligned_encoding_matrix(
                pair_by_id=pair_by_id_from_trainers(trainers),
                encoder_summary=encoder_summary,
                feature_classes=feature_classes,
                aggregate=aggregate,
                alignment=resolved_alignment,
                column_ids=column_ids,
                source_by_id=source_by_id_from_trainers(trainers),
            )
        if normalize:
            comparison_data = normalize_cross_encoding_rows_by_diagonal(comparison_data)
        resolved_order = list(feature_classes) if order == "input" else order
        default_ylabel, default_xlabel = _oriented_cross_encoding_axis_labels(
            resolved_alignment,
        )
        if alignment_was_auto and resolved_alignment == "counterfactual":
            default_xlabel = ORIENTED_CROSS_ENCODING_XLABEL_FACTUAL
        return plot_comparison_heatmap(
            comparison_data,
            order=resolved_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_fontsize=cbar_fontsize,
            cbar_shrink=cbar_shrink,
            cbar_label=cbar_label,
            percentages=percentages,
            row_label_mapping=row_label_mapping,
            column_label_mapping=column_label_mapping,
            xlabel=default_xlabel if xlabel is None else xlabel,
            ylabel=default_ylabel if ylabel is None else ylabel,
            dispersion_display=dispersion_display,
            seed_annotation=seed_annotation,
            models=trainers,
            highlight_non_convergence=highlight_non_convergence,
            **filter_comparison_heatmap_plot_kwargs(plot_kwargs),
        )

    comparison_data = compute_trainer_pair_encoding_matrix(
        trainers,
        split=split,
        max_size=max_size,
        use_cache=use_cache,
        metric=metric,
        full_eval=full_eval,
        encoder_eval=encoder_eval,
        allow_incomplete=allow_incomplete,
        seed_selection=seed_selection,
        seed_aggregate=seed_aggregate,
        dispersion=dispersion,
    )
    if normalize:
        comparison_data = normalize_cross_encoding_rows_by_diagonal(comparison_data)
    return plot_comparison_heatmap(
        comparison_data,
        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_fontsize=cbar_fontsize,
        cbar_shrink=cbar_shrink,
        cbar_label=cbar_label,
        percentages=percentages,
        row_label_mapping=row_label_mapping,
        column_label_mapping=column_label_mapping,
        xlabel=xlabel,
        ylabel=ylabel,
        dispersion_display=dispersion_display,
        seed_annotation=seed_annotation,
        models=trainers,
        highlight_non_convergence=highlight_non_convergence,
        **filter_comparison_heatmap_plot_kwargs(plot_kwargs),
    )