python - 笨拙的nditer以节省内存？

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

背景

我正在尝试为3D数组中的每个点计算3x3对称矩阵的特征值。
我的数据是一个形状为[6，x，y，z]的4D数组，其中6个值是在float32的〜500x500x500立方体上的点x，y，z处的矩阵值。
我首先使用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，该nditer在迭代时的性能要优于常规的嵌套循环（一旦成功，我将更改函数调用，缓冲区迭代...，但到目前为止，我只希望它能工作^^）

最佳答案

您实际上并不需要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阶（列大），则在第一个维度上进行迭代可能比在后三个维度上进行迭代要快得多，因为那么您将访问存储器地址的相邻块。因此，您可能需要考虑切换到Fortran顺序。

我不会期望这两种方法都能带来巨大的性能提升，因为到了最后，您仍将在Python中进行所有循环并且仅在标量上运行，而不是利用矢量化。

关于python - 笨拙的nditer以节省内存？，我们在Stack Overflow上找到一个类似的问题：https://stackoverflow.com/questions/31180797/