Add deterministic metrics to eval script, fix 1 multistep support (no prediction_timedelta dim)

Author: 14renus

geoarches/dataloaders/era5.py

@@ -17,11 +17,14 @@
     last_train=lambda x: ("2018" in x),
     last_train_z0012=lambda x: ("2018" in x and ("0h" in x or "12h" in x)),
     train=lambda x: not ("2019" in x or "2020" in x or "2021" in x),
+    # Before and after 2000. Need to load timestamp after to account for offset..
+    train_before_2000=lambda x: any([str(y) in x for y in range(1979, 2001)]),  # 1979-1999
+    train_after_2000=lambda x: any([str(y) in x for y in range(2000, 2020)]),  # 2000-2018
     # Splits val and test  are from 2019 and 2020 respectively, but
     # we read the years before and after to account for offsets when
     # loading previous and future timestamps for an example.
-    val=lambda x: ("2018" in x or "2019" in x or "2020" in x),
-    test=lambda x: ("2019" in x or "2020" in x or "2021" in x),
+    val=lambda x: ("2018" in x or "2019" in x or "2020" in x),  # 2019
+    test=lambda x: ("2019" in x or "2020" in x or "2021" in x),  # 2020
     test_z0012=lambda x: ("2019" in x or "2020" in x or "2021" in x) and ("0h" in x or "12h" in x),
     test2022_z0012=lambda x: ("2022" in x) and ("0h" in x or "12h" in x),  # check if that works ?
     recent2=lambda x: any([str(y) in x for y in range(2007, 2019)]),
@@ -274,7 +277,6 @@ def __init__(
         )
 
         # depending on domain, re-set timestamp bounds
-
         if domain in ("val", "test", "test_z0012"):
             # re-select timestamps
             year = 2019 if domain.startswith("val") else 2020

geoarches/evaluation/eval_multistep.py

@@ -151,11 +151,17 @@ def main():
         required=True,
         help="Directory or file path to read groundtruth.",
     )
+    parser.add_argument(
+        "--groundtruth_dataset_domain",
+        type=str,
+        default="test_z0012",
+        help="Domain (all, train, val, test) for groundtruth dataset. Should be a key in filename_filters. Determines filename_filter used.",
+    )
     parser.add_argument(
         "--multistep",
         default=10,
         type=int,
-        help="Number of future timesteps model is rolled out for evaluation. In days "
+        help="Number of future timesteps model is rolled out for evaluation. Set to 1 if just one step."
         "(This script assumes lead time is 24 hours).",
     )
     parser.add_argument(
@@ -198,6 +204,16 @@ def main():
         action="store_true",
         help="Whether to evaluate climatology.",
     )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        help="Whether to  print more verbose debug logs.",
+    )
+    parser.add_argument(
+        "--breakpoint",
+        action="store_true",
+        help="Whether to add breakpoint for debug.",
+    )
 
     args = parser.parse_args()
 
@@ -231,7 +247,7 @@ def main():
                 surface_variables=args.surface_vars,
                 level_variables=args.level_vars,
                 pressure_levels=[500, 700, 850],
-                lead_time_hours=24 if args.multistep else None,
+                lead_time_hours=24 if args.multistep and args.multistep > 1 else None,
                 rollout_iterations=args.multistep,
             ).to(device)
     print(f"Computing: {metrics.keys()}")
@@ -240,7 +256,7 @@ def main():
     ds_test = era5.Era5Forecast(
         path=args.groundtruth_path,
         # filename_filter=lambda x: ("2020" in x) and ("0h" in x or "12h" in x),
-        domain="test_z0012",
+        domain=args.groundtruth_dataset_domain,
         lead_time_hours=24,
         multistep=args.multistep,
         load_prev=False,
@@ -251,30 +267,36 @@ def main():
     )
 
     print(f"Reading {len(ds_test.files)} files from groundtruth path: {args.groundtruth_path}.")
+    if args.verbose:
+        print(ds_test.files)
 
     # Predictions.
     def _pred_filename_filter(filename):
         if "metric" in filename:
             return False
         if args.pred_filename_filter is None:
             return True
-        for substring in args.pred_filename_filter:
-            if substring not in filename:
-                return False
-        return True
+        return any([str(y) in filename for y in args.pred_filename_filter])
 
     if not args.eval_clim:
+        dimension_indexers = dict(level=[500, 700, 850])
+        if args.multistep > 1:
+            dimension_indexers["prediction_timedelta"] = [
+                timedelta(days=i) for i in range(1, args.multistep + 1)
+            ]
+
         ds_pred = era5.Era5Dataset(
             path=args.pred_path,
             filename_filter=_pred_filename_filter,  # Update filename_filter to filter within pred_path.
             variables=variables,
             return_timestamp=True,
-            dimension_indexers=dict(
-                prediction_timedelta=[timedelta(days=i) for i in range(1, args.multistep + 1)],
-                level=[500, 700, 850],
-            ),
+            dimension_indexers=dimension_indexers,
         )
         print(f"Reading {len(ds_pred.files)} files from pred_path: {args.pred_path}.")
+        if args.verbose:
+            print(ds_pred.files)
+            print("# prediction examples:", len(ds_pred))
+            print("# test examples:", len(ds_test))
 
         if reloaded_timestamp is not None:
             # Don't include the reloaded timestamp.
@@ -315,8 +337,13 @@ def __getitem__(self, idx):
             collate_fn=_custom_collate_fn,
         )
 
