| from __future__ import annotations |
|
|
| import os |
| import sys |
| from pathlib import Path |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent / "src")) |
|
|
| import joblib |
| import pandas as pd |
| from batteryswap_public.utils import iterate_scenarios, load_dataset |
|
|
| from batteryswapai.competition_features import scenario_history_snapshot |
|
|
|
|
| def main() -> None: |
| dataset_path = Path(os.environ.get("BATTERYSWAP_DATASET_PATH", "/tmp/data")) |
| artifact_path = Path( |
| os.environ.get("BATTERYSWAP_PLANNER_PATH", "submission_artifacts/planner.joblib") |
| ) |
| splits = os.environ.get("BATTERYSWAP_SPLITS", "public,private").split(",") |
| output_path = Path(os.environ.get("BATTERYSWAP_SUBMISSION_PATH", "submission.csv")) |
|
|
| planner = joblib.load(artifact_path) |
|
|
| plans = [] |
| for split in splits: |
| locations, timeseries, hidden_eol_times, scenarios = load_dataset(dataset_path / split) |
|
|
| for scenario, locs, visible_history, _ in iterate_scenarios( |
| locations, |
| timeseries, |
| hidden_eol_times, |
| scenarios, |
| ): |
| snapshot = scenario_history_snapshot( |
| visible_history, |
| locs, |
| scenario["name"], |
| scenario["start_time"], |
| ) |
|
|
| plan = planner.plan_snapshot( |
| snapshot, |
| locs, |
| scenario["travel_costs"], |
| scenario["settings"], |
| scenario["start_time"], |
| ) |
|
|
| plan["split"] = split |
| plan["scenario"] = scenario["name"] |
| plans.append(plan) |
|
|
| submission = pd.concat(plans, ignore_index=True) |
| submission.to_csv(output_path, index=False) |
|
|
| if not output_path.exists(): |
| raise RuntimeError(f"Submission was not created: {output_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
