如何有效地填充 numpy 二維數組？

Question

我想創建一個大小為 ( N_r * N_z ) 的二維 numpy 數組。

跨列，應根據值r_thresh[j]創建 1 個特定列（例如j ）的元素。 因此，numpy 2D 數組中的N_z列總數中的 1 列（例如j ）被創建為：

(np.arange(N_r) + 0.5) * r_thresh[j] # this gives an array of size (1, N_r)

當然，列j + 1應創建為：

(np.arange(N_r) + 0.5) * r_thresh[j+1] # this gives an array of size (1, N_r)

r_thresh是一個大小為(1, N_z)的 numpy 數組，在我想創建二維數組之前已經填充了值。

我想問你如何進一步使用 go 並使用創建 numpy 二維數組的每個元素的“規則”，並以最有效的方式（速度方面）實際創建整個數組。

我最初使用 2 個嵌套的 for 循環和普通的 python 列表編寫了所有代碼，並且代碼工作正常，但運行時間很長。

更有經驗的程序員告訴我要避免 for 循環並使用 numpy 因為它是最好的。

我現在了解如何使用 numpy np.arange()指令創建一維 arrays，但我缺乏如何將其外推到二維的知識。

謝謝！

Answer 1

最簡單的方法是使用einsum 。 對於形狀為(N_z,) r_thresh r_thresh ，您可以使用以下代碼：

res = np.einsum("i,j->ij", np.arange(N_r) + 0.5, r_thresh)

此外，您可以將np.arange(N_r) + 0.5重塑為形狀(N_r,1)並將r_thresh為形狀(1,N_z) 。 因此，您可以使用點積（對於 Python 版本 > 3.5）：

res = (np.arange(N_r) + 0.5).reshape(N_r,1) @ r_thresh.reshape(1,N_z)

或按照hpaulj的評論：

res = (np.arange(N_r) + 0.5)[:,None] @ r_thresh[None,:]

編輯1

hpaulj的評論也很有幫助（我將其粘貼到我的答案中以便更好地查看）：

res = (np.arange(N_r) + 0.5)[:,None] * r_thresh
res = np.outer(np.arange(N_r) + 0.5, r_thresh)

此外

您還可以使用tensordot ：

res = np.tensordot((np.arange(N_r) + 0.5)[:,None], r_thresh[:,None], axes=[[-1],[-1]])