Fix MILP never solving

batteryswap_example/planners/best.pickle

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6508510aaf0ee366dac78ac2743e6899e8365508af79777929fa3e417d8d25a8
-size 201707
+oid sha256:098ee847de3b9cbe57b2ccdadc65213625cd06585bd2fcb492bb8f8eed56693c
+size 201765

batteryswap_example/train.py

@@ -1,18 +1,11 @@
-
-import pandas
-import numpy
 import pickle
-from typing import Optional, Sequence
-from pathlib import Path
-import pathlib
 import os
+from pathlib import Path
+from typing import Optional, Sequence
 
-
+import pandas
 import pandas as pd
-import numpy
 import numpy as np
-import pandas
-from pydantic import Field
 from pydantic_settings import BaseSettings, SettingsConfigDict
 import structlog
 from ortools.sat.python import cp_model
@@ -23,13 +16,6 @@ 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
 
@@ -54,7 +40,9 @@ def normalize_timeseries(timeseries):
     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)
+    return frame.sort_values([DEVICE_COLUMN, TIME_COLUMN], kind="stable").reset_index(
+        drop=True
+    )
 
 
 def _setting(settings, name, default):
@@ -66,7 +54,9 @@ def _setting(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})
+    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:
@@ -82,7 +72,10 @@ def _travel_lookup(travel_costs):
     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()}
+    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):
@@ -98,7 +91,9 @@ def order_daily_route(batteries, locations, travel_costs, base_building):
         next_building = min(
             buildings,
             key=lambda building: (
-                travel.get((current, building), 0.0 if current == building else float("inf")),
+                travel.get(
+                    (current, building), 0.0 if current == building else float("inf")
+                ),
                 building,
             ),
         )
@@ -117,10 +112,23 @@ def order_daily_route(batteries, locations, travel_costs, base_building):
 
 
 class MilpPlanner(Planner):
-    def __init__(self, rul_estimator, solver_time_limit_seconds=20.0, late_risk_multiplier=1.0):
+    def __init__(
+        self,
+        rul_estimator,
+        solver_time_limit_seconds=30.0,
+        late_risk_multiplier=1.0,
+        candidate_benefit_threshold=30.0,
+        solver_workers=8,
+    ):
         self.rul_estimator = rul_estimator
         self.solver_time_limit_seconds = float(solver_time_limit_seconds)
         self.late_risk_multiplier = float(late_risk_multiplier)
+        # Only batteries where swapping beats skipping by more than this margin
+        # get decision variables. With ~450 batteries but only ~2-4% actually
+        # due in the window, this shrinks the model ~20x and lets CP-SAT reach
+        # optimality instead of timing out at a 99.9% gap.
+        self.candidate_benefit_threshold = float(candidate_benefit_threshold)
+        self.solver_workers = int(solver_workers)
 
     def _expected_costs(self, timeseries, batteries, settings):
         horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
@@ -131,15 +139,23 @@ class MilpPlanner(Planner):
         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))
+            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,
         )
         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
+        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):
@@ -166,13 +182,22 @@ class MilpPlanner(Planner):
         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
+            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]
+            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}
+        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
@@ -184,39 +209,61 @@ class MilpPlanner(Planner):
                 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))
+                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)
+                    building_visit[building, day]
+                    <= sum(assignment[battery, day] for battery in members)
                 )
 
-        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)
+        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)
+            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())
+            len(batteries) * battery_minutes
+            + len(rooms) * room_minutes
+            + sum(building_work.values())
         )
         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)
+        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)
+                + 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
+                )
             )
             model.add(work == expression)
             daily_work[day] = work
@@ -230,20 +277,25 @@ class MilpPlanner(Planner):
             daily_limit_hit[day] = hit
 
         weekly_limit_hit = {}
