diff options
| author | Stefan van der Walt <stefan@sun.ac.za> | 2012-08-27 09:24:59 -0700 |
|---|---|---|
| committer | Stefan van der Walt <stefan@sun.ac.za> | 2012-09-02 07:55:59 -0700 |
| commit | 942bdb26bee477a8ac2c6b1c1c963c52c7c08403 (patch) | |
| tree | 8a9d4ad6f9c47d4397dcb37255bef6a3b20b1a36 /numpy/lib/tests/test_index_tricks.py | |
| parent | e60c70d7ca5dbe45860c44673ddab02d47770155 (diff) | |
| download | numpy-942bdb26bee477a8ac2c6b1c1c963c52c7c08403.tar.gz | |
Improve ndindex execution speed.
Diffstat (limited to 'numpy/lib/tests/test_index_tricks.py')
| -rw-r--r-- | numpy/lib/tests/test_index_tricks.py | 8 |
1 files changed, 7 insertions, 1 deletions
diff --git a/numpy/lib/tests/test_index_tricks.py b/numpy/lib/tests/test_index_tricks.py index beda2d146..0ede40d5a 100644 --- a/numpy/lib/tests/test_index_tricks.py +++ b/numpy/lib/tests/test_index_tricks.py @@ -2,7 +2,7 @@ from numpy.testing import * import numpy as np from numpy import ( array, ones, r_, mgrid, unravel_index, zeros, where, ndenumerate, fill_diagonal, diag_indices, - diag_indices_from, s_, index_exp ) + diag_indices_from, s_, index_exp, ndindex ) class TestRavelUnravelIndex(TestCase): def test_basic(self): @@ -237,5 +237,11 @@ def test_diag_indices_from(): assert_array_equal(c, np.arange(4)) +def test_ndindex(): + x = list(np.ndindex(1, 2, 3)) + expected = [ix for ix, e in np.ndenumerate(np.zeros((1, 2, 3)))] + assert_array_equal(x, expected) + + if __name__ == "__main__": run_module_suite() |