codekingpro/portable-devtools
114k
1"""
2Collection of utilities to manipulate structured arrays.
3
4Most of these functions were initially implemented by John Hunter for
5matplotlib. They have been rewritten and extended for convenience.
6
7"""
8import itertools
9
10import numpy as np
11import numpy.ma as ma
12import numpy.ma.mrecords as mrec
13from numpy._core.overrides import array_function_dispatch
14from numpy.lib._iotools import _is_string_like
15
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
26
27def _recursive_fill_fields_dispatcher(input, output):
28 return (input, output)
29
30
31@array_function_dispatch(_recursive_fill_fields_dispatcher)
32def recursive_fill_fields(input, output):
33 """
34 Fills fields from output with fields from input,
35 with support for nested structures.
36
37 Parameters
38 ----------
39 input : ndarray
40 Input array.
41 output : ndarray
42 Output array.
43
44 Notes
45 -----
46 * `output` should be at least the same size as `input`
47
48 Examples
49 --------
50 >>> import numpy as np
51 >>> from numpy.lib import recfunctions as rfn
52 >>> a = np.array([(1, 10.), (2, 20.)], dtype=[('A', np.int64), ('B', np.float64)])
53 >>> b = np.zeros((3,), dtype=a.dtype)
54 >>> rfn.recursive_fill_fields(a, b)
55 array([(1, 10.), (2, 20.), (0, 0.)], dtype=[('A', '<i8'), ('B', '<f8')])
56
57 """
58 newdtype = output.dtype
59 for field in newdtype.names:
60 try:
61 current = input[field]
62 except ValueError:
63 continue
64 if current.dtype.names is not None:
65 recursive_fill_fields(current, output[field])
66 else:
67 output[field][:len(current)] = current
68 return output
69
70
71def _get_fieldspec(dtype):
72 """
73 Produce a list of name/dtype pairs corresponding to the dtype fields
74
75 Similar to dtype.descr, but the second item of each tuple is a dtype, not a
76 string. As a result, this handles subarray dtypes
77
78 Can be passed to the dtype constructor to reconstruct the dtype, noting that
79 this (deliberately) discards field offsets.
80
81 Examples
82 --------
83 >>> import numpy as np
84 >>> dt = np.dtype([(('a', 'A'), np.int64), ('b', np.double, 3)])
85 >>> dt.descr
86 [(('a', 'A'), '<i8'), ('b', '<f8', (3,))]
87 >>> _get_fieldspec(dt)
88 [(('a', 'A'), dtype('int64')), ('b', dtype(('<f8', (3,))))]
89
90 """
91 if dtype.names is None:
92 # .descr returns a nameless field, so we should too
93 return [('', dtype)]
94 else:
95 fields = ((name, dtype.fields[name]) for name in dtype.names)
96 # keep any titles, if present
97 return [
98 (name if len(f) == 2 else (f[2], name), f[0])
99 for name, f in fields
100 ]
101
102
103def get_names(adtype):
104 """
105 Returns the field names of the input datatype as a tuple. Input datatype
106 must have fields otherwise error is raised.
107
108 Parameters
109 ----------
110 adtype : dtype
111 Input datatype
112
113 Examples
114 --------
115 >>> import numpy as np
116 >>> from numpy.lib import recfunctions as rfn
117 >>> rfn.get_names(np.empty((1,), dtype=[('A', int)]).dtype)
118 ('A',)
119 >>> rfn.get_names(np.empty((1,), dtype=[('A',int), ('B', float)]).dtype)
120 ('A', 'B')
121 >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])])
122 >>> rfn.get_names(adtype)
123 ('a', ('b', ('ba', 'bb')))
124 """
125 listnames = []
126 names = adtype.names
127 for name in names:
128 current = adtype[name]
129 if current.names is not None:
130 listnames.append((name, tuple(get_names(current))))
131 else:
132 listnames.append(name)
133 return tuple(listnames)
134
135
136def get_names_flat(adtype):
137 """
138 Returns the field names of the input datatype as a tuple. Input datatype
139 must have fields otherwise error is raised.
140 Nested structure are flattened beforehand.
141
142 Parameters
143 ----------
144 adtype : dtype
145 Input datatype
146
147 Examples
148 --------
149 >>> import numpy as np
150 >>> from numpy.lib import recfunctions as rfn
151 >>> rfn.get_names_flat(np.empty((1,), dtype=[('A', int)]).dtype) is None
152 False
153 >>> rfn.get_names_flat(np.empty((1,), dtype=[('A',int), ('B', str)]).dtype)
154 ('A', 'B')
155 >>> adtype = np.dtype([('a', int), ('b', [('ba', int), ('bb', int)])])
156 >>> rfn.get_names_flat(adtype)
157 ('a', 'b', 'ba', 'bb')
158 """
159 listnames = []
160 names = adtype.names
161 for name in names:
162 listnames.append(name)
163 current = adtype[name]
164 if current.names is not None:
165 listnames.extend(get_names_flat(current))
166 return tuple(listnames)
167
168
169def flatten_descr(ndtype):
170 """
171 Flatten a structured data-type description.
172
173 Examples
174 --------
175 >>> import numpy as np
176 >>> from numpy.lib import recfunctions as rfn
177 >>> ndtype = np.dtype([('a', '<i4'), ('b', [('ba', '<f8'), ('bb', '<i4')])])
178 >>> rfn.flatten_descr(ndtype)
179 (('a', dtype('int32')), ('ba', dtype('float64')), ('bb', dtype('int32')))
180
181 """
182 names = ndtype.names
183 if names is None:
184 return (('', ndtype),)
185 else:
186 descr = []
187 for field in names:
188 (typ, _) = ndtype.fields[field]
189 if typ.names is not None:
190 descr.extend(flatten_descr(typ))
191 else:
192 descr.append((field, typ))
193 return tuple(descr)
194
195
196def _zip_dtype(seqarrays, flatten=False):
197 newdtype = []
198 if flatten:
199 for a in seqarrays:
200 newdtype.extend(flatten_descr(a.dtype))
201 else:
202 for a in seqarrays:
203 current = a.dtype
204 if current.names is not None and len(current.names) == 1:
205 # special case - dtypes of 1 field are flattened
206 newdtype.extend(_get_fieldspec(current))
207 else:
208 newdtype.append(('', current))
209 return np.dtype(newdtype)
210
211
212def _zip_descr(seqarrays, flatten=False):
213 """
214 Combine the dtype description of a series of arrays.
215
216 Parameters
217 ----------
218 seqarrays : sequence of arrays
219 Sequence of arrays
220 flatten : {boolean}, optional
221 Whether to collapse nested descriptions.
222 """
223 return _zip_dtype(seqarrays, flatten=flatten).descr
224
225
226def get_fieldstructure(adtype, lastname=None, parents=None,):
227 """
228 Returns a dictionary with fields indexing lists of their parent fields.
229
230 This function is used to simplify access to fields nested in other fields.
231
232 Parameters
233 ----------
234 adtype : np.dtype
235 Input datatype
236 lastname : optional
237 Last processed field name (used internally during recursion).
238 parents : dictionary
239 Dictionary of parent fields (used internally during recursion).
240
241 Examples
242 --------
243 >>> import numpy as np
244 >>> from numpy.lib import recfunctions as rfn
245 >>> ndtype = np.dtype([('A', int),
246 ... ('B', [('BA', int),
247 ... ('BB', [('BBA', int), ('BBB', int)])])])
248 >>> rfn.get_fieldstructure(ndtype)
249 ... # XXX: possible regression, order of BBA and BBB is swapped
250 {'A': [], 'B': [], 'BA': ['B'], 'BB': ['B'], 'BBA': ['B', 'BB'], 'BBB': ['B', 'BB']}
251
252 """
253 if parents is None:
254 parents = {}
255 names = adtype.names
256 for name in names:
257 current = adtype[name]
258 if current.names is not None:
259 if lastname:
260 parents[name] = [lastname, ]
261 else:
262 parents[name] = []
263 parents.update(get_fieldstructure(current, name, parents))
264 else:
265 lastparent = list(parents.get(lastname, []) or [])
266 if lastparent:
267 lastparent.append(lastname)
268 elif lastname:
269 lastparent = [lastname, ]
270 parents[name] = lastparent or []
271 return parents
272
273
274def _izip_fields_flat(iterable):
275 """
276 Returns an iterator of concatenated fields from a sequence of arrays,
277 collapsing any nested structure.
278
279 """
280 for element in iterable:
281 if isinstance(element, np.void):
282 yield from _izip_fields_flat(tuple(element))
283 else:
284 yield element
285
286
287def _izip_fields(iterable):
288 """
289 Returns an iterator of concatenated fields from a sequence of arrays.
290
291 """
292 for element in iterable:
293 if (hasattr(element, '__iter__') and
294 not isinstance(element, str)):
295 yield from _izip_fields(element)
296 elif isinstance(element, np.void) and len(tuple(element)) == 1:
297 # this statement is the same from the previous expression
298 yield from _izip_fields(element)
299 else:
300 yield element
301
302
303def _izip_records(seqarrays, fill_value=None, flatten=True):
304 """
305 Returns an iterator of concatenated items from a sequence of arrays.
306
307 Parameters
308 ----------
309 seqarrays : sequence of arrays
310 Sequence of arrays.
311 fill_value : {None, integer}
312 Value used to pad shorter iterables.
313 flatten : {True, False},
314 Whether to
315 """
316
317 # Should we flatten the items, or just use a nested approach
318 if flatten:
319 zipfunc = _izip_fields_flat
320 else:
321 zipfunc = _izip_fields
322
323 for tup in itertools.zip_longest(*seqarrays, fillvalue=fill_value):
324 yield tuple(zipfunc(tup))
325
326
327def _fix_output(output, usemask=True, asrecarray=False):
328 """
329 Private function: return a recarray, a ndarray, a MaskedArray
330 or a MaskedRecords depending on the input parameters
331 """
332 if not isinstance(output, ma.MaskedArray):
333 usemask = False
334 if usemask:
335 if asrecarray:
336 output = output.view(mrec.MaskedRecords)
337 else:
338 output = ma.filled(output)
339 if asrecarray:
340 output = output.view(np.recarray)
341 return output
342
343
344def _fix_defaults(output, defaults=None):
345 """
346 Update the fill_value and masked data of `output`
347 from the default given in a dictionary defaults.
348 """
349 names = output.dtype.names
350 (data, mask, fill_value) = (output.data, output.mask, output.fill_value)
351 for (k, v) in (defaults or {}).items():
352 if k in names:
353 fill_value[k] = v
354 data[k][mask[k]] = v
355 return output
356
357
358def _merge_arrays_dispatcher(seqarrays, fill_value=None, flatten=None,
359 usemask=None, asrecarray=None):
360 return seqarrays
361
362
363@array_function_dispatch(_merge_arrays_dispatcher)
364def merge_arrays(seqarrays, fill_value=-1, flatten=False,
365 usemask=False, asrecarray=False):
366 """
367 Merge arrays field by field.
368
369 Parameters
370 ----------
371 seqarrays : sequence of ndarrays
372 Sequence of arrays
373 fill_value : {float}, optional
374 Filling value used to pad missing data on the shorter arrays.
375 flatten : {False, True}, optional
376 Whether to collapse nested fields.
377 usemask : {False, True}, optional
378 Whether to return a masked array or not.
379 asrecarray : {False, True}, optional
380 Whether to return a recarray (MaskedRecords) or not.
381
382 Examples
383 --------
384 >>> import numpy as np
385 >>> from numpy.lib import recfunctions as rfn
386 >>> rfn.merge_arrays((np.array([1, 2]), np.array([10., 20., 30.])))
387 array([( 1, 10.), ( 2, 20.), (-1, 30.)],
388 dtype=[('f0', '<i8'), ('f1', '<f8')])
389
390 >>> rfn.merge_arrays((np.array([1, 2], dtype=np.int64),
391 ... np.array([10., 20., 30.])), usemask=False)
392 array([(1, 10.0), (2, 20.0), (-1, 30.0)],
393 dtype=[('f0', '<i8'), ('f1', '<f8')])
394 >>> rfn.merge_arrays((np.array([1, 2]).view([('a', np.int64)]),
395 ... np.array([10., 20., 30.])),
396 ... usemask=False, asrecarray=True)
397 rec.array([( 1, 10.), ( 2, 20.), (-1, 30.)],
398 dtype=[('a', '<i8'), ('f1', '<f8')])
399
400 Notes
401 -----
402 * Without a mask, the missing value will be filled with something,
403 depending on what its corresponding type:
404
405 * ``-1`` for integers
406 * ``-1.0`` for floating point numbers
407 * ``'-'`` for characters
408 * ``'-1'`` for strings
409 * ``True`` for boolean values
410 * XXX: I just obtained these values empirically
411 """
412 # Only one item in the input sequence ?
413 if (len(seqarrays) == 1):
414 seqarrays = np.asanyarray(seqarrays[0])
415 # Do we have a single ndarray as input ?
416 if isinstance(seqarrays, (np.ndarray, np.void)):
417 seqdtype = seqarrays.dtype
418 # Make sure we have named fields
419 if seqdtype.names is None:
420 seqdtype = np.dtype([('', seqdtype)])
421 if not flatten or _zip_dtype((seqarrays,), flatten=True) == seqdtype:
422 # Minimal processing needed: just make sure everything's a-ok
423 seqarrays = seqarrays.ravel()
424 # Find what type of array we must return
425 if usemask:
426 if asrecarray:
427 seqtype = mrec.MaskedRecords
428 else:
429 seqtype = ma.MaskedArray
430 elif asrecarray:
431 seqtype = np.recarray
432 else:
433 seqtype = np.ndarray
434 return seqarrays.view(dtype=seqdtype, type=seqtype)
435 else:
436 seqarrays = (seqarrays,)
437 else:
438 # Make sure we have arrays in the input sequence
439 seqarrays = [np.asanyarray(_m) for _m in seqarrays]
440 # Find the sizes of the inputs and their maximum
441 sizes = tuple(a.size for a in seqarrays)
442 maxlength = max(sizes)
443 # Get the dtype of the output (flattening if needed)
444 newdtype = _zip_dtype(seqarrays, flatten=flatten)
445 # Initialize the sequences for data and mask
446 seqdata = []
447 seqmask = []
448 # If we expect some kind of MaskedArray, make a special loop.
449 if usemask:
450 for (a, n) in zip(seqarrays, sizes):
451 nbmissing = (maxlength - n)
452 # Get the data and mask
453 data = a.ravel().__array__()
454 mask = ma.getmaskarray(a).ravel()
455 # Get the filling value (if needed)
456 if nbmissing:
457 fval = mrec._check_fill_value(fill_value, a.dtype)
458 if isinstance(fval, (np.ndarray, np.void)):
459 if len(fval.dtype) == 1:
460 fval = fval.item()[0]
461 fmsk = True
462 else:
463 fval = np.array(fval, dtype=a.dtype, ndmin=1)
464 fmsk = np.ones((1,), dtype=mask.dtype)
465 else:
466 fval = None
467 fmsk = True
468 # Store an iterator padding the input to the expected length
469 seqdata.append(itertools.chain(data, [fval] * nbmissing))
470 seqmask.append(itertools.chain(mask, [fmsk] * nbmissing))
471 # Create an iterator for the data
472 data = tuple(_izip_records(seqdata, flatten=flatten))
473 output = ma.array(np.fromiter(data, dtype=newdtype, count=maxlength),
474 mask=list(_izip_records(seqmask, flatten=flatten)))
475 if asrecarray:
476 output = output.view(mrec.MaskedRecords)
477 else:
478 # Same as before, without the mask we don't need...
479 for (a, n) in zip(seqarrays, sizes):
480 nbmissing = (maxlength - n)
481 data = a.ravel().__array__()
482 if nbmissing:
483 fval = mrec._check_fill_value(fill_value, a.dtype)
484 if isinstance(fval, (np.ndarray, np.void)):
485 if len(fval.dtype) == 1:
486 fval = fval.item()[0]
487 else:
488 fval = np.array(fval, dtype=a.dtype, ndmin=1)
489 else:
490 fval = None
491 seqdata.append(itertools.chain(data, [fval] * nbmissing))
492 output = np.fromiter(tuple(_izip_records(seqdata, flatten=flatten)),
493 dtype=newdtype, count=maxlength)
494 if asrecarray:
495 output = output.view(np.recarray)
496 # And we're done...
497 return output
498
499
500def _drop_fields_dispatcher(base, drop_names, usemask=None, asrecarray=None):
501 return (base,)
502
503
504@array_function_dispatch(_drop_fields_dispatcher)
505def drop_fields(base, drop_names, usemask=True, asrecarray=False):
506 """
507 Return a new array with fields in `drop_names` dropped.
508
509 Nested fields are supported.
510
511 Parameters
512 ----------
513 base : array
514 Input array
515 drop_names : string or sequence
516 String or sequence of strings corresponding to the names of the
517 fields to drop.
518 usemask : {False, True}, optional
519 Whether to return a masked array or not.
520 asrecarray : string or sequence, optional
521 Whether to return a recarray or a mrecarray (`asrecarray=True`) or
522 a plain ndarray or masked array with flexible dtype. The default
523 is False.
524
525 Examples
526 --------
527 >>> import numpy as np
528 >>> from numpy.lib import recfunctions as rfn
529 >>> a = np.array([(1, (2, 3.0)), (4, (5, 6.0))],
530 ... dtype=[('a', np.int64), ('b', [('ba', np.double), ('bb', np.int64)])])
531 >>> rfn.drop_fields(a, 'a')
532 array([((2., 3),), ((5., 6),)],
533 dtype=[('b', [('ba', '<f8'), ('bb', '<i8')])])
534 >>> rfn.drop_fields(a, 'ba')
535 array([(1, (3,)), (4, (6,))], dtype=[('a', '<i8'), ('b', [('bb', '<i8')])])
536 >>> rfn.drop_fields(a, ['ba', 'bb'])
537 array([(1,), (4,)], dtype=[('a', '<i8')])
538 """
539 if _is_string_like(drop_names):
540 drop_names = [drop_names]
541 else:
542 drop_names = set(drop_names)
543
544 def _drop_descr(ndtype, drop_names):
545 names = ndtype.names
546 newdtype = []
547 for name in names:
548 current = ndtype[name]
549 if name in drop_names:
550 continue
551 if current.names is not None:
552 descr = _drop_descr(current, drop_names)
553 if descr:
554 newdtype.append((name, descr))
555 else:
556 newdtype.append((name, current))
557 return newdtype
558
559 newdtype = _drop_descr(base.dtype, drop_names)
560
561 output = np.empty(base.shape, dtype=newdtype)
562 output = recursive_fill_fields(base, output)
563 return _fix_output(output, usemask=usemask, asrecarray=asrecarray)
564
565
566def _keep_fields(base, keep_names, usemask=True, asrecarray=False):
567 """
568 Return a new array keeping only the fields in `keep_names`,
569 and preserving the order of those fields.
570
571 Parameters
572 ----------
573 base : array
574 Input array
575 keep_names : string or sequence
576 String or sequence of strings corresponding to the names of the
577 fields to keep. Order of the names will be preserved.
578 usemask : {False, True}, optional
579 Whether to return a masked array or not.
580 asrecarray : string or sequence, optional
581 Whether to return a recarray or a mrecarray (`asrecarray=True`) or
582 a plain ndarray or masked array with flexible dtype. The default
583 is False.
584 """
585 newdtype = [(n, base.dtype[n]) for n in keep_names]
586 output = np.empty(base.shape, dtype=newdtype)
587 output = recursive_fill_fields(base, output)
588 return _fix_output(output, usemask=usemask, asrecarray=asrecarray)
589
590
591def _rec_drop_fields_dispatcher(base, drop_names):
592 return (base,)
593
594
595@array_function_dispatch(_rec_drop_fields_dispatcher)
596def rec_drop_fields(base, drop_names):
597 """
598 Returns a new numpy.recarray with fields in `drop_names` dropped.
599 """
600 return drop_fields(base, drop_names, usemask=False, asrecarray=True)
601
602
603def _rename_fields_dispatcher(base, namemapper):
604 return (base,)
605
606
607@array_function_dispatch(_rename_fields_dispatcher)
608def rename_fields(base, namemapper):
609 """
610 Rename the fields from a flexible-datatype ndarray or recarray.
611
612 Nested fields are supported.
613
614 Parameters
615 ----------
616 base : ndarray
617 Input array whose fields must be modified.
618 namemapper : dictionary
619 Dictionary mapping old field names to their new version.
620
621 Examples
622 --------
623 >>> import numpy as np
624 >>> from numpy.lib import recfunctions as rfn
625 >>> a = np.array([(1, (2, [3.0, 30.])), (4, (5, [6.0, 60.]))],
626 ... dtype=[('a', int),('b', [('ba', float), ('bb', (float, 2))])])
627 >>> rfn.rename_fields(a, {'a':'A', 'bb':'BB'})
628 array([(1, (2., [ 3., 30.])), (4, (5., [ 6., 60.]))],
629 dtype=[('A', '<i8'), ('b', [('ba', '<f8'), ('BB', '<f8', (2,))])])
630
631 """
632 def _recursive_rename_fields(ndtype, namemapper):
633 newdtype = []
634 for name in ndtype.names:
635 newname = namemapper.get(name, name)
636 current = ndtype[name]
637 if current.names is not None:
638 newdtype.append(
639 (newname, _recursive_rename_fields(current, namemapper))
640 )
641 else:
642 newdtype.append((newname, current))
643 return newdtype
644 newdtype = _recursive_rename_fields(base.dtype, namemapper)
645 return base.view(newdtype)
646
647
648def _append_fields_dispatcher(base, names, data, dtypes=None,
649 fill_value=None, usemask=None, asrecarray=None):
650 yield base
651 yield from data
652
653
654@array_function_dispatch(_append_fields_dispatcher)
655def append_fields(base, names, data, dtypes=None,
656 fill_value=-1, usemask=True, asrecarray=False):
657 """
658 Add new fields to an existing array.
659
660 The names of the fields are given with the `names` arguments,
661 the corresponding values with the `data` arguments.
662 If a single field is appended, `names`, `data` and `dtypes` do not have
663 to be lists but just values.
664
665 Parameters
666 ----------
667 base : array
668 Input array to extend.
669 names : string, sequence
670 String or sequence of strings corresponding to the names
671 of the new fields.
672 data : array or sequence of arrays
673 Array or sequence of arrays storing the fields to add to the base.
674 dtypes : sequence of datatypes, optional
675 Datatype or sequence of datatypes.
676 If None, the datatypes are estimated from the `data`.
677 fill_value : {float}, optional
678 Filling value used to pad missing data on the shorter arrays.
679 usemask : {False, True}, optional
680 Whether to return a masked array or not.
681 asrecarray : {False, True}, optional
682 Whether to return a recarray (MaskedRecords) or not.
683
684 """
685 # Check the names
686 if isinstance(names, (tuple, list)):
687 if len(names) != len(data):
688 msg = "The number of arrays does not match the number of names"
689 raise ValueError(msg)
690 elif isinstance(names, str):
691 names = [names, ]
692 data = [data, ]
693 #
694 if dtypes is None:
695 data = [np.array(a, copy=None, subok=True) for a in data]
696 data = [a.view([(name, a.dtype)]) for (name, a) in zip(names, data)]
697 else:
698 if not isinstance(dtypes, (tuple, list)):
699 dtypes = [dtypes, ]
700 if len(data) != len(dtypes):
701 if len(dtypes) == 1:
702 dtypes = dtypes * len(data)
703 else:
704 msg = "The dtypes argument must be None, a dtype, or a list."
705 raise ValueError(msg)
706 data = [np.array(a, copy=None, subok=True, dtype=d).view([(n, d)])
707 for (a, n, d) in zip(data, names, dtypes)]
708 #
709 base = merge_arrays(base, usemask=usemask, fill_value=fill_value)
710 if len(data) > 1:
711 data = merge_arrays(data, flatten=True, usemask=usemask,
712 fill_value=fill_value)
713 else:
714 data = data.pop()
715 #
716 output = ma.masked_all(
717 max(len(base), len(data)),
718 dtype=_get_fieldspec(base.dtype) + _get_fieldspec(data.dtype))
719 output = recursive_fill_fields(base, output)
720 output = recursive_fill_fields(data, output)
721 #
722 return _fix_output(output, usemask=usemask, asrecarray=asrecarray)
723
724
725def _rec_append_fields_dispatcher(base, names, data, dtypes=None):
726 yield base
727 yield from data
728
729
730@array_function_dispatch(_rec_append_fields_dispatcher)
731def rec_append_fields(base, names, data, dtypes=None):
732 """
733 Add new fields to an existing array.
734
735 The names of the fields are given with the `names` arguments,
736 the corresponding values with the `data` arguments.
737 If a single field is appended, `names`, `data` and `dtypes` do not have
738 to be lists but just values.
739
740 Parameters
741 ----------
742 base : array
743 Input array to extend.
744 names : string, sequence
745 String or sequence of strings corresponding to the names
746 of the new fields.
747 data : array or sequence of arrays
748 Array or sequence of arrays storing the fields to add to the base.
749 dtypes : sequence of datatypes, optional
750 Datatype or sequence of datatypes.
751 If None, the datatypes are estimated from the `data`.
752
753 See Also
754 --------
755 append_fields
756
757 Returns
758 -------
759 appended_array : np.recarray
760 """
761 return append_fields(base, names, data=data, dtypes=dtypes,
762 asrecarray=True, usemask=False)
763
764
765def _repack_fields_dispatcher(a, align=None, recurse=None):
766 return (a,)
767
768
769@array_function_dispatch(_repack_fields_dispatcher)
770def repack_fields(a, align=False, recurse=False):
771 """
772 Re-pack the fields of a structured array or dtype in memory.
773
774 The memory layout of structured datatypes allows fields at arbitrary
775 byte offsets. This means the fields can be separated by padding bytes,
776 their offsets can be non-monotonically increasing, and they can overlap.
777
778 This method removes any overlaps and reorders the fields in memory so they
779 have increasing byte offsets, and adds or removes padding bytes depending
780 on the `align` option, which behaves like the `align` option to
781 `numpy.dtype`.
782
783 If `align=False`, this method produces a "packed" memory layout in which
784 each field starts at the byte the previous field ended, and any padding
785 bytes are removed.
786
787 If `align=True`, this methods produces an "aligned" memory layout in which
788 each field's offset is a multiple of its alignment, and the total itemsize
789 is a multiple of the largest alignment, by adding padding bytes as needed.
790
791 Parameters
792 ----------
793 a : ndarray or dtype
794 array or dtype for which to repack the fields.
795 align : boolean
796 If true, use an "aligned" memory layout, otherwise use a "packed" layout.
797 recurse : boolean
798 If True, also repack nested structures.
799
800 Returns
801 -------
802 repacked : ndarray or dtype
803 Copy of `a` with fields repacked, or `a` itself if no repacking was
804 needed.
805
806 Examples
807 --------
808 >>> import numpy as np
809
810 >>> from numpy.lib import recfunctions as rfn
811 >>> def print_offsets(d):
812 ... print("offsets:", [d.fields[name][1] for name in d.names])
813 ... print("itemsize:", d.itemsize)
814 ...
815 >>> dt = np.dtype('u1, <i8, <f8', align=True)
816 >>> dt
817 dtype({'names': ['f0', 'f1', 'f2'], 'formats': ['u1', '<i8', '<f8'], \
818'offsets': [0, 8, 16], 'itemsize': 24}, align=True)
819 >>> print_offsets(dt)
820 offsets: [0, 8, 16]
821 itemsize: 24
822 >>> packed_dt = rfn.repack_fields(dt)
823 >>> packed_dt
824 dtype([('f0', 'u1'), ('f1', '<i8'), ('f2', '<f8')])
825 >>> print_offsets(packed_dt)
826 offsets: [0, 1, 9]
827 itemsize: 17
828
829 """
830 if not isinstance(a, np.dtype):
831 dt = repack_fields(a.dtype, align=align, recurse=recurse)
832 return a.astype(dt, copy=False)
833
834 if a.names is None:
835 return a
836
837 fieldinfo = []
838 for name in a.names:
839 tup = a.fields[name]
840 if recurse:
841 fmt = repack_fields(tup[0], align=align, recurse=True)
842 else:
843 fmt = tup[0]
844
845 if len(tup) == 3:
846 name = (tup[2], name)
847
848 fieldinfo.append((name, fmt))
849
850 dt = np.dtype(fieldinfo, align=align)
851 return np.dtype((a.type, dt))
852
853def _get_fields_and_offsets(dt, offset=0):
854 """
855 Returns a flat list of (dtype, count, offset) tuples of all the
856 scalar fields in the dtype "dt", including nested fields, in left
857 to right order.
858 """
859
860 # counts up elements in subarrays, including nested subarrays, and returns
861 # base dtype and count
862 def count_elem(dt):
863 count = 1
864 while dt.shape != ():
865 for size in dt.shape:
866 count *= size
867 dt = dt.base
868 return dt, count
869
870 fields = []
871 for name in dt.names:
872 field = dt.fields[name]
873 f_dt, f_offset = field[0], field[1]
874 f_dt, n = count_elem(f_dt)
875
876 if f_dt.names is None:
877 fields.append((np.dtype((f_dt, (n,))), n, f_offset + offset))
878 else:
879 subfields = _get_fields_and_offsets(f_dt, f_offset + offset)
880 size = f_dt.itemsize
881
882 for i in range(n):
883 if i == 0:
884 # optimization: avoid list comprehension if no subarray
885 fields.extend(subfields)
886 else:
887 fields.extend([(d, c, o + i * size) for d, c, o in subfields])
888 return fields
889
890def _common_stride(offsets, counts, itemsize):
891 """
892 Returns the stride between the fields, or None if the stride is not
893 constant. The values in "counts" designate the lengths of
894 subarrays. Subarrays are treated as many contiguous fields, with
895 always positive stride.
896 """
897 if len(offsets) <= 1:
898 return itemsize
899
900 negative = offsets[1] < offsets[0] # negative stride
901 if negative:
902 # reverse, so offsets will be ascending
903 it = zip(reversed(offsets), reversed(counts))
904 else:
905 it = zip(offsets, counts)
906
907 prev_offset = None
908 stride = None
909 for offset, count in it:
910 if count != 1: # subarray: always c-contiguous
911 if negative:
912 return None # subarrays can never have a negative stride
913 if stride is None:
914 stride = itemsize
915 if stride != itemsize:
916 return None
917 end_offset = offset + (count - 1) * itemsize
918 else:
919 end_offset = offset
920
921 if prev_offset is not None:
922 new_stride = offset - prev_offset
923 if stride is None:
924 stride = new_stride
925 if stride != new_stride:
926 return None
927
928 prev_offset = end_offset
929
930 if negative:
931 return -stride
932 return stride
933
934
935def _structured_to_unstructured_dispatcher(arr, dtype=None, copy=None,
936 casting=None):
937 return (arr,)
938
939@array_function_dispatch(_structured_to_unstructured_dispatcher)
940def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'):
941 """
942 Converts an n-D structured array into an (n+1)-D unstructured array.
943
944 The new array will have a new last dimension equal in size to the
945 number of field-elements of the input array. If not supplied, the output
946 datatype is determined from the numpy type promotion rules applied to all
947 the field datatypes.
948
949 Nested fields, as well as each element of any subarray fields, all count
950 as a single field-elements.
951
952 Parameters
953 ----------
954 arr : ndarray
955 Structured array or dtype to convert. Cannot contain object datatype.
956 dtype : dtype, optional
957 The dtype of the output unstructured array.
958 copy : bool, optional
959 If true, always return a copy. If false, a view is returned if
960 possible, such as when the `dtype` and strides of the fields are
961 suitable and the array subtype is one of `numpy.ndarray`,
962 `numpy.recarray` or `numpy.memmap`.
963
964 .. versionchanged:: 1.25.0
965 A view can now be returned if the fields are separated by a
966 uniform stride.
967
968 casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
969 See casting argument of `numpy.ndarray.astype`. Controls what kind of
970 data casting may occur.
971
972 Returns
973 -------
974 unstructured : ndarray
975 Unstructured array with one more dimension.
976
977 Examples
978 --------
979 >>> import numpy as np
980
981 >>> from numpy.lib import recfunctions as rfn
982 >>> a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
983 >>> a
984 array([(0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.]),
985 (0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.])],
986 dtype=[('a', '<i4'), ('b', [('f0', '<f4'), ('f1', '<u2')]), ('c', '<f4', (2,))])
987 >>> rfn.structured_to_unstructured(a)
988 array([[0., 0., 0., 0., 0.],
989 [0., 0., 0., 0., 0.],
990 [0., 0., 0., 0., 0.],
991 [0., 0., 0., 0., 0.]])
992
993 >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
994 ... dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
995 >>> np.mean(rfn.structured_to_unstructured(b[['x', 'z']]), axis=-1)
996 array([ 3. , 5.5, 9. , 11. ])
997
998 """ # noqa: E501
999 if arr.dtype.names is None:
1000 raise ValueError('arr must be a structured array')
1001
1002 fields = _get_fields_and_offsets(arr.dtype)
1003 n_fields = len(fields)
1004 if n_fields == 0 and dtype is None:
1005 raise ValueError("arr has no fields. Unable to guess dtype")
1006 elif n_fields == 0:
1007 # too many bugs elsewhere for this to work now
1008 raise NotImplementedError("arr with no fields is not supported")
1009
1010 dts, counts, offsets = zip(*fields)
1011 names = [f'f{n}' for n in range(n_fields)]
1012
1013 if dtype is None:
1014 out_dtype = np.result_type(*[dt.base for dt in dts])
1015 else:
1016 out_dtype = np.dtype(dtype)
1017
1018 # Use a series of views and casts to convert to an unstructured array:
1019
1020 # first view using flattened fields (doesn't work for object arrays)
1021 # Note: dts may include a shape for subarrays
1022 flattened_fields = np.dtype({'names': names,
1023 'formats': dts,
1024 'offsets': offsets,
1025 'itemsize': arr.dtype.itemsize})
1026 arr = arr.view(flattened_fields)
1027
1028 # we only allow a few types to be unstructured by manipulating the
1029 # strides, because we know it won't work with, for example, np.matrix nor
1030 # np.ma.MaskedArray.
1031 can_view = type(arr) in (np.ndarray, np.recarray, np.memmap)
1032 if (not copy) and can_view and all(dt.base == out_dtype for dt in dts):
1033 # all elements have the right dtype already; if they have a common
1034 # stride, we can just return a view
1035 common_stride = _common_stride(offsets, counts, out_dtype.itemsize)
1036 if common_stride is not None:
1037 wrap = arr.__array_wrap__
1038
1039 new_shape = arr.shape + (sum(counts), out_dtype.itemsize)
1040 new_strides = arr.strides + (abs(common_stride), 1)
1041
1042 arr = arr[..., np.newaxis].view(np.uint8) # view as bytes
1043 arr = arr[..., min(offsets):] # remove the leading unused data
1044 arr = np.lib.stride_tricks.as_strided(arr,
1045 new_shape,
1046 new_strides,
1047 subok=True)
1048
1049 # cast and drop the last dimension again
1050 arr = arr.view(out_dtype)[..., 0]
1051
1052 if common_stride < 0:
1053 arr = arr[..., ::-1] # reverse, if the stride was negative
1054 if type(arr) is not type(wrap.__self__):
1055 # Some types (e.g. recarray) turn into an ndarray along the
1056 # way, so we have to wrap it again in order to match the
1057 # behavior with copy=True.
1058 arr = wrap(arr)
1059 return arr
1060
1061 # next cast to a packed format with all fields converted to new dtype
1062 packed_fields = np.dtype({'names': names,
1063 'formats': [(out_dtype, dt.shape) for dt in dts]})
1064 arr = arr.astype(packed_fields, copy=copy, casting=casting)
1065
1066 # finally is it safe to view the packed fields as the unstructured type
1067 return arr.view((out_dtype, (sum(counts),)))
1068
1069
1070def _unstructured_to_structured_dispatcher(arr, dtype=None, names=None,
1071 align=None, copy=None, casting=None):
1072 return (arr,)
1073
1074@array_function_dispatch(_unstructured_to_structured_dispatcher)
1075def unstructured_to_structured(arr, dtype=None, names=None, align=False,
1076 copy=False, casting='unsafe'):
1077 """
1078 Converts an n-D unstructured array into an (n-1)-D structured array.
1079
1080 The last dimension of the input array is converted into a structure, with
1081 number of field-elements equal to the size of the last dimension of the
1082 input array. By default all output fields have the input array's dtype, but
1083 an output structured dtype with an equal number of fields-elements can be
1084 supplied instead.
1085
1086 Nested fields, as well as each element of any subarray fields, all count
1087 towards the number of field-elements.
1088
1089 Parameters
1090 ----------
1091 arr : ndarray
1092 Unstructured array or dtype to convert.
1093 dtype : dtype, optional
1094 The structured dtype of the output array
1095 names : list of strings, optional
1096 If dtype is not supplied, this specifies the field names for the output
1097 dtype, in order. The field dtypes will be the same as the input array.
1098 align : boolean, optional
1099 Whether to create an aligned memory layout.
1100 copy : bool, optional
1101 See copy argument to `numpy.ndarray.astype`. If true, always return a
1102 copy. If false, and `dtype` requirements are satisfied, a view is
1103 returned.
1104 casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
1105 See casting argument of `numpy.ndarray.astype`. Controls what kind of
1106 data casting may occur.
1107
1108 Returns
1109 -------
1110 structured : ndarray
1111 Structured array with fewer dimensions.
1112
1113 Examples
1114 --------
1115 >>> import numpy as np
1116
1117 >>> from numpy.lib import recfunctions as rfn
1118 >>> dt = np.dtype([('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
1119 >>> a = np.arange(20).reshape((4,5))
1120 >>> a
1121 array([[ 0, 1, 2, 3, 4],
1122 [ 5, 6, 7, 8, 9],
1123 [10, 11, 12, 13, 14],
1124 [15, 16, 17, 18, 19]])
1125 >>> rfn.unstructured_to_structured(a, dt)
1126 array([( 0, ( 1., 2), [ 3., 4.]), ( 5, ( 6., 7), [ 8., 9.]),
1127 (10, (11., 12), [13., 14.]), (15, (16., 17), [18., 19.])],
1128 dtype=[('a', '<i4'), ('b', [('f0', '<f4'), ('f1', '<u2')]), ('c', '<f4', (2,))])
1129
1130 """ # noqa: E501
1131 if arr.shape == ():
1132 raise ValueError('arr must have at least one dimension')
1133 n_elem = arr.shape[-1]
1134 if n_elem == 0:
1135 # too many bugs elsewhere for this to work now
1136 raise NotImplementedError("last axis with size 0 is not supported")
1137
1138 if dtype is None:
1139 if names is None:
1140 names = [f'f{n}' for n in range(n_elem)]
1141 out_dtype = np.dtype([(n, arr.dtype) for n in names], align=align)
1142 fields = _get_fields_and_offsets(out_dtype)
1143 dts, counts, offsets = zip(*fields)
1144 else:
1145 if names is not None:
1146 raise ValueError("don't supply both dtype and names")
1147 # if dtype is the args of np.dtype, construct it
1148 dtype = np.dtype(dtype)
1149 # sanity check of the input dtype
1150 fields = _get_fields_and_offsets(dtype)
1151 if len(fields) == 0:
1152 dts, counts, offsets = [], [], []
1153 else:
1154 dts, counts, offsets = zip(*fields)
1155
1156 if n_elem != sum(counts):
1157 raise ValueError('The length of the last dimension of arr must '
1158 'be equal to the number of fields in dtype')
1159 out_dtype = dtype
1160 if align and not out_dtype.isalignedstruct:
1161 raise ValueError("align was True but dtype is not aligned")
1162
1163 names = [f'f{n}' for n in range(len(fields))]
1164
1165 # Use a series of views and casts to convert to a structured array:
1166
1167 # first view as a packed structured array of one dtype
1168 packed_fields = np.dtype({'names': names,
1169 'formats': [(arr.dtype, dt.shape) for dt in dts]})
1170 arr = np.ascontiguousarray(arr).view(packed_fields)
1171
1172 # next cast to an unpacked but flattened format with varied dtypes
1173 flattened_fields = np.dtype({'names': names,
1174 'formats': dts,
1175 'offsets': offsets,
1176 'itemsize': out_dtype.itemsize})
1177 arr = arr.astype(flattened_fields, copy=copy, casting=casting)
1178
1179 # finally view as the final nested dtype and remove the last axis
1180 return arr.view(out_dtype)[..., 0]
1181
1182def _apply_along_fields_dispatcher(func, arr):
1183 return (arr,)
1184
1185@array_function_dispatch(_apply_along_fields_dispatcher)
1186def apply_along_fields(func, arr):
1187 """
1188 Apply function 'func' as a reduction across fields of a structured array.
1189
1190 This is similar to `numpy.apply_along_axis`, but treats the fields of a
1191 structured array as an extra axis. The fields are all first cast to a
1192 common type following the type-promotion rules from `numpy.result_type`
1193 applied to the field's dtypes.
1194
1195 Parameters
1196 ----------
1197 func : function
1198 Function to apply on the "field" dimension. This function must
1199 support an `axis` argument, like `numpy.mean`, `numpy.sum`, etc.
1200 arr : ndarray