Merge pull request #31 from Chilipp/master

fix plotting of curvilinear data with bounds

Author: Chilipp

CHANGELOG.rst

@@ -13,6 +13,8 @@ Fixed
   `#29 <https://github.com/psyplot/psy-maps/pull/29>`__).
 - fixed plotting of data with 3D bounds (see
   `#30 <https://github.com/psyplot/psy-maps/pull/30>`__)
+- fixed plotting of curvilinear data with 3D bounds (see
+  `#31 <https://github.com/psyplot/psy-maps/pull/31>`__)
 
 v1.3.0
 ======

psy_maps/plotters.py

@@ -1679,6 +1679,10 @@ def _polycolor(self):
             # However, in order to wrap the boundaries correctly, we have to
             # identify the corresponding grid cells and then use the standard
             # matplotlib transform.
+            if xb.ndim > 2:
+                xb = xb.reshape((-1, xb.shape[-1]))
+                yb = yb.reshape((-1, yb.shape[-1]))
+                arr = arr.reshape(-1)
             if isinstance(t, wrap_proj_types) and 'lon_0' in proj.proj4_params:
                 # We adopt and copy some code from the methodology of cartopy
                 # _pcolormesh_patched method of the geoaxes. As such, we
@@ -1717,9 +1721,6 @@ def _polycolor(self):
                     yb_wrap = yb[mask]
                     xb = xb[~mask]
                     yb = yb[~mask]
-            if xb.ndim > 2:
-                xb = xb.reshape((-1, xb.shape[-1]))
-                yb = yb.reshape((-1, yb.shape[-1]))
             self.logger.debug('Making plot with %i cells', arr.size)
             transformed = proj.transform_points(
                 t, xb.ravel(), yb.ravel())[..., :2].reshape(xb.shape + (2,))
@@ -1791,6 +1792,11 @@ def update(self, value):
         proj = self.ax.projection
         # HACK: We remove the cells at the boundary of the map projection
         xb_wrap = None
+
+        if xb.ndim > 2:
+            xb = xb.reshape((-1, xb.shape[-1]))
+            yb = yb.reshape((-1, yb.shape[-1]))
+
         if isinstance(t, wrap_proj_types) and 'lon_0' in proj.proj4_params:
             # See the :meth:`MapPlot2D._polycolor` method for a documentation
             # of the steps
@@ -1822,20 +1828,48 @@ def update(self, value):
         transformed = proj.transform_points(t, xb.ravel(), yb.ravel())
         xb = transformed[..., 0].reshape(orig_shape)
         yb = transformed[..., 1].reshape(orig_shape)
-        if xb.ndim > 2:
-            xb = xb.reshape((-1, xb.shape[-1]))
-            yb = yb.reshape((-1, yb.shape[-1]))
         # We insert nan values in the flattened edges arrays rather than
         # plotting the grid cells separately as it considerably speeds-up code
         # execution.
         n = len(xb)
         xb = np.c_[xb, xb[:, :1], [[np.nan]] * n].ravel()
         yb = np.c_[yb, yb[:, :1], [[np.nan]] * n].ravel()
+
         if isinstance(value, dict):
             self._artists = self.ax.plot(xb, yb, **value)
         else:
             self._artists = self.ax.plot(xb, yb, value)
+
         if xb_wrap is not None:
+            # first we drop all grid cells that have any NaN in their bounds
+            # as we do not know, how to handle these
+            valid = (~np.isnan(xb_wrap).any(-1)) & (~np.isnan(yb_wrap).any(-1))
+            xb_wrap = xb_wrap[valid]
+            yb_wrap = yb_wrap[valid]
+
+            if isinstance(t, ccrs.PlateCarree):
+
+                # identify the grid cells at the boundary
+                xdiff2min = xb_wrap - xb_wrap.min(axis=-1, keepdims=True)
+                cross_world_mask = np.any(np.abs(xdiff2min) > 180, -1)
+                if cross_world_mask.any():
+                    cross_world_x = xb_wrap[cross_world_mask]
+                    cross_world_y = yb_wrap[cross_world_mask]
+                    cross_world_diff = xdiff2min[cross_world_mask]
+                    xdiff2max = cross_world_x - cross_world_x.max(
+                        axis=-1, keepdims=True)
+
+                    leftx = cross_world_x.copy()
+                    leftx[cross_world_diff > 180] -= 360
+
+                    rightx = cross_world_x.copy()
+                    rightx[xdiff2max < -180] += 360
+
+                    xb_wrap = np.r_[xb_wrap[~cross_world_mask], leftx, rightx]
+                    yb_wrap = np.r_[
+                        yb_wrap[~cross_world_mask], cross_world_y, cross_world_y
+                    ]
+
             n = len(xb_wrap)
             xb_wrap = np.c_[xb_wrap, xb_wrap[:, :1], [[np.nan]] * n].ravel()
             yb_wrap = np.c_[yb_wrap, yb_wrap[:, :1], [[np.nan]] * n].ravel()

tests/curvilinear-with-bounds.nc

@@ -1774,5 +1774,30 @@ def test_datagrid_3D_bounds():
         assert abs(ymax - ymin - 52) < 2
 
 
+def test_plot_curvilinear_datagrid(tmpdir):
+    """Test if the there is no datagrid plotted over land
+
+    This implicitly checks, if grid cells at the boundary are warped correctly.
+    The file ``'curvilinear-with-bounds.nc'`` contains a variable on a
+    curvilinear grid that is only defined over the ocean (derived from MPI-OM).
+    Within this test, we focus on a region over land far away from
+    the ocean (Czech Republic) where there are no grid cells. If the datagrid
+    is plotted correctly, it should be all white.
+    """
+    from matplotlib.testing.compare import compare_images
+    fname = os.path.join(bt.test_dir, 'curvilinear-with-bounds.nc')
+    # make a white plot without datagrid
+    kws = dict(plot=None, xgrid=False, ygrid=False, map_extent='Czech Republic')
+    with psy.plot.mapplot(fname, **kws) as sp:
+        sp.export(str(tmpdir / "ref.png"))  # this is all white
+    # now draw the datagrid, it should still be empty (as the input file only
+    # defines the data over the ocean)
+    with psy.plot.mapplot(fname, datagrid='k-', **kws) as sp:
+        sp.export(str(tmpdir / "test.png"))  # this should be all white, too
+    results = compare_images(
+        str(tmpdir / "ref.png"), str(tmpdir / "test.png"), tol=1)
+    assert results is None, results
+
+
 if __name__ == '__main__':
     bt.RefTestProgram()