issue #213, fix and tests for the 2d case

Author: guziy

lib/mpl_toolkits/basemap/__init__.py

@@ -4777,10 +4777,15 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
             lonsin1 = lonsin[0,:]
             lonsin1 = np.where(lonsin1 > lon_0+180, lonsin1-360 ,lonsin1)
             lonsin1 = np.where(lonsin1 < lon_0-180, lonsin1+360 ,lonsin1)
-            londiff = np.abs(lonsin1[0:-1]-lonsin1[1:])
-            londiff_sort = np.sort(londiff)
-            thresh = 360.-londiff_sort[-2]
-            itemindex = nlons-np.where(londiff>=thresh)[0]
+            if nlons > 1:
+                londiff = np.abs(lonsin1[0:-1]-lonsin1[1:])
+                londiff_sort = np.sort(londiff)
+                thresh = 360.-londiff_sort[-2] if nlons > 2 else 360.-londiff_sort[-1]
+                itemindex = nlons-np.where(londiff>=thresh)[0]
+            else:
+                lonsin[0, :] = lonsin1
+                itemindex = 0
+
             # if no shift necessary, itemindex will be
             # empty, so don't do anything
             if itemindex:

lib/mpl_toolkits/basemap/test.py

@@ -98,6 +98,13 @@ def test_no_cyc2(self):
 
 class TestShiftdata(TestCase):
 
+    def _get_2d_lons(self, lons1d):
+        """
+        Generate a 2d grid
+        """
+        lats = [10, ] * len(lons1d)
+        return np.meshgrid(lons1d, lats)[0]
+
     def test_2_points_should_work(self):
         """
         Shiftdata should work with 2 points
@@ -122,7 +129,24 @@ def test_1_point_should_work(self):
         lonsout = bm.shiftdata([10])
         assert_almost_equal(lonsout, [10.0,])
 
+        lonsin = np.array([361.0])
+        lonsin.shape = (1, 1)
+        lonsout = bm.shiftdata(lonsin)
+        assert_almost_equal(lonsout.squeeze(), [1.0,])
+
+    def test_less_than_n_by_3_points_should_work(self):
+        bm = Basemap(llcrnrlon=0, llcrnrlat=-80, urcrnrlon=360, urcrnrlat=80, projection='mill')
+        lons_expected = self._get_2d_lons([10, 15, 20])
+
+        # nothing should change
+        lonsout = bm.shiftdata(lons_expected)
+        assert_almost_equal(lons_expected, lonsout)
 
+        # shift n x 3 and n x 2 grids and compare results over overlapping region
+        lonsin = self._get_2d_lons([10, 361, 362])
+        lonsout_expected = bm.shiftdata(lonsin)[:, :2]
+        lonsout = bm.shiftdata(lonsin[:, :2])
+        assert_almost_equal(lonsout_expected, lonsout)
 
 
 def test():