diff --git a/README.md b/README.md index 2489a90..9c67cc4 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,7 @@ xarray-ecmwf is a Python library and [Xarray](https://docs.xarray.dev) backend w - the [ECMWF Open data](https://www.ecmwf.int/en/forecasts/datasets/open-data) via [ecmwf-opendata](https://github.com/ecmwf/ecmwf-opendata): High resolution forecasts, ensemble forecast - allows lazy loading the data and well integrated with [Dask](https://www.dask.org) and [Dask.distributed](https://distributed.dask.org) - allows chunking the input request according to a configurable splitting strategy. Allowed strategies: + - by N years - by one month - by one day - supports requests returning a single GRIB file, via [cfgrib](https://github.com/ecmwf/cfgrib) @@ -54,6 +55,25 @@ Attributes: ``` +## `force_end_date` + +Pass `force_end_date` in `request_chunker_kwargs`. The keyword uses underscores. + +For a request built from `year`, `month`, `day`, and `time`, the time coordinate is every combination of those lists. Timestamps after `force_end_date` are dropped. A timestamp equal to it stays. The dataset coordinate ends at the last remaining timestamp, also when only one chunk is left. With a monthly series on day `01` and `force_end_date` `1950-08-17`, the last time is `1950-08-01`. + +```python +>>> ds = xr.open_dataset( +... request, +... engine="ecmwf", +... request_chunks={"year": 1}, +... request_chunker_kwargs={"force_end_date": "1950-08-17"}, +... ) +``` + +The option applies to requests that contain `year`. A request that uses `date` is left unchanged. + +The CDS request for a year still carries the original `month` list. The coordinate ends at the last timestamp that is not after `force_end_date`. + ## License ``` diff --git a/tests/test_10_client_common.py b/tests/test_10_client_common.py index 07c80b7..a19cadc 100644 --- a/tests/test_10_client_common.py +++ b/tests/test_10_client_common.py @@ -2,10 +2,11 @@ import itertools from typing import Any +import numpy as np import pandas as pd import pytest -from xarray_ecmwf import client_common +from xarray_ecmwf import client_cdsapi, client_common ALL_MONTHS = [f"{m:02}" for m in range(1, 13)] ALL_DAYS = [f"{d:02}" for d in range(1, 32)] @@ -188,6 +189,141 @@ def test_build_chunk_ymd_year_requests_groups_years( assert time_chunk == tuple(len(group) for group in expected_year_groups) +def _monthly_year_request() -> dict[str, Any]: + return { + "year": ["2024", "2025", "2026"], + "month": [f"{month:02}" for month in range(1, 13)], + "day": ["01"], + "time": ["00:00"], + } + + +def test_force_end_date_trims_partial_last_year() -> None: + time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests( + _monthly_year_request(), + {"year": 1}, + force_end_date="2026-07-01", + ) + assert len(time) == 12 + 12 + 7 + assert time[-1] == np.datetime64("2026-07-01T00:00", "ns") + assert time_chunk == (12, 12, 7) + assert [chunk_request["year"] for _, chunk_request in time_chunk_requests] == [ + ["2024"], + ["2025"], + ["2026"], + ] + + +def test_force_end_date_drops_years_after_the_cutoff() -> None: + time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests( + _monthly_year_request(), + {"year": 1}, + force_end_date="2025-12-01", + ) + assert len(time) == 24 + assert time[-1] == np.datetime64("2025-12-01T00:00", "ns") + assert time_chunk == (12, 12) + assert [chunk_request["year"] for _, chunk_request in time_chunk_requests] == [ + ["2024"], + ["2025"], + ] + + +def test_force_end_date_shrinks_a_multi_year_chunk() -> None: + time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests( + _monthly_year_request(), + {"year": 2}, + force_end_date="2025-06-01", + ) + assert len(time) == 12 + 6 + assert time[-1] == np.datetime64("2025-06-01T00:00", "ns") + assert time_chunk == (18,) + assert time_chunk_requests == [(0, {"year": ["2024", "2025"]})] + + +def test_force_end_date_on_a_chunk_boundary_keeps_equal_day_chunks() -> None: + request = { + "year": ["2024"], + "month": ["01", "02"], + "day": ["01", "02"], + "time": ["00:00"], + } + time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests( + request, + {"day": 1}, + force_end_date="2024-01-02", + ) + assert list(time) == [ + np.datetime64("2024-01-01T00:00", "ns"), + np.datetime64("2024-01-02T00:00", "ns"), + ] + assert time_chunk == 1 + assert len(time_chunk_requests) == 2 + + +def test_force_end_date_shrinks_the_last_equal_sized_chunk() -> None: + request = { + "year": ["2024"], + "month": ["01"], + "day": ["01", "02"], + "time": ["00:00", "12:00"], + } + time, time_chunk, time_chunk_requests = client_common.build_time_chunk_requests( + request, + {"day": 1}, + force_end_date="2024-01-02T00:00", + ) + assert list(time) == [ + np.datetime64("2024-01-01T00:00", "ns"), + np.datetime64("2024-01-01T12:00", "ns"), + np.datetime64("2024-01-02T00:00", "ns"), + ] + assert time_chunk == (2, 1) + assert [chunk_request["day"] for _, chunk_request in time_chunk_requests] == [ + "01", + "02", + ] + + +def test_force_end_date_past_the_last_timestamp_is_a_no_op() -> None: + full, full_chunk, full_requests = client_common.build_time_chunk_requests( + _monthly_year_request(), {"year": 1} + ) + trimmed, trimmed_chunk, trimmed_requests = client_common.build_time_chunk_requests( + _monthly_year_request(), {"year": 1}, force_end_date="2027-01-01" + ) + assert list(trimmed) == list(full) + assert trimmed_chunk == full_chunk + assert trimmed_requests == full_requests + + +def test_force_end_date_stays_when_one_chunk_remains() -> None: + chunker = client_cdsapi.CdsapiRequestChunker( + request={ + "year": ["1950", "1951"], + "month": [f"{month:02}" for month in range(1, 13)], + "day": ["01"], + "time": ["00:00"], + }, + request_chunks={"year": 1}, + time_dim="valid_time", + force_end_date="1950-08-17", + ) + valid_time = chunker.compute_chunked_request_coords()["valid_time"] + assert valid_time.size == 8 + assert valid_time.values[-1] == np.datetime64("1950-08-01T00:00", "ns") + assert chunker.first_chunk_request()["year"] == ["1950"] + + +def test_force_end_date_before_the_first_timestamp_raises() -> None: + with pytest.raises(ValueError, match="before the first timestamp"): + client_common.build_time_chunk_requests( + _monthly_year_request(), + {"year": 1}, + force_end_date="2020-01-01", + ) + + def test_build_chunk_ymd_year_requests_invalid() -> None: request = { "year": ["2022"], diff --git a/xarray_ecmwf/client_cdsapi.py b/xarray_ecmwf/client_cdsapi.py index db56c4c..aa50055 100644 --- a/xarray_ecmwf/client_cdsapi.py +++ b/xarray_ecmwf/client_cdsapi.py @@ -52,6 +52,7 @@ class CdsapiRequestChunker: merge_date_time: bool = True time_dim: str = "time" time_sep: str = "/" + force_end_date: str | None = None def get_request_dimensions(self) -> dict[str, list[Any]]: request_dimensions: dict[str, list[Any]] = {} @@ -116,9 +117,12 @@ def compute_chunked_request_coords(self) -> dict[str, Any]: time_chunk, time_chunk_requests, ) = client_common.build_time_chunk_requests( - self.request | override_time, self.request_chunks, self.time_sep + self.request | override_time, + self.request_chunks, + self.time_sep, + force_end_date=self.force_end_date, ) - if len(time_chunk_requests) > 1: + if self.force_end_date is not None or len(time_chunk_requests) > 1: self.chunks[self.time_dim] = time_chunk self.chunk_requests[self.time_dim] = time_chunk_requests self.chunked_coords[self.time_dim] = xr.IndexVariable( # type: ignore diff --git a/xarray_ecmwf/client_common.py b/xarray_ecmwf/client_common.py index abb3b72..93239cc 100644 --- a/xarray_ecmwf/client_common.py +++ b/xarray_ecmwf/client_common.py @@ -1,10 +1,13 @@ import calendar +import logging from typing import Any, ContextManager, Protocol import numpy as np import pandas as pd import xarray as xr +LOGGER = logging.getLogger(__name__) + class RequestClientProtocol(Protocol): def __init__(self, client_kwargs: dict[str, Any]) -> None: ... @@ -263,8 +266,62 @@ def build_chunk_ymd_year_requests( return np.array(datetimes), tuple(chunks), chunk_requests +def trim_ymd_time_to_force_end_date( + time: np.typing.NDArray[np.datetime64], + time_chunk: int | tuple[int, ...], + time_chunk_requests: list[tuple[int, dict[str, Any]]], + force_end_date: str, +) -> tuple[ + np.typing.NDArray[np.datetime64], + int | tuple[int, ...], + list[tuple[int, dict[str, Any]]], +]: + """Cut a year/month/day time coordinate at ``force_end_date``. + + ``build_chunk_ymd_requests`` builds every timestamp in the cartesian + product, so the last year is a full set of months even when CDS has not + published them yet. Timestamps after the cutoff are dropped, and the + chunk sizes are resized so they still add up to the coordinate that + remains. The cutoff is inclusive: a timestamp equal to + ``force_end_date`` stays. A chunk that lies entirely past the cutoff + is removed. A chunk that crosses it keeps the same CDS request and only + its length changes. + """ + if len(time) == 0: + raise ValueError("cannot apply force_end_date to an empty time coordinate") + + end = pd.Timestamp(force_end_date).to_datetime64().astype("datetime64[ns]") + cut = int(np.searchsorted(time, end, side="right")) + if cut == 0: + raise ValueError( + f"end_date {force_end_date} is before the first timestamp {time[0]}" + ) + if cut >= len(time): + return time, time_chunk, time_chunk_requests + + starts = [start for start, _ in time_chunk_requests] + new_requests: list[tuple[int, dict[str, Any]]] = [] + new_sizes: list[int] = [] + for index, (start, chunk_request) in enumerate(time_chunk_requests): + if start >= cut: + break + next_start = starts[index + 1] if index + 1 < len(starts) else len(time) + new_sizes.append(min(next_start, cut) - start) + new_requests.append((start, chunk_request)) + + trimmed = time[:cut] + if isinstance(time_chunk, int) and all(size == time_chunk for size in new_sizes): + return trimmed, time_chunk, new_requests + if isinstance(time_chunk, int) and len(new_sizes) == 1: + return trimmed, new_sizes[0], new_requests + return trimmed, tuple(new_sizes), new_requests + + def build_time_chunk_requests( - request: dict[str, Any], request_chunks: dict[str, int], sep: str = "/" + request: dict[str, Any], + request_chunks: dict[str, int], + sep: str = "/", + force_end_date: str | None = None, ) -> tuple[ np.typing.NDArray[np.datetime64], int | tuple[int, ...], @@ -274,10 +331,17 @@ def build_time_chunk_requests( time, time_chunk, time_chunk_requests = build_chunk_ymd_requests( request, request_chunks ) + if force_end_date is not None: + time, time_chunk, time_chunk_requests = trim_ymd_time_to_force_end_date( + time, time_chunk, time_chunk_requests, force_end_date + ) elif "date" in request: time, time_chunk, time_chunk_requests = build_chunk_date_requests( request, request_chunks, sep ) + if force_end_date is not None: + msg = "forced_end_date param is ignored for date-interval requests." + LOGGER.warning(msg) else: raise ValueError("request must contain either 'year' or 'date'")