From 70594c58991961e0a4519647ad358820100d37da Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Fri, 13 Oct 2023 08:50:05 +0200 Subject: [PATCH 01/10] standardize requests to optimize caching --- earthkit/data/sources/cds.py | 20 +++++++++++++++++++- tests/sources/test_cds.py | 20 ++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index 657ebce44..52ff1b54a 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -7,6 +7,8 @@ # nor does it submit to any jurisdiction. # +import collections + import cdsapi import yaml @@ -18,6 +20,12 @@ from .prompt import APIKeyPrompt +def ensure_iterable(obj): + if isinstance(obj, str) or not isinstance(obj, collections.abc.Iterable): + return [obj] + return obj + + class CDSAPIKeyPrompt(APIKeyPrompt): register_or_sign_in_url = "https://cds.climate.copernicus.eu/" retrieve_api_key_url = "https://cds.climate.copernicus.eu/api-how-to" @@ -103,7 +111,7 @@ def retrieve(target, args): return self.cache_file( retrieve, - (dataset, request), + (dataset, self.normalized_request(**request)), extension=EXTENSIONS.get(request.get("format"), ".cache"), ) @@ -123,5 +131,15 @@ def requests(self, **kwargs): return result + @staticmethod + def normalized_request(**request): + result = {} + for k, v in sorted(request.items()): + v = ensure_iterable(v) + if k not in ("area", "grid"): + v = sorted(v) + result[k] = v[0] if len(v) == 1 else v + return result + source = CdsRetriever diff --git a/tests/sources/test_cds.py b/tests/sources/test_cds.py index 25c230d44..0768f426c 100644 --- a/tests/sources/test_cds.py +++ b/tests/sources/test_cds.py @@ -64,6 +64,26 @@ def test_cds_grib_3(): assert len(s) == 8 +@pytest.mark.long_test +@pytest.mark.download +@pytest.mark.skipif(NO_CDS, reason="No access to CDS") +def test_cds_normalized_request(): + variables = ["2t", "msl"] + base_request = dict( + product_type="reanalysis", + area=[50, -50, 20, 50], + date="2012-12-12", + time="12:00", + ) + s0 = from_source( + "cds", "reanalysis-era5-single-levels", variable=variables, **base_request + ) + s1 = from_source( + "cds", "reanalysis-era5-single-levels", variable=variables[::-1], **base_request + ) + assert s0.path == s1.path + + @pytest.mark.long_test @pytest.mark.download @pytest.mark.skipif(NO_CDS, reason="No access to CDS") From e9676344a2fe47f0e8d9695e81fa9ae0c86a340d Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Fri, 13 Oct 2023 09:02:17 +0200 Subject: [PATCH 02/10] py38 compatibility --- earthkit/data/sources/cds.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index 52ff1b54a..21ff853e5 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -7,7 +7,7 @@ # nor does it submit to any jurisdiction. # -import collections +import collections.abc import cdsapi import yaml From aca87f7ff27b43f5194441eb194d7d093a082cfa Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Mon, 16 Oct 2023 08:16:57 +0200 Subject: [PATCH 03/10] align with PR#228 --- earthkit/data/sources/cds.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index d69b862a6..9b9486107 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -111,7 +111,7 @@ def retrieve(target, args): return self.cache_file( retrieve, - (dataset, self.normalized_request(**request)), + (dataset, self._normalize_request(**request)), extension=EXTENSIONS.get(request.get("format"), ".cache"), ) @@ -125,14 +125,14 @@ def requests(self, **kwargs): result = [] for values in itertools.product( - *[ensure_iterable(kwargs[k]) for k in split_on] + *[sorted(ensure_iterable(kwargs[k])) for k in split_on] ): subrequest = dict(zip(split_on, values)) result.append(kwargs | subrequest) return result or [kwargs] @staticmethod - def normalized_request(**request): + def _normalize_request(**request): result = {} for k, v in sorted(request.items()): v = ensure_iterable(v) From 04954ed1f1ba0d29841d659e8b47052eda1c905b Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Tue, 17 Oct 2023 12:11:28 +0200 Subject: [PATCH 04/10] add test decorator --- tests/sources/test_cds.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/sources/test_cds.py b/tests/sources/test_cds.py index b81c38ce7..564c9fde5 100644 --- a/tests/sources/test_cds.py +++ b/tests/sources/test_cds.py @@ -84,6 +84,9 @@ def test_cds_normalized_request(): assert s0.path == s1.path +@pytest.mark.long_test +@pytest.mark.download +@pytest.mark.skipif(NO_CDS, reason="No access to CDS") @pytest.mark.parametrize( "split_on,expected_len", ( From 323d25304d35df8c6a5cf31bdcf905905e7cae9f Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Tue, 17 Oct 2023 12:26:05 +0200 Subject: [PATCH 05/10] fix normalizer --- earthkit/data/sources/cds.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index 4b4e77023..3de3ce2c3 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -125,7 +125,7 @@ def retrieve(target, args): return self.cache_file( retrieve, - (dataset, self._normalize_request(**request)), + (dataset, request), extension=EXTENSIONS.get(request.get("format"), ".cache"), ) @@ -133,13 +133,13 @@ def retrieve(target, args): @normalize("date", "date-list(%Y-%m-%d)") @normalize("area", "bounding-box(list)") def _normalize_request(**kwargs): - kwargs = {} + request = {} for k, v in sorted(kwargs.items()): v = ensure_iterable(v) if k not in ("area", "grid"): v = sorted(v) - kwargs[k] = v[0] if len(v) == 1 else v - return kwargs + request[k] = v[0] if len(v) == 1 else v + return request @cached_property def requests(self): From 9db4591b47b9dfa7c7b850bbaa4d21858b6e46d6 Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Tue, 17 Oct 2023 12:34:47 +0200 Subject: [PATCH 06/10] standardize requests list --- earthkit/data/sources/cds.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index 3de3ce2c3..cfe25d1c4 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -144,8 +144,7 @@ def _normalize_request(**kwargs): @cached_property def requests(self): requests = [] - for arg in self._args: - request = self._normalize_request(**arg) + for request in self._args: split_on = request.pop("split_on", None) if split_on is None: requests.append(request) @@ -153,12 +152,13 @@ def requests(self): if not isinstance(split_on, dict): split_on = {k: 1 for k in ensure_iterable(split_on)} + request = self._normalize_request(**request) for values in itertools.product( *[batched(ensure_iterable(request[k]), v) for k, v in split_on.items()] ): subrequest = dict(zip(split_on, values)) requests.append(request | subrequest) - return requests + return [self._normalize_request(**request) for request in requests] source = CdsRetriever From 2aebe0692429fa5a33b792acac09a0cf5860d20d Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Tue, 17 Oct 2023 12:44:06 +0200 Subject: [PATCH 07/10] fix empy iterables --- earthkit/data/sources/cds.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index cfe25d1c4..14cbf5969 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -146,12 +146,14 @@ def requests(self): requests = [] for request in self._args: split_on = request.pop("split_on", None) - if split_on is None: + if split_on is not None: + split_on = ensure_iterable(split_on) + if not split_on: requests.append(request) continue if not isinstance(split_on, dict): - split_on = {k: 1 for k in ensure_iterable(split_on)} + split_on = {k: 1 for k in split_on} request = self._normalize_request(**request) for values in itertools.product( *[batched(ensure_iterable(request[k]), v) for k, v in split_on.items()] From f452e6ebe7f56afe70dda5515679412a44e1ff7c Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Tue, 17 Oct 2023 12:53:22 +0200 Subject: [PATCH 08/10] cleanup --- earthkit/data/sources/cds.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/earthkit/data/sources/cds.py b/earthkit/data/sources/cds.py index 14cbf5969..a3bbc650e 100644 --- a/earthkit/data/sources/cds.py +++ b/earthkit/data/sources/cds.py @@ -145,15 +145,13 @@ def _normalize_request(**kwargs): def requests(self): requests = [] for request in self._args: - split_on = request.pop("split_on", None) - if split_on is not None: - split_on = ensure_iterable(split_on) + split_on = request.pop("split_on", {}) + if not isinstance(split_on, dict): + split_on = {k: 1 for k in ensure_iterable(split_on) if k is not None} if not split_on: requests.append(request) continue - if not isinstance(split_on, dict): - split_on = {k: 1 for k in split_on} request = self._normalize_request(**request) for values in itertools.product( *[batched(ensure_iterable(request[k]), v) for k, v in split_on.items()] From bfcaedbdf3e3d48c5750b431a58224149b6a8551 Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Wed, 18 Oct 2023 09:24:59 +0200 Subject: [PATCH 09/10] improve tests --- tests/sources/test_cds.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/tests/sources/test_cds.py b/tests/sources/test_cds.py index 8be932a10..79801619c 100644 --- a/tests/sources/test_cds.py +++ b/tests/sources/test_cds.py @@ -67,21 +67,30 @@ def test_cds_grib_3(): @pytest.mark.long_test @pytest.mark.download @pytest.mark.skipif(NO_CDS, reason="No access to CDS") -def test_cds_normalized_request(): - variables = ["2t", "msl"] +@pytest.mark.parametrize( + "variable1,variable2", + ( + ("2t", ["2t"]), + (["2t", "msl"], ["msl", "2t"]), + ), +) +def test_cds_normalized_request(variable1, variable2): base_request = dict( product_type="reanalysis", area=[50, -50, 20, 50], + grid=[2, 1], date="2012-12-12", time="12:00", ) - s0 = from_source( - "cds", "reanalysis-era5-single-levels", variable=variables, **base_request - ) s1 = from_source( - "cds", "reanalysis-era5-single-levels", variable=variables[::-1], **base_request + "cds", "reanalysis-era5-single-levels", variable=variable1, **base_request + ) + s2 = from_source( + "cds", "reanalysis-era5-single-levels", variable=variable2, **base_request ) - assert s0.path == s1.path + assert s1.path == s2.path + assert s1.to_xarray()["longitude"].values.tolist() == list(range(-50, 51, 2)) + assert s1.to_xarray()["latitude"].values.tolist() == list(range(50, 19, -1)) @pytest.mark.long_test From c51da96c98ea71797a44ff076238e7ea8b597154 Mon Sep 17 00:00:00 2001 From: Mattia Almansi Date: Wed, 18 Oct 2023 09:29:38 +0200 Subject: [PATCH 10/10] improve tests --- tests/sources/test_cds.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/sources/test_cds.py b/tests/sources/test_cds.py index 79801619c..e963cf465 100644 --- a/tests/sources/test_cds.py +++ b/tests/sources/test_cds.py @@ -68,13 +68,13 @@ def test_cds_grib_3(): @pytest.mark.download @pytest.mark.skipif(NO_CDS, reason="No access to CDS") @pytest.mark.parametrize( - "variable1,variable2", + "variable1,variable2,expected_vars", ( - ("2t", ["2t"]), - (["2t", "msl"], ["msl", "2t"]), + ("2t", ["2t"], {"t2m"}), + (["2t", "msl"], ["msl", "2t"], {"t2m", "msl"}), ), ) -def test_cds_normalized_request(variable1, variable2): +def test_cds_normalized_request(variable1, variable2, expected_vars): base_request = dict( product_type="reanalysis", area=[50, -50, 20, 50], @@ -89,6 +89,7 @@ def test_cds_normalized_request(variable1, variable2): "cds", "reanalysis-era5-single-levels", variable=variable2, **base_request ) assert s1.path == s2.path + assert set(s1.to_xarray().data_vars) == expected_vars assert s1.to_xarray()["longitude"].values.tolist() == list(range(-50, 51, 2)) assert s1.to_xarray()["latitude"].values.tolist() == list(range(50, 19, -1))