recfunctions.py 55 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594
  1. """
  2. Collection of utilities to manipulate structured arrays.
  3. Most of these functions were initially implemented by John Hunter for
  4. matplotlib. They have been rewritten and extended for convenience.
  5. """
  6. import itertools
  7. import numpy as np
  8. import numpy.ma as ma
  9. from numpy import ndarray, recarray
  10. from numpy.ma import MaskedArray
  11. from numpy.ma.mrecords import MaskedRecords
  12. from numpy.core.overrides import array_function_dispatch
  13. from numpy.lib._iotools import _is_string_like
  14. from numpy.testing import suppress_warnings
  15. _check_fill_value = np.ma.core._check_fill_value
  16. __all__ = [
  17. 'append_fields', 'apply_along_fields', 'assign_fields_by_name',
  18. 'drop_fields', 'find_duplicates', 'flatten_descr',
  19. 'get_fieldstructure', 'get_names', 'get_names_flat',
  20. 'join_by', 'merge_arrays', 'rec_append_fields',
  21. 'rec_drop_fields', 'rec_join', 'recursive_fill_fields',
  22. 'rename_fields', 'repack_fields', 'require_fields',
  23. 'stack_arrays', 'structured_to_unstructured', 'unstructured_to_structured',
  24. ]
  25. def _recursive_fill_fields_dispatcher(input, output):
  26. return (input, output)
  27. @array_function_dispatch(_recursive_fill_fields_dispatcher)
  28. def recursive_fill_fields(input, output):
  29. """
  30. Fills fields from output with fields from input,
  31. with support for nested structures.
  32. Parameters
  33. ----------
  34. input : ndarray
  35. Input array.
  36. output : ndarray
  37. Output array.
  38. Notes
  39. -----
  40. * `output` should be at least the same size as `input`
  41. Examples
  42. --------
  43. >>> from numpy.lib import recfunctions as rfn
  44. >>> a = np.array([(1, 10.), (2, 20.)], dtype=[('A', np.int64), ('B', np.float64)])
  45. >>> b = np.zeros((3,), dtype=a.dtype)
  46. >>> rfn.recursive_fill_fields(a, b)
  47. array([(1, 10.), (2, 20.), (0, 0.)], dtype=[('A', '<i8'), ('B', '<f8')])
  48. """
  49. newdtype = output.dtype
  50. for field in newdtype.names:
  51. try:
  52. current = input[field]
  53. except ValueError:
  54. continue
  55. if current.dtype.names is not None:
  56. recursive_fill_fields(current, output[field])
  57. else:
  58. output[field][:len(current)] = current
  59. return output
  60. def _get_fieldspec(dtype):
  61. """
  62. Produce a list of name/dtype pairs corresponding to the dtype fields
  63. Similar to dtype.descr, but the second item of each tuple is a dtype, not a
  64. string. As a result, this handles subarray dtypes
  65. Can be passed to the dtype constructor to reconstruct the dtype, noting that
  66. this (deliberately) discards field offsets.
  67. Examples
  68. --------
  69. >>> dt = np.dtype([(('a', 'A'), np.int64), ('b', np.double, 3)])
  70. >>> dt.descr
  71. [(('a', 'A'), '<i8'), ('b', '<f8', (3,))]
  72. >>> _get_fieldspec(dt)
  73. [(('a', 'A'), dtype('int64')), ('b', dtype(('<f8', (3,))))]
  74. """
  75. if dtype.names is None:
  76. # .descr returns a nameless field, so we should too
  77. return [('', dtype)]
  78. else:
  79. fields = ((name, dtype.fields[name]) for name in dtype.names)
  80. # keep any titles, if present
  81. return [
  82. (name if len(f) == 2 else (f[2], name), f[0])
  83. for name, f in fields
  84. ]
  85. def get_names(adtype):
  86. """
  87. Returns the field names of the input datatype as a tuple.
  88. Parameters
  89. ----------
  90. adtype : dtype
  91. Input datatype
  92. Examples
  93. --------
  94. >>> from numpy.lib import recfunctions as rfn
  95. >>> rfn.get_names(np.empty((1,), dtype=int))
  96. Traceback (most recent call last):
  97. ...
  98. AttributeError: 'numpy.ndarray' object has no attribute 'names'
  99. >>> rfn.get_names(np.empty((1,), dtype=[('A',int), ('B', float)]))
  100. Traceback (most recent call last):
  101. ...
  102. AttributeError: 'numpy.ndarray' object has no attribute 'names'
  103. >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])])
  104. >>> rfn.get_names(adtype)
  105. ('a', ('b', ('ba', 'bb')))
  106. """
  107. listnames = []
  108. names = adtype.names
  109. for name in names:
  110. current = adtype[name]
  111. if current.names is not None:
  112. listnames.append((name, tuple(get_names(current))))
  113. else:
  114. listnames.append(name)
  115. return tuple(listnames)
  116. def get_names_flat(adtype):
  117. """
  118. Returns the field names of the input datatype as a tuple. Nested structure
  119. are flattened beforehand.
  120. Parameters
  121. ----------
  122. adtype : dtype
  123. Input datatype
  124. Examples
  125. --------
  126. >>> from numpy.lib import recfunctions as rfn
  127. >>> rfn.get_names_flat(np.empty((1,), dtype=int)) is None
  128. Traceback (most recent call last):
  129. ...
  130. AttributeError: 'numpy.ndarray' object has no attribute 'names'
  131. >>> rfn.get_names_flat(np.empty((1,), dtype=[('A',int), ('B', float)]))
  132. Traceback (most recent call last):
  133. ...
  134. AttributeError: 'numpy.ndarray' object has no attribute 'names'
  135. >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])])
  136. >>> rfn.get_names_flat(adtype)
  137. ('a', 'b', 'ba', 'bb')
  138. """
  139. listnames = []
  140. names = adtype.names
  141. for name in names:
  142. listnames.append(name)
  143. current = adtype[name]
  144. if current.names is not None:
  145. listnames.extend(get_names_flat(current))
  146. return tuple(listnames)
  147. def flatten_descr(ndtype):
  148. """
  149. Flatten a structured data-type description.
  150. Examples
  151. --------
  152. >>> from numpy.lib import recfunctions as rfn
  153. >>> ndtype = np.dtype([('a', '<i4'), ('b', [('ba', '<f8'), ('bb', '<i4')])])
  154. >>> rfn.flatten_descr(ndtype)
  155. (('a', dtype('int32')), ('ba', dtype('float64')), ('bb', dtype('int32')))
  156. """
  157. names = ndtype.names
  158. if names is None:
  159. return (('', ndtype),)
  160. else:
  161. descr = []
  162. for field in names:
  163. (typ, _) = ndtype.fields[field]
  164. if typ.names is not None:
  165. descr.extend(flatten_descr(typ))
  166. else:
  167. descr.append((field, typ))
  168. return tuple(descr)
  169. def _zip_dtype(seqarrays, flatten=False):
  170. newdtype = []
  171. if flatten:
  172. for a in seqarrays:
  173. newdtype.extend(flatten_descr(a.dtype))
  174. else:
  175. for a in seqarrays:
  176. current = a.dtype
  177. if current.names is not None and len(current.names) == 1:
  178. # special case - dtypes of 1 field are flattened
  179. newdtype.extend(_get_fieldspec(current))
  180. else:
  181. newdtype.append(('', current))
  182. return np.dtype(newdtype)
  183. def _zip_descr(seqarrays, flatten=False):
  184. """
  185. Combine the dtype description of a series of arrays.
  186. Parameters
  187. ----------
  188. seqarrays : sequence of arrays
  189. Sequence of arrays
  190. flatten : {boolean}, optional
  191. Whether to collapse nested descriptions.
  192. """
  193. return _zip_dtype(seqarrays, flatten=flatten).descr
  194. def get_fieldstructure(adtype, lastname=None, parents=None,):
  195. """
  196. Returns a dictionary with fields indexing lists of their parent fields.
  197. This function is used to simplify access to fields nested in other fields.
  198. Parameters
  199. ----------
  200. adtype : np.dtype
  201. Input datatype
  202. lastname : optional
  203. Last processed field name (used internally during recursion).
  204. parents : dictionary
  205. Dictionary of parent fields (used interbally during recursion).
  206. Examples
  207. --------
  208. >>> from numpy.lib import recfunctions as rfn
  209. >>> ndtype = np.dtype([('A', int),
  210. ... ('B', [('BA', int),
  211. ... ('BB', [('BBA', int), ('BBB', int)])])])
  212. >>> rfn.get_fieldstructure(ndtype)
  213. ... # XXX: possible regression, order of BBA and BBB is swapped
  214. {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'], 'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']}
  215. """
  216. if parents is None:
  217. parents = {}
  218. names = adtype.names
  219. for name in names:
  220. current = adtype[name]
  221. if current.names is not None:
  222. if lastname:
  223. parents[name] = [lastname, ]
  224. else:
  225. parents[name] = []
  226. parents.update(get_fieldstructure(current, name, parents))
  227. else:
  228. lastparent = [_ for _ in (parents.get(lastname, []) or [])]
  229. if lastparent:
  230. lastparent.append(lastname)
  231. elif lastname:
  232. lastparent = [lastname, ]
  233. parents[name] = lastparent or []
  234. return parents
  235. def _izip_fields_flat(iterable):
  236. """
  237. Returns an iterator of concatenated fields from a sequence of arrays,
  238. collapsing any nested structure.
  239. """
  240. for element in iterable:
  241. if isinstance(element, np.void):
  242. yield from _izip_fields_flat(tuple(element))
  243. else:
  244. yield element
  245. def _izip_fields(iterable):
  246. """
  247. Returns an iterator of concatenated fields from a sequence of arrays.
  248. """
  249. for element in iterable:
  250. if (hasattr(element, '__iter__') and
  251. not isinstance(element, str)):
  252. yield from _izip_fields(element)
  253. elif isinstance(element, np.void) and len(tuple(element)) == 1:
  254. # this statement is the same from the previous expression
  255. yield from _izip_fields(element)
  256. else:
  257. yield element
  258. def _izip_records(seqarrays, fill_value=None, flatten=True):
  259. """
  260. Returns an iterator of concatenated items from a sequence of arrays.
  261. Parameters
  262. ----------
  263. seqarrays : sequence of arrays
  264. Sequence of arrays.
  265. fill_value : {None, integer}
  266. Value used to pad shorter iterables.
  267. flatten : {True, False},
  268. Whether to
  269. """
  270. # Should we flatten the items, or just use a nested approach
  271. if flatten:
  272. zipfunc = _izip_fields_flat
  273. else:
  274. zipfunc = _izip_fields
  275. for tup in itertools.zip_longest(*seqarrays, fillvalue=fill_value):
  276. yield tuple(zipfunc(tup))
  277. def _fix_output(output, usemask=True, asrecarray=False):
  278. """
  279. Private function: return a recarray, a ndarray, a MaskedArray
  280. or a MaskedRecords depending on the input parameters
  281. """
  282. if not isinstance(output, MaskedArray):
  283. usemask = False
  284. if usemask:
  285. if asrecarray:
  286. output = output.view(MaskedRecords)
  287. else:
  288. output = ma.filled(output)
  289. if asrecarray:
  290. output = output.view(recarray)
  291. return output
  292. def _fix_defaults(output, defaults=None):
  293. """
  294. Update the fill_value and masked data of `output`
  295. from the default given in a dictionary defaults.
  296. """
  297. names = output.dtype.names
  298. (data, mask, fill_value) = (output.data, output.mask, output.fill_value)
  299. for (k, v) in (defaults or {}).items():
  300. if k in names:
  301. fill_value[k] = v
  302. data[k][mask[k]] = v
  303. return output
  304. def _merge_arrays_dispatcher(seqarrays, fill_value=None, flatten=None,
  305. usemask=None, asrecarray=None):
  306. return seqarrays
  307. @array_function_dispatch(_merge_arrays_dispatcher)
  308. def merge_arrays(seqarrays, fill_value=-1, flatten=False,
  309. usemask=False, asrecarray=False):
  310. """
  311. Merge arrays field by field.
  312. Parameters
  313. ----------
  314. seqarrays : sequence of ndarrays
  315. Sequence of arrays
  316. fill_value : {float}, optional
  317. Filling value used to pad missing data on the shorter arrays.
  318. flatten : {False, True}, optional
  319. Whether to collapse nested fields.
  320. usemask : {False, True}, optional
  321. Whether to return a masked array or not.
  322. asrecarray : {False, True}, optional
  323. Whether to return a recarray (MaskedRecords) or not.
  324. Examples
  325. --------
  326. >>> from numpy.lib import recfunctions as rfn
  327. >>> rfn.merge_arrays((np.array([1, 2]), np.array([10., 20., 30.])))
  328. array([( 1, 10.), ( 2, 20.), (-1, 30.)],
  329. dtype=[('f0', '<i8'), ('f1', '<f8')])
  330. >>> rfn.merge_arrays((np.array([1, 2], dtype=np.int64),
  331. ... np.array([10., 20., 30.])), usemask=False)
  332. array([(1, 10.0), (2, 20.0), (-1, 30.0)],
  333. dtype=[('f0', '<i8'), ('f1', '<f8')])
  334. >>> rfn.merge_arrays((np.array([1, 2]).view([('a', np.int64)]),
  335. ... np.array([10., 20., 30.])),
  336. ... usemask=False, asrecarray=True)
  337. rec.array([( 1, 10.), ( 2, 20.), (-1, 30.)],
  338. dtype=[('a', '<i8'), ('f1', '<f8')])
  339. Notes
  340. -----
  341. * Without a mask, the missing value will be filled with something,
  342. depending on what its corresponding type:
  343. * ``-1`` for integers
  344. * ``-1.0`` for floating point numbers
  345. * ``'-'`` for characters
  346. * ``'-1'`` for strings
  347. * ``True`` for boolean values
  348. * XXX: I just obtained these values empirically
  349. """
  350. # Only one item in the input sequence ?
  351. if (len(seqarrays) == 1):
  352. seqarrays = np.asanyarray(seqarrays[0])
  353. # Do we have a single ndarray as input ?
  354. if isinstance(seqarrays, (ndarray, np.void)):
  355. seqdtype = seqarrays.dtype
  356. # Make sure we have named fields
  357. if seqdtype.names is None:
  358. seqdtype = np.dtype([('', seqdtype)])
  359. if not flatten or _zip_dtype((seqarrays,), flatten=True) == seqdtype:
  360. # Minimal processing needed: just make sure everything's a-ok
  361. seqarrays = seqarrays.ravel()
  362. # Find what type of array we must return
  363. if usemask:
  364. if asrecarray:
  365. seqtype = MaskedRecords
  366. else:
  367. seqtype = MaskedArray
  368. elif asrecarray:
  369. seqtype = recarray
  370. else:
  371. seqtype = ndarray
  372. return seqarrays.view(dtype=seqdtype, type=seqtype)
  373. else:
  374. seqarrays = (seqarrays,)
  375. else:
  376. # Make sure we have arrays in the input sequence
  377. seqarrays = [np.asanyarray(_m) for _m in seqarrays]
  378. # Find the sizes of the inputs and their maximum
  379. sizes = tuple(a.size for a in seqarrays)
  380. maxlength = max(sizes)
  381. # Get the dtype of the output (flattening if needed)
  382. newdtype = _zip_dtype(seqarrays, flatten=flatten)
  383. # Initialize the sequences for data and mask
  384. seqdata = []
  385. seqmask = []
  386. # If we expect some kind of MaskedArray, make a special loop.
  387. if usemask:
  388. for (a, n) in zip(seqarrays, sizes):
  389. nbmissing = (maxlength - n)
  390. # Get the data and mask
  391. data = a.ravel().__array__()
  392. mask = ma.getmaskarray(a).ravel()
  393. # Get the filling value (if needed)
  394. if nbmissing:
  395. fval = _check_fill_value(fill_value, a.dtype)
  396. if isinstance(fval, (ndarray, np.void)):
  397. if len(fval.dtype) == 1:
  398. fval = fval.item()[0]
  399. fmsk = True
  400. else:
  401. fval = np.array(fval, dtype=a.dtype, ndmin=1)
  402. fmsk = np.ones((1,), dtype=mask.dtype)
  403. else:
  404. fval = None
  405. fmsk = True
  406. # Store an iterator padding the input to the expected length
  407. seqdata.append(itertools.chain(data, [fval] * nbmissing))
  408. seqmask.append(itertools.chain(mask, [fmsk] * nbmissing))
  409. # Create an iterator for the data
  410. data = tuple(_izip_records(seqdata, flatten=flatten))
  411. output = ma.array(np.fromiter(data, dtype=newdtype, count=maxlength),
  412. mask=list(_izip_records(seqmask, flatten=flatten)))
  413. if asrecarray:
  414. output = output.view(MaskedRecords)
  415. else:
  416. # Same as before, without the mask we don't need...
  417. for (a, n) in zip(seqarrays, sizes):
  418. nbmissing = (maxlength - n)
  419. data = a.ravel().__array__()
  420. if nbmissing:
  421. fval = _check_fill_value(fill_value, a.dtype)
  422. if isinstance(fval, (ndarray, np.void)):
  423. if len(fval.dtype) == 1:
  424. fval = fval.item()[0]
  425. else:
  426. fval = np.array(fval, dtype=a.dtype, ndmin=1)
  427. else:
  428. fval = None
  429. seqdata.append(itertools.chain(data, [fval] * nbmissing))
  430. output = np.fromiter(tuple(_izip_records(seqdata, flatten=flatten)),
  431. dtype=newdtype, count=maxlength)
  432. if asrecarray:
  433. output = output.view(recarray)
  434. # And we're done...
  435. return output
  436. def _drop_fields_dispatcher(base, drop_names, usemask=None, asrecarray=None):
  437. return (base,)
  438. @array_function_dispatch(_drop_fields_dispatcher)
  439. def drop_fields(base, drop_names, usemask=True, asrecarray=False):
  440. """
  441. Return a new array with fields in `drop_names` dropped.
  442. Nested fields are supported.
  443. .. versionchanged:: 1.18.0
  444. `drop_fields` returns an array with 0 fields if all fields are dropped,
  445. rather than returning ``None`` as it did previously.
  446. Parameters
  447. ----------
  448. base : array
  449. Input array
  450. drop_names : string or sequence
  451. String or sequence of strings corresponding to the names of the
  452. fields to drop.
  453. usemask : {False, True}, optional
  454. Whether to return a masked array or not.
  455. asrecarray : string or sequence, optional
  456. Whether to return a recarray or a mrecarray (`asrecarray=True`) or
  457. a plain ndarray or masked array with flexible dtype. The default
  458. is False.
  459. Examples
  460. --------
  461. >>> from numpy.lib import recfunctions as rfn
  462. >>> a = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
  463. ... dtype=[('a', np.int64), ('b', [('ba', np.double), ('bb', np.int64)])])
  464. >>> rfn.drop_fields(a, 'a')
  465. array([((2., 3),), ((5., 6),)],
  466. dtype=[('b', [('ba', '<f8'), ('bb', '<i8')])])
  467. >>> rfn.drop_fields(a, 'ba')
  468. array([(1, (3,)), (4, (6,))], dtype=[('a', '<i8'), ('b', [('bb', '<i8')])])
  469. >>> rfn.drop_fields(a, ['ba', 'bb'])
  470. array([(1,), (4,)], dtype=[('a', '<i8')])
  471. """
  472. if _is_string_like(drop_names):
  473. drop_names = [drop_names]
  474. else:
  475. drop_names = set(drop_names)
  476. def _drop_descr(ndtype, drop_names):
  477. names = ndtype.names
  478. newdtype = []
  479. for name in names:
  480. current = ndtype[name]
  481. if name in drop_names:
  482. continue
  483. if current.names is not None:
  484. descr = _drop_descr(current, drop_names)
  485. if descr:
  486. newdtype.append((name, descr))
  487. else:
  488. newdtype.append((name, current))
  489. return newdtype
  490. newdtype = _drop_descr(base.dtype, drop_names)
  491. output = np.empty(base.shape, dtype=newdtype)
  492. output = recursive_fill_fields(base, output)
  493. return _fix_output(output, usemask=usemask, asrecarray=asrecarray)
  494. def _keep_fields(base, keep_names, usemask=True, asrecarray=False):
  495. """
  496. Return a new array keeping only the fields in `keep_names`,
  497. and preserving the order of those fields.
  498. Parameters
  499. ----------
  500. base : array
  501. Input array
  502. keep_names : string or sequence
  503. String or sequence of strings corresponding to the names of the
  504. fields to keep. Order of the names will be preserved.
  505. usemask : {False, True}, optional
  506. Whether to return a masked array or not.
  507. asrecarray : string or sequence, optional
  508. Whether to return a recarray or a mrecarray (`asrecarray=True`) or
  509. a plain ndarray or masked array with flexible dtype. The default
  510. is False.
  511. """
  512. newdtype = [(n, base.dtype[n]) for n in keep_names]
  513. output = np.empty(base.shape, dtype=newdtype)
  514. output = recursive_fill_fields(base, output)
  515. return _fix_output(output, usemask=usemask, asrecarray=asrecarray)
  516. def _rec_drop_fields_dispatcher(base, drop_names):
  517. return (base,)
  518. @array_function_dispatch(_rec_drop_fields_dispatcher)
  519. def rec_drop_fields(base, drop_names):
  520. """
  521. Returns a new numpy.recarray with fields in `drop_names` dropped.
  522. """
  523. return drop_fields(base, drop_names, usemask=False, asrecarray=True)
  524. def _rename_fields_dispatcher(base, namemapper):
  525. return (base,)
  526. @array_function_dispatch(_rename_fields_dispatcher)
  527. def rename_fields(base, namemapper):
  528. """
  529. Rename the fields from a flexible-datatype ndarray or recarray.
  530. Nested fields are supported.
  531. Parameters
  532. ----------
  533. base : ndarray
  534. Input array whose fields must be modified.
  535. namemapper : dictionary
  536. Dictionary mapping old field names to their new version.
  537. Examples
  538. --------
  539. >>> from numpy.lib import recfunctions as rfn
  540. >>> a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))],
  541. ... dtype=[('a', int),('b', [('ba', float), ('bb', (float, 2))])])
  542. >>> rfn.rename_fields(a, {'a':'A', 'bb':'BB'})
  543. array([(1, (2., [ 3., 30.])), (4, (5., [ 6., 60.]))],
  544. dtype=[('A', '<i8'), ('b', [('ba', '<f8'), ('BB', '<f8', (2,))])])
  545. """
  546. def _recursive_rename_fields(ndtype, namemapper):
  547. newdtype = []
  548. for name in ndtype.names:
  549. newname = namemapper.get(name, name)
  550. current = ndtype[name]
  551. if current.names is not None:
  552. newdtype.append(
  553. (newname, _recursive_rename_fields(current, namemapper))
  554. )
  555. else:
  556. newdtype.append((newname, current))
  557. return newdtype
  558. newdtype = _recursive_rename_fields(base.dtype, namemapper)
  559. return base.view(newdtype)
  560. def _append_fields_dispatcher(base, names, data, dtypes=None,
  561. fill_value=None, usemask=None, asrecarray=None):
  562. yield base
  563. yield from data
  564. @array_function_dispatch(_append_fields_dispatcher)
  565. def append_fields(base, names, data, dtypes=None,
  566. fill_value=-1, usemask=True, asrecarray=False):
  567. """
  568. Add new fields to an existing array.
  569. The names of the fields are given with the `names` arguments,
  570. the corresponding values with the `data` arguments.
  571. If a single field is appended, `names`, `data` and `dtypes` do not have
  572. to be lists but just values.
  573. Parameters
  574. ----------
  575. base : array
  576. Input array to extend.
  577. names : string, sequence
  578. String or sequence of strings corresponding to the names
  579. of the new fields.
  580. data : array or sequence of arrays
  581. Array or sequence of arrays storing the fields to add to the base.
  582. dtypes : sequence of datatypes, optional
  583. Datatype or sequence of datatypes.
  584. If None, the datatypes are estimated from the `data`.
  585. fill_value : {float}, optional
  586. Filling value used to pad missing data on the shorter arrays.
  587. usemask : {False, True}, optional
  588. Whether to return a masked array or not.
  589. asrecarray : {False, True}, optional
  590. Whether to return a recarray (MaskedRecords) or not.
  591. """
  592. # Check the names
  593. if isinstance(names, (tuple, list)):
  594. if len(names) != len(data):
  595. msg = "The number of arrays does not match the number of names"
  596. raise ValueError(msg)
  597. elif isinstance(names, str):
  598. names = [names, ]
  599. data = [data, ]
  600. #
  601. if dtypes is None:
  602. data = [np.array(a, copy=False, subok=True) for a in data]
  603. data = [a.view([(name, a.dtype)]) for (name, a) in zip(names, data)]
  604. else:
  605. if not isinstance(dtypes, (tuple, list)):
  606. dtypes = [dtypes, ]
  607. if len(data) != len(dtypes):
  608. if len(dtypes) == 1:
  609. dtypes = dtypes * len(data)
  610. else:
  611. msg = "The dtypes argument must be None, a dtype, or a list."
  612. raise ValueError(msg)
  613. data = [np.array(a, copy=False, subok=True, dtype=d).view([(n, d)])
  614. for (a, n, d) in zip(data, names, dtypes)]
  615. #
  616. base = merge_arrays(base, usemask=usemask, fill_value=fill_value)
  617. if len(data) > 1:
  618. data = merge_arrays(data, flatten=True, usemask=usemask,
  619. fill_value=fill_value)
  620. else:
  621. data = data.pop()
  622. #
  623. output = ma.masked_all(
  624. max(len(base), len(data)),
  625. dtype=_get_fieldspec(base.dtype) + _get_fieldspec(data.dtype))
  626. output = recursive_fill_fields(base, output)
  627. output = recursive_fill_fields(data, output)
  628. #
  629. return _fix_output(output, usemask=usemask, asrecarray=asrecarray)
  630. def _rec_append_fields_dispatcher(base, names, data, dtypes=None):
  631. yield base
  632. yield from data
  633. @array_function_dispatch(_rec_append_fields_dispatcher)
  634. def rec_append_fields(base, names, data, dtypes=None):
  635. """
  636. Add new fields to an existing array.
  637. The names of the fields are given with the `names` arguments,
  638. the corresponding values with the `data` arguments.
  639. If a single field is appended, `names`, `data` and `dtypes` do not have
  640. to be lists but just values.
  641. Parameters
  642. ----------
  643. base : array
  644. Input array to extend.
  645. names : string, sequence
  646. String or sequence of strings corresponding to the names
  647. of the new fields.
  648. data : array or sequence of arrays
  649. Array or sequence of arrays storing the fields to add to the base.
  650. dtypes : sequence of datatypes, optional
  651. Datatype or sequence of datatypes.
  652. If None, the datatypes are estimated from the `data`.
  653. See Also
  654. --------
  655. append_fields
  656. Returns
  657. -------
  658. appended_array : np.recarray
  659. """
  660. return append_fields(base, names, data=data, dtypes=dtypes,
  661. asrecarray=True, usemask=False)
  662. def _repack_fields_dispatcher(a, align=None, recurse=None):
  663. return (a,)
  664. @array_function_dispatch(_repack_fields_dispatcher)
  665. def repack_fields(a, align=False, recurse=False):
  666. """
  667. Re-pack the fields of a structured array or dtype in memory.
  668. The memory layout of structured datatypes allows fields at arbitrary
  669. byte offsets. This means the fields can be separated by padding bytes,
  670. their offsets can be non-monotonically increasing, and they can overlap.
  671. This method removes any overlaps and reorders the fields in memory so they
  672. have increasing byte offsets, and adds or removes padding bytes depending
  673. on the `align` option, which behaves like the `align` option to `np.dtype`.
  674. If `align=False`, this method produces a "packed" memory layout in which
  675. each field starts at the byte the previous field ended, and any padding
  676. bytes are removed.
  677. If `align=True`, this methods produces an "aligned" memory layout in which
  678. each field's offset is a multiple of its alignment, and the total itemsize
  679. is a multiple of the largest alignment, by adding padding bytes as needed.
  680. Parameters
  681. ----------
  682. a : ndarray or dtype
  683. array or dtype for which to repack the fields.
  684. align : boolean
  685. If true, use an "aligned" memory layout, otherwise use a "packed" layout.
  686. recurse : boolean
  687. If True, also repack nested structures.
  688. Returns
  689. -------
  690. repacked : ndarray or dtype
  691. Copy of `a` with fields repacked, or `a` itself if no repacking was
  692. needed.
  693. Examples
  694. --------
  695. >>> from numpy.lib import recfunctions as rfn
  696. >>> def print_offsets(d):
  697. ... print("offsets:", [d.fields[name][1] for name in d.names])
  698. ... print("itemsize:", d.itemsize)
  699. ...
  700. >>> dt = np.dtype('u1, <i8, <f8', align=True)
  701. >>> dt
  702. dtype({'names':['f0','f1','f2'], 'formats':['u1','<i8','<f8'], 'offsets':[0,8,16], 'itemsize':24}, align=True)
  703. >>> print_offsets(dt)
  704. offsets: [0, 8, 16]
  705. itemsize: 24
  706. >>> packed_dt = rfn.repack_fields(dt)
  707. >>> packed_dt
  708. dtype([('f0', 'u1'), ('f1', '<i8'), ('f2', '<f8')])
  709. >>> print_offsets(packed_dt)
  710. offsets: [0, 1, 9]
  711. itemsize: 17
  712. """
  713. if not isinstance(a, np.dtype):
  714. dt = repack_fields(a.dtype, align=align, recurse=recurse)
  715. return a.astype(dt, copy=False)
  716. if a.names is None:
  717. return a
  718. fieldinfo = []
  719. for name in a.names:
  720. tup = a.fields[name]
  721. if recurse:
  722. fmt = repack_fields(tup[0], align=align, recurse=True)
  723. else:
  724. fmt = tup[0]
  725. if len(tup) == 3:
  726. name = (tup[2], name)
  727. fieldinfo.append((name, fmt))
  728. dt = np.dtype(fieldinfo, align=align)
  729. return np.dtype((a.type, dt))
  730. def _get_fields_and_offsets(dt, offset=0):
  731. """
  732. Returns a flat list of (dtype, count, offset) tuples of all the
  733. scalar fields in the dtype "dt", including nested fields, in left
  734. to right order.
  735. """
  736. # counts up elements in subarrays, including nested subarrays, and returns
  737. # base dtype and count
  738. def count_elem(dt):
  739. count = 1
  740. while dt.shape != ():
  741. for size in dt.shape:
  742. count *= size
  743. dt = dt.base
  744. return dt, count
  745. fields = []
  746. for name in dt.names:
  747. field = dt.fields[name]
  748. f_dt, f_offset = field[0], field[1]
  749. f_dt, n = count_elem(f_dt)
  750. if f_dt.names is None:
  751. fields.append((np.dtype((f_dt, (n,))), n, f_offset + offset))
  752. else:
  753. subfields = _get_fields_and_offsets(f_dt, f_offset + offset)
  754. size = f_dt.itemsize
  755. for i in range(n):
  756. if i == 0:
  757. # optimization: avoid list comprehension if no subarray
  758. fields.extend(subfields)
  759. else:
  760. fields.extend([(d, c, o + i*size) for d, c, o in subfields])
  761. return fields
  762. def _structured_to_unstructured_dispatcher(arr, dtype=None, copy=None,
  763. casting=None):
  764. return (arr,)
  765. @array_function_dispatch(_structured_to_unstructured_dispatcher)
  766. def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'):
  767. """
  768. Converts an n-D structured array into an (n+1)-D unstructured array.
  769. The new array will have a new last dimension equal in size to the
  770. number of field-elements of the input array. If not supplied, the output
  771. datatype is determined from the numpy type promotion rules applied to all
  772. the field datatypes.
  773. Nested fields, as well as each element of any subarray fields, all count
  774. as a single field-elements.
  775. Parameters
  776. ----------
  777. arr : ndarray
  778. Structured array or dtype to convert. Cannot contain object datatype.
  779. dtype : dtype, optional
  780. The dtype of the output unstructured array.
  781. copy : bool, optional
  782. See copy argument to `ndarray.astype`. If true, always return a copy.
  783. If false, and `dtype` requirements are satisfied, a view is returned.
  784. casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
  785. See casting argument of `ndarray.astype`. Controls what kind of data
  786. casting may occur.
  787. Returns
  788. -------
  789. unstructured : ndarray
  790. Unstructured array with one more dimension.
  791. Examples
  792. --------
  793. >>> from numpy.lib import recfunctions as rfn
  794. >>> a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
  795. >>> a
  796. array([(0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.]),
  797. (0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.])],
  798. dtype=[('a', '<i4'), ('b', [('f0', '<f4'), ('f1', '<u2')]), ('c', '<f4', (2,))])
  799. >>> rfn.structured_to_unstructured(a)
  800. array([[0., 0., 0., 0., 0.],
  801. [0., 0., 0., 0., 0.],
  802. [0., 0., 0., 0., 0.],
  803. [0., 0., 0., 0., 0.]])
  804. >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
  805. ... dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
  806. >>> np.mean(rfn.structured_to_unstructured(b[['x', 'z']]), axis=-1)
  807. array([ 3. , 5.5, 9. , 11. ])
  808. """
  809. if arr.dtype.names is None:
  810. raise ValueError('arr must be a structured array')
  811. fields = _get_fields_and_offsets(arr.dtype)
  812. n_fields = len(fields)
  813. if n_fields == 0 and dtype is None:
  814. raise ValueError("arr has no fields. Unable to guess dtype")
  815. elif n_fields == 0:
  816. # too many bugs elsewhere for this to work now
  817. raise NotImplementedError("arr with no fields is not supported")
  818. dts, counts, offsets = zip(*fields)
  819. names = ['f{}'.format(n) for n in range(n_fields)]
  820. if dtype is None:
  821. out_dtype = np.result_type(*[dt.base for dt in dts])
  822. else:
  823. out_dtype = dtype
  824. # Use a series of views and casts to convert to an unstructured array:
  825. # first view using flattened fields (doesn't work for object arrays)
  826. # Note: dts may include a shape for subarrays
  827. flattened_fields = np.dtype({'names': names,
  828. 'formats': dts,
  829. 'offsets': offsets,
  830. 'itemsize': arr.dtype.itemsize})
  831. with suppress_warnings() as sup: # until 1.16 (gh-12447)
  832. sup.filter(FutureWarning, "Numpy has detected")
  833. arr = arr.view(flattened_fields)
  834. # next cast to a packed format with all fields converted to new dtype
  835. packed_fields = np.dtype({'names': names,
  836. 'formats': [(out_dtype, dt.shape) for dt in dts]})
  837. arr = arr.astype(packed_fields, copy=copy, casting=casting)
  838. # finally is it safe to view the packed fields as the unstructured type
  839. return arr.view((out_dtype, (sum(counts),)))
  840. def _unstructured_to_structured_dispatcher(arr, dtype=None, names=None,
  841. align=None, copy=None, casting=None):
  842. return (arr,)
  843. @array_function_dispatch(_unstructured_to_structured_dispatcher)
  844. def unstructured_to_structured(arr, dtype=None, names=None, align=False,
  845. copy=False, casting='unsafe'):
  846. """
  847. Converts an n-D unstructured array into an (n-1)-D structured array.
  848. The last dimension of the input array is converted into a structure, with
  849. number of field-elements equal to the size of the last dimension of the
  850. input array. By default all output fields have the input array's dtype, but
  851. an output structured dtype with an equal number of fields-elements can be
  852. supplied instead.
  853. Nested fields, as well as each element of any subarray fields, all count
  854. towards the number of field-elements.
  855. Parameters
  856. ----------
  857. arr : ndarray
  858. Unstructured array or dtype to convert.
  859. dtype : dtype, optional
  860. The structured dtype of the output array
  861. names : list of strings, optional
  862. If dtype is not supplied, this specifies the field names for the output
  863. dtype, in order. The field dtypes will be the same as the input array.
  864. align : boolean, optional
  865. Whether to create an aligned memory layout.
  866. copy : bool, optional
  867. See copy argument to `ndarray.astype`. If true, always return a copy.
  868. If false, and `dtype` requirements are satisfied, a view is returned.
  869. casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
  870. See casting argument of `ndarray.astype`. Controls what kind of data
  871. casting may occur.
  872. Returns
  873. -------
  874. structured : ndarray
  875. Structured array with fewer dimensions.
  876. Examples
  877. --------
  878. >>> from numpy.lib import recfunctions as rfn
  879. >>> dt = np.dtype([('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
  880. >>> a = np.arange(20).reshape((4,5))
  881. >>> a
  882. array([[ 0, 1, 2, 3, 4],
  883. [ 5, 6, 7, 8, 9],
  884. [10, 11, 12, 13, 14],
  885. [15, 16, 17, 18, 19]])
  886. >>> rfn.unstructured_to_structured(a, dt)
  887. array([( 0, ( 1., 2), [ 3., 4.]), ( 5, ( 6., 7), [ 8., 9.]),
  888. (10, (11., 12), [13., 14.]), (15, (16., 17), [18., 19.])],
  889. dtype=[('a', '<i4'), ('b', [('f0', '<f4'), ('f1', '<u2')]), ('c', '<f4', (2,))])
  890. """
  891. if arr.shape == ():
  892. raise ValueError('arr must have at least one dimension')
  893. n_elem = arr.shape[-1]
  894. if n_elem == 0:
  895. # too many bugs elsewhere for this to work now
  896. raise NotImplementedError("last axis with size 0 is not supported")
  897. if dtype is None:
  898. if names is None:
  899. names = ['f{}'.format(n) for n in range(n_elem)]
  900. out_dtype = np.dtype([(n, arr.dtype) for n in names], align=align)
  901. fields = _get_fields_and_offsets(out_dtype)
  902. dts, counts, offsets = zip(*fields)
  903. else:
  904. if names is not None:
  905. raise ValueError("don't supply both dtype and names")
  906. # sanity check of the input dtype
  907. fields = _get_fields_and_offsets(dtype)
  908. if len(fields) == 0:
  909. dts, counts, offsets = [], [], []
  910. else:
  911. dts, counts, offsets = zip(*fields)
  912. if n_elem != sum(counts):
  913. raise ValueError('The length of the last dimension of arr must '
  914. 'be equal to the number of fields in dtype')
  915. out_dtype = dtype
  916. if align and not out_dtype.isalignedstruct:
  917. raise ValueError("align was True but dtype is not aligned")
  918. names = ['f{}'.format(n) for n in range(len(fields))]
  919. # Use a series of views and casts to convert to a structured array:
  920. # first view as a packed structured array of one dtype
  921. packed_fields = np.dtype({'names': names,
  922. 'formats': [(arr.dtype, dt.shape) for dt in dts]})
  923. arr = np.ascontiguousarray(arr).view(packed_fields)
  924. # next cast to an unpacked but flattened format with varied dtypes
  925. flattened_fields = np.dtype({'names': names,
  926. 'formats': dts,
  927. 'offsets': offsets,
  928. 'itemsize': out_dtype.itemsize})
  929. arr = arr.astype(flattened_fields, copy=copy, casting=casting)
  930. # finally view as the final nested dtype and remove the last axis
  931. return arr.view(out_dtype)[..., 0]
  932. def _apply_along_fields_dispatcher(func, arr):
  933. return (arr,)
  934. @array_function_dispatch(_apply_along_fields_dispatcher)
  935. def apply_along_fields(func, arr):
  936. """
  937. Apply function 'func' as a reduction across fields of a structured array.
  938. This is similar to `apply_along_axis`, but treats the fields of a
  939. structured array as an extra axis. The fields are all first cast to a
  940. common type following the type-promotion rules from `numpy.result_type`
  941. applied to the field's dtypes.
  942. Parameters
  943. ----------
  944. func : function
  945. Function to apply on the "field" dimension. This function must
  946. support an `axis` argument, like np.mean, np.sum, etc.
  947. arr : ndarray
  948. Structured array for which to apply func.
  949. Returns
  950. -------
  951. out : ndarray
  952. Result of the recution operation
  953. Examples
  954. --------
  955. >>> from numpy.lib import recfunctions as rfn
  956. >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
  957. ... dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
  958. >>> rfn.apply_along_fields(np.mean, b)
  959. array([ 2.66666667, 5.33333333, 8.66666667, 11. ])
  960. >>> rfn.apply_along_fields(np.mean, b[['x', 'z']])
  961. array([ 3. , 5.5, 9. , 11. ])
  962. """
  963. if arr.dtype.names is None:
  964. raise ValueError('arr must be a structured array')
  965. uarr = structured_to_unstructured(arr)
  966. return func(uarr, axis=-1)
  967. # works and avoids axis requirement, but very, very slow:
  968. #return np.apply_along_axis(func, -1, uarr)
  969. def _assign_fields_by_name_dispatcher(dst, src, zero_unassigned=None):
  970. return dst, src
  971. @array_function_dispatch(_assign_fields_by_name_dispatcher)
  972. def assign_fields_by_name(dst, src, zero_unassigned=True):
  973. """
  974. Assigns values from one structured array to another by field name.
  975. Normally in numpy >= 1.14, assignment of one structured array to another
  976. copies fields "by position", meaning that the first field from the src is
  977. copied to the first field of the dst, and so on, regardless of field name.
  978. This function instead copies "by field name", such that fields in the dst
  979. are assigned from the identically named field in the src. This applies
  980. recursively for nested structures. This is how structure assignment worked
  981. in numpy >= 1.6 to <= 1.13.
  982. Parameters
  983. ----------
  984. dst : ndarray
  985. src : ndarray
  986. The source and destination arrays during assignment.
  987. zero_unassigned : bool, optional
  988. If True, fields in the dst for which there was no matching
  989. field in the src are filled with the value 0 (zero). This
  990. was the behavior of numpy <= 1.13. If False, those fields
  991. are not modified.
  992. """
  993. if dst.dtype.names is None:
  994. dst[...] = src
  995. return
  996. for name in dst.dtype.names:
  997. if name not in src.dtype.names:
  998. if zero_unassigned:
  999. dst[name] = 0
  1000. else:
  1001. assign_fields_by_name(dst[name], src[name],
  1002. zero_unassigned)
  1003. def _require_fields_dispatcher(array, required_dtype):
  1004. return (array,)
  1005. @array_function_dispatch(_require_fields_dispatcher)
  1006. def require_fields(array, required_dtype):
  1007. """
  1008. Casts a structured array to a new dtype using assignment by field-name.
  1009. This function assigns from the old to the new array by name, so the
  1010. value of a field in the output array is the value of the field with the
  1011. same name in the source array. This has the effect of creating a new
  1012. ndarray containing only the fields "required" by the required_dtype.
  1013. If a field name in the required_dtype does not exist in the
  1014. input array, that field is created and set to 0 in the output array.
  1015. Parameters
  1016. ----------
  1017. a : ndarray
  1018. array to cast
  1019. required_dtype : dtype
  1020. datatype for output array
  1021. Returns
  1022. -------
  1023. out : ndarray
  1024. array with the new dtype, with field values copied from the fields in
  1025. the input array with the same name
  1026. Examples
  1027. --------
  1028. >>> from numpy.lib import recfunctions as rfn
  1029. >>> a = np.ones(4, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')])
  1030. >>> rfn.require_fields(a, [('b', 'f4'), ('c', 'u1')])
  1031. array([(1., 1), (1., 1), (1., 1), (1., 1)],
  1032. dtype=[('b', '<f4'), ('c', 'u1')])
  1033. >>> rfn.require_fields(a, [('b', 'f4'), ('newf', 'u1')])
  1034. array([(1., 0), (1., 0), (1., 0), (1., 0)],
  1035. dtype=[('b', '<f4'), ('newf', 'u1')])
  1036. """
  1037. out = np.empty(array.shape, dtype=required_dtype)
  1038. assign_fields_by_name(out, array)
  1039. return out
  1040. def _stack_arrays_dispatcher(arrays, defaults=None, usemask=None,
  1041. asrecarray=None, autoconvert=None):
  1042. return arrays
  1043. @array_function_dispatch(_stack_arrays_dispatcher)
  1044. def stack_arrays(arrays, defaults=None, usemask=True, asrecarray=False,
  1045. autoconvert=False):
  1046. """
  1047. Superposes arrays fields by fields
  1048. Parameters
  1049. ----------
  1050. arrays : array or sequence
  1051. Sequence of input arrays.
  1052. defaults : dictionary, optional
  1053. Dictionary mapping field names to the corresponding default values.
  1054. usemask : {True, False}, optional
  1055. Whether to return a MaskedArray (or MaskedRecords is
  1056. `asrecarray==True`) or a ndarray.
  1057. asrecarray : {False, True}, optional
  1058. Whether to return a recarray (or MaskedRecords if `usemask==True`)
  1059. or just a flexible-type ndarray.
  1060. autoconvert : {False, True}, optional
  1061. Whether automatically cast the type of the field to the maximum.
  1062. Examples
  1063. --------
  1064. >>> from numpy.lib import recfunctions as rfn
  1065. >>> x = np.array([1, 2,])
  1066. >>> rfn.stack_arrays(x) is x
  1067. True
  1068. >>> z = np.array([('A', 1), ('B', 2)], dtype=[('A', '|S3'), ('B', float)])
  1069. >>> zz = np.array([('a', 10., 100.), ('b', 20., 200.), ('c', 30., 300.)],
  1070. ... dtype=[('A', '|S3'), ('B', np.double), ('C', np.double)])
  1071. >>> test = rfn.stack_arrays((z,zz))
  1072. >>> test
  1073. masked_array(data=[(b'A', 1.0, --), (b'B', 2.0, --), (b'a', 10.0, 100.0),
  1074. (b'b', 20.0, 200.0), (b'c', 30.0, 300.0)],
  1075. mask=[(False, False, True), (False, False, True),
  1076. (False, False, False), (False, False, False),
  1077. (False, False, False)],
  1078. fill_value=(b'N/A', 1.e+20, 1.e+20),
  1079. dtype=[('A', 'S3'), ('B', '<f8'), ('C', '<f8')])
  1080. """
  1081. if isinstance(arrays, ndarray):
  1082. return arrays
  1083. elif len(arrays) == 1:
  1084. return arrays[0]
  1085. seqarrays = [np.asanyarray(a).ravel() for a in arrays]
  1086. nrecords = [len(a) for a in seqarrays]
  1087. ndtype = [a.dtype for a in seqarrays]
  1088. fldnames = [d.names for d in ndtype]
  1089. #
  1090. dtype_l = ndtype[0]
  1091. newdescr = _get_fieldspec(dtype_l)
  1092. names = [n for n, d in newdescr]
  1093. for dtype_n in ndtype[1:]:
  1094. for fname, fdtype in _get_fieldspec(dtype_n):
  1095. if fname not in names:
  1096. newdescr.append((fname, fdtype))
  1097. names.append(fname)
  1098. else:
  1099. nameidx = names.index(fname)
  1100. _, cdtype = newdescr[nameidx]
  1101. if autoconvert:
  1102. newdescr[nameidx] = (fname, max(fdtype, cdtype))
  1103. elif fdtype != cdtype:
  1104. raise TypeError("Incompatible type '%s' <> '%s'" %
  1105. (cdtype, fdtype))
  1106. # Only one field: use concatenate
  1107. if len(newdescr) == 1:
  1108. output = ma.concatenate(seqarrays)
  1109. else:
  1110. #
  1111. output = ma.masked_all((np.sum(nrecords),), newdescr)
  1112. offset = np.cumsum(np.r_[0, nrecords])
  1113. seen = []
  1114. for (a, n, i, j) in zip(seqarrays, fldnames, offset[:-1], offset[1:]):
  1115. names = a.dtype.names
  1116. if names is None:
  1117. output['f%i' % len(seen)][i:j] = a
  1118. else:
  1119. for name in n:
  1120. output[name][i:j] = a[name]
  1121. if name not in seen:
  1122. seen.append(name)
  1123. #
  1124. return _fix_output(_fix_defaults(output, defaults),
  1125. usemask=usemask, asrecarray=asrecarray)
  1126. def _find_duplicates_dispatcher(
  1127. a, key=None, ignoremask=None, return_index=None):
  1128. return (a,)
  1129. @array_function_dispatch(_find_duplicates_dispatcher)
  1130. def find_duplicates(a, key=None, ignoremask=True, return_index=False):
  1131. """
  1132. Find the duplicates in a structured array along a given key
  1133. Parameters
  1134. ----------
  1135. a : array-like
  1136. Input array
  1137. key : {string, None}, optional
  1138. Name of the fields along which to check the duplicates.
  1139. If None, the search is performed by records
  1140. ignoremask : {True, False}, optional
  1141. Whether masked data should be discarded or considered as duplicates.
  1142. return_index : {False, True}, optional
  1143. Whether to return the indices of the duplicated values.
  1144. Examples
  1145. --------
  1146. >>> from numpy.lib import recfunctions as rfn
  1147. >>> ndtype = [('a', int)]
  1148. >>> a = np.ma.array([1, 1, 1, 2, 2, 3, 3],
  1149. ... mask=[0, 0, 1, 0, 0, 0, 1]).view(ndtype)
  1150. >>> rfn.find_duplicates(a, ignoremask=True, return_index=True)
  1151. (masked_array(data=[(1,), (1,), (2,), (2,)],
  1152. mask=[(False,), (False,), (False,), (False,)],
  1153. fill_value=(999999,),
  1154. dtype=[('a', '<i8')]), array([0, 1, 3, 4]))
  1155. """
  1156. a = np.asanyarray(a).ravel()
  1157. # Get a dictionary of fields
  1158. fields = get_fieldstructure(a.dtype)
  1159. # Get the sorting data (by selecting the corresponding field)
  1160. base = a
  1161. if key:
  1162. for f in fields[key]:
  1163. base = base[f]
  1164. base = base[key]
  1165. # Get the sorting indices and the sorted data
  1166. sortidx = base.argsort()
  1167. sortedbase = base[sortidx]
  1168. sorteddata = sortedbase.filled()
  1169. # Compare the sorting data
  1170. flag = (sorteddata[:-1] == sorteddata[1:])
  1171. # If masked data must be ignored, set the flag to false where needed
  1172. if ignoremask:
  1173. sortedmask = sortedbase.recordmask
  1174. flag[sortedmask[1:]] = False
  1175. flag = np.concatenate(([False], flag))
  1176. # We need to take the point on the left as well (else we're missing it)
  1177. flag[:-1] = flag[:-1] + flag[1:]
  1178. duplicates = a[sortidx][flag]
  1179. if return_index:
  1180. return (duplicates, sortidx[flag])
  1181. else:
  1182. return duplicates
  1183. def _join_by_dispatcher(
  1184. key, r1, r2, jointype=None, r1postfix=None, r2postfix=None,
  1185. defaults=None, usemask=None, asrecarray=None):
  1186. return (r1, r2)
  1187. @array_function_dispatch(_join_by_dispatcher)
  1188. def join_by(key, r1, r2, jointype='inner', r1postfix='1', r2postfix='2',
  1189. defaults=None, usemask=True, asrecarray=False):
  1190. """
  1191. Join arrays `r1` and `r2` on key `key`.
  1192. The key should be either a string or a sequence of string corresponding
  1193. to the fields used to join the array. An exception is raised if the
  1194. `key` field cannot be found in the two input arrays. Neither `r1` nor
  1195. `r2` should have any duplicates along `key`: the presence of duplicates
  1196. will make the output quite unreliable. Note that duplicates are not
  1197. looked for by the algorithm.
  1198. Parameters
  1199. ----------
  1200. key : {string, sequence}
  1201. A string or a sequence of strings corresponding to the fields used
  1202. for comparison.
  1203. r1, r2 : arrays
  1204. Structured arrays.
  1205. jointype : {'inner', 'outer', 'leftouter'}, optional
  1206. If 'inner', returns the elements common to both r1 and r2.
  1207. If 'outer', returns the common elements as well as the elements of
  1208. r1 not in r2 and the elements of not in r2.
  1209. If 'leftouter', returns the common elements and the elements of r1
  1210. not in r2.
  1211. r1postfix : string, optional
  1212. String appended to the names of the fields of r1 that are present
  1213. in r2 but absent of the key.
  1214. r2postfix : string, optional
  1215. String appended to the names of the fields of r2 that are present
  1216. in r1 but absent of the key.
  1217. defaults : {dictionary}, optional
  1218. Dictionary mapping field names to the corresponding default values.
  1219. usemask : {True, False}, optional
  1220. Whether to return a MaskedArray (or MaskedRecords is
  1221. `asrecarray==True`) or a ndarray.
  1222. asrecarray : {False, True}, optional
  1223. Whether to return a recarray (or MaskedRecords if `usemask==True`)
  1224. or just a flexible-type ndarray.
  1225. Notes
  1226. -----
  1227. * The output is sorted along the key.
  1228. * A temporary array is formed by dropping the fields not in the key for
  1229. the two arrays and concatenating the result. This array is then
  1230. sorted, and the common entries selected. The output is constructed by
  1231. filling the fields with the selected entries. Matching is not
  1232. preserved if there are some duplicates...
  1233. """
  1234. # Check jointype
  1235. if jointype not in ('inner', 'outer', 'leftouter'):
  1236. raise ValueError(
  1237. "The 'jointype' argument should be in 'inner', "
  1238. "'outer' or 'leftouter' (got '%s' instead)" % jointype
  1239. )
  1240. # If we have a single key, put it in a tuple
  1241. if isinstance(key, str):
  1242. key = (key,)
  1243. # Check the keys
  1244. if len(set(key)) != len(key):
  1245. dup = next(x for n,x in enumerate(key) if x in key[n+1:])
  1246. raise ValueError("duplicate join key %r" % dup)
  1247. for name in key:
  1248. if name not in r1.dtype.names:
  1249. raise ValueError('r1 does not have key field %r' % name)
  1250. if name not in r2.dtype.names:
  1251. raise ValueError('r2 does not have key field %r' % name)
  1252. # Make sure we work with ravelled arrays
  1253. r1 = r1.ravel()
  1254. r2 = r2.ravel()
  1255. # Fixme: nb2 below is never used. Commenting out for pyflakes.
  1256. # (nb1, nb2) = (len(r1), len(r2))
  1257. nb1 = len(r1)
  1258. (r1names, r2names) = (r1.dtype.names, r2.dtype.names)
  1259. # Check the names for collision
  1260. collisions = (set(r1names) & set(r2names)) - set(key)
  1261. if collisions and not (r1postfix or r2postfix):
  1262. msg = "r1 and r2 contain common names, r1postfix and r2postfix "
  1263. msg += "can't both be empty"
  1264. raise ValueError(msg)
  1265. # Make temporary arrays of just the keys
  1266. # (use order of keys in `r1` for back-compatibility)
  1267. key1 = [ n for n in r1names if n in key ]
  1268. r1k = _keep_fields(r1, key1)
  1269. r2k = _keep_fields(r2, key1)
  1270. # Concatenate the two arrays for comparison
  1271. aux = ma.concatenate((r1k, r2k))
  1272. idx_sort = aux.argsort(order=key)
  1273. aux = aux[idx_sort]
  1274. #
  1275. # Get the common keys
  1276. flag_in = ma.concatenate(([False], aux[1:] == aux[:-1]))
  1277. flag_in[:-1] = flag_in[1:] + flag_in[:-1]
  1278. idx_in = idx_sort[flag_in]
  1279. idx_1 = idx_in[(idx_in < nb1)]
  1280. idx_2 = idx_in[(idx_in >= nb1)] - nb1
  1281. (r1cmn, r2cmn) = (len(idx_1), len(idx_2))
  1282. if jointype == 'inner':
  1283. (r1spc, r2spc) = (0, 0)
  1284. elif jointype == 'outer':
  1285. idx_out = idx_sort[~flag_in]
  1286. idx_1 = np.concatenate((idx_1, idx_out[(idx_out < nb1)]))
  1287. idx_2 = np.concatenate((idx_2, idx_out[(idx_out >= nb1)] - nb1))
  1288. (r1spc, r2spc) = (len(idx_1) - r1cmn, len(idx_2) - r2cmn)
  1289. elif jointype == 'leftouter':
  1290. idx_out = idx_sort[~flag_in]
  1291. idx_1 = np.concatenate((idx_1, idx_out[(idx_out < nb1)]))
  1292. (r1spc, r2spc) = (len(idx_1) - r1cmn, 0)
  1293. # Select the entries from each input
  1294. (s1, s2) = (r1[idx_1], r2[idx_2])
  1295. #
  1296. # Build the new description of the output array .......
  1297. # Start with the key fields
  1298. ndtype = _get_fieldspec(r1k.dtype)
  1299. # Add the fields from r1
  1300. for fname, fdtype in _get_fieldspec(r1.dtype):
  1301. if fname not in key:
  1302. ndtype.append((fname, fdtype))
  1303. # Add the fields from r2
  1304. for fname, fdtype in _get_fieldspec(r2.dtype):
  1305. # Have we seen the current name already ?
  1306. # we need to rebuild this list every time
  1307. names = list(name for name, dtype in ndtype)
  1308. try:
  1309. nameidx = names.index(fname)
  1310. except ValueError:
  1311. #... we haven't: just add the description to the current list
  1312. ndtype.append((fname, fdtype))
  1313. else:
  1314. # collision
  1315. _, cdtype = ndtype[nameidx]
  1316. if fname in key:
  1317. # The current field is part of the key: take the largest dtype
  1318. ndtype[nameidx] = (fname, max(fdtype, cdtype))
  1319. else:
  1320. # The current field is not part of the key: add the suffixes,
  1321. # and place the new field adjacent to the old one
  1322. ndtype[nameidx:nameidx + 1] = [
  1323. (fname + r1postfix, cdtype),
  1324. (fname + r2postfix, fdtype)
  1325. ]
  1326. # Rebuild a dtype from the new fields
  1327. ndtype = np.dtype(ndtype)
  1328. # Find the largest nb of common fields :
  1329. # r1cmn and r2cmn should be equal, but...
  1330. cmn = max(r1cmn, r2cmn)
  1331. # Construct an empty array
  1332. output = ma.masked_all((cmn + r1spc + r2spc,), dtype=ndtype)
  1333. names = output.dtype.names
  1334. for f in r1names:
  1335. selected = s1[f]
  1336. if f not in names or (f in r2names and not r2postfix and f not in key):
  1337. f += r1postfix
  1338. current = output[f]
  1339. current[:r1cmn] = selected[:r1cmn]
  1340. if jointype in ('outer', 'leftouter'):
  1341. current[cmn:cmn + r1spc] = selected[r1cmn:]
  1342. for f in r2names:
  1343. selected = s2[f]
  1344. if f not in names or (f in r1names and not r1postfix and f not in key):
  1345. f += r2postfix
  1346. current = output[f]
  1347. current[:r2cmn] = selected[:r2cmn]
  1348. if (jointype == 'outer') and r2spc:
  1349. current[-r2spc:] = selected[r2cmn:]
  1350. # Sort and finalize the output
  1351. output.sort(order=key)
  1352. kwargs = dict(usemask=usemask, asrecarray=asrecarray)
  1353. return _fix_output(_fix_defaults(output, defaults), **kwargs)
  1354. def _rec_join_dispatcher(
  1355. key, r1, r2, jointype=None, r1postfix=None, r2postfix=None,
  1356. defaults=None):
  1357. return (r1, r2)
  1358. @array_function_dispatch(_rec_join_dispatcher)
  1359. def rec_join(key, r1, r2, jointype='inner', r1postfix='1', r2postfix='2',
  1360. defaults=None):
  1361. """
  1362. Join arrays `r1` and `r2` on keys.
  1363. Alternative to join_by, that always returns a np.recarray.
  1364. See Also
  1365. --------
  1366. join_by : equivalent function
  1367. """
  1368. kwargs = dict(jointype=jointype, r1postfix=r1postfix, r2postfix=r2postfix,
  1369. defaults=defaults, usemask=False, asrecarray=True)
  1370. return join_by(key, r1, r2, **kwargs)