import os
from logging import info, warning
import numpy as np
import xarray as xr
from .. import unstructured_NetCDF as unstructured
from .....utils.check.errclass import CifKeyError
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 gridded NetCDF file.
Builds a dataset with ``time``/``lat``/``lon`` coordinates (plus
``lat_bnds``/``lon_bnds``) from the domain in ``metadata``, and writes it
to `flx_file` (creating it, or appending if it already exists). If the
domain is unstructured, delegates to the ``unstructured_NetCDF`` plugin's
`write` function instead.
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; squeezed along ``lev`` if it has a single level.
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 contain a ``'domain'`` key with the
`Domain` object used to derive lat/lon coordinates and bounds.
Raises:
CifKeyError: If `metadata` has no ``'domain'`` key.
"""
domain = metadata.get('domain', None)
if domain is None:
raise CifKeyError("'metadata' has no 'domain' key")
if getattr(domain, "unstructured_domain", False):
warning(
"Calling flux/gridded_netcdf/std plugin write method on with an unstructured "
"domain, falling back to flux/gridded_netcdf/unstructured plugin write method."
)
unstructured.write(self, name, flx_file, flx, mode, metadata, **kwargs)
return
lat = domain.zlat[:, 0]
lon = domain.zlon[0, :]
lat_corners = domain.zlatc[:, 0]
lon_corners = domain.zlonc[0, :]
lat_bnds = np.concatenate([lat_corners[:-1, np.newaxis],
lat_corners[1:, np.newaxis]], axis=1)
lon_bnds = np.concatenate([lon_corners[:-1, np.newaxis],
lon_corners[1:, np.newaxis]], axis=1)
if flx.sizes["lev"] == 1:
flx = flx.squeeze('lev')
ds = xr.Dataset(
{name: (flx.dims, flx.data)},
coords={
'time': (['time'], flx.time.data, {
'standard_name': "time",
'long_name': "time",
'axis': "T",
}),
'lat': (['lat'], lat, {
'standard_name': "latitude",
'long_name': "latitude",
'units': "degrees_north",
'axis': "Y",
'bounds': "lat_bnds"
}),
'lon': (['lon'], lon, {
'standard_name': "longitude",
'long_name': "longitude",
'units': "degrees_east",
'axis': "X",
'bounds': "lon_bnds"
}),
'lat_bnds': (['lat', 'bnds'], lat_bnds, {
'standard_name': "latitude_bounds",
'long_name': "latitude bounds",
'units': "degrees_north"
}),
'lon_bnds': (['lon', 'bnds'], lon_bnds, {
'standard_name': "longitude_bounds",
'long_name': "longitude bounds",
'units': "degrees_east"
})
}
)
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')
