test_recfunctions.py 43 KB

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