def test_3d_interp(self):
f = get_demo_file("wrf_d01_allvars_cropped.nc")
ds = sio.open_wrf_dataset(f)
out = ds.salem.wrf_zlevel("Z", levels=6000.0)
ref_2d = out * 0.0 + 6000.0
assert_allclose(out, ref_2d)
# this used to raise an error
_ = out.isel(time=1)
out = ds.salem.wrf_zlevel("Z", levels=[6000.0, 7000.0])
assert_allclose(out.sel(z=6000.0), ref_2d)
assert_allclose(out.sel(z=7000.0), ref_2d * 0.0 + 7000.0)
out = ds.salem.wrf_zlevel("Z")
assert_allclose(out.sel(z=7500.0), ref_2d * 0.0 + 7500.0)
out = ds.salem.wrf_plevel("PRESSURE", levels=400.0)
ref_2d = out * 0.0 + 400.0
assert_allclose(out, ref_2d)
out = ds.salem.wrf_plevel("PRESSURE", levels=[400.0, 300.0])
assert_allclose(out.sel(p=400.0), ref_2d)
assert_allclose(out.sel(p=300.0), ref_2d * 0.0 + 300.0)
out = ds.salem.wrf_plevel("PRESSURE")
assert_allclose(out.sel(p=300.0), ref_2d * 0.0 + 300.0)
ds = sio.open_wrf_dataset(get_demo_file("wrfout_d01.nc"))
ws_h = ds.isel(time=1).salem.wrf_zlevel("WS", levels=8000.0, use_multiprocessing=False)
assert np.all(np.isfinite(ws_h))
ws_h2 = ds.isel(time=1).salem.wrf_zlevel("WS", levels=8000.0)
assert_allclose(ws_h, ws_h2)
def test_transform_logic(self):
# This is just for the naming and dim logic, the rest is tested elsewh
ds1 = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc')).chunk()
ds2 = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc')).chunk()
# 2darray case
t2 = ds2.T2.isel(time=1)
with pytest.raises(ValueError):
ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid)
ds1.salem.transform_and_add(t2.values,
grid=t2.salem.grid,
name='t2_2darr')
assert 't2_2darr' in ds1
assert_allclose(ds1.t2_2darr.coords['south_north'],
t2.coords['south_north'])
assert_allclose(ds1.t2_2darr.coords['west_east'],
t2.coords['west_east'])
assert ds1.salem.grid == ds1.t2_2darr.salem.grid
# 3darray case
t2 = ds2.T2
ds1.salem.transform_and_add(t2.values,
grid=t2.salem.grid,
name='t2_3darr')
assert 't2_3darr' in ds1
assert_allclose(ds1.t2_3darr.coords['south_north'],
t2.coords['south_north'])
assert_allclose(ds1.t2_3darr.coords['west_east'],
t2.coords['west_east'])
assert 'time' in ds1.t2_3darr.coords
# dataarray case
ds1.salem.transform_and_add(t2, name='NEWT2')
assert 'NEWT2' in ds1
assert_allclose(ds1.NEWT2, ds1.T2)
assert_allclose(ds1.t2_3darr.coords['south_north'],
t2.coords['south_north'])
assert_allclose(ds1.t2_3darr.coords['west_east'],
t2.coords['west_east'])
assert 'time' in ds1.t2_3darr.coords
# dataset case
ds1.salem.transform_and_add(ds2[['RAINC', 'RAINNC']],
name={
'RAINC': 'PRCPC',
'RAINNC': 'PRCPNC'
})
assert 'PRCPC' in ds1
assert_allclose(ds1.PRCPC, ds1.RAINC)
assert 'time' in ds1.PRCPNC.coords
# what happens with external data?
dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
out = ds1.salem.transform(dse.t2m, interp='linear')
assert_allclose(out.coords['south_north'], t2.coords['south_north'])
assert_allclose(out.coords['west_east'], t2.coords['west_east'])
def test_full_wrf_wfile(self):
from salem.wrftools import var_classes
# TODO: these tests are qualitative and should be compared against ncl
f = get_demo_file('wrf_d01_allvars_cropped.nc')
ds = sio.open_wrf_dataset(f).chunk()
# making a repr was causing trouble because of the small chunks
_ = ds.__repr__()
# just check that the data is here
var_classes = copy.deepcopy(var_classes)
for vn in var_classes:
_ = ds[vn].values
dss = ds.isel(west_east=slice(2, 6),
south_north=slice(2, 5),
bottom_top=slice(0, 15))
_ = dss[vn].values
dss = ds.isel(west_east=1, south_north=2, bottom_top=3, time=2)
_ = dss[vn].values
# some chunking experiments
v = ds.WS.chunk((2, 1, 4, 5))
assert_allclose(v.mean(), ds.WS.mean(), atol=1e-3)
ds = ds.isel(time=slice(1, 4))
v = ds.PRCP.chunk((1, 2, 2))
assert_allclose(v.mean(), ds.PRCP.mean())
assert_allclose(v.max(), ds.PRCP.max())
def test_full_wrf_wfile(self):
from salem.wrftools import var_classes
# TODO: these tests are qualitative and should be compared against ncl
f = get_demo_file("wrf_d01_allvars_cropped.nc")
ds = sio.open_wrf_dataset(f)
# making a repr was causing trouble because of the small chunks
_ = ds.__repr__()
# just check that the data is here
var_classes = copy.deepcopy(var_classes)
for vn in var_classes:
_ = ds[vn].values
dss = ds.isel(west_east=slice(2, 6), south_north=slice(2, 5), bottom_top=slice(0, 15))
_ = dss[vn].values
dss = ds.isel(west_east=1, south_north=2, bottom_top=3, time=2)
_ = dss[vn].values
# some chunking experiments
v = ds.WS.chunk((2, 1, 4, 5))
assert_allclose(v.mean(), ds.WS.mean(), atol=1e-3)
ds = ds.isel(time=slice(1, 4))
v = ds.PRCP.chunk((1, 2, 2))
assert_allclose(v.mean(), ds.PRCP.mean())
assert_allclose(v.max(), ds.PRCP.max())
def test_wrf(self):
import xarray as xr
ds = sio.open_wrf_dataset(get_demo_file('wrf_tip_d1.nc')).chunk()
# this is because read_dataset changes some stuff, let's see if
# georef still ok
dsxr = xr.open_dataset(get_demo_file('wrf_tip_d1.nc'))
assert ds.salem.grid == dsxr.salem.grid
lon, lat = ds.salem.grid.ll_coordinates
assert_allclose(lon, ds['lon'], atol=1e-4)
assert_allclose(lat, ds['lat'], atol=1e-4)
# then something strange happened
assert ds.isel(time=0).salem.grid == ds.salem.grid
assert ds.isel(time=0).T2.salem.grid == ds.salem.grid
nlon, nlat = ds.isel(time=0).T2.salem.grid.ll_coordinates
assert_allclose(nlon, ds['lon'], atol=1e-4)
assert_allclose(nlat, ds['lat'], atol=1e-4)
# the grid should not be missunderstood as lonlat
t2 = ds.T2.isel(time=0) - 273.15
with pytest.raises(RuntimeError):
g = t2.salem.grid
def test_wrf(self):
import xarray as xr
ds = sio.open_wrf_dataset(get_demo_file("wrf_tip_d1.nc"))
# this is because read_dataset changes some stuff, let's see if
# georef still ok
dsxr = xr.open_dataset(get_demo_file("wrf_tip_d1.nc"))
assert ds.salem.grid == dsxr.salem.grid
lon, lat = ds.salem.grid.ll_coordinates
assert_allclose(lon, ds["lon"], atol=1e-4)
assert_allclose(lat, ds["lat"], atol=1e-4)
# then something strange happened
assert ds.isel(time=0).salem.grid == ds.salem.grid
assert ds.isel(time=0).T2.salem.grid == ds.salem.grid
nlon, nlat = ds.isel(time=0).T2.salem.grid.ll_coordinates
assert_allclose(nlon, ds["lon"], atol=1e-4)
assert_allclose(nlat, ds["lat"], atol=1e-4)
# the grid should not be missunderstood as lonlat
t2 = ds.T2.isel(time=0) - 273.15
with pytest.raises(RuntimeError):
g = t2.salem.grid
def test_mf_datasets(self):
import xarray as xr
if not os.path.exists(testdir):
os.makedirs(testdir)
# prepare the data
f = get_demo_file('wrf_d01_allvars_cropped.nc')
ds = xr.open_dataset(f)
for i in range(4):
dss = ds.isel(Time=[i])
dss.to_netcdf(os.path.join(testdir, 'wrf_slice_{}.nc'.format(i)))
dss.close()
ds = sio.open_wrf_dataset(f)
dsm = sio.open_mf_wrf_dataset(os.path.join(testdir, 'wrf_slice_*.nc'))
assert_allclose(ds['RAINNC'], dsm['RAINNC'])
assert_allclose(ds['GEOPOTENTIAL'], dsm['GEOPOTENTIAL'])
assert_allclose(ds['T2C'], dsm['T2C'])
assert 'PRCP' not in dsm.variables
prcp_nc_r = dsm.RAINNC.salem.deacc(as_rate=False)
self.assertEqual(prcp_nc_r.units, 'mm step-1')
self.assertEqual(prcp_nc_r.description,
'TOTAL GRID SCALE PRECIPITATION')
prcp_nc = dsm.RAINNC.salem.deacc()
self.assertEqual(prcp_nc.units, 'mm h-1')
self.assertEqual(prcp_nc.description, 'TOTAL GRID SCALE PRECIPITATION')
assert_allclose(prcp_nc_r / 3, prcp_nc)
# note that this is needed because there are variables which just
# can't be computed lazily (i.e. prcp)
fo = os.path.join(testdir, 'wrf_merged.nc')
if os.path.exists(fo):
os.remove(fo)
dsm = dsm[['RAINNC', 'RAINC']].load()
dsm.to_netcdf(fo)
dsm.close()
dsm = sio.open_wrf_dataset(fo)
assert_allclose(ds['PRCP'], dsm['PRCP'])
assert_allclose(prcp_nc,
dsm['PRCP_NC'].isel(time=slice(1, 4)),
rtol=1e-6)
def test_lookup_transform(self):
dsw = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len)
# qualitative tests (quantitative testing done elsewhere)
assert out[0, 0] == 0
assert out.mean() > 1
dsw = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
dse = sio.open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
_, lut = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len,
return_lut=True)
out2 = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len,
lut=lut)
# qualitative tests (quantitative testing done elsewhere)
assert_allclose(out, out2)
def test_geo_em(self):
for i in [1, 2, 3]:
fg = get_demo_file('geo_em_d0{}_lambert.nc'.format(i))
ds = sio.open_wrf_dataset(fg).chunk()
self.assertFalse('Time' in ds.dims)
self.assertTrue('time' in ds.dims)
self.assertTrue('south_north' in ds.dims)
self.assertTrue('south_north' in ds.coords)
def test_prcp(self):
wf = get_demo_file("wrfout_d01.nc")
w = sio.open_wrf_dataset(wf)
nc = sio.open_xr_dataset(wf)
nc["REF_PRCP_NC"] = nc["RAINNC"] * 0.0
uns = (
nc["RAINNC"].isel(Time=slice(1, len(nc.bottom_top_stag))).values
- nc["RAINNC"].isel(Time=slice(0, -1)).values
)
nc["REF_PRCP_NC"].values[1:, ...] = uns * 60 / 180.0 # for three hours
nc["REF_PRCP_NC"].values[0, ...] = np.NaN
nc["REF_PRCP_C"] = nc["RAINC"] * 0.0
uns = (
nc["RAINC"].isel(Time=slice(1, len(nc.bottom_top_stag))).values - nc["RAINC"].isel(Time=slice(0, -1)).values
)
nc["REF_PRCP_C"].values[1:, ...] = uns * 60 / 180.0 # for three hours
nc["REF_PRCP_C"].values[0, ...] = np.NaN
nc["REF_PRCP"] = nc["REF_PRCP_C"] + nc["REF_PRCP_NC"]
for suf in ["_NC", "_C", ""]:
assert_allclose(w["PRCP" + suf], nc["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=slice(1, 3))
ncn = nc.isel(Time=slice(1, 3))
assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=2)
ncn = nc.isel(Time=2)
assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=1)
ncn = nc.isel(Time=1)
assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=0)
self.assertTrue(~np.any(np.isfinite(wn["PRCP" + suf].values)))
wn = w.isel(time=slice(1, 3), south_north=slice(50, -1))
ncn = nc.isel(Time=slice(1, 3), south_north=slice(50, -1))
assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=2, south_north=slice(50, -1))
ncn = nc.isel(Time=2, south_north=slice(50, -1))
assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=1, south_north=slice(50, -1))
ncn = nc.isel(Time=1, south_north=slice(50, -1))
assert_allclose(wn["PRCP" + suf], ncn["REF_PRCP" + suf], rtol=1e-5)
wn = w.isel(time=0, south_north=slice(50, -1))
self.assertTrue(~np.any(np.isfinite(wn["PRCP" + suf].values)))
def test_mf_datasets(self):
import xarray as xr
if not os.path.exists(testdir):
os.makedirs(testdir)
# prepare the data
f = get_demo_file("wrf_d01_allvars_cropped.nc")
ds = xr.open_dataset(f)
for i in range(4):
dss = ds.isel(Time=[i])
dss.to_netcdf(os.path.join(testdir, "wrf_slice_{}.nc".format(i)))
dss.close()
ds = sio.open_wrf_dataset(f)
dsm = sio.open_mf_wrf_dataset(os.path.join(testdir, "wrf_slice_*.nc"))
assert_allclose(ds["RAINNC"], dsm["RAINNC"])
assert_allclose(ds["GEOPOTENTIAL"], dsm["GEOPOTENTIAL"])
assert_allclose(ds["T2C"], dsm["T2C"])
assert "PRCP" not in dsm.variables
prcp_nc_r = dsm.RAINNC.salem.deacc(as_rate=False)
self.assertEqual(prcp_nc_r.units, "mm step-1")
self.assertEqual(prcp_nc_r.description, "TOTAL GRID SCALE PRECIPITATION")
prcp_nc = dsm.RAINNC.salem.deacc()
self.assertEqual(prcp_nc.units, "mm h-1")
self.assertEqual(prcp_nc.description, "TOTAL GRID SCALE PRECIPITATION")
assert_allclose(prcp_nc_r / 3, prcp_nc)
# note that this is needed because there are variables which just
# can't be computed lazily (i.e. prcp)
fo = os.path.join(testdir, "wrf_merged.nc")
if os.path.exists(fo):
os.remove(fo)
dsm = dsm[["RAINNC", "RAINC"]].load()
dsm.to_netcdf(fo)
dsm.close()
dsm = sio.open_wrf_dataset(fo)
assert_allclose(ds["PRCP"], dsm["PRCP"])
assert_allclose(prcp_nc, dsm["PRCP_NC"].isel(time=slice(1, 4)), rtol=1e-6)
def test_transform_logic(self):
# This is just for the naming and dim logic, the rest is tested elsewh
ds1 = sio.open_wrf_dataset(get_demo_file("wrfout_d01.nc"))
ds2 = sio.open_wrf_dataset(get_demo_file("wrfout_d01.nc"))
# 2darray case
t2 = ds2.T2.isel(time=1)
with pytest.raises(ValueError):
ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid)
ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid, name="t2_2darr")
assert "t2_2darr" in ds1
assert_allclose(ds1.t2_2darr.coords["south_north"], t2.coords["south_north"])
assert_allclose(ds1.t2_2darr.coords["west_east"], t2.coords["west_east"])
assert ds1.salem.grid == ds1.t2_2darr.salem.grid
# 3darray case
t2 = ds2.T2
ds1.salem.transform_and_add(t2.values, grid=t2.salem.grid, name="t2_3darr")
assert "t2_3darr" in ds1
assert_allclose(ds1.t2_3darr.coords["south_north"], t2.coords["south_north"])
assert_allclose(ds1.t2_3darr.coords["west_east"], t2.coords["west_east"])
assert "time" in ds1.t2_3darr.coords
# dataarray case
ds1.salem.transform_and_add(t2, name="NEWT2")
assert "NEWT2" in ds1
assert_allclose(ds1.NEWT2, ds1.T2)
assert_allclose(ds1.t2_3darr.coords["south_north"], t2.coords["south_north"])
assert_allclose(ds1.t2_3darr.coords["west_east"], t2.coords["west_east"])
assert "time" in ds1.t2_3darr.coords
# dataset case
ds1.salem.transform_and_add(ds2[["RAINC", "RAINNC"]], name={"RAINC": "PRCPC", "RAINNC": "PRCPNC"})
assert "PRCPC" in ds1
assert_allclose(ds1.PRCPC, ds1.RAINC)
assert "time" in ds1.PRCPNC.coords
# what happens with external data?
dse = sio.open_xr_dataset(get_demo_file("era_interim_tibet.nc"))
out = ds1.salem.transform(dse.t2m, interp="linear")
assert_allclose(out.coords["south_north"], t2.coords["south_north"])
assert_allclose(out.coords["west_east"], t2.coords["west_east"])
def test_diagvars(self):
wf = get_demo_file("wrf_d01_allvars_cropped.nc")
w = sio.open_wrf_dataset(wf)
# ws
w["ws_ref"] = np.sqrt(w["U"] ** 2 + w["V"] ** 2)
assert_allclose(w["ws_ref"], w["WS"])
wcrop = w.isel(west_east=slice(4, 8), bottom_top=4)
assert_allclose(wcrop["ws_ref"], wcrop["WS"])
def test_3d_interp(self):
f = get_demo_file('wrf_d01_allvars_cropped.nc')
ds = sio.open_wrf_dataset(f).chunk()
out = ds.salem.wrf_zlevel('Z', levels=6000.)
ref_2d = out * 0. + 6000.
assert_allclose(out, ref_2d)
# this used to raise an error
_ = out.isel(time=1)
out = ds.salem.wrf_zlevel('Z', levels=[6000., 7000.])
assert_allclose(out.sel(z=6000.), ref_2d)
assert_allclose(out.sel(z=7000.), ref_2d * 0. + 7000.)
assert np.all(np.isfinite(out))
out = ds.salem.wrf_zlevel('Z')
assert_allclose(out.sel(z=7500.), ref_2d * 0. + 7500.)
out = ds.salem.wrf_plevel('PRESSURE', levels=400.)
ref_2d = out * 0. + 400.
assert_allclose(out, ref_2d)
out = ds.salem.wrf_plevel('PRESSURE', levels=[400., 300.])
assert_allclose(out.sel(p=400.), ref_2d)
assert_allclose(out.sel(p=300.), ref_2d * 0. + 300.)
out = ds.salem.wrf_plevel('PRESSURE')
assert_allclose(out.sel(p=300.), ref_2d * 0. + 300.)
assert np.any(~np.isfinite(out))
out = ds.salem.wrf_plevel('PRESSURE', fill_value='extrapolate')
assert_allclose(out.sel(p=300.), ref_2d * 0. + 300.)
assert np.all(np.isfinite(out))
ds = sio.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
ws_h = ds.isel(time=1).salem.wrf_zlevel('WS',
levels=8000.,
use_multiprocessing=False)
assert np.all(np.isfinite(ws_h))
ws_h2 = ds.isel(time=1).salem.wrf_zlevel('WS', levels=8000.)
assert_allclose(ws_h, ws_h2)
def test_diagvars(self):
wf = get_demo_file('wrf_d01_allvars_cropped.nc')
w = sio.open_wrf_dataset(wf).chunk()
# ws
w['ws_ref'] = np.sqrt(w['U']**2 + w['V']**2)
assert_allclose(w['ws_ref'], w['WS'])
wcrop = w.isel(west_east=slice(4, 8), bottom_top=4)
assert_allclose(wcrop['ws_ref'], wcrop['WS'])
def test_prcp(self):
wf = get_demo_file('wrfout_d01.nc')
w = sio.open_wrf_dataset(wf).chunk()
nc = sio.open_xr_dataset(wf)
nc['REF_PRCP_NC'] = nc['RAINNC'] * 0.
uns = nc['RAINNC'].isel(Time=slice(1, len(nc.bottom_top_stag))).values - \
nc['RAINNC'].isel(Time=slice(0, -1)).values
nc['REF_PRCP_NC'].values[1:, ...] = uns * 60 / 180. # for three hours
nc['REF_PRCP_NC'].values[0, ...] = np.NaN
nc['REF_PRCP_C'] = nc['RAINC'] * 0.
uns = nc['RAINC'].isel(Time=slice(1, len(nc.bottom_top_stag))).values - \
nc['RAINC'].isel(Time=slice(0, -1)).values
nc['REF_PRCP_C'].values[1:, ...] = uns * 60 / 180. # for three hours
nc['REF_PRCP_C'].values[0, ...] = np.NaN
nc['REF_PRCP'] = nc['REF_PRCP_C'] + nc['REF_PRCP_NC']
for suf in ['_NC', '_C', '']:
assert_allclose(w['PRCP' + suf], nc['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=slice(1, 3))
ncn = nc.isel(Time=slice(1, 3))
assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=2)
ncn = nc.isel(Time=2)
assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=1)
ncn = nc.isel(Time=1)
assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=0)
self.assertTrue(~np.any(np.isfinite(wn['PRCP' + suf].values)))
wn = w.isel(time=slice(1, 3), south_north=slice(50, -1))
ncn = nc.isel(Time=slice(1, 3), south_north=slice(50, -1))
assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=2, south_north=slice(50, -1))
ncn = nc.isel(Time=2, south_north=slice(50, -1))
assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=1, south_north=slice(50, -1))
ncn = nc.isel(Time=1, south_north=slice(50, -1))
assert_allclose(wn['PRCP' + suf], ncn['REF_PRCP' + suf], rtol=1e-5)
wn = w.isel(time=0, south_north=slice(50, -1))
self.assertTrue(~np.any(np.isfinite(wn['PRCP' + suf].values)))
def test_unstagger(self):
wf = get_demo_file("wrf_cropped.nc")
w = sio.open_wrf_dataset(wf)
nc = sio.open_xr_dataset(wf)
nc["PH_UNSTAGG"] = nc["P"] * 0.0
uns = (
nc["PH"].isel(bottom_top_stag=slice(0, -1)).values
+ nc["PH"].isel(bottom_top_stag=slice(1, len(nc.bottom_top_stag))).values
)
nc["PH_UNSTAGG"].values = uns * 0.5
assert_allclose(w["PH"], nc["PH_UNSTAGG"])
# chunk
v = w["PH"].chunk((1, 6, 13, 13))
assert_allclose(v.mean(), nc["PH_UNSTAGG"].mean(), atol=1e-2)
wn = w.isel(west_east=slice(4, 8))
ncn = nc.isel(west_east=slice(4, 8))
assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])
wn = w.isel(south_north=slice(4, 8), time=1)
ncn = nc.isel(south_north=slice(4, 8), Time=1)
assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])
wn = w.isel(west_east=4)
ncn = nc.isel(west_east=4)
assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])
wn = w.isel(bottom_top=4)
ncn = nc.isel(bottom_top=4)
assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])
wn = w.isel(bottom_top=0)
ncn = nc.isel(bottom_top=0)
assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])
wn = w.isel(bottom_top=-1)
ncn = nc.isel(bottom_top=-1)
assert_allclose(wn["PH"], ncn["PH_UNSTAGG"])
w["PH"].chunk()
def test_unstagger(self):
wf = get_demo_file('wrf_cropped.nc')
w = sio.open_wrf_dataset(wf).chunk()
nc = sio.open_xr_dataset(wf).chunk()
nc['PH_UNSTAGG'] = nc['P'] * 0.
uns = nc['PH'].isel(bottom_top_stag=slice(0, -1)).values + \
nc['PH'].isel(bottom_top_stag=slice(1, len(nc.bottom_top_stag))).values
nc['PH_UNSTAGG'].values = uns * 0.5
assert_allclose(w['PH'], nc['PH_UNSTAGG'])
# chunk
v = w['PH'].chunk((1, 6, 13, 13))
assert_allclose(v.mean(), nc['PH_UNSTAGG'].mean(), atol=1e-2)
wn = w.isel(west_east=slice(4, 8))
ncn = nc.isel(west_east=slice(4, 8))
assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])
wn = w.isel(south_north=slice(4, 8), time=1)
ncn = nc.isel(south_north=slice(4, 8), Time=1)
assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])
wn = w.isel(west_east=4)
ncn = nc.isel(west_east=4)
assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])
wn = w.isel(bottom_top=4)
ncn = nc.isel(bottom_top=4)
assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])
wn = w.isel(bottom_top=0)
ncn = nc.isel(bottom_top=0)
assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])
wn = w.isel(bottom_top=-1)
ncn = nc.isel(bottom_top=-1)
assert_allclose(wn['PH'], ncn['PH_UNSTAGG'])
w['PH'].chunk()
def test_ncl_diagvars(self):
import xarray as xr
wf = get_demo_file('wrf_cropped.nc')
ncl_out = get_demo_file('wrf_cropped_ncl.nc')
with warnings.catch_warnings(record=True) as war:
w = sio.open_wrf_dataset(wf).chunk()
self.assertEqual(len(war), 0)
nc = xr.open_dataset(ncl_out)
ref = nc['TK']
tot = w['TK']
assert_allclose(ref, tot, rtol=1e-6)
ref = nc['SLP']
tot = w['SLP']
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
w = w.isel(time=1, south_north=slice(12, 16), west_east=slice(9, 16))
nc = nc.isel(Time=1, south_north=slice(12, 16), west_east=slice(9, 16))
ref = nc['TK']
tot = w['TK']
assert_allclose(ref, tot, rtol=1e-6)
ref = nc['SLP']
tot = w['SLP']
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
w = w.isel(bottom_top=slice(3, 5))
nc = nc.isel(bottom_top=slice(3, 5))
ref = nc['TK']
tot = w['TK']
assert_allclose(ref, tot, rtol=1e-6)
ref = nc['SLP']
tot = w['SLP']
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
def test_ncl_diagvars(self):
wf = get_demo_file('wrf_cropped.nc')
ncl_out = get_demo_file('wrf_cropped_ncl.nc')
w = sio.open_wrf_dataset(wf)
nc = sio.open_xr_dataset(ncl_out)
ref = nc['TK']
tot = w['TK']
assert_allclose(ref, tot, rtol=1e-6)
ref = nc['SLP']
tot = w['SLP']
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
w = w.isel(time=1, south_north=slice(12, 16), west_east=slice(9, 16))
nc = nc.isel(Time=1, south_north=slice(12, 16), west_east=slice(9, 16))
ref = nc['TK']
tot = w['TK']
assert_allclose(ref, tot, rtol=1e-6)
ref = nc['SLP']
tot = w['SLP']
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
w = w.isel(bottom_top=slice(3, 5))
nc = nc.isel(bottom_top=slice(3, 5))
ref = nc['TK']
tot = w['TK']
assert_allclose(ref, tot, rtol=1e-6)
ref = nc['SLP']
tot = w['SLP']
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
def test_ncl_diagvars(self):
wf = get_demo_file("wrf_cropped.nc")
ncl_out = get_demo_file("wrf_cropped_ncl.nc")
w = sio.open_wrf_dataset(wf)
nc = sio.open_xr_dataset(ncl_out)
ref = nc["TK"]
tot = w["TK"]
assert_allclose(ref, tot, rtol=1e-6)
ref = nc["SLP"]
tot = w["SLP"]
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
w = w.isel(time=1, south_north=slice(12, 16), west_east=slice(9, 16))
nc = nc.isel(Time=1, south_north=slice(12, 16), west_east=slice(9, 16))
ref = nc["TK"]
tot = w["TK"]
assert_allclose(ref, tot, rtol=1e-6)
ref = nc["SLP"]
tot = w["SLP"]
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
w = w.isel(bottom_top=slice(3, 5))
nc = nc.isel(bottom_top=slice(3, 5))
ref = nc["TK"]
tot = w["TK"]
assert_allclose(ref, tot, rtol=1e-6)
ref = nc["SLP"]
tot = w["SLP"]
tot = tot.values
assert_allclose(ref, tot, rtol=1e-6)
