| 1 | import unittest2 as unittest |
|---|
| 2 | import numpy as np |
|---|
| 3 | |
|---|
| 4 | from vsm import * |
|---|
| 5 | from vsm.structarr import * |
|---|
| 6 | |
|---|
| 7 | class TestCore(unittest.TestCase): |
|---|
| 8 | |
|---|
| 9 | def test_arr_add_field(self): |
|---|
| 10 | |
|---|
| 11 | arr = np.array([(1, '1'), (2, '2'), (3, '3')], |
|---|
| 12 | dtype=[('i', np.int), ('c', '|S1')]) |
|---|
| 13 | new_arr = np.array([(1, '1', 0), (2, '2', 0), (3, '3', 0)], |
|---|
| 14 | dtype=[('i', np.int), ('c', '|S1'), ('new', np.int)]) |
|---|
| 15 | |
|---|
| 16 | new_field = 'new' |
|---|
| 17 | vals = np.zeros(3, dtype=np.int) |
|---|
| 18 | |
|---|
| 19 | test_arr = arr_add_field(arr, new_field, vals) |
|---|
| 20 | |
|---|
| 21 | self.assertTrue((new_arr==test_arr).all()) |
|---|
| 22 | self.assertTrue(new_arr.dtype==test_arr.dtype) |
|---|
| 23 | |
|---|
| 24 | def test_enum_matrix(self): |
|---|
| 25 | |
|---|
| 26 | arr = np.array([[6,3,7], [2,0,4]]) |
|---|
| 27 | em1 = enum_matrix(arr) |
|---|
| 28 | em2 = enum_matrix(arr, indices=[10,20,30], field_name='tens') |
|---|
| 29 | |
|---|
| 30 | self.assertTrue((em1 == np.array([[(2,7), (0,6), (1, 3)],[(2,4), (0,2), (1,0)]], |
|---|
| 31 | dtype=[('i', '<i8'), ('value', '<i8')])).all()) |
|---|
| 32 | self.assertTrue((em2 == np.array([[(30,7), (10,6), (20, 3)],[(30,4), (10,2), (20,0)]], |
|---|
| 33 | dtype=[('tens', '<i8'), ('value', '<i8')])).all()) |
|---|
| 34 | |
|---|
| 35 | |
|---|
| 36 | |
|---|
| 37 | def test_enum_sort(self): |
|---|
| 38 | |
|---|
| 39 | arr = np.array([7,3,1,8,2]) |
|---|
| 40 | sorted_arr = enum_sort(arr) |
|---|
| 41 | sorted_arr1 = enum_sort(arr, indices=[10,20,30,40,50]) |
|---|
| 42 | |
|---|
| 43 | self.assertTrue((sorted_arr == |
|---|
| 44 | np.array([(3, 8), (0, 7), (1, 3), (4, 2), (2, 1)], |
|---|
| 45 | dtype=[('i', '<i8'), ('value', '<i8')])).all()) |
|---|
| 46 | |
|---|
| 47 | self.assertTrue((sorted_arr1 == |
|---|
| 48 | np.array([(40, 8), (10, 7), (20, 3), (50, 2), (30, 1)], |
|---|
| 49 | dtype=[('i', '<i8'), ('value', '<i8')])).all()) |
|---|
| 50 | |
|---|
| 51 | |
|---|
| 52 | def test_enum_array(self): |
|---|
| 53 | |
|---|
| 54 | arr1 = np.array([7,3,1,8,2]) |
|---|
| 55 | ea1 = enum_array(arr1) |
|---|
| 56 | arr2 = np.array([6,3,7,2,0,4]) |
|---|
| 57 | ea2 = enum_array(arr2) |
|---|
| 58 | |
|---|
| 59 | self.assertTrue((ea1 == |
|---|
| 60 | np.array([(0,7), (1,3), (2,1), (3,8), (4,2)], |
|---|
| 61 | dtype=[('i', '<i8'), ('value', '<i8')])).all()) |
|---|
| 62 | self.assertTrue((ea2 == |
|---|
| 63 | np.array([(0,6), (1,3), (2,7), (3,2), (4,0), (5,4)], |
|---|
| 64 | dtype=[('i', '<i8'), ('value', '<i8')])).all()) |
|---|
| 65 | |
|---|
| 66 | |
|---|
| 67 | def test_zip_arr(self): |
|---|
| 68 | |
|---|
| 69 | arr1 = np.array([[2,4], [6,8]]) |
|---|
| 70 | arr2 = np.array([[1,3], [5,7]]) |
|---|
| 71 | |
|---|
| 72 | zipped = zip_arr(arr1, arr2, field_names=['even', 'odd']) |
|---|
| 73 | self.assertTrue((zipped == np.array([[(2,1), (4,3)], [(6,5), (8,7)]], |
|---|
| 74 | dtype=[('even', '<i8'), ('odd', '<i8')])).all()) |
|---|
| 75 | |
|---|
| 76 | |
|---|
| 77 | def test_map_strarr(self): |
|---|
| 78 | |
|---|
| 79 | arr = np.array([(0, 1.), (1, 2.)], |
|---|
| 80 | dtype=[('i', 'i4'), ('v', 'f4')]) |
|---|
| 81 | m = ['foo', 'bar'] |
|---|
| 82 | arr = map_strarr(arr, m, 'i', new_k='str') |
|---|
| 83 | |
|---|
| 84 | self.assertTrue((arr['str'] == np.array(m, |
|---|
| 85 | dtype=np.array(m).dtype)).all()) |
|---|
| 86 | self.assertTrue((arr['v'] == np.array([1., 2.], dtype='f4')).all()) |
|---|
| 87 | |
|---|
| 88 | |
|---|
| 89 | suite = unittest.TestLoader().loadTestsFromTestCase(TestCore) |
|---|
| 90 | unittest.TextTestRunner(verbosity=2).run(suite) |
|---|
