import numpy as np
import plotly.graph_objects as go
from numpy.typing import NDArray
from .plotting_constants import (
DEFAULT_FONT_CONFIG,
DEFAULT_LAYOUT_CONFIG,
GRID_COLOR,
GRID_WIDTH,
MARKER_SIZE,
TEXT_COLOR_DARK,
FontConfig,
LayoutConfig,
)
def _quantile_binned_stats(
x: NDArray[np.float64],
y: NDArray[np.float64],
n_bins: int,
min_per_bin: int,
) -> tuple[NDArray[np.float64], NDArray[np.float64], NDArray[np.float64]]:
"""Bin x into quantile bins and return each bin's center x, mean y, and std of y."""
if len(x) < min_per_bin + 1:
empty = np.array([], dtype=np.float64)
return empty, empty, empty
edges = np.unique(np.quantile(x, np.linspace(0.0, 1.0, n_bins + 1)))
if edges.size < 3:
empty = np.array([], dtype=np.float64)
return empty, empty, empty
centers, means, stds = [], [], []
for lo, hi in zip(edges[:-1], edges[1:]):
mask = (x >= lo) & (x <= hi)
if int(mask.sum()) < min_per_bin:
continue
centers.append(0.5 * (lo + hi))
means.append(float(y[mask].mean()))
stds.append(float(y[mask].std()))
return np.asarray(centers), np.asarray(means), np.asarray(stds)
[docs]
def plot_binned_scatter(
x: NDArray[np.float64],
y: NDArray[np.float64],
*,
color_values: NDArray[np.float64] | None = None,
color_label: str = "group",
title: str = "",
x_label: str = "",
y_label: str = "",
n_bins: int = 24,
min_per_bin: int = 5,
reference_y: float | None = 0.0,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
) -> go.Figure:
"""Scatter plot of (x, y) points, grouped into quantiles on x. Displays
mean and +/- 1 std of y for each bin
Parameters
----------
x : NDArray[float64]
y : NDArray[float64]
color_values : NDArray[float64], optional
Per-point values used to color the scatter (e.g. a group id). Defaults to None (single color).
color_label : str, optional
Colorbar title when ``color_values`` is given.
title, x_label, y_label : str, optional
Figure title and axis labels.
n_bins : int, optional
Number of x-bins for the mean/sigma overlay. Default is 24.
min_per_bin : int, optional
Minimum points a bin needs to be shown. Default is 5.
reference_y : float, optional
Y value for a dashed horizontal reference line. Default is 0.0;
pass None to omit it.
font_config : FontConfig, optional
layout_config : LayoutConfig, optional
Returns
-------
go.Figure
"""
x = np.asarray(x, dtype=np.float64)
y = np.asarray(y, dtype=np.float64)
marker: dict[str, object] = dict(size=MARKER_SIZE - 1, opacity=0.45)
if color_values is not None:
marker.update(
color=np.asarray(color_values),
colorscale="Viridis",
colorbar=dict(title=color_label),
)
fig = go.Figure()
fig.add_trace(
go.Scattergl(
x=x,
y=y,
mode="markers",
marker=marker,
name="samples",
)
)
bx, bm, bs = _quantile_binned_stats(x, y, n_bins, min_per_bin)
if bx.size:
fig.add_trace(
go.Scattergl(
x=np.concatenate([bx, bx[::-1]]),
y=np.concatenate([bm + bs, (bm - bs)[::-1]]),
fill="toself",
fillcolor="rgba(220,40,40,0.30)",
line=dict(color="rgba(180,20,20,0.65)", width=1),
hoverinfo="skip",
name="+/- 1 sigma",
)
)
fig.add_trace(
go.Scattergl(
x=bx,
y=bm,
mode="lines+markers",
line=dict(color="black", width=2),
marker=dict(size=MARKER_SIZE + 2, color="black"),
name="binned mean",
)
)
if reference_y is not None:
fig.add_hline(
y=reference_y,
line=dict(color="#1a73e8", width=1, dash="dash"),
)
fig.update_layout(
title={
"text": title,
"font": dict(size=font_config.large, color=TEXT_COLOR_DARK),
"x": layout_config.title_x,
"xanchor": layout_config.title_xanchor,
},
xaxis_title=dict(text=x_label, font=dict(size=font_config.medium)),
yaxis_title=dict(text=y_label, font=dict(size=font_config.medium)),
height=layout_config.height,
width=layout_config.width,
margin=layout_config.margin,
plot_bgcolor=layout_config.plot_bgcolor,
legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="center", x=0.5),
)
fig.update_xaxes(
showgrid=True,
gridwidth=GRID_WIDTH,
gridcolor=GRID_COLOR,
tickfont=dict(size=font_config.small),
)
fig.update_yaxes(
showgrid=True,
gridwidth=GRID_WIDTH,
gridcolor=GRID_COLOR,
tickfont=dict(size=font_config.small),
)
return fig