用于节省内存的 Numpy nditer？

Question

使用 nditer 迭代 ndarray 时，我迷路了。

背景

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

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

问题是现在在我的矢量化中，每个操作都会创建一个包含我的数据大小的内部数组，最终导致使用过多的 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 中进行所有循环，并且仅对标量进行操作，而不是利用矢量化。