import os
import numpy as np
import xarray as xr
[docs]
def write(self, name, path, data, metadata=None, **kwargs) -> None:
"""Write prescribed species files for LMDZ
Args:
self: this plugin
path (str): path to the file to write
data (xarray.Dataset or xarray.DataArray): Data to write
"""
plg = self.plugin
_, ext = os.path.splitext(path)
if ext == '.bin':
if not isinstance(data, xr.Dataset):
raise TypeError(
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 KeyError(
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 TypeError(
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)
if not os.path.exists(path):
ds.to_netcdf(path, mode='w')
else:
ds.to_netcdf(path, mode='a')
else:
raise ValueError(f"Unsupported extension '{ext}', supported extensions "
"are '.bin' (Fortran binary) and .'nc' (NetCDF)")
