import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
from numpy.typing import NDArray
from .color_themes import get_theme
from .plotting_constants import (
CONTOUR_LABEL_COLOR,
DEFAULT_COLORBAR_CONFIG,
DEFAULT_FONT_CONFIG,
DEFAULT_LAYOUT_CONFIG,
DEFAULT_SCENE_CONFIG,
DEFAULT_SMOOTHING_CONFIG,
FLOAT_PRECISION,
GRID_COLOR,
LINE_WIDTH,
MARKER_SIZE,
NUM_CONTOURS,
TEXT_COLOR_DARK,
TEXT_COLOR_LIGHT,
ColorbarConfig,
FontConfig,
LayoutConfig,
SceneConfig,
SmoothingConfig,
)
from .utils import (
DataType3D,
prepare_smooth_data_3D,
prepare_smooth_data_3D_scatter,
validate_data,
)
# Set default plotly template for better aesthetics
pio.templates.default = "plotly_white"
[docs]
def plot3D_surface(
x_list: NDArray[np.float64],
y_list: NDArray[np.float64],
z_list: NDArray[np.float64],
title: str,
x_axis: str,
y_axis: str,
z_axis: str,
subtitle: str | None = None,
theme: str | None = None,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
colorbar_config: ColorbarConfig = DEFAULT_COLORBAR_CONFIG,
smoothing_config: SmoothingConfig = DEFAULT_SMOOTHING_CONFIG,
scene_config: SceneConfig = DEFAULT_SCENE_CONFIG,
) -> go.Figure:
"""
Generic 3D surface plotting utility.
Handles both grid data (2D z array) and scatter data (1D x, y, z arrays).
For grid data: Creates interpolated surface plot with contour lines.
For scatter data: Creates 3D scatter plot with points colored by z values.
Parameters
----------
x_list : numpy.ndarray
X coordinates (1D for grid, 1D for scatter)
y_list : numpy.ndarray
Y coordinates (1D for grid, 1D for scatter)
z_list : numpy.ndarray
Z values (2D for grid, 1D for scatter)
title : str
Plot title
x_axis : str
X-axis label
y_axis : str
Y-axis label
z_axis : str
Z-axis label
subtitle : str, optional
Optional subtitle
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
smoothing_config : SmoothingConfig, optional
SmoothingConfig object for smoothing settings
scene_config : SceneConfig, optional
SceneConfig object for 3D scene settings
Returns
-------
go.Figure
Plotly figure object
"""
# Detect data type
data_type = validate_data(x_list, y_list, z_list)
# Create figure and add trace
fig = go.Figure()
theme_colors = get_theme(theme)
# Prepare smoothed/interpolated data based on input type
match data_type:
case DataType3D.GridInput:
x_dense, y_dense, z_dense = prepare_smooth_data_3D(
x_list, y_list, z_list, smoothing_config=smoothing_config
)
case DataType3D.ScatterInput:
x_dense, y_dense, z_dense = prepare_smooth_data_3D_scatter(
x_list, y_list, z_list, smoothing_config=smoothing_config
)
# Calculate contour size (use nanmin/nanmax to handle potential NaN values)
z_min, z_max = np.nanmin(z_dense), np.nanmax(z_dense)
contour_size = (z_max - z_min) / NUM_CONTOURS if z_max > z_min else 1
# Add surface plot
fig.add_trace(
go.Surface(
x=np.array([x_dense] * len(x_dense)),
y=np.array([y_dense] * len(x_dense)).T,
z=z_dense,
colorscale=theme_colors["colorscale"],
contours={
"z": {
"show": True,
"start": z_min,
"end": z_max,
"size": contour_size,
"width": LINE_WIDTH,
"color": CONTOUR_LABEL_COLOR,
}
},
colorbar=dict(
title=dict(
text=z_axis,
font=dict(size=font_config.medium),
),
thickness=colorbar_config.thickness,
len=colorbar_config.length,
tickfont=dict(size=font_config.small),
),
hovertemplate=(
f"{x_axis}: %{{x:{FLOAT_PRECISION}}}<br>"
f"{y_axis}: %{{y:{FLOAT_PRECISION}}}<br>"
f"{z_axis}: %{{z:{FLOAT_PRECISION}}}<extra></extra>"
),
)
)
# Overlay original scatter points if input was scatter data
if data_type == DataType3D.ScatterInput:
fig.add_trace(
go.Scatter3d(
x=x_list,
y=y_list,
z=z_list,
mode="markers",
marker=dict(
color="black",
size=MARKER_SIZE / 2,
opacity=0.5,
),
showlegend=False,
hovertemplate=(
f"Original point<br>"
f"{x_axis}: %{{x:{FLOAT_PRECISION}}}<br>"
f"{y_axis}: %{{y:{FLOAT_PRECISION}}}<br>"
f"{z_axis}: %{{z:{FLOAT_PRECISION}}}<extra></extra>"
),
)
)
# Configure layout
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,
"yanchor": layout_config.title_yanchor,
},
scene=dict(
xaxis_title=x_axis,
yaxis_title=y_axis,
zaxis_title=z_axis,
xaxis=dict(
gridcolor=GRID_COLOR,
showbackground=True,
backgroundcolor=layout_config.scene_bgcolor,
tickformat=FLOAT_PRECISION,
),
yaxis=dict(
gridcolor=GRID_COLOR,
showbackground=True,
backgroundcolor=layout_config.scene_bgcolor,
tickformat=FLOAT_PRECISION,
),
zaxis=dict(
gridcolor=GRID_COLOR,
showbackground=True,
backgroundcolor=layout_config.scene_bgcolor,
tickformat=FLOAT_PRECISION,
),
aspectratio=dict(
x=scene_config.aspect_ratio_x,
y=scene_config.aspect_ratio_y,
z=scene_config.aspect_ratio_z,
),
),
width=layout_config.width,
height=layout_config.height,
margin=layout_config.margin,
)
fig.update_layout(
scene_camera=dict(
center=dict(x=0, y=0, z=0),
eye=dict(
x=scene_config.camera_distance
* np.cos(np.radians(scene_config.default_view_angle)),
y=scene_config.camera_distance
* np.sin(np.radians(scene_config.default_view_angle)),
z=scene_config.camera_z,
),
up=dict(
x=scene_config.camera_up_x,
y=scene_config.camera_up_y,
z=scene_config.camera_up_z,
),
)
)
return fig
[docs]
def plot3D_contour(
x_list: NDArray[np.float64],
y_list: NDArray[np.float64],
z_list: NDArray[np.float64],
title: str,
x_axis: str,
y_axis: str,
z_axis: str,
subtitle: str | None = None,
theme: str | None = None,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
colorbar_config: ColorbarConfig = DEFAULT_COLORBAR_CONFIG,
smoothing_config: SmoothingConfig = DEFAULT_SMOOTHING_CONFIG,
) -> go.Figure:
"""
Generic 3D contour plotting utility.
Handles both grid data (2D z array) and scatter data (1D x, y, z arrays).
For grid data: Creates interpolated contour plot with contour lines.
For scatter data: Creates scatter plot with points colored by z values.
Parameters
----------
x_list : numpy.ndarray
X coordinates (1D for grid, 1D for scatter)
y_list : numpy.ndarray
Y coordinates (1D for grid, 1D for scatter)
z_list : numpy.ndarray
Z values (2D for grid, 1D for scatter)
title : str
Plot title
x_axis : str
X-axis label
y_axis : str
Y-axis label
z_axis : str
Z-axis label
subtitle : str, optional
Optional subtitle
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
smoothing_config : SmoothingConfig, optional
SmoothingConfig object for smoothing settings
Returns
-------
go.Figure
Plotly figure object
"""
# Detect data type
data_type = validate_data(x_list, y_list, z_list)
theme_colors = get_theme(theme)
fig = go.Figure()
# Prepare smoothed/interpolated data based on input type
match data_type:
case DataType3D.GridInput:
x_dense, y_dense, z_dense = prepare_smooth_data_3D(
x_list, y_list, z_list, smoothing_config=smoothing_config
)
case DataType3D.ScatterInput:
x_dense, y_dense, z_dense = prepare_smooth_data_3D_scatter(
x_list, y_list, z_list, smoothing_config=smoothing_config
)
fig.add_trace(
go.Scatter(
x=x_list,
y=y_list,
mode="markers",
marker=dict(
color="black",
size=MARKER_SIZE,
opacity=0.5,
symbol="x",
),
showlegend=False,
hovertemplate=(
f"Original point<br>"
f"{x_axis}: %{{x:{FLOAT_PRECISION}}}<br>"
f"{y_axis}: %{{y:{FLOAT_PRECISION}}}<br>"
f"{z_axis}: %{{customdata:{FLOAT_PRECISION}}}<extra></extra>"
),
customdata=z_list,
)
)
# Calculate contour size (use nanmin/nanmax to handle potential NaN values)
z_min, z_max = np.nanmin(z_dense), np.nanmax(z_dense)
contour_size = int((z_max - z_min) / NUM_CONTOURS) if z_max > z_min else 1
# Add contour plot
fig.add_trace(
go.Contour(
x=x_dense,
y=y_dense,
z=z_dense,
colorscale=theme_colors["colorscale"],
contours=dict(
showlabels=True,
labelfont=dict(
size=font_config.small,
color=CONTOUR_LABEL_COLOR,
),
start=int(z_min),
end=int(z_max),
size=contour_size,
labelformat=FLOAT_PRECISION,
),
colorbar=dict(
title=dict(
text=z_axis,
font=dict(size=font_config.medium),
),
thickness=colorbar_config.thickness,
len=colorbar_config.length,
tickfont=dict(size=font_config.small),
tickformat=FLOAT_PRECISION,
),
hovertemplate=(
f"{x_axis}: %{{x:{FLOAT_PRECISION}}}<br>"
f"{y_axis}: %{{y:{FLOAT_PRECISION}}}<br>"
f"{z_axis}: %{{z:{FLOAT_PRECISION}}}<extra></extra>"
),
)
)
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_axis,
"font": dict(size=font_config.medium, color=TEXT_COLOR_DARK),
},
yaxis_title={
"text": y_axis,
"font": dict(size=font_config.medium, color=TEXT_COLOR_DARK),
},
width=layout_config.width,
height=layout_config.height,
margin=layout_config.margin,
plot_bgcolor=layout_config.plot_bgcolor,
)
fig.update_xaxes(
tickfont=dict(size=font_config.small),
tickformat=FLOAT_PRECISION,
)
fig.update_yaxes(
tickfont=dict(size=font_config.small),
tickformat=FLOAT_PRECISION,
)
return fig