from typing import Optional
import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
from ..plotting.grid_plot_3d import PlotType, plot_grid_3D
from ..plotting.plot_3d import plot3D_contour, plot3D_surface
from ..plotting.plotting_constants import *
from .energy_grid_models import EnergyGridConfig, EnergyGridData
from .types import *
pio.templates.default = "plotly_white"
# When plotting metric X, show metric Y in the hover tooltip as a companion.
ENERGY_GRID_COMPANION_METRICS = {
EnergyGridMetric.COAST_TRIGGER: EnergyGridMetric.TOTAL_PTS,
EnergyGridMetric.ENDURANCE_T: EnergyGridMetric.TOTAL_PTS,
EnergyGridMetric.NET_ENERGY: EnergyGridMetric.TOTAL_PTS,
EnergyGridMetric.CAPACITY: EnergyGridMetric.TOTAL_PTS,
EnergyGridMetric.VEHICLE_MASS: EnergyGridMetric.TOTAL_PTS,
}
[docs]
class EnergyGridResults:
"""Results and visualization for energy-constrained grid characterization."""
def __init__(self, data: EnergyGridData, config: EnergyGridConfig):
self.data = data
self.config = config
def _resolve_axes(self, show_capacity_kwh: bool):
"""Return (x_vals, y_vals, x_label, y_label) with optional kg→kWh conversion.
When show_capacity_kwh is True, any axis that sweeps accum.pack_weight
is converted to pack capacity in kWh using the stored energy_density.
"""
def _convert(vals, name):
if show_capacity_kwh and name == PACK_WEIGHT_PARAM:
delta = vals - self.data.nominal_pack_weight
cap = self.data.nominal_capacity + delta * self.data.energy_density
return cap, "Pack Capacity (kWh)"
return vals, name
x_vals, x_label = _convert(self.data.var_1_list, self.data.var_1_name)
y_vals, y_label = _convert(self.data.var_2_list, self.data.var_2_name)
return x_vals, y_vals, x_label, y_label
def _build_hover(
self,
metric: EnergyGridMetric,
x_label: str,
y_label: str,
mask_infeasible: bool,
):
"""Build (text_2d, hovertemplate) for a contour trace.
Returns a 2D string array containing the companion metric value,
and a hovertemplate referencing x/y/z with labels plus the companion line.
"""
companion = ENERGY_GRID_COMPANION_METRICS.get(
metric, EnergyGridMetric.COAST_TRIGGER
)
comp_label = ENERGY_GRID_METRIC_LABELS[companion]
comp_data = getattr(self.data, companion.value).copy().astype(float)
if mask_infeasible:
comp_data[~self.data.feasible] = np.nan
text_2d = np.empty(comp_data.shape, dtype=object)
for i in range(comp_data.shape[0]):
for j in range(comp_data.shape[1]):
v = comp_data[i, j]
text_2d[i, j] = f"{v:.2f}" if np.isfinite(v) else "—"
z_label = ENERGY_GRID_METRIC_LABELS[metric]
ht = (
f"{x_label}: %{{x:.2f}}<br>"
f"{y_label}: %{{y:.2f}}<br>"
f"{z_label}: %{{z:.2f}}<br>"
f"{comp_label}: %{{text}}"
"<extra></extra>"
)
return text_2d.T, ht
[docs]
def plot_contour(
self,
metric: EnergyGridMetric = EnergyGridMetric.TOTAL_PTS,
*,
mask_infeasible: bool = True,
show_capacity_kwh: bool = False,
title: Optional[str] = None,
theme: Optional[str] = "accumulator",
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
) -> go.Figure:
"""
2D filled contour of any metric over the grid.
Parameters
----------
metric : EnergyGridMetric
Which z-variable to plot.
mask_infeasible : bool
Replace infeasible cells with NaN so they appear blank.
show_capacity_kwh : bool
Show pack capacity (kWh) instead of pack weight (kg) on axes.
title : str, optional
Custom title.
theme : str, optional
Color theme name. Default "accumulator".
font_config : FontConfig, optional
layout_config : LayoutConfig, optional
Returns
-------
go.Figure
"""
x_vals, y_vals, x_label, y_label = self._resolve_axes(show_capacity_kwh)
z = getattr(self.data, metric.value).copy().astype(float)
if mask_infeasible:
z[~self.data.feasible] = np.nan
label = ENERGY_GRID_METRIC_LABELS[metric]
final_title = title or f"Energy Grid: {label}"
text_2d, hovertemplate = self._build_hover(
metric, x_label, y_label, mask_infeasible
)
fig = plot3D_contour(
x_vals,
y_vals,
z,
final_title,
x_label,
y_label,
label,
theme=theme,
font_config=font_config,
layout_config=layout_config,
colorbar_config=DEFAULT_COLORBAR_CONFIG,
smoothing_config=SmoothingConfig(interp_factor=1, smoothing_sigma=0),
)
fig.data[0].update(
text=text_2d,
hovertemplate=hovertemplate,
contours_coloring="heatmap",
line_smoothing=0.85,
ncontours=NUM_CONTOURS,
contours=dict(showlabels=False),
)
fig.update_layout(hovermode=HOVER_MODE)
return fig
[docs]
def plot_surface(
self,
metric: EnergyGridMetric = EnergyGridMetric.TOTAL_PTS,
*,
mask_infeasible: bool = True,
show_capacity_kwh: bool = False,
title: Optional[str] = None,
theme: Optional[str] = "accumulator",
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
scene_config: SceneConfig = DEFAULT_SCENE_CONFIG,
) -> go.Figure:
"""
3D surface of any metric over the grid.
Parameters
----------
metric : EnergyGridMetric
Which z-variable to plot.
mask_infeasible : bool, optional
Replace infeasible cells with NaN so they appear blank. Default True.
show_capacity_kwh : bool, optional
Show pack capacity (kWh) instead of pack weight (kg) on axes. Default False.
title : str, optional
Custom title for the plot. Default None (auto-generated).
theme : str, optional
Theme for the surface colors. Default "accumulator".
layout_config : LayoutConfig, optional
font_config : FontConfig, optional
scene_config : SceneConfig, optional
Returns
-------
go.Figure
"""
x_vals, y_vals, x_label, y_label = self._resolve_axes(show_capacity_kwh)
z = getattr(self.data, metric.value).copy().astype(float)
if mask_infeasible:
z[~self.data.feasible] = np.nan
label = ENERGY_GRID_METRIC_LABELS[metric]
final_title = title or f"Energy Grid: {label}"
return plot3D_surface(
x_vals,
y_vals,
z,
final_title,
x_label,
y_label,
label,
theme=theme,
font_config=font_config,
layout_config=layout_config,
colorbar_config=DEFAULT_COLORBAR_CONFIG,
smoothing_config=SmoothingConfig(interp_factor=1, smoothing_sigma=0),
scene_config=scene_config,
)
[docs]
def plot_feasibility(
self,
*,
show_capacity_kwh: bool = False,
title: Optional[str] = None,
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
) -> go.Figure:
"""
Binary heatmap: green = feasible, red = infeasible.
Parameters
----------
show_capacity_kwh : bool, optional
Show pack capacity (kWh) instead of pack weight (kg) on axes. Default False.
title : str, optional
Custom title for the plot. Default None (auto-generated).
font_config : FontConfig, optional
layout_config : LayoutConfig, optional
Returns
-------
go.Figure
"""
x_vals, y_vals, x_label, y_label = self._resolve_axes(show_capacity_kwh)
z = self.data.feasible.astype(float).T
fig = go.Figure(
data=go.Heatmap(
x=x_vals,
y=y_vals,
z=z,
colorscale=[[0, "rgba(220,60,60,0.7)"], [1, "rgba(60,180,75,0.7)"]],
zmin=0,
zmax=1,
showscale=False,
)
)
fig.update_layout(
title=dict(
text=title or "Feasibility Map",
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 fig
[docs]
def grid_plot(
self,
*,
mask_infeasible: bool = True,
show_capacity_kwh: bool = False,
title: Optional[str] = None,
theme: Optional[str] = "accumulator",
layout_config: LayoutConfig = DEFAULT_LAYOUT_CONFIG,
font_config: FontConfig = DEFAULT_FONT_CONFIG,
) -> go.Figure:
"""
2x3 contour grid: total pts, endurance pts, efficiency pts,
coast trigger, net energy, endurance time.
Parameters
----------
mask_infeasible : bool, optional
Replace infeasible cells with NaN so they appear blank. Default True.
show_capacity_kwh : bool, optional
Show pack capacity (kWh) instead of pack weight (kg) on axes. Default False.
title : str, optional
Custom title for the entire grid. Default None (auto-generated).
theme : str, optional
Theme for the contour colors. Default "accumulator".
font_config : FontConfig, optional
layout_config : LayoutConfig, optional
Returns
-------
go.Figure
"""
x_vals, y_vals, x_label, y_label = self._resolve_axes(show_capacity_kwh)
panels = [
EnergyGridMetric.TOTAL_PTS,
EnergyGridMetric.ENDURANCE_PTS,
EnergyGridMetric.EFFICIENCY_PTS,
EnergyGridMetric.COAST_TRIGGER,
EnergyGridMetric.NET_ENERGY,
EnergyGridMetric.ENDURANCE_T,
]
z_data_dict = {}
z_label_dict = {}
hover_data = {}
for metric in panels:
label = ENERGY_GRID_METRIC_LABELS[metric]
z = getattr(self.data, metric.value).copy().astype(float)
if mask_infeasible:
z[~self.data.feasible] = np.nan
z_data_dict[metric.value] = z
z_label_dict[metric.value] = label
hover_data[metric.value] = self._build_hover(
metric, x_label, y_label, mask_infeasible
)
subplot_titles = [ENERGY_GRID_METRIC_LABELS[m] for m in panels]
final_title = title or "Energy Grid: All Metrics"
fig = plot_grid_3D(
x_vals,
y_vals,
z_data_dict,
subplot_titles,
final_title,
x_label,
y_label,
z_label_dict,
rows=2,
cols=3,
plot_type=PlotType.CONTOUR,
theme=theme,
font_config=font_config,
layout_config=layout_config,
smoothing_config=SmoothingConfig(interp_factor=1, smoothing_sigma=0),
)
for i, metric in enumerate(panels):
text_2d, hovertemplate = hover_data[metric.value]
fig.data[i].update(
text=text_2d,
hovertemplate=hovertemplate,
contours_coloring="heatmap",
line_smoothing=0.85,
contours=dict(showlabels=False),
)
return fig
[docs]
def summarize(self) -> None:
"""Print a concise summary of the grid sweep results."""
SEP = "=" * 55
INDENT = " "
KEY_W = 28
feasible_mask = self.data.feasible
n_total = feasible_mask.size
n_feasible = int(feasible_mask.sum())
print(SEP)
print("ENERGY GRID SWEEP SUMMARY")
print(SEP)
print(
f"{INDENT}{'Grid:':{KEY_W}}{self.data.var_1_name} x {self.data.var_2_name}"
)
print(
f"{INDENT}{'Dimensions:':{KEY_W}}{len(self.data.var_1_list)} x {len(self.data.var_2_list)} = {n_total} points"
)
print(
f"{INDENT}{'Feasible:':{KEY_W}}{n_feasible}/{n_total} ({100*n_feasible/n_total:.1f}%)"
)
print(f"{INDENT}{'Buffer:':{KEY_W}}{self.config.capacity_buffer_kwh} kWh")
if n_feasible == 0:
print(f"{INDENT}No feasible grid points found.")
print(SEP)
return
pts = self.data.total_pts.copy()
pts[~feasible_mask] = -np.inf
best_idx = np.unravel_index(pts.argmax(), pts.shape)
print(f"\n{INDENT}Best feasible point:")
print(
f"{INDENT*2}{self.data.var_1_name:{KEY_W}} = {self.data.var_1_list[best_idx[0]]:.4f}"
)
print(
f"{INDENT*2}{self.data.var_2_name:{KEY_W}} = {self.data.var_2_list[best_idx[1]]:.4f}"
)
print(
f"{INDENT*2}{'coast_trigger':{KEY_W}} = {self.data.optimal_coast_trigger[best_idx]:.2f} m/s"
)
print(
f"{INDENT*2}{'total_points':{KEY_W}} = {self.data.total_pts[best_idx]:.2f}"
)
print(
f"{INDENT*2}{'endurance_time':{KEY_W}} = {self.data.endurance_t[best_idx]:.2f} s"
)
print(
f"{INDENT*2}{'net_energy':{KEY_W}} = {self.data.net_energy_kwh[best_idx]:.3f} kWh"
)
print(
f"{INDENT*2}{'pack_capacity':{KEY_W}} = {self.data.capacity_kwh[best_idx]:.3f} kWh"
)
print(
f"{INDENT*2}{'vehicle_mass':{KEY_W}} = {self.data.vehicle_mass[best_idx]:.2f} kg"
)
print(SEP)