from __future__ import annotations import argparse import itertools import json from pathlib import Path import numpy as np import pandas as pd from batteryswap_public.evaluate import evaluate_plan from batteryswap_public.utils import iterate_scenarios, load_dataset from sklearn.metrics import average_precision_score, roc_auc_score from sklearn.model_selection import GroupKFold from batteryswapai.competition_features import ( build_trajectory_matrix, scenario_history_snapshot, attach_training_targets, build_daily_features, scenario_snapshot, ) from batteryswapai.competition_model import ( blend_risk_components, blend_trajectory, fit_event_time_model, ) from batteryswapai.competition_planner import CompetitionPlanner, PlannerPolicy from train_submission import DATASET_REVISION def _numbers(value: str) -> list[float]: return [float(item.strip()) for item in value.split(",") if item.strip()] def _integers(value: str) -> list[int]: return [int(item.strip()) for item in value.split(",") if item.strip()] def _schedule_plans( fractions: str, quotas: str, bands: str ) -> list[tuple[float | None, int, int | None]]: plans: list[tuple[float | None, int, int | None]] = [ (fraction, 8, 24) for fraction in _numbers(fractions) ] plans.extend((None, minimum, maximum) for minimum, maximum in _schedule_bands(quotas, bands)) return plans def _schedule_bands(quotas: str, bands: str) -> list[tuple[int, int]]: values = [(quota, quota) for quota in _integers(quotas)] for item in bands.split(","): item = item.strip() if not item: continue minimum, maximum = item.split(":", maxsplit=1) values.append((int(minimum), int(maximum))) return list(dict.fromkeys(values)) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--dataset-path", type=Path, default=Path("data/raw/train")) parser.add_argument("--folds", type=int, default=5) parser.add_argument("--quantile", type=float, default=0.05) parser.add_argument("--inner-calibration-folds", type=int, default=3) parser.add_argument( "--outer-group", choices=("battery", "building"), default="building", help=( "Hold out entire buildings by default to emulate hidden-domain " "generalization; use battery only to reproduce legacy reports." ), ) parser.add_argument("--risk-scales", default="1.50") parser.add_argument("--gain-margins", default="10") parser.add_argument("--service-costs", default="2") parser.add_argument("--emergency-operational-scales", default="0") parser.add_argument("--offsets", default="-5") parser.add_argument("--schedule-fractions", default="0.038") parser.add_argument("--schedule-quotas", default="") parser.add_argument("--schedule-bands", default="") parser.add_argument("--capacity-lookback-days", type=int, default=42) parser.add_argument("--weekly-guard-fractions", default="0.95") parser.add_argument("--hard-limit-penalty-multipliers", default="1.5") parser.add_argument( "--predictions-csv", type=Path, help="Reuse previously generated OOF predictions for policy-only sweeps.", ) parser.add_argument( "--output", type=Path, default=Path("artifacts/submission_cv_final.json") ) return parser.parse_args() def main() -> None: args = parse_args() locations, timeseries, eol_times, scenarios = load_dataset(args.dataset_path) daily = build_daily_features(timeseries) snapshots = [] scenario_inputs = {} for scenario, locs, visible, not_dead in iterate_scenarios( locations, timeseries, eol_times, scenarios ): snapshot = scenario_history_snapshot( visible, locs, scenario["name"], scenario["start_time"] ) snapshot = attach_training_targets( snapshot, eol_times, unobserved_eol_days=float(scenario["settings"].unobserved_eol_days), ) snapshots.append(snapshot) scenario_inputs[scenario["name"]] = (scenario, locs, not_dead) training = pd.concat(snapshots, ignore_index=True) oof_risk = np.full(len(training), np.nan) oof_rul = np.full(len(training), np.nan) oof_survivor = np.full(len(training), np.nan) oof_components: dict[str, np.ndarray] = {} oof_residual = np.full(len(training), np.nan) oof_residual_weight = np.zeros(len(training)) blend_weights = (1.0, 0.0, 0.0, 0.0) outer_fold_audits: list[dict[str, int]] = [] if args.predictions_csv is not None: cached = pd.read_csv(args.predictions_csv).set_index(["scenario", "battery"]) keys = pd.MultiIndex.from_frame(training[["scenario", "battery"]]) aligned = cached.reindex(keys) if aligned[["oof_event_risk", "oof_rul_days"]].isna().any().any(): raise ValueError("Cached OOF predictions do not cover the validation rows") oof_risk[:] = aligned["oof_event_risk"].to_numpy(dtype=float) oof_rul[:] = aligned["oof_rul_days"].to_numpy(dtype=float) if "oof_survivor_rul" in aligned.columns: oof_survivor[:] = aligned["oof_survivor_rul"].to_numpy(dtype=float) else: splitter = GroupKFold(n_splits=args.folds) groups = training[args.outer_group].astype(str) for fold_number, (train_index, valid_index) in enumerate( splitter.split(training, groups=groups), start=1 ): train_rows = training.iloc[train_index] valid_rows = training.iloc[valid_index] train_batteries = set(train_rows["battery"].astype(str)) valid_batteries = set(valid_rows["battery"].astype(str)) train_buildings = set(train_rows["building"].astype(str)) valid_buildings = set(valid_rows["building"].astype(str)) if train_batteries & valid_batteries: raise AssertionError(f"outer fold {fold_number} leaks a battery") if args.outer_group == "building" and train_buildings & valid_buildings: raise AssertionError(f"outer fold {fold_number} leaks a building") # The dense trajectory forecasters have future-voltage and EOL targets. # Their matrix and labels must therefore obey the same outer split as the # snapshot model. Passing the full matrix here made the old OOF report # optimistic even though hidden inference itself remained train-only. fold_daily = daily.loc[ daily["device_id"].astype(str).isin(train_batteries) ].copy() fold_trajectory = build_trajectory_matrix(fold_daily) trajectory_batteries = set(fold_trajectory["index"]) if trajectory_batteries & valid_batteries: raise AssertionError( f"outer fold {fold_number} leaks a trajectory battery" ) fold_eol_times = eol_times.reindex(sorted(train_batteries)).copy() if set(fold_eol_times.index.astype(str)) & valid_batteries: raise AssertionError(f"outer fold {fold_number} leaks an EOL label") model = fit_event_time_model( train_rows, quantile=args.quantile, dataset_revision=DATASET_REVISION, random_state=2026 + fold_number, calibration_folds=args.inner_calibration_folds, trajectory=fold_trajectory, eol_times=fold_eol_times, ) valid = valid_rows for name, values in model.predict_risk_components(valid).items(): oof_components.setdefault( name, np.full(len(training), np.nan) )[valid_index] = values blend_weights = model.blend_weights residual, residual_weight = model._residual_risk(valid) if residual is not None: oof_residual[valid_index] = residual oof_residual_weight[valid_index] = residual_weight oof_rul[valid_index] = model.predict_rul(valid) survivor = model.predict_survivor_rul(valid) if survivor is not None: oof_survivor[valid_index] = survivor outer_fold_audits.append( { "fold": fold_number, "train_batteries": len(train_batteries), "valid_batteries": len(valid_batteries), "train_buildings": len(train_buildings), "valid_buildings": len(valid_buildings), "trajectory_batteries": len(trajectory_batteries), "battery_overlap": 0, "building_overlap": 0 if args.outer_group == "building" else -1, } ) scenario_groups = training["scenario"].astype(str).to_numpy() if oof_components: oof_risk[:] = blend_risk_components(oof_components, blend_weights, scenario_groups) # Mirror inference: the trajectory nudge is applied over complete scenarios. if np.any(oof_residual_weight > 0.0): oof_risk[:] = blend_trajectory( oof_risk, np.nan_to_num(oof_residual), oof_residual_weight, scenario_groups ) due = ( training["event_observed"].astype(bool) & training["target_rul_days"].between(0.0, 42.0) ).astype(int) classification = { "positive_rows": int(due.sum()), "total_rows": len(due), "roc_auc": float(roc_auc_score(due, oof_risk)), "average_precision": float(average_precision_score(due, oof_risk)), } adjusted_oof_risk = oof_risk.copy() gaps = pd.to_numeric(training["data_gap_days"], errors="coerce").to_numpy(dtype=float) recent_gap = (gaps > 0.0) & (gaps <= 7.0) stale = (gaps > 7.0) | ~np.isfinite(gaps) adjusted_oof_risk[recent_gap] *= 0.75 adjusted_oof_risk[stale] *= 0.10 classification["freshness_adjusted_roc_auc"] = float( roc_auc_score(due, adjusted_oof_risk) ) classification["freshness_adjusted_average_precision"] = float( average_precision_score(due, adjusted_oof_risk) ) observed = training["event_observed"].astype(bool) observed_rul_error = np.abs( oof_rul[observed] - training.loc[observed, "target_rul_days"].to_numpy() ) due_mask = due.astype(bool).to_numpy() due_rul_error = np.abs( oof_rul[due_mask] - training.loc[due_mask, "target_rul_days"].to_numpy() ) classification["all_observed_rul_mae_days"] = float(np.mean(observed_rul_error)) classification["due_within_horizon_rul_mae_days"] = float(np.mean(due_rul_error)) policy_results = [] case_results = [] placeholder_model = None grid = itertools.product( _numbers(args.risk_scales), _numbers(args.gain_margins), _numbers(args.service_costs), _numbers(args.emergency_operational_scales), _numbers(args.weekly_guard_fractions), _numbers(args.hard_limit_penalty_multipliers), _numbers(args.offsets), _schedule_plans( args.schedule_fractions, args.schedule_quotas, args.schedule_bands, ), ) for ( risk_scale, gain_margin, service_cost, emergency_scale, weekly_guard, hard_multiplier, offset, (fraction, minimum, maximum), ) in grid: policy_id = f"p{len(policy_results):03d}" planner = CompetitionPlanner( placeholder_model, PlannerPolicy( event_risk_threshold=0.50, prediction_offset_days=offset, use_expected_cost=True, risk_calibration_scale=risk_scale, expected_service_cost_hours=service_cost, expected_gain_margin=gain_margin, emergency_operational_scale=emergency_scale, capacity_lookahead_days=21, capacity_lookback_days=args.capacity_lookback_days, capacity_weekly_limit_fraction=weekly_guard, capacity_limit_penalty_multiplier=hard_multiplier, scheduled_fraction=fraction, minimum_scheduled_batteries=minimum, maximum_scheduled_batteries=maximum, ), ) scores = [] scheduled_counts = [] for scenario_name, (scenario, locs, not_dead) in scenario_inputs.items(): mask = training["scenario"].eq(scenario_name).to_numpy() snapshot = training.loc[mask] risk = oof_risk[mask] rul = oof_rul[mask] survivor = oof_survivor[mask] plan = planner.plan_snapshot( snapshot, locs, scenario["travel_costs"], scenario["settings"], scenario["start_time"], predicted_rul=rul, predicted_risk=risk, predicted_survivor_rul=( None if np.isnan(survivor).all() else survivor ), ) start = pd.Timestamp(scenario["start_time"]) horizon_end = start + pd.Timedelta( days=scenario["settings"].planning_window_days ) scheduled = plan["day"].le(horizon_end) scheduled_count = int(scheduled.sum()) scheduled_counts.append(scheduled_count) _, _, score = evaluate_plan( plan, locs, scenario["travel_costs"], scenario["settings"], eol_times=not_dead, start_time=start, verbose=0, ) scores.append(score) required = pd.to_datetime(not_dead).between( start, horizon_end, inclusive="right", ) required_ids = set(not_dead.index[required].astype(str)) scheduled_ids = set(plan.loc[scheduled, "battery"].astype(str)) case_results.append( { "policy_id": policy_id, "scenario": scenario_name, "start_time": start.isoformat(), "scheduled_count": scheduled_count, "required_count": len(required_ids), "true_positive_count": len(required_ids & scheduled_ids), "missed_count": len(required_ids - scheduled_ids), **{key: float(value) for key, value in score.items()}, } ) mean_score = pd.concat(scores, axis=1).mean(axis=1) policy_results.append( { "policy_id": policy_id, "risk_calibration_scale": risk_scale, "expected_gain_margin": gain_margin, "expected_service_cost_hours": service_cost, "emergency_operational_scale": emergency_scale, "prediction_offset_days": offset, "scheduled_fraction": fraction, "minimum_scheduled_batteries": minimum, "maximum_scheduled_batteries": maximum, "capacity_weekly_limit_fraction": weekly_guard, "capacity_limit_penalty_multiplier": hard_multiplier, "stale_risk_cutoff_days": planner.policy.stale_risk_cutoff_days, "recent_gap_risk_factor": planner.policy.recent_gap_risk_factor, "stale_risk_factor": planner.policy.stale_risk_factor, "building_batch_window_days": planner.policy.building_batch_window_days, "capacity_operational_cost_weight": planner.policy.capacity_operational_cost_weight, "mean_scheduled_batteries": float(np.mean(scheduled_counts)), **{key: float(value) for key, value in mean_score.items()}, } ) policy_results.sort(key=lambda item: item["total_cost"]) report = { "dataset_revision": DATASET_REVISION, "folds": args.folds, "inner_calibration_folds": args.inner_calibration_folds, "outer_group": args.outer_group, "prediction_source": ( args.predictions_csv.as_posix() if args.predictions_csv is not None else "generated by this run" ), "outer_fold_trajectory_scope": ( "caller-supplied predictions; not refit" if args.predictions_csv is not None else "outer-train batteries only" ), "outer_fold_eol_scope": ( "caller-supplied predictions; not refit" if args.predictions_csv is not None else "outer-train batteries only" ), "outer_fold_audits": outer_fold_audits, "quantile": args.quantile, "classification": classification, "policies": policy_results, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(report, indent=2), encoding="utf-8") predictions = training[["scenario", "battery", "target_rul_days", "event_observed"]].copy() predictions["oof_event_risk"] = oof_risk predictions["oof_rul_days"] = oof_rul predictions["oof_survivor_rul"] = oof_survivor predictions.to_csv(args.output.with_suffix(".csv"), index=False) pd.DataFrame(case_results).to_csv( args.output.with_suffix(".cases.csv"), index=False ) print(json.dumps(report, indent=2)) if __name__ == "__main__": main()