Use simulated annealing planner for submission

batteryswap_example/planners/best.pickle

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e8a24a181db4723084582e1f526305bde6be6496dbe305af888ac39259611783
-size 6519559
+oid sha256:3cca90ca40103c1e80c3251f185a99ca2668019c80bf076551352315aadaeb91
+size 6525989

batteryswap_example/train.py

@@ -16,7 +16,6 @@ import pandas
 from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 import structlog
-from ortools.sat.python import cp_model
 
 from batteryswap_public.interfaces import Planner, RULModel
 from batteryswap_public.utils import load_dataset, iterate_scenarios
@@ -25,287 +24,365 @@ from batteryswap_public.evaluate import evaluate_plan, check_plan_valid
 log = structlog.get_logger()
 
 
-# --------------------------------------------------------------------------
-# Planner: CP-SAT (MILP-style) assignment of batteries to swap-days,
-# accounting for room/building visit costs and daily/weekly worker-hour
-# limits. Replaces the OrderedPlanner heuristic.
-# --------------------------------------------------------------------------
-
-COST_SCALE = 600
-MINUTES_PER_HOUR = 60
-
-DEVICE_COLUMN = "device_id"
-TIME_COLUMN = "end_time"
-VALUE_COLUMNS = ("voltage", "temperature")
-
-
-def normalize_timeseries(timeseries):
-    frame = timeseries.copy()
-    missing_identity = {DEVICE_COLUMN, TIME_COLUMN} - set(frame.columns)
-    if missing_identity:
-        frame = frame.reset_index()
-
-    required = {DEVICE_COLUMN, TIME_COLUMN, *VALUE_COLUMNS}
-    missing = required - set(frame.columns)
-    if missing:
-        raise ValueError(f"Timeseries is missing required columns: {sorted(missing)}")
-
-    frame = frame.loc[:, [DEVICE_COLUMN, TIME_COLUMN, *VALUE_COLUMNS]].copy()
-    frame[TIME_COLUMN] = pandas.to_datetime(frame[TIME_COLUMN])
-    frame[DEVICE_COLUMN] = frame[DEVICE_COLUMN].astype(str)
-    for column in VALUE_COLUMNS:
-        frame[column] = pandas.to_numeric(frame[column], errors="coerce")
-    return frame.sort_values([DEVICE_COLUMN, TIME_COLUMN], kind="stable").reset_index(drop=True)
-
-
-def _setting(settings, name, default):
-    if isinstance(settings, dict):
-        return settings.get(name, default)
-    return getattr(settings, name, default)
-
-
-def _normalize_locations(locations):
-    frame = locations.copy().reset_index(drop=True)
-    aliases = {"device_id": "battery", "building_id": "building", "room_id": "room"}
-    frame = frame.rename(columns={old: new for old, new in aliases.items() if new not in frame})
-    required = {"battery", "building", "room"}
-    missing = required - set(frame.columns)
-    if missing:
-        raise ValueError(f"Locations is missing required columns: {sorted(missing)}")
-    if frame["battery"].duplicated().any():
-        raise ValueError("Each battery must have exactly one location")
-    return frame
-
-
-def _travel_lookup(travel_costs):
-    frame = travel_costs.copy()
-    required = {"from", "to", "hours"}
-    missing = required - set(frame.columns)
-    if missing:
-        raise ValueError(f"Travel costs is missing required columns: {sorted(missing)}")
-    return {(str(row["from"]), str(row["to"])): float(row["hours"]) for _, row in frame.iterrows()}
-
-
-def order_daily_route(batteries, locations, travel_costs, base_building):
-    selected = set(str(value) for value in batteries)
-    if not selected:
-        return []
-    loc = _normalize_locations(locations).set_index("battery")
-    travel = _travel_lookup(travel_costs)
-    buildings = set(loc.loc[list(selected), "building"].astype(str))
-    current = str(base_building)
-    building_order = []
-    while buildings:
-        next_building = min(
-            buildings,
-            key=lambda building: (
-                travel.get((current, building), 0.0 if current == building else float("inf")),
-                building,
-            ),
-        )
-        building_order.append(next_building)
-        buildings.remove(next_building)
-        current = next_building
-
-    ordered = []
-    for building in building_order:
-        subset = loc.loc[list(selected)]
-        subset = subset.loc[subset["building"].astype(str) == building].copy()
-        subset["battery_key"] = subset.index.astype(str)
-        subset = subset.sort_values(["room", "battery_key"], kind="stable")
-        ordered.extend(subset.index.astype(str).tolist())
-    return ordered
-
-
-class MilpPlanner(Planner):
-    def __init__(self, rul_estimator, solver_time_limit_seconds=20.0, late_risk_multiplier=1.0):
+class OrderedPlanner(Planner):
+    def __init__(
+        self,
+        rul_estimator,
+        safety_days=14,
+        max_swaps_per_day=6,
+        sa_iterations=1000,
+        initial_temperature=500.0,
+        cooling_rate=0.995,
+        random_seed=42,
+    ):
         self.rul_estimator = rul_estimator
