def test_unstagger(self):
wf = get_demo_file('wrf_cropped.nc')
with netCDF4.Dataset(wf) as nc:
nc.set_auto_mask(False)
ref = nc['PH'][:]
ref = 0.5 * (ref[:, :-1, ...] + ref[:, 1:, ...])
# Own constructor
v = wrftools.Unstaggerer(nc['PH'])
assert_allclose(v[:], ref)
assert_allclose(v[0:2, 2:12, ...], ref[0:2, 2:12, ...])
assert_allclose(v[:, 2:12, ...], ref[:, 2:12, ...])
assert_allclose(v[0:2, 2:12, 5:10, 15:17], ref[0:2, 2:12, 5:10,
15:17])
assert_allclose(v[1:2, 2:, 5:10, 15:17], ref[1:2, 2:, 5:10, 15:17])
assert_allclose(v[1:2, :-2, 5:10, 15:17], ref[1:2, :-2, 5:10,
15:17])
assert_allclose(v[1:2, 2:-4, 5:10, 15:17], ref[1:2, 2:-4, 5:10,
15:17])
assert_allclose(v[[0, 2], ...], ref[[0, 2], ...])
assert_allclose(v[..., [0, 2]], ref[..., [0, 2]])
assert_allclose(v[0, ...], ref[0, ...])
# Under WRF
nc = WRF(wf)
assert_allclose(nc.get_vardata('PH'), ref)
nc.set_period(1, 2)
assert_allclose(nc.get_vardata('PH'), ref[1:3, ...])
def __init__(self, file, grid=None, time=None):
GeoNetcdf.__init__(self, file, grid=grid, time=time)
# Change staggered variables to unstaggered ones
for vn, v in self.variables.items():
if wrftools.Unstaggerer.can_do(v):
self.variables[vn] = wrftools.Unstaggerer(v)
# Check if we can add diagnostic variables to the pot
for vn in wrftools.var_classes:
cl = getattr(wrftools, vn)
if cl.can_do(self._nc):
self.variables[vn] = cl(self._nc)
def open_wrf_dataset(file, **kwargs):
"""Wrapper around xarray's open_dataset to make WRF files a bit better.
This is needed because variables often have not enough georef attrs
to be understood alone, and datasets tend to loose their attrs with
operations...
Parameters
----------
file : str
the path to the WRF file
**kwargs : optional
Additional arguments passed on to ``xarray.open_dataset``.
Returns
-------
an xarray Dataset
"""
nc = netCDF4.Dataset(file)
nc.set_auto_mask(False)
# Change staggered variables to unstaggered ones
for vn, v in nc.variables.items():
if wrftools.Unstaggerer.can_do(v):
nc.variables[vn] = wrftools.Unstaggerer(v)
# Check if we can add diagnostic variables to the pot
for vn in wrftools.var_classes:
cl = getattr(wrftools, vn)
if vn not in nc.variables and cl.can_do(nc):
nc.variables[vn] = cl(nc)
# trick xarray with our custom netcdf
ds = xr.open_dataset(NetCDF4DataStore(nc), **kwargs)
# remove time dimension to lon lat
for vn in ['XLONG', 'XLAT']:
try:
v = ds[vn].isel(Time=0)
ds[vn] = xr.DataArray(v.values, dims=['south_north', 'west_east'])
except (ValueError, KeyError):
pass
# Convert time (if necessary)
if 'Time' in ds.dims:
time = netcdf_time(ds)
if time is not None:
ds['Time'] = time
ds = ds.rename({'Time': 'time'})
tr = {'Time': 'time', 'XLAT': 'lat', 'XLONG': 'lon', 'XTIME': 'xtime'}
tr = {k: tr[k] for k in tr.keys() if k in ds.variables}
ds = ds.rename(tr)
# drop ugly vars
vns = ['Times', 'XLAT_V', 'XLAT_U', 'XLONG_U', 'XLONG_V']
vns = [vn for vn in vns if vn in ds.variables]
try:
ds = ds.drop_vars(vns)
except AttributeError:
ds = ds.drop(vns)
# add cartesian coords
ds['west_east'] = ds.salem.grid.x_coord
ds['south_north'] = ds.salem.grid.y_coord
# add pyproj string everywhere
ds.attrs['pyproj_srs'] = ds.salem.grid.proj.srs
for v in ds.data_vars:
ds[v].attrs['pyproj_srs'] = ds.salem.grid.proj.srs
return ds
