Fixed broken evaluate (ecmwf#131)

* Fixed broken evaluate:
- evaluate() used all channels infos for model
- file name for latest checkpoint was incorrect (introduced in pathlib changes)

* Fix path handling to make it backward compatible. (ecmwf#133)

* Fixed broken evaluate:
- evaluate() used all channels infos for model
- file name for latest checkpoint was incorrect (introduced in pathlib changes)

* Fixed issues with evaluate.

* Ruffed.

Author: clessig

src/weathergen/__init__.py

@@ -25,19 +25,6 @@ def evaluate():
     """
     Evaluation function for WeatherGenerator model.
     Entry point for calling the evaluation code from the command line.
-
-    Args:
-        run_id (str): Run/model id of pretrained WeatherGenerator model.
-        start_date (str): Start date for evaluation. Format must be parsable with pd.to_datetime.
-        end_date (str): End date for evaluation. Format must be parsable with pd.to_datetime.
-        epoch (int, optional): Epoch of pretrained WeatherGenerator model used for evaluation (-1 corresponds to last epoch). Defaults to -1.
-        masking_mode (str, optional): Masking mode for evaluation. Defaults to None.
-        forecast_steps (int, optional): Number of forecast steps for evaluation. Defaults to None.
-        samples (int, optional): Number of samples for evaluation. Defaults to 10000000.
-        shuffle (bool, optional): Shuffle samples for evaluation. Defaults to False.
-        save_samples (bool, optional): Save samples for evaluation. Defaults to True.
-        analysis_streams_output (list, optional): Analysis output streams during evaluation. Defaults to ['ERA5'].
-        gridded_output_streams(list, optional): Currently unused and threrefore omitted here
     """
     parser = argparse.ArgumentParser()
 
@@ -52,13 +39,15 @@ def evaluate():
         "-start",
         type=str,
         required=False,
+        default="2022-10-01",
         help="Start date for evaluation. Format must be parsable with pd.to_datetime.",
     )
     parser.add_argument(
         "--end_date",
         "-end",
         type=str,
         required=False,
+        default="2022-12-01",
         help="End date for evaluation. Format must be parsable with pd.to_datetime.",
     )
     parser.add_argument(
@@ -104,7 +93,8 @@ def evaluate():
     init_loggers()
 
     # load config: if run_id is full path, it loads from there
-    cf = Config.load(args.run_id, args.epoch, private_cf["model_path"])
+    model_path = private_cf["model_path"] if hasattr(private_cf, "model_path") else "./models"
+    cf = Config.load(args.run_id, args.epoch, model_path)
 
     # add parameters from private (paths) config
     for k, v in private_cf.items():
@@ -117,16 +107,8 @@ def evaluate():
 
     start_date, end_date = pd.to_datetime(args.start_date), pd.to_datetime(args.end_date)
 
-    cf.start_date_val = start_date.strftime(
-        "%Y%m%d%H%M"
-    )  # ML: would be better to use datetime-objects
+    cf.start_date_val = start_date.strftime("%Y%m%d%H%M")
     cf.end_date_val = end_date.strftime("%Y%m%d%H%M")
-    # # Oct-Nov 2022
-    # cf.start_date_val = 202210011600
-    # cf.end_date_val = 202212010400
-    # # 2022
-    # cf.start_date_val = 202201010400
-    # cf.end_date_val = 202301010400
 
     cf.shuffle = args.shuffle
 

src/weathergen/train/trainer.py

@@ -120,10 +120,12 @@ def evaluate(self, cf, run_id_trained, epoch, run_id_new=False):
             self.dataset_val, **loader_params, sampler=None
         )
 
-        num_channels = self.dataset_val.get_num_chs()
-        self.geoinfo_sizes = self.dataset_val.get_geoinfo_sizes()
+        sources_size = self.dataset_val.get_sources_size()
+        targets_num_channels = self.dataset_val.get_targets_num_channels()
+        targets_coords_size = self.dataset_val.get_targets_coords_size()
 
-        self.model = Model(cf, num_channels, self.geoinfo_sizes).create().to(self.devices[0])
+        self.model = Model(cf, sources_size, targets_num_channels, targets_coords_size).create()
+        self.model = self.model.to(self.devices[0])
         self.model.load(run_id_trained, epoch)
         print(f"Loaded model {run_id_trained} at epoch {epoch}.")
         self.ddp_model = self.model