From ac1d42bdc5c096aaddeec7a320045403bfc5e471 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20Mari=C4=87?=
 <5569474+zoltanmaric@users.noreply.github.com>
Date: Fri, 30 Sep 2022 14:25:06 +0200
Subject: [PATCH] Combine ERA5 and ERA5T data

Requesting cutout data spanning recent (ERA5T)
and data older than ~3 months (ERA5) results in an
additional dimension in `cutout.data`, called `expver`,
which `atlite` currently cannot handle gracefully.

This change collapses the two dimensions into a single dimension.

See discussion in https://github.com/PyPSA/atlite/issues/190
---
 atlite/data.py                          |  9 ++++++++-
 atlite/datasets/era5.py                 |  7 +++++++
 test/test_preparation_and_conversion.py | 20 +++++++++++++++++---
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/atlite/data.py b/atlite/data.py
index 60dadcad..5280e56d 100644
--- a/atlite/data.py
+++ b/atlite/data.py
@@ -108,7 +108,9 @@ def wrapper(*args, **kwargs):
 
 
 @maybe_remove_tmpdir
-def cutout_prepare(cutout, features=None, tmpdir=None, overwrite=False):
+def cutout_prepare(
+    cutout, features=None, tmpdir=None, overwrite=False, flatten_expver=True
+):
     """
     Prepare all or a selection of features in a cutout.
 
@@ -133,6 +135,11 @@ def cutout_prepare(cutout, features=None, tmpdir=None, overwrite=False):
     overwrite : bool, optional
         Whether to overwrite variables which are already included in the
         cutout. The default is False.
+    flatten_expver: bool, default True
+        If true, data for time periods spanning both ERA5 and ERA5T data
+        will be combined into a single dimension.
+        See https://confluence.ecmwf.int/pages/viewpage.action?pageId=173385064
+        for details.
 
     Returns
     -------
diff --git a/atlite/datasets/era5.py b/atlite/datasets/era5.py
index c65ac71e..227d34f8 100644
--- a/atlite/datasets/era5.py
+++ b/atlite/datasets/era5.py
@@ -91,6 +91,13 @@ def _rename_and_clean_coords(ds, add_lon_lat=True):
     ds = maybe_swap_spatial_dims(ds)
     if add_lon_lat:
         ds = ds.assign_coords(lon=ds.coords["x"], lat=ds.coords["y"])
+
+    # Combine ERA5 and ERA5T data into a single dimension.
+    # See https://github.com/PyPSA/atlite/issues/190
+    if "expver" in ds.dims.keys():
+        ds = ds.sel(expver=1).combine_first(ds.sel(expver=5))
+        # ds = ds.reduce(np.nansum, 'expver')
+
     return ds
 
 
diff --git a/test/test_preparation_and_conversion.py b/test/test_preparation_and_conversion.py
index 522a8be8..ea555fca 100755
--- a/test/test_preparation_and_conversion.py
+++ b/test/test_preparation_and_conversion.py
@@ -29,7 +29,7 @@
 
 # %% Predefine tests for cutout
 
-TIME = "2013-01-01"
+TIME = "2022-07-31"
 BOUNDS = (-4, 56, 1.5, 62)
 SARAH_DIR = os.getenv("SARAH_DIR", "/home/vres/climate-data/sarah_v2")
 GEBCO_PATH = os.getenv("GEBCO_PATH", "/home/vres/climate-data/GEBCO_2014_2D.nc")
@@ -42,7 +42,7 @@ def all_notnull_test(cutout):
 
 def prepared_features_test(cutout):
     """
-    The prepared features series should contain all variables in cuttout.data
+    The prepared features series should contain all variables in cutout.data
     """
     assert set(cutout.prepared_features) == set(cutout.data)
 
@@ -306,6 +306,16 @@ def cutout_era5(tmp_path_factory):
     return cutout
 
 
+@pytest.fixture(scope="session")
+def cutout_mixed_era5_and_era5t(tmp_path_factory):
+    # Try just getting Berlin
+    # Try just requesting era5 and era5t
+    tmp_path = tmp_path_factory.mktemp("era5_era5t")
+    cutout = Cutout(path=tmp_path / "era5_era5t", module="era5", bounds=BOUNDS, time=slice("2022-07-31", "2022-08-01"))
+    cutout.prepare()
+    return cutout
+
+
 @pytest.fixture(scope="session")
 def cutout_era5_3h_sampling(tmp_path_factory):
     tmp_path = tmp_path_factory.mktemp("era5")
@@ -429,7 +439,7 @@ class TestERA5:
     @staticmethod
     def test_data_module_arguments_era5(cutout_era5):
         """
-        All data variables should have an attribute to which module thay belong
+        All data variables should have an attribute to which module they belong
         """
         for v in cutout_era5.data:
             assert cutout_era5.data.attrs["module"] == "era5"
@@ -526,6 +536,10 @@ def test_pv_era5_3h_sampling(cutout_era5_3h_sampling):
         assert pd.infer_freq(cutout_era5_3h_sampling.data.time) == "3H"
         return pv_test(cutout_era5_3h_sampling)
 
+    @staticmethod
+    def test_pv_era5(cutout_mixed_era5_and_era5t):
+        return pv_test(cutout_mixed_era5_and_era5t)
+
     @staticmethod
     def test_wind_era5(cutout_era5):
         return wind_test(cutout_era5)