+    if args.breakpoint:
+        breakpoint()
+
     # iterable = tqdm(dl_test) if args.eval_clim else tqdm(zip(dl_test, dl_pred))
     for next_batch in tqdm(dl_test) if args.eval_clim else tqdm(zip(dl_test, dl_pred)):
+        if args.verbose:
+            print(f"{nbatches} batch")
         nbatches += 1
 
         if args.eval_clim:
@@ -333,7 +360,7 @@ def __getitem__(self, idx):
             pred = pred.apply(
                 lambda tensor: rearrange(
                     tensor,
-                    "batch var mem ... lev lat lon -> batch mem ... var lev lat lon",
+                    "batch var ... lev lat lon -> batch ... var lev lat lon",
                 )
             )
         timestamps = target["timestamp"]
@@ -344,9 +371,14 @@ def __getitem__(self, idx):
         else:
             target = target["future_states"]
 
+        if args.breakpoint:
+            breakpoint()
+
         # Update metrics.
         for metric in metrics.values():
             metric.update(target.to(device), pred.to(device))
+            if args.breakpoint:
+                breakpoint()
 
         if args.cache_metrics_every_nbatches and nbatches % args.cache_metrics_every_nbatches == 0:
             print(f"Processed {nbatches} batches.")
@@ -370,26 +402,35 @@ def __getitem__(self, idx):
         else:
             output_filename = f"test-multistep={args.multistep}-{metric_name}"
 
-        # Get xr dataset.
         if isinstance(labelled_metric_output, dict):
             labelled_dict = {
                 k: (v.cpu() if hasattr(v, "cpu") else v) for k, v in labelled_metric_output.items()
             }
-            extra_dimensions = ["prediction_timedelta"]
-            if "brier" in metric_name:
-                extra_dimensions = ["quantile", "prediction_timedelta"]
-            if "rankhist" in metric_name or "rank_hist" in metric_name:
-                extra_dimensions = ["bins", "prediction_timedelta"]
-            ds = convert_metric_dict_to_xarray(labelled_dict, extra_dimensions)
-
             # Write labeled dict.
             labelled_dict["metadata"] = dict(
                 groundtruth_path=args.groundtruth_path, predictions_path=args.pred_path
             )
             torch.save(labelled_dict, Path(output_dir).joinpath(f"{output_filename}.pt"))
+
+            # Convert to xr dataset.
+            extra_dimensions = []
+            if args.multistep > 1:
+                extra_dimensions = ["prediction_timedelta"]
+            if "brier" in metric_name:
+                extra_dimensions.insert(0, "quantile")  # ["quantile", "prediction_timedelta"]
+            if "rankhist" in metric_name or "rank_hist" in metric_name:
+                extra_dimensions.insert(0, "bins")  # ["bins", "prediction_timedelta"]
+            if "spatial" in metric_name:
+                # Does not yet handle extra lat/lon dims.
+                continue
+
+            ds = convert_metric_dict_to_xarray(labelled_dict, extra_dimensions)
         else:
             ds = labelled_metric_output
         # Write xr dataset.
+        ds.attrs["groundtruth_path"] = args.groundtruth_path
+        ds.attrs["predictions_path"] = args.args.pred_path
+        ds.attrs["groundtruth_dataset_domain"] = args.groundtruth_dataset_domain
         ds.to_netcdf(Path(output_dir).joinpath(f"{output_filename}.nc"))
 
 

geoarches/evaluation/metric_registry.py

@@ -5,6 +5,7 @@
 import torchmetrics
 
 from geoarches.metrics.brier_skill_score import Era5BrierSkillScore
+from geoarches.metrics.deterministic_metrics import Era5DeterministicMetrics
 from geoarches.metrics.ensemble_metrics import Era5EnsembleMetrics
 from geoarches.metrics.rank_histogram import Era5RankHistogram
 from geoarches.metrics.spherical_power_spectrum import Era5PowerSpectrum
@@ -38,6 +39,13 @@ def instantiate_metric(metric_name: str, **extra_kwargs):
 #######################################################
 ###### Registering classes with their arguments. ######
 #######################################################
+register_metric(
+    "era5_deterministic_metrics",
+    Era5DeterministicMetrics,
+)
+register_metric(
+    "era5_deterministic_metrics_with_spatial", Era5DeterministicMetrics, compute_per_gridpoint=True
+)
 register_metric(
     "era5_ensemble_metrics",
     Era5EnsembleMetrics,

geoarches/metrics/label_wrapper.py

@@ -59,7 +59,11 @@ def _convert(self, raw_metric_dict: Dict[str, Tensor]):
         labeled_dict = dict()
         for var, index in self.variable_indices.items():
             for metric_name, metric in raw_metric_dict.items():
-                labeled_dict[f"{metric_name}_{var}"] = metric.__getitem__((..., *index))
+                if any(s in metric_name for s in ["spatial", "per_gridpt", "per_gridpoint"]):
+                    # Account for lat, lon dims
+                    labeled_dict[f"{metric_name}_{var}"] = metric[..., *index, :, :]
+                else:
+                    labeled_dict[f"{metric_name}_{var}"] = metric[..., *index]
         return labeled_dict
 
     def update(self, *args: Any, **kwargs: Any) -> None:
@@ -134,6 +138,7 @@ def _convert_coord(name, value):
         labels = label.split("_")
         if len(labels) - 2 != len(extra_dimensions):
             raise ValueError(
+                f"Assumes metric name {label} is in format <metric>_<var>_<dim1>...."
                 f"Expected length of extra_dimensions for key {label} to be: {len(labels) - 2}. Got extra_dimensions={extra_dimensions}."
             )
         metrics.add(labels[0])