|
| 1 | +import logging |
| 2 | +from pathlib import Path |
| 3 | + |
| 4 | +from imas_streams import StreamingIMASMetadata |
| 5 | +from imas_streams.xarray_consumers import StreamingXArrayConsumer |
| 6 | + |
| 7 | +logger = logging.getLogger(__name__) |
| 8 | +try: |
| 9 | + import netCDF4 |
| 10 | + import xarray |
| 11 | +except ImportError: |
| 12 | + logger.error("Optional dependency 'netCDF4' or 'xarray' is not installed.") |
| 13 | + raise |
| 14 | + |
| 15 | + |
| 16 | +class NetCDFConsumer: |
| 17 | + """Consumer of streaming IMAS data which stores the data in a netCDF file. |
| 18 | +
|
| 19 | + Note: writes to the filesystem are batched for better performance. The first data |
| 20 | + will be stored once the first batch is complete. |
| 21 | +
|
| 22 | + Example: |
| 23 | + .. code-block:: python |
| 24 | +
|
| 25 | + # Create metadata (from JSON) |
| 26 | + metadata = StreamingIMASMetadata.model_validate_json(json_metadata) |
| 27 | + # Create reader |
| 28 | + reader = NetCDFConsumer(metadata, filename="output.nc", batch_size=1024) |
| 29 | +
|
| 30 | + # Consume dynamic data |
| 31 | + for dynamic_data in dynamic_data_stream: |
| 32 | + reader.process_message(dynamic_data) |
| 33 | + reader.finalize() |
| 34 | + """ |
| 35 | + |
| 36 | + def __init__( |
| 37 | + self, metadata: StreamingIMASMetadata, *, filename: Path, batch_size: int = 1024 |
| 38 | + ) -> None: |
| 39 | + self._metadata = metadata |
| 40 | + self._groupname = f"{metadata.ids_name}/0" |
| 41 | + self._filename = filename |
| 42 | + self._batch_size = batch_size |
| 43 | + |
| 44 | + # Touch the file so we know that it exists |
| 45 | + try: |
| 46 | + self._filename.touch(exist_ok=False) |
| 47 | + except FileExistsError as exc: |
| 48 | + exc.add_note( |
| 49 | + "NetCDFConsumer will not overwrite existing files. Please rename or " |
| 50 | + f"remove {filename} and try again." |
| 51 | + ) |
| 52 | + raise |
| 53 | + self._netcdf_file = None |
| 54 | + self._time_variables = [] |
| 55 | + |
| 56 | + # Let the xarray consumer handle all buffering and tensorization: |
| 57 | + self._xarray_consumer = StreamingXArrayConsumer(metadata, batch_size=batch_size) |
| 58 | + |
| 59 | + def _store_data(self, ds: xarray.Dataset) -> None: |
| 60 | + """Store data to netCDF file""" |
| 61 | + if self._netcdf_file is None: |
| 62 | + # Check which variables are time-dependent |
| 63 | + for varname in ds.variables: |
| 64 | + if "time" in ds[varname].dims: |
| 65 | + self._time_variables.append(varname) |
| 66 | + |
| 67 | + # Compress dynamic data |
| 68 | + encoding = { |
| 69 | + name: { |
| 70 | + "compression": "zlib", |
| 71 | + "complevel": 1, |
| 72 | + "chunksizes": self._chunksize(ds[name]), |
| 73 | + } |
| 74 | + for name in self._time_variables |
| 75 | + if ds[name].dtype.kind in "if" |
| 76 | + } |
| 77 | + |
| 78 | + # Use xarray API to create the netCDF file |
| 79 | + ds.to_netcdf( |
| 80 | + self._filename, |
| 81 | + "w", |
| 82 | + format="NETCDF4", |
| 83 | + engine="netcdf4", |
| 84 | + group=self._groupname, |
| 85 | + auto_complex=True, |
| 86 | + unlimited_dims=["time"], |
| 87 | + encoding=encoding, |
| 88 | + ) |
| 89 | + self._netcdf_file = netCDF4.Dataset(self._filename, "r+", auto_complex=True) |
| 90 | + self._netcdf_file.set_fill_off() # Don't fill on resize |
| 91 | + # Set mandatory global metadata attributes |
| 92 | + self._netcdf_file.Conventions = "IMAS" |
| 93 | + self._netcdf_file.data_dictionary_version = ( |
| 94 | + self._metadata.data_dictionary_version |
| 95 | + ) |
| 96 | + |
| 97 | + else: |
| 98 | + # Add time slices to netCDF file |
| 99 | + group = self._netcdf_file[self._groupname] |
| 100 | + timeslice = slice(len(group.dimensions["time"]), None) |
| 101 | + allslice = slice(None) |
| 102 | + for varname in self._time_variables: |
| 103 | + xrvar = ds[varname] |
| 104 | + index = tuple( |
| 105 | + timeslice if dim == "time" else allslice for dim in xrvar.dims |
| 106 | + ) |
| 107 | + group[varname][index] = xrvar.data |
| 108 | + |
| 109 | + def _chunksize(self, array: xarray.DataArray) -> tuple[int, ...]: |
| 110 | + """Heuristic to determine chunksizes for storing the provided data array""" |
| 111 | + if array.nbytes > 4_194_304: # 4 MB |
| 112 | + logger.warning( |
| 113 | + "Huge chunk detected due to batchsize for %s: %s (=%d bytes)", |
| 114 | + array.name, |
| 115 | + array.shape, |
| 116 | + array.nbytes, |
| 117 | + ) |
| 118 | + return array.shape |
| 119 | + |
| 120 | + def process_message(self, data: bytes | bytearray) -> None: |
| 121 | + """Process a dynamic data message and store it (when a batch is full).""" |
| 122 | + result = self._xarray_consumer.process_message(data) |
| 123 | + if result is not None: |
| 124 | + self._store_data(result) |
| 125 | + |
| 126 | + def finalize(self) -> None: |
| 127 | + """Indicate that the final message is received and store any remaining data.""" |
| 128 | + result = self._xarray_consumer.finalize() |
| 129 | + if result is not None: |
| 130 | + self._store_data(result) |
| 131 | + if self._netcdf_file is not None: |
| 132 | + self._netcdf_file.close() |
0 commit comments