from pathlib import Path from types import SimpleNamespace import pytest import numpy as np import pandas as pd from batteryswap_public.evaluate import EvaluationSettings, check_plan_valid from batteryswapai.competition_planner import CompetitionPlanner, PlannerPolicy class DummyEventTimeModel: def predict_event_risk(self, snapshot): return np.array([0.9, 0.8, 0.01]) def predict_rul(self, snapshot): return np.array([20.0, 22.0, 5.0]) def _inputs(): locations = pd.DataFrame( { "battery": ["d_a", "d_b", "d_c"], "building": ["b_remote", "b_remote", "b_base"], "room": ["r_1", "r_1", "r_base"], "start_time": pd.to_datetime(["2025-01-01"] * 3), "end_time": pd.to_datetime(["2026-12-01"] * 3), } ) travel = pd.DataFrame( [ {"from": "b_base", "to": "b_base", "hours": 0.0}, {"from": "b_base", "to": "b_remote", "hours": 1.0}, {"from": "b_remote", "to": "b_base", "hours": 1.0}, {"from": "b_remote", "to": "b_remote", "hours": 0.0}, ] ) settings = EvaluationSettings(base_location="b_base", base_room="r_base") return locations, travel, settings def _scheduler_settings(**overrides): values = { "base_location": "b_base", "base_room": "r_base", "worker_limit_daily_hours": 4.0, "worker_limit_weekly_hours": 100.0, "worker_limit_daily_penalty": 100.0, "worker_limit_weekly_penalty": 100.0, "overtime_start": 4.0, "overtime_penalty_factor": 10.0, "time_per_building_change_hours": 0.0, "time_per_room_change_hours": 0.0, "time_per_battery_hours": 1.0, "early_replacement_penalty_daily": 0.5, "late_replacement_penalty_daily": 2.0, } values.update(overrides) return SimpleNamespace(**values) def test_plan_is_complete_and_batches_same_building(): locations, travel, settings = _inputs() planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy( event_risk_threshold=0.1, prediction_offset_days=0.0, building_batch_window_days=3, ), ) plan = planner.plan_snapshot( locations, locations, travel, settings, "2026-01-01", ) check_plan_valid(plan, locations, start_time=pd.Timestamp("2026-01-01")) days = plan.set_index("battery")["day"] assert days["d_a"] == days["d_b"] assert days["d_c"] > pd.Timestamp("2026-02-12") def test_prediction_arrays_can_be_supplied_for_oof_validation(): locations, travel, settings = _inputs() planner = CompetitionPlanner(DummyEventTimeModel()) plan = planner.plan_snapshot( locations, locations, travel, settings, "2026-01-01", predicted_rul=np.array([5.0, 6.0, 7.0]), predicted_risk=np.array([0.9, 0.01, 0.01]), ) inside = plan[plan["day"] <= pd.Timestamp("2026-02-12")] assert inside["battery"].tolist() == ["d_a"] def test_stale_sensor_history_reduces_proactive_risk(): locations, travel, settings = _inputs() snapshot = locations.copy() snapshot["data_gap_days"] = [0.0, 8.0, 0.0] planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy(event_risk_threshold=0.5, prediction_offset_days=0.0), ) plan = planner.plan_snapshot( snapshot, locations, travel, settings, "2026-01-01", predicted_rul=np.array([5.0, 5.0, 5.0]), predicted_risk=np.array([0.6, 0.6, 0.01]), ) inside = plan[plan["day"] <= pd.Timestamp("2026-02-12")] assert inside["battery"].tolist() == ["d_a"] def test_schedule_quota_uses_risk_ranking_when_threshold_drifts(): locations, travel, settings = _inputs() planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy( event_risk_threshold=0.99, prediction_offset_days=0.0, minimum_scheduled_batteries=2, maximum_scheduled_batteries=2, ), ) plan = planner.plan_snapshot( locations, locations, travel, settings, "2026-01-01", predicted_rul=np.array([20.0, 22.0, 5.0]), predicted_risk=np.array([0.9, 0.8, 0.01]), ) inside = plan[plan["day"] <= pd.Timestamp("2026-02-12")] assert set(inside["battery"]) == {"d_a", "d_b"} def test_capacity_balancing_accounts_for_marginal_route_cost(): locations = pd.DataFrame( { "battery": ["d_a", "d_b"], "building": ["b_a", "b_b"], "room": ["r_a", "r_b"], "start_time": pd.to_datetime(["2025-01-01"] * 2), "end_time": pd.to_datetime(["2026-12-01"] * 2), } ) travel = pd.DataFrame( [ { "from": left, "to": right, "hours": 0.0 if left == right else 0.1 if "b_base" not in (left, right) else 1.0, } for left in ("b_base", "b_a", "b_b") for right in ("b_base", "b_a", "b_b") ] ) settings = EvaluationSettings(base_location="b_base", base_room="r_base") planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy( prediction_offset_days=0.0, building_batch_window_days=0, ), ) plan = planner.plan_snapshot( locations, locations, travel, settings, "2026-01-01", predicted_rul=np.array([1.0, 2.0]), predicted_risk=np.array([0.9, 0.9]), ) days = plan.set_index("battery")["day"] assert days["d_a"] == days["d_b"] def test_scheduler_splits_visit_before_daily_limit_penalty(): candidates = pd.DataFrame( { "battery": [f"d_{index}" for index in range(6)], "building": ["b_remote"] * 6, "room": ["r_1"] * 6, "predicted_rul": [0.0] * 6, "predicted_risk": [0.9] * 6, "target_day": [0] * 6, } ) travel = pd.DataFrame( [ {"from": left, "to": right, "hours": 0.0} for left in ("b_base", "b_remote") for right in ("b_base", "b_remote") ] ) planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy(capacity_lookahead_days=7), ) scheduled = planner._schedule_candidates( candidates, travel, _scheduler_settings(), horizon=7, ) counts = scheduled.groupby("assigned_day").size() assert counts.max() <= 4 assert len(counts) >= 2 def test_scheduler_batches_trip_when_savings_exceed_early_cost(): candidates = pd.DataFrame( { "battery": ["d_a", "d_b"], "building": ["b_remote", "b_remote"], "room": ["r_1", "r_1"], "predicted_rul": [0.0, 1.0], "predicted_risk": [0.9, 0.9], "target_day": [0, 1], } ) travel = pd.DataFrame( [ {"from": "b_base", "to": "b_base", "hours": 0.0}, {"from": "b_base", "to": "b_remote", "hours": 2.0}, {"from": "b_remote", "to": "b_base", "hours": 2.0}, {"from": "b_remote", "to": "b_remote", "hours": 0.0}, ] ) planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy(capacity_lookahead_days=7), ) scheduled = planner._schedule_candidates( candidates, travel, _scheduler_settings( worker_limit_daily_hours=20.0, overtime_start=20.0, time_per_battery_hours=0.25, ), horizon=7, ) days = scheduled.set_index("battery")["assigned_day"] assert days["d_a"] == days["d_b"] == 0 def test_exact_building_route_beats_greedy_trap(): buildings = ["b_a", "b_b", "b_c"] nodes = ["b_base", *buildings] values = {} for left in nodes: for right in nodes: values[(left, right)] = 0.0 if left == right else 10.0 values[("b_base", "b_a")] = 1.0 values[("b_a", "b_c")] = 1.0 values[("b_c", "b_b")] = 1.0 values[("b_b", "b_base")] = 1.0 distances = pd.Series(values) route = CompetitionPlanner._exact_building_route( buildings, distances, "b_base", ) assert route == ["b_a", "b_c", "b_b"] def test_official_planner_interface_has_a_working_fallback(): locations, travel, settings = _inputs() battery_data = pd.DataFrame( { "device_id": ["d_a", "d_b", "d_c"], "end_time": pd.to_datetime(["2026-01-01"] * 3), "voltage": [2.7, 2.8, 3.0], "temperature": [20.0, 20.0, 20.0], } ).set_index(["device_id", "end_time"]) planner = CompetitionPlanner( DummyEventTimeModel(), PlannerPolicy(prediction_offset_days=0.0), ) plan = planner.plan(battery_data, locations, travel, settings) check_plan_valid(plan, locations, start_time=pd.Timestamp("2026-01-01")) def test_return_leg_is_charged_to_the_next_working_day(): """Characterises the evaluator rule the capacity model exists to track.""" from batteryswap_public.evaluate import evaluate_plan travel = pd.DataFrame( [ ("b_base", "b_base", 0.0), ("b_base", "b_far", 5.0), ("b_far", "b_base", 5.0), ("b_base", "b_near", 0.1), ("b_near", "b_base", 0.1), ("b_far", "b_near", 5.0), ("b_near", "b_far", 5.0), ("b_far", "b_far", 0.0), ("b_near", "b_near", 0.0), ], columns=["from", "to", "hours"], ) near = [f"d_near_{index}" for index in range(30)] batteries = ["d_far", *near] locations = pd.DataFrame( { "battery": batteries, "building": ["b_far"] + ["b_near"] * len(near), "room": ["r_far"] + ["r_near"] * len(near), "start_time": pd.to_datetime(["2025-01-01"] * len(batteries)), "end_time": pd.to_datetime(["2026-12-01"] * len(batteries)), } ) settings = EvaluationSettings(base_location="b_base", base_room="r_base") start_time = pd.Timestamp("2025-09-01") eol = pd.Series({battery: pd.NaT for battery in batteries}) def overtime(first_day_batteries, second_day_batteries): rows = [(1, battery) for battery in first_day_batteries] rows += [(2, battery) for battery in second_day_batteries] plan = pd.DataFrame( { "day": [start_time + pd.Timedelta(days=day) for day, _ in rows], "battery": [battery for _, battery in rows], } ) _, _, scores = evaluate_plan( plan, locations, travel, settings, eol_times=eol, start_time=start_time, verbose=0, ) return float(scores.overtime) far_first = overtime(["d_far"], near) far_last = overtime(near, ["d_far"]) # Both days run into overtime, so the extra return leg is billed twice. assert far_first - far_last == pytest.approx( settings.overtime_penalty_factor * (5.0 - 0.1), abs=1e-6 ) def test_route_finishes_at_the_building_nearest_base(): distances = pd.DataFrame( [ ("b_base", "b_base", 0.0), ("b_base", "b_a", 1.0), ("b_a", "b_base", 0.2), ("b_base", "b_b", 1.5), ("b_b", "b_base", 1.0), ("b_a", "b_b", 0.5), ("b_b", "b_a", 0.8), ("b_a", "b_a", 0.0), ("b_b", "b_b", 0.0), ], columns=["from", "to", "hours"], ).set_index(["from", "to"])["hours"] # Both orders cost 2.5 in travel; only the final leg differs. plain, plain_cycle, plain_last = CompetitionPlanner._order_buildings( ["b_a", "b_b"], distances, "b_base", last_leg_weight=0.0 ) weighted, weighted_cycle, weighted_last = CompetitionPlanner._order_buildings( ["b_a", "b_b"], distances, "b_base", last_leg_weight=2.0 ) assert plain_cycle == pytest.approx(weighted_cycle) == pytest.approx(2.5) assert plain[-1] == "b_b" and plain_last == pytest.approx(1.0) assert weighted[-1] == "b_a" and weighted_last == pytest.approx(0.2) def test_scheduler_merges_far_trips_that_would_breach_the_carried_daily_limit(): from batteryswap_public.evaluate import evaluate_plan travel = pd.DataFrame( [ ("b_base", "b_base", 0.0), ("b_base", "b_f1", 9.0), ("b_f1", "b_base", 9.0), ("b_base", "b_f2", 8.0), ("b_f2", "b_base", 8.0), ("b_f1", "b_f2", 0.5), ("b_f2", "b_f1", 0.5), ("b_f1", "b_f1", 0.0), ("b_f2", "b_f2", 0.0), ], columns=["from", "to", "hours"], ) locations = pd.DataFrame( { "battery": ["d_1", "d_2"], "building": ["b_f1", "b_f2"], "room": ["r_1", "r_2"], "start_time": pd.to_datetime(["2025-01-01"] * 2), "end_time": pd.to_datetime(["2026-12-01"] * 2), } ) settings = EvaluationSettings(base_location="b_base", base_room="r_base") snapshot = locations.assign(data_gap_days=0.0) planner = CompetitionPlanner(None, PlannerPolicy(event_risk_threshold=0.1)) plan = planner.plan_snapshot( snapshot, locations, travel, settings, "2025-09-01", predicted_rul=np.array([12.0, 17.0]), predicted_risk=np.array([0.9, 0.9]), ) scheduled = plan[plan["day"] <= pd.Timestamp("2025-09-01") + pd.Timedelta(days=42)] assert len(scheduled) == 2 assert scheduled["day"].nunique() == 1 _, _, scores = evaluate_plan( plan, locations, travel, settings, eol_times=pd.Series({"d_1": pd.NaT, "d_2": pd.NaT}), start_time=pd.Timestamp("2025-09-01"), verbose=0, ) assert float(scores.daily_limit) == 0.0 def test_survivor_horizon_defaults_to_the_window_and_never_falls_below_it(): planner = CompetitionPlanner(None, PlannerPolicy(survivor_horizon_scale=0.1)) snapshot = pd.DataFrame({"battery": ["d_a", "d_b"]}) assert list(planner._survivor_rul(snapshot, 42, None)) == [42.0, 42.0] values = planner._survivor_rul(snapshot, 42, np.array([500.0, np.nan])) assert values[0] == pytest.approx(50.0) assert values[1] == pytest.approx(42.0) def test_scheduled_fraction_tracks_fleet_size_and_ignores_risk_rescaling(): size = 40 locations = pd.DataFrame( { "battery": [f"d_{index}" for index in range(size)], "building": ["b_base"] * size, "room": ["r_base"] * size, "start_time": pd.to_datetime(["2025-01-01"] * size), "end_time": pd.to_datetime(["2026-12-01"] * size), } ) travel = pd.DataFrame([("b_base", "b_base", 0.0)], columns=["from", "to", "hours"]) settings = EvaluationSettings(base_location="b_base", base_room="r_base") snapshot = locations.assign(data_gap_days=0.0) risk = np.linspace(0.02, 0.6, size) rul = np.linspace(5.0, 40.0, size) horizon_end = pd.Timestamp("2025-09-01") + pd.Timedelta(days=42) def count(scale): planner = CompetitionPlanner( None, PlannerPolicy( use_expected_cost=True, risk_calibration_scale=scale, scheduled_fraction=0.25, minimum_scheduled_batteries=2, maximum_scheduled_batteries=30, ), ) plan = planner.plan_snapshot( snapshot, locations, travel, settings, "2025-09-01", predicted_rul=rul, predicted_risk=risk, ) return int(plan["day"].le(horizon_end).sum()) assert count(1.0) == 10 assert count(1.0) == count(1.9) == count(0.5) def test_plan_ignores_label_columns_and_is_scenario_independent(): """Guards rule 10: a plan must depend only on one scenario's own inputs.""" locations, travel, settings = _inputs() snapshot = locations.assign(data_gap_days=0.0) planner = CompetitionPlanner(DummyEventTimeModel(), PlannerPolicy()) first = planner.plan_snapshot(snapshot, locations, travel, settings, "2025-09-01") other = snapshot.assign(battery=["d_x", "d_y", "d_z"]) planner.plan_snapshot(other, other, travel, settings, "2025-11-01") repeat = planner.plan_snapshot(snapshot, locations, travel, settings, "2025-09-01") pd.testing.assert_frame_equal(first, repeat) leaked = snapshot.assign( effective_eol=pd.to_datetime(["2025-09-05", "2025-09-06", "2025-09-07"]), target_rul_days=[4.0, 5.0, 6.0], event_observed=[1, 1, 1], ) pd.testing.assert_frame_equal( first, planner.plan_snapshot(leaked, locations, travel, settings, "2025-09-01"), ) def test_selection_ignores_the_dataset_boundary(): """V1/V2 priced healthy swaps off locations.end_time, so workload moved with the split.""" size = 24 locations = pd.DataFrame( { "battery": [f"d_{index}" for index in range(size)], "building": ["b_base"] * size, "room": ["r_base"] * size, "start_time": pd.to_datetime(["2025-01-01"] * size), "end_time": pd.to_datetime(["2026-08-01"] * size), } ) travel = pd.DataFrame([("b_base", "b_base", 0.0)], columns=["from", "to", "hours"]) settings = EvaluationSettings(base_location="b_base", base_room="r_base") risk = np.linspace(0.05, 0.7, size) rul = np.linspace(4.0, 40.0, size) survivor = np.full(size, 180.0) horizon_end = pd.Timestamp("2025-09-01") + pd.Timedelta(days=42) planner = CompetitionPlanner( None, PlannerPolicy( use_expected_cost=True, scheduled_fraction=0.25, minimum_scheduled_batteries=2, maximum_scheduled_batteries=20, ), ) plans = [] for shift in (-180, 0, 365): snapshot = locations.assign( data_gap_days=0.0, end_time=locations["end_time"] + pd.Timedelta(days=shift), censor_proxy_rul_days=float(shift) + 365.0, ) plan = planner.plan_snapshot( snapshot, locations, travel, settings, "2025-09-01", predicted_rul=rul, predicted_risk=risk, predicted_survivor_rul=survivor, ) plans.append(plan[plan["day"].le(horizon_end)].reset_index(drop=True)) for plan in plans[1:]: pd.testing.assert_frame_equal(plans[0], plan) def test_planner_runs_with_an_artifact_that_has_no_survivor_head(): locations, travel, settings = _inputs() snapshot = locations.assign(data_gap_days=0.0) planner = CompetitionPlanner(DummyEventTimeModel(), PlannerPolicy(use_expected_cost=True)) plan = planner.plan_snapshot(snapshot, locations, travel, settings, "2025-09-01") check_plan_valid(plan, locations, start_time=pd.Timestamp("2025-09-01")) def test_weekly_guard_band_reserves_headroom_below_the_hard_limit(): """The evaluator charges the weekly penalty at >= 24h, so the planner keeps a margin.""" policy = PlannerPolicy(capacity_weekly_limit_fraction=0.95, capacity_limit_penalty_multiplier=1.5) settings = EvaluationSettings(base_location="b_base", base_room="r_base") guarded = float(settings.worker_limit_weekly_hours) * policy.capacity_weekly_limit_fraction assert guarded < float(settings.worker_limit_weekly_hours) assert guarded == pytest.approx(22.8) assert policy.capacity_limit_penalty_multiplier > 1.0 def _train_split_available() -> bool: from pathlib import Path return Path("data/raw/train/scenarios.json").exists() and Path( "submission_artifacts/planner.joblib" ).exists() @pytest.mark.skipif(not _train_split_available(), reason="train split or artifact absent") def test_official_interface_matches_our_direct_scenario_path(): """make_submissions -> Planner.plan must reproduce our iterate_scenarios path exactly.""" import joblib from batteryswap_public.utils import iterate_scenarios, load_dataset from batteryswapai.competition_features import scenario_history_snapshot locations, timeseries, eol_times, scenarios = load_dataset(Path("data/raw/train")) subset = scenarios[:3] planner = joblib.load("submission_artifacts/planner.joblib") for scenario, locs, cut, _ in iterate_scenarios(locations, timeseries, eol_times, subset): start = pd.Timestamp(scenario["start_time"]) rows = cut.reset_index() # the planner is handed exactly the scenario-visible history assert rows["end_time"].max() == start assert not (rows["end_time"] > start).any() official = planner.plan(cut, locs, scenario["travel_costs"], scenario["settings"]) snapshot = scenario_history_snapshot(cut, locs, scenario["name"], start) direct = planner.plan_snapshot( snapshot, locs, scenario["travel_costs"], scenario["settings"], start ) pd.testing.assert_frame_equal(official, direct) # ranking is deterministic across repeated calls again = planner.plan_snapshot( snapshot, locs, scenario["travel_costs"], scenario["settings"], start ) pd.testing.assert_frame_equal(direct, again)