如何对NumPy结构化数组应用旋转变换？

问题场景 #

使用numpy结构化数组存储LAMMPS原子数据时，尝试直接对位置字段做旋转变换出现类型不匹配错误：

示例代码：

import numpy as np

transform = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
dtype = np.dtype([("x", np.float64), ("y", np.float64), ("z", np.float64)])

atoms = np.array(
    [
        (0.0, 0.0, 0.0),
        (1.0, 0.0, 0.0),
        (0.0, 1.0, 0.0),
        (1.0, 1.0, 1.0),
    ],
    dtype=dtype,
)

atoms[["x", "y", "z"]] = atoms[["x", "y", "z"]] @ transform.T

报错信息：

Traceback (most recent call last):
  File "c:\Users\acgc99\Desktop\rotation.py", line 16, in <module>
    atoms[["x", "y", "z"]] = atoms[["x", "y", "z"]] @ transform.T
                             ~~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~
numpy._core._exceptions._UFuncNoLoopError: ufunc 'matmul' did not contain a loop with signature matching types (dtype([('x', '<f8'), ('y', '<f8'), ('z', '<f8')]), dtype('float64')) -> None

核心原因 #

numpy的@（矩阵乘法）不支持复合dtype的结构化数组与普通数值数组直接运算，因为结构化数组的类型无法被通用函数（ufunc）识别处理。

高效解决方案（无数据复制） #

利用numpy的**视图（view）**功能，将结构化数组的位置字段转为连续的普通数值数组视图，无需复制数据，操作完成后直接同步回原结构化数组，完全适配千万级原子的高效处理：

import numpy as np

transform = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64)
dtype = np.dtype([("x", np.float64), ("y", np.float64), ("z", np.float64)])

atoms = np.array(
    [
        (0.0, 0.0, 0.0),
        (1.0, 0.0, 0.0),
        (0.0, 1.0, 0.0),
        (1.0, 1.0, 1.0),
    ],
    dtype=dtype,
)

# 将结构化数组的x/y/z字段转为(N,3)的float64数组视图（无数据复制）
pos_view = atoms.view(np.float64).reshape(-1, 3)
# 执行旋转变换，结果直接写入原数组内存
pos_view[:] = pos_view @ transform.T

# 验证结果
print(atoms)

方案说明 #

1. 视图原理：由于结构化数组中x/y/z字段均为float64类型，内存布局是连续的，atoms.view(np.float64)仅改变数组的类型解释方式，不复制任何数据，时间复杂度为O(1)。
2. 效率优势：针对千万级原子，该方案的运算效率与操作普通(N,3)数组完全一致，避免了结构化数组与非结构化数组转换的额外内存开销和时间损耗。
3. 兼容性：修改视图的操作会直接同步到原结构化数组，无需额外赋值步骤，保证数据一致性。

备选方案（手动展开运算） #

如果不想使用视图，也可以手动展开矩阵乘法公式，直接对字段进行运算，但代码冗余度较高：

x_new = atoms['x'] * transform[0,0] + atoms['y'] * transform[0,1] + atoms['z'] * transform[0,2]
y_new = atoms['x'] * transform[1,0] + atoms['y'] * transform[1,1] + atoms['z'] * transform[1,2]
z_new = atoms['x'] * transform[2,0] + atoms['y'] * transform[2,1] + atoms['z'] * transform[2,2]
atoms['x'], atoms['y'], atoms['z'] = x_new, y_new, z_new

内容的提问来源于stack exchange，提问作者Abel Gutiérrez