Merge pull request numpy#4454 from jurnix/namedargs

juliantaylor · juliantaylor · commit 10098daf387b · 2014-06-08T16:29:06.000+02:00
ENH: apply_along_axis accepts named arguments
diff --git a/numpy/lib/shape_base.py b/numpy/lib/shape_base.py
@@ -12,7 +12,7 @@
 from numpy.core.fromnumeric import product, reshape
 from numpy.core import hstack, vstack, atleast_3d
 
-def apply_along_axis(func1d,axis,arr,*args):
+def apply_along_axis(func1d, axis, arr, *args, **kwargs):
     """
     Apply a function to 1-D slices along the given axis.
 
@@ -30,6 +30,11 @@ def apply_along_axis(func1d,axis,arr,*args):
         Input array.
     args : any
         Additional arguments to `func1d`.
+    kwargs: any
+        Additional named arguments to `func1d`.
+
+        .. versionadded:: 1.9.0
+
 
     Returns
     -------
@@ -78,7 +83,7 @@ def apply_along_axis(func1d,axis,arr,*args):
     i[axis] = slice(None, None)
     outshape = asarray(arr.shape).take(indlist)
     i.put(indlist, ind)
-    res = func1d(arr[tuple(i.tolist())],*args)
+    res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
     #  if res is a number, then we have a smaller output array
     if isscalar(res):
         outarr = zeros(outshape, asarray(res).dtype)
@@ -94,7 +99,7 @@ def apply_along_axis(func1d,axis,arr,*args):
                 ind[n] = 0
                 n -= 1
             i.put(indlist, ind)
-            res = func1d(arr[tuple(i.tolist())],*args)
+            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
             outarr[tuple(ind)] = res
             k += 1
         return outarr
@@ -115,7 +120,7 @@ def apply_along_axis(func1d,axis,arr,*args):
                 ind[n] = 0
                 n -= 1
             i.put(indlist, ind)
-            res = func1d(arr[tuple(i.tolist())],*args)
+            res = func1d(arr[tuple(i.tolist())], *args, **kwargs)
             outarr[tuple(i.tolist())] = res
             k += 1
         return outarr
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
@@ -479,6 +479,16 @@ def myfunc(b):
         xa = apply_along_axis(myfunc, 2, a)
         assert_equal(xa, [[1, 4], [7, 10]])
 
+   # Tests kwargs functions
+    def test_3d_kwargs(self):
+        a = arange(12).reshape(2, 2, 3)
+
+        def myfunc(b, offset=0):
+            return b[1+offset]
+
+        xa = apply_along_axis(myfunc, 2, a, offset=1)
+        assert_equal(xa, [[2, 5], [8, 11]])
+
 
 class TestApplyOverAxes(TestCase):
     # Tests apply_over_axes
