import multiprocessing as mp
import sys
import traceback
from io import StringIO
from multiprocessing import Pool
from scipy.optimize import brentq
from tqdm import tqdm
from ..compsim import energy_data
from ..compsim.competition_factory import from_data
from ..compsim.models import CompetitionData
from ..compsim.utils import compute_efficiency_points
from .constants import ROUNDING_PRECISION
from .energy_grid_models import (
EnergyGridConfig,
EnergyGridData,
EnergyGridProcessInput,
EnergyGridProcessOutput,
)
from .energy_grid_results import EnergyGridResults
from .types import (
COAST_TRIGGER_PARAM,
PACK_WEIGHT_PARAM,
)
from .utils import create_steps
def _process_energy_grid_item(
inp: EnergyGridProcessInput,
) -> EnergyGridProcessOutput:
"""
Run one grid point: apply params, run static events, bisect coast_trigger on endurance.
Non-endurance events are only run once per grid point because coast_trigger has no effect on them
Parameters
----------
inp : EnergyGridProcessInput
All input data for this grid point, including comp_data and sweep variable values.
Returns
-------
EnergyGridProcessOutput
"""
try:
buf = StringIO()
old_stdout = sys.stdout
sys.stdout = buf
try:
comp = from_data(inp.comp_data)
# Apply grid variables
capacity = comp.mycar.accum.nominal_capacity
for name, value in [
(inp.var_1_name, inp.var_1_value),
(inp.var_2_name, inp.var_2_value),
]:
if name == PACK_WEIGHT_PARAM:
capacity = comp.mycar.accum.apply_pack_weight(value)
else:
comp.mycar.modify_params(name, value)
energy_budget = capacity - inp.capacity_buffer_kwh
# Run coast-independent events once
accel_res = comp.accel_event()
skidpad_res = comp.skidpad_event()
autoX_res = comp.autoX_event()
# Bisect coast_trigger on endurance
def _energy_residual(ct: float) -> float:
"""f(ct) = 22-lap net energy - budget. Root is exactly at budget."""
comp.mycar.modify_params(COAST_TRIGGER_PARAM, ct)
endu = comp.endurance_event(use_coast=True)
_, _, net_kwh_per_lap = energy_data(
endu.lapsim_results.lap_t,
endu.lapsim_results.lap_powers,
)
return net_kwh_per_lap * 22 - energy_budget
f_lo = _energy_residual(inp.coast_trigger_lo)
f_hi = _energy_residual(inp.coast_trigger_hi)
iters_used = 2 # the two bound evaluations
if f_hi <= 0:
# Even max coast_trigger stays within budget, so use it (max perf)
optimal_ct = inp.coast_trigger_hi
feasible = True
elif f_lo > 0:
# Even minimum coast_trigger exceeds budget, so it's infeasible
optimal_ct = inp.coast_trigger_lo
feasible = False
else:
# Normal case: root exists in (lo, hi)
optimal_ct = brentq(
_energy_residual,
inp.coast_trigger_lo,
inp.coast_trigger_hi,
xtol=inp.bisection_tol,
maxiter=inp.bisection_max_iter,
)
feasible = True
iters_used += inp.bisection_max_iter # upper bound; brentq converges fast
# Final endurance run at converged coast_trigger ────
comp.mycar.modify_params(COAST_TRIGGER_PARAM, optimal_ct)
endu_res = comp.endurance_event(use_coast=True)
_, _, net_kwh_per_lap = energy_data(
endu_res.lapsim_results.lap_t,
endu_res.lapsim_results.lap_powers,
)
net_energy_22 = net_kwh_per_lap * 22
eff_pts = compute_efficiency_points(
comp.scoring.efficiency,
net_kwh_per_lap, # per-lap energy
endu_res.tyour, # total endurance time (22 × lap_t)
)
total = (
accel_res.points + skidpad_res.points + autoX_res.points + endu_res.points + eff_pts
)
warnings = buf.getvalue()
return EnergyGridProcessOutput(
x_idx=inp.x_idx,
y_idx=inp.y_idx,
optimal_coast_trigger=round(optimal_ct, ROUNDING_PRECISION),
feasible=feasible,
total_points=round(total, ROUNDING_PRECISION),
accel_pts=round(accel_res.points, ROUNDING_PRECISION),
skidpad_pts=round(skidpad_res.points, ROUNDING_PRECISION),
autoX_pts=round(autoX_res.points, ROUNDING_PRECISION),
endurance_pts=round(endu_res.points, ROUNDING_PRECISION),
efficiency_pts=round(eff_pts, ROUNDING_PRECISION),
accel_t=round(accel_res.tyour, ROUNDING_PRECISION),
skidpad_t=round(skidpad_res.tyour, ROUNDING_PRECISION),
autoX_t=round(autoX_res.tyour, ROUNDING_PRECISION),
endurance_t=round(endu_res.tyour, ROUNDING_PRECISION),
net_energy_kwh=round(net_energy_22, ROUNDING_PRECISION),
capacity_kwh=round(capacity, ROUNDING_PRECISION),
vehicle_mass=round(comp.mycar.params.mass, ROUNDING_PRECISION),
bisection_iters=iters_used,
warnings=warnings,
)
finally:
sys.stdout = old_stdout
except Exception as e:
return EnergyGridProcessOutput(
x_idx=inp.x_idx,
y_idx=inp.y_idx,
optimal_coast_trigger=0,
feasible=False,
total_points=0,
net_energy_kwh=0,
capacity_kwh=0,
vehicle_mass=0,
error=f"Process error: {e}\n{traceback.format_exc()}",
)
[docs]
class EnergyGridSweeper:
"""
Energy-constrained grid characterizer.
Sweeps two design variables on a regular grid. At each grid point the
optimal coast_trigger is found via bisection so that 22-lap endurance
energy equals (pack capacity - buffer).
"""
def __init__(self, comp_data: CompetitionData, config: EnergyGridConfig) -> None:
self.comp_data = comp_data
self.config = config
self.var_1_list = create_steps(config.var_1.min, config.var_1.max, config.var_1.steps)
self.var_2_list = create_steps(config.var_2.min, config.var_2.max, config.var_2.steps)
self.grid_data = EnergyGridData.create(
config.var_1.name,
self.var_1_list,
config.var_2.name,
self.var_2_list,
nominal_pack_weight=self.comp_data.vehicle_model.accum.pack_weight,
nominal_capacity=self.comp_data.vehicle_model.accum.capacity,
energy_density=self.comp_data.vehicle_model.accum.energy_density,
)
[docs]
def sweep(
self,
verbose: bool = False,
num_processes: int | None = None,
) -> EnergyGridResults:
"""Run the full grid sweep with multiprocessing.
Parameters
----------
verbose : bool
Print per-point progress info.
num_processes : int, optional
Worker count. Defaults to CPU count.
Returns
-------
EnergyGridResults
"""
if num_processes is None:
num_processes = mp.cpu_count()
total = len(self.var_1_list) * len(self.var_2_list)
ct_lo, ct_hi = self.config.coast_trigger_bounds
print(
f"Energy Grid Sweep: {total} grid points, "
f"{num_processes} processes, "
f"coast_trigger ∈ [{ct_lo}, {ct_hi}] m/s, "
f"buffer = {self.config.capacity_buffer_kwh} kWh"
)
# Build flattened input list
inputs: list[EnergyGridProcessInput] = []
for x_idx, x_val in enumerate(self.var_1_list):
for y_idx, y_val in enumerate(self.var_2_list):
inputs.append(
EnergyGridProcessInput(
comp_data=self.comp_data,
var_1_name=self.config.var_1.name,
var_1_value=float(x_val),
var_2_name=self.config.var_2.name,
var_2_value=float(y_val),
coast_trigger_lo=ct_lo,
coast_trigger_hi=ct_hi,
capacity_buffer_kwh=self.config.capacity_buffer_kwh,
bisection_tol=self.config.bisection_tol,
bisection_max_iter=self.config.bisection_max_iter,
x_idx=x_idx,
y_idx=y_idx,
)
)
pbar = tqdm(
total=total,
desc="Energy grid",
unit="pt",
dynamic_ncols=True,
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
)
errors: list[list[str | None]] = [
[None] * len(self.var_2_list) for _ in range(len(self.var_1_list))
]
try:
with Pool(processes=num_processes) as pool:
for r in pool.imap_unordered(_process_energy_grid_item, inputs):
xi, yi = r.x_idx, r.y_idx
self.grid_data.total_pts[xi][yi] = r.total_points
self.grid_data.accel_pts[xi][yi] = r.accel_pts
self.grid_data.skidpad_pts[xi][yi] = r.skidpad_pts
self.grid_data.autoX_pts[xi][yi] = r.autoX_pts
self.grid_data.endurance_pts[xi][yi] = r.endurance_pts
self.grid_data.efficiency_pts[xi][yi] = r.efficiency_pts
self.grid_data.accel_t[xi][yi] = r.accel_t
self.grid_data.skidpad_t[xi][yi] = r.skidpad_t
self.grid_data.autoX_t[xi][yi] = r.autoX_t
self.grid_data.endurance_t[xi][yi] = r.endurance_t
self.grid_data.optimal_coast_trigger[xi][yi] = r.optimal_coast_trigger
self.grid_data.net_energy_kwh[xi][yi] = r.net_energy_kwh
self.grid_data.capacity_kwh[xi][yi] = r.capacity_kwh
self.grid_data.vehicle_mass[xi][yi] = r.vehicle_mass
self.grid_data.feasible[xi][yi] = r.feasible
self.grid_data.bisection_iters[xi][yi] = r.bisection_iters
errors[xi][yi] = r.error
if r.warnings:
pbar.write(r.warnings.strip())
if verbose:
tag = "OK" if r.feasible else "INFEASIBLE"
pbar.set_postfix_str(
f"{self.config.var_1.name}={self.var_1_list[xi]:.2f}, "
f"{self.config.var_2.name}={self.var_2_list[yi]:.2f} → "
f"ct={r.optimal_coast_trigger:.1f} [{tag}]"
)
pbar.update(1)
# Report errors
for xi in range(len(self.var_1_list)):
for yi in range(len(self.var_2_list)):
if errors[xi][yi] is not None:
print(
f"Warning: error at {self.config.var_1.name}="
f"{self.var_1_list[xi]:.4f}, "
f"{self.config.var_2.name}="
f"{self.var_2_list[yi]:.4f}:\n"
f"{errors[xi][yi]}"
)
return EnergyGridResults(self.grid_data, self.config)
except Exception as e:
pbar.write(f"Energy grid sweep failed: {e}")
raise
finally:
pbar.close()