Update for datetime plans

.gitignore

@@ -1,3 +1,5 @@
 submission.csv
 data/
 dataset/
+
+venv/

README.md

@@ -2,9 +2,25 @@
 license: mit
 ---
 
+# Example Model for BatterySwapAI 2026
 
+This repository shows how to create a solution for the
+[BatterySwapAI 2026 challenge](https://www.nora.ai/competitions/batteryswapai/batteryswapai2026.html).
 
-# Setting up virtual environment
+It contains working example code that can be submitted as-is,
+along with tools to check the solution before submitting.
+
+# Setup
+
+## Prerequisites
+
+You need to have the following installed
+
+- git
+- Python 3.10+
+- Docker
+
+## Setting up virtual environment
 
 On Linux / Mac OS / Windows Subsystem for Linux (WSL)
 ```
@@ -17,7 +33,10 @@ Install dependencies
 pip install -r requirements.txt -r requirements.dev.txt
 ```
 
-# Develop in virtual environment
+
+# Developing
+
+## Develop in virtual environment
 
 
 Run the training
@@ -25,14 +44,16 @@ Run the training
 python batteryswap_example/train.py
 ```
 
-# Submitting trained models
+# Making submissions
+
+## Submitting trained models
 
 Trained models or other output used by submission processing (`script.py`),
 must be committed in the git repository.
 An example is `batteryswap_example/planners/best.pickle`.
 
 
-# Test submissions in Docker - recommended
+## Test submissions in Docker - recommended
  
 Using Docker allows to have exactly the same software versions as the submissions system.
 
@@ -55,4 +76,10 @@ Copy submission.csv out of container
 docker cp batteryswapai:/app/submission.csv ./submission.csv
 ```
 
+## Create new submission
+
+NOTE: remember to commit and push your changes to the HuggingFace model repository.
+
+Use `New submission` in the competition application to submit your current code for evaluation.
+
 

batteryswap_example/train.py

@@ -2,34 +2,221 @@
 import pandas
 import numpy
 import pickle
+from typing import Optional, Sequence
+from pathlib import Path
+import pathlib
 
-class OrderedPlanner():
-    def __init__(self):
-        pass
+
+from sklearn.dummy import DummyRegressor
+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 batteryswap_public.interfaces import Planner, RULModel
+from batteryswap_public.utils import load_dataset, iterate_scenarios
+from batteryswap_public.evaluate import evaluate_plan, check_plan_valid
+
+log = structlog.get_logger()
+
+
+class OrderedPlanner(Planner):
+    def __init__(self, rul_estimator):
+        self.rul_estimator = rul_estimator
 
     def plan(self, battery_data, locations, travel_costs, settings):
 
-        # Order by voltage
-        latest_voltage = battery_data.groupby(['device_id'])['voltage'].last()
-        locations = locations.copy().set_index('battery')
-        locations['latest_voltage'] = latest_voltage
-        order = locations.sort_values('latest_voltage')
+        # Remaining Useful Life estimation
+        # FIXME: consider balance of over/under-estimation
+        percentile = 'p50'
+        rul = self.rul_estimator.predict(battery_data)
+        rul_days = rul[percentile]
+        # convert to a date
+        start_time = battery_data.reset_index()['end_time'].max()
+        predict_eol = start_time + pandas.to_timedelta(rul_days, unit='D')
 
-        # Do one swap per day
-        days = numpy.arange(len(order))
+        loc = locations.copy().set_index('battery')
+        loc['eol_time'] = predict_eol
+        order = loc.sort_values('eol_time', ascending=True)
+
+        # Planner
+        # Stupid heuristic: Do one swap per day
+        # FIXME: take travel distances into account
+        # FIXME: take co-location into account
+        # FIXME: take daily and weekly limits into account
+        days = start_time + pandas.to_timedelta(numpy.arange(len(order)), unit='D')
         plan = pandas.DataFrame({
             'day': days,
             'battery': order.index,
         })
+
+        check_plan_valid(plan, locations, start_time=start_time)
+
         return plan
 
+
+class DummyRULModel(RULModel):
+    # RUL model that predicts (no information rate)
+    # FIXME: make a model that actually uses the data to improve predictions
+
+    def __init__(
+        self,
+        time_col: str = 'end_time',
+        group_col: str = 'device_id',
+        value_cols: Sequence[str] = ('voltage', 'temperature'),
+        quantiles: Sequence[float] = (0.5, ),
+    ):
+        self.time_col = time_col
+        self.group_col = group_col
+        self.value_cols = list(value_cols)
+        self.quantiles = sorted(quantiles)
+        self.quantile_cols = [f"p{round(q * 100):02d}" for q in self.quantiles]
+
+        self.model = None
+        self.use_total_elapsed_days = True
+
+    def _compute_features(self, unit_df: pd.DataFrame) -> np.ndarray:
+        unit_df = unit_df.sort_values(self.time_col)
+
+        all_stats = {}
+        # FIXME: actually compute some features 
+
+
+        feature_names = sorted(all_stats.keys())
+        self._feature_names_ = feature_names  # stored for inspection/debugging
+        return np.array([all_stats[k] for k in feature_names], dtype=float)
+
+    def _build_feature_matrix(self, timeseries: pd.DataFrame) -> tuple[np.ndarray, list]:
+        rows, ids = [], []
+        for unit_id, unit_df in timeseries.groupby(self.group_col):
+            if len(unit_df) < 2:
+                continue
+            rows.append(self._compute_features(unit_df))
+            ids.append(unit_id)
+        return np.vstack(rows), ids
+
+    def fit(self, timeseries: pd.DataFrame, rul: pd.Series):
+        X, ids = self._build_feature_matrix(timeseries)
+        y = np.array([rul[unit_id] for unit_id in ids])
+
+        # FIXME: actually use an estimator that learns
+        self.model = DummyRegressor(strategy='median')
+        self.model.fit(X, y)
+        return self
+
+    def predict(self, timeseries: pd.DataFrame) -> pd.DataFrame:
+        rows, ids = [], []
+        for unit_id, unit_df in timeseries.groupby(self.group_col):
+            if len(unit_df) < 2:
+                raise ValueError(f"too short unit={unit_id} len={len(unit_df)}")
+            rows.append(self._compute_features(unit_df))
+            ids.append(unit_id)
+
+        X = np.vstack(rows)
+
+        # every quantile column just gets the single point prediction.
+        point_pred = self.model.predict(X)
+        preds = {col: point_pred for col in self.quantile_cols}
+
+        out = pd.DataFrame(preds, index=pd.Index(ids, name=self.group_col))
+        return out[self.quantile_cols]
+
+
+def train_rul_model(locations, timeseries, eol_times, scenarios):
+
+    # Collect training data
+    # FIXME: train/validate/test split to estimate generalized predictive performance
+    gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
+    cut_dfs = []
+    cut_ruls = []
+    for scenario, locs, cut, eol in gen:
+
+        cut = cut.reset_index()
+        cut['device_id'] = cut['device_id'].astype(str) + scenario['name']
+        cut = cut.set_index(['device_id', 'end_time'])
+        cut_dfs.append(cut)
+        plan_start = pandas.Timestamp(scenario['start_time'])
+
+        # Convert EOL datetimes to RUL in days relative to planning time
+        rul_days = (eol - plan_start) / pandas.Timedelta(days=1)
+        rul_days.index = pandas.Series(rul_days.index) + scenario['name'] 
+        cut_ruls.append(rul_days)
+
+        assert set(rul_days.index) == set(cut.index.get_level_values('device_id'))
+
+
+        
+    # Train a RUL prediction model
+    rul_model = DummyRULModel()
+    X = pandas.concat(cut_dfs)
+    Y = pandas.concat(cut_ruls)
+
+    print(X.head())
+    print(Y.head())
+
+    rul_model.fit(X, Y)
+
+    # FIXME: do model selection
+
+    return rul_model
+
+
+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: Path = Path("dataset")
+    split : str = 'train'
+
 def main():
-    planner = OrderedPlanner()
-    planner_path = 'planner.pickle' 
+    cfg = Config()
+
+    dataset_path = pathlib.Path(cfg.dataset_path)
 
+    split_path = dataset_path / cfg.split
+    locations, timeseries, eol_times, scenarios  = load_dataset(split_path)
+
+    log.info('evaluate-load-data', path=dataset_path)
+
+
+    # Prediction model training
+    rul_model = train_rul_model(locations, timeseries, eol_times, scenarios)
+
+
+    # Evaluate on the planning scenarios
+    gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
+    for scenario, locs, cut, eol in gen:
+        scenario_name = scenario['name']
+        travel_costs = scenario['travel_costs']
+        settings = scenario['settings']
+
+        planner = OrderedPlanner(rul_model)
+        plan = planner.plan(cut, locs, travel_costs, settings)
+
+        start_time = pandas.Timestamp(scenario['start_time'])
+
+        transitions, daily, overall = evaluate_plan(plan, locs, travel_costs, settings, eol_times=eol, start_time=start_time)
+
+        print('scores', scenario_name, overall)
+
+
+    # Save best planner
+    planner = OrderedPlanner(rul_model)
+
+    planner_path = 'batteryswap_example/planners/best.pickle' 
     with open(planner_path, "wb") as f:
         pickle.dump(planner, f)
         print('planner-save', planner_path)
 
+
 if __name__ == '__main__':
     main()