import os
from logging import debug, info
import xarray as xr
from .....utils.check.errclass import CifKeyError, CifTypeError
from .....utils.hdf5 import _hdf5_lock
[docs]
def write(self, name, flx_file, flx, mode="a", metadata=None, **kwargs):
"""Write a flux DataArray to a CF-style unstructured NetCDF file.
Builds a dataset with a ``cell`` dimension (plus ``time``) and
``lat``/``lon``/``lat_bnds``/``lon_bnds`` variables from the domain's
cell coordinates and vertices, and writes it to `flx_file` (creating
it, or appending if it already exists).
Args:
self: The flux tracer plugin instance.
name (str): Name of the flux variable to write.
flx_file (str): Path of the NetCDF file to write or append to.
flx (xr.DataArray): Flux data with ``time``, ``lev``, ``lat``,
``lon`` dimensions; ``lat`` and ``lev`` are squeezed and
``lon`` is renamed to ``cell``.
mode (str): Unused directly; existence of `flx_file` determines
whether the dataset is written (``mode='w'``) or appended
(``mode='a'``) to the NetCDF file.
metadata (dict, optional): Must be a dict containing a ``'domain'``
key with the `Domain` object used to derive cell coordinates
and vertices.
Raises:
CifTypeError: If `metadata` is not a dict.
CifKeyError: If `metadata` has no ``'domain'`` key.
"""
if not isinstance(metadata, dict):
raise CifTypeError("'metadata' must be a dict")
domain = metadata.get('domain', None)
if domain is None:
raise CifKeyError("'metadata' has no 'domain' key")
lat = domain.zlat[0, :]
lon = domain.zlon[0, :]
lat_vertices = domain.zlatc.T
lon_vertices = domain.zlonc.T
flx = flx.squeeze(["lat", "lev"])
flx = flx.rename({"lon": "cell"})
ds = xr.Dataset(
{name: (flx.dims, flx.data)},
# fmt: off
coords={
"time": (["time"], flx.time.data, {
"standard_name": "time",
"long_name": "time",
"axis": "T",
}),
"lat": (["cell"], lat, {
"standard_name": "latitude",
"long_name": "latitude",
"units": "degrees_north",
"bounds": "lat_bnds"
}),
"lon": (["cell"], lon, {
"standard_name": "longitude",
"long_name": "longitude",
"units": "degrees_east",
"bounds": "lon_bnds"
}),
"lat_bnds": (["cell", "vertex"], lat_vertices, {
"standard_name": "latitude_bounds",
"long_name": "latitude bounds",
"units": "degrees_north"
}),
"lon_bnds": (["cell", "vertex"], lon_vertices, {
"standard_name": "longitude_bounds",
"long_name": "longitude bounds",
"units": "degrees_east"
})
}
# fmt: on
)
if not os.path.exists(flx_file):
info(f"writing gridded NetCDF flux '{name}' fluxes to '{flx_file}'")
with _hdf5_lock:
ds.to_netcdf(flx_file, mode='w')
else:
info(f"appending gridded NetCDF flux '{name}' fluxes to '{flx_file}'")
with _hdf5_lock:
ds.to_netcdf(flx_file, mode='a')
