heatmap#
Plot a heatmap of any pandas DataFrame with optional axis labels, annotations, and custom color-scaling.
π₯ Arguments#
Name |
Type |
Required |
Description |
|---|---|---|---|
data |
pd.DataFrame |
β |
Pandas DataFrame containing the matrix to visualize; index and columns are used as y- and x-labels by default. |
x_labels |
List[str] |
β |
Labels for columns (overrides DataFrame columns). |
y_labels |
List[str] |
β |
Labels for rows (overrides DataFrame index). |
x_category |
str |
β |
Category for x-axis labels (overrides DataFrame columns). |
y_category |
str |
β |
Category for y-axis labels (overrides DataFrame index). |
figsize |
tuple |
β |
Figure size (width, height), e.g., (8, 6). |
cmap |
str |
β |
Name of matplotlib colormap, e.g., βviridisβ. |
vmin |
float |
β |
Minimum data value for color scaling. |
vmax |
float |
β |
Maximum data value for color scaling. |
annot |
bool |
β |
If True, overlay numeric annotations on each cell. |
fmt |
str |
β |
Format string for annotations, e.g. β.2fβ or βdβ. |
cbar |
bool |
β |
Whether to display the colorbar. |
title |
str |
β |
Title of the heatmap. |
save |
str |
β |
Base filename to save PNG and PDF (without extension). |
cbar_title |
str |
β |
Title for the colorbar. |
fontsize |
float |
β |
Base font size for axis labels, tick labels, title, cell annotations, and colorbar label. |
ax |
matplotlib.axes.Axes |
β |
Matplotlib Axes object to plot on. If None, a new figure is created. |
π¦ Example Output#
Click to show example code
import pandas as pd
import matplotlib.pyplot as plt
from swizz import plot
# 1) Create the exact DataFrame from your example image
data = pd.DataFrame({
10: {6: 0.93, 9: 0.86, 12: 0.86},
20: {6: 0.92, 9: 0.88, 12: 0.86},
50: {6: 0.92, 9: 0.90, 12: 0.75},
100:{6: 0.75, 9: 0.90, 12: 0.78},
})
data.index.name = "Radius"
data.columns.name = "Team Reward"
# 2) Plot with annotations, custom colormap, cbar label, axis labels, and save output
fig, ax = plot(
"heatmap",
data=data,
figsize=(7, 5),
cmap="Blues",
vmin=0.70,
vmax=0.95,
annot=True,
fmt=".2f",
cbar=True,
cbar_title="Lone Wolf Rate",
save="lone_wolf_rate_heatmap",
x_category="Team Reward",
y_category="Radius",
)
# 3) Show the figure
plt.show()