train: Fix NaN for unobserved, respect envvar for dataset

batteryswap_example/train.py

@@ -123,28 +123,35 @@ class DummyRULModel(RULModel):
         return out[self.quantile_cols]
 
 
-def train_rul_model(locations, timeseries, eol_times, scenarios):
+def train_rul_model(locations, timeseries, eol_times, scenarios, limit_scenarios=None):
 
     # 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 = []
+    scenarios_loaded = 0
     for scenario, locs, cut, eol in gen:
-
+        print('load-scenario')
         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'])
+        unobserved_rul = 120
 
         # 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'] 
+        rul_days.index = pandas.Series(rul_days.index) + scenario['name']
+        rul_days = rul_days.fillna(unobserved_rul)
         cut_ruls.append(rul_days)
 
         assert set(rul_days.index) == set(cut.index.get_level_values('device_id'))
 
+        if limit_scenarios is not None:
+            if scenarios_loaded > limit_scenarios:
+                break
+        scenarios_loaded += 1
 
         
     # Train a RUL prediction model
@@ -172,13 +179,18 @@ class Config(BaseSettings):
         cli_ignore_unknown_args=True,
     )
  
-    dataset_path: Path = Path("dataset")
+    dataset_path: Optional[Path] = None
     split : str = 'train'
 
 def main():
     cfg = Config()
 
-    dataset_path = pathlib.Path(cfg.dataset_path)
+    if cfg.dataset_path is 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)
@@ -187,9 +199,11 @@ def main():
 
 
     # Prediction model training
-    rul_model = train_rul_model(locations, timeseries, eol_times, scenarios)
-
+    #rul_model = DummyRULModel()
+    rul_model = train_rul_model(locations, timeseries, eol_times, scenarios, limit_scenarios=1)
+    log.info('train-done')
 
+    log.info('evaluate')
     # Evaluate on the planning scenarios
     gen = iterate_scenarios(locations, timeseries, eol_times, scenarios)
     for scenario, locs, cut, eol in gen: