用於節省內存的 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 中進行所有循環，並且僅對標量進行操作，而不是利用矢量化。