-        weekly_limit = round(float(_setting(settings, "worker_limit_weekly_hours", 24.0)) * 60)
+        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
+                sum(daily_work[day] for day in week_days)
+                <= max(weekly_limit - 1, -1) + weekly_maximum * hit
             )
             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 = int(
+                    round(float(expected_costs.loc[battery, action]) * COST_SCALE)
+                )
                 coefficient += action_index + battery_index % 3
                 objective_terms.append(coefficient * assignment[battery, action])
 
@@ -251,25 +303,60 @@ class MilpPlanner(Planner):
         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())
+        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.num_search_workers = max(1, int(getattr(self, "solver_workers", 8)))
         solver.parameters.random_seed = 0
         status = solver.solve(model)
+        # Never silently accept a non-solve: an UNKNOWN status here previously
+        # fell through to the greedy fallback and cost us most of the late_swap.
+        log.info(
+            "cpsat-solve",
+            status=solver.status_name(status),
+            batteries=len(batteries),
+            wall_seconds=round(solver.wall_time, 2),
+        )
         if status not in (cp_model.OPTIMAL, cp_model.FEASIBLE):
+            log.warning("cpsat-no-solution-using-greedy-fallback", status=solver.status_name(status))
             return {
-                battery: min(actions, key=lambda action: (expected_costs.loc[battery, action], str(action)))
+                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]))
+            battery: next(
+                action
+                for action in actions
+                if solver.value(assignment[battery, action])
+            )
             for battery in batteries
         }
 
@@ -285,33 +372,251 @@ class MilpPlanner(Planner):
             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)
+
+        # Restrict the solve to batteries where swapping actually beats skipping.
+        # Only ~2-4% of batteries are due in any window; giving all ~450 of them
+        # decision variables made the model too big to solve (CP-SAT returned
+        # UNKNOWN at a 99.9% gap). Everything below the threshold is pinned to
+        # no_swap, which its own expected costs already say is optimal.
+        day_costs = expected_costs.drop(columns=["no_swap"])
+        benefit = expected_costs["no_swap"] - day_costs.min(axis=1)
+        # getattr keeps planners pickled before this attribute existed loadable
+        threshold = float(getattr(self, "candidate_benefit_threshold", 30.0))
+        candidates = sorted(benefit[benefit > threshold].index.astype(str))
+        log.info("candidate-filter", total=len(batteries), candidates=len(candidates))
+
+        assignments = {battery: "no_swap" for battery in batteries}
+        if candidates:
+            solved = self._solve_assignments(
+                expected_costs.loc[candidates], loc, travel_costs, settings
+            )
+            assignments.update(solved)
+
         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})
+                records.append(
+                    {"day": start_time + pandas.Timedelta(days=day), "battery": battery}
+                )
 
         no_swap_day = start_time + pandas.Timedelta(days=horizon + 1)
-        for battery in sorted(battery for battery in batteries if assignments[battery] == "no_swap"):
+        for battery in sorted(
+            battery for battery in batteries if assignments[battery] == "no_swap"
+        ):
             records.append({"day": no_swap_day, "battery": battery})
 
