diff --git a/README.md b/README.md index 553fbb6..87040da 100644 --- a/README.md +++ b/README.md @@ -21,11 +21,19 @@ uv run python -m climatebenchpress.data_loader.datasets.cams uv run python -m climatebenchpress.data_loader.datasets.ifs_uncompressed uv run python -m climatebenchpress.data_loader.datasets.ifs_humidity uv run python -m climatebenchpress.data_loader.datasets.nextgems -uv run python -m climatebenchpress.data_loader.datasets.cmip6.access_ta -uv run python -m climatebenchpress.data_loader.datasets.cmip6.access_tos +uv run python -m climatebenchpress.data_loader.datasets.cmip6.access_atmos +uv run python -m climatebenchpress.data_loader.datasets.cmip6.access_ocean ``` This will download the data into a sub-directory named `datasets` within this repository. If you want to store the data in a different directory you can use the `--basepath=${path/to/dir}` command line argument for the scripts which will store the data at `${path/to/dir}/datasets` instead. +By default, the already-processed datasets are downloaded from the object store at `s3://esiwacebucket/ClimateBenchPress/`. This is much faster than reprocessing them, and it means that the `[data]` extras are only needed for the reprocessing described below. + +### Reprocessing the data + +To rebuild a dataset from its original data source instead, pass `--reprocess-data`: +```bash +uv run python -m climatebenchpress.data_loader.datasets.cams --reprocess-data\ + ## Citation If you find this work useful, please consider citing the following paper: diff --git a/mkdocs.yml b/mkdocs.yml index 6309d8e..60bd3e9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -40,7 +40,7 @@ plugins: source_dirs: - nav_heading: [Documentation] base: src - ignore: ["cf.py", "download.py", "canon.py", "monitor.py", "all.py"] + ignore: ["cf.py", "download.py", "canon.py", "monitor.py", "all.py", "s3.py"] - mkdocstrings: enable_inventory: true handlers: diff --git a/pyproject.toml b/pyproject.toml index 9a3f7cb..e940dbb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,6 +17,7 @@ dependencies = [ "intake==0.7.0", "intake-xarray==0.7.0", "requests~=2.32.3", + "s3fs>=2024.10.0,<2025.4", # Setuptools is added to ensure compatibility with intake-xarray version 0.7.0. # intake-xarray is relying on distutils which was removed from Python 3.12. "setuptools~=75.8.2", @@ -54,5 +55,5 @@ where = ["src"] addopts = ["--import-mode=importlib"] [[tool.mypy.overrides]] -module = ["fsspec.*", "intake.*", "healpy.*", "earthkit.*"] +module = ["fsspec.*", "intake.*", "healpy.*", "earthkit.*", "s3fs.*", "zarr.*"] follow_untyped_imports = true diff --git a/scripts/check_s3_mirror.py b/scripts/check_s3_mirror.py new file mode 100644 index 0000000..6c37667 --- /dev/null +++ b/scripts/check_s3_mirror.py @@ -0,0 +1,201 @@ +"""Check that the datasets published in the ClimateBenchPress object store are +identical to a local reference copy of the processed datasets. + +Every dataset that is published under `s3://esiwacebucket/ClimateBenchPress/` is +downloaded with `climatebenchpress.data_loader.s3.fetch_standardized` -- i.e. the +very code path that the data loader itself uses -- and then compared against +`//standardized.zarr` both byte-for-byte and semantically. + +The script exits non-zero if any dataset differs. +""" + +import argparse +import hashlib +import sys +import tempfile +from pathlib import Path + +import xarray as xr +from climatebenchpress.data_loader import s3 + +DEFAULT_REFERENCE = Path( + "/gws/ssde/j25a/aopp/treichelt/ClimateBenchPress/data-loader/datasets" +) + +# Directories in the reference tree that are stale leftovers and must not be +# compared against the object store. +IGNORED_REFERENCES = frozenset({"ifs-humidity-old", "ifs-uncompressed-old"}) + +_HASH_BUFFER_SIZE = 4 * 1024 * 1024 + + +def _digest(path: Path) -> str: + h = hashlib.md5() + with path.open("rb") as f: + while chunk := f.read(_HASH_BUFFER_SIZE): + h.update(chunk) + return h.hexdigest() + + +def _relative_files(root: Path) -> dict[str, Path]: + return { + str(path.relative_to(root)): path for path in root.rglob("*") if path.is_file() + } + + +def compare_bytes(fetched: Path, reference: Path) -> list[str]: + """Compare the two Zarr stores object by object.""" + ours, theirs = _relative_files(fetched), _relative_files(reference) + + problems = [] + + for key in sorted(set(ours) - set(theirs)): + problems.append(f"only in the object store: {key}") + for key in sorted(set(theirs) - set(ours)): + problems.append(f"only in the reference: {key}") + + for key in sorted(set(ours) & set(theirs)): + mine, yours = ours[key], theirs[key] + + if mine.stat().st_size != yours.stat().st_size: + problems.append( + f"size differs: {key} " + f"({mine.stat().st_size} vs {yours.stat().st_size} bytes)" + ) + elif _digest(mine) != _digest(yours): + problems.append(f"content differs: {key}") + + return problems + + +def _differs(mine: object, yours: object) -> bool: + # NaN fill values are not equal to themselves, so compare them by repr. + return mine != yours and repr(mine) != repr(yours) + + +def compare_semantics(fetched: Path, reference: Path) -> list[str]: + """Compare the two Zarr stores as xarray datasets, including their encoding.""" + problems = [] + + ours = xr.open_zarr(fetched) + theirs = xr.open_zarr(reference) + + try: + xr.testing.assert_identical(ours, theirs) + except AssertionError as err: + problems.append(f"datasets are not identical: {err}") + + for name in sorted(set(ours.variables) & set(theirs.variables)): + mine, yours = ours[name], theirs[name] + + if mine.dtype != yours.dtype: + problems.append(f"{name}: dtype {mine.dtype} vs {yours.dtype}") + + for key in ("chunks", "compressor", "filters", "dtype", "_FillValue"): + if _differs(mine.encoding.get(key), yours.encoding.get(key)): + problems.append( + f"{name}: encoding[{key!r}] " + f"{mine.encoding.get(key)!r} vs {yours.encoding.get(key)!r}" + ) + + return problems + + +def check(name: str, cache: Path, reference: Path, progress: bool) -> str: + """Check a single dataset. Returns "ok", "unverified" or "MISMATCH".""" + print(f"\n=== {name}") + + expected = reference / name / s3.STANDARDIZED + if not expected.exists(): + print(f" UNVERIFIED: no reference copy at {expected}") + return "unverified" + + fetched = cache / name / s3.STANDARDIZED + if not fetched.exists(): + if not s3.fetch_standardized(name, fetched, progress=progress): + print(f" UNVERIFIED: {name} is not published in the object store") + return "unverified" + + problems = compare_bytes(fetched, expected) + if problems: + print(f" byte comparison: {len(problems)} problem(s)") + for problem in problems[:20]: + print(f" - {problem}") + if len(problems) > 20: + print(f" ... and {len(problems) - 20} more") + else: + print(" byte comparison: identical") + + semantic = compare_semantics(fetched, expected) + if semantic: + print(f" semantic comparison: {len(semantic)} problem(s)") + for problem in semantic[:20]: + print(f" - {problem}") + return "MISMATCH" + + print(" semantic comparison: identical") + + # A pure metadata difference is worth reporting but is not a mismatch. + return "ok (byte differences)" if problems else "ok" + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--reference", type=Path, default=DEFAULT_REFERENCE) + parser.add_argument( + "--cache", + type=Path, + default=None, + help="where to download the object store copies to (default: a temp dir)", + ) + parser.add_argument( + "--datasets", + nargs="*", + default=None, + help="only check these dataset names (default: everything in the bucket)", + ) + parser.add_argument("--no-progress", action="store_true") + args = parser.parse_args() + + published = s3.list_published() + names = args.datasets if args.datasets is not None else published + + print(f"published in the object store: {', '.join(published)}") + + local = { + path.name + for path in args.reference.iterdir() + if path.is_dir() and path.name not in IGNORED_REFERENCES + } + print(f"reference copies: {', '.join(sorted(local))}") + + missing_from_bucket = sorted(local - set(published)) + if missing_from_bucket: + print( + f"\nNOT published (will be reprocessed): {', '.join(missing_from_bucket)}" + ) + + with tempfile.TemporaryDirectory() as tmp: + cache = args.cache if args.cache is not None else Path(tmp) + cache.mkdir(parents=True, exist_ok=True) + + results = { + name: check(name, cache, args.reference, not args.no_progress) + for name in names + } + + print("\n=== summary") + for name, result in results.items(): + print(f" {result:24} {name}") + + failed = [name for name, result in results.items() if result == "MISMATCH"] + if failed: + print(f"\nFAILED: {len(failed)} dataset(s) differ: {', '.join(failed)}") + return 1 + + print("\nAll compared datasets match the reference copies.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/climatebenchpress/data_loader/__init__.py b/src/climatebenchpress/data_loader/__init__.py index 1a363c6..748bc95 100644 --- a/src/climatebenchpress/data_loader/__init__.py +++ b/src/climatebenchpress/data_loader/__init__.py @@ -1,26 +1,37 @@ __all__ = [ "canon", "datasets", + "s3", "open_downloaded_canonicalized_dataset", "open_downloaded_tiny_canonicalized_dataset", ] +import logging +from collections.abc import Callable, Mapping from pathlib import Path -from typing import Optional +from typing import Optional, Union import xarray as xr +import zarr -from . import canon, datasets, monitor +from . import canon, datasets, monitor, s3 from .datasets.abc import Dataset +Chunks = Union[int, Mapping[str, int]] + def open_downloaded_canonicalized_dataset( cls: type[Dataset], basepath: Path = Path(), progress: bool = True, + reprocess: bool = False, ) -> xr.Dataset: """Download a given dataset and canonicalize it, i.e. ensure that all the axes names are consistent between different datasets. + By default, the pre-processed dataset is downloaded from the S3 ESIWACE + object store. If it is not published there, or if `reprocess` is set, it is + instead reprocessed from the original data source. + Parameters ---------- cls : type[Dataset] @@ -29,6 +40,9 @@ def open_downloaded_canonicalized_dataset( The base path where the dataset should be stored, by default Path() progress : bool, optional Whether to show a progress bar during the download, by default True + reprocess : bool, optional + Whether to reprocess the dataset from the original data source instead of + downloading the pre-processed dataset, by default False Returns ------- @@ -37,16 +51,23 @@ def open_downloaded_canonicalized_dataset( """ datasets = basepath / "datasets" - download = datasets / cls.name / "download" - if not download.exists(): - download.mkdir(parents=True, exist_ok=True) - # The download function is responsible for checking whether the download is - # complete or not. If the previous download was interrupt it will resume the download. - # If the download is complete it will skip the download. - cls.download(download, progress) - standardized = datasets / cls.name / "standardized.zarr" - if not standardized.exists(): + if not standardized.exists() and not _download_standardized( + cls.name, + standardized, + chunks=cls.chunks, + consolidated=None, + progress=progress, + reprocess=reprocess, + ): + download = datasets / cls.name / "download" + if not download.exists(): + download.mkdir(parents=True, exist_ok=True) + # The download function is responsible for checking whether the download is + # complete or not. If the previous download was interrupt it will resume the download. + # If the download is complete it will skip the download. + cls.download(download, progress) + ds = cls.open(download) ds = canon.canonicalize_dataset(ds) @@ -61,6 +82,7 @@ def open_downloaded_tiny_canonicalized_dataset( basepath: Path = Path(), progress: bool = True, slices: Optional[dict[str, slice]] = None, + reprocess: bool = False, ) -> xr.Dataset: """Same as `open_downloaded_canonicalized_dataset`, but returns a subset of the dataset. @@ -76,6 +98,9 @@ def open_downloaded_tiny_canonicalized_dataset( Whether to show a progress bar during the download, by default True slices : Optional[dict[str, slice]], optional A dictionary of slices to apply to the dataset, by default None + reprocess : bool, optional + Whether to reprocess the dataset from the original data source instead of + downloading the pre-processed dataset, by default False Returns ------- @@ -84,13 +109,22 @@ def open_downloaded_tiny_canonicalized_dataset( """ datasets = basepath / "datasets" - download = datasets / f"{cls.name}" / "download" - if not download.exists(): - download.mkdir(parents=True, exist_ok=True) - cls.download(download, progress) - standardized = datasets / f"{cls.name}-tiny" / "standardized.zarr" - if not standardized.exists(): + if not standardized.exists() and not _download_standardized( + f"{cls.name}-tiny", + standardized, + # "Tiny-fication" can lead to inconsistent or suboptimal chunking, so the + # tiny datasets are always stored as a single chunk per array. + chunks=lambda ds: -1, + consolidated=True, + progress=progress, + reprocess=reprocess, + ): + download = datasets / f"{cls.name}" / "download" + if not download.exists(): + download.mkdir(parents=True, exist_ok=True) + cls.download(download, progress) + ds = cls.open(download) ds = canon.canonicalize_dataset(ds) ds = canon.canonical_tiny_dataset(ds, slices=slices) @@ -102,3 +136,121 @@ def open_downloaded_tiny_canonicalized_dataset( ds.to_zarr(standardized, compute=False, consolidated=True).compute() return xr.open_dataset(standardized, chunks=dict(), engine="zarr") + + +def _download_standardized( + name: str, + standardized: Path, + chunks: Callable[[xr.Dataset], Chunks], + consolidated: Optional[bool], + progress: bool, + reprocess: bool, +) -> bool: + """Download the pre-processed dataset from the S3 ESIWACE object store. + + If the published dataset was written by an older version of the pipeline, and + thus uses an outdated chunking or compressor, it is reformatted locally. Note + that only the encoding of the dataset is checked, not its contents. + + Returns + ------- + bool + Whether the pre-processed dataset is now available at `standardized` + """ + if reprocess: + return False + + if not s3.fetch_standardized(name, standardized, progress): + logging.warning( + f"{name} is not published in the S3 ESIWACE object store, " + f"reprocessing it from the original data source instead" + ) + return False + + ds = xr.open_zarr(standardized) + spec = chunks(ds) + + problems = _verify_encoding(ds, spec) + if problems: + logging.warning( + f"the published {name} uses an outdated encoding " + f"({'; '.join(problems)}), reformatting it locally" + ) + _reformat(standardized, ds, spec, consolidated, progress) + + return True + + +def _verify_encoding(ds: xr.Dataset, chunks: Chunks) -> list[str]: + """Check that a canonicalized dataset is stored the way the pipeline writes it. + + Returns + ------- + list[str] + The encoding problems that were found, empty if the dataset is up to date + """ + problems = [] + + for name, var in ds.variables.items(): + # The pipeline drops the encoding of the original data before writing, + # so every array uses Zarr's default compressor and no filters. + compressor = var.encoding.get("compressor") + if compressor != zarr.storage.default_compressor: + problems.append(f"{name} is compressed with {compressor}") + + filters = var.encoding.get("filters") + if filters: + problems.append(f"{name} is filtered with {filters}") + + # Only the data variables are chunked, the coordinates are always stored + # as a single chunk. + expected = _expected_chunks(var, chunks if name in ds.data_vars else -1) + actual = var.encoding.get("chunks") + + if actual is not None: + # A chunk that is larger than the array still stores the array in a + # single chunk, so only the effective chunk grid is compared. The + # published CMIP6 datasets, for instance, inherited the time + # chunking of the much longer upstream time series. + effective = tuple( + min(chunk, size) for chunk, size in zip(actual, var.shape) + ) + if effective != expected: + problems.append(f"{name} is chunked {effective} instead of {expected}") + + return problems + + +def _expected_chunks(var: xr.Variable, chunks: Chunks) -> tuple[int, ...]: + if not isinstance(chunks, Mapping): + chunks = {str(dim): chunks for dim in var.dims} + + return tuple( + size if (chunk := chunks.get(str(dim), -1)) in (-1, None) else min(chunk, size) + for dim, size in zip(var.dims, var.shape) + ) + + +def _reformat( + standardized: Path, + ds: xr.Dataset, + chunks: Chunks, + consolidated: Optional[bool], + progress: bool, +) -> None: + """Rewrite a dataset with the chunking and compressor that the pipeline uses. + + The dataset is written to a sibling directory first, so that it stays readable + while it is being reformatted. + """ + reformatted = standardized.parent / (standardized.name + ".reformatted") + if reformatted.exists(): + s3.remove_tree(reformatted) + + ds = ds.drop_encoding().chunk(chunks) + + with monitor.progress_bar(progress): + ds.to_zarr(reformatted, compute=False, consolidated=consolidated).compute() + + s3.remove_tree(standardized) + s3.move_tree(reformatted, standardized) diff --git a/src/climatebenchpress/data_loader/cli.py b/src/climatebenchpress/data_loader/cli.py new file mode 100644 index 0000000..6043123 --- /dev/null +++ b/src/climatebenchpress/data_loader/cli.py @@ -0,0 +1,65 @@ +__all__ = ["main"] + +import argparse +from collections.abc import Callable +from pathlib import Path +from typing import Optional + +import xarray as xr + +from . import ( + open_downloaded_canonicalized_dataset, + open_downloaded_tiny_canonicalized_dataset, +) +from .datasets.abc import Dataset + + +def main( + cls: type[Dataset], + tiny_slices: Optional[Callable[[xr.Dataset], dict[str, slice]]] = None, +) -> None: + """Download a dataset from the command line. + + This is the shared entry point for all `python -m climatebenchpress.data_loader.datasets.*` + scripts. By default the pre-processed dataset is downloaded from the + ClimateBenchPress object store, unless `--reprocess-data` is passed. + + Parameters + ---------- + cls : type[Dataset] + The dataset class to download and open + tiny_slices : Optional[Callable[[xr.Dataset], dict[str, slice]]], optional + Computes the slices that are used to build the tiny variant of the dataset + from the full dataset, by default None, i.e. the default slices are used + """ + parser = argparse.ArgumentParser(description=cls.__doc__) + parser.add_argument( + "--basepath", + type=Path, + default=Path(), + help="the directory that the `datasets` directory is created in", + ) + parser.add_argument( + "--reprocess-data", + action="store_true", + help=( + "reprocess the dataset from its original data source instead of " + "downloading the pre-processed dataset from the ClimateBenchPress " + "object store; note that an already-processed dataset is never " + "rebuilt, delete it first to force a rebuild" + ), + ) + args = parser.parse_args() + + ds = open_downloaded_canonicalized_dataset( + cls, basepath=args.basepath, reprocess=args.reprocess_data + ) + open_downloaded_tiny_canonicalized_dataset( + cls, + basepath=args.basepath, + slices=None if tiny_slices is None else tiny_slices(ds), + reprocess=args.reprocess_data, + ) + + for v, da in ds.items(): + print(f"- {v}: {da.dims}") diff --git a/src/climatebenchpress/data_loader/datasets/abc.py b/src/climatebenchpress/data_loader/datasets/abc.py index 6b83387..7cb3098 100644 --- a/src/climatebenchpress/data_loader/datasets/abc.py +++ b/src/climatebenchpress/data_loader/datasets/abc.py @@ -40,6 +40,29 @@ def download(download_path: Path, progress: bool = True): """ pass + @staticmethod + def chunks(ds: xr.Dataset) -> int | Mapping[str, int]: + """The chunking that the canonicalized dataset should be stored with. + + The return value is passed to `xarray.Dataset.chunk`. By default, each + array is stored as a single chunk. + + This is the single source of truth for the chunking of a dataset: it is + applied by `open` and is also used to check whether a pre-processed + dataset that was downloaded from the object store is still up to date. + + Parameters + ---------- + ds : xr.Dataset + The dataset that is about to be chunked + + Returns + ------- + int | Mapping[str, int] + The chunk sizes, either for all dimensions or per dimension + """ + return -1 + @staticmethod @abstractmethod def open(download_path: Path) -> xr.Dataset: diff --git a/src/climatebenchpress/data_loader/datasets/cams.py b/src/climatebenchpress/data_loader/datasets/cams.py index c1e2e87..ca32409 100644 --- a/src/climatebenchpress/data_loader/datasets/cams.py +++ b/src/climatebenchpress/data_loader/datasets/cams.py @@ -1,15 +1,11 @@ __all__ = ["CamsNitrogenDioxideDataset"] -import argparse import logging from pathlib import Path import xarray as xr -from .. import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ..cli import main from ..download import _download_netcdf from .abc import Dataset @@ -43,11 +39,8 @@ def download(download_path: Path, progress: bool = True): @staticmethod def open(download_path: Path) -> xr.Dataset: - ds = ( - xr.open_dataset(download_path / Path(NO2_FILE).name) - .drop_encoding() - .chunk(-1) - ) + ds = xr.open_dataset(download_path / Path(NO2_FILE).name).drop_encoding() + ds = ds.chunk(CamsNitrogenDioxideDataset.chunks(ds)) # valid_time contains actual dates, whereas step is the seconds (in simulated time) # since the model as been initialised. @@ -63,16 +56,4 @@ def open(download_path: Path) -> xr.Dataset: if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset( - CamsNitrogenDioxideDataset, basepath=args.basepath - ) - open_downloaded_tiny_canonicalized_dataset( - CamsNitrogenDioxideDataset, basepath=args.basepath - ) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(CamsNitrogenDioxideDataset) diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/abc.py b/src/climatebenchpress/data_loader/datasets/cmip6/abc.py index 61bdc1c..d0eedc2 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/abc.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/abc.py @@ -68,7 +68,8 @@ def download_with( @staticmethod def open(download_path: Path) -> xr.Dataset: - return xr.open_zarr(download_path / "download.zarr").drop_encoding().chunk(-1) + ds = xr.open_zarr(download_path / "download.zarr").drop_encoding() + return ds.chunk(Cmip6Dataset.chunks(ds)) @lru_cache @staticmethod diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/access_atmos.py b/src/climatebenchpress/data_loader/datasets/cmip6/access_atmos.py index fda28d4..cc8f5a3 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/access_atmos.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/access_atmos.py @@ -1,12 +1,8 @@ __all__ = ["Cmip6AtmosphereAccessDataset"] -import argparse from pathlib import Path -from ... import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ...cli import main from .abc import Cmip6AtmosphereDataset, Cmip6Dataset @@ -31,16 +27,4 @@ def download(download_path: Path, progress: bool = True): if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset( - Cmip6AtmosphereAccessDataset, basepath=args.basepath - ) - open_downloaded_tiny_canonicalized_dataset( - Cmip6AtmosphereAccessDataset, basepath=args.basepath - ) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Cmip6AtmosphereAccessDataset) diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/access_ocean.py b/src/climatebenchpress/data_loader/datasets/cmip6/access_ocean.py index 95aee04..d82d28f 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/access_ocean.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/access_ocean.py @@ -1,12 +1,8 @@ __all__ = ["Cmip6OceanAccessDataset"] -import argparse from pathlib import Path -from ... import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ...cli import main from .abc import Cmip6Dataset, Cmip6OceanDataset @@ -33,16 +29,4 @@ def download(download_path: Path, progress: bool = True): if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset( - Cmip6OceanAccessDataset, basepath=args.basepath - ) - open_downloaded_tiny_canonicalized_dataset( - Cmip6OceanAccessDataset, basepath=args.basepath - ) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Cmip6OceanAccessDataset) diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_atmos.py b/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_atmos.py index 600ef42..120335b 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_atmos.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_atmos.py @@ -2,10 +2,7 @@ from pathlib import Path -from ... import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ...cli import main from .abc import Cmip6AtmosphereDataset, Cmip6Dataset @@ -28,8 +25,4 @@ def download(download_path: Path, progress: bool = True): if __name__ == "__main__": - ds = open_downloaded_canonicalized_dataset(Cmip6AtmosphereCanEsm5Dataset) - open_downloaded_tiny_canonicalized_dataset(Cmip6AtmosphereCanEsm5Dataset) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Cmip6AtmosphereCanEsm5Dataset) diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_ocean.py b/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_ocean.py index 87ddcb0..ecede85 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_ocean.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/canesm5_ocean.py @@ -2,10 +2,7 @@ from pathlib import Path -from ... import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ...cli import main from .abc import Cmip6Dataset, Cmip6OceanDataset @@ -30,8 +27,4 @@ def download(download_path: Path, progress: bool = True): if __name__ == "__main__": - ds = open_downloaded_canonicalized_dataset(Cmip6OceanCanEsm5Dataset) - open_downloaded_tiny_canonicalized_dataset(Cmip6OceanCanEsm5Dataset) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Cmip6OceanCanEsm5Dataset) diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_atmos.py b/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_atmos.py index 5b3ed2f..13e485a 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_atmos.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_atmos.py @@ -2,10 +2,7 @@ from pathlib import Path -from ... import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ...cli import main from .abc import Cmip6AtmosphereDataset, Cmip6Dataset @@ -28,8 +25,4 @@ def download(download_path: Path, progress: bool = True): if __name__ == "__main__": - ds = open_downloaded_canonicalized_dataset(Cmip6AtmosphereUkEsmDataset) - open_downloaded_tiny_canonicalized_dataset(Cmip6AtmosphereUkEsmDataset) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Cmip6AtmosphereUkEsmDataset) diff --git a/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_ocean.py b/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_ocean.py index b306388..b791617 100644 --- a/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_ocean.py +++ b/src/climatebenchpress/data_loader/datasets/cmip6/ukesm_ocean.py @@ -2,10 +2,7 @@ from pathlib import Path -from ... import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ...cli import main from .abc import Cmip6Dataset, Cmip6OceanDataset @@ -30,8 +27,4 @@ def download(download_path: Path, progress: bool = True): if __name__ == "__main__": - ds = open_downloaded_canonicalized_dataset(Cmip6OceanUkEsmDataset) - open_downloaded_tiny_canonicalized_dataset(Cmip6OceanUkEsmDataset) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Cmip6OceanUkEsmDataset) diff --git a/src/climatebenchpress/data_loader/datasets/era5.py b/src/climatebenchpress/data_loader/datasets/era5.py index caed1fb..77e298c 100644 --- a/src/climatebenchpress/data_loader/datasets/era5.py +++ b/src/climatebenchpress/data_loader/datasets/era5.py @@ -1,15 +1,11 @@ __all__ = ["Era5Dataset"] -import argparse from pathlib import Path import xarray as xr -from .. import ( - monitor, - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from .. import monitor +from ..cli import main from .abc import Dataset ERA5_GCP_PATH = "https://storage.googleapis.com/gcp-public-data-arco-era5/ar/1959-2022-full_37-1h-0p25deg-chunk-1.zarr-v2" @@ -55,16 +51,9 @@ def download(download_path: Path, progress: bool = True): @staticmethod def open(download_path: Path) -> xr.Dataset: - return xr.open_zarr(download_path / "download.zarr").drop_encoding().chunk(-1) + ds = xr.open_zarr(download_path / "download.zarr").drop_encoding() + return ds.chunk(Era5Dataset.chunks(ds)) if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset(Era5Dataset, basepath=args.basepath) - open_downloaded_tiny_canonicalized_dataset(Era5Dataset, basepath=args.basepath) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(Era5Dataset) diff --git a/src/climatebenchpress/data_loader/datasets/esa_biomass_cci.py b/src/climatebenchpress/data_loader/datasets/esa_biomass_cci.py index 12460d6..75cd97f 100644 --- a/src/climatebenchpress/data_loader/datasets/esa_biomass_cci.py +++ b/src/climatebenchpress/data_loader/datasets/esa_biomass_cci.py @@ -1,15 +1,11 @@ __all__ = ["EsaBiomassCciDataset"] -import argparse import logging from pathlib import Path import xarray as xr -from .. import ( - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from ..cli import main from ..download import _download_netcdf from .abc import Dataset @@ -58,28 +54,17 @@ def open(download_path: Path) -> xr.Dataset: ds = ds.sel( lon=slice(FRANCE_BBOX[0], FRANCE_BBOX[2]), lat=slice(FRANCE_BBOX[3], FRANCE_BBOX[1]), - ).chunk(-1) + ) + ds = ds.chunk(EsaBiomassCciDataset.chunks(ds)) return ds[["agb"]] if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset( - EsaBiomassCciDataset, basepath=args.basepath - ) - num_lon, num_lat = ds.lon.size, ds.lat.size - open_downloaded_tiny_canonicalized_dataset( + main( EsaBiomassCciDataset, # Use a smaller spatial subset for the tiny dataset. - slices={ - "X": slice(num_lon // 2, (num_lon // 2) + 500), - "Y": slice(num_lat // 2, (num_lat // 2) + 500), + tiny_slices=lambda ds: { + "X": slice(ds.lon.size // 2, (ds.lon.size // 2) + 500), + "Y": slice(ds.lat.size // 2, (ds.lat.size // 2) + 500), }, - basepath=args.basepath, ) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") diff --git a/src/climatebenchpress/data_loader/datasets/ifs_humidity.py b/src/climatebenchpress/data_loader/datasets/ifs_humidity.py index be9a18e..b6e7702 100644 --- a/src/climatebenchpress/data_loader/datasets/ifs_humidity.py +++ b/src/climatebenchpress/data_loader/datasets/ifs_humidity.py @@ -1,15 +1,12 @@ __all__ = ["IFSHumidityDataset"] -import argparse +from collections.abc import Mapping from pathlib import Path import xarray as xr -from .. import ( - monitor, - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from .. import monitor +from ..cli import main from .abc import Dataset from .ifs_uncompressed import load_hplp_data, regrid_to_regular @@ -41,18 +38,23 @@ def download(download_path: Path, progress: bool = True): with monitor.progress_bar(progress): ds_regridded.to_zarr(downloadfile, mode="w", compute=False).compute() + @staticmethod + def chunks(ds: xr.Dataset) -> int | Mapping[str, int]: + # Split the vertical levels across two chunks. + num_levels = ds["level"].size + + return { + "latitude": -1, + "longitude": -1, + "time": -1, + "level": (num_levels // 2) + 1, + } + @staticmethod def open(download_path: Path) -> xr.Dataset: ds = xr.open_zarr(download_path / "ifs_humidity.zarr").drop_encoding() - num_levels = ds["level"].size - ds = ds.isel(time=slice(0, 1)).chunk( - { - "latitude": -1, - "longitude": -1, - "time": -1, - "level": (num_levels // 2) + 1, - } - ) + ds = ds.isel(time=slice(0, 1)) + ds = ds.chunk(IFSHumidityDataset.chunks(ds)) # Needed to make the dataset CF-compliant. ds.longitude.attrs["axis"] = "X" @@ -63,16 +65,4 @@ def open(download_path: Path) -> xr.Dataset: if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset( - IFSHumidityDataset, basepath=args.basepath - ) - open_downloaded_tiny_canonicalized_dataset( - IFSHumidityDataset, basepath=args.basepath - ) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(IFSHumidityDataset) diff --git a/src/climatebenchpress/data_loader/datasets/ifs_uncompressed.py b/src/climatebenchpress/data_loader/datasets/ifs_uncompressed.py index 9e60e50..bc7f8a3 100644 --- a/src/climatebenchpress/data_loader/datasets/ifs_uncompressed.py +++ b/src/climatebenchpress/data_loader/datasets/ifs_uncompressed.py @@ -1,6 +1,5 @@ __all__ = ["IFSUncompressedDataset"] -import argparse from pathlib import Path import earthkit.regrid @@ -8,11 +7,8 @@ import requests import xarray as xr -from .. import ( - monitor, - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from .. import monitor +from ..cli import main from .abc import Dataset BASE_URL = "https://object-store.os-api.cci1.ecmwf.int/esiwacebucket" @@ -43,11 +39,8 @@ def download(download_path: Path, progress: bool = True): @staticmethod def open(download_path: Path) -> xr.Dataset: - ds = ( - xr.open_dataset(download_path / "ifs_uncompressed.zarr") - .drop_encoding() - .chunk(-1) - ) + ds = xr.open_dataset(download_path / "ifs_uncompressed.zarr").drop_encoding() + ds = ds.chunk(IFSUncompressedDataset.chunks(ds)) # Needed to make the dataset CF-compliant. ds.longitude.attrs["axis"] = "X" @@ -179,16 +172,4 @@ def regrid_to_regular(ds, in_grid, out_grid): if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset( - IFSUncompressedDataset, basepath=args.basepath - ) - open_downloaded_tiny_canonicalized_dataset( - IFSUncompressedDataset, basepath=args.basepath - ) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(IFSUncompressedDataset) diff --git a/src/climatebenchpress/data_loader/datasets/nextgems.py b/src/climatebenchpress/data_loader/datasets/nextgems.py index 0b9a522..87b66fd 100644 --- a/src/climatebenchpress/data_loader/datasets/nextgems.py +++ b/src/climatebenchpress/data_loader/datasets/nextgems.py @@ -1,6 +1,5 @@ __all__ = ["NextGemsDataset"] -import argparse from pathlib import Path import healpy @@ -8,11 +7,8 @@ import numpy as np import xarray as xr -from .. import ( - monitor, - open_downloaded_canonicalized_dataset, - open_downloaded_tiny_canonicalized_dataset, -) +from .. import monitor +from ..cli import main from .abc import Dataset NEXTGEMS_CATALOG = "https://data.nextgems-h2020.eu/online.yaml" @@ -76,7 +72,8 @@ def download(download_path: Path, progress: bool = True): @staticmethod def open(download_path: Path) -> xr.Dataset: - return xr.open_zarr(download_path / "download.zarr").drop_encoding().chunk(-1) + ds = xr.open_zarr(download_path / "download.zarr").drop_encoding() + return ds.chunk(NextGemsDataset.chunks(ds)) def _get_nn_lon_lat_index(nside, lons, lats): @@ -104,12 +101,4 @@ def _get_nn_lon_lat_index(nside, lons, lats): if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("--basepath", type=Path, default=Path()) - args = parser.parse_args() - - ds = open_downloaded_canonicalized_dataset(NextGemsDataset, basepath=args.basepath) - open_downloaded_tiny_canonicalized_dataset(NextGemsDataset, basepath=args.basepath) - - for v, da in ds.items(): - print(f"- {v}: {da.dims}") + main(NextGemsDataset) diff --git a/src/climatebenchpress/data_loader/s3.py b/src/climatebenchpress/data_loader/s3.py new file mode 100644 index 0000000..ebce8fb --- /dev/null +++ b/src/climatebenchpress/data_loader/s3.py @@ -0,0 +1,201 @@ +__all__ = [ + "fetch_standardized", + "is_published", + "list_published", + "move_tree", + "remove_tree", +] + +import logging +import shutil +from concurrent.futures import ThreadPoolExecutor +from functools import lru_cache +from pathlib import Path +from typing import Any + +import fsspec +from tqdm import tqdm + +# The object store that hosts the pre-processed ClimateBenchPress datasets. +ENDPOINT_URL = "https://object-store.os-api.cci1.ecmwf.int" +BUCKET_PREFIX = "esiwacebucket/ClimateBenchPress" + +STANDARDIZED = "standardized.zarr" + +# Number of objects that are copied in parallel. The Zarr stores consist of +# many small-ish chunk files, for which the round-trip latency dominates. +_NUM_WORKERS = 8 + +_COPY_BUFFER_SIZE = 4 * 1024 * 1024 + + +@lru_cache +def _filesystem() -> Any: + # The bucket is publicly readable, so no credentials are required. + return fsspec.filesystem("s3", anon=True, endpoint_url=ENDPOINT_URL) + + +def _source_root(name: str) -> str: + return f"{BUCKET_PREFIX}/{name}/{STANDARDIZED}" + + +def list_published() -> list[str]: + """List the names of all datasets that are published in the object store. + + Returns + ------- + list[str] + The sorted dataset names, e.g. ``["cams-nitrogen-dioxide", ...]``. + """ + fs = _filesystem() + + return sorted( + entry.rstrip("/").rsplit("/", 1)[-1] for entry in fs.ls(BUCKET_PREFIX) + ) + + +def is_published(name: str) -> bool: + """Check whether a dataset is published in the object store. + + Parameters + ---------- + name : str + The name of the dataset, e.g. ``"cams-nitrogen-dioxide"``. + + Returns + ------- + bool + Whether the pre-processed dataset is available for download. + """ + return bool(_filesystem().exists(_source_root(name))) + + +def fetch_standardized(name: str, standardized: Path, progress: bool = True) -> bool: + """Download the pre-processed dataset from the ClimateBenchPress object store. + + The Zarr store is copied byte-for-byte so that the chunking, the compressors + and the encoding are exactly as published. + + The store is first assembled in a sibling `.partial` directory and only moved + into place once every object has been copied, such that an interrupted + download never leaves behind a store that looks complete. + + Parameters + ---------- + name : str + The name of the dataset, e.g. ``"cams-nitrogen-dioxide"``. Note that the + tiny variants are published under their own name, e.g. + ``"cams-nitrogen-dioxide-tiny"``. + standardized : Path + The path that the `standardized.zarr` store is downloaded to. + progress : bool, optional + Whether to show a progress bar during the download, by default True + + Returns + ------- + bool + True if the dataset was downloaded, False if it is not published, in + which case nothing has been written. + """ + fs = _filesystem() + + source = _source_root(name) + if not fs.exists(source): + return False + + # Fetch the sizes alongside the listing so that the progress bar does not + # need a HEAD request per object. + sizes = { + key: info.get("size", 0) + for key, info in fs.find(source, detail=True).items() + if not key.endswith("/") + } + + partial = standardized.parent / (standardized.name + ".partial") + if partial.exists(): + remove_tree(partial) + partial.mkdir(parents=True, exist_ok=True) + + logging.debug(f"Downloading {len(sizes)} objects from {source} to {standardized}") + + with tqdm( + total=sum(sizes.values()), + unit="B", + unit_scale=True, + desc=f"{name}/{STANDARDIZED}", + ascii=True, + disable=not progress, + ) as pbar: + + def copy(key: str) -> None: + target = partial + for part in key[len(source) + 1 :].split("/"): + target = target / part + target.parent.mkdir(parents=True, exist_ok=True) + + with fs.open(key, "rb") as fsrc, target.open("wb") as fdst: + shutil.copyfileobj(fsrc, fdst, _COPY_BUFFER_SIZE) + + pbar.update(sizes[key]) + + with ThreadPoolExecutor(max_workers=_NUM_WORKERS) as pool: + # Consume the iterator so that any exception is re-raised here. + for _ in pool.map(copy, sizes): + pass + + move_tree(partial, standardized) + + return True + + +def move_tree(src: Path, dst: Path) -> None: + """Move a directory and everything inside it. + + Falls back to a recursive copy for the filesystems that `basepath` may point + at which cannot rename a non-empty directory, e.g. fsspec's in-memory + filesystem. Such a rename fails without moving anything, so the fallback + always starts from an intact `src`. + + Parameters + ---------- + src : Path + The directory to move + dst : Path + The path to move it to, which must not exist yet + """ + try: + src.rename(dst) + return + except OSError: + pass + + dst.mkdir(parents=True, exist_ok=True) + + for child in src.iterdir(): + if child.is_dir(): + move_tree(child, dst / child.name) + else: + with child.open("rb") as fsrc, (dst / child.name).open("wb") as fdst: + shutil.copyfileobj(fsrc, fdst, _COPY_BUFFER_SIZE) + child.unlink() + + src.rmdir() + + +def remove_tree(path: Path) -> None: + """Recursively remove a directory and everything inside it. + + Unlike `shutil.rmtree`, this only uses the `pathlib.Path` interface and thus + also works for the non-local paths that `basepath` may point at. + + Parameters + ---------- + path : Path + The directory to remove + """ + for child in path.iterdir(): + if child.is_dir(): + remove_tree(child) + else: + child.unlink() + path.rmdir() diff --git a/tests/test_virtual.py b/tests/test_virtual.py index c3cc59b..613133b 100644 --- a/tests/test_virtual.py +++ b/tests/test_virtual.py @@ -1,20 +1,29 @@ +import logging from pathlib import Path import climatebenchpress.data_loader import climatebenchpress.data_loader.datasets.abc import climatebenchpress.data_loader.monitor import fsspec +import numpy as np +import pytest import xarray as xr +import zarr from upath import UPath -def test_virtual_download(): +@pytest.fixture +def basepath(request): fs = fsspec.filesystem("memory") - basepath = UPath(fs.unstrip_protocol("root"), fs=fs) + # Each test needs its own root, since the memory filesystem is global. + return UPath(fs.unstrip_protocol(request.node.name), fs=fs) + +def test_virtual_download(basepath): ds = climatebenchpress.data_loader.open_downloaded_canonicalized_dataset( - TestDataset, + VirtualDataset, basepath=basepath, + reprocess=True, ) assert ds.t.shape == (1, 1, 1, 2, 2) @@ -22,22 +31,155 @@ def test_virtual_download(): assert (basepath / "datasets" / "test" / "standardized.zarr").exists() -class TestDataset(climatebenchpress.data_loader.datasets.abc.Dataset): +def test_falls_back_to_reprocessing_when_unpublished(basepath, monkeypatch, caplog): + monkeypatch.setattr( + climatebenchpress.data_loader.s3, + "fetch_standardized", + lambda name, standardized, progress=True: False, + ) + + with caplog.at_level(logging.WARNING): + ds = climatebenchpress.data_loader.open_downloaded_canonicalized_dataset( + VirtualDataset, basepath=basepath + ) + + assert ds.t.shape == (1, 1, 1, 2, 2) + assert "not published" in caplog.text + # The dataset had to be reprocessed from the original data source. + assert (basepath / "datasets" / "test" / "download" / "download.zarr").exists() + + +def test_skips_reprocessing_when_published(basepath, monkeypatch): + monkeypatch.setattr( + climatebenchpress.data_loader.s3, + "fetch_standardized", + _publish(chunks=-1, compressor=zarr.storage.default_compressor), + ) + + ds = climatebenchpress.data_loader.open_downloaded_canonicalized_dataset( + VirtualDataset, basepath=basepath + ) + + assert ds.t.shape == (1, 1, 1, 2, 2) + # No raw data was downloaded, the pre-processed dataset was used as-is. + assert not (basepath / "datasets" / "test" / "download").exists() + + +def test_reformats_outdated_encoding(basepath, monkeypatch, caplog): + # Mimics the published IFS datasets, which were uploaded before the + # chunking of the pipeline was fixed. + monkeypatch.setattr( + climatebenchpress.data_loader.s3, + "fetch_standardized", + _publish(chunks=1, compressor=zarr.Blosc(cname="zstd")), + ) + + with caplog.at_level(logging.WARNING): + ds = climatebenchpress.data_loader.open_downloaded_canonicalized_dataset( + VirtualDataset, basepath=basepath + ) + + assert "outdated encoding" in caplog.text + # The raw data was still not needed to repair the dataset. + assert not (basepath / "datasets" / "test" / "download").exists() + + xr.testing.assert_identical(ds, _canonical_dataset()) + for var in ds.variables.values(): + assert var.encoding["compressor"] == zarr.storage.default_compressor + assert ds.t.encoding["chunks"] == (1, 1, 1, 2, 2) + + +def test_reformats_to_the_declared_chunking(basepath, monkeypatch): + monkeypatch.setattr( + climatebenchpress.data_loader.s3, + "fetch_standardized", + _publish(chunks=1, compressor=zarr.storage.default_compressor), + ) + + ds = climatebenchpress.data_loader.open_downloaded_canonicalized_dataset( + ChunkedVirtualDataset, basepath=basepath + ) + + # The lat dimension is split across two chunks, the others are not. + assert ds.t.encoding["chunks"] == (1, 1, 1, 1, 2) + + +def test_verify_encoding(): + ds = _canonical_dataset() + ds.t.encoding.update( + chunks=(1, 1, 1, 2, 2), compressor=zarr.storage.default_compressor + ) + for coord in ds.coords.values(): + coord.encoding.update(compressor=zarr.storage.default_compressor) + + assert climatebenchpress.data_loader._verify_encoding(ds, -1) == [] + + # A wrong chunking of a data variable is reported ... + (problem,) = climatebenchpress.data_loader._verify_encoding(ds, 1) + assert problem == "t is chunked (1, 1, 1, 2, 2) instead of (1, 1, 1, 1, 1)" + + # ... and so is a wrong compressor on a coordinate alone, which is how the + # published ifs-humidity dataset is out of date. + ds.lat.encoding.update(compressor=zarr.Blosc(cname="zstd")) + (problem,) = climatebenchpress.data_loader._verify_encoding(ds, -1) + assert problem.startswith("lat is compressed with Blosc(cname='zstd'") + + +def test_verify_encoding_allows_oversized_chunks(): + # The published CMIP6 datasets inherited the chunking of the much longer + # upstream time series. Such a chunk still stores the array as a single + # chunk, so it must not trigger a needless reformat. + ds = _canonical_dataset() + ds.t.encoding.update( + chunks=(1, 1, 1, 3432, 3432), compressor=zarr.storage.default_compressor + ) + for coord in ds.coords.values(): + coord.encoding.update( + chunks=(3432,), compressor=zarr.storage.default_compressor + ) + + assert climatebenchpress.data_loader._verify_encoding(ds, -1) == [] + + +def _dataset() -> xr.Dataset: + return xr.Dataset( + { + "t": (("lat", "lon"), np.asarray([[1, 2], [3, 4]])), + }, + coords={ + "lat": ("lat", [-45, 45], {"standard_name": "latitude", "axis": "Y"}), + "lon": ("lon", [0, 180], {"standard_name": "longitude", "axis": "X"}), + }, + ) + + +def _canonical_dataset() -> xr.Dataset: + return climatebenchpress.data_loader.canon.canonicalize_dataset(_dataset()) + + +def _publish(chunks, compressor): + """Build a `fetch_standardized` that publishes the dataset with a given encoding.""" + + def fetch_standardized( + name: str, standardized: Path, progress: bool = True + ) -> bool: + ds = _canonical_dataset().chunk(chunks) + ds.to_zarr( + standardized, + encoding={v: dict(compressor=compressor) for v in ds.variables}, + ) + return True + + return fetch_standardized + + +class VirtualDataset(climatebenchpress.data_loader.datasets.abc.Dataset): name = "test" @staticmethod def download(download_path: Path, progress: bool = True): - ds = xr.Dataset( - { - "t": (("lat", "lon"), [[1, 2], [3, 4]]), - }, - coords={ - "lat": ("lat", [-45, 45], {"standard_name": "latitude", "axis": "Y"}), - "lon": ("lon", [0, 180], {"standard_name": "longitude", "axis": "X"}), - }, - ) with climatebenchpress.data_loader.monitor.progress_bar(progress): - ds.to_zarr( + _dataset().to_zarr( download_path / "download.zarr", mode="w", compute=False, @@ -46,3 +188,11 @@ def download(download_path: Path, progress: bool = True): @staticmethod def open(download_path: Path) -> xr.Dataset: return xr.open_zarr(download_path / "download.zarr") + + +class ChunkedVirtualDataset(VirtualDataset): + name = "test-chunked" + + @staticmethod + def chunks(ds: xr.Dataset) -> int | dict[str, int]: + return {"lat": 1, "lon": -1}