Merge pull request #215 from guziy/shiftdata_issue2013

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

Author: WeatherGod

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:
@@ -4822,10 +4827,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,59 @@ def test_no_cyc2(self):
         assert (grid==gridout).all()
 
 
+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
+        """
+        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,])
+
+        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)
+
+
 class TestProjectCoords(TestCase):
     def get_data(self):
         lons, lats = np.arange(-180, 180, 20), np.arange(-90, 90, 10)
@@ -130,7 +183,6 @@ def test_results_should_be_same_for_c_and_f_order_arrays(self):
 
 
 
-
 def test():
     """
     Run some tests.