如何有效地填充 numpy 二维数组？

Question

我想创建一个大小为 (N_r * N_z) 的 2D numpy 数组。

跨列，1 个特定列（比如j）的元素应基于值r_thresh[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 数组，在我想创建二维数组之前已经填充了值。

我想问你我该如何走得更远，使用这个“规则”来创建 numpy 二维数组的每个元素，并以最有效的方式（速度方面）实际创建整个数组。

我最初使用 2 个嵌套的 for 循环和普通的 python 列表编写了所有代码，并且代码可以运行，但运行时间很长。

更有经验的程序员告诉我要避免 for 循环并使用 numpy，因为它是最好的。

我现在了解如何使用 numpy np.arange() 指令创建一维数组，但我缺乏如何将其推断为 2 维的知识。

谢谢！

Answer 1

最简单的方法是使用einsum。在r_thresh 的形状为(N_z,) 的情况下，您可以使用此代码：

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,:]

EDIT1

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]])