Merge pull request #332 from guziy/shiftdata_fix

Shiftdata fix

Author: WeatherGod

lib/mpl_toolkits/basemap/__init__.py

@@ -526,7 +526,8 @@ def with_transform(self,x,y,data,*args,**kwargs):
             # shift data to map projection region for
             # cylindrical and pseudo-cylindrical projections.
             if self.projection in _cylproj or self.projection in _pseudocyl:
-                x, data = self.shiftdata(x, data)
+                x, data = self.shiftdata(x, data,
+                                         fix_wrap_around=plotfunc.__name__ not in ["scatter"])
             # convert lat/lon coords to map projection coords.
             x, y = self(x,y)
         return plotfunc(self,x,y,data,*args,**kwargs)
@@ -544,7 +545,7 @@ def with_transform(self,x,y,*args,**kwargs):
             # cylindrical and pseudo-cylindrical projections.
             if self.projection in _cylproj or self.projection in _pseudocyl:
                 if x.ndim == 1:
-                    x = self.shiftdata(x)
+                    x = self.shiftdata(x, fix_wrap_around=plotfunc.__name__ not in ["scatter"])
                 elif x.ndim == 0:
                     if x > 180:
                         x = x - 360.
@@ -4723,7 +4724,7 @@ def _ax_plt_from_kw(self, kw):
                 _ax = plt.gca()
         return _ax, plt
 
-    def shiftdata(self,lonsin,datain=None,lon_0=None):
+    def shiftdata(self,lonsin,datain=None,lon_0=None,fix_wrap_around=True):
         """
         Shift longitudes (and optionally data) so that they match map projection region.
         Only valid for cylindrical/pseudo-cylindrical global projections and data
@@ -4746,6 +4747,13 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
         datain           original 1-d or 2-d data. Default None.
         lon_0            center of map projection region. Defaut None,
                          given by current map projection.
+        fix_wrap_around  if True reindex (if required) longitudes (and data) to
+                         avoid jumps caused by remapping of longitudes of
+                         points from outside of the [lon_0-180, lon_0+180]
+                         interval back into the interval.
+                         If False do not reindex longitudes and data, but do
+                         make sure that longitudes are in the
+                         [lon_0-180, lon_0+180] range.
         ==============   ====================================================
 
         if datain given, returns ``dataout,lonsout`` (data and longitudes shifted to fit in interval
@@ -4784,7 +4792,7 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
 
             # if no shift necessary, itemindex will be
             # empty, so don't do anything
-            if itemindex:
+            if fix_wrap_around and itemindex:
                 # check to see if cyclic (wraparound) point included
                 # if so, remove it.
                 if np.abs(lonsin1[0]-lonsin1[-1]) < 1.e-4:
@@ -4811,13 +4819,7 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
                         datain_save[:,1:] = datain
                         datain_save[:,0] = datain[:,-1]
                         datain = datain_save
-                # mask points outside
-                # map region so they don't wrap back in the domain.
-                mask = np.logical_or(lonsin<lon_0-180,lonsin>lon_0+180)
-                lonsin = np.where(mask,1.e30,lonsin)
-                if datain is not None and mask.any():
-                    # superimpose on existing mask
-                    datain = ma.masked_where(mask, datain)
+
         # 1-d data.
         elif lonsin.ndim == 1:
             nlons = len(lonsin)
@@ -4832,7 +4834,7 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
             else:
                 itemindex = 0
 
-            if itemindex:
+            if fix_wrap_around and itemindex:
                 # check to see if cyclic (wraparound) point included
                 # if so, remove it.
                 if np.abs(lonsin[0]-lonsin[-1]) < 1.e-4:
@@ -4856,12 +4858,14 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
                         datain_save[1:] = datain
                         datain_save[0] = datain[-1]
                         datain = datain_save
-                # mask points outside
-                # map region so they don't wrap back in the domain.
-                mask = np.logical_or(lonsin<lon_0-180,lonsin>lon_0+180)
-                lonsin = np.where(mask,1.e30,lonsin)
-                if datain is not None and mask.any():
-                    datain = ma.masked_where(mask, datain)
+
+        # mask points outside
+        # map region so they don't wrap back in the domain.
+        mask = np.logical_or(lonsin<lon_0-180,lonsin>lon_0+180)
+        lonsin = np.where(mask,1.e30,lonsin)
+        if datain is not None and mask.any():
+            datain = ma.masked_where(mask, datain)
+
         if datain is not None:
             return lonsin, datain
         else:

lib/mpl_toolkits/basemap/test.py

@@ -38,7 +38,7 @@ def test_cylindrical(self):
     def test_nan(self):
         B = Basemap()
         u,v,lat,lon=self.make_array()
-        # Set one element to 0, so that the vector magnitude is 0. 
+        # Set one element to 0, so that the vector magnitude is 0.
         u[1,1] = 0.
         ru, rv = B.rotate_vector(u,v, lon, lat)
         assert not np.isnan(ru).any()
@@ -117,6 +117,42 @@ def _get_2d_lons(self, lons1d):
         lats = [10, ] * len(lons1d)
         return np.meshgrid(lons1d, lats)[0]
 
+    def test_non_monotonous_longitudes(self):
+        """
+        when called for scatter, the longitudes passed to shiftdata are
+        not necessarily monotonous...
+        """
+        lons = [179, 180, 180, 0, 290, 10, 320, -150, 350, -250, 250]
+        bm = Basemap(lon_0=0)
+
+        # before, having several break points would cause the exception,
+        # inside the shiftdata method called from scatter method.
+        self.assertRaises(ValueError, bm.shiftdata, lons, fix_wrap_around=True)
+
+        lons_new = bm.shiftdata(lons, fix_wrap_around=False)
+
+        # Check if the modified longitudes are inside of the projection region
+        for lon in lons_new:
+            assert lon >= bm.projparams["lon_0"] - 180
+            assert lon <= bm.projparams["lon_0"] + 180
+
+
+    def test_shiftdata_on_monotonous_lons(self):
+        """
+        Test that shiftdata with fix_wrap_around keyword added works as before,
+        when it is True
+        """
+
+        bm = Basemap(lon_0=0)
+
+        lons_in = [120, 140, 160, 180, 200, 220]
+        lons_out_expect = [-160, -140,  120,  140,  160,  180]
+        lons_out = bm.shiftdata(lons_in, fix_wrap_around=True)
+
+        assert_almost_equal(lons_out, lons_out_expect)
+
+
+
     def test_2_points_should_work(self):
         """
         Shiftdata should work with 2 points
@@ -160,7 +196,7 @@ def test_less_than_n_by_3_points_should_work(self):
         lonsout = bm.shiftdata(lonsin[:, :2])
         assert_almost_equal(lonsout_expected, lonsout)
 
-@skipIf(PY3 and LooseVersion(pyproj.__version__) <= LooseVersion("1.9.4"), 
+@skipIf(PY3 and LooseVersion(pyproj.__version__) <= LooseVersion("1.9.4"),
         "Test skipped in Python 3.x with pyproj version 1.9.4 and below.")
 class TestProjectCoords(TestCase):
     def get_data(self):
@@ -224,7 +260,7 @@ def test():
     import unittest
 
     from mpl_toolkits.basemap.diagnostic import package_versions
-    
+
     if '--verbose' in sys.argv or '-v' in sys.argv:
         pkg_vers = package_versions()
         print('Basemaps installed package versions:')