import os
from logging import info, warning
import numpy as np
import xarray as xr
from .. import unstructured_NetCDF as unstructured
[docs]
def write(self, name, flx_file, flx, mode="a", metadata=None, **kwargs):
domain = metadata.get('domain', None)
if domain is None:
raise KeyError("'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)
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}'")
ds.to_netcdf(flx_file, mode='w')
else:
info(f"appending gridded NetCDF flux '{name}' fluxes to '{flx_file}'")
ds.to_netcdf(flx_file, mode='a')
