import json
from pathlib import Path
from typing import Optional
import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
from ..compsim.utils import pretty_print_results
from ..plotting.color_themes import *
from ..plotting.plotting_constants import *
from .pack_optimizer_models import *
pio.templates.default = "plotly_white"
[docs]
class PackOptimizerResults:
"""Results for a pack optimization run, with visualization"""
def __init__(self, data: PackOptimizationData, config: PackOptimizerConfig):
self.data = data
self.config = config
[docs]
def plot_3D(
self,
x_var: str,
y_var: str,
show_infeasible: bool = True,
title: Optional[str] = None,
subtitle: Optional[str] = None,
theme: Optional[str] = None,
font_config: Optional[FontConfig] = DEFAULT_FONT_CONFIG,
layout_config: Optional[LayoutConfig] = DEFAULT_LAYOUT_CONFIG,
) -> go.Figure:
"""Create a 3D Plotly figure of the optimization results, including
all intermediate evaluations. Plots two parameters (e.g., coast_trigger and gear_ratio) vs total points.
Parameters
----------
x_var : str
Name of the parameter to plot on the x-axis (e.g., "coast_trigger")
y_var : str
Name of the parameter to plot on the y-axis (e.g., "gear_ratio")
show_infeasible : bool
Whether to include energy-infeasible points in the plot.
subtitle : str, optional
Optional subtitle displayed below the title.
theme : str, optional
Color theme name, see plotting.color_themes.
font_config : FontConfig, optional
FontConfig object for font settings.
layout_config : LayoutConfig, optional
LayoutConfig object for layout settings.
Returns
-------
go.Figure
Interactive Plotly 3D scatter figure.
"""
colors = get_theme(theme)
# Resolve variable names to indices
eval_0 = self.data.evaluation_history[0]
x_idx = eval_0.param_names.index(x_var)
y_idx = eval_0.param_names.index(y_var)
# Get data for plotting
x_list = np.array([r.params[x_idx] for r in self.data.evaluation_history])
y_list = np.array([r.params[y_idx] for r in self.data.evaluation_history])
z_list = np.array(
[
r.competition_results.total_points if r.competition_results else 0
for r in self.data.evaluation_history
]
)
feasible = np.array(
[r.feasible for r in self.data.evaluation_history], dtype=bool
)
full_title = title or "Optimal Pack Setup"
if subtitle is not None:
full_title += f"<br><span style='font-size: {font_config.medium}px; color: {TEXT_COLOR_LIGHT};'>{subtitle}</span>"
# Create 3D scatter plot
fig = go.Figure()
if show_infeasible and np.any(~feasible):
fig.add_trace(
go.Scatter3d(
x=x_list[~feasible],
y=y_list[~feasible],
z=z_list[~feasible],
mode="markers",
name="Infeasible (over energy)",
marker=dict(size=MARKER_SIZE, color="red", symbol="x", opacity=0.5),
text=[
f"{x_var}: {x:{FLOAT_PRECISION}}<br>{y_var}: {y:{FLOAT_PRECISION}}<br>Points: {z:{FLOAT_PRECISION}}<br>(Constraint violated)"
for x, y, z in zip(
x_list[~feasible], y_list[~feasible], z_list[~feasible]
)
],
hovertemplate="%{text}<extra></extra>",
)
)
if np.any(feasible):
fig.add_trace(
go.Scatter3d(
x=x_list[feasible],
y=y_list[feasible],
z=z_list[feasible],
mode="markers",
name="Feasible",
marker=dict(
size=MARKER_SIZE,
color=z_list[feasible],
colorscale=colors["colorscale"],
showscale=True,
colorbar=dict(
title=dict(
text="Total Points",
font=dict(size=font_config.medium),
),
thickness=DEFAULT_COLORBAR_CONFIG.thickness,
len=DEFAULT_COLORBAR_CONFIG.length,
tickfont=dict(size=font_config.small),
tickformat=FLOAT_PRECISION,
),
),
text=[
f"{x_var}: {x:{FLOAT_PRECISION}}<br>{y_var}: {y:{FLOAT_PRECISION}}<br>Points: {z:{FLOAT_PRECISION}}"
for x, y, z in zip(
x_list[feasible], y_list[feasible], z_list[feasible]
)
],
hovertemplate="%{text}<extra></extra>",
)
)
# Overlay optimal point
opt = self.data.optimal_evaluation
if opt is not None:
opt_x = opt.params[x_idx]
opt_y = opt.params[y_idx]
opt_z = (
opt.competition_results.total_points if opt.competition_results else 0
)
fig.add_trace(
go.Scatter3d(
x=[opt_x],
y=[opt_y],
z=[opt_z],
mode="markers",
name="Optimal",
marker=dict(
size=MARKER_SIZE_LARGE,
color="green",
),
text=[
f"<b>OPTIMAL</b><br>{x_var}: {opt_x:{FLOAT_PRECISION}}<br>{y_var}: {opt_y:{FLOAT_PRECISION}}<br>Points: {opt_z:{FLOAT_PRECISION}}"
],
hovertemplate="%{text}<extra></extra>",
)
)
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_var,
yaxis_title=y_var,
zaxis_title="Total Points",
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,
),
),
legend_title="Legend",
legend_title_font=dict(size=font_config.medium),
legend_font=dict(size=font_config.small),
hovermode=HOVER_MODE,
width=layout_config.width,
height=layout_config.height,
margin=layout_config.margin,
)
return fig
[docs]
def summarize(self) -> None:
"""Print a concise summary of the optimization results."""
opt = self.data.optimal_evaluation
print("=" * 50)
print("PACK OPTIMIZER SUMMARY")
print("=" * 50)
print(f" Evaluations: {len(self.data.evaluation_history)}")
if opt is None:
print(" No feasible solution found.")
print("=" * 50)
return
print(f" Parameters:")
for name, val in zip(opt.param_names, opt.params):
print(f" {name:28s} = {val:.4f}")
print(f" Mass: {opt.mass:.2f} kg")
print(f" Capacity: {opt.capacity:.3f} kWh")
comp = opt.competition_results
if comp is not None:
pretty_print_results(comp)
print("=" * 50)
[docs]
def save(self, path: str) -> None:
"""Serialize results to a JSON file using Pydantic's native serialization.
Parameters
----------
path : str
Destination file path (e.g., "results/run.json").
"""
payload = {
"data": self.data.model_dump(mode="json"),
"config": self.config.model_dump(mode="json"),
}
out = Path(path)
out.parent.mkdir(parents=True, exist_ok=True)
out.write_text(json.dumps(payload, indent=2))
[docs]
def load(cls, path: str) -> PackOptimizerResults:
"""Deserialize results from a JSON file.
Parameters
----------
path : str
Path to a JSON file previously written by ``save()``.
Returns
-------
PackOptimizerResults
"""
payload = json.loads(Path(path).read_text())
data = PackOptimizationData.model_validate(payload["data"])
# Reconstruct the correct config subtype
if "capacity" in payload["config"]:
config = FixedPackConfig.model_validate(payload["config"])
else:
config = PackOptimizerConfig.model_validate(payload["config"])
return cls(data=data, config=config)