Numpy nditer 用于节省内存？

Question

使用 nditer 遍历 ndarray 时我迷路了。

背景

我正在尝试为 3D 数组中的每个点计算 3x3 对称矩阵的特征值。 我的数据是一个形状为 [6,x,y,z] 的 4D 数组，其中 6 个值是点 x、y、z 处的矩阵值，位于 float32 的 ~500x500x500 立方体上。 我首先使用了 numpy 的 eigvalsh，但它针对大型矩阵进行了优化，而我可以对 3x3 对称矩阵使用解析简化。

然后我实现了 wikipedia's simplification ，既是一个函数，它接受单个矩阵并计算特征值（然后用嵌套的 for 循环天真地迭代），然后使用 numpy 进行向量化。

问题是现在在我的矢量化中，每个操作都会创建一个数据大小的内部数组，最终导致使用过多的 RAM 和 PC 冻结。

我尝试使用 numexpr 等，它仍然在 10G 左右。

我要做什么

我想通过我的数组迭代（使用 numpy 的 nditer），以便为每个矩阵计算我的特征值。这将消除分配大量中间数组的需要，因为我们一次只计算约 10 个浮点数。 基本上试图将嵌套的for 循环替换为一个迭代器。

我正在寻找这样的东西：

for a,b,c,d,e,f in np.nditer([symMatrix,eigenOut]): # for each matrix in x,y,z

    # computing my output for this matrix
    eigenOut[...] = myLovelyEigenvalue(a,b,c,d,e,f)

到目前为止，我最好的是：

for i in np.nditer([derived],[],[['readonly']],op_axes=[[1,2,3]]):

但这意味着i 采用 4D 数组的所有值，而不是长度为 6 的元组。 我似乎无法掌握 nditer 文档的窍门。

我做错了什么？对于迭代“除一个”轴之外的所有轴，您有任何提示和技巧吗？

关键是要有一个在迭代时优于常规嵌套循环的 nditer（一旦它工作，我将更改函数调用，缓冲区迭代......但到目前为止我只希望它工作 ^^）

Answer 1

您实际上并不需要 np.nditer。一种更简单的迭代除第一个轴之外的所有轴的方法是重塑为[6, 500 ** 3] 数组，将其转置为[500 ** 3, 6]，然后遍历行：

for (a, b, c, d, e, f) in (symMatrix.reshape(6, -1).T):
    # do something involving a, b, c, d, e, f...

如果你真的想使用np.nditer，那么你可以这样做：

for (a, b, c, d, e, f) in np.nditer(x, flags=['external_loop'], order='F'):
    # do something involving a, b, c, d, e, f...

需要考虑的一个潜在重要的事情是，如果 symMatrix 是 C 顺序（行主要）而不是 Fortran 顺序（列主要），那么迭代第一个维度可能比迭代最后 3 个维度快得多尺寸，从那时起您将访问相邻的内存地址块。因此，您可能需要考虑切换到 Fortran 顺序。

我不认为这两种方法都会带来巨大的性能提升，因为归根结底，您仍然在 Python 中执行所有循环并且仅在标量上运行，而不是利用矢量化。