improvement of basemap.addcyclic

Changelog

@@ -6,6 +6,7 @@ version 1.0.8 (not yet released)
   ability to fill counties with specified matplotlib color argument.
 * fix drawgreatcircle bug so that lines exiting and reentering a projection
   region don't draw horizontally across the map. 
+* change addcyclic so that it works for ND-arrays with an optional 'axis' parameter.
 
 version 1.0.7 (git tag v1.0.7rel)
 ---------------------------------

README.md

@@ -108,4 +108,4 @@ Jeff Whitaker <jeffrey.s.whitaker@noaa.gov>
 
 ##Thanks
 
-Special thanks to John Hunter, Andrew Straw, Eric Firing, Rob Hetland, Scott Sinclair, Ivan Lima, Erik Andersen, Michael Hearne, Jesper Larsen, Ryan May, David Huard, Mauro Cavalcanti, Chris Murphy, Pierre Gerard-Marchant, Christoph Gohlke, Eric Bruning, Stephane Raynaud, Tom Loredo, Patrick Marsh, Phil Elson, and Henry Hammond for valuable contributions.
+Special thanks to John Hunter, Andrew Straw, Eric Firing, Rob Hetland, Scott Sinclair, Ivan Lima, Erik Andersen, Michael Hearne, Jesper Larsen, Ryan May, David Huard, Mauro Cavalcanti, Chris Murphy, Pierre Gerard-Marchant, Christoph Gohlke, Eric Bruning, Stephane Raynaud, Tom Loredo, Patrick Marsh, Jonas Bluethgen, Phil Elson, and Henry Hammond for valuable contributions.

lib/mpl_toolkits/basemap/__init__.py

@@ -5087,28 +5087,25 @@ def shiftgrid(lon0,datain,lonsin,start=True,cyclic=360.0):
     dataout[...,i0_shift:] = datain[...,start_idx:i0+start_idx]
     return dataout,lonsout
 
-def addcyclic(arrin,lonsin):
+def addcyclic(*arr,**axiskwarg):
     """
-    ``arrout, lonsout = addcyclic(arrin, lonsin)``
-    adds cyclic (wraparound) point in longitude to ``arrin`` and ``lonsin``,
-    assumes longitude is the right-most dimension of ``arrin``.
+    ``arrout, lonsout = addcyclic(arrin,lonsin,axis=-1)``
+    adds cyclic (wraparound) points in longitude to one or several arrays, 
+    (e.g. ``arrin`` and ``lonsin``),
+    where ``axis`` sets the dimension longitude is in (optional, default: right-most).
     """
-    nlons = arrin.shape[-1]
-    newshape = list(arrin.shape)
-    newshape[-1] += 1
-    if ma.isMA(arrin):
-        arrout  = ma.zeros(newshape,arrin.dtype)
+    # get axis keyword argument (default: -1)
+    axis = axiskwarg.get('axis',-1)
+    # define function for a single grid array
+    def _addcyclic(a):
+        aT = np.swapaxes(a,0,axis)
+        idx = np.append(np.arange(aT.shape[0]),0)
+        return np.swapaxes(aT[idx],axis,0)
+    # process array(s)
+    if len(arr) == 1:
+        return _addcyclic(arr[0])
     else:
-        arrout  = np.zeros(newshape,arrin.dtype)
-    arrout[...,0:nlons] = arrin[:]
-    arrout[...,nlons] = arrin[...,0]
-    if ma.isMA(lonsin):
-        lonsout = ma.zeros(nlons+1,lonsin.dtype)
-    else:
-        lonsout = np.zeros(nlons+1,lonsin.dtype)
-    lonsout[0:nlons] = lonsin[:]
-    lonsout[nlons]  = lonsin[-1] + lonsin[1]-lonsin[0]
-    return arrout,lonsout
+        return map(_addcyclic,arr)
 
 def _choosecorners(width,height,**kwargs):
     """