diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index 687d5bce..48128bec 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -50,6 +50,10 @@ Upcoming Release **Features** * Add support for Python 3.14. +* Add ``sarah_from_api`` option to ``Cutout`` to download SARAH data directly from + the EUMETSAT Data Store (optional dependency ``eumdac``). The raw files are + stored temporarily and removed after the cutout is prepared. The number of + parallel downloads is set with ``sarah_api_workers``. **Bug fixes** diff --git a/atlite/cutout.py b/atlite/cutout.py index 5b66cf46..eb68f5ae 100644 --- a/atlite/cutout.py +++ b/atlite/cutout.py @@ -143,7 +143,18 @@ def __init__(self, path: PathLike, **cutoutparams: Any) -> None: effect for 'era5' data loading. sarah_dir : str, Path Directory of on-disk sarah data. This must be given when using the - sarah module. + sarah module, unless ``sarah_from_api=True``. + sarah_from_api : bool, default False + Download the SARAH data from the EUMETSAT Data Store API instead of + reading it from ``sarah_dir`` (which is then ignored). The time range + is taken from ``time``. The raw files are stored in a temporary + directory and deleted once the cutout is prepared. Requires the + ``eumdac`` package and credentials set once via + ``eumdac set-credentials `` via CLI. The + estimated download size is logged before the download starts. + sarah_api_workers : int, default 15 + Number of parallel downloads when ``sarah_from_api=True``. + See https://user.eumetsat.int/resources/user-guides/data-store-detailed-guide#ID-Throttling for maximum number. sarah_interpolate : bool, default True Whether to interpolate NaN's in the SARAH data. This takes effect for sarah data which has missing data for areas where dawn and diff --git a/atlite/data.py b/atlite/data.py index b01d049d..8b02a92c 100644 --- a/atlite/data.py +++ b/atlite/data.py @@ -279,51 +279,59 @@ def cutout_prepare( available_features(modules_list).loc[:, features_normalized].drop_duplicates() ) - for module in target.index.unique("module"): - missing_vars = target[module] - if not overwrite: - missing_vars = missing_vars[lambda v: ~v.isin(cutout.data)] - if missing_vars.empty: - continue - logger.info("Calculating and writing with module %s:", module) - missing_features = missing_vars.index.unique("feature") - ds = get_features( - cutout, - module, - missing_features, - data_format=data_format, - tmpdir=tmpdir, - monthly_requests=monthly_requests, - concurrent_requests=concurrent_requests, - ) - prepared |= set(missing_features) - - cutout.data.attrs.update({"prepared_features": list(prepared)}) - attrs = non_bool_dict(cutout.data.attrs) - attrs.update(ds.attrs) - - if compression: - for v in missing_vars: - ds[v].encoding.update(compression) - - ds = cutout.data.merge(ds[missing_vars.values]).assign_attrs(**attrs) - - directory, filename = os.path.split(str(cutout.path)) - fd, tmp = mkstemp(suffix=filename, dir=directory) - os.close(fd) - - logger.debug("Writing cutout to file...") - write_job = ds.to_netcdf(tmp, compute=False) - if show_progress: - with ProgressBar(minimum=2): + try: + for module in target.index.unique("module"): + missing_vars = target[module] + if not overwrite: + missing_vars = missing_vars[lambda v: ~v.isin(cutout.data)] + if missing_vars.empty: + continue + logger.info("Calculating and writing with module %s:", module) + missing_features = missing_vars.index.unique("feature") + ds = get_features( + cutout, + module, + missing_features, + data_format=data_format, + tmpdir=tmpdir, + monthly_requests=monthly_requests, + concurrent_requests=concurrent_requests, + ) + prepared |= set(missing_features) + + cutout.data.attrs.update({"prepared_features": list(prepared)}) + attrs = non_bool_dict(cutout.data.attrs) + attrs.update(ds.attrs) + + if compression: + for v in missing_vars: + ds[v].encoding.update(compression) + + ds = cutout.data.merge(ds[missing_vars.values]).assign_attrs(**attrs) + + directory, filename = os.path.split(str(cutout.path)) + fd, tmp = mkstemp(suffix=filename, dir=directory) + os.close(fd) + + logger.debug("Writing cutout to file...") + write_job = ds.to_netcdf(tmp, compute=False) + if show_progress: + with ProgressBar(minimum=2): + write_job.compute(**dask_kwargs) + else: write_job.compute(**dask_kwargs) - else: - write_job.compute(**dask_kwargs) - if cutout.path.exists(): - cutout.data.close() - cutout.path.unlink() - Path(tmp).rename(cutout.path) - - cutout.data = xr.open_dataset(cutout.path, chunks="auto") + if cutout.path.exists(): + cutout.data.close() + cutout.path.unlink() + Path(tmp).rename(cutout.path) + + cutout.data = xr.open_dataset(cutout.path, chunks="auto") + del ds, write_job + finally: + # Remove raw files that modules downloaded into tmpdir (e.g. SARAH API) + for module in target.index.unique("module"): + module_cleanup = getattr(datamodules[module], "cleanup", None) + if module_cleanup is not None: + module_cleanup(tmpdir) return cutout diff --git a/atlite/datasets/sarah.py b/atlite/datasets/sarah.py index 4c43ceee..995603aa 100644 --- a/atlite/datasets/sarah.py +++ b/atlite/datasets/sarah.py @@ -5,8 +5,13 @@ from __future__ import annotations +import gc import logging +import shutil +import time import warnings +import zipfile +from concurrent.futures import ThreadPoolExecutor, as_completed from functools import partial from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -40,6 +45,173 @@ } static_features: dict[str, list[str]] = {} +# EUMETSAT Data Store collection holding the SARAH-3 products +SARAH_COLLECTION_ID = "EO:EUM:DAT:0863" +SARAH_API_SUBDIR = "sarah_raw" +SARAH_API_PATTERNS = {"SID": "SIDin*", "SIS": "SISin*"} +SARAH_API_RETRIES = 3 +SARAH_API_DEFAULT_WORKERS = 15 + + +def _api_time_range(coords: dict[str, Any]) -> tuple[str, str]: + """Return the ISO start/end strings covering the days of the cutout.""" + index = coords["time"].to_index() + start = index[0].floor("D") + end = index[-1].floor("D") + pd.Timedelta(hours=23, minutes=30) + fmt = "%Y-%m-%dT%H:%M:%SZ" + return start.strftime(fmt), end.strftime(fmt) + + +def _download_nc(product: Any, target: Path, retries: int = SARAH_API_RETRIES) -> str: + """Download a product ZIP, extract its NetCDF file and remove the ZIP.""" + zip_path = target.with_suffix(".zip") + # Extract to a temporary name so an interrupted run never leaves a + # partial ``.nc`` file that looks valid. + tmp = target.with_suffix(".part") + + for attempt in range(1, retries + 1): + try: + with product.open() as src, zip_path.open("wb") as dst: + shutil.copyfileobj(src, dst) + + with zipfile.ZipFile(zip_path) as zf: + member = next( + name for name in zf.namelist() if name.lower().endswith(".nc") + ) + with zf.open(member) as src, tmp.open("wb") as dst: + shutil.copyfileobj(src, dst) + + tmp.replace(target) + except Exception as exc: + logger.warning( + "%s attempt %d/%d failed: %s", product, attempt, retries, exc + ) + if attempt < retries: + time.sleep(5 * attempt) + else: + return "file extracted" + finally: + zip_path.unlink(missing_ok=True) + + tmp.unlink(missing_ok=True) + return "failed" + + +def download_from_api( + start: str, + end: str, + target_dir: PathLike, + max_workers: int = SARAH_API_DEFAULT_WORKERS, +) -> Path: + """ + Download SARAH SID and SIS NetCDF files from the EUMETSAT Data Store. + + Credentials are read by the EUMDAC library from the file written by + ``eumdac set-credentials ``. + + Parameters + ---------- + start, end : str + ISO timestamps (``YYYY-MM-DDTHH:MM:SSZ``) delimiting the search. + target_dir : PathLike + Directory in which the ``SID`` and ``SIS`` sub-directories are created. + max_workers : int, default 15 + Number of parallel downloads. + + Returns + ------- + pathlib.Path + The directory containing the downloaded files. + + Raises + ------ + ImportError + If the optional ``eumdac`` package is not installed. + RuntimeError + If no credentials are found, nothing matches the time range or any + product fails to download. + """ + try: + import eumdac + from eumdac.cli_helpers import ( + CredentialsFileNotFoundError, + load_credentials, + ) + except ImportError as exc: + raise ImportError( + "Downloading SARAH data requires the 'eumdac' package. " + "Install it with `pip install eumdac`." + ) from exc + + try: + credentials = load_credentials() + except CredentialsFileNotFoundError as exc: + raise RuntimeError( + "No EUMETSAT credentials found. Set them once with " + "`eumdac set-credentials `." + ) from exc + + datastore = eumdac.DataStore(eumdac.AccessToken(credentials)) + collection = datastore.get_collection(SARAH_COLLECTION_ID) + + target_dir = Path(target_dir) + pending = [ + (var, product, product.size) + for var, pattern in SARAH_API_PATTERNS.items() + for product in collection.search(dtstart=start, dtend=end, title=pattern) + ] + if not pending: + raise RuntimeError(f"No SARAH products found between {start} and {end}.") + + # product.size is in kilobytes + size_gib = sum(size for _, _, size in pending) / 1024**2 + logger.warning( + "SARAH download: %d products, estimated size %.2f GiB.", + len(pending), + size_gib, + ) + + for var in SARAH_API_PATTERNS: + (target_dir / var).mkdir(parents=True, exist_ok=True) + + failed = [] + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit( + _download_nc, product, target_dir / var / f"{product}.nc" + ): ( + var, + product, + ) + for var, product, _ in pending + } + for i, future in enumerate(as_completed(futures), 1): + var, product = futures[future] + try: + status = future.result() + except Exception as exc: + logger.warning("%s failed: %s", product, exc) + status = "failed" + if status == "failed": + failed.append(str(product)) + logger.info("[%d/%d] %s %s -> %s", i, len(pending), var, product, status) + + if failed: + raise RuntimeError(f"{len(failed)} SARAH products failed to download: {failed}") + return target_dir + + +def cleanup(tmpdir: PathLike) -> None: + """ + Remove the raw SARAH files downloaded from the API into ``tmpdir``. + + Called by :func:`atlite.data.cutout_prepare` once the cutout is written. + """ + # Drop lingering references so file handles of the lazily opened NetCDF + # files are released (required on Windows before deleting). + gc.collect() + shutil.rmtree(Path(tmpdir) / SARAH_API_SUBDIR, ignore_errors=True) + def get_filenames(sarah_dir: str | PathLike, coords: dict[str, Any]) -> pd.DataFrame: """ @@ -218,8 +390,11 @@ def get_data( monthly_requests : bool, default False Whether to split requests by month. **creation_parameters - Must include ``sarah_dir``. Optional keys: ``parallel`` (bool), - ``sarah_interpolate`` (bool). + Either ``sarah_dir`` or ``sarah_from_api=True`` must be given. + Optional keys: ``parallel`` (bool), ``sarah_interpolate`` (bool), + ``sarah_from_api`` (bool, default False; downloads the files from the + EUMETSAT API into ``tmpdir`` and ignores ``sarah_dir``), + ``sarah_api_workers`` (int, default 15). Returns ------- @@ -232,10 +407,22 @@ def get_data( coords = cutout.coords chunks = cutout.chunks - sarah_dir = creation_parameters["sarah_dir"] creation_parameters.setdefault("parallel", False) creation_parameters.setdefault("sarah_interpolate", True) + if creation_parameters.get("sarah_from_api", False): + start, end = _api_time_range(coords) + sarah_dir = download_from_api( + start, + end, + Path(tmpdir) / SARAH_API_SUBDIR, + max_workers=int( + creation_parameters.get("sarah_api_workers", SARAH_API_DEFAULT_WORKERS) + ), + ) + else: + sarah_dir = creation_parameters["sarah_dir"] + files = get_filenames(sarah_dir, coords) open_kwargs = {"chunks": chunks, "parallel": creation_parameters["parallel"]} ds_sis = xr.open_mfdataset(files.sis, combine="by_coords", **open_kwargs)[["SIS"]] diff --git a/pyproject.toml b/pyproject.toml index 73ea9844..a36131e1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -60,6 +60,8 @@ Documentation = "https://atlite.readthedocs.io/en/latest/" [project.optional-dependencies] +sarah-api = ["eumdac"] + dev = ["pre-commit", "pytest", "pytest-cov", "matplotlib", "ruff", "mypy", "types-PyYAML", "types-requests"] docs = [ diff --git a/test/test_sarah_api.py b/test/test_sarah_api.py new file mode 100644 index 00000000..9ad86ab4 --- /dev/null +++ b/test/test_sarah_api.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Contributors to atlite +# +# SPDX-License-Identifier: MIT +"""Tests for downloading SARAH data from the EUMETSAT API (mocked).""" + +import io +import sys +import types +import zipfile + +import pandas as pd +import pytest +import xarray as xr + +from atlite.datasets import sarah + + +class FakeProduct: + size = 1024 # kilobytes + + def __init__(self, name): + self.name = name + + def __str__(self): + return self.name + + def open(self): + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + zf.writestr(f"{self.name}.nc", b"netcdf-bytes") + buf.seek(0) + return buf + + +class FakeCollection: + def search(self, dtstart, dtend, title): + prefix = title.rstrip("*") + return [FakeProduct(f"{prefix}20190101"), FakeProduct(f"{prefix}20190102")] + + +@pytest.fixture +def fake_eumdac(monkeypatch): + class CredentialsFileNotFoundError(Exception): + pass + + eumdac = types.ModuleType("eumdac") + eumdac.AccessToken = lambda credentials: credentials + eumdac.DataStore = lambda token: types.SimpleNamespace( + get_collection=lambda collection_id: FakeCollection() + ) + helpers = types.ModuleType("eumdac.cli_helpers") + helpers.CredentialsFileNotFoundError = CredentialsFileNotFoundError + helpers.load_credentials = lambda: ("key", "secret") + eumdac.cli_helpers = helpers + monkeypatch.setitem(sys.modules, "eumdac", eumdac) + monkeypatch.setitem(sys.modules, "eumdac.cli_helpers", helpers) + return helpers + + +def test_api_time_range(): + time = pd.date_range("2019-01-01", "2019-01-03 12:00", freq="h") + coords = xr.Dataset(coords={"time": time}).coords + assert sarah._api_time_range(coords) == ( + "2019-01-01T00:00:00Z", + "2019-01-03T23:30:00Z", + ) + + +def test_download_from_api(fake_eumdac, tmp_path, caplog): + with caplog.at_level("INFO", logger="atlite.datasets.sarah"): + out = sarah.download_from_api( + "2019-01-01T00:00:00Z", + "2019-01-02T23:30:00Z", + tmp_path / "raw", + max_workers=2, + ) + assert "estimated size 0.00 GiB" in caplog.text + assert "4 products" in caplog.text + assert sorted(p.name for p in (out / "SID").glob("*.nc")) == [ + "SIDin20190101.nc", + "SIDin20190102.nc", + ] + assert len(list((out / "SIS").glob("*.nc"))) == 2 + assert not list(out.rglob("*.zip")) + assert not list(out.rglob("*.part")) + + +def test_missing_credentials(fake_eumdac, tmp_path): + def raise_missing(): + raise fake_eumdac.CredentialsFileNotFoundError + + fake_eumdac.load_credentials = raise_missing + with pytest.raises(RuntimeError, match="set-credentials"): + sarah.download_from_api("a", "b", tmp_path / "raw") + + +def test_cleanup_removes_raw_files(tmp_path): + raw = tmp_path / sarah.SARAH_API_SUBDIR / "SID" + raw.mkdir(parents=True) + (raw / "x.nc").write_bytes(b"x") + sarah.cleanup(tmp_path) + assert not (tmp_path / sarah.SARAH_API_SUBDIR).exists()