test_recfunctions.py 41 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987
  1. import pytest
  2. import numpy as np
  3. import numpy.ma as ma
  4. from numpy.ma.mrecords import MaskedRecords
  5. from numpy.ma.testutils import assert_equal
  6. from numpy.testing import assert_, assert_raises
  7. from numpy.lib.recfunctions import (
  8. drop_fields, rename_fields, get_fieldstructure, recursive_fill_fields,
  9. find_duplicates, merge_arrays, append_fields, stack_arrays, join_by,
  10. repack_fields, unstructured_to_structured, structured_to_unstructured,
  11. apply_along_fields, require_fields, assign_fields_by_name)
  12. get_fieldspec = np.lib.recfunctions._get_fieldspec
  13. get_names = np.lib.recfunctions.get_names
  14. get_names_flat = np.lib.recfunctions.get_names_flat
  15. zip_descr = np.lib.recfunctions._zip_descr
  16. zip_dtype = np.lib.recfunctions._zip_dtype
  17. class TestRecFunctions:
  18. # Misc tests
  19. def setup_method(self):
  20. x = np.array([1, 2, ])
  21. y = np.array([10, 20, 30])
  22. z = np.array([('A', 1.), ('B', 2.)],
  23. dtype=[('A', '|S3'), ('B', float)])
  24. w = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
  25. dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
  26. self.data = (w, x, y, z)
  27. def test_zip_descr(self):
  28. # Test zip_descr
  29. (w, x, y, z) = self.data
  30. # Std array
  31. test = zip_descr((x, x), flatten=True)
  32. assert_equal(test,
  33. np.dtype([('', int), ('', int)]))
  34. test = zip_descr((x, x), flatten=False)
  35. assert_equal(test,
  36. np.dtype([('', int), ('', int)]))
  37. # Std & flexible-dtype
  38. test = zip_descr((x, z), flatten=True)
  39. assert_equal(test,
  40. np.dtype([('', int), ('A', '|S3'), ('B', float)]))
  41. test = zip_descr((x, z), flatten=False)
  42. assert_equal(test,
  43. np.dtype([('', int),
  44. ('', [('A', '|S3'), ('B', float)])]))
  45. # Standard & nested dtype
  46. test = zip_descr((x, w), flatten=True)
  47. assert_equal(test,
  48. np.dtype([('', int),
  49. ('a', int),
  50. ('ba', float), ('bb', int)]))
  51. test = zip_descr((x, w), flatten=False)
  52. assert_equal(test,
  53. np.dtype([('', int),
  54. ('', [('a', int),
  55. ('b', [('ba', float), ('bb', int)])])]))
  56. def test_drop_fields(self):
  57. # Test drop_fields
  58. a = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
  59. dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
  60. # A basic field
  61. test = drop_fields(a, 'a')
  62. control = np.array([((2, 3.0),), ((5, 6.0),)],
  63. dtype=[('b', [('ba', float), ('bb', int)])])
  64. assert_equal(test, control)
  65. # Another basic field (but nesting two fields)
  66. test = drop_fields(a, 'b')
  67. control = np.array([(1,), (4,)], dtype=[('a', int)])
  68. assert_equal(test, control)
  69. # A nested sub-field
  70. test = drop_fields(a, ['ba', ])
  71. control = np.array([(1, (3.0,)), (4, (6.0,))],
  72. dtype=[('a', int), ('b', [('bb', int)])])
  73. assert_equal(test, control)
  74. # All the nested sub-field from a field: zap that field
  75. test = drop_fields(a, ['ba', 'bb'])
  76. control = np.array([(1,), (4,)], dtype=[('a', int)])
  77. assert_equal(test, control)
  78. # dropping all fields results in an array with no fields
  79. test = drop_fields(a, ['a', 'b'])
  80. control = np.array([(), ()], dtype=[])
  81. assert_equal(test, control)
  82. def test_rename_fields(self):
  83. # Test rename fields
  84. a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))],
  85. dtype=[('a', int),
  86. ('b', [('ba', float), ('bb', (float, 2))])])
  87. test = rename_fields(a, {'a': 'A', 'bb': 'BB'})
  88. newdtype = [('A', int), ('b', [('ba', float), ('BB', (float, 2))])]
  89. control = a.view(newdtype)
  90. assert_equal(test.dtype, newdtype)
  91. assert_equal(test, control)
  92. def test_get_names(self):
  93. # Test get_names
  94. ndtype = np.dtype([('A', '|S3'), ('B', float)])
  95. test = get_names(ndtype)
  96. assert_equal(test, ('A', 'B'))
  97. ndtype = np.dtype([('a', int), ('b', [('ba', float), ('bb', int)])])
  98. test = get_names(ndtype)
  99. assert_equal(test, ('a', ('b', ('ba', 'bb'))))
  100. ndtype = np.dtype([('a', int), ('b', [])])
  101. test = get_names(ndtype)
  102. assert_equal(test, ('a', ('b', ())))
  103. ndtype = np.dtype([])
  104. test = get_names(ndtype)
  105. assert_equal(test, ())
  106. def test_get_names_flat(self):
  107. # Test get_names_flat
  108. ndtype = np.dtype([('A', '|S3'), ('B', float)])
  109. test = get_names_flat(ndtype)
  110. assert_equal(test, ('A', 'B'))
  111. ndtype = np.dtype([('a', int), ('b', [('ba', float), ('bb', int)])])
  112. test = get_names_flat(ndtype)
  113. assert_equal(test, ('a', 'b', 'ba', 'bb'))
  114. ndtype = np.dtype([('a', int), ('b', [])])
  115. test = get_names_flat(ndtype)
  116. assert_equal(test, ('a', 'b'))
  117. ndtype = np.dtype([])
  118. test = get_names_flat(ndtype)
  119. assert_equal(test, ())
  120. def test_get_fieldstructure(self):
  121. # Test get_fieldstructure
  122. # No nested fields
  123. ndtype = np.dtype([('A', '|S3'), ('B', float)])
  124. test = get_fieldstructure(ndtype)
  125. assert_equal(test, {'A': [], 'B': []})
  126. # One 1-nested field
  127. ndtype = np.dtype([('A', int), ('B', [('BA', float), ('BB', '|S1')])])
  128. test = get_fieldstructure(ndtype)
  129. assert_equal(test, {'A': [], 'B': [], 'BA': ['B', ], 'BB': ['B']})
  130. # One 2-nested fields
  131. ndtype = np.dtype([('A', int),
  132. ('B', [('BA', int),
  133. ('BB', [('BBA', int), ('BBB', int)])])])
  134. test = get_fieldstructure(ndtype)
  135. control = {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'],
  136. 'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']}
  137. assert_equal(test, control)
  138. # 0 fields
  139. ndtype = np.dtype([])
  140. test = get_fieldstructure(ndtype)
  141. assert_equal(test, {})
  142. def test_find_duplicates(self):
  143. # Test find_duplicates
  144. a = ma.array([(2, (2., 'B')), (1, (2., 'B')), (2, (2., 'B')),
  145. (1, (1., 'B')), (2, (2., 'B')), (2, (2., 'C'))],
  146. mask=[(0, (0, 0)), (0, (0, 0)), (0, (0, 0)),
  147. (0, (0, 0)), (1, (0, 0)), (0, (1, 0))],
  148. dtype=[('A', int), ('B', [('BA', float), ('BB', '|S1')])])
  149. test = find_duplicates(a, ignoremask=False, return_index=True)
  150. control = [0, 2]
  151. assert_equal(sorted(test[-1]), control)
  152. assert_equal(test[0], a[test[-1]])
  153. test = find_duplicates(a, key='A', return_index=True)
  154. control = [0, 1, 2, 3, 5]
  155. assert_equal(sorted(test[-1]), control)
  156. assert_equal(test[0], a[test[-1]])
  157. test = find_duplicates(a, key='B', return_index=True)
  158. control = [0, 1, 2, 4]
  159. assert_equal(sorted(test[-1]), control)
  160. assert_equal(test[0], a[test[-1]])
  161. test = find_duplicates(a, key='BA', return_index=True)
  162. control = [0, 1, 2, 4]
  163. assert_equal(sorted(test[-1]), control)
  164. assert_equal(test[0], a[test[-1]])
  165. test = find_duplicates(a, key='BB', return_index=True)
  166. control = [0, 1, 2, 3, 4]
  167. assert_equal(sorted(test[-1]), control)
  168. assert_equal(test[0], a[test[-1]])
  169. def test_find_duplicates_ignoremask(self):
  170. # Test the ignoremask option of find_duplicates
  171. ndtype = [('a', int)]
  172. a = ma.array([1, 1, 1, 2, 2, 3, 3],
  173. mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype)
  174. test = find_duplicates(a, ignoremask=True, return_index=True)
  175. control = [0, 1, 3, 4]
  176. assert_equal(sorted(test[-1]), control)
  177. assert_equal(test[0], a[test[-1]])
  178. test = find_duplicates(a, ignoremask=False, return_index=True)
  179. control = [0, 1, 2, 3, 4, 6]
  180. assert_equal(sorted(test[-1]), control)
  181. assert_equal(test[0], a[test[-1]])
  182. def test_repack_fields(self):
  183. dt = np.dtype('u1,f4,i8', align=True)
  184. a = np.zeros(2, dtype=dt)
  185. assert_equal(repack_fields(dt), np.dtype('u1,f4,i8'))
  186. assert_equal(repack_fields(a).itemsize, 13)
  187. assert_equal(repack_fields(repack_fields(dt), align=True), dt)
  188. # make sure type is preserved
  189. dt = np.dtype((np.record, dt))
  190. assert_(repack_fields(dt).type is np.record)
  191. def test_structured_to_unstructured(self):
  192. a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
  193. out = structured_to_unstructured(a)
  194. assert_equal(out, np.zeros((4,5), dtype='f8'))
  195. b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
  196. dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
  197. out = np.mean(structured_to_unstructured(b[['x', 'z']]), axis=-1)
  198. assert_equal(out, np.array([ 3. , 5.5, 9. , 11. ]))
  199. out = np.mean(structured_to_unstructured(b[['x']]), axis=-1)
  200. assert_equal(out, np.array([ 1. , 4. , 7. , 10. ]))
  201. c = np.arange(20).reshape((4,5))
  202. out = unstructured_to_structured(c, a.dtype)
  203. want = np.array([( 0, ( 1., 2), [ 3., 4.]),
  204. ( 5, ( 6., 7), [ 8., 9.]),
  205. (10, (11., 12), [13., 14.]),
  206. (15, (16., 17), [18., 19.])],
  207. dtype=[('a', 'i4'),
  208. ('b', [('f0', 'f4'), ('f1', 'u2')]),
  209. ('c', 'f4', (2,))])
  210. assert_equal(out, want)
  211. d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
  212. dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
  213. assert_equal(apply_along_fields(np.mean, d),
  214. np.array([ 8.0/3, 16.0/3, 26.0/3, 11. ]))
  215. assert_equal(apply_along_fields(np.mean, d[['x', 'z']]),
  216. np.array([ 3. , 5.5, 9. , 11. ]))
  217. # check that for uniform field dtypes we get a view, not a copy:
  218. d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
  219. dtype=[('x', 'i4'), ('y', 'i4'), ('z', 'i4')])
  220. dd = structured_to_unstructured(d)
  221. ddd = unstructured_to_structured(dd, d.dtype)
  222. assert_(dd.base is d)
  223. assert_(ddd.base is d)
  224. # including uniform fields with subarrays unpacked
  225. d = np.array([(1, [2, 3], [[ 4, 5], [ 6, 7]]),
  226. (8, [9, 10], [[11, 12], [13, 14]])],
  227. dtype=[('x0', 'i4'), ('x1', ('i4', 2)),
  228. ('x2', ('i4', (2, 2)))])
  229. dd = structured_to_unstructured(d)
  230. ddd = unstructured_to_structured(dd, d.dtype)
  231. assert_(dd.base is d)
  232. assert_(ddd.base is d)
  233. # test that nested fields with identical names don't break anything
  234. point = np.dtype([('x', int), ('y', int)])
  235. triangle = np.dtype([('a', point), ('b', point), ('c', point)])
  236. arr = np.zeros(10, triangle)
  237. res = structured_to_unstructured(arr, dtype=int)
  238. assert_equal(res, np.zeros((10, 6), dtype=int))
  239. # test nested combinations of subarrays and structured arrays, gh-13333
  240. def subarray(dt, shape):
  241. return np.dtype((dt, shape))
  242. def structured(*dts):
  243. return np.dtype([('x{}'.format(i), dt) for i, dt in enumerate(dts)])
  244. def inspect(dt, dtype=None):
  245. arr = np.zeros((), dt)
  246. ret = structured_to_unstructured(arr, dtype=dtype)
  247. backarr = unstructured_to_structured(ret, dt)
  248. return ret.shape, ret.dtype, backarr.dtype
  249. dt = structured(subarray(structured(np.int32, np.int32), 3))
  250. assert_equal(inspect(dt), ((6,), np.int32, dt))
  251. dt = structured(subarray(subarray(np.int32, 2), 2))
  252. assert_equal(inspect(dt), ((4,), np.int32, dt))
  253. dt = structured(np.int32)
  254. assert_equal(inspect(dt), ((1,), np.int32, dt))
  255. dt = structured(np.int32, subarray(subarray(np.int32, 2), 2))
  256. assert_equal(inspect(dt), ((5,), np.int32, dt))
  257. dt = structured()
  258. assert_raises(ValueError, structured_to_unstructured, np.zeros(3, dt))
  259. # these currently don't work, but we may make it work in the future
  260. assert_raises(NotImplementedError, structured_to_unstructured,
  261. np.zeros(3, dt), dtype=np.int32)
  262. assert_raises(NotImplementedError, unstructured_to_structured,
  263. np.zeros((3,0), dtype=np.int32))
  264. def test_unstructured_to_structured(self):
  265. # test if dtype is the args of np.dtype
  266. a = np.zeros((20, 2))
  267. test_dtype_args = [('x', float), ('y', float)]
  268. test_dtype = np.dtype(test_dtype_args)
  269. field1 = unstructured_to_structured(a, dtype=test_dtype_args) # now
  270. field2 = unstructured_to_structured(a, dtype=test_dtype) # before
  271. assert_equal(field1, field2)
  272. def test_field_assignment_by_name(self):
  273. a = np.ones(2, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')])
  274. newdt = [('b', 'f4'), ('c', 'u1')]
  275. assert_equal(require_fields(a, newdt), np.ones(2, newdt))
  276. b = np.array([(1,2), (3,4)], dtype=newdt)
  277. assign_fields_by_name(a, b, zero_unassigned=False)
  278. assert_equal(a, np.array([(1,1,2),(1,3,4)], dtype=a.dtype))
  279. assign_fields_by_name(a, b)
  280. assert_equal(a, np.array([(0,1,2),(0,3,4)], dtype=a.dtype))
  281. # test nested fields
  282. a = np.ones(2, dtype=[('a', [('b', 'f8'), ('c', 'u1')])])
  283. newdt = [('a', [('c', 'u1')])]
  284. assert_equal(require_fields(a, newdt), np.ones(2, newdt))
  285. b = np.array([((2,),), ((3,),)], dtype=newdt)
  286. assign_fields_by_name(a, b, zero_unassigned=False)
  287. assert_equal(a, np.array([((1,2),), ((1,3),)], dtype=a.dtype))
  288. assign_fields_by_name(a, b)
  289. assert_equal(a, np.array([((0,2),), ((0,3),)], dtype=a.dtype))
  290. # test unstructured code path for 0d arrays
  291. a, b = np.array(3), np.array(0)
  292. assign_fields_by_name(b, a)
  293. assert_equal(b[()], 3)
  294. class TestRecursiveFillFields:
  295. # Test recursive_fill_fields.
  296. def test_simple_flexible(self):
  297. # Test recursive_fill_fields on flexible-array
  298. a = np.array([(1, 10.), (2, 20.)], dtype=[('A', int), ('B', float)])
  299. b = np.zeros((3,), dtype=a.dtype)
  300. test = recursive_fill_fields(a, b)
  301. control = np.array([(1, 10.), (2, 20.), (0, 0.)],
  302. dtype=[('A', int), ('B', float)])
  303. assert_equal(test, control)
  304. def test_masked_flexible(self):
  305. # Test recursive_fill_fields on masked flexible-array
  306. a = ma.array([(1, 10.), (2, 20.)], mask=[(0, 1), (1, 0)],
  307. dtype=[('A', int), ('B', float)])
  308. b = ma.zeros((3,), dtype=a.dtype)
  309. test = recursive_fill_fields(a, b)
  310. control = ma.array([(1, 10.), (2, 20.), (0, 0.)],
  311. mask=[(0, 1), (1, 0), (0, 0)],
  312. dtype=[('A', int), ('B', float)])
  313. assert_equal(test, control)
  314. class TestMergeArrays:
  315. # Test merge_arrays
  316. def setup_method(self):
  317. x = np.array([1, 2, ])
  318. y = np.array([10, 20, 30])
  319. z = np.array(
  320. [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)])
  321. w = np.array(
  322. [(1, (2, 3.0, ())), (4, (5, 6.0, ()))],
  323. dtype=[('a', int), ('b', [('ba', float), ('bb', int), ('bc', [])])])
  324. self.data = (w, x, y, z)
  325. def test_solo(self):
  326. # Test merge_arrays on a single array.
  327. (_, x, _, z) = self.data
  328. test = merge_arrays(x)
  329. control = np.array([(1,), (2,)], dtype=[('f0', int)])
  330. assert_equal(test, control)
  331. test = merge_arrays((x,))
  332. assert_equal(test, control)
  333. test = merge_arrays(z, flatten=False)
  334. assert_equal(test, z)
  335. test = merge_arrays(z, flatten=True)
  336. assert_equal(test, z)
  337. def test_solo_w_flatten(self):
  338. # Test merge_arrays on a single array w & w/o flattening
  339. w = self.data[0]
  340. test = merge_arrays(w, flatten=False)
  341. assert_equal(test, w)
  342. test = merge_arrays(w, flatten=True)
  343. control = np.array([(1, 2, 3.0), (4, 5, 6.0)],
  344. dtype=[('a', int), ('ba', float), ('bb', int)])
  345. assert_equal(test, control)
  346. def test_standard(self):
  347. # Test standard & standard
  348. # Test merge arrays
  349. (_, x, y, _) = self.data
  350. test = merge_arrays((x, y), usemask=False)
  351. control = np.array([(1, 10), (2, 20), (-1, 30)],
  352. dtype=[('f0', int), ('f1', int)])
  353. assert_equal(test, control)
  354. test = merge_arrays((x, y), usemask=True)
  355. control = ma.array([(1, 10), (2, 20), (-1, 30)],
  356. mask=[(0, 0), (0, 0), (1, 0)],
  357. dtype=[('f0', int), ('f1', int)])
  358. assert_equal(test, control)
  359. assert_equal(test.mask, control.mask)
  360. def test_flatten(self):
  361. # Test standard & flexible
  362. (_, x, _, z) = self.data
  363. test = merge_arrays((x, z), flatten=True)
  364. control = np.array([(1, 'A', 1.), (2, 'B', 2.)],
  365. dtype=[('f0', int), ('A', '|S3'), ('B', float)])
  366. assert_equal(test, control)
  367. test = merge_arrays((x, z), flatten=False)
  368. control = np.array([(1, ('A', 1.)), (2, ('B', 2.))],
  369. dtype=[('f0', int),
  370. ('f1', [('A', '|S3'), ('B', float)])])
  371. assert_equal(test, control)
  372. def test_flatten_wflexible(self):
  373. # Test flatten standard & nested
  374. (w, x, _, _) = self.data
  375. test = merge_arrays((x, w), flatten=True)
  376. control = np.array([(1, 1, 2, 3.0), (2, 4, 5, 6.0)],
  377. dtype=[('f0', int),
  378. ('a', int), ('ba', float), ('bb', int)])
  379. assert_equal(test, control)
  380. test = merge_arrays((x, w), flatten=False)
  381. controldtype = [('f0', int),
  382. ('f1', [('a', int),
  383. ('b', [('ba', float), ('bb', int), ('bc', [])])])]
  384. control = np.array([(1., (1, (2, 3.0, ()))), (2, (4, (5, 6.0, ())))],
  385. dtype=controldtype)
  386. assert_equal(test, control)
  387. def test_wmasked_arrays(self):
  388. # Test merge_arrays masked arrays
  389. (_, x, _, _) = self.data
  390. mx = ma.array([1, 2, 3], mask=[1, 0, 0])
  391. test = merge_arrays((x, mx), usemask=True)
  392. control = ma.array([(1, 1), (2, 2), (-1, 3)],
  393. mask=[(0, 1), (0, 0), (1, 0)],
  394. dtype=[('f0', int), ('f1', int)])
  395. assert_equal(test, control)
  396. test = merge_arrays((x, mx), usemask=True, asrecarray=True)
  397. assert_equal(test, control)
  398. assert_(isinstance(test, MaskedRecords))
  399. def test_w_singlefield(self):
  400. # Test single field
  401. test = merge_arrays((np.array([1, 2]).view([('a', int)]),
  402. np.array([10., 20., 30.])),)
  403. control = ma.array([(1, 10.), (2, 20.), (-1, 30.)],
  404. mask=[(0, 0), (0, 0), (1, 0)],
  405. dtype=[('a', int), ('f1', float)])
  406. assert_equal(test, control)
  407. def test_w_shorter_flex(self):
  408. # Test merge_arrays w/ a shorter flexndarray.
  409. z = self.data[-1]
  410. # Fixme, this test looks incomplete and broken
  411. #test = merge_arrays((z, np.array([10, 20, 30]).view([('C', int)])))
  412. #control = np.array([('A', 1., 10), ('B', 2., 20), ('-1', -1, 20)],
  413. # dtype=[('A', '|S3'), ('B', float), ('C', int)])
  414. #assert_equal(test, control)
  415. # Hack to avoid pyflakes warnings about unused variables
  416. merge_arrays((z, np.array([10, 20, 30]).view([('C', int)])))
  417. np.array([('A', 1., 10), ('B', 2., 20), ('-1', -1, 20)],
  418. dtype=[('A', '|S3'), ('B', float), ('C', int)])
  419. def test_singlerecord(self):
  420. (_, x, y, z) = self.data
  421. test = merge_arrays((x[0], y[0], z[0]), usemask=False)
  422. control = np.array([(1, 10, ('A', 1))],
  423. dtype=[('f0', int),
  424. ('f1', int),
  425. ('f2', [('A', '|S3'), ('B', float)])])
  426. assert_equal(test, control)
  427. class TestAppendFields:
  428. # Test append_fields
  429. def setup_method(self):
  430. x = np.array([1, 2, ])
  431. y = np.array([10, 20, 30])
  432. z = np.array(
  433. [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)])
  434. w = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
  435. dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
  436. self.data = (w, x, y, z)
  437. def test_append_single(self):
  438. # Test simple case
  439. (_, x, _, _) = self.data
  440. test = append_fields(x, 'A', data=[10, 20, 30])
  441. control = ma.array([(1, 10), (2, 20), (-1, 30)],
  442. mask=[(0, 0), (0, 0), (1, 0)],
  443. dtype=[('f0', int), ('A', int)],)
  444. assert_equal(test, control)
  445. def test_append_double(self):
  446. # Test simple case
  447. (_, x, _, _) = self.data
  448. test = append_fields(x, ('A', 'B'), data=[[10, 20, 30], [100, 200]])
  449. control = ma.array([(1, 10, 100), (2, 20, 200), (-1, 30, -1)],
  450. mask=[(0, 0, 0), (0, 0, 0), (1, 0, 1)],
  451. dtype=[('f0', int), ('A', int), ('B', int)],)
  452. assert_equal(test, control)
  453. def test_append_on_flex(self):
  454. # Test append_fields on flexible type arrays
  455. z = self.data[-1]
  456. test = append_fields(z, 'C', data=[10, 20, 30])
  457. control = ma.array([('A', 1., 10), ('B', 2., 20), (-1, -1., 30)],
  458. mask=[(0, 0, 0), (0, 0, 0), (1, 1, 0)],
  459. dtype=[('A', '|S3'), ('B', float), ('C', int)],)
  460. assert_equal(test, control)
  461. def test_append_on_nested(self):
  462. # Test append_fields on nested fields
  463. w = self.data[0]
  464. test = append_fields(w, 'C', data=[10, 20, 30])
  465. control = ma.array([(1, (2, 3.0), 10),
  466. (4, (5, 6.0), 20),
  467. (-1, (-1, -1.), 30)],
  468. mask=[(
  469. 0, (0, 0), 0), (0, (0, 0), 0), (1, (1, 1), 0)],
  470. dtype=[('a', int),
  471. ('b', [('ba', float), ('bb', int)]),
  472. ('C', int)],)
  473. assert_equal(test, control)
  474. class TestStackArrays:
  475. # Test stack_arrays
  476. def setup_method(self):
  477. x = np.array([1, 2, ])
  478. y = np.array([10, 20, 30])
  479. z = np.array(
  480. [('A', 1.), ('B', 2.)], dtype=[('A', '|S3'), ('B', float)])
  481. w = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
  482. dtype=[('a', int), ('b', [('ba', float), ('bb', int)])])
  483. self.data = (w, x, y, z)
  484. def test_solo(self):
  485. # Test stack_arrays on single arrays
  486. (_, x, _, _) = self.data
  487. test = stack_arrays((x,))
  488. assert_equal(test, x)
  489. assert_(test is x)
  490. test = stack_arrays(x)
  491. assert_equal(test, x)
  492. assert_(test is x)
  493. def test_unnamed_fields(self):
  494. # Tests combinations of arrays w/o named fields
  495. (_, x, y, _) = self.data
  496. test = stack_arrays((x, x), usemask=False)
  497. control = np.array([1, 2, 1, 2])
  498. assert_equal(test, control)
  499. test = stack_arrays((x, y), usemask=False)
  500. control = np.array([1, 2, 10, 20, 30])
  501. assert_equal(test, control)
  502. test = stack_arrays((y, x), usemask=False)
  503. control = np.array([10, 20, 30, 1, 2])
  504. assert_equal(test, control)
  505. def test_unnamed_and_named_fields(self):
  506. # Test combination of arrays w/ & w/o named fields
  507. (_, x, _, z) = self.data
  508. test = stack_arrays((x, z))
  509. control = ma.array([(1, -1, -1), (2, -1, -1),
  510. (-1, 'A', 1), (-1, 'B', 2)],
  511. mask=[(0, 1, 1), (0, 1, 1),
  512. (1, 0, 0), (1, 0, 0)],
  513. dtype=[('f0', int), ('A', '|S3'), ('B', float)])
  514. assert_equal(test, control)
  515. assert_equal(test.mask, control.mask)
  516. test = stack_arrays((z, x))
  517. control = ma.array([('A', 1, -1), ('B', 2, -1),
  518. (-1, -1, 1), (-1, -1, 2), ],
  519. mask=[(0, 0, 1), (0, 0, 1),
  520. (1, 1, 0), (1, 1, 0)],
  521. dtype=[('A', '|S3'), ('B', float), ('f2', int)])
  522. assert_equal(test, control)
  523. assert_equal(test.mask, control.mask)
  524. test = stack_arrays((z, z, x))
  525. control = ma.array([('A', 1, -1), ('B', 2, -1),
  526. ('A', 1, -1), ('B', 2, -1),
  527. (-1, -1, 1), (-1, -1, 2), ],
  528. mask=[(0, 0, 1), (0, 0, 1),
  529. (0, 0, 1), (0, 0, 1),
  530. (1, 1, 0), (1, 1, 0)],
  531. dtype=[('A', '|S3'), ('B', float), ('f2', int)])
  532. assert_equal(test, control)
  533. def test_matching_named_fields(self):
  534. # Test combination of arrays w/ matching field names
  535. (_, x, _, z) = self.data
  536. zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
  537. dtype=[('A', '|S3'), ('B', float), ('C', float)])
  538. test = stack_arrays((z, zz))
  539. control = ma.array([('A', 1, -1), ('B', 2, -1),
  540. (
  541. 'a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
  542. dtype=[('A', '|S3'), ('B', float), ('C', float)],
  543. mask=[(0, 0, 1), (0, 0, 1),
  544. (0, 0, 0), (0, 0, 0), (0, 0, 0)])
  545. assert_equal(test, control)
  546. assert_equal(test.mask, control.mask)
  547. test = stack_arrays((z, zz, x))
  548. ndtype = [('A', '|S3'), ('B', float), ('C', float), ('f3', int)]
  549. control = ma.array([('A', 1, -1, -1), ('B', 2, -1, -1),
  550. ('a', 10., 100., -1), ('b', 20., 200., -1),
  551. ('c', 30., 300., -1),
  552. (-1, -1, -1, 1), (-1, -1, -1, 2)],
  553. dtype=ndtype,
  554. mask=[(0, 0, 1, 1), (0, 0, 1, 1),
  555. (0, 0, 0, 1), (0, 0, 0, 1), (0, 0, 0, 1),
  556. (1, 1, 1, 0), (1, 1, 1, 0)])
  557. assert_equal(test, control)
  558. assert_equal(test.mask, control.mask)
  559. def test_defaults(self):
  560. # Test defaults: no exception raised if keys of defaults are not fields.
  561. (_, _, _, z) = self.data
  562. zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
  563. dtype=[('A', '|S3'), ('B', float), ('C', float)])
  564. defaults = {'A': '???', 'B': -999., 'C': -9999., 'D': -99999.}
  565. test = stack_arrays((z, zz), defaults=defaults)
  566. control = ma.array([('A', 1, -9999.), ('B', 2, -9999.),
  567. (
  568. 'a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
  569. dtype=[('A', '|S3'), ('B', float), ('C', float)],
  570. mask=[(0, 0, 1), (0, 0, 1),
  571. (0, 0, 0), (0, 0, 0), (0, 0, 0)])
  572. assert_equal(test, control)
  573. assert_equal(test.data, control.data)
  574. assert_equal(test.mask, control.mask)
  575. def test_autoconversion(self):
  576. # Tests autoconversion
  577. adtype = [('A', int), ('B', bool), ('C', float)]
  578. a = ma.array([(1, 2, 3)], mask=[(0, 1, 0)], dtype=adtype)
  579. bdtype = [('A', int), ('B', float), ('C', float)]
  580. b = ma.array([(4, 5, 6)], dtype=bdtype)
  581. control = ma.array([(1, 2, 3), (4, 5, 6)], mask=[(0, 1, 0), (0, 0, 0)],
  582. dtype=bdtype)
  583. test = stack_arrays((a, b), autoconvert=True)
  584. assert_equal(test, control)
  585. assert_equal(test.mask, control.mask)
  586. with assert_raises(TypeError):
  587. stack_arrays((a, b), autoconvert=False)
  588. def test_checktitles(self):
  589. # Test using titles in the field names
  590. adtype = [(('a', 'A'), int), (('b', 'B'), bool), (('c', 'C'), float)]
  591. a = ma.array([(1, 2, 3)], mask=[(0, 1, 0)], dtype=adtype)
  592. bdtype = [(('a', 'A'), int), (('b', 'B'), bool), (('c', 'C'), float)]
  593. b = ma.array([(4, 5, 6)], dtype=bdtype)
  594. test = stack_arrays((a, b))
  595. control = ma.array([(1, 2, 3), (4, 5, 6)], mask=[(0, 1, 0), (0, 0, 0)],
  596. dtype=bdtype)
  597. assert_equal(test, control)
  598. assert_equal(test.mask, control.mask)
  599. def test_subdtype(self):
  600. z = np.array([
  601. ('A', 1), ('B', 2)
  602. ], dtype=[('A', '|S3'), ('B', float, (1,))])
  603. zz = np.array([
  604. ('a', [10.], 100.), ('b', [20.], 200.), ('c', [30.], 300.)
  605. ], dtype=[('A', '|S3'), ('B', float, (1,)), ('C', float)])
  606. res = stack_arrays((z, zz))
  607. expected = ma.array(
  608. data=[
  609. (b'A', [1.0], 0),
  610. (b'B', [2.0], 0),
  611. (b'a', [10.0], 100.0),
  612. (b'b', [20.0], 200.0),
  613. (b'c', [30.0], 300.0)],
  614. mask=[
  615. (False, [False], True),
  616. (False, [False], True),
  617. (False, [False], False),
  618. (False, [False], False),
  619. (False, [False], False)
  620. ],
  621. dtype=zz.dtype
  622. )
  623. assert_equal(res.dtype, expected.dtype)
  624. assert_equal(res, expected)
  625. assert_equal(res.mask, expected.mask)
  626. class TestJoinBy:
  627. def setup_method(self):
  628. self.a = np.array(list(zip(np.arange(10), np.arange(50, 60),
  629. np.arange(100, 110))),
  630. dtype=[('a', int), ('b', int), ('c', int)])
  631. self.b = np.array(list(zip(np.arange(5, 15), np.arange(65, 75),
  632. np.arange(100, 110))),
  633. dtype=[('a', int), ('b', int), ('d', int)])
  634. def test_inner_join(self):
  635. # Basic test of join_by
  636. a, b = self.a, self.b
  637. test = join_by('a', a, b, jointype='inner')
  638. control = np.array([(5, 55, 65, 105, 100), (6, 56, 66, 106, 101),
  639. (7, 57, 67, 107, 102), (8, 58, 68, 108, 103),
  640. (9, 59, 69, 109, 104)],
  641. dtype=[('a', int), ('b1', int), ('b2', int),
  642. ('c', int), ('d', int)])
  643. assert_equal(test, control)
  644. def test_join(self):
  645. a, b = self.a, self.b
  646. # Fixme, this test is broken
  647. #test = join_by(('a', 'b'), a, b)
  648. #control = np.array([(5, 55, 105, 100), (6, 56, 106, 101),
  649. # (7, 57, 107, 102), (8, 58, 108, 103),
  650. # (9, 59, 109, 104)],
  651. # dtype=[('a', int), ('b', int),
  652. # ('c', int), ('d', int)])
  653. #assert_equal(test, control)
  654. # Hack to avoid pyflakes unused variable warnings
  655. join_by(('a', 'b'), a, b)
  656. np.array([(5, 55, 105, 100), (6, 56, 106, 101),
  657. (7, 57, 107, 102), (8, 58, 108, 103),
  658. (9, 59, 109, 104)],
  659. dtype=[('a', int), ('b', int),
  660. ('c', int), ('d', int)])
  661. def test_join_subdtype(self):
  662. # tests the bug in https://stackoverflow.com/q/44769632/102441
  663. foo = np.array([(1,)],
  664. dtype=[('key', int)])
  665. bar = np.array([(1, np.array([1,2,3]))],
  666. dtype=[('key', int), ('value', 'uint16', 3)])
  667. res = join_by('key', foo, bar)
  668. assert_equal(res, bar.view(ma.MaskedArray))
  669. def test_outer_join(self):
  670. a, b = self.a, self.b
  671. test = join_by(('a', 'b'), a, b, 'outer')
  672. control = ma.array([(0, 50, 100, -1), (1, 51, 101, -1),
  673. (2, 52, 102, -1), (3, 53, 103, -1),
  674. (4, 54, 104, -1), (5, 55, 105, -1),
  675. (5, 65, -1, 100), (6, 56, 106, -1),
  676. (6, 66, -1, 101), (7, 57, 107, -1),
  677. (7, 67, -1, 102), (8, 58, 108, -1),
  678. (8, 68, -1, 103), (9, 59, 109, -1),
  679. (9, 69, -1, 104), (10, 70, -1, 105),
  680. (11, 71, -1, 106), (12, 72, -1, 107),
  681. (13, 73, -1, 108), (14, 74, -1, 109)],
  682. mask=[(0, 0, 0, 1), (0, 0, 0, 1),
  683. (0, 0, 0, 1), (0, 0, 0, 1),
  684. (0, 0, 0, 1), (0, 0, 0, 1),
  685. (0, 0, 1, 0), (0, 0, 0, 1),
  686. (0, 0, 1, 0), (0, 0, 0, 1),
  687. (0, 0, 1, 0), (0, 0, 0, 1),
  688. (0, 0, 1, 0), (0, 0, 0, 1),
  689. (0, 0, 1, 0), (0, 0, 1, 0),
  690. (0, 0, 1, 0), (0, 0, 1, 0),
  691. (0, 0, 1, 0), (0, 0, 1, 0)],
  692. dtype=[('a', int), ('b', int),
  693. ('c', int), ('d', int)])
  694. assert_equal(test, control)
  695. def test_leftouter_join(self):
  696. a, b = self.a, self.b
  697. test = join_by(('a', 'b'), a, b, 'leftouter')
  698. control = ma.array([(0, 50, 100, -1), (1, 51, 101, -1),
  699. (2, 52, 102, -1), (3, 53, 103, -1),
  700. (4, 54, 104, -1), (5, 55, 105, -1),
  701. (6, 56, 106, -1), (7, 57, 107, -1),
  702. (8, 58, 108, -1), (9, 59, 109, -1)],
  703. mask=[(0, 0, 0, 1), (0, 0, 0, 1),
  704. (0, 0, 0, 1), (0, 0, 0, 1),
  705. (0, 0, 0, 1), (0, 0, 0, 1),
  706. (0, 0, 0, 1), (0, 0, 0, 1),
  707. (0, 0, 0, 1), (0, 0, 0, 1)],
  708. dtype=[('a', int), ('b', int), ('c', int), ('d', int)])
  709. assert_equal(test, control)
  710. def test_different_field_order(self):
  711. # gh-8940
  712. a = np.zeros(3, dtype=[('a', 'i4'), ('b', 'f4'), ('c', 'u1')])
  713. b = np.ones(3, dtype=[('c', 'u1'), ('b', 'f4'), ('a', 'i4')])
  714. # this should not give a FutureWarning:
  715. j = join_by(['c', 'b'], a, b, jointype='inner', usemask=False)
  716. assert_equal(j.dtype.names, ['b', 'c', 'a1', 'a2'])
  717. def test_duplicate_keys(self):
  718. a = np.zeros(3, dtype=[('a', 'i4'), ('b', 'f4'), ('c', 'u1')])
  719. b = np.ones(3, dtype=[('c', 'u1'), ('b', 'f4'), ('a', 'i4')])
  720. assert_raises(ValueError, join_by, ['a', 'b', 'b'], a, b)
  721. def test_same_name_different_dtypes_key(self):
  722. a_dtype = np.dtype([('key', 'S5'), ('value', '<f4')])
  723. b_dtype = np.dtype([('key', 'S10'), ('value', '<f4')])
  724. expected_dtype = np.dtype([
  725. ('key', 'S10'), ('value1', '<f4'), ('value2', '<f4')])
  726. a = np.array([('Sarah', 8.0), ('John', 6.0)], dtype=a_dtype)
  727. b = np.array([('Sarah', 10.0), ('John', 7.0)], dtype=b_dtype)
  728. res = join_by('key', a, b)
  729. assert_equal(res.dtype, expected_dtype)
  730. def test_same_name_different_dtypes(self):
  731. # gh-9338
  732. a_dtype = np.dtype([('key', 'S10'), ('value', '<f4')])
  733. b_dtype = np.dtype([('key', 'S10'), ('value', '<f8')])
  734. expected_dtype = np.dtype([
  735. ('key', '|S10'), ('value1', '<f4'), ('value2', '<f8')])
  736. a = np.array([('Sarah', 8.0), ('John', 6.0)], dtype=a_dtype)
  737. b = np.array([('Sarah', 10.0), ('John', 7.0)], dtype=b_dtype)
  738. res = join_by('key', a, b)
  739. assert_equal(res.dtype, expected_dtype)
  740. def test_subarray_key(self):
  741. a_dtype = np.dtype([('pos', int, 3), ('f', '<f4')])
  742. a = np.array([([1, 1, 1], np.pi), ([1, 2, 3], 0.0)], dtype=a_dtype)
  743. b_dtype = np.dtype([('pos', int, 3), ('g', '<f4')])
  744. b = np.array([([1, 1, 1], 3), ([3, 2, 1], 0.0)], dtype=b_dtype)
  745. expected_dtype = np.dtype([('pos', int, 3), ('f', '<f4'), ('g', '<f4')])
  746. expected = np.array([([1, 1, 1], np.pi, 3)], dtype=expected_dtype)
  747. res = join_by('pos', a, b)
  748. assert_equal(res.dtype, expected_dtype)
  749. assert_equal(res, expected)
  750. def test_padded_dtype(self):
  751. dt = np.dtype('i1,f4', align=True)
  752. dt.names = ('k', 'v')
  753. assert_(len(dt.descr), 3) # padding field is inserted
  754. a = np.array([(1, 3), (3, 2)], dt)
  755. b = np.array([(1, 1), (2, 2)], dt)
  756. res = join_by('k', a, b)
  757. # no padding fields remain
  758. expected_dtype = np.dtype([
  759. ('k', 'i1'), ('v1', 'f4'), ('v2', 'f4')
  760. ])
  761. assert_equal(res.dtype, expected_dtype)
  762. class TestJoinBy2:
  763. @classmethod
  764. def setup_method(cls):
  765. cls.a = np.array(list(zip(np.arange(10), np.arange(50, 60),
  766. np.arange(100, 110))),
  767. dtype=[('a', int), ('b', int), ('c', int)])
  768. cls.b = np.array(list(zip(np.arange(10), np.arange(65, 75),
  769. np.arange(100, 110))),
  770. dtype=[('a', int), ('b', int), ('d', int)])
  771. def test_no_r1postfix(self):
  772. # Basic test of join_by no_r1postfix
  773. a, b = self.a, self.b
  774. test = join_by(
  775. 'a', a, b, r1postfix='', r2postfix='2', jointype='inner')
  776. control = np.array([(0, 50, 65, 100, 100), (1, 51, 66, 101, 101),
  777. (2, 52, 67, 102, 102), (3, 53, 68, 103, 103),
  778. (4, 54, 69, 104, 104), (5, 55, 70, 105, 105),
  779. (6, 56, 71, 106, 106), (7, 57, 72, 107, 107),
  780. (8, 58, 73, 108, 108), (9, 59, 74, 109, 109)],
  781. dtype=[('a', int), ('b', int), ('b2', int),
  782. ('c', int), ('d', int)])
  783. assert_equal(test, control)
  784. def test_no_postfix(self):
  785. assert_raises(ValueError, join_by, 'a', self.a, self.b,
  786. r1postfix='', r2postfix='')
  787. def test_no_r2postfix(self):
  788. # Basic test of join_by no_r2postfix
  789. a, b = self.a, self.b
  790. test = join_by(
  791. 'a', a, b, r1postfix='1', r2postfix='', jointype='inner')
  792. control = np.array([(0, 50, 65, 100, 100), (1, 51, 66, 101, 101),
  793. (2, 52, 67, 102, 102), (3, 53, 68, 103, 103),
  794. (4, 54, 69, 104, 104), (5, 55, 70, 105, 105),
  795. (6, 56, 71, 106, 106), (7, 57, 72, 107, 107),
  796. (8, 58, 73, 108, 108), (9, 59, 74, 109, 109)],
  797. dtype=[('a', int), ('b1', int), ('b', int),
  798. ('c', int), ('d', int)])
  799. assert_equal(test, control)
  800. def test_two_keys_two_vars(self):
  801. a = np.array(list(zip(np.tile([10, 11], 5), np.repeat(np.arange(5), 2),
  802. np.arange(50, 60), np.arange(10, 20))),
  803. dtype=[('k', int), ('a', int), ('b', int), ('c', int)])
  804. b = np.array(list(zip(np.tile([10, 11], 5), np.repeat(np.arange(5), 2),
  805. np.arange(65, 75), np.arange(0, 10))),
  806. dtype=[('k', int), ('a', int), ('b', int), ('c', int)])
  807. control = np.array([(10, 0, 50, 65, 10, 0), (11, 0, 51, 66, 11, 1),
  808. (10, 1, 52, 67, 12, 2), (11, 1, 53, 68, 13, 3),
  809. (10, 2, 54, 69, 14, 4), (11, 2, 55, 70, 15, 5),
  810. (10, 3, 56, 71, 16, 6), (11, 3, 57, 72, 17, 7),
  811. (10, 4, 58, 73, 18, 8), (11, 4, 59, 74, 19, 9)],
  812. dtype=[('k', int), ('a', int), ('b1', int),
  813. ('b2', int), ('c1', int), ('c2', int)])
  814. test = join_by(
  815. ['a', 'k'], a, b, r1postfix='1', r2postfix='2', jointype='inner')
  816. assert_equal(test.dtype, control.dtype)
  817. assert_equal(test, control)
  818. class TestAppendFieldsObj:
  819. """
  820. Test append_fields with arrays containing objects
  821. """
  822. # https://github.com/numpy/numpy/issues/2346
  823. def setup_method(self):
  824. from datetime import date
  825. self.data = dict(obj=date(2000, 1, 1))
  826. def test_append_to_objects(self):
  827. "Test append_fields when the base array contains objects"
  828. obj = self.data['obj']
  829. x = np.array([(obj, 1.), (obj, 2.)],
  830. dtype=[('A', object), ('B', float)])
  831. y = np.array([10, 20], dtype=int)
  832. test = append_fields(x, 'C', data=y, usemask=False)
  833. control = np.array([(obj, 1.0, 10), (obj, 2.0, 20)],
  834. dtype=[('A', object), ('B', float), ('C', int)])
  835. assert_equal(test, control)