Merge branch 'vectorize'

2 files changed, 162 insertions, 57 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 3b8ddf099..3f293ae4c 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -117,15 +117,16 @@ def histogram(a, bins=10, range=None, normed=False, weights=None):
     if weights is not None:
         weights = asarray(weights)
         if np.any(weights.shape != a.shape):
-            raise ValueError, 'weights should have the same shape as a.'
+            raise ValueError(
+                    'weights should have the same shape as a.')
         weights = weights.ravel()
     a =  a.ravel()
 
     if (range is not None):
         mn, mx = range
         if (mn > mx):
-            raise AttributeError, \
-                'max must be larger than min in range parameter.'
+            raise AttributeError(
+                'max must be larger than min in range parameter.')
 
     if not iterable(bins):
         if range is None:
@@ -138,7 +139,8 @@ def histogram(a, bins=10, range=None, normed=False, weights=None):
     else:
         bins = asarray(bins)
         if (np.diff(bins) < 0).any():
-            raise AttributeError, 'bins must increase monotonically.'
+            raise AttributeError(
+                    'bins must increase monotonically.')
 
     # Histogram is an integer or a float array depending on the weights.
     if weights is None:
@@ -244,8 +246,9 @@ def histogramdd(sample, bins=10, range=None, normed=False, weights=None):
     try:
         M = len(bins)
         if M != D:
-            raise AttributeError, 'The dimension of bins must be equal ' \
-                                  'to the dimension of the sample x.'
+            raise AttributeError(
+                    'The dimension of bins must be equal'\
+                    ' to the dimension of the sample x.')
     except TypeError:
         bins = D*[bins]
 
@@ -334,12 +337,13 @@ def histogramdd(sample, bins=10, range=None, normed=False, weights=None):
         s = hist.sum()
         for i in arange(D):
             shape = ones(D, int)
-            shape[i] = nbin[i]-2
+            shape[i] = nbin[i] - 2
             hist = hist / dedges[i].reshape(shape)
         hist /= s
 
-    if (hist.shape != nbin-2).any():
-        raise RuntimeError('Internal Shape Error')
+    if (hist.shape != nbin - 2).any():
+        raise RuntimeError(
+                "Internal Shape Error")
     return hist, edges
 
 
@@ -429,23 +433,30 @@ def average(a, axis=None, weights=None, returned=False):
         # Sanity checks
         if a.shape != wgt.shape :
             if axis is None :
-                raise TypeError, "Axis must be specified when shapes of a and weights differ."
+                raise TypeError(
+                        "Axis must be specified when shapes of a "\
+                        "and weights differ.")
             if wgt.ndim != 1 :
-                raise TypeError, "1D weights expected when shapes of a and weights differ."
+                raise TypeError(
+                        "1D weights expected when shapes of a and "\
+                        "weights differ.")
             if wgt.shape[0] != a.shape[axis] :
-                raise ValueError, "Length of weights not compatible with specified axis."
+                raise ValueError(
+                        "Length of weights not compatible with "\
+                        "specified axis.")
 
             # setup wgt to broadcast along axis
-            wgt = np.array(wgt, copy=0, ndmin=a.ndim).swapaxes(-1,axis)
+            wgt = np.array(wgt, copy=0, ndmin=a.ndim).swapaxes(-1, axis)
 
         scl = wgt.sum(axis=axis)
         if (scl == 0.0).any():
-            raise ZeroDivisionError, "Weights sum to zero, can't be normalized"
+            raise ZeroDivisionError(
+                    "Weights sum to zero, can't be normalized")
 
-        avg = np.multiply(a,wgt).sum(axis)/scl
+        avg = np.multiply(a, wgt).sum(axis)/scl
 
     if returned:
-        scl = np.multiply(avg,0) + scl
+        scl = np.multiply(avg, 0) + scl
         return avg, scl
     else:
         return avg
@@ -513,7 +524,8 @@ def asarray_chkfinite(a):
     a = asarray(a)
     if (a.dtype.char in typecodes['AllFloat']) \
            and (_nx.isnan(a).any() or _nx.isinf(a).any()):
-        raise ValueError, "array must not contain infs or NaNs"
+        raise ValueError(
+                "array must not contain infs or NaNs")
     return a
 
 def piecewise(x, condlist, funclist, *args, **kw):
@@ -612,8 +624,8 @@ def piecewise(x, condlist, funclist, *args, **kw):
         condlist.append(~totlist)
         n += 1
     if (n != n2):
-        raise ValueError, "function list and condition list " \
-                          "must be the same"
+        raise ValueError(
+                "function list and condition list must be the same")
     zerod = False
     # This is a hack to work around problems with NumPy's
     #  handling of 0-d arrays and boolean indexing with
@@ -683,7 +695,8 @@ def select(condlist, choicelist, default=0):
     n = len(condlist)
     n2 = len(choicelist)
     if n2 != n:
-        raise ValueError, "list of cases must be same length as list of conditions"
+        raise ValueError(
+                "list of cases must be same length as list of conditions")
     choicelist = [default] + choicelist
     S = 0
     pfac = 1
@@ -791,7 +804,8 @@ def gradient(f, *varargs):
     elif n == N:
         dx = list(varargs)
     else:
-        raise SyntaxError, "invalid number of arguments"
+        raise SyntaxError(
+                "invalid number of arguments")
 
     # use central differences on interior and first differences on endpoints
 
@@ -885,7 +899,8 @@ def diff(a, n=1, axis=-1):
     if n == 0:
         return a
     if n < 0:
-        raise ValueError, 'order must be non-negative but got ' + repr(n)
+        raise ValueError(
+                "order must be non-negative but got " + repr(n))
     a = asanyarray(a)
     nd = len(a.shape)
     slice1 = [slice(None)]*nd
@@ -1596,22 +1611,49 @@ def disp(mesg, device=None, linefeed=True):
 
 # return number of input arguments and
 #  number of default arguments
-import re
+
 def _get_nargs(obj):
+    import re
+
+    terr = re.compile(r'.*? takes (exactly|at least) (?P<exargs>(\d+)|(\w+))' +
+            r' argument(s|) \((?P<gargs>(\d+)|(\w+)) given\)')
+    def _convert_to_int(strval):
+        try:
+            result = int(strval)
+        except ValueError:
+            if strval=='zero':
+                result = 0
+            elif strval=='one':
+                result = 1
+            elif strval=='two':
+                result = 2
+            # How high to go? English only?
+            else:
+                raise
+        return result
+
     if not callable(obj):
-        raise TypeError, "Object is not callable."
+        raise TypeError(
+                "Object is not callable.")
     if sys.version_info[0] >= 3:
+        # inspect currently fails for binary extensions
+        # like math.cos. So fall back to other methods if
+        # it fails.
         import inspect
-        spec = inspect.getargspec(obj)
-        nargs = len(spec.args)
-        if spec.defaults:
-            ndefaults = len(spec.defaults)
-        else:
-            ndefaults = 0
-        if inspect.ismethod(obj):
-            nargs -= 1
-        return nargs, ndefaults
-    elif hasattr(obj,'func_code'):
+        try:
+            spec = inspect.getargspec(obj)
+            nargs = len(spec.args)
+            if spec.defaults:
+                ndefaults = len(spec.defaults)
+            else:
+                ndefaults = 0
+            if inspect.ismethod(obj):
+                nargs -= 1
+            return nargs, ndefaults
+        except:
+            pass
+
+    if hasattr(obj,'func_code'):
         fcode = obj.func_code
         nargs = fcode.co_argcount
         if obj.func_defaults is not None:
@@ -1621,19 +1663,21 @@ def _get_nargs(obj):
         if isinstance(obj, types.MethodType):
             nargs -= 1
         return nargs, ndefaults
-    terr = re.compile(r'.*? takes exactly (?P<exargs>\d+) argument(s|) \((?P<gargs>\d+) given\)')
+
     try:
         obj()
         return 0, 0
     except TypeError, msg:
         m = terr.match(str(msg))
         if m:
-            nargs = int(m.group('exargs'))
-            ndefaults = int(m.group('gargs'))
+            nargs = _convert_to_int(m.group('exargs'))
+            ndefaults = _convert_to_int(m.group('gargs'))
             if isinstance(obj, types.MethodType):
                 nargs -= 1
             return nargs, ndefaults
-    raise ValueError, 'failed to determine the number of arguments for %s' % (obj)
+
+    raise ValueError(
+            "failed to determine the number of arguments for %s" % (obj))
 
 
 class vectorize(object):
@@ -1717,11 +1761,13 @@ class vectorize(object):
             self.otypes = otypes
             for char in self.otypes:
                 if char not in typecodes['All']:
-                    raise ValueError, "invalid otype specified"
+                    raise ValueError(
+                            "invalid otype specified")
         elif iterable(otypes):
             self.otypes = ''.join([_nx.dtype(x).char for x in otypes])
         else:
-            raise ValueError, "output types must be a string of typecode characters or a list of data-types"
+            raise ValueError(
+                    "Invalid otype specification")
         self.lastcallargs = 0
 
     def __call__(self, *args):
@@ -1730,8 +1776,8 @@ class vectorize(object):
         nargs = len(args)
         if self.nin:
             if (nargs > self.nin) or (nargs < self.nin_wo_defaults):
-                raise ValueError, "mismatch between python function inputs"\
-                      " and received arguments"
+                raise ValueError(
+                        "Invalid number of arguments")
 
         # we need a new ufunc if this is being called with more arguments.
         if (self.lastcallargs != nargs):
@@ -3052,7 +3098,8 @@ def delete(arr, obj, axis=None):
     if isinstance(obj, (int, long, integer)):
         if (obj < 0): obj += N
         if (obj < 0 or obj >=N):
-            raise ValueError, "invalid entry"
+            raise ValueError(
+                    "invalid entry")
         newshape[axis]-=1;
         new = empty(newshape, arr.dtype, arr.flags.fnc)
         slobj[axis] = slice(None, obj)
@@ -3197,8 +3244,9 @@ def insert(arr, obj, values, axis=None):
     if isinstance(obj, (int, long, integer)):
         if (obj < 0): obj += N
         if obj < 0 or obj > N:
-            raise ValueError, "index (%d) out of range (0<=index<=%d) "\
-                  "in dimension %d" % (obj, N, axis)
+            raise ValueError(
+                    "index (%d) out of range (0<=index<=%d) "\
+                    "in dimension %d" % (obj, N, axis))
         newshape[axis] += 1;
         new = empty(newshape, arr.dtype, arr.flags.fnc)
         slobj[axis] = slice(None, obj)
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index 330a1a51f..6381c0a72 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -23,6 +23,7 @@ class TestAny(TestCase):
         assert_array_equal(sometrue(y1, axis=0), [1, 1, 0])
         assert_array_equal(sometrue(y1, axis=1), [0, 1, 1])
 
+
 class TestAll(TestCase):
     def test_basic(self):
         y1 = [0, 1, 1, 0]
@@ -39,6 +40,7 @@ class TestAll(TestCase):
         assert_array_equal(alltrue(y1, axis=0), [0, 0, 1])
         assert_array_equal(alltrue(y1, axis=1), [0, 0, 1])
 
+
 class TestAverage(TestCase):
     def test_basic(self):
         y1 = array([1, 2, 3])
@@ -64,31 +66,30 @@ class TestAverage(TestCase):
     def test_weights(self):
         y = arange(10)
         w = arange(10)
-        assert_almost_equal(average(y, weights=w), (arange(10) ** 2).sum()*1. / arange(10).sum())
-
+        actual = average(y, weights=w)
+        desired = (arange(10) ** 2).sum()*1. / arange(10).sum()
+        assert_almost_equal(actual, desired)
+        
         y1 = array([[1, 2, 3], [4, 5, 6]])
         w0 = [1, 2]
         actual = average(y1, weights=w0, axis=0)
         desired = array([3., 4., 5.])
         assert_almost_equal(actual, desired)
 
-
         w1 = [0, 0, 1]
+        actual = average(y1, weights=w1, axis=1)
         desired = array([3., 6.])
-        assert_almost_equal(average(y1, weights=w1, axis=1), desired)
+        assert_almost_equal(actual, desired)
 
         # This should raise an error. Can we test for that ?
         # assert_equal(average(y1, weights=w1), 9./2.)
 
-
         # 2D Case
         w2 = [[0, 0, 1], [0, 0, 2]]
         desired = array([3., 6.])
         assert_array_equal(average(y1, weights=w2, axis=1), desired)
-
         assert_equal(average(y1, weights=w2), 5.)
 
-
     def test_returned(self):
         y = array([[1, 2, 3], [4, 5, 6]])
 
@@ -136,6 +137,7 @@ class TestSelect(TestCase):
         assert_equal(len(choices), 3)
         assert_equal(len(conditions), 3)
 
+
 class TestInsert(TestCase):
     def test_basic(self):
         a = [1, 2, 3]
@@ -143,6 +145,7 @@ class TestInsert(TestCase):
         assert_equal(insert(a, 3, 1), [1, 2, 3, 1])
         assert_equal(insert(a, [1, 1, 1], [1, 2, 3]), [1, 1, 2, 3, 2, 3])
 
+
 class TestAmax(TestCase):
     def test_basic(self):
         a = [3, 4, 5, 10, -3, -5, 6.0]
@@ -153,6 +156,7 @@ class TestAmax(TestCase):
         assert_equal(amax(b, axis=0), [8.0, 10.0, 9.0])
         assert_equal(amax(b, axis=1), [9.0, 10.0, 8.0])
 
+
 class TestAmin(TestCase):
     def test_basic(self):
         a = [3, 4, 5, 10, -3, -5, 6.0]
@@ -163,6 +167,7 @@ class TestAmin(TestCase):
         assert_equal(amin(b, axis=0), [3.0, 3.0, 2.0])
         assert_equal(amin(b, axis=1), [3.0, 4.0, 2.0])
 
+
 class TestPtp(TestCase):
     def test_basic(self):
         a = [3, 4, 5, 10, -3, -5, 6.0]
@@ -173,6 +178,7 @@ class TestPtp(TestCase):
         assert_equal(ptp(b, axis=0), [5.0, 7.0, 7.0])
         assert_equal(ptp(b, axis= -1), [6.0, 6.0, 6.0])
 
+
 class TestCumsum(TestCase):
     def test_basic(self):
         ba = [1, 2, 10, 11, 6, 5, 4]
@@ -189,6 +195,7 @@ class TestCumsum(TestCase):
                                       [5, 11, 18, 27],
                                       [10, 13, 17, 22]], ctype))
 
+
 class TestProd(TestCase):
     def test_basic(self):
         ba = [1, 2, 10, 11, 6, 5, 4]
@@ -207,6 +214,7 @@ class TestProd(TestCase):
                                    array([50, 36, 84, 180], ctype))
                 assert_array_equal(prod(a2, axis= -1), array([24, 1890, 600], ctype))
 
+
 class TestCumprod(TestCase):
     def test_basic(self):
         ba = [1, 2, 10, 11, 6, 5, 4]
@@ -232,6 +240,7 @@ class TestCumprod(TestCase):
                                           [ 5, 30, 210, 1890],
                                           [10, 30, 120, 600]], ctype))
 
+
 class TestDiff(TestCase):
     def test_basic(self):
         x = [1, 4, 6, 7, 12]
@@ -253,6 +262,7 @@ class TestDiff(TestCase):
         assert_array_equal(diff(x, axis=0), out3)
         assert_array_equal(diff(x, n=2, axis=0), out4)
 
+
 class TestGradient(TestCase):
     def test_basic(self):
         x = array([[1, 1], [3, 4]])
@@ -271,6 +281,7 @@ class TestGradient(TestCase):
         x = np.ma.array([[1, 1], [3, 4]])
         assert_equal(type(gradient(x)[0]), type(x))
 
+
 class TestAngle(TestCase):
     def test_basic(self):
         x = [1 + 3j, sqrt(2) / 2.0 + 1j * sqrt(2) / 2, 1, 1j, -1, -1j, 1 - 3j, -1 + 3j]
@@ -282,6 +293,7 @@ class TestAngle(TestCase):
         assert_array_almost_equal(y, yo, 11)
         assert_array_almost_equal(z, zo, 11)
 
+
 class TestTrimZeros(TestCase):
     """ only testing for integer splits.
     """
@@ -289,10 +301,12 @@ class TestTrimZeros(TestCase):
         a = array([0, 0, 1, 2, 3, 4, 0])
         res = trim_zeros(a)
         assert_array_equal(res, array([1, 2, 3, 4]))
+
     def test_leading_skip(self):
         a = array([0, 0, 1, 0, 2, 3, 4, 0])
         res = trim_zeros(a)
         assert_array_equal(res, array([1, 0, 2, 3, 4]))
+
     def test_trailing_skip(self):
         a = array([0, 0, 1, 0, 2, 3, 0, 4, 0])
         res = trim_zeros(a)
@@ -304,10 +318,12 @@ class TestExtins(TestCase):
         a = array([1, 3, 2, 1, 2, 3, 3])
         b = extract(a > 1, a)
         assert_array_equal(b, [3, 2, 2, 3, 3])
+
     def test_place(self):
         a = array([1, 4, 3, 2, 5, 8, 7])
         place(a, [0, 1, 0, 1, 0, 1, 0], [2, 4, 6])
         assert_array_equal(a, [1, 2, 3, 4, 5, 6, 7])
+
     def test_both(self):
         a = rand(10)
         mask = a > 0.5
@@ -317,6 +333,7 @@ class TestExtins(TestCase):
         place(a, mask, c)
         assert_array_equal(a, ac)
 
+
 class TestVectorize(TestCase):
     def test_simple(self):
         def addsubtract(a, b):
@@ -327,6 +344,7 @@ class TestVectorize(TestCase):
         f = vectorize(addsubtract)
         r = f([0, 3, 6, 9], [1, 3, 5, 7])
         assert_array_equal(r, [1, 6, 1, 2])
+
     def test_scalar(self):
         def addsubtract(a, b):
             if a > b:
@@ -336,12 +354,45 @@ class TestVectorize(TestCase):
         f = vectorize(addsubtract)
         r = f([0, 3, 6, 9], 5)
         assert_array_equal(r, [5, 8, 1, 4])
+
     def test_large(self):
         x = linspace(-3, 2, 10000)
         f = vectorize(lambda x: x)
         y = f(x)
         assert_array_equal(y, x)
 
+    def test_ufunc(self):
+        import math
+        f = vectorize(math.cos)
+        args = array([0, 0.5*pi, pi, 1.5*pi, 2*pi])
+        r1 = f(args)
+        r2 = cos(args)
+        assert_array_equal(r1, r2)
+
+    def test_keywords(self):
+        import math
+        def foo(a, b=1):
+            return a + b
+        f = vectorize(foo)
+        args = array([1,2,3])
+        r1 = f(args)
+        r2 = array([2,3,4])
+        assert_array_equal(r1, r2)
+        r1 = f(args, 2)
+        r2 = array([3,4,5])
+        assert_array_equal(r1, r2)
+
+    def test_keywords_no_func_code(self):
+        # This needs to test a function that has keywords but
+        # no func_code attribute, since otherwise vectorize will
+        # inspect the func_code.
+        import random
+        try:
+            f = vectorize(random.randrange)
+        except:
+            raise AssertionError()
+
+
 class TestDigitize(TestCase):
     def test_forward(self):
         x = arange(-6, 5)
@@ -358,6 +409,7 @@ class TestDigitize(TestCase):
         bin = linspace(x.min(), x.max(), 10)
         assert all(digitize(x, bin) != 0)
 
+
 class TestUnwrap(TestCase):
     def test_simple(self):
                 #check that unwrap removes jumps greather that 2*pi
@@ -447,6 +499,7 @@ class TestSinc(TestCase):
         #check symmetry
         assert_array_almost_equal(w, flipud(w), 7)
 
+
 class TestHistogram(TestCase):
     def setUp(self):
         pass
@@ -486,7 +539,6 @@ class TestHistogram(TestCase):
         area = sum(a * diff(b))
         assert_almost_equal(area, 1)
 
