如何对Numpy结构化数组执行按列操作？

问题代码 #

import numpy as np

a_type = np.dtype([("x", int), ("y", float)])
a_list = []

for i in range(0, 8, 2):
    entry = np.zeros((1,), dtype=a_type)
    entry["x"][0] = i
    entry["y"][0] = i + 1.0
    a_list.append(entry)
a_array = np.array(a_list, dtype=a_type)
a_array_flat = a_array.reshape(-1)
print(a_array_flat["x"])
print(np.sum(a_array_flat["x"]))

运行报错 #

[0 2 4 6]
Traceback (most recent call last):
  File "/home/andreas/src/masiri/booking_algorythm/demo_structured_aarray_flatten.py", line 14, in <module>
    print(np.sum(a_array_flat["x"]))
  File "<array_function__ internals>", line 180, in sum
  File "/home/andreas/src/masiri/venv/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 2298, in sum
    return _wrapreduction(a, np.add, 'sum', axis, dtype, out, keepdims=keepdims,
  File "/home/andreas/src/masiri/venv/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 86, in _wrapreduction
    return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
numpy.core._exceptions._UFuncNoLoopError: ufunc 'add' did not contain a loop with signature matching types (dtype({'names': ['x'], 'formats': ['<i8'], 'offsets': [0], 'itemsize': 16}), dtype({'names': ['x'], 'formats': ['<i8'], 'offsets': [0], 'itemsize': 16})) -> None

问题原因 #

循环创建的每个entry是形状为(1,)的结构化数组，将这些数组存入列表后用np.array转换，得到的是二维结构化数组（形状(4,1)）。即便用reshape(-1)展平，取字段"x"后得到的依然是嵌套的单字段结构化数组，而非普通整数数组，Numpy的sum无法直接对这种嵌套结构执行加法操作。

解决方案 #

方案1：直接创建一维结构化数组（推荐）

跳过列表拼接的步骤，直接初始化目标长度的一维数组，避免嵌套结构：

import numpy as np

a_type = np.dtype([("x", int), ("y", float)])
# 直接创建长度为4的一维结构化数组
a_array = np.zeros(4, dtype=a_type)

for i in range(0, 8, 2):
    idx = i // 2
    a_array["x"][idx] = i
    a_array["y"][idx] = i + 1.0

print(a_array["x"])       # 输出 [0 2 4 6]
print(np.sum(a_array["x"]))  # 输出 12

方案2：转换嵌套结构为基础类型数组

如果必须保留列表拼接的流程，展平后用view()将单字段结构化数组转换为对应基础类型：

import numpy as np

a_type = np.dtype([("x", int), ("y", float)])
a_list = []

for i in range(0, 8, 2):
    entry = np.zeros((1,), dtype=a_type)
    entry["x"][0] = i
    entry["y"][0] = i + 1.0
    a_list.append(entry)

a_array = np.array(a_list, dtype=a_type)
a_array_flat = a_array.reshape(-1)
# 将单字段结构化数组转换为普通整数数组
x_array = a_array_flat["x"].view(np.int64)

print(x_array)            # 输出 [0 2 4 6]
print(np.sum(x_array))    # 输出 12

说明 #

两种方案都能让你正常执行按字段的求和操作，同时保留结构化数组对特殊类型（timedelta64、datetime64）的支持，满足按列快速操作的需求。

内容的提问来源于stack exchange，提问作者Andreas Schuldei