diff --git a/docs/changes/newsfragments/8356.improved b/docs/changes/newsfragments/8356.improved new file mode 100644 index 000000000000..fc6e0cac9f9f --- /dev/null +++ b/docs/changes/newsfragments/8356.improved @@ -0,0 +1,10 @@ +Datasets exported to NetCDF are now compressed using the deflate (gzip) compression +built into NetCDF4/HDF5. Depending on the data this reduces the size of the exported +files by a factor of roughly 1.3 to 3.3 with no measurable slowdown of the export, and +no loss of precision since the compression is lossless. Compressed files remain regular +NetCDF files that can be read by any NetCDF client without any special handling. + +The compression level can be controlled with the new +``qcodes.config.dataset.export_netcdf_compression_level`` config option. It defaults to +``4`` and can be set to ``0`` to write uncompressed files as before. Higher levels are +significantly slower while giving almost no additional size reduction. diff --git a/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb b/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb index f9175d0da1dc..4c330ef7881b 100644 --- a/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb +++ b/docs/examples/DataSet/Exporting-data-to-other-file-formats.ipynb @@ -178,7 +178,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The dataset can be exported using the `export` method. Currently exporting to netcdf and csv is supported." + "The dataset can be exported using the `export` method. Currently exporting to netcdf and csv is supported.\n", + "\n", + "NetCDF files are written with the deflate (gzip) compression that is built into NetCDF4/HDF5 enabled. This is lossless and typically reduces the file size by a factor of 1.3 to 3 depending on the data. The compression level can be changed with the `qcodes.config.dataset.export_netcdf_compression_level` config option and set to `0` to write uncompressed files.\n" ] }, { diff --git a/src/qcodes/configuration/qcodesrc.json b/src/qcodes/configuration/qcodesrc.json index 583f36bf8326..c501cdcee004 100644 --- a/src/qcodes/configuration/qcodesrc.json +++ b/src/qcodes/configuration/qcodesrc.json @@ -78,6 +78,7 @@ "export_name_elements": ["captured_run_id", "guid"], "export_chunked_export_of_large_files_enabled": false, "export_chunked_threshold": 1000, + "export_netcdf_compression_level": 4, "in_memory_cache": true, "load_from_exported_file": false }, diff --git a/src/qcodes/configuration/qcodesrc_schema.json b/src/qcodes/configuration/qcodesrc_schema.json index ddd4929be22c..bcb41807fa66 100644 --- a/src/qcodes/configuration/qcodesrc_schema.json +++ b/src/qcodes/configuration/qcodesrc_schema.json @@ -373,6 +373,13 @@ "default": 1000, "description": "Estimated size in MB above which the dataset will be exported in chuncks and recombined." }, + "export_netcdf_compression_level": { + "type": "integer", + "minimum": 0, + "maximum": 9, + "default": 4, + "description": "Deflate (gzip) compression level used when exporting datasets to netcdf. Set to 0 to write uncompressed files. Higher levels are slower with almost no additional size reduction, so the default of 4 is recommended." + }, "load_from_exported_file": { "description": "Flag to load metadata and raw data from exported file of type specified in export_type. If set to true, qcodes will try to import from file first, if it exists.", "type": "boolean", diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py index f36ad008772f..5b8fd43eb0e2 100644 --- a/src/qcodes/dataset/data_set.py +++ b/src/qcodes/dataset/data_set.py @@ -1522,6 +1522,9 @@ def _export_as_netcdf(self, path: Path, file_name: str) -> Path: xarray_to_h5netcdf_with_complex_numbers( self.to_xarray_dataset(start=i + 1, end=i + 1), temp_path / file_name_template.format(i), + # these files are temporary and immediately recombined + # into the final file so compressing them only costs time + compression_level=0, ) files = tuple(temp_path.glob("*.nc")) data = xr.open_mfdataset(files) diff --git a/src/qcodes/dataset/exporters/export_to_xarray.py b/src/qcodes/dataset/exporters/export_to_xarray.py index 1734831dd820..ffa5604957e3 100644 --- a/src/qcodes/dataset/exporters/export_to_xarray.py +++ b/src/qcodes/dataset/exporters/export_to_xarray.py @@ -1,14 +1,16 @@ from __future__ import annotations +import itertools import logging import warnings from importlib.metadata import version from math import prod -from typing import TYPE_CHECKING, Literal +from typing import TYPE_CHECKING, Any, Literal import numpy as np from packaging import version as p_version +import qcodes from qcodes.dataset.linked_datasets.links import links_to_str from ..descriptions.versioning import serialization as serial @@ -29,6 +31,107 @@ _LOG = logging.getLogger(__name__) +# Target size in bytes of a single HDF5 chunk when compression is enabled. +# The chunk size matters a lot for how well the data compresses since each chunk +# is compressed independently. The auto chunking performed by h5py picks chunks +# that are significantly smaller than this which results in noticeably worse +# compression ratios for typical QCoDeS data. +_TARGET_CHUNK_SIZE_BYTES = 4 * 1024**2 + + +def _variable_chunksizes( + variable: xr.DataArray, target_bytes: int = _TARGET_CHUNK_SIZE_BYTES +) -> tuple[int, ...] | None: + """ + Calculate the HDF5 chunk shape to use for a given variable. + + The trailing dimensions are kept whole and the leading dimension is sized + such that a chunk is approximately ``target_bytes`` large. Returns None for + scalars and for variables with a zero sized dimension since those cannot be + chunked. + """ + shape = variable.shape + if len(shape) == 0 or 0 in shape: + return None + + bytes_per_leading_element = max(prod(shape[1:]), 1) * variable.dtype.itemsize + leading_chunk = target_bytes // max(bytes_per_leading_element, 1) + leading_chunk = max(1, min(int(leading_chunk), shape[0])) + return (leading_chunk, *shape[1:]) + + +def _netcdf_compression_encoding( + dataset: xr.Dataset, complevel: int +) -> dict[str, dict[str, Any]]: + """ + Build a per variable netcdf encoding dict enabling deflate compression. + + The shuffle filter is only enabled for non complex data. QCoDeS stores + complex numbers as an HDF5 compound type of two floats and byte shuffling + such a compound type destroys the byte patterns that deflate relies on, + roughly halving the compression ratio. + + Args: + dataset: The dataset that is about to be written. + complevel: Deflate compression level between 1 and 9. + + Returns: + A mapping from variable name to netcdf encoding options. Variables that + cannot be compressed, such as variable length strings, are omitted. + + """ + encoding: dict[str, dict[str, Any]] = {} + + for name in itertools.chain(dataset.data_vars, dataset.coords): + variable = dataset[name] + # Compression filters can only be applied to numeric and boolean data. + # Variable length strings and object arrays are written without filters. + if variable.dtype.kind not in "buifc": + continue + + variable_encoding: dict[str, Any] = { + "zlib": True, + "complevel": complevel, + "shuffle": variable.dtype.kind != "c", + } + chunksizes = _variable_chunksizes(variable) + if chunksizes is not None: + variable_encoding["chunksizes"] = chunksizes + encoding[str(name)] = variable_encoding + + return encoding + + +def _rechunk_to_match_encoding( + dataset: xr.Dataset, encoding: Mapping[str, Mapping[str, Any]] +) -> xr.Dataset: + """ + Align the dask chunks of a dataset with the HDF5 chunks it will be written to. + + Writing a dask backed dataset is only performed one dask block at a time. If + a block covers less than a full HDF5 chunk, that chunk has to be read, + decompressed, updated and recompressed again for every block, which is very + slow. Datasets that are not dask backed are returned unmodified. + """ + is_dask_backed = any( + getattr(dataset[name].data, "chunks", None) is not None + for name in dataset.variables + ) + if not is_dask_backed: + return dataset + + dim_chunks: dict[Hashable, int] = {} + for name, variable_encoding in encoding.items(): + chunksizes = variable_encoding.get("chunksizes") + if chunksizes is None: + continue + for dim, chunksize in zip(dataset[name].dims, chunksizes, strict=True): + dim_chunks[dim] = min(dim_chunks.get(dim, chunksize), chunksize) + + if not dim_chunks: + return dataset + return dataset.chunk(dim_chunks) + def _calculate_index_shape(idx: pd.Index | pd.MultiIndex) -> dict[Hashable, int]: # heavily inspired by xarray.core.dataset.from_dataframe @@ -415,11 +518,30 @@ def _paramspec_dict_with_extras( def xarray_to_h5netcdf_with_complex_numbers( - xarray_dataset: xr.Dataset, file_path: str | Path, compute: bool = True + xarray_dataset: xr.Dataset, + file_path: str | Path, + compute: bool = True, + compression_level: int | None = None, ) -> None: + """ + Write an xarray dataset to a netcdf file using the h5netcdf engine. + + Args: + xarray_dataset: The dataset to write. + file_path: Path of the netcdf file to write. + compute: If False the write is returned as a Dask delayed job which is + computed with a progress bar rather than written eagerly. + compression_level: Deflate compression level between 0 and 9 where 0 + means no compression. If None the level is read from + ``qcodes.config.dataset.export_netcdf_compression_level``. + + """ import cf_xarray as cf_xr from pandas import MultiIndex + if compression_level is None: + compression_level = int(qcodes.config.dataset.export_netcdf_compression_level) + has_multi_index = any( isinstance(xarray_dataset.indexes[index_name], MultiIndex) for index_name in xarray_dataset.indexes @@ -453,6 +575,12 @@ def xarray_to_h5netcdf_with_complex_numbers( xarray_too_old or h5netcdf_too_old ) + if compression_level > 0: + encoding = _netcdf_compression_encoding(internal_ds, compression_level) + internal_ds = _rechunk_to_match_encoding(internal_ds, encoding) + else: + encoding = {} + with warnings.catch_warnings(): # see http://xarray.pydata.org/en/stable/howdoi.html # for how to export complex numbers @@ -468,6 +596,7 @@ def xarray_to_h5netcdf_with_complex_numbers( engine="h5netcdf", invalid_netcdf=allow_invalid_netcdf, compute=compute, + encoding=encoding, ) if not compute and maybe_write_job is not None: # Dask and therefor tqdm.dask is slow to diff --git a/tests/dataset/test_dataset_export.py b/tests/dataset/test_dataset_export.py index bf1036e348f3..c3c328686da2 100644 --- a/tests/dataset/test_dataset_export.py +++ b/tests/dataset/test_dataset_export.py @@ -34,7 +34,11 @@ from qcodes.dataset.descriptions.versioning import serialization as serial from qcodes.dataset.export_config import DataExportType from qcodes.dataset.exporters.export_to_pandas import _generate_pandas_index -from qcodes.dataset.exporters.export_to_xarray import _calculate_index_shape +from qcodes.dataset.exporters.export_to_xarray import ( + _TARGET_CHUNK_SIZE_BYTES, + _calculate_index_shape, + _netcdf_compression_encoding, +) from qcodes.dataset.linked_datasets.links import links_to_str from qcodes.parameters import ManualParameter, Parameter, ParamSpecBase @@ -2326,3 +2330,130 @@ def test_incomplete_measurement_with_shared_setpoint( assert "signal_1d" in xr_ds.data_vars assert "signal_2d" in xr_ds.data_vars assert "x" in xr_ds.coords + + +def test_netcdf_compression_encoding_helper() -> None: + """Compression is applied per variable with shuffle disabled for complex data.""" + ds = xr.Dataset( + { + "real": (("x", "y"), np.zeros((100, 7))), + "complex": (("x", "y"), np.zeros((100, 7), dtype=np.complex128)), + "text": ("x", np.array(["a"] * 100, dtype=object)), + "scalar": ((), 1.0), + }, + coords={"x": np.arange(100.0), "y": np.arange(7.0)}, + ) + encoding = _netcdf_compression_encoding(ds, complevel=4) + + # object/string variables cannot be compressed and are left out + assert "text" not in encoding + assert set(encoding) == {"real", "complex", "scalar", "x", "y"} + + assert encoding["real"] == { + "zlib": True, + "complevel": 4, + "shuffle": True, + "chunksizes": (100, 7), + } + # shuffle destroys the byte layout of the complex compound type + assert encoding["complex"]["shuffle"] is False + assert encoding["complex"]["chunksizes"] == (100, 7) + # scalars cannot be chunked + assert "chunksizes" not in encoding["scalar"] + + +def test_netcdf_compression_encoding_chunks_are_capped_to_target_size() -> None: + """The leading dimension is sized to give chunks of about the target size.""" + n_columns = 128 + ds = xr.Dataset({"z": (("x", "y"), np.zeros((10000, n_columns)))}) + encoding = _netcdf_compression_encoding(ds, complevel=4) + + chunksizes = encoding["z"]["chunksizes"] + assert chunksizes[1] == n_columns + assert chunksizes[0] * n_columns * 8 <= _TARGET_CHUNK_SIZE_BYTES + assert (chunksizes[0] + 1) * n_columns * 8 > _TARGET_CHUNK_SIZE_BYTES + + +@pytest.mark.parametrize("complevel", [0, 4]) +def test_export_netcdf_compression_config( + tmp_path_factory: TempPathFactory, mock_dataset_grid: DataSet, complevel: int +) -> None: + """The compression level from the config is applied to the exported file.""" + h5py = pytest.importorskip("h5py") + tmp_path = tmp_path_factory.mktemp("export_netcdf_compression") + qcodes.config.dataset.export_netcdf_compression_level = complevel + + mock_dataset_grid.export(export_type="netcdf", path=tmp_path, prefix="qcodes_") + file_path = mock_dataset_grid.export_info.export_paths["nc"] + + with h5py.File(file_path, "r") as file: + z = file["z"] + if complevel == 0: + assert z.compression is None + else: + assert z.compression == "gzip" + assert z.compression_opts == complevel + assert z.shuffle is True + assert z.chunks == (10, 5) + + loaded_ds = xr.load_dataset(file_path, engine="h5netcdf") + assert_allclose(loaded_ds.z, mock_dataset_grid.to_xarray_dataset().z) + + +def test_export_netcdf_compression_of_complex_data( + tmp_path_factory: TempPathFactory, mock_dataset_numpy_complex: DataSet +) -> None: + """Complex data is compressed but not shuffled and round trips exactly.""" + h5py = pytest.importorskip("h5py") + tmp_path = tmp_path_factory.mktemp("export_netcdf_compression_complex") + qcodes.config.dataset.export_netcdf_compression_level = 4 + + mock_dataset_numpy_complex.export( + export_type="netcdf", path=tmp_path, prefix="qcodes_" + ) + file_path = mock_dataset_numpy_complex.export_info.export_paths["nc"] + + with h5py.File(file_path, "r") as file: + z = file["z"] + assert z.compression == "gzip" + assert z.shuffle is False + + original = mock_dataset_numpy_complex.to_xarray_dataset() + loaded_ds = xr.load_dataset(file_path, engine="h5netcdf") + assert_array_equal(loaded_ds.z, original.z) + + +def test_export_netcdf_compression_reduces_file_size( + tmp_path_factory: TempPathFactory, experiment: Experiment +) -> None: + """Compression gives a significantly smaller file for realistic data.""" + dataset = new_data_set("compressible") + xparam = ParamSpecBase("x", "numeric") + yparam = ParamSpecBase("y", "numeric") + zparam = ParamSpecBase("z", "numeric") + dataset.set_interdependencies( + InterDependencies_(dependencies={zparam: (xparam, yparam)}) + ) + dataset.mark_started() + y_values = np.linspace(-1, 1, 200) + for i in range(100): + # quantized as a real instrument would return it + z_values = np.round(np.sin(i / 10) * np.exp(-(y_values**2)) * 32767) / 32767 + dataset.add_results( + [{"x": i, "y": y, "z": z} for y, z in zip(y_values, z_values, strict=True)] + ) + dataset.mark_completed() + + sizes = {} + for complevel in (0, 4): + qcodes.config.dataset.export_netcdf_compression_level = complevel + path = tmp_path_factory.mktemp(f"compression_size_{complevel}") + dataset.export(export_type="netcdf", path=path, prefix="qcodes_") + sizes[complevel] = Path(dataset.export_info.export_paths["nc"]).stat().st_size + + assert sizes[4] < sizes[0] / 1.3 + + loaded_ds = xr.load_dataset( + dataset.export_info.export_paths["nc"], engine="h5netcdf" + ) + assert_array_equal(loaded_ds.z, dataset.to_xarray_dataset().z)