import os
import numpy as np
import xarray as xr
from .....utils.check.errclass import CifKeyError, CifTypeError, CifValueError
from .....utils.hdf5 import _hdf5_lock
[docs]
def write(self, name, path, data, metadata=None, **kwargs) -> None:
"""Write an LMDz chemical field to a Fortran binary or NetCDF file.
The output format is selected from ``path``'s extension: ``.bin``
writes a Fortran binary file from a ``Dataset`` with ``fwd`` and ``tl``
(forward and tangent-linear) variables, transposed and dumped via
``numpy.ndarray.tofile``; ``.nc`` writes (or appends to) a NetCDF file,
dropping the cyclic-closure longitude, building ``lat``/``lon``
coordinates (and bounds) from ``metadata['domain']`` if given, and
renaming dimensions to LMDZ's ``time_counter``/``presnivs`` convention.
Args:
self: This plugin; ``self.plugin`` is used for error messages.
name: Name of the variable to write (used only for ``.nc`` output).
path: Path to the file to write; its extension (``.bin`` or
``.nc``) selects the output format.
data: Data to write: an ``xarray.Dataset`` with ``fwd``/``tl``
variables for ``.bin`` output, or an ``xarray.DataArray`` for
``.nc`` output.
metadata: Optional dict with a ``"domain"`` entry, used to build
the ``lat``/``lon`` coordinates when writing NetCDF.
**kwargs: Unused.
Raises:
CifTypeError: If ``data`` is not of the type expected for the
target extension.
CifKeyError: If writing ``.bin`` and ``data`` lacks a ``fwd`` or
``tl`` variable.
CifValueError: If ``path``'s extension is neither ``.bin`` nor
``.nc``.
"""
plg = self.plugin
_, ext = os.path.splitext(path)
if ext == '.bin':
if not isinstance(data, xr.Dataset):
raise CifTypeError(
f"Plugin {plg.name} / {plg.version} / {plg.type} "
"write method 'data' argument should be a xarray.Dataset when "
"writing a '.bin' file. "
f"Passed 'data' argument type: {type(data)}"
)
elif 'fwd' not in data or 'tl' not in data:
raise CifKeyError(
f"Plugin {plg.name} / {plg.version} / {plg.type} "
"write method 'data' argument should contain a 'fwd' and a "
"'tl' variable when writing a '.bin' file."
)
prescr_fwd = data['fwd'].values
prescr_tl = data['tl'].values
out = np.transpose([prescr_fwd, prescr_tl], axes=(0, 4, 3, 2, 1)).T
out.tofile(path)
elif ext == '.nc':
if not isinstance(data, xr.DataArray):
raise CifTypeError(
f"Plugin {plg.name} / {plg.version} / {plg.type} "
"write method 'data' argument should be a xarray.DataArray when "
"writing a '.nc' file. "
f"Passed 'data' argument type: {type(data)}"
)
# Droping looping lon coordinate
data = data.isel(lon=slice(None, -1))
coords = {}
if 'time' in data.coords:
coords.update({
'time_counter': (['time_counter'], data.time.values, {
'standard_name': "time",
'long_name': "Time axis",
'axis': "T",
})
})
# Reading domain if it exists
if metadata is not None and 'domain' in metadata:
domain = metadata['domain']
lat = domain.zlat[:, 0]
lat_bnds = domain.zlatc[:, 0]
lat_bnds = np.concatenate([lat_bnds[:-1, np.newaxis],
lat_bnds[1:, np.newaxis]], axis=1)
lon = domain.zlon[0, :-1]
lon_bnds = domain.zlonc[0, :-1]
lon_bnds = np.concatenate([lon_bnds[:-1, np.newaxis],
lon_bnds[1:, np.newaxis]], axis=1)
coords.update({
'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"
}),
})
# Formating dataset
data = data.rename({'time': "time_counter", 'lev': "presnivs"})
ds = xr.Dataset({name: (data.dims, data.values)}, coords=coords)
with _hdf5_lock:
if not os.path.exists(path):
ds.to_netcdf(path, mode='w')
else:
ds.to_netcdf(path, mode='a')
else:
raise CifValueError(f"Unsupported extension '{ext}', supported extensions "
"are '.bin' (Fortran binary) and .'nc' (NetCDF)")
