test_recfunctions.py 43 KB

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