-        plan = pandas.DataFrame.from_records(records, columns=["day", "battery"]).reset_index(drop=True)
+        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
 
 
+class SearchPlanner(Planner):
+    """Greedy construction + local search scored by the *real* evaluate_plan.
+
+    The MILP approximates the official cost model; this scores candidate plans
+    with the actual evaluator (using predicted EOL as surrogate truth), so it
+    has zero modeling error -- it sees the emergency-visit mechanics, weekly
+    limit accounting and end-of-day travel exactly as the scorer does.
+
+    Only batteries that plausibly fail inside the window get scheduled; with
+    ~2-4% due per scenario the search space is small enough to explore well.
+    """
+
+    def __init__(
+        self,
+        rul_estimator,
+        candidate_benefit_threshold=5.0,
+        iterations=400,
+        random_seed=0,
+        late_risk_multiplier=1.0,
+    ):
+        self.rul_estimator = rul_estimator
+        self.candidate_benefit_threshold = float(candidate_benefit_threshold)
+        self.iterations = int(iterations)
+        self.random_seed = int(random_seed)
+        self.late_risk_multiplier = float(late_risk_multiplier)
+
+    def _expected_costs(self, timeseries, batteries, settings):
+        return MilpPlanner._expected_costs(self, timeseries, batteries, settings)
+
+    @staticmethod
+    def _build_plan(assignment, all_batteries, start_time, park_day):
+        """assignment: battery -> day offset (int). Others parked past window."""
+        records = [
+            {"day": start_time + pandas.Timedelta(days=int(d)), "battery": b}
+            for b, d in assignment.items()
+        ]
+        assigned = set(assignment)
+        records.extend(
+            {"day": park_day, "battery": b} for b in all_batteries if b not in assigned
+        )
+        plan = pandas.DataFrame.from_records(records, columns=["day", "battery"])
+        # stable ordering: by day, then grouped by building/room handled by caller
+        plan = plan.sort_values(["day", "battery"], kind="stable").reset_index(drop=True)
+        plan["day"] = pandas.to_datetime(plan["day"])
+        return plan
+
+    @staticmethod
+    def _route_day(day_batteries, building_of, room_of, travel, base):
+        """Nearest-building route for one day, using precomputed lookups.
+
+        Same logic as order_daily_route but without rebuilding the location
+        frame and travel dict on every call -- this runs inside the search loop.
+        """
+        remaining = set(building_of[b] for b in day_batteries)
+        current = str(base)
+        building_order = []
+        while remaining:
+            nxt = min(
+                remaining,
+                key=lambda bl: (
+                    travel.get((current, bl), 0.0 if current == bl else float("inf")),
+                    bl,
+                ),
+            )
+            building_order.append(nxt)
+            remaining.discard(nxt)
+            current = nxt
+        ordered = []
+        for building in building_order:
+            here = [b for b in day_batteries if building_of[b] == building]
+            here.sort(key=lambda b: (room_of[b], b))
+            ordered.extend(here)
+        return ordered
+
+    def plan(self, timeseries, locations, travel_costs, settings):
+        import random
+
+        loc = _normalize_locations(locations)
+        batteries = sorted(loc["battery"].astype(str).tolist())
+        normalized = normalize_timeseries(timeseries)
+        if normalized.empty:
+            start_time = pandas.to_datetime(loc["end_time"]).max().normalize()
+        else:
+            start_time = normalized["end_time"].max().normalize()
+
+        horizon = int(round(float(_setting(settings, "planning_window_days", 42))))
+        horizon_end = start_time + pandas.Timedelta(days=horizon)
+        base = str(_setting(settings, "base_location", ""))
+        park_day = start_time + pandas.Timedelta(days=horizon + 1)
+
+        expected_costs = self._expected_costs(timeseries, batteries, settings)
+        day_costs = expected_costs.drop(columns=["no_swap"])
+        benefit = expected_costs["no_swap"] - day_costs.min(axis=1)
+        candidates = sorted(
+            benefit[benefit > self.candidate_benefit_threshold].index.astype(str)
+        )
+        log.info("candidate-filter", total=len(batteries), candidates=len(candidates))
+        if not candidates:
+            plan = self._build_plan({}, batteries, start_time, park_day)
+            check_plan_valid(plan, loc, start_time=start_time)
+            return plan
+
+        # Surrogate EOL for scoring: the cost-minimising swap day per battery
+        # already encodes the newsvendor trade-off (early 0.5/day vs late 10/day).
+        target_day = {b: int(day_costs.loc[b].idxmin()) for b in candidates}
+        surrogate_eol = pandas.Series(
+            {
+                b: (start_time + pandas.Timedelta(days=int(target_day[b])))
+                if b in target_day
+                else park_day + pandas.Timedelta(days=365)
+                for b in batteries
+            }
+        )
+
+        # Precomputed lookups so the search loop never rebuilds these.
+        indexed = loc.set_index("battery")
+        building_of = indexed["building"].astype(str).to_dict()
+        room_of = indexed["room"].astype(str).to_dict()
+        travel = _travel_lookup(travel_costs)
+        parked = [b for b in batteries]
+
+        def make_plan(assignment):
+            by_day = {}
+            for b, d in assignment.items():
+                by_day.setdefault(int(d), []).append(b)
+            rows_day, rows_bat = [], []
+            for d in sorted(by_day):
+                ordered = self._route_day(
+                    sorted(by_day[d]), building_of, room_of, travel, base
+                )
+                stamp = start_time + pandas.Timedelta(days=d)
+                rows_day.extend([stamp] * len(ordered))
+                rows_bat.extend(ordered)
+            assigned = set(assignment)
+            rest = [b for b in parked if b not in assigned]
+            rows_day.extend([park_day] * len(rest))
+            rows_bat.extend(rest)
+            plan = pandas.DataFrame({"day": rows_day, "battery": rows_bat})
+            plan["day"] = pandas.to_datetime(plan["day"])
+            return plan
+
+        def score(assignment):
+            try:
+                _, _, overall = evaluate_plan(
+                    make_plan(assignment),
+                    loc,
+                    travel_costs,
+                    settings,
+                    eol_times=surrogate_eol,
+                    start_time=start_time,
+                    verbose=0,
+                )
+                return float(overall["total_cost"])
+            except Exception:
+                return float("inf")
+
+        # Greedy start: everyone at their own cost-minimising day.
+        current = dict(target_day)
+        current_cost = score(current)
+
+        # Local search. Moves: shift a day, batch onto another candidate's day,
+        # drop a battery, or restore a dropped one.
+        rng = random.Random(self.random_seed)
+        best, best_cost = dict(current), current_cost
+        used_days = sorted(set(target_day.values()))
+        for _ in range(self.iterations):
+            trial = dict(current)
+            battery = rng.choice(candidates)
+            move = rng.random()
+            if move < 0.4 and trial:
+                # batch: move onto a day already being worked (saves a trip)
+                if used_days:
+                    trial[battery] = rng.choice(
+                        sorted(set(trial.values())) or used_days
+                    )
+            elif move < 0.75:
+                # shift, biased earlier (late costs 20x more than early)
+                base_day = trial.get(battery, target_day[battery])
+                shift = -rng.randint(1, 7) if rng.random() < 0.7 else rng.randint(1, 4)
+                trial[battery] = int(min(max(base_day + shift, 0), horizon))
+            elif move < 0.9:
+                trial.pop(battery, None)  # drop
+            else:
+                trial[battery] = target_day[battery]  # restore
+
+            trial_cost = score(trial)
+            if trial_cost <= current_cost:
+                current, current_cost = trial, trial_cost
+                if trial_cost < best_cost:
+                    best, best_cost = dict(trial), trial_cost
+
+        log.info(
+            "search-planner",
+            candidates=len(candidates),
+            scheduled=len(best),
+            surrogate_cost=round(best_cost, 1),
+        )
+        plan = make_plan(best)
+        check_plan_valid(plan, loc, start_time=start_time)
+        return plan
 
-# --------------------------------------------------------------------------
-# RUL model: hand-crafted battery snapshot features feeding a discrete-time
-# hazard model (see DiscreteHazardRULModel below). Gives a failure-probability
-# distribution (not a single point estimate), so MilpPlanner's expected-cost
-# objective can properly weigh early vs. late risk instead of committing to
-# one predicted day.
-# --------------------------------------------------------------------------
 
 RUL_DEVICE_COLUMN = "device_id"
 RUL_TIME_COLUMN = "end_time"
@@ -345,7 +650,11 @@ def extract_snapshot_features(timeseries, reference_time=None, windows=RUL_WINDO
     if frame.empty:
         return pd.DataFrame(index=pd.Index([], name=RUL_DEVICE_COLUMN))
 
-    reference = pd.Timestamp(reference_time) if reference_time is not None else frame[RUL_TIME_COLUMN].max()
+    reference = (
+        pd.Timestamp(reference_time)
+        if reference_time is not None
+        else frame[RUL_TIME_COLUMN].max()
+    )
     frame = frame.loc[frame[RUL_TIME_COLUMN] <= reference].copy()
     rows = []
 
@@ -355,8 +664,12 @@ def extract_snapshot_features(timeseries, reference_time=None, windows=RUL_WINDO
         last_time = group[RUL_TIME_COLUMN].iloc[-1]
         row = {
             RUL_DEVICE_COLUMN: device_id,
-            "device_age_days": max((reference - first_time).total_seconds() / 86400.0, 0.0),
-            "history_span_days": max((last_time - first_time).total_seconds() / 86400.0, 0.0),
+            "device_age_days": max(
+                (reference - first_time).total_seconds() / 86400.0, 0.0
+            ),
+            "history_span_days": max(
+                (last_time - first_time).total_seconds() / 86400.0, 0.0
+            ),
             "days_since_last_observation": max(
                 (reference - last_time).total_seconds() / 86400.0, 0.0
             ),
@@ -370,10 +683,14 @@ def extract_snapshot_features(timeseries, reference_time=None, windows=RUL_WINDO
             row[f"{value_column}_std"] = _rul_finite(values.std(ddof=0))
             row[f"{value_column}_min"] = _rul_finite(values.min())
             row[f"{value_column}_max"] = _rul_finite(values.max())
-            row[f"{value_column}_slope"] = _rul_finite(_rul_slope(values, group[RUL_TIME_COLUMN]))
+            row[f"{value_column}_slope"] = _rul_finite(
+                _rul_slope(values, group[RUL_TIME_COLUMN])
+            )
 
         for days in windows:
-            window = group.loc[group[RUL_TIME_COLUMN] >= reference - pd.Timedelta(days=int(days))]
+            window = group.loc[
+                group[RUL_TIME_COLUMN] >= reference - pd.Timedelta(days=int(days))
+            ]
             row[f"observation_count_{days}d"] = float(len(window))
             for value_column in RUL_VALUE_COLUMNS:
                 values = window[value_column]
@@ -382,7 +699,9 @@ def extract_snapshot_features(timeseries, reference_time=None, windows=RUL_WINDO
                 row[f"{prefix}_std"] = _rul_finite(values.std(ddof=0))
                 row[f"{prefix}_min"] = _rul_finite(values.min())
                 row[f"{prefix}_max"] = _rul_finite(values.max())
-                row[f"{value_column}_slope_{days}d"] = _rul_finite(_rul_slope(values, window[RUL_TIME_COLUMN]))
+                row[f"{value_column}_slope_{days}d"] = _rul_finite(
+                    _rul_slope(values, window[RUL_TIME_COLUMN])
+                )
 
         rows.append(row)
 
@@ -390,7 +709,9 @@ def extract_snapshot_features(timeseries, reference_time=None, windows=RUL_WINDO
     return features.astype(float)
 
 
-def expected_costs_from_failure_distribution(failure_probability, replacement_days, early_penalty, late_penalty):
+def expected_costs_from_failure_distribution(
+    failure_probability, replacement_days, early_penalty, late_penalty
+):
     probabilities = np.asarray(failure_probability, dtype=float)
     probabilities = np.clip(probabilities, 0.0, None)
     total_probability = probabilities.sum()
@@ -406,7 +727,9 @@ def expected_costs_from_failure_distribution(failure_probability, replacement_da
     return costs @ probabilities
 
 
-def expected_no_swap_cost(failure_probability, horizon_day, emergency_day, late_penalty):
+def expected_no_swap_cost(
+    failure_probability, horizon_day, emergency_day, late_penalty
+):
     probabilities = np.asarray(failure_probability, dtype=float)
     probabilities = np.clip(probabilities, 0.0, None)
     total_probability = probabilities.sum()
@@ -416,17 +739,13 @@ def expected_no_swap_cost(failure_probability, horizon_day, emergency_day, late_
     failure_days = np.arange(len(probabilities), dtype=float)
     due_inside_window = failure_days <= int(horizon_day)
     late_days = np.maximum(float(emergency_day) - failure_days, 0.0)
-    return float(np.sum(probabilities[due_inside_window] * late_days[due_inside_window]) * late_penalty)
+    return float(
+        np.sum(probabilities[due_inside_window] * late_days[due_inside_window])
+        * late_penalty
+    )
 
 
 class DiscreteHazardRULModel(RULModel):
-    """Fully non-parametric discrete-time hazard model: bins time into
-    weekly periods and trains one sklearn classifier to predict
-    P(fail in period | survived to it, features, elapsed periods). No
-    Weibull-shape or proportional-hazards assumption -- can capture a sharp
-    voltage-threshold-crossing failure pattern that parametric models can't.
-    """
-
     quantile_cols = ["p10", "p50", "p90"]
     period_days = 7
     horizon_cap_days = 126  # ~18 weekly periods; covers the 42-day planning window plus emergency margin
@@ -440,7 +759,9 @@ class DiscreteHazardRULModel(RULModel):
         self.fallback_hazard_ = 0.01
 
     def _prepare_features(self, features, fitting=False):
-        numeric = features.apply(pd.to_numeric, errors="coerce").replace([np.inf, -np.inf], np.nan)
+        numeric = features.apply(pd.to_numeric, errors="coerce").replace(
+            [np.inf, -np.inf], np.nan
+        )
         if fitting:
             self.feature_columns_ = list(numeric.columns)
             self.feature_medians_ = numeric.median().fillna(0.0)
@@ -459,7 +780,9 @@ class DiscreteHazardRULModel(RULModel):
             capped_duration = min(duration, float(self.horizon_cap_days))
             failure_period = None
             if event and duration <= self.horizon_cap_days:
-                failure_period = min(int(capped_duration // self.period_days), n_periods - 1)
+                failure_period = min(
+                    int(capped_duration // self.period_days), n_periods - 1
+                )
             max_period = int(np.ceil(capped_duration / self.period_days))
             if failure_period is not None:
                 max_period = max(max_period, failure_period + 1)
@@ -476,15 +799,21 @@ class DiscreteHazardRULModel(RULModel):
         return period_features, np.array(labels, dtype=int)
 
     def fit_snapshots(self, snapshot_features, durations, events):
-        common = snapshot_features.index.intersection(durations.index).intersection(events.index)
+        common = snapshot_features.index.intersection(durations.index).intersection(
+            events.index
+        )
         if common.empty:
             raise ValueError("No aligned snapshot labels were provided")
         features = self._prepare_features(snapshot_features.loc[common], fitting=True)
-        duration = pd.to_numeric(durations.loc[common], errors="coerce").clip(lower=0.25)
+        duration = pd.to_numeric(durations.loc[common], errors="coerce").clip(
+            lower=0.25
+        )
         event = events.loc[common].fillna(False).astype(bool)
 
         period_features, labels = self._expand_person_periods(features, duration, event)
-        self.fallback_hazard_ = float(np.clip(labels.mean(), 1e-3, 0.5)) if len(labels) else 0.01
+        self.fallback_hazard_ = (
+            float(np.clip(labels.mean(), 1e-3, 0.5)) if len(labels) else 0.01
+        )
 
         self.model_ = None
         if labels.sum() >= 2 and len(labels) >= 10:
@@ -521,7 +850,9 @@ class DiscreteHazardRULModel(RULModel):
         hazards = self._period_hazards(features)
         n = hazards.shape[0]
         period_survival = np.cumprod(1.0 - hazards, axis=1)
-        period_survival = np.hstack([np.ones((n, 1)), period_survival])  # prepend day-0 survival = 1
+        period_survival = np.hstack(
+            [np.ones((n, 1)), period_survival]
+        )  # prepend day-0 survival = 1
 
         times = np.asarray(times, dtype=float)
         result = np.ones((len(times), n), dtype=float)
@@ -532,14 +863,20 @@ class DiscreteHazardRULModel(RULModel):
 
     def predict(self, timeseries):
         features = extract_snapshot_features(timeseries)
-        times = np.arange(0, self.horizon_cap_days + self.period_days, self.period_days, dtype=float)
+        times = np.arange(
+            0, self.horizon_cap_days + self.period_days, self.period_days, dtype=float
+        )
         survival = self._survival(features, times)
         quantile_days = {}
         for q_col, target in zip(self.quantile_cols, (0.9, 0.5, 0.1)):
             days = []
             for j in range(survival.shape[1]):
                 below = np.where(survival[:, j] <= target)[0]
-                days.append(float(times[below[0]]) if len(below) else float(self.horizon_cap_days))
+                days.append(
+                    float(times[below[0]])
+                    if len(below)
+                    else float(self.horizon_cap_days)
+                )
             quantile_days[q_col] = days
         return pd.DataFrame(quantile_days, index=features.index)[self.quantile_cols]
 
@@ -551,11 +888,23 @@ class DiscreteHazardRULModel(RULModel):
         probabilities = np.vstack([interval_mass, survival[-1:]]).T
         row_sums = probabilities.sum(axis=1, keepdims=True)
         probabilities = np.divide(
-            probabilities, row_sums, out=np.zeros_like(probabilities), where=row_sums > 0
+            probabilities,
+            row_sums,
+            out=np.zeros_like(probabilities),
+            where=row_sums > 0,
+        )
+        return pd.DataFrame(
+            probabilities, index=features.index, columns=range(max_day + 2)
         )
-        return pd.DataFrame(probabilities, index=features.index, columns=range(max_day + 2))
 
-    def expected_replacement_costs(self, timeseries, horizon_days, early_penalty, late_penalty, no_swap_extension_days):
+    def expected_replacement_costs(
+        self,
+        timeseries,
+        horizon_days,
+        early_penalty,
+        late_penalty,
+        no_swap_extension_days,
+    ):
         emergency_day = int(horizon_days + no_swap_extension_days)
         max_failure_day = emergency_day + max(int(horizon_days), 30)
         probabilities = self.failure_probabilities(timeseries, max_day=max_failure_day)
@@ -563,10 +912,16 @@ class DiscreteHazardRULModel(RULModel):
         rows = [
             np.append(
                 expected_costs_from_failure_distribution(
-                    row, replacement_days, early_penalty=float(early_penalty), late_penalty=float(late_penalty)
+                    row,
+                    replacement_days,
+                    early_penalty=float(early_penalty),
+                    late_penalty=float(late_penalty),
                 ),
                 expected_no_swap_cost(
-                    row, horizon_day=int(horizon_days), emergency_day=emergency_day, late_penalty=float(late_penalty)
+                    row,
+                    horizon_day=int(horizon_days),
+                    emergency_day=emergency_day,
+                    late_penalty=float(late_penalty),
                 ),
             )
             for row in probabilities.to_numpy(dtype=float)
@@ -587,7 +942,9 @@ def split_scenarios(scenarios, val_fraction=0.25, seed=0):
     return shuffled[n_val:], shuffled[:n_val]
 
 
-def build_training_snapshots(locations, timeseries, eol_times, scenarios, limit_scenarios=None):
+def build_training_snapshots(
+    locations, timeseries, eol_times, scenarios, limit_scenarios=None
+):
     feature_parts = []
     duration_parts = []
     event_parts = []
@@ -647,59 +1004,54 @@ class Config(BaseSettings):
     """
     Automatically provides command-line argument support for specified fields
     """
+
     model_config = SettingsConfigDict(
         env_prefix="",
         cli_parse_args=True,
         cli_ignore_unknown_args=True,
     )
- 
+
     dataset_path: Optional[Path] = None
-    split : str = 'train'
+    split: str = "train"
     solver_time_limit_seconds: float = 20.0
-    # NOTE: fixed at 1.0 (neutral). early/late penalties are identical across
-    # every scenario in the dataset (0.5 / 10.0) -- there is no real per-scenario
-    # risk knob to tune here, so this multiplier must not be used to game the
-    # planner's internal cost objective away from the real evaluator.
     late_risk_multiplier: float = 1.0
     val_fraction: float = 0.25
     split_seed: int = 0
 
+
 def main():
     cfg = Config()
 
     if cfg.dataset_path is None:
-        dataset_path = os.environ.get('BATTERYSWAP_DATASET_PATH', None)
+        dataset_path = os.environ.get("BATTERYSWAP_DATASET_PATH", None)
         assert dataset_path
         dataset_path = Path(dataset_path)
     else:
         dataset_path = cfg.dataset_path
 
     split_path = dataset_path / cfg.split
-    locations, timeseries, eol_times, scenarios  = load_dataset(split_path)
+    locations, timeseries, eol_times, scenarios = load_dataset(split_path)
 
-    log.info('evaluate-load-data', path=dataset_path)
+    log.info("evaluate-load-data", path=dataset_path)
 
-    # Honest train/val split across scenarios -- never evaluate on the
-    # scenarios the RUL model was fit on.
     train_scenarios, val_scenarios = split_scenarios(
         scenarios, val_fraction=cfg.val_fraction, seed=cfg.split_seed
     )
-    log.info('scenario-split', total=len(scenarios), train=len(train_scenarios), val=len(val_scenarios))
-
-    # Prediction model training (train fold only -- for the held-out metric below)
-    # DiscreteHazardRULModel (non-parametric discrete-time hazard model) beat both
-    # a Weibull AFT and a Cox proportional-hazards model by a wide margin on
-    # held-out total_cost, confirmed across two different train/val splits.
+    log.info(
+        "scenario-split",
+        total=len(scenarios),
+        train=len(train_scenarios),
+        val=len(val_scenarios),
+    )
     rul_model = train_rul_model(locations, timeseries, eol_times, train_scenarios)
-    log.info('train-done')
+    log.info("train-done")
 
-    log.info('evaluate-held-out')
-    # Evaluate ONLY on the held-out validation scenarios
+    log.info("evaluate-held-out")
     gen = iterate_scenarios(locations, timeseries, eol_times, val_scenarios)
     for scenario, locs, cut, eol in gen:
-        scenario_name = scenario['name']
-        travel_costs = scenario['travel_costs']
-        settings = scenario['settings']
+        scenario_name = scenario["name"]
+        travel_costs = scenario["travel_costs"]
+        settings = scenario["settings"]
 
         planner = MilpPlanner(
             rul_model,
@@ -708,15 +1060,15 @@ def main():
         )
         plan = planner.plan(cut, locs, travel_costs, settings)
 
-        start_time = pandas.Timestamp(scenario['start_time'])
+        start_time = pandas.Timestamp(scenario["start_time"])
 
-        transitions, daily, overall = evaluate_plan(plan, locs, travel_costs, settings, eol_times=eol, start_time=start_time)
+        transitions, daily, overall = evaluate_plan(
+            plan, locs, travel_costs, settings, eol_times=eol, start_time=start_time
+        )
 
-        print('scores', scenario_name, overall)
+        print("scores", scenario_name, overall)
 
-    log.info('refit-on-full-data-for-submission')
-    # For the actual submitted planner, refit on ALL scenarios (train + held-out)
-    # now that the held-out metric above has told us how good the model really is.
+    log.info("refit-on-full-data-for-submission")
     rul_model_full = train_rul_model(locations, timeseries, eol_times, scenarios)
 
     # Save best planner