def comparison_heatmap(
vals: pd.Series,
ax: plt.Axes,
log: bool = False,
fmt: Callable[[float], str] = lambda x: f"{x:.2f}",
cbar_kws: Optional[Dict[str, Any]] = None,
cmap: str = "GnBu",
robust: bool = False,
preserve_order: bool = False,
label_fstr: Optional[str] = None,
mask: Optional[pd.Series] = None,
yaxis: bool = True,
**kwargs,
) -> None:
"""Plot a heatmap, with target_reward_type as x-axis and remainder as y-axis.
This is intended for plotting reward function distances, comparing reward functions
on the y-axis to those on the x-axis.
Args:
vals: The values to visualize.
log: log-10 scale for the values if True.
fmt: format string for annotated values.
cmap: color map.
robust: If true, set vmin and vmax to 25th and 75th quantiles.
This makes the color scale robust to outliers, but will compress it
more than is desirable if the data does not contain outliers.
preserve_order: If True, retains the same order as the input index
after rewriting the index values for readability. If false,
sorts the rewritten values alphabetically.
label_fstr: Format string to use for the label for the colorbar legend.` {args}` is
replaced with arguments to distance and `{transform_start}` and `{transform_end}`
is replaced with any transformations of the distance (e.g. log).
mask: If provided, only display cells where mask is True.
yaxis: Whether to plot labels for y-axis.
**kwargs: passed through to sns.heatmap.
"""
vals = transformations.index_reformat(vals, preserve_order)
if mask is not None:
mask = transformations.index_reformat(mask, preserve_order)
data = np.log10(vals) if log else vals
annot = vals.applymap(fmt)
cbar_kws = dict(cbar_kws or {})
if label_fstr is None:
label_fstr = "{transform_start}D({args}){transform_end}"
transform_start = r"\log_{10}\left(" if log else ""
transform_end = r"\right)" if log else ""
label = label_fstr.format(transform_start=transform_start,
args="R_A,R_B",
transform_end=transform_end)
cbar_kws.setdefault("label", f"${label}$")
if robust:
flat = data.values.flatten()
kwargs["vmin"], kwargs["vmax"] = np.quantile(flat, [0.25, 0.75])
yticklabels = "auto" if yaxis else False
sns.heatmap(
data,
annot=annot,
fmt="s",
cmap=cmap,
cbar_kws=cbar_kws,
mask=mask,
ax=ax,
yticklabels=yticklabels,
**kwargs,
)
ax.set_xlabel(r"$R_B$")
ax.set_ylabel(r"$R_A$" if yaxis else "")
