Merge pull request #111 from xcube-dev/konstntokas-xxx-add_drought_index_ensemble

Drought index ensemble added

Author: konstntokas

CHANGES.md

@@ -2,6 +2,7 @@
 
  - Adjust time range in ERA5 demo notebook. (#106)
  - Add data access for [monthly drought indices from 1940 to present derived from ERA5 reanalysis](https://cds.climate.copernicus.eu/datasets/derived-drought-historical-monthly?tab=overview). (#109)
+ - Add data access for [monthly drought indices from 1940 to present derived from ERA5 ensemble members](https://cds.climate.copernicus.eu/datasets/derived-drought-historical-monthly?tab=overview). (#110)
  - Support chunking of returned datasets. (#110)
 
 ## Changes in 1.1.0

examples/notebooks/Ex4-CDS-store-drought-indices.ipynb

@@ -0,0 +1 @@
+{"_dataset_name": "derived-drought-historical-monthly", "variable": ["standardised_precipitation_index", "test_for_normality_spi"], "accumulation_period": ["1", "3"], "area": [0.875, -0.875, -0.875, 0.875], "version": "1_0", "product_type": ["ensemble_members"], "dataset_type": "consolidated_dataset", "year": ["2015", "2016"], "month": ["01", "02", "10", "11", "12"]}

test/mock_results/test_drought_indices_open_data_ensemble/request.json

@@ -41,10 +41,11 @@ class CDSDroughtIndicesDatasetHandlerTest(unittest.TestCase):
     def setUp(self) -> None:
         self.drought_idx_handler = DroughtIndicesDatasetHandler()
         self.data_id_reanalysis = "derived-drought-historical-monthly:reanalysis"
+        self.data_id_ensemble = "derived-drought-historical-monthly:ensemble_members"
 
     def test_get_supported_data_ids(self):
         ids = self.drought_idx_handler.get_supported_data_ids()
-        self.assertCountEqual([self.data_id_reanalysis], ids)
+        self.assertCountEqual([self.data_id_reanalysis, self.data_id_ensemble], ids)
 
     def test_get_human_readable_data_id(self):
         self.assertEqual(
@@ -67,7 +68,7 @@ def test_get_open_data_params_schema(self):
             schema.required,
         )
 
-    def test_describe_data(self):
+    def test_describe_data_reanalysis(self):
         descriptor = self.drought_idx_handler.describe_data(self.data_id_reanalysis)
         self.assertEqual(self.data_id_reanalysis, descriptor.data_id)
         self.assertEqual("EPSG:4326", descriptor.crs)
@@ -76,8 +77,22 @@ def test_describe_data(self):
         self.assertEqual("1940-01-01", descriptor.time_range[0])
         self.assertEqual("2025-12-31", descriptor.time_range[1])
         self.assertEqual("1M", descriptor.time_period)
+        self.assertEqual(("time", "lat", "lon"), descriptor.data_vars["spi1"].dims)
 
-    def test_open_data(self):
+    def test_describe_data_ensemble(self):
+        descriptor = self.drought_idx_handler.describe_data(self.data_id_ensemble)
+        self.assertEqual(self.data_id_ensemble, descriptor.data_id)
+        self.assertEqual("EPSG:4326", descriptor.crs)
+        self.assertEqual((-180.0, -90.0, 180.0, 90.0), descriptor.bbox)
+        self.assertEqual(0.25, descriptor.spatial_res)
+        self.assertEqual("1940-01-01", descriptor.time_range[0])
+        self.assertEqual("2025-12-31", descriptor.time_range[1])
+        self.assertEqual("1M", descriptor.time_period)
+        self.assertEqual(
+            ("time", "number", "lat", "lon"), descriptor.data_vars["spi1"].dims
+        )
+
+    def test_open_data_reanalysis(self):
         opener = CDSDataOpener(
             client_class=get_cds_client(),
             endpoint_url=_CDS_API_URL,
@@ -105,6 +120,39 @@ def test_open_data(self):
             dataset.data_vars,
         )
 
+    def test_open_data_ensemble(self):
+        opener = CDSDataOpener(
+            client_class=get_cds_client(),
+            endpoint_url=_CDS_API_URL,
+            cds_api_key=_CDS_API_KEY,
+        )
+        dataset = opener.open_data(
+            self.data_id_ensemble,
+            variable_names=[
+                "standardised_precipitation_index",
+                "test_for_normality_spi",
+            ],
+            accumulation_periods=[1, 3],
+            bbox=[-1, -1, 1, 1],
+            time_range=["2015-10-15", "2016-02-02"],
+        )
+        self.assertIsNotNone(dataset)
+        # Monthly data is timestamped at the first of the month, so we expect
+        # four time co-ordinates (November to February inclusive).
+        self.assertEqual(
+            [4, 10, 7, 7],
+            [
+                dataset.sizes["time"],
+                dataset.sizes["number"],
+                dataset.sizes["lat"],
+                dataset.sizes["lon"],
+            ],
+        )
+        self.assertCountEqual(
+            ["spi1", "spi3", "spi1_significance", "spi3_significance"],
+            dataset.data_vars,
+        )
+
     def test_get_filepath_pattern_raises(self):
         with self.assertRaises(ValueError) as cm:
             self.drought_idx_handler._get_filepath_pattern("not_a_valid_key", 12)

test/mock_results/test_drought_indices_open_data_ensemble/result

@@ -24,6 +24,7 @@
 import os
 import pathlib
 
+import numpy as np
 import pandas as pd
 import xarray as xr
 from xcube.core.store import DatasetDescriptor, VariableDescriptor
@@ -45,6 +46,10 @@ def __init__(self):
                 "Monthly drought indices from 1940–present derived "
                 "from ERA5 reanalysis (main run)"
             ),
+            "derived-drought-historical-monthly:ensemble_members": (
+                "Monthly drought indices from 1940–present derived "
+                "from ERA5 ensemble (10 members)"
+            ),
         }
         self._variable_names = [
             "standardised_precipitation_index",
@@ -114,13 +119,18 @@ def describe_data(self, data_id: str) -> DatasetDescriptor:
                     var_name, accum_period
                 )
 
+        if data_id.endswith("reanalysis"):
+            dims = ("time", "lat", "lon")
+        else:
+            dims = ("time", "number", "lat", "lon")
+
         variable_descriptors = []
         for var_name, attrs in mapping_varname_attrs.items():
             variable_descriptors.append(
                 VariableDescriptor(
                     name=var_name,
                     dtype="float64",
-                    dims=("time", "lat", "lon"),
+                    dims=dims,
                     attrs=attrs,
                 )
             )
@@ -177,33 +187,72 @@ def read_file(
             zip_ref.extractall(path_temp)
         file_paths = glob.glob(f"{path_temp}/*")
         dss = []
-        for var_name in open_params["variable_names"]:
-            for accum_period in open_params["accumulation_periods"]:
-                pattern = self._get_filepath_pattern(var_name, accum_period)
-                file_sel = [path for path in file_paths if pattern in path]
-                file_sel = sorted(file_sel)
-                ds = xr.open_mfdataset(
-                    file_sel,
-                    engine="netcdf4",
-                    chunks="auto",
-                    combine_attrs="drop_conflicts",
-                )
-                if "standardised_precipitation" in var_name:
-                    ds = ds.sel(
-                        time=slice(
-                            open_params["time_range"][0], open_params["time_range"][1]
-                        )
+
+        if cds_api_params["product_type"] == ["reanalysis"]:
+            for var_name in open_params["variable_names"]:
+                for accum_period in open_params["accumulation_periods"]:
+                    pattern = self._get_filepath_pattern(var_name, accum_period)
+                    file_sel = [path for path in file_paths if pattern in path]
+                    file_sel = sorted(file_sel)
+                    ds = xr.open_mfdataset(
+                        file_sel,
+                        engine="netcdf4",
+                        chunks="auto",
+                        combine_attrs="drop_conflicts",
                     )
-                else:
-                    ds = self._resample_quality_ds(ds, open_params["time_range"])
-                assert len(ds.data_vars) == 1
-                ds_varname = self._get_varname(var_name, accum_period)
-                ds = ds.rename({list(ds.data_vars.keys())[0]: ds_varname})
-                dss.append(ds)
-        ds_final = xr.merge(dss, join="outer", combine_attrs="drop_conflicts")
-        ds_final = ds_final.sel(
-            time=slice(open_params["time_range"][0], open_params["time_range"][1])
-        )
+                    if "standardised_precipitation" in var_name:
+                        ds = ds.sel(
+                            time=slice(
+                                open_params["time_range"][0],
+                                open_params["time_range"][1],
+                            )
+                        )
+                    else:
+                        ds = self._resample_quality_ds(ds, open_params["time_range"])
+                    assert len(ds.data_vars) == 1
+                    ds_varname = self._get_varname(var_name, accum_period)
+                    ds = ds.rename({list(ds.data_vars.keys())[0]: ds_varname})
+                    dss.append(ds)
+            ds_final = xr.merge(dss, join="outer", combine_attrs="drop_conflicts")
+            ds_final = ds_final.sel(
+                time=slice(open_params["time_range"][0], open_params["time_range"][1])
+            )
+        else:
+            for var_name in open_params["variable_names"]:
+                for accum_period in open_params["accumulation_periods"]:
+                    pattern = self._get_filepath_pattern(var_name, accum_period)
+                    file_sel = [path for path in file_paths if pattern in path]
+                    file_sel = sorted(file_sel)
+                    dss_inner = []
+                    for path in file_sel:
+                        ds = xr.open_dataset(path, engine="netcdf4", chunks="auto")
+                        time_axis = ds.time
+                        # The data from the backend uses the confusing name `time` for the
+                        # ensemble member index. We rename it to `number` to be consistent
+                        # with other ERA5 datasets, and to free up the name `time` for the actual
+                        # time.
+                        ds = ds.rename({"time": "number"})
+                        ds = ds.assign_coords(number=np.arange(10))
+                        ds = ds.expand_dims(time=[time_axis[0].values])
+                        dss_inner.append(ds)
+                    ds = xr.concat(dss_inner, "time", combine_attrs="drop_conflicts")
+                    if "standardised_precipitation" in var_name:
+                        ds = ds.sel(
+                            time=slice(
+                                open_params["time_range"][0],
+                                open_params["time_range"][1],
+                            )
+                        )
+                    else:
+                        ds = self._resample_quality_ds(ds, open_params["time_range"])
+                    assert len(ds.data_vars) == 1
+                    ds_varname = self._get_varname(var_name, accum_period)
+                    ds = ds.rename({list(ds.data_vars.keys())[0]: ds_varname})
+                    dss.append(ds)
+            ds_final = xr.merge(dss, join="outer", combine_attrs="drop_conflicts")
+            ds_final = ds_final.sel(
+                time=slice(open_params["time_range"][0], open_params["time_range"][1])
+            )
         return ds_final
 
     @staticmethod