Support for polar stereographic WRF files (#78)

docs/examples/plot_polarstereo.py

@@ -0,0 +1,50 @@
+# -*- coding: utf-8 -*-
+"""
+===============================
+Polar stereographic projections
+===============================
+
+This example illustrates that salem graphics have a small issue at the pole
+and proposes an alternative based on cartopy.
+
+The only problem at the pole is the plotting of the longitude contours. 
+There is probably no easy way so solve this problem in salem,
+hence the suggestion to use cartopy in this case. Note that if the pole isn't
+in the domain, salem works just fine.
+"""
+
+from salem import open_xr_dataset, get_demo_file
+import matplotlib.pyplot as plt
+
+# prepare the figure
+f = plt.figure(figsize=(8, 7))
+
+# WRF polar file(s)
+d1 = open_xr_dataset(get_demo_file('geo_em_d01_polarstereo.nc'))
+d2 = open_xr_dataset(get_demo_file('geo_em_d02_polarstereo.nc'))
+
+# Plot with salem
+ax = plt.subplot(2, 2, 1)
+d1.HGT_M.salem.quick_map(ax=ax, cmap='Oranges')
+ax = plt.subplot(2, 2, 3)
+d2.HGT_M.salem.quick_map(ax=ax, cmap='Oranges')
+
+# Now with cartopy
+proj = d1.salem.cartopy()
+ax = plt.subplot(2, 2, 2, projection=proj)
+ax.coastlines()
+ax.gridlines()
+d1.HGT_M.plot(ax=ax, transform=proj, cmap='Oranges')
+ax.set_extent(d1.salem.grid.extent, crs=proj)
+
+# D2 can use a higher resolution coastline
+proj = d2.salem.cartopy()
+ax = plt.subplot(2, 2, 4, projection=proj)
+ax.coastlines(resolution='50m')
+ax.gridlines()
+d2.HGT_M.plot(ax=ax, transform=proj, cmap='Oranges')
+ax.set_extent(d2.salem.grid.extent, crs=proj)
+
+# make it nice
+plt.tight_layout()
+plt.show()

docs/whats-new.rst

@@ -15,6 +15,7 @@ Enhancements
   grids to lower-res grids (:ref:`sphx_glr_auto_examples_plot_subgrid_mask.py`)
 - new :py:func:`~Grid.to_geometry` method, useful to compute precise
   vector to raster masks (TODO: example showing its use)
+- new projection for WRF files: polar stereographic
 
 
 Bug fixes

salem/datasets.py

@@ -413,7 +413,7 @@ def get_vardata(self, var_id=0, as_xarray=False):
         item = []
         for d in v.dimensions:
             it = slice(None)
-            if d == self.t_dim:
+            if d == self.t_dim and self.sub_t is not None:
                 it = slice(self.sub_t[0], self.sub_t[1]+1)
             elif d == self.y_dim:
                 it = slice(self.sub_y[0], self.sub_y[1]+1)

salem/gis.py

@@ -1355,6 +1355,8 @@ def proj_to_cartopy(proj):
                 cl = ccrs.Mercator
             if v == 'utm':
                 cl = ccrs.UTM
+            if v == 'stere':
+                cl = ccrs.Stereographic
         if k in km_proj:
             kw_proj[km_proj[k]] = v
         if k in km_globe:
@@ -1374,6 +1376,11 @@ def proj_to_cartopy(proj):
         kw_proj.pop('false_northing', None)
         if LooseVersion(cartopy.__version__) < LooseVersion('0.15'):
             kw_proj.pop('latitude_true_scale', None)
+    elif cl.__name__ == 'Stereographic':
+        kw_proj.pop('scale_factor', None)
+        if 'latitude_true_scale' in kw_proj:
+            kw_proj['true_scale_latitude'] = kw_proj['latitude_true_scale']
+            kw_proj.pop('latitude_true_scale', None)
     else:
         kw_proj.pop('latitude_true_scale', None)
 

salem/sio.py

@@ -160,12 +160,20 @@ def _wrf_grid_from_dataset(ds):
         pargs['center_lon'] = ds.CEN_LON
         proj_id = ds.MAP_PROJ
 
+    atol = 1e-4
     if proj_id == 1:
         # Lambert
         p4 = '+proj=lcc +lat_1={lat_1} +lat_2={lat_2} ' \
              '+lat_0={lat_0} +lon_0={lon_0} ' \
              '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
         p4 = p4.format(**pargs)
+    elif proj_id == 2:
+        # Polar stereo
+        p4 = '+proj=stere +lat_ts={lat_1} +lon_0={lon_0} +lat_0=90.0' \
+             '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
+        p4 = p4.format(**pargs)
+        # pyproj and WRF do not agree well close to the pole
+        atol = 5e-3
     elif proj_id == 3:
         # Mercator
         p4 = '+proj=merc +lat_ts={lat_1} ' \
@@ -203,37 +211,26 @@ def _wrf_grid_from_dataset(ds):
         #  Temporary asserts
         if 'XLONG' in ds.variables:
             # Normal WRF
-            mylon, mylat = grid.ll_coordinates
             reflon = ds.variables['XLONG']
             reflat = ds.variables['XLAT']
-            if len(reflon.shape) == 3:
-                reflon = reflon[0, :, :]
-                reflat = reflat[0, :, :]
-            assert np.allclose(reflon, mylon, atol=1e-4)
-            assert np.allclose(reflat, mylat, atol=1e-4)
         elif 'XLONG_M' in ds.variables:
             # geo_em
-            mylon, mylat = grid.ll_coordinates
             reflon = ds.variables['XLONG_M']
             reflat = ds.variables['XLAT_M']
-            if len(reflon.shape) == 3:
-                reflon = reflon[0, :, :]
-                reflat = reflat[0, :, :]
-            assert np.allclose(reflon, mylon, atol=1e-4)
-            assert np.allclose(reflat, mylat, atol=1e-4)
         elif 'lon' in ds.variables:
             # HAR
-            mylon, mylat = grid.ll_coordinates
             reflon = ds.variables['lon']
             reflat = ds.variables['lat']
-            if len(reflon.shape) == 3:
-                reflon = reflon[0, :, :]
-                reflat = reflat[0, :, :]
-            assert np.allclose(reflon, mylon, atol=1e-4)
-            assert np.allclose(reflat, mylat, atol=1e-4)
         else:
             raise RuntimeError("couldn't test for correct WRF lon-lat")
 
+        if len(reflon.shape) == 3:
+            reflon = reflon[0, :, :]
+            reflat = reflat[0, :, :]
+        mylon, mylat = grid.ll_coordinates
+        np.testing.assert_allclose(reflon, mylon, atol=atol)
+        np.testing.assert_allclose(reflat, mylat, atol=atol)
+
     return grid
 
 

salem/tests/__init__.py

@@ -1,5 +1,6 @@
 from __future__ import division
 import unittest
+from distutils.version import LooseVersion
 import os
 from salem import python_version
 from six.moves.urllib.request import urlopen
@@ -52,8 +53,10 @@ def has_internet():
 try:
     import matplotlib
     has_matplotlib = True
+    mpl_version = LooseVersion(matplotlib.__version__)
 except ImportError:
     has_matplotlib = False
+    mpl_version = LooseVersion('0.0.0')
 
 try:
     import rasterio
@@ -84,6 +87,12 @@ def requires_matplotlib(test):
     return test if has_matplotlib else unittest.skip(msg)(test)
 
 
+def requires_matplotlibv2(test):
+    msg = "requires matplotlib v2+"
+    return test if mpl_version >= LooseVersion('2.0.0') \
+        else unittest.skip(msg)(test)
+
+
 def requires_motionless(test):
     msg = "requires motionless"
     return test if has_motionless else unittest.skip(msg)(test)

salem/tests/test_datasets.py

@@ -369,6 +369,26 @@ def test_wrf(self):
         np.testing.assert_allclose(reflon, mylon, atol=1e-4)
         np.testing.assert_allclose(reflat, mylat, atol=1e-4)
 
+    @requires_pandas
+    @requires_geopandas
+    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)
+
     @requires_pandas
     def test_longtime(self):
         """There was a bug with time"""

