test_recfunctions.py 43 KB

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