-        self.solver_time_limit_seconds = float(solver_time_limit_seconds)
-        self.late_risk_multiplier = float(late_risk_multiplier)
-
-    def _expected_costs(self, timeseries, batteries, settings):
-        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
-        normalized = normalize_timeseries(timeseries)
-        scenario_start = normalized["end_time"].max().normalize()
-        horizon_end = scenario_start + pandas.Timedelta(days=horizon)
-        emergency_delay = 6 - horizon_end.weekday()
-        costs = self.rul_estimator.expected_replacement_costs(
-            timeseries,
-            horizon_days=horizon,
-            early_penalty=float(_setting(settings, "early_replacement_penalty_daily", 0.5)),
-            late_penalty=float(_setting(settings, "late_replacement_penalty_daily", 10.0))
-            * float(getattr(self, "late_risk_multiplier", 1.0)),
-            no_swap_extension_days=emergency_delay,
+        self.safety_days = safety_days
+        self.max_swaps_per_day = max_swaps_per_day
+
+        # Simulated annealing parameters
+        self.sa_iterations = sa_iterations
+        self.initial_temperature = initial_temperature
+        self.cooling_rate = cooling_rate
+        self.random_seed = random_seed
+
+    def _score_plan(
+        self,
+        plan,
+        locations,
+        travel_costs,
+        settings,
+        predicted_eol,
+        start_time,
+    ):
+        """
+        Evaluate a candidate schedule.
+
+        True EOL is unavailable at prediction time, so predicted EOL
+        is used as a surrogate inside the official cost function.
+        """
+
+        try:
+            check_plan_valid(
+                plan,
+                locations,
+                start_time=start_time,
+            )
+
+            _, _, overall = evaluate_plan(
+                plan,
+                locations,
+                travel_costs,
+                settings,
+                eol_times=predicted_eol,
+                start_time=start_time,
+            )
+
+            return float(overall["total_cost"])
+
+        except Exception:
+            # Invalid schedules should never be selected
+            return float("inf")
+
+    def _make_initial_plan(
+        self,
+        order,
+        start_time,
+    ):
+        """
+        Build the 14-day heuristic schedule that we already know works.
+        """
+
+        planned_days = []
+
+        current_day = start_time
+        swaps_today = 0
+
+        for battery, row in order.iterrows():
+
+            desired_day = max(
+                start_time,
+                row["target_day"],
+            )
+
+            if desired_day > current_day:
+                current_day = desired_day
+                swaps_today = 0
+
+            if swaps_today >= self.max_swaps_per_day:
+                current_day += pandas.Timedelta(days=1)
+                swaps_today = 0
+
+            planned_days.append(current_day)
+            swaps_today += 1
+
+        planned_days = pandas.to_datetime(
+            planned_days
+        ).normalize()
+
+        return pandas.DataFrame({
+            "day": planned_days,
+            "battery": order.index,
+        }).reset_index(drop=True)
+
+    def _neighbor(
+        self,
+        plan,
+        start_time,
+        rng,
+    ):
+        """
+        Produce a nearby schedule.
+
+        Two types of moves:
+        1. Move one battery a few days earlier/later.
+        2. Swap the scheduled days of two batteries.
+        """
+
+        candidate = plan.copy()
+
+        n = len(candidate)
+
+        if n < 2:
+            return candidate
+
+        move_type = rng.integers(0, 2)
+
+        
+        #Step 1: shift one battery
+        if move_type == 0:
+
+            idx = int(rng.integers(0, n))
+
+            # Bias moves toward earlier scheduling
+            if rng.random() < 0.7:
+                shift = -int(rng.integers(1, 15))
+            else:
+                shift = int(rng.integers(1, 8))
+
+            new_day = (
+                candidate.loc[idx, "day"]
+                + pandas.Timedelta(days=shift)
+            )
+
+            # Never schedule before planning begins
+            if new_day < start_time:
+                new_day = start_time
+
+            candidate.loc[idx, "day"] = new_day
+
+        
+        # Step 2: swap two batteries' days
+        else:
+
+            i, j = rng.choice(
+                n,
+                size=2,
+                replace=False,
+            )
+
+            day_i = candidate.loc[i, "day"]
+            day_j = candidate.loc[j, "day"]
+
+            candidate.loc[i, "day"] = day_j
+            candidate.loc[j, "day"] = day_i
+
+        candidate["day"] = pandas.to_datetime(
+            candidate["day"]
+        ).dt.normalize()
+
+        return candidate
+
+    def plan(
+        self,
+        battery_data,
+        locations,
+        travel_costs,
+        settings,
+    ):
+
+        
+        #Predict Remaining Useful Life
+        rul = self.rul_estimator.predict(
+            battery_data
+        )
+
+        rul_days = rul["p50"]
+
+        start_time = (
+            battery_data
+            .reset_index()["end_time"]
+            .max()
+            .normalize()
+        )
+
+        predicted_eol = (
+            start_time
+            + pandas.to_timedelta(
+                rul_days,
+                unit="D",
+            )
+        )
+
+        # Official evaluation uses calendar days
+        predicted_eol = predicted_eol.dt.normalize()
+
+        
+        #Prepare batteries and target dates
+       
+
+        loc = (
+            locations
+            .copy()
+            .set_index("battery")
         )
-        expected_columns = list(range(horizon + 1)) + ["no_swap"]
-        costs = costs.reindex(index=batteries, columns=expected_columns)
-        finite = costs.to_numpy(dtype=float)
-        fallback = float(np.nanmax(finite[np.isfinite(finite)])) if np.isfinite(finite).any() else 1e6
-        return costs.replace([np.inf, -np.inf], np.nan).fillna(fallback + 1e3)
-
-    def _solve_assignments(self, expected_costs, locations, travel_costs, settings):
-        loc = _normalize_locations(locations).set_index("battery")
-        batteries = list(expected_costs.index.astype(str))
-        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
-        real_days = list(range(horizon + 1))
-        actions = real_days + ["no_swap"]
-        base = str(_setting(settings, "base_location", ""))
-        base_room = str(_setting(settings, "base_room", ""))
-        travel = _travel_lookup(travel_costs)
-
-        model = cp_model.CpModel()
-        assignment = {
-            (battery, action): model.new_bool_var(f"x_{index}_{action}")
-            for index, battery in enumerate(batteries)
-            for action in actions
-        }
-        for battery in batteries:
-            model.add_exactly_one(assignment[battery, action] for action in actions)
-
-        rooms = sorted(loc.loc[batteries, "room"].astype(str).unique())
-        buildings = sorted(loc.loc[batteries, "building"].astype(str).unique())
-        battery_rooms = loc.loc[batteries, "room"].astype(str).to_dict()
-        battery_buildings = loc.loc[batteries, "building"].astype(str).to_dict()
-        room_members = {
-            room: [battery for battery in batteries if battery_rooms[battery] == room] for room in rooms
-        }
-        building_members = {
-            building: [battery for battery in batteries if battery_buildings[battery] == building]
-            for building in buildings
-        }
-        room_visit = {(room, day): model.new_bool_var(f"room_{room}_{day}") for room in rooms for day in real_days}
-        building_visit = {
-            (building, day): model.new_bool_var(f"building_{building}_{day}")
-            for building in buildings
-            for day in real_days
-        }
-
-        for day in real_days:
-            for room in rooms:
-                members = room_members[room]
-                for battery in members:
-                    model.add(assignment[battery, day] <= room_visit[room, day])
-                model.add(room_visit[room, day] <= sum(assignment[battery, day] for battery in members))
-            for building in buildings:
-                members = building_members[building]
-                for battery in members:
-                    model.add(assignment[battery, day] <= building_visit[building, day])
-                model.add(
-                    building_visit[building, day] <= sum(assignment[battery, day] for battery in members)
+
+        # Align prediction order explicitly
+        loc["predicted_eol"] = predicted_eol.reindex(
+            loc.index
+        )
+
+        loc["target_day"] = (
+            loc["predicted_eol"]
+            - pandas.to_timedelta(
+                self.safety_days,
+                unit="D",
+            )
+        )
+
+        loc["target_day"] = (
+            loc["target_day"]
+            .dt.normalize()
+            .clip(lower=start_time)
+        )
+
+        
+        #Initial urgency/location ordering
+        
+        location_columns = []
+
+        for candidate_column in [
+            "building",
+            "building_id",
+            "room",
+            "room_id",
+        ]:
+            if candidate_column in loc.columns:
+                location_columns.append(
+                    candidate_column
                 )
 
-        battery_minutes = round(float(_setting(settings, "time_per_battery_hours", 0.25)) * 60)
-        room_minutes = round(float(_setting(settings, "time_per_room_change_hours", 0.5)) * 60)
-        building_minutes = round(float(_setting(settings, "time_per_building_change_hours", 1.0)) * 60)
-        building_work = {}
-        for building in buildings:
-            round_trip = travel.get((base, building), 0.0 if base == building else 24.0)
-            round_trip += travel.get((building, base), 0.0 if base == building else 24.0)
-            building_work[building] = round(round_trip * 60) + (0 if building == base else building_minutes)
-
-        maximum_daily = (
-            len(batteries) * battery_minutes + len(rooms) * room_minutes + sum(building_work.values())
+        sort_columns = [
+            "target_day"
+        ] + location_columns
+
+        order = loc.sort_values(
+            sort_columns,
+            ascending=True,
+        )
+
+       
+        #Start from our V3 14-day heuristic
+        current_plan = self._make_initial_plan(
+            order,
+            start_time,
+        )
+
+        # Predicted EOL series must use battery IDs
+        predicted_eol_for_score = (
+            loc["predicted_eol"].copy()
         )
-        daily_work = {}
-        daily_overtime = {}
-        daily_limit_hit = {}
-        overtime_start = round(float(_setting(settings, "overtime_start", 8.0)) * 60)
-        daily_limit = round(float(_setting(settings, "worker_limit_daily_hours", 24.0)) * 60)
-
-        for day in real_days:
-            work = model.new_int_var(0, maximum_daily, f"work_{day}")
-            expression = (
-                battery_minutes * sum(assignment[battery, day] for battery in batteries)
-                + room_minutes * sum(room_visit[room, day] for room in rooms if room != base_room)
-                + sum(building_work[building] * building_visit[building, day] for building in buildings)
+
+        current_cost = self._score_plan(
+            current_plan,
+            locations,
+            travel_costs,
+            settings,
+            predicted_eol_for_score,
+            start_time,
+        )
+
+        best_plan = current_plan.copy()
+        best_cost = current_cost
+
+        
+        #Simulated annealing
+        rng = numpy.random.default_rng(
+            self.random_seed
+        )
+
+        temperature = self.initial_temperature
+
+        for iteration in range(
+            self.sa_iterations
+        ):
+
+            candidate_plan = self._neighbor(
+                current_plan,
+                start_time,
+                rng,
             )
-            model.add(work == expression)
-            daily_work[day] = work
-
-            overtime = model.new_int_var(0, maximum_daily, f"overtime_{day}")
-            model.add(overtime >= work - overtime_start)
-            daily_overtime[day] = overtime
-
-            hit = model.new_bool_var(f"daily_limit_hit_{day}")
-            model.add(work <= daily_limit + maximum_daily * hit)
-            daily_limit_hit[day] = hit
-
-        weekly_limit_hit = {}
-        weekly_limit = round(float(_setting(settings, "worker_limit_weekly_hours", 24.0)) * 60)
-        for week_start in range(0, len(real_days), 7):
-            week_days = real_days[week_start : week_start + 7]
-            hit = model.new_bool_var(f"weekly_limit_hit_{week_start // 7}")
-            weekly_maximum = maximum_daily * len(week_days)
-            model.add(
-                sum(daily_work[day] for day in week_days) <= max(weekly_limit - 1, -1) + weekly_maximum * hit
+
+            candidate_cost = self._score_plan(
+                candidate_plan,
+                locations,
+                travel_costs,
+                settings,
+                predicted_eol_for_score,
+                start_time,
             )
-            weekly_limit_hit[week_start] = hit
-
-        objective_terms = []
-        for battery_index, battery in enumerate(batteries):
-            for action_index, action in enumerate(actions):
-                coefficient = int(round(float(expected_costs.loc[battery, action]) * COST_SCALE))
-                coefficient += action_index + battery_index % 3
-                objective_terms.append(coefficient * assignment[battery, action])
-
-        minute_cost = COST_SCALE // MINUTES_PER_HOUR
-        objective_terms.extend(minute_cost * daily_work[day] for day in real_days)
-        overtime_factor = float(_setting(settings, "overtime_penalty_factor", 2.0))
-        overtime_minute_cost = int(round(overtime_factor * minute_cost))
-        objective_terms.extend(overtime_minute_cost * daily_overtime[day] for day in real_days)
-        daily_penalty = int(round(float(_setting(settings, "worker_limit_daily_penalty", 100.0)) * COST_SCALE))
-        weekly_penalty = int(round(float(_setting(settings, "worker_limit_weekly_penalty", 100.0)) * COST_SCALE))
-        objective_terms.extend(daily_penalty * value for value in daily_limit_hit.values())
-        objective_terms.extend(weekly_penalty * value for value in weekly_limit_hit.values())
-        model.minimize(sum(objective_terms))
-
-        solver = cp_model.CpSolver()
-        solver.parameters.max_time_in_seconds = self.solver_time_limit_seconds
-        solver.parameters.num_search_workers = 1
-        solver.parameters.random_seed = 0
-        status = solver.solve(model)
-        if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
-            return {
-                battery: min(actions, key=lambda action: (expected_costs.loc[battery, action], str(action)))
-                for battery in batteries
-            }
-        return {
-            battery: next(action for action in actions if solver.value(assignment[battery, action]))
-            for battery in batteries
-        }
-
-    def plan(self, timeseries, locations, travel_costs, settings):
-        loc = _normalize_locations(locations)
-        batteries = sorted(loc["battery"].astype(str).tolist())
-        normalized_timeseries = normalize_timeseries(timeseries)
-        if normalized_timeseries.empty:
-            if "end_time" not in loc:
-                raise ValueError("Cannot determine scenario start time")
-            start_time = pandas.to_datetime(loc["end_time"]).max().normalize()
-        else:
-            start_time = normalized_timeseries["end_time"].max().normalize()
 
-        expected_costs = self._expected_costs(timeseries, batteries, settings)
-        assignments = self._solve_assignments(expected_costs, loc, travel_costs, settings)
-        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
-        base = str(_setting(settings, "base_location", ""))
-        records = []
-        for day in range(horizon + 1):
-            selected = [battery for battery in batteries if assignments[battery] == day]
-            for battery in order_daily_route(selected, loc, travel_costs, base):
-                records.append({"day": start_time + pandas.Timedelta(days=day), "battery": battery})
+            delta = (
+                candidate_cost
+                - current_cost
+            )
+
+            # Always accept improvements
+            accept = delta <= 0
+
+            # Sometimes accept worse solutions
+            # so SA can escape local minima
+            if (
+                not accept
+                and numpy.isfinite(candidate_cost)
+                and temperature > 1e-8
+            ):
+
+                probability = numpy.exp(
+                    -delta / temperature
+                )
+
+                if rng.random() < probability:
+                    accept = True
+
+            if accept:
+                current_plan = candidate_plan
+                current_cost = candidate_cost
 
-        no_swap_day = start_time + pandas.Timedelta(days=horizon + 1)
-        for battery in sorted(battery for battery in batteries if assignments[battery] == "no_swap"):
-            records.append({"day": no_swap_day, "battery": battery})
+            # Remember the best schedule ever seen
+            if current_cost < best_cost:
+                best_plan = current_plan.copy()
+                best_cost = current_cost
 
-        plan = pandas.DataFrame.from_records(records, columns=["day", "battery"]).reset_index(drop=True)
-        plan["day"] = pandas.to_datetime(plan["day"])
-        check_plan_valid(plan, loc, start_time=start_time)
-        return plan
+            temperature *= self.cooling_rate
 
+        
+        #Final validation
+        best_plan["day"] = (
+            pandas.to_datetime(
+                best_plan["day"]
+            )
+            .dt.normalize()
+        )
 
+        check_plan_valid(
+            best_plan,
+            locations,
+            start_time=start_time,
+        )
 
+        return best_plan
+    
+    
 class DummyRULModel(RULModel):
     # RUL model that predicts (no information rate)
     # FIXME: make a model that actually uses the data to improve predictions
@@ -486,38 +563,6 @@ class DummyRULModel(RULModel):
         out = pd.DataFrame(preds, index=pd.Index(ids, name=self.group_col))
         return out[self.quantile_cols]
 
-    def expected_replacement_costs(
-        self,
-        timeseries: pd.DataFrame,
-        horizon_days,
-        early_penalty,
-        late_penalty,
-        no_swap_extension_days,
-    ) -> pd.DataFrame:
-        """Adapter so MilpPlanner can consume this point-forecast RUL model.
-
-        Treats the p50 prediction as a deterministic failure day (no
-        distributional information available) and prices each candidate
-        swap day with a linear early/late penalty around that day.
-        """
-        horizon = int(round(float(horizon_days)))
-        emergency_day = horizon + int(round(float(no_swap_extension_days)))
-        predicted_days = self.predict(timeseries)['p50'].clip(lower=0.0)
-
-        replacement_days = np.arange(horizon + 1, dtype=float)
-        rows = {}
-        for battery, day in predicted_days.items():
-            early = np.maximum(day - replacement_days, 0.0)
-            late = np.maximum(replacement_days - day, 0.0)
-            costs = float(early_penalty) * early + float(late_penalty) * late
-            no_swap_cost = (
-                float(late_penalty) * max(emergency_day - day, 0.0) if day <= horizon else 0.0
-            )
-            rows[battery] = np.append(costs, no_swap_cost)
-
-        columns = list(range(horizon + 1)) + ["no_swap"]
-        return pd.DataFrame.from_dict(rows, orient='index', columns=columns)
-
 
 def train_rul_model(locations, timeseries, eol_times, scenarios, limit_scenarios=None):
 
@@ -577,8 +622,6 @@ class Config(BaseSettings):
  
     dataset_path: Optional[Path] = None
     split : str = 'train'
-    solver_time_limit_seconds: float = 20.0
-    late_risk_multiplier: float = 1.0
 
 def main():
     cfg = Config()
@@ -609,11 +652,7 @@ def main():
         travel_costs = scenario['travel_costs']
         settings = scenario['settings']
 
-        planner = MilpPlanner(
-            rul_model,
-            solver_time_limit_seconds=cfg.solver_time_limit_seconds,
-            late_risk_multiplier=cfg.late_risk_multiplier,
-        )
+        planner = OrderedPlanner(rul_model)
         plan = planner.plan(cut, locs, travel_costs, settings)
 
         start_time = pandas.Timestamp(scenario['start_time'])
@@ -624,11 +663,7 @@ def main():
 
 
     # Save best planner
-    planner = MilpPlanner(
-        rul_model,
-        solver_time_limit_seconds=cfg.solver_time_limit_seconds,
-        late_risk_multiplier=cfg.late_risk_multiplier,
-    )
+    planner = OrderedPlanner(rul_model)
 
     planner_path = 'batteryswap_example/planners/best.pickle' 
     with open(planner_path, "wb") as f: