import multiprocessing as mp
import sys
import traceback
from io import StringIO
from multiprocessing import Pool
import numpy as np
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, compute_event_points
from .constants import ROUNDING_PRECISION
from .energy_grid_models import (
EnergyGridConfig,
EnergyGridData,
EnergyGridProcessInput,
EnergyGridProcessOutput,
)
from .energy_grid_results import EnergyGridResults
from .models import SweepParamConfig
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):
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 = []
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 = [[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()