import numpy as np
import plotly.graph_objects as go
from numpy.typing import NDArray
from .color_themes import get_theme
from .plotting_constants import (
DEFAULT_COLORBAR_CONFIG,
DEFAULT_FONT_CONFIG,
DEFAULT_LAYOUT_CONFIG,
FLOAT_PRECISION,
TEXT_COLOR_DARK,
ColorbarConfig,
FontConfig,
LayoutConfig,
)
[docs]
def plot_heatmap(
x_values: NDArray[np.float64],
y_values: NDArray[np.float64],
z_values: NDArray[np.float64],
title: str,
x_label: str,
y_label: str,
z_label: str,
*,
theme: str | None = None,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
colorbar_config: ColorbarConfig = DEFAULT_COLORBAR_CONFIG,
) -> go.Figure:
"""
Generic 2D heatmap plotting utility.
Parameters
----------
x_values : numpy.ndarray
Coordinates along the x axis, one per column of `z_values`.
y_values : numpy.ndarray
Coordinates along the y axis, one per row of `z_values`.
z_values : numpy.ndarray
Cell values, shaped (len(y_values), len(x_values)).
title : str
Plot title.
x_label : str
X-axis label.
y_label : str
Y-axis label.
z_label : str
Colorbar label describing the cell values.
theme : str, optional
Color theme name
font_config : FontConfig, optional
FontConfig object for font settings
layout_config : LayoutConfig, optional
LayoutConfig object for layout settings
colorbar_config : ColorbarConfig, optional
ColorbarConfig object for colorbar settings
Returns
-------
go.Figure
Plotly figure object
"""
theme_colors = get_theme(theme)
figure = go.Figure(
data=go.Heatmap(
x=x_values,
y=y_values,
z=z_values,
colorscale=theme_colors["colorscale"],
colorbar=dict(
title=dict(
text=z_label,
font=dict(size=font_config.medium),
),
thickness=colorbar_config.thickness,
len=colorbar_config.length,
tickfont=dict(size=font_config.small),
),
hovertemplate=(
f"{x_label}: %{{x:{FLOAT_PRECISION}}}<br>"
f"{y_label}: %{{y:{FLOAT_PRECISION}}}<br>"
f"{z_label}: %{{z:{FLOAT_PRECISION}}}<extra></extra>"
),
)
)
figure.update_layout(
title=dict(
text=title,
font=dict(size=font_config.large, color=TEXT_COLOR_DARK),
x=layout_config.title_x,
xanchor=layout_config.title_xanchor,
yanchor=layout_config.title_yanchor,
),
xaxis_title=x_label,
yaxis_title=y_label,
width=layout_config.width,
height=layout_config.height,
margin=layout_config.margin,
)
return figure