from typing import Optional
import numpy as np
import numpy.typing as npt
import plotly.graph_objects as go
from .plotting_constants import *
def format_polynomial_equation(coeffs: npt.NDArray[np.float64]) -> str:
"""Format a descending-order coefficient array as a human-readable polynomial equation."""
degree = len(coeffs) - 1
parts: list[str] = []
for i, c in enumerate(coeffs):
p = degree - i
if p == 0:
term = f"{abs(c):.2f}"
elif p == 1:
term = f"{abs(c):.2f}x"
else:
term = f"{abs(c):.2f}x^{p}"
sign = ("-" if c < 0 else "") if not parts else ("- " if c < 0 else "+ ")
parts.append(sign + term)
return "y = " + " ".join(parts)
[docs]
def plot_polynomial_fit(
x: npt.NDArray[np.float64],
y: npt.NDArray[np.float64],
coeffs: npt.NDArray[np.float64],
title: str,
x_label: str,
y_label: str,
*,
subtitle: Optional[str] = None,
timestamps: Optional[npt.NDArray[np.float64]] = None,
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
) -> go.Figure:
"""Generic scatter plot with a polynomial fit line overlaid.
Parameters
----------
x : NDArray[np.float64]
X-axis data for scatter points.
y : NDArray[np.float64]
Y-axis data for scatter points.
coeffs : NDArray[np.float64]
Polynomial coefficients in descending power order (compatible with ``np.polyval``).
title : str
Plot title.
x_label : str
X-axis label.
y_label : str
Y-axis label.
subtitle : str, optional
Subtitle shown below the title in smaller grey text.
layout_config : LayoutConfig
Layout dimensions and spacing configuration.
font_config : FontConfig
Font size configuration.
Returns
-------
go.Figure
"""
eq_label = format_polynomial_equation(coeffs)
fig = go.Figure()
hover_template = (
f"{x_label}: %{{x:{FLOAT_PRECISION}}}<br>"
f"{y_label}: %{{y:{FLOAT_PRECISION}}}<br>"
"%{text}<extra></extra>"
)
fig.add_trace(
go.Scattergl(
x=x,
y=y,
mode="markers",
marker=dict(size=MARKER_SIZE, opacity=0.4),
name="Data",
text=(
[f"t = {t:.2f} s" for t in timestamps]
if timestamps is not None
else None
),
hovertemplate=hover_template,
)
)
x_fit = np.linspace(x.min(), x.max(), 200)
fig.add_trace(
go.Scatter(
x=x_fit,
y=np.polyval(coeffs, x_fit),
mode="lines",
line=dict(color="red", width=LINE_WIDTH),
name=eq_label,
)
)
full_title = title
if subtitle is not None:
full_title += f"<br><span style='font-size: {font_config.medium}px; color: {TEXT_COLOR_LIGHT};'>{subtitle}</span>"
fig.update_layout(
title={
"text": full_title,
"font": dict(size=font_config.large, color=TEXT_COLOR_DARK),
"x": layout_config.title_x,
"xanchor": layout_config.title_xanchor,
},
xaxis_title={
"text": x_label,
"font": dict(size=font_config.medium, color=TEXT_COLOR_DARK),
},
yaxis_title={
"text": y_label,
"font": dict(size=font_config.medium, color=TEXT_COLOR_DARK),
},
legend_title="Legend",
showlegend=True,
legend_title_font=dict(size=font_config.medium),
legend_font=dict(size=font_config.small),
hovermode=HOVER_MODE,
plot_bgcolor=layout_config.plot_bgcolor,
width=layout_config.width,
height=layout_config.height,
margin=layout_config.margin,
)
fig.update_xaxes(
showgrid=True,
gridwidth=GRID_WIDTH,
gridcolor=GRID_COLOR,
zeroline=True,
zerolinewidth=ZEROLINE_WIDTH,
zerolinecolor=ZEROLINE_COLOR,
tickfont=dict(size=font_config.small),
tickformat=FLOAT_PRECISION,
)
fig.update_yaxes(
showgrid=True,
gridwidth=GRID_WIDTH,
gridcolor=GRID_COLOR,
zeroline=True,
zerolinewidth=ZEROLINE_WIDTH,
zerolinecolor=ZEROLINE_COLOR,
tickfont=dict(size=font_config.small),
tickformat=FLOAT_PRECISION,
)
padding = 1 + RANGE_PADDING
fig.update_xaxes(range=[(1 / padding) * x.min(), padding * x.max()])
fig.update_yaxes(range=[(1 / padding) * y.min(), padding * y.max()])
return fig