change shiftdata so 1 and 2 points work, add tests.. (#213)

Author: guziy

lib/mpl_toolkits/basemap/__init__.py

@@ -4822,10 +4822,15 @@ def shiftdata(self,lonsin,datain=None,lon_0=None):
             nlons = len(lonsin)
             lonsin = np.where(lonsin > lon_0+180, lonsin-360 ,lonsin)
             lonsin = np.where(lonsin < lon_0-180, lonsin+360 ,lonsin)
-            londiff = np.abs(lonsin[0:-1]-lonsin[1:])
-            londiff_sort = np.sort(londiff)
-            thresh = 360.-londiff_sort[-2]
-            itemindex = len(lonsin)-np.where(londiff>=thresh)[0]
+
+            if nlons > 1:
+                londiff = np.abs(lonsin[0:-1]-lonsin[1:])
+                londiff_sort = np.sort(londiff)
+                thresh = 360.-londiff_sort[-2] if nlons > 2 else 360.0 - londiff_sort[-1]
+                itemindex = len(lonsin)-np.where(londiff>=thresh)[0]
+            else:
+                itemindex = 0
+
             if itemindex:
                 # check to see if cyclic (wraparound) point included
                 # if so, remove it.

lib/mpl_toolkits/basemap/test.py

@@ -13,7 +13,7 @@ def make_array(self):
         u = np.ones((len(lat), len(lon)))
         v = np.zeros((len(lat), len(lon)))
         return u,v,lat,lon
-        
+
     def test_cylindrical(self):
         # Cylindrical case
         B = Basemap()
@@ -22,7 +22,7 @@ def test_cylindrical(self):
         # Check that the vectors are identical.
         assert_almost_equal(ru, u)
         assert_almost_equal(rv, v)
-        
+
     def test_nan(self):
         B = Basemap()
         u,v,lat,lon=self.make_array()
@@ -32,12 +32,12 @@ def test_nan(self):
         assert not np.isnan(ru).any()
         assert_almost_equal(u, ru)
         assert_almost_equal(v, rv)
-        
+
     def test_npstere(self):
         # NP Stereographic case
         B=Basemap(projection='npstere', boundinglat=50., lon_0=0.)
         u,v,lat,lon=self.make_array()
-        v = np.ones((len(lat), len(lon)))    
+        v = np.ones((len(lat), len(lon)))
         ru, rv = B.rotate_vector(u,v, lon, lat)
         assert_almost_equal(ru[2, :],[1,-1,-1,1], 6)
         assert_almost_equal(rv[2, :],[1,1,-1,-1], 6)
@@ -96,6 +96,35 @@ def test_no_cyc2(self):
         assert (grid==gridout).all()
 
 
+class TestShiftdata(TestCase):
+
+    def test_2_points_should_work(self):
+        """
+        Shiftdata should work with 2 points
+        """
+        bm = Basemap(llcrnrlon=0, llcrnrlat=-80, urcrnrlon=360, urcrnrlat=80, projection='mill')
+
+        lons_expected = [10, 15, 20]
+        lonsout = bm.shiftdata(lons_expected)
+        assert_almost_equal(lons_expected, lonsout)
+
+        lonsout_expected = bm.shiftdata([10, 361, 362])
+        lonsout = bm.shiftdata([10, 361])
+        assert_almost_equal(lonsout_expected[:len(lonsout)], lonsout)
+
+    def test_1_point_should_work(self):
+        bm = Basemap(llcrnrlon=0, llcrnrlat=-80, urcrnrlon=360, urcrnrlat=80, projection='mill')
+
+        # should not fail
+        lonsout = bm.shiftdata([361])
+        assert_almost_equal(lonsout, [1.0,])
+
+        lonsout = bm.shiftdata([10])
+        assert_almost_equal(lonsout, [10.0,])
+
+
+
+
 def test():
     """
     Run some tests.