From 1a93b79b79ae459852f769f094048557ab0bfe81 Mon Sep 17 00:00:00 2001 From: Fabian Neumann Date: Wed, 7 Oct 2026 14:21:32 +0200 Subject: [PATCH 1/2] switch default cutout compression to zstd with BitRound quantization Default compression in Cutout.prepare is now zstd (complevel 3) after BitRound quantization to 14 mantissa bits (max relative error 2**-15). The encoding applies to all variables, because xarray < 2026.9 drops the zstd setting of existing variables when the cutout file is rewritten. Require netCDF4>=1.7.4, whose wheels bundle the zstd plugin. Closes #508 --- RELEASE_NOTES.rst | 12 +++++ atlite/data.py | 29 +++++++---- examples/create_cutout.ipynb | 9 ++-- examples/large_cutouts_with_era5_edh.ipynb | 11 +--- pyproject.toml | 2 +- test/test_compression.py | 60 ++++++++++++++++++++++ 6 files changed, 100 insertions(+), 23 deletions(-) create mode 100644 test/test_compression.py diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index acbb91bd..34223339 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -17,6 +17,18 @@ Upcoming Release To use the features already you have to install the ``master`` branch, e.g. ``pip install git+https://github.com/pypsa/atlite``. +**Breaking** + +* ``Cutout.prepare`` now compresses cutouts with zstd after BitRound + quantization to 14 mantissa bits (relative error below 3.1e-5) instead of + lossless zlib. Writing is about 10 times faster and files are about half the + size. The setting now applies to all variables of the cutout, also to + variables prepared earlier. Reading the new cutouts requires ``netCDF4>=1.7.4`` + (pip) or a ``libnetcdf`` with the zstd plugin (conda-forge). To keep the old + behaviour, pass ``compression={"compression": "zlib", "complevel": 9, + "shuffle": True}`` (https://github.com/PyPSA/atlite/issues/508). +* Raise the minimum version of ``netCDF4`` to 1.7.4. + `v0.7.0 `__ (3rd October 2026) ======================================================================================= diff --git a/atlite/data.py b/atlite/data.py index 2064a8dd..7f853092 100644 --- a/atlite/data.py +++ b/atlite/data.py @@ -222,12 +222,18 @@ def cutout_prepare( Whether to overwrite variables which are already included in the cutout. The default is False. compression : None/dict, optional - Compression level to use for all features which are being prepared. - The compression is handled via xarray.Dataset.to_netcdf(...), for details see: + Variable encoding for compression, applied to all variables of the + cutout when the file is written. The compression is handled via + xarray.Dataset.to_netcdf(...), for details see: https://docs.xarray.dev/en/stable/generated/xarray.Dataset.to_netcdf.html . - To efficiently reduce cutout sizes, specify the number of 'least_significant_digits': n here. - To disable compression, set "complevel" to None. - Default is {'zlib': True, 'complevel': 9, 'shuffle': True}. + The default uses lossless zstd compression after BitRound + quantization to 14 mantissa bits (relative error below 2**-15 = 3.1e-5), + see Klöwer et al. (2021), https://doi.org/10.1038/s43588-021-00156-2 . + For lossless compression, remove 'quantize_mode' and + 'significant_digits'. For the zlib compression of atlite < v0.8, use + {'compression': 'zlib', 'complevel': 9, 'shuffle': True}. + Default is {'compression': 'zstd', 'complevel': 3, + 'quantize_mode': 'BitRound', 'significant_digits': 14}. show_progress : bool, optional If True, a progress bar is shown. The default is False. dask_kwargs : dict, default {} @@ -258,7 +264,12 @@ def cutout_prepare( dask_kwargs = {} if compression is None: - compression = {"zlib": True, "complevel": 9, "shuffle": True} + compression = { + "compression": "zstd", + "complevel": 3, + "quantize_mode": "BitRound", + "significant_digits": 14, + } if cutout.prepared and not overwrite: logger.info("Cutout already prepared.") @@ -302,12 +313,12 @@ def cutout_prepare( attrs = non_bool_dict(cutout.data.attrs) attrs.update(ds.attrs) + ds = cutout.data.merge(ds[missing_vars.values]).assign_attrs(**attrs) + if compression: - for v in missing_vars: + for v in ds.data_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) diff --git a/examples/create_cutout.ipynb b/examples/create_cutout.ipynb index d5c5925e..5a59a923 100644 --- a/examples/create_cutout.ipynb +++ b/examples/create_cutout.ipynb @@ -1475,18 +1475,19 @@ "## Reducing Cutout file sizes\n", "\n", "Cutouts can become quite large, depending on the spatial and temporal scope they cover.\n", - "By default `atlite` uses a trade-off between speed and compression to reduce the file size of cutouts.\n", + "By default `atlite` compresses cutouts with `zstd` after `BitRound` quantization to 14 mantissa bits (relative error below 3.1e-5).\n", + "This is fast and gives about half the file size of lossless `zlib` compression.\n", "\n", - "Stronger compression can be selected when creating a new cutout by choosing a higher `complevel` (`1` to `9`, default: `4`)\n", + "For lossless compression, remove the quantization:\n", "```\n", - "cutout.prepare(compression={\"zlib\": True, \"complevel\": 9})\n", + "cutout.prepare(compression={\"compression\": \"zstd\", \"complevel\": 3})\n", "```\n", "\n", "To change the compression for an existing cutout:\n", "```\n", "cutout = atlite.Cutout(\"cutout-path.nc\")\n", "\n", - "compression = {\"zlib\": True, \"complevel\": 9}\n", + "compression = {\"compression\": \"zlib\", \"complevel\": 9, \"shuffle\": True}\n", "for var in cutout.data.data_vars:\n", " cutout.data[var].encoding.update(compression)\n", "\n", diff --git a/examples/large_cutouts_with_era5_edh.ipynb b/examples/large_cutouts_with_era5_edh.ipynb index 8736a02b..52d32c33 100644 --- a/examples/large_cutouts_with_era5_edh.ipynb +++ b/examples/large_cutouts_with_era5_edh.ipynb @@ -116,7 +116,7 @@ "id": "28c5e310-938e-49de-9ad2-68acef9e566b", "metadata": {}, "source": [ - "Now let's set up the cutout preparation. This is where the data will actually be downloaded and saved to disk. We pass a compression config to `cutout.prepare()`. Here we use the `zstd` compression algorithm (which is modern and fast) and the `GranularBitRound` filter. This filter will save only the first 5 significant digits of each value we generate, which generates a relative error of 1.8e-5 and saves a ton of drive space. To learn more about this filter and it's precision tradeoffs, look [here](https://docs.unidata.ucar.edu/netcdf-c/4.9.0/md_quantize.html).\n", + "Now let's set up the cutout preparation. This is where the data will actually be downloaded and saved to disk. By default, `cutout.prepare()` uses the `zstd` compression algorithm (which is modern and fast) and the `BitRound` filter. This filter keeps only the first 14 mantissa bits of each value we generate, which gives a relative error below 3.1e-5 and saves a ton of drive space. To learn more about this filter and its precision tradeoffs, look [here](https://docs.unidata.ucar.edu/netcdf-c/4.9.0/md_quantize.html). To change the compression, pass a `compression` dict to `cutout.prepare()`.\n", "\n", "We also use a memory sampler to demonstrate the type of statistics that you can get with a `distributed` scheduler. If you chose to install the `bokeh` library, after you run the next cell you should be able to open a [dashboard](https://docs.dask.org/en/stable/dashboard.html) under [http://localhost:8787/](http://localhost:8787/) where you can monitor the computation." ] @@ -130,17 +130,10 @@ }, "outputs": [], "source": [ - "COMPRESSION = {\n", - " \"compression\": \"zstd\",\n", - " \"complevel\": 3,\n", - " \"significant_digits\": 5,\n", - " \"quantize_mode\": \"GranularBitRound\",\n", - "}\n", - "\n", "ms = MemorySampler()\n", "\n", "with ms.sample(\"prepare cutout\", interval=5):\n", - " cutout.prepare(compression=COMPRESSION)" + " cutout.prepare()" ] }, { diff --git a/pyproject.toml b/pyproject.toml index e176e0a2..54b88adb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ dependencies = [ "bottleneck", "numexpr", "xarray>=2024.03.0", - "netcdf4", + "netcdf4>=1.7.4", "dask>=2021.10.0", "toolz", "requests", diff --git a/test/test_compression.py b/test/test_compression.py new file mode 100644 index 00000000..1bcd0759 --- /dev/null +++ b/test/test_compression.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Contributors to atlite +# +# SPDX-License-Identifier: MIT + +"""Tests for the compression of prepared cutouts (#508).""" + +from pathlib import Path +from typing import Any, Literal + +import netCDF4 +import numpy as np +import pytest +import xarray as xr + +import atlite.data +from atlite import Cutout +from atlite.data import available_features + +ZLIB = {"compression": "zlib", "complevel": 9, "shuffle": True} + + +def _fake_features(cutout, module, features, **kwargs): + variables = available_features(module).loc[module].loc[list(features)] + shape = tuple(cutout.data.sizes[d] for d in ("time", "y", "x")) + values = np.random.default_rng(0).uniform(250, 300, shape) + values[0, 0, 0] = np.nan + data = { + v: xr.Variable(("time", "y", "x"), values, {"module": module, "feature": f}) + for f, v in variables.items() + } + return xr.Dataset(data, coords=cutout.data.coords) + + +@pytest.mark.parametrize( + ("compression", "codec", "quantization"), + [(None, "zstd", (14, "BitRound")), (ZLIB, "zlib", None)], +) +def test_prepare_compresses_all_variables( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + compression: dict[str, Any] | None, + codec: Literal["zstd", "zlib"], + quantization: tuple[int, str] | None, +) -> None: + monkeypatch.setattr(atlite.data, "get_features", _fake_features) + cutout = Cutout( + tmp_path / "cutout.nc", module="era5", bounds=(0, 50, 1, 51), time="2013-01-01" + ) + for feature in ["height", "temperature"]: + cutout.prepare(feature, tmpdir=tmp_path, compression=compression) + + with netCDF4.Dataset(cutout.path) as nc: + assert nc.variables["height"].filters()[codec] + assert nc.variables["temperature"].filters()[codec] + assert nc.variables["temperature"].quantization() == quantization + + expected = _fake_features(cutout, "era5", ["temperature"]) + xr.testing.assert_allclose( + cutout.data["temperature"], expected["temperature"], rtol=2**-15 + ) From ce3361ced02dca29737a7933a72af83a81a0d28b Mon Sep 17 00:00:00 2001 From: Fabian Neumann Date: Wed, 7 Oct 2026 14:22:09 +0200 Subject: [PATCH 2/2] add PR link to release notes --- RELEASE_NOTES.rst | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/RELEASE_NOTES.rst b/RELEASE_NOTES.rst index 34223339..499a781a 100644 --- a/RELEASE_NOTES.rst +++ b/RELEASE_NOTES.rst @@ -26,8 +26,9 @@ Upcoming Release variables prepared earlier. Reading the new cutouts requires ``netCDF4>=1.7.4`` (pip) or a ``libnetcdf`` with the zstd plugin (conda-forge). To keep the old behaviour, pass ``compression={"compression": "zlib", "complevel": 9, - "shuffle": True}`` (https://github.com/PyPSA/atlite/issues/508). -* Raise the minimum version of ``netCDF4`` to 1.7.4. + "shuffle": True}`` (https://github.com/PyPSA/atlite/pull/529). +* Raise the minimum version of ``netCDF4`` to 1.7.4 + (https://github.com/PyPSA/atlite/pull/529). `v0.7.0 `__ (3rd October 2026) =======================================================================================