salem/tests/test_graphics.py

@@ -25,7 +25,8 @@
     read_shapefile_to_grid, GeoTiff, GoogleCenterMap, GoogleVisibleMap, \
     open_wrf_dataset, open_xr_dataset, python_version, cache_dir
 from salem.utils import get_demo_file
-from salem.tests import requires_matplotlib, requires_cartopy
+from salem.tests import (requires_matplotlib, requires_cartopy,
+                         requires_matplotlibv2)
 
 # Globals
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -704,3 +705,42 @@ def test_cartopy():
     ax.scatter(ds.lon, ds.lat, transform=cartopy.crs.PlateCarree())
 
     return fig
+
+
+@requires_matplotlibv2
+@requires_cartopy
+@pytest.mark.mpl_image_compare(baseline_dir=baseline_dir, tolerance=5)
+def test_cartopy_polar():
+
+    import cartopy
+
+    fig = plt.figure(figsize=(8, 8))
+
+    ods = open_wrf_dataset(get_demo_file('geo_em_d02_polarstereo.nc')).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)
+
+    plt.tight_layout()
+    return fig

salem/tests/test_misc.py

@@ -982,4 +982,20 @@ def test_mercator(self):
             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)
+            assert_allclose(lat, ds['XLAT_M'][0, ...], atol=1e-4)
+
+    @requires_geopandas
+    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)

salem/wrftools.py

@@ -639,13 +639,16 @@ def geogrid_simulator(fpath, do_maps=True):
     # define projection
     if map_proj == 'LAMBERT':
         pwrf = '+proj=lcc +lat_1={lat_1} +lat_2={lat_2} ' \
-             '+lat_0={lat_0} +lon_0={lon_0} ' \
-             '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
+               '+lat_0={lat_0} +lon_0={lon_0} ' \
+               '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
         pwrf = pwrf.format(**pargs)
     elif map_proj == 'MERCATOR':
-        pwrf = '+proj=merc +lat_ts={lat_1} ' \
-             '+lon_0={lon_0} ' \
-             '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
+        pwrf = '+proj=merc +lat_ts={lat_1} +lon_0={lon_0} ' \
+               '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
+        pwrf = pwrf.format(**pargs)
+    elif map_proj == 'POLAR':
+        pwrf = '+proj=stere +lat_ts={lat_1} +lat_0=90.0 +lon_0={lon_0} ' \
+               '+x_0=0 +y_0=0 +a=6370000 +b=6370000'
         pwrf = pwrf.format(**pargs)
     else:
         raise NotImplementedError('WRF proj not implemented yet: '