Question

我在用nditer迭代ndarray时迷路了。

背景

我正在尝试计算3D阵列中每个点的3x3对称矩阵的特征值。我的数据是形状为[6，x，y，z]的4D阵列，其中6个值是点x，y，z处的矩阵值，在浮点数的~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（一旦这个工作我将改变函数调用，缓冲迭代......但到目前为止我只是想让它工作^^）

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-order（列主要），那么迭代第一维可能比迭代迭代更快最后3个维度，从那时起您将访问相邻的内存地址块。因此，您可能需要考虑切换到Fortran-order。

我不希望其中任何一个获得巨大的性能提升，因为在一天结束时你仍然在使用Python进行所有循环并仅在标量上运行而不是利用向量化

Numpy nditer可以节省内存吗？

1 个答案: