Merge pull request #3999 from WarrenWeckesser/vander

ENH: lib: Rewrite vander

1 files changed, 45 insertions, 5 deletions

diff --git a/numpy/lib/tests/test_twodim_base.py b/numpy/lib/tests/test_twodim_base.py
index 4c19eba58..8d0275a25 100644
--- a/numpy/lib/tests/test_twodim_base.py
+++ b/numpy/lib/tests/test_twodim_base.py
@@ -3,11 +3,16 @@
 """
 from __future__ import division, absolute_import, print_function
 
-from numpy.testing import *
+from numpy.testing import (
+    TestCase, run_module_suite, assert_equal, assert_array_equal,
+    assert_array_max_ulp, assert_array_almost_equal, assert_raises, rand,
+    )
 
-from numpy import (arange, rot90, add, fliplr, flipud, zeros, ones, eye,
-                   array, diag, histogram2d, tri, mask_indices, triu_indices,
-                   triu_indices_from, tril_indices, tril_indices_from)
+from numpy import (
+    arange, rot90, add, fliplr, flipud, zeros, ones, eye, array, diag,
+    histogram2d, tri, mask_indices, triu_indices, triu_indices_from,
+    tril_indices, tril_indices_from, vander,
+    )
 
 import numpy as np
 from numpy.compat import asbytes, asbytes_nested
@@ -345,7 +350,8 @@ class TestTriuIndices(object):
                       [9,  10, -1, -1],
                       [13, 14, 15, -1]]))
 
-        # These cover almost the whole array (two diagonals right of the main one):
+        # These cover almost the whole array (two diagonals right of the
+        # main one):
         a[iu2] = -10
         yield (assert_array_equal, a,
                array([[-1,  -1, -10, -10],
@@ -368,5 +374,39 @@ class TestTriuIndicesFrom(object):
         assert_raises(ValueError, triu_indices_from, np.ones((2, 3)))
 
 
+class TestVander(object):
+    def test_basic(self):
+        c = np.array([0, 1, -2, 3])
+        v = vander(c)
+        powers = np.array([[ 0,  0, 0,  0, 1],
+                           [ 1,  1, 1,  1, 1],
+                           [16, -8, 4, -2, 1],
+                           [81, 27, 9,  3, 1]])
+        # Check default value of N:
+        yield (assert_array_equal, v, powers[:, 1:])
+        # Check a range of N values, including 0 and 5 (greater than default)
+        m = powers.shape[1]
+        for n in range(6):
+            v = vander(c, N=n)
+            yield (assert_array_equal, v, powers[:, m-n:m])
+
+    def test_dtypes(self):
+        c = array([11, -12, 13], dtype=np.int8)
+        v = vander(c)
+        expected = np.array([[121,  11, 1],
+                             [144, -12, 1],
+                             [169,  13, 1]])
+        yield (assert_array_equal, v, expected)
+
+        c = array([1.0+1j, 1.0-1j])
+        v = vander(c, N=3)
+        expected = np.array([[ 2j, 1+1j, 1],
+                             [-2j, 1-1j, 1]])
+        # The data is floating point, but the values are small integers,
+        # so assert_array_equal *should* be safe here (rather than, say,
+        # assert_array_almost_equal).
+        yield (assert_array_equal, v, expected)
+
+
 if __name__ == "__main__":
     run_module_suite()