Numpy: 合并结构化数组？

Question

输入

我有很多 numpy 结构化数组，它们被放在一个列表里，像这个例子：

import numpy

a1 = numpy.array([(1, 2), (3, 4), (5, 6)], dtype=[('x', int), ('y', int)])

a2 = numpy.array([(7,10), (8,11), (9,12)], dtype=[('z', int), ('w', float)])

arrays = [a1, a2]

期望的输出

我想知道，怎样才能把它们都合并在一起，创建一个统一的结构化数组，像下面这样？

desired_result = numpy.array([(1, 2, 7, 10), (3, 4, 8, 11), (5, 6, 9, 12)],
                             dtype=[('x', int), ('y', int), ('z', int), ('w', float)])

当前的方法

这是我现在使用的方法，但速度非常慢，所以我怀疑一定有更高效的做法。

from numpy.lib.recfunctions import append_fields

def join_struct_arrays(arrays):
    for array in arrays:
        try:
            result = append_fields(result, array.dtype.names, [array[name] for name in array.dtype.names], usemask=False)
        except NameError:
            result = array

    return result

Answer 1

还有另外一种方法，这种写法更容易读懂，而且我觉得速度也快很多：

def join_struct_arrays(arrays):
    newdtype = []
    for a in arrays:
        descr = []
        for field in a.dtype.names:
            (typ, _) = a.dtype.fields[field]
            descr.append((field, typ))
        newdtype.extend(tuple(descr))
    newrecarray = np.zeros(len(arrays[0]), dtype = newdtype)
    for a in arrays:
        for name in a.dtype.names:
            newrecarray[name] = a[name]
    return newrecarray

编辑：根据Sven的建议，变成了这样（虽然稍微慢一点，但其实挺容易理解的）：

def join_struct_arrays2(arrays):
    newdtype = sum((a.dtype.descr for a in arrays), [])
    newrecarray = np.empty(len(arrays[0]), dtype = newdtype)
    for a in arrays:
        for name in a.dtype.names:
            newrecarray[name] = a[name]
    return newrecarray

Answer 2

你还可以使用 numpy.lib.recfunctions 里的 merge_arrays 函数：

import numpy.lib.recfunctions as rfn
rfn.merge_arrays(arrays, flatten = True, usemask = False)

Out[52]: 
array([(1, 2, 7, 10.0), (3, 4, 8, 11.0), (5, 6, 9, 12.0)], 
     dtype=[('x', '<i4'), ('y', '<i4'), ('z', '<i4'), ('w', '<f8')])

Answer 3

这里有一个实现方式，应该会更快。它把所有东西都转换成了 numpy.uint8 的数组，并且没有使用任何临时变量。

def join_struct_arrays(arrays):
    sizes = numpy.array([a.itemsize for a in arrays])
    offsets = numpy.r_[0, sizes.cumsum()]
    n = len(arrays[0])
    joint = numpy.empty((n, offsets[-1]), dtype=numpy.uint8)
    for a, size, offset in zip(arrays, sizes, offsets):
        joint[:,offset:offset+size] = a.view(numpy.uint8).reshape(n,size)
    dtype = sum((a.dtype.descr for a in arrays), [])
    return joint.ravel().view(dtype)

编辑: 简化了代码，避免了不必要的 as_strided()。