summaryrefslogtreecommitdiff
path: root/numpy/lib/tests
diff options
context:
space:
mode:
authorAllan Haldane <allan.haldane@gmail.com>2018-01-30 19:15:54 -0500
committerAllan Haldane <allan.haldane@gmail.com>2018-10-31 12:22:21 -0400
commitf1fba70edd1829c64e3290fa6b1a20d01e9d9674 (patch)
tree9d60664d69bb1fec9fae6a01dbdfe8d47919e4e7 /numpy/lib/tests
parente9c631825df7b0a9c969dbfcbaf19d975e4afdde (diff)
downloadnumpy-f1fba70edd1829c64e3290fa6b1a20d01e9d9674.tar.gz
ENH: add multi-field assignment helpers in np.lib.recfunctions
Adds helper functions for the copy->view transition for multi-field indexes. Adds `structured_to_unstructured`, `apply_along_fields`, `assign_fields_by_name`, `require_fields`.
Diffstat (limited to 'numpy/lib/tests')
-rw-r--r--numpy/lib/tests/test_recfunctions.py53
1 files changed, 52 insertions, 1 deletions
diff --git a/numpy/lib/tests/test_recfunctions.py b/numpy/lib/tests/test_recfunctions.py
index 5585a95f9..8b61ba0bf 100644
--- a/numpy/lib/tests/test_recfunctions.py
+++ b/numpy/lib/tests/test_recfunctions.py
@@ -10,7 +10,8 @@ from numpy.testing import assert_, assert_raises
from numpy.lib.recfunctions import (
drop_fields, rename_fields, get_fieldstructure, recursive_fill_fields,
find_duplicates, merge_arrays, append_fields, stack_arrays, join_by,
- repack_fields)
+ repack_fields, unstructured_to_structured, structured_to_unstructured,
+ apply_along_fields, require_fields, assign_fields_by_name)
get_names = np.lib.recfunctions.get_names
get_names_flat = np.lib.recfunctions.get_names_flat
zip_descr = np.lib.recfunctions.zip_descr
@@ -204,6 +205,56 @@ class TestRecFunctions(object):
dt = np.dtype((np.record, dt))
assert_(repack_fields(dt).type is np.record)
+ def test_structured_to_unstructured(self):
+ a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
+ out = structured_to_unstructured(a)
+ assert_equal(out, np.zeros((4,5), dtype='f8'))
+
+ b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+ dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
+ out = np.mean(structured_to_unstructured(b[['x', 'z']]), axis=-1)
+ assert_equal(out, np.array([ 3. , 5.5, 9. , 11. ]))
+
+ c = np.arange(20).reshape((4,5))
+ out = unstructured_to_structured(c, a.dtype)
+ want = np.array([( 0, ( 1., 2), [ 3., 4.]),
+ ( 5, ( 6., 7), [ 8., 9.]),
+ (10, (11., 12), [13., 14.]),
+ (15, (16., 17), [18., 19.])],
+ dtype=[('a', '<i4'),
+ ('b', [('f0', '<f4'), ('f1', '<u2')]),
+ ('c', '<f4', (2,))])
+ assert_equal(out, want)
+
+ d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+ dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
+ assert_equal(apply_along_fields(np.mean, d),
+ np.array([ 8.0/3, 16.0/3, 26.0/3, 11. ]))
+ assert_equal(apply_along_fields(np.mean, d[['x', 'z']]),
+ np.array([ 3. , 5.5, 9. , 11. ]))
+
+ def test_field_assignment_by_name(self):
+ a = np.ones(2, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')])
+ newdt = [('b', 'f4'), ('c', 'u1')]
+
+ assert_equal(require_fields(a, newdt), np.ones(2, newdt))
+
+ b = np.array([(1,2), (3,4)], dtype=newdt)
+ assign_fields_by_name(a, b, zero_unassigned=False)
+ assert_equal(a, np.array([(1,1,2),(1,3,4)], dtype=a.dtype))
+ assign_fields_by_name(a, b)
+ assert_equal(a, np.array([(0,1,2),(0,3,4)], dtype=a.dtype))
+
+ # test nested fields
+ a = np.ones(2, dtype=[('a', [('b', 'f8'), ('c', 'u1')])])
+ newdt = [('a', [('c', 'u1')])]
+ assert_equal(require_fields(a, newdt), np.ones(2, newdt))
+ b = np.array([((2,),), ((3,),)], dtype=newdt)
+ assign_fields_by_name(a, b, zero_unassigned=False)
+ assert_equal(a, np.array([((1,2),), ((1,3),)], dtype=a.dtype))
+ assign_fields_by_name(a, b)
+ assert_equal(a, np.array([((0,2),), ((0,3),)], dtype=a.dtype))
+
class TestRecursiveFillFields(object):
# Test recursive_fill_fields.