def test_example_docs():
import salem
from salem.utils import get_demo_file
ds = salem.open_xr_dataset(get_demo_file('wrfout_d01.nc'))
t2 = ds.T2.isel(Time=2)
t2_sub = t2.salem.subset(corners=((77., 20.), (97., 35.)), crs=salem.wgs84)
shdf = salem.read_shapefile(get_demo_file('world_borders.shp'))
shdf = shdf.loc[shdf['CNTRY_NAME'].isin(['Nepal', 'Bhutan'
])] # GeoPandas' GeoDataFrame
t2_sub = t2_sub.salem.subset(shape=shdf, margin=2) # add 2 grid points
t2_roi = t2_sub.salem.roi(shape=shdf)
smap = t2_roi.salem.get_map(data=t2_roi - 273.15,
cmap='RdYlBu_r',
vmin=-14,
vmax=18)
_ = smap.set_topography(get_demo_file('himalaya.tif'))
smap.set_shapefile(shape=shdf, color='grey', linewidth=3, zorder=5)
smap.set_points(91.1, 29.6)
smap.set_text(91.2, 29.7, 'Lhasa', fontsize=17)
smap.set_data(ds.T2.isel(Time=1) - 273.15, crs=ds.salem.grid)
fig, ax = plt.subplots(1, 1)
smap.visualize(ax=ax)
plt.tight_layout()
return fig
def test_center(self):
gm = GoogleCenterMap(center_ll=(10.762660, 46.794221),
zoom=13,
size_x=500,
size_y=500,
use_cache=False)
gm.set_roi(shape=get_demo_file('Hintereisferner.shp'))
gm.set_subset(toroi=True, margin=10)
img = gm.get_vardata()
img[np.nonzero(gm.roi == 0)] /= 2.
# from scipy.misc import toimage
# toimage(img).save(get_demo_file('hef_google_roi.png'))
ref = mpl.image.imread(get_demo_file('hef_google_roi.png'))
rmsd = np.sqrt(np.mean((ref - img)**2))
self.assertTrue(rmsd < 0.1)
# assert_allclose(ref, img, atol=2e-2)
gm = GoogleCenterMap(center_ll=(10.762660, 46.794221),
zoom=13,
size_x=500,
size_y=500)
gm.set_roi(shape=get_demo_file('Hintereisferner.shp'))
gm.set_subset(toroi=True, margin=10)
img = gm.get_vardata()
img[np.nonzero(gm.roi == 0)] /= 2.
rmsd = np.sqrt(np.mean((ref - img)**2))
self.assertTrue(rmsd < 0.1)
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_gmap_transformed():
dem = GeoTiff(get_demo_file('hef_srtm.tif'))
dem.set_subset(margin=-100)
dem = mercator_grid(center_ll=(10.76, 46.798444),
extent=(10000, 7000))
i, j = dem.ij_coordinates
g = GoogleVisibleMap(x=i, y=j, crs=dem, size_x=500, size_y=400)
img = g.get_vardata()
m = Map(dem, countries=False)
with pytest.raises(ValueError):
m.set_data(img)
m.set_lonlat_contours(interval=0.025)
m.set_shapefile(get_demo_file('Hintereisferner.shp'),
linewidths=2, edgecolor='darkred')
m.set_rgb(img, g.grid)
fig, ax = plt.subplots(1, 1)
m.visualize(ax=ax, addcbar=False)
plt.tight_layout()
return fig
def test_gmap_transformed():
dem = GeoTiff(get_demo_file('hef_srtm.tif'))
dem.set_subset(margin=-100)
dem = mercator_grid(center_ll=(10.76, 46.798444), extent=(10000, 7000))
i, j = dem.ij_coordinates
g = GoogleVisibleMap(x=i, y=j, crs=dem, size_x=500, size_y=400)
img = g.get_vardata()
m = Map(dem, countries=False)
with pytest.raises(ValueError):
m.set_data(img)
m.set_lonlat_contours(interval=0.025)
m.set_shapefile(get_demo_file('Hintereisferner.shp'),
linewidths=2,
edgecolor='darkred')
m.set_rgb(img, g.grid)
fig, ax = plt.subplots(1, 1)
m.visualize(ax=ax, addcbar=False)
plt.tight_layout()
return fig
def test_read_to_grid(self):
g = GeoTiff(utils.get_demo_file('hef_srtm.tif'))
sf = utils.get_demo_file('Hintereisferner_UTM.shp')
df1 = read_shapefile_to_grid(sf, g.grid)
df2 = transform_geopandas(read_shapefile(sf), to_crs=g.grid)
assert_allclose(df1.geometry[0].exterior.coords,
df2.geometry[0].exterior.coords)
# test for caching
d = g.grid.to_dict()
# change key ordering by chance
d2 = dict((k, v) for k, v in d.items())
from salem.sio import _memory_shapefile_to_grid, cached_shapefile_path
shape_cpath = cached_shapefile_path(sf)
res = _memory_shapefile_to_grid.call_and_shelve(shape_cpath,
grid=g.grid,
**d)
h1 = res.argument_hash
res = _memory_shapefile_to_grid.call_and_shelve(shape_cpath,
grid=g.grid,
**d2)
h2 = res.argument_hash
self.assertEqual(h1, h2)
def test_subset(self):
"""Open geotiff, do subsets and stuff"""
go = get_demo_file('hef_srtm.tif')
gs = get_demo_file('hef_srtm_subset.tif')
go = GeoTiff(go)
gs = GeoTiff(gs)
go.set_roi(grid=gs.grid)
go.set_subset(toroi=True)
ref = gs.get_vardata()
totest = go.get_vardata()
np.testing.assert_array_equal(ref.shape, (go.grid.ny, go.grid.nx))
np.testing.assert_array_equal(ref.shape, totest.shape)
np.testing.assert_array_equal(ref, totest)
go.set_roi()
go.set_subset()
eps = 1e-5
ex = gs.grid.extent_in_crs(crs=wgs84) # [left, right, bot, top
go.set_subset(corners=((ex[0], ex[2] + eps), (ex[1], ex[3] - eps)),
crs=wgs84,
margin=-2)
ref = gs.get_vardata()[2:-2, 2:-2]
totest = go.get_vardata()
np.testing.assert_array_equal(ref.shape, totest.shape)
np.testing.assert_array_equal(ref, totest)
go.set_roi()
go.set_subset()
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_gmap():
g = GoogleCenterMap(center_ll=(10.762660, 46.794221), zoom=13,
size_x=640, size_y=640)
m = Map(g.grid, countries=False, nx=640)
m.set_lonlat_countours(interval=0.025)
m.set_shapefile(get_demo_file('Hintereisferner.shp'),
linewidths=2, edgecolor='darkred')
m.set_rgb(g.get_vardata())
m.visualize(addcbar=False)
dem = salem.GeoTiff(get_demo_file('hef_srtm.tif'))
dem.set_subset(margin=-100)
dem = salem.grids.local_mercator_grid(center_ll=(10.76, 46.798444),
extent=(10000, 7000))
i, j = dem.ij_coordinates
g = GoogleVisibleMap(x=i, y=j, src=dem, size_x=500, size_y=400)
img = g.get_vardata()
m = Map(dem, countries=False)
assert_raises(ValueError, m.set_data, img)
m.set_lonlat_countours(interval=0.025)
m.set_shapefile(get_demo_file('Hintereisferner.shp'),
linewidths=2, edgecolor='darkred')
m.set_rgb(img, g.grid)
m.visualize(addcbar=False)
def test_shape(self):
"""Is the transformation doing well?"""
from salem import read_shapefile
so = read_shapefile(get_demo_file('Hintereisferner.shp'))
sref = read_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
st = gis.transform_geopandas(so, to_crs=sref.crs)
self.assertFalse(st is so)
assert_allclose(st.geometry[0].exterior.coords,
sref.geometry[0].exterior.coords)
sti = gis.transform_geopandas(so, to_crs=sref.crs, inplace=True)
self.assertTrue(sti is so)
assert_allclose(so.geometry[0].exterior.coords,
sref.geometry[0].exterior.coords)
assert_allclose(sti.geometry[0].exterior.coords,
sref.geometry[0].exterior.coords)
g = Grid(nxny=(1, 1), dxdy=(1, 1), ll_corner=(10., 46.), proj=wgs84)
so = read_shapefile(get_demo_file('Hintereisferner.shp'))
st = gis.transform_geopandas(so, to_crs=g)
ref = np.array(so.geometry[0].exterior.coords)
ref = ref - np.floor(ref)
assert_allclose(ref, st.geometry[0].exterior.coords)
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_shape(self):
"""Is the transformation doing well?"""
from salem import read_shapefile
so = read_shapefile(get_demo_file('Hintereisferner.shp'))
sref = read_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
st = gis.transform_geopandas(so, to_crs=sref.crs)
self.assertFalse(st is so)
assert_allclose(st.geometry[0].exterior.coords,
sref.geometry[0].exterior.coords)
sti = gis.transform_geopandas(so, to_crs=sref.crs, inplace=True)
self.assertTrue(sti is so)
assert_allclose(so.geometry[0].exterior.coords,
sref.geometry[0].exterior.coords)
assert_allclose(sti.geometry[0].exterior.coords,
sref.geometry[0].exterior.coords)
g = Grid(nxny=(1, 1), dxdy=(1, 1), x0y0=(10., 46.), proj=wgs84)
so = read_shapefile(get_demo_file('Hintereisferner.shp'))
st = gis.transform_geopandas(so, to_crs=g)
ref = np.array(so.geometry[0].exterior.coords)
ref = ref - np.floor(ref)
assert_allclose(ref, st.geometry[0].exterior.coords)
# round trip
so_back = gis.transform_geopandas(st, to_crs=so.crs)
assert_allclose(so_back.geometry[0].exterior.coords,
so.geometry[0].exterior.coords)
def test_diagnostic_vars(self):
d = WRF(get_demo_file('wrf_tip_d1.nc'))
d2 = GeoNetcdf(get_demo_file('wrf_tip_d2.nc'))
self.assertTrue('T2C' in d.variables)
ref = d.get_vardata('T2')
tot = d.get_vardata('T2C') + 273.15
np.testing.assert_allclose(ref, tot)
d.set_roi(grid=d2.grid)
d.set_subset(toroi=True)
ref = d.get_vardata('T2')
tot = d.get_vardata('T2C') + 273.15
self.assertEqual(tot.shape[-1] * 3, d2.grid.nx)
self.assertEqual(tot.shape[-2] * 3, d2.grid.ny)
np.testing.assert_allclose(ref, tot)
d = WRF(get_demo_file('wrf_tip_d1.nc'))
ref = d.variables['T2'][:]
d.set_subset(margin=-5)
tot = d.get_vardata('T2')
assert_array_equal(ref.shape[1] - 10, tot.shape[1])
assert_array_equal(ref.shape[2] - 10, tot.shape[2])
assert_array_equal(ref[:, 5:-5, 5:-5], tot)
def test_diagnostic_vars(self):
d = WRF(get_demo_file("wrf_tip_d1.nc"))
d2 = GeoNetcdf(get_demo_file("wrf_tip_d2.nc"))
self.assertTrue("T2C" in d.variables)
ref = d.get_vardata("T2")
tot = d.get_vardata("T2C") + 273.15
np.testing.assert_allclose(ref, tot)
d.set_roi(grid=d2.grid)
d.set_subset(toroi=True)
ref = d.get_vardata("T2")
tot = d.get_vardata("T2C") + 273.15
self.assertEqual(tot.shape[-1] * 3, d2.grid.nx)
self.assertEqual(tot.shape[-2] * 3, d2.grid.ny)
np.testing.assert_allclose(ref, tot)
d = WRF(get_demo_file("wrf_tip_d1.nc"))
ref = d.variables["T2"][:]
d.set_subset(margin=-5)
tot = d.get_vardata("T2")
assert_array_equal(ref.shape[1] - 10, tot.shape[1])
assert_array_equal(ref.shape[2] - 10, tot.shape[2])
assert_array_equal(ref[:, 5:-5, 5:-5], tot)
def test_subset(self):
"""Open geotiff, do subsets and stuff"""
go = get_demo_file("hef_srtm.tif")
gs = get_demo_file("hef_srtm_subset.tif")
go = GeoTiff(go)
gs = GeoTiff(gs)
go.set_roi(grid=gs.grid)
go.set_subset(toroi=True)
ref = gs.get_vardata()
totest = go.get_vardata()
np.testing.assert_array_equal(ref.shape, (go.grid.ny, go.grid.nx))
np.testing.assert_array_equal(ref.shape, totest.shape)
np.testing.assert_array_equal(ref, totest)
go.set_roi()
go.set_subset()
eps = 1e-5
ex = gs.grid.extent_in_crs(crs=wgs84) # [left, right, bot, top
go.set_subset(corners=((ex[0], ex[2] + eps), (ex[1], ex[3] - eps)), crs=wgs84, margin=-2)
ref = gs.get_vardata()[2:-2, 2:-2]
totest = go.get_vardata()
np.testing.assert_array_equal(ref.shape, totest.shape)
np.testing.assert_array_equal(ref, totest)
go.set_roi()
go.set_subset()
def test_lookup_transform():
dsw = open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
dse = open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len)
fig, ax = plt.subplots(1, 1)
out.salem.quick_map(ax=ax)
return fig
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_visible(self):
import matplotlib as mpl
x = [91.176036, 92.05, 88.880927]
y = [29.649702, 31.483333, 29.264956]
g = GoogleVisibleMap(x=x,
y=y,
size_x=400,
size_y=400,
maptype='terrain')
img = g.get_vardata()[..., :3]
i, j = g.grid.transform(x, y, nearest=True)
for _i, _j in zip(i, j):
img[_j - 3:_j + 4, _i - 3:_i + 4, 0] = 1
img[_j - 3:_j + 4, _i - 3:_i + 4, 1:] = 0
# from PIL import Image
# Image.fromarray((img * 255).astype(np.uint8)).save(
# get_demo_file('hef_google_visible.png'))
ref = mpl.image.imread(get_demo_file('hef_google_visible.png'))
rmsd = np.sqrt(np.mean((ref - img)**2))
self.assertTrue(rmsd < 1e-1)
self.assertRaises(ValueError, GoogleVisibleMap, x=x, y=y, zoom=12)
fw = get_demo_file('wrf_tip_d1.nc')
d = GeoNetcdf(fw)
i, j = d.grid.ij_coordinates
g = GoogleVisibleMap(x=i, y=j, crs=d.grid, size_x=500, size_y=500)
img = g.get_vardata()[..., :3]
mask = g.grid.map_gridded_data(i * 0 + 1, d.grid)
img[np.nonzero(mask)] = np.clip(img[np.nonzero(mask)] + 0.3, 0, 1)
# from PIL import Image
# Image.fromarray((img * 255).astype(np.uint8)).save(
# get_demo_file('hef_google_visible_grid.png'))
ref = mpl.image.imread(get_demo_file('hef_google_visible_grid.png'))
rmsd = np.sqrt(np.mean((ref - img)**2))
self.assertTrue(rmsd < 5e-1)
gm = GoogleVisibleMap(x=i, y=j, crs=d.grid, size_x=500, size_y=500)
gm2 = GoogleVisibleMap(x=i,
y=j,
crs=d.grid,
scale=2,
size_x=500,
size_y=500)
assert (gm.grid.nx * 2) == gm2.grid.nx
assert gm.grid.extent == gm2.grid.extent
# Test regression for non array inputs
grid = mercator_grid(center_ll=(72.5, 30.), extent=(2.0e6, 2.0e6))
GoogleVisibleMap(x=[0, grid.nx - 1], y=[0, grid.ny - 1], crs=grid)
def test_read_to_grid(self):
g = GeoTiff(utils.get_demo_file("hef_srtm.tif"))
sf = utils.get_demo_file("Hintereisferner_UTM.shp")
df1 = read_shapefile_to_grid(sf, g.grid)
df2 = transform_geopandas(read_shapefile(sf), to_crs=g.grid)
assert_allclose(df1.geometry[0].exterior.coords, df2.geometry[0].exterior.coords)
def test_read_to_grid(self):
g = GeoTiff(utils.get_demo_file('hef_srtm.tif'))
sf = utils.get_demo_file('Hintereisferner_UTM.shp')
df1 = read_shapefile_to_grid(sf, g.grid)
df2 = transform_geopandas(read_shapefile(sf), to_crs=g.grid)
assert_allclose(df1.geometry[0].exterior.coords,
df2.geometry[0].exterior.coords)
def test_hef_default_spline():
grid = mercator_grid(center_ll=(10.76, 46.798444), extent=(10000, 7000))
c = Map(grid, countries=False)
c.set_lonlat_contours(interval=0)
c.set_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
c.set_topography(get_demo_file('hef_srtm.tif'))
fig, ax = plt.subplots(1, 1)
c.visualize(ax=ax, addcbar=False, title='Default: spline deg 3')
plt.tight_layout()
return fig
def test_hef_default_spline():
grid = mercator_grid(center_ll=(10.76, 46.798444),
extent=(10000, 7000))
c = Map(grid, countries=False)
c.set_lonlat_contours(interval=0)
c.set_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
c.set_topography(get_demo_file('hef_srtm.tif'))
fig, ax = plt.subplots(1, 1)
c.visualize(ax=ax, addcbar=False, title='Default: spline deg 3')
plt.tight_layout()
return fig
def test_unstagger_compressed(self):
wf = get_demo_file('wrf_cropped.nc')
wfc = get_demo_file('wrf_cropped_compressed.nc')
# Under WRF
nc = WRF(wf)
ncc = WRF(wfc)
assert_allclose(nc.get_vardata('PH'), ncc.get_vardata('PH'), rtol=.003)
nc.set_period(1, 2)
ncc.set_period(1, 2)
assert_allclose(nc.get_vardata('PH'), ncc.get_vardata('PH'), rtol=.003)
def test_lookup_transform():
dsw = open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
dse = open_xr_dataset(get_demo_file('era_interim_tibet.nc'))
out = dse.salem.lookup_transform(dsw.T2C.isel(time=0), method=len)
fig, ax = plt.subplots(1, 1)
sm = out.salem.get_map()
sm.set_data(out)
sm.set_geometry(dsw.salem.grid.extent_as_polygon(), edgecolor='r',
linewidth=2)
sm.visualize(ax=ax)
return fig
def test_lambert_tuto(self):
from salem.wrftools import geogrid_simulator
g, m = geogrid_simulator(get_demo_file('namelist_tutorial.wps'))
assert len(g) == 1
fg = get_demo_file('geo_em.d01_tutorial.nc')
with netCDF4.Dataset(fg) as nc:
nc.set_auto_mask(False)
lon, lat = g[0].ll_coordinates
assert_allclose(lon, nc['XLONG_M'][0, ...], atol=1e-4)
assert_allclose(lat, nc['XLAT_M'][0, ...], atol=1e-4)
def test_plot_on_map():
import salem
from salem.utils import get_demo_file
ds = salem.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
t2_sub = ds.salem.subset(corners=((77., 20.), (97., 35.)), crs=salem.wgs84).T2.isel(time=2)
shdf = salem.read_shapefile(get_demo_file('world_borders.shp'))
shdf = shdf.loc[shdf['CNTRY_NAME'].isin(
['Nepal', 'Bhutan'])] # GeoPandas' GeoDataFrame
t2_sub = t2_sub.salem.subset(shape=shdf, margin=2) # add 2 grid points
t2_roi = t2_sub.salem.roi(shape=shdf)
fig, ax = plt.subplots(1, 1)
t2_roi.salem.quick_map(ax=ax)
plt.tight_layout()
return fig
def test_polar(self):
from salem.wrftools import geogrid_simulator
g, m = geogrid_simulator(get_demo_file('namelist_polar.wps'))
assert len(g) == 2
for i in [1, 2]:
fg = get_demo_file('geo_em_d0{}_polarstereo.nc'.format(i))
ds = netCDF4.Dataset(fg)
lon, lat = g[i - 1].ll_coordinates
assert_allclose(lon, ds['XLONG_M'][0, ...], atol=5e-3)
assert_allclose(lat, ds['XLAT_M'][0, ...], atol=5e-3)
def test_hef_from_array():
grid = mercator_grid(center_ll=(10.76, 46.798444), extent=(10000, 7000))
c = Map(grid, countries=False)
c.set_lonlat_contours(interval=0)
c.set_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
dem = GeoTiff(get_demo_file('hef_srtm.tif'))
mytopo = dem.get_vardata()
c.set_topography(mytopo, crs=dem.grid, interp='spline')
fig, ax = plt.subplots(1, 1)
c.visualize(ax=ax, addcbar=False, title='From array')
plt.tight_layout()
return fig
def test_mercator(self):
from salem.wrftools import geogrid_simulator
g, m = geogrid_simulator(get_demo_file("namelist_mercator.wps"))
assert len(g) == 4
for i in [1, 2, 3, 4]:
fg = get_demo_file("geo_em_d0{}_mercator.nc".format(i))
ds = netCDF4.Dataset(fg)
lon, lat = g[i - 1].ll_coordinates
assert_allclose(lon, ds["XLONG_M"][0, ...], atol=1e-4)
assert_allclose(lat, ds["XLAT_M"][0, ...], atol=1e-4)
def test_plot_on_map():
import salem
from salem.utils import get_demo_file
ds = salem.open_wrf_dataset(get_demo_file('wrfout_d01.nc'))
t2_sub = ds.salem.subset(corners=((77., 20.), (97., 35.)), crs=salem.wgs84).T2.isel(time=2)
shdf = salem.read_shapefile(get_demo_file('world_borders.shp'))
shdf = shdf.loc[shdf['CNTRY_NAME'].isin(
['Nepal', 'Bhutan'])] # GeoPandas' GeoDataFrame
t2_sub = t2_sub.salem.subset(shape=shdf, margin=2) # add 2 grid points
t2_roi = t2_sub.salem.roi(shape=shdf)
fig, ax = plt.subplots(1, 1)
t2_roi.salem.quick_map(ax=ax)
plt.tight_layout()
return fig
def test_ncl_diagvars(self):
wf = get_demo_file("wrf_cropped.nc")
ncl_out = get_demo_file("wrf_cropped_ncl.nc")
w = WRF(wf)
nc = netCDF4.Dataset(ncl_out)
ref = nc.variables["TK"][:]
tot = w.get_vardata("TK")
assert_allclose(ref, tot, rtol=1e-6)
ref = nc.variables["SLP"][:]
tot = w.get_vardata("SLP")
assert_allclose(ref, tot, rtol=1e-6)
def test_mercator(self):
from salem.wrftools import geogrid_simulator
g, m = geogrid_simulator(get_demo_file('namelist_mercator.wps'))
assert len(g) == 4
for i in [1, 2, 3, 4]:
fg = get_demo_file('geo_em_d0{}_mercator.nc'.format(i))
with netCDF4.Dataset(fg) as nc:
nc.set_auto_mask(False)
lon, lat = g[i - 1].ll_coordinates
assert_allclose(lon, nc['XLONG_M'][0, ...], atol=1e-4)
assert_allclose(lat, nc['XLAT_M'][0, ...], atol=1e-4)
def test_ncl_diagvars(self):
wf = get_demo_file('wrf_cropped.nc')
ncl_out = get_demo_file('wrf_cropped_ncl.nc')
w = WRF(wf)
nc = netCDF4.Dataset(ncl_out)
ref = nc.variables['TK'][:]
tot = w.get_vardata('TK')
assert_allclose(ref, tot, rtol=1e-6)
ref = nc.variables['SLP'][:]
tot = w.get_vardata('SLP')
assert_allclose(ref, tot, rtol=1e-6)
def test_hef_from_array():
grid = mercator_grid(center_ll=(10.76, 46.798444),
extent=(10000, 7000))
c = Map(grid, countries=False)
c.set_lonlat_contours(interval=0)
c.set_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
dem = GeoTiff(get_demo_file('hef_srtm.tif'))
mytopo = dem.get_vardata()
c.set_topography(mytopo, crs=dem.grid, interp='spline')
fig, ax = plt.subplots(1, 1)
c.visualize(ax=ax, addcbar=False, title='From array')
plt.tight_layout()
return fig
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_cartopy_polar():
import cartopy
fig = plt.figure(figsize=(8, 8))
ods = open_wrf_dataset(get_demo_file('geo_em_d02_polarstereo.nc'))
ods = ods.isel(time=0)
ax = plt.subplot(2, 2, 1)
smap = ods.salem.get_map()
smap.plot(ax=ax)
p = ods.salem.cartopy()
ax = plt.subplot(2, 2, 2, projection=p)
ax.coastlines(resolution='50m')
ax.gridlines()
ax.set_extent(ods.salem.grid.extent, crs=p)
ds = ods.salem.subset(corners=((-52.8, 70.11), (-52.8, 70.11)), margin=12)
ax = plt.subplot(2, 2, 3)
smap = ds.HGT_M.salem.quick_map(ax=ax, cmap='Oranges')
ax.scatter(ds.XLONG_M, ds.XLAT_M, s=5, transform=smap.transform(ax=ax))
p = ds.salem.cartopy()
ax = plt.subplot(2, 2, 4, projection=p)
ds.HGT_M.plot.imshow(ax=ax, transform=p, cmap='Oranges')
ax.coastlines(resolution='50m')
ax.gridlines()
ax.scatter(ds.XLONG_M, ds.XLAT_M, transform=cartopy.crs.PlateCarree(), s=5)
return fig
def test_longtime(self):
"""There was a bug with time"""
fs = get_demo_file("test_longtime.nc")
c = GeoNetcdf(fs)
self.assertEqual(len(c.time), 2424)
assert_array_equal(c.time[0:2], pd.to_datetime([datetime(1801, 10, 1), datetime(1801, 11, 1)]))
def test_unstagger(self):
wf = get_demo_file("wrf_cropped.nc")
nc = netCDF4.Dataset(wf)
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 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_eraint(self):
f = get_demo_file('era_interim_tibet.nc')
d = GeoNetcdf(f)
assert d.grid.origin == 'upper-left'
stat_lon = 91.1
stat_lat = 31.1
with netCDF4.Dataset(f) as nc:
nc.set_auto_mask(False)
flon = nc.variables['longitude'][:]
flat = nc.variables['latitude'][:]
alon = np.argmin(np.abs(flon - stat_lon))
alat = np.argmin(np.abs(flat - stat_lat))
d.set_subset(corners=((stat_lon, stat_lat), (stat_lon, stat_lat)))
slon, slat = d.grid.ll_coordinates
assert_array_equal(flon[alon], slon)
assert_allclose(flat[alat], slat)
# Exotic subset
assert_array_equal(flon[alon], d.get_vardata('longitude'))
assert_allclose(flat[alat], d.get_vardata('latitude'))
assert_allclose(nc.variables['t2m'][:, alat, alon],
np.squeeze(d.get_vardata('t2m')))
d.set_period(t0='2012-06-01 06:00:00', t1='2012-06-01 12:00:00')
assert_allclose(nc.variables['t2m'][1:3, alat, alon],
np.squeeze(d.get_vardata('t2m')))
def test_metum(self):
ds = sio.open_metum_dataset(get_demo_file('rotated_grid.nc'))
# One way
mylons, mylats = ds.salem.grid.ll_coordinates
assert_allclose(mylons, ds.longitude_t, atol=1e-7)
assert_allclose(mylats, ds.latitude_t, atol=1e-7)
# Round trip
i, j = ds.salem.grid.transform(mylons, mylats)
ii, jj = ds.salem.grid.ij_coordinates
assert_allclose(i, ii, atol=1e-7)
assert_allclose(j, jj, atol=1e-7)
# Cartopy
from salem.gis import proj_to_cartopy
from cartopy.crs import PlateCarree
cp = proj_to_cartopy(ds.salem.grid.proj)
xx, yy = ds.salem.grid.xy_coordinates
out = PlateCarree().transform_points(cp, xx.flatten(), yy.flatten())
assert_allclose(out[:, 0].reshape(ii.shape), ds.longitude_t, atol=1e-7)
assert_allclose(out[:, 1].reshape(ii.shape), ds.latitude_t, atol=1e-7)
# Round trip
out = cp.transform_points(PlateCarree(),
ds.longitude_t.values.flatten(),
ds.latitude_t.values.flatten())
assert_allclose(out[:, 0].reshape(ii.shape), xx, atol=1e-7)
assert_allclose(out[:, 1].reshape(ii.shape), yy, atol=1e-7)
def test_eraint(self):
f = get_demo_file("era_interim_tibet.nc")
d = GeoNetcdf(f)
assert d.grid.order == "ul"
stat_lon = 91.1
stat_lat = 31.1
nc = netCDF4.Dataset(f)
flon = nc.variables["longitude"][:]
flat = nc.variables["latitude"][:]
alon = np.argmin(np.abs(flon - stat_lon))
alat = np.argmin(np.abs(flat - stat_lat))
d.set_subset(corners=((stat_lon, stat_lat), (stat_lon, stat_lat)))
slon, slat = d.grid.ll_coordinates
assert_array_equal(flon[alon], slon)
assert_allclose(flat[alat], slat)
# Exotic subset
assert_array_equal(flon[alon], d.get_vardata("longitude"))
assert_allclose(flat[alat], d.get_vardata("latitude"))
assert_allclose(nc.variables["t2m"][:, alat, alon], np.squeeze(d.get_vardata("t2m")))
d.set_period(t0="2012-06-01 06:00:00", t1="2012-06-01 12:00:00")
assert_allclose(nc.variables["t2m"][1:3, alat, alon], np.squeeze(d.get_vardata("t2m")))
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 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_ncl_diagvars_compressed(self):
wf = get_demo_file('wrf_cropped_compressed.nc')
ncl_out = get_demo_file('wrf_cropped_ncl.nc')
w = WRF(wf)
with netCDF4.Dataset(ncl_out) as nc:
nc.set_auto_mask(False)
ref = nc.variables['TK'][:]
tot = w.get_vardata('TK')
assert_allclose(ref, tot, rtol=1e-5)
ref = nc.variables['SLP'][:]
tot = w.get_vardata('SLP')
assert_allclose(ref, tot, rtol=1e-4)
def test_wrf_polar(self):
d = GeoNetcdf(get_demo_file('geo_em_d01_polarstereo.nc'))
mylon, mylat = d.grid.ll_coordinates
reflon = np.squeeze(d.get_vardata('XLONG_M'))
reflat = np.squeeze(d.get_vardata('XLAT_M'))
np.testing.assert_allclose(reflon, mylon, atol=5e-3)
np.testing.assert_allclose(reflat, mylat, atol=5e-3)
d = GeoNetcdf(get_demo_file('geo_em_d02_polarstereo.nc'))
mylon, mylat = d.grid.ll_coordinates
reflon = np.squeeze(d.get_vardata('XLONG_M'))
reflat = np.squeeze(d.get_vardata('XLAT_M'))
np.testing.assert_allclose(reflon, mylon, atol=1e-4)
np.testing.assert_allclose(reflat, mylat, atol=1e-4)
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_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_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_staggeredcoords(self):
wf = get_demo_file("wrf_cropped.nc")
nc = GeoNetcdf(wf)
lon, lat = nc.grid.xstagg_ll_coordinates
assert_allclose(np.squeeze(nc.variables["XLONG_U"][0, ...]), lon, atol=1e-4)
assert_allclose(np.squeeze(nc.variables["XLAT_U"][0, ...]), lat, atol=1e-4)
lon, lat = nc.grid.ystagg_ll_coordinates
assert_allclose(np.squeeze(nc.variables["XLONG_V"][0, ...]), lon, atol=1e-4)
assert_allclose(np.squeeze(nc.variables["XLAT_V"][0, ...]), lat, atol=1e-4)
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_as_xarray(self):
f = get_demo_file("era_interim_tibet.nc")
d = GeoNetcdf(f)
t2 = d.get_vardata("t2m", as_xarray=True)
stat_lon = 91.1
stat_lat = 31.1
d.set_subset(corners=((stat_lon, stat_lat), (stat_lon, stat_lat)))
t2_sub = d.get_vardata("t2m", as_xarray=True)
np.testing.assert_allclose(t2_sub - t2, np.zeros(4).reshape((4, 1, 1)))
d.set_period(t0="2012-06-01 06:00:00", t1="2012-06-01 12:00:00")
t2_sub = d.get_vardata("t2m", as_xarray=True)
np.testing.assert_allclose(t2_sub - t2, np.zeros(2).reshape((2, 1, 1)))
wf = get_demo_file("wrf_cropped.nc")
d = WRF(wf)
tk = d.get_vardata("TK", as_xarray=True)
def test_har(self):
# HAR
hf = get_demo_file("har_d30km_y_2d_t2_2000.nc")
d = GeoNetcdf(hf)
reflon = np.squeeze(d.get_vardata("lon"))
reflat = np.squeeze(d.get_vardata("lat"))
mylon, mylat = d.grid.ll_coordinates
np.testing.assert_allclose(reflon, mylon, atol=1e-5)
np.testing.assert_allclose(reflat, mylat, atol=1e-5)
def test_oceans():
f = os.path.join(get_demo_file('wrf_tip_d1.nc'))
grid = GeoNetcdf(f).grid
m = Map(grid, countries=False)
m.set_shapefile(rivers=True, linewidths=2)
m.set_shapefile(oceans=True, edgecolor='k', linewidth=3)
fig, ax = plt.subplots(1, 1)
m.visualize(ax=ax, addcbar=False)
plt.tight_layout()
def test_hef_topo_withnan():
grid = mercator_grid(center_ll=(10.76, 46.798444),
extent=(10000, 7000))
c = Map(grid, countries=False)
c.set_lonlat_contours(interval=10)
c.set_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
dem = GeoTiff(get_demo_file('hef_srtm.tif'))
mytopo = dem.get_vardata()
h = c.set_topography(mytopo, crs=dem.grid, interp='spline')
c.set_lonlat_contours()
c.set_cmap(get_cmap('topo'))
c.set_plot_params(nlevels=256)
# Try with nan data
h[-100:, -100:] = np.NaN
c.set_data(h)
fig, ax = plt.subplots(1, 1)
c.visualize(ax=ax, title='color with NaN')
plt.tight_layout()
return fig
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_to_cartopy(self):
import cartopy.crs as ccrs
from salem import GeoNetcdf, GeoTiff
grid = gis.mercator_grid(center_ll=(11.38, 47.26),
extent=(2000000, 2000000))
p = gis.proj_to_cartopy(grid.proj)
assert isinstance(p, ccrs.TransverseMercator)
fuzzy_proj_tester(grid.proj, pyproj.Proj(p.proj4_params))
ds = GeoNetcdf(get_demo_file('wrfout_d01.nc'))
p = gis.proj_to_cartopy(ds.grid.proj)
assert isinstance(p, ccrs.LambertConformal)
fuzzy_proj_tester(ds.grid.proj, pyproj.Proj(p.proj4_params))
ds = GeoNetcdf(get_demo_file('wrf_mercator.nc'))
p = gis.proj_to_cartopy(ds.grid.proj)
assert isinstance(p, ccrs.Mercator)
fuzzy_proj_tester(ds.grid.proj, pyproj.Proj(p.proj4_params))
ds = GeoTiff(get_demo_file('himalaya.tif'))
p = gis.proj_to_cartopy(ds.grid.proj)
assert isinstance(p, ccrs.PlateCarree)
ds = GeoTiff(get_demo_file('hef_roi.tif'))
p = gis.proj_to_cartopy(ds.grid.proj)
assert isinstance(p, ccrs.PlateCarree)
fuzzy_proj_tester(ds.grid.proj, pyproj.Proj(p.proj4_params))
p = gis.proj_to_cartopy(wgs84)
assert isinstance(p, ccrs.PlateCarree)
p = gis.proj_to_cartopy(pyproj.Proj('+proj=utm +zone=15'))
assert isinstance(p, ccrs.UTM)
# this needs gdal
if gis.has_gdal:
p = gis.proj_to_cartopy(pyproj.Proj(init='epsg:26915'))
assert isinstance(p, ccrs.UTM)
def test_hef():
grid = salem.grids.local_mercator_grid(center_ll=(10.76, 46.798444),
extent=(10000, 7000))
c = Map(grid, countries=False)
c.set_lonlat_countours(interval=10)
c.set_shapefile(get_demo_file('Hintereisferner_UTM.shp'))
c.set_topography(get_demo_file('hef_srtm.tif'),
interp='linear')
c.visualize(addcbar=False, title='linear')
c.set_topography(get_demo_file('hef_srtm.tif'),
interp='spline', ks=2)
c.visualize(addcbar=False, title='spline deg 2')
c.set_topography(get_demo_file('hef_srtm.tif'))
c.visualize(addcbar=False, title='Default: spline deg 3')
h = c.set_topography(get_demo_file('hef_srtm.tif'),
interp='spline', ks=5)
c.visualize(addcbar=False, title='spline deg 5')
dem = salem.GeoTiff(get_demo_file('hef_srtm.tif'))
mytopo = dem.get_vardata()
c.set_topography(mytopo, crs=dem.grid, interp='spline')
c.visualize(addcbar=False, title='From array')
c.set_lonlat_countours()
c.set_cmap(cleo.get_cm('topo'))
c.set_plot_params(nlevels=256)
c.set_data(h)
c.visualize()
def test_xarray(self):
from salem import open_xr_dataset
go = get_demo_file("hef_srtm.tif")
gs = get_demo_file("hef_srtm_subset.tif")
geo = GeoTiff(go)
go = open_xr_dataset(go)
gs = open_xr_dataset(gs)
gos = go.salem.subset(grid=gs.salem.grid)
ref = gs["data"]
totest = gos["data"]
np.testing.assert_array_equal(ref.shape, (gos.salem.grid.ny, gos.salem.grid.nx))
np.testing.assert_array_equal(ref.shape, totest.shape)
np.testing.assert_array_equal(ref, totest)
rlon, rlat = geo.grid.ll_coordinates
tlon, tlat = go.salem.grid.ll_coordinates
assert_allclose(rlon, tlon)
assert_allclose(rlat, tlat)
def test_center(self):
gm = GoogleCenterMap(center_ll=(10.762660, 46.794221), zoom=13, size_x=500, size_y=500, use_cache=False)
gm.set_roi(shape=get_demo_file("Hintereisferner.shp"))
gm.set_subset(toroi=True, margin=10)
img = gm.get_vardata()
img[np.nonzero(gm.roi == 0)] /= 2.0
# from scipy.misc import toimage
# toimage(img).save(get_demo_file('hef_google_roi.png'))
ref = mpl.image.imread(get_demo_file("hef_google_roi.png"))
rmsd = np.sqrt(np.mean((ref - img) ** 2))
self.assertTrue(rmsd < 0.1)
# assert_allclose(ref, img, atol=2e-2)
gm = GoogleCenterMap(center_ll=(10.762660, 46.794221), zoom=13, size_x=500, size_y=500)
gm.set_roi(shape=get_demo_file("Hintereisferner.shp"))
gm.set_subset(toroi=True, margin=10)
img = gm.get_vardata()
img[np.nonzero(gm.roi == 0)] /= 2.0
rmsd = np.sqrt(np.mean((ref - img) ** 2))
self.assertTrue(rmsd < 0.1)