-
     def test_outliers(self):
         # Check that outliers are not tallied
         a = arange(10) + .5
@@ -511,7 +563,6 @@ class TestHistogram(TestCase):
         h, b = histogram(a, bins=8, range=[1, 9], weights=w)
         assert_equal(h, w[1:-1])
 
-
     def test_type(self):
         # Check the type of the returned histogram
         a = arange(10) + .5
@@ -527,7 +578,6 @@ class TestHistogram(TestCase):
         h, b = histogram(a, weights=ones(10, float))
         assert(issubdtype(h.dtype, float))
 
-
     def test_weights(self):
         v = rand(100)
         w = ones(100) * 5
@@ -550,6 +600,7 @@ class TestHistogram(TestCase):
         wa, wb = histogram([1, 2, 2, 4], bins=4, weights=[4, 3, 2, 1], normed=True)
         assert_array_equal(wa, array([4, 5, 0, 1]) / 10. / 3. * 4)
 
+
 class TestHistogramdd(TestCase):
     def test_simple(self):
         x = array([[-.5, .5, 1.5], [-.5, 1.5, 2.5], [-.5, 2.5, .5], \
@@ -622,6 +673,7 @@ class TestHistogramdd(TestCase):
         hist, edges = histogramdd(x, bins=2)
         assert_array_equal(edges[0], array([-0.5, 0. , 0.5]))
 
+
 class TestUnique(TestCase):
     def test_simple(self):
         x = array([4, 3, 2, 1, 1, 2, 3, 4, 0])
@@ -642,6 +694,7 @@ class TestCheckFinite(TestCase):
         assert_raises(ValueError, numpy.lib.asarray_chkfinite, b)
         assert_raises(ValueError, numpy.lib.asarray_chkfinite, c)
 
+
 class TestNaNFuncts(TestCase):
     def setUp(self):
         self.A = array([[[ nan, 0.01319214, 0.01620964],
@@ -728,7 +781,6 @@ class TestNaNFuncts(TestCase):
         assert_equal(np.isinf(a), np.zeros((2, 4), dtype=bool))
 
 
-
 class TestCorrCoef(TestCase):
     def test_simple(self):
         A = array([[ 0.15391142, 0.18045767, 0.14197213],
@@ -756,7 +808,6 @@ class TestCorrCoef(TestCase):
                                      - 0.66173113, 0.98317823, 1.        ]]))
 
 
-
 class Test_i0(TestCase):
     def test_simple(self):
         assert_almost_equal(i0(0.5), array(1.0634833707413234))
@@ -775,6 +826,7 @@ class Test_i0(TestCase):
                                    [ 1.03352052, 1.13557954],
                                    [ 1.0588429 , 1.06432317]]))
 
+
 class TestKaiser(TestCase):
     def test_simple(self):
         assert_almost_equal(kaiser(0, 1.0), array([]))
@@ -790,6 +842,7 @@ class TestKaiser(TestCase):
     def test_int_beta(self):
         kaiser(3, 4)
 
+
 class TestMsort(TestCase):
     def test_simple(self):
         A = array([[ 0.44567325, 0.79115165, 0.5490053 ],
@@ -800,6 +853,7 @@ class TestMsort(TestCase):
                                    [ 0.44567325, 0.52929049, 0.5490053 ],
                                    [ 0.64864341, 0.79115165, 0.96098397]]))
 
+
 class TestMeshgrid(TestCase):
     def test_simple(self):
         [X, Y] = meshgrid([1, 2, 3], [4, 5, 6, 7])
@@ -857,6 +911,7 @@ class TestPiecewise(TestCase):
         assert y.ndim == 0
         assert y == 0
 
+
 class TestBincount(TestCase):
     def test_simple(self):
         y = np.bincount(np.arange(4))
@@ -878,9 +933,11 @@ class TestBincount(TestCase):
         y = np.bincount(x, w)
         assert_array_equal(y, np.array([0, 0.2, 0.5, 0, 0.5, 0.1]))
 
+
 def compare_results(res, desired):
     for i in range(len(desired)):
         assert_array_equal(res[i], desired[i])
 
+
 if __name__ == "__main__":
     run_module_suite()