python - 正确写一个numpy meshgrid的函数

Question

下面的代码有效。但是，计算 rez 的行可能不是按照 numpy 开发人员预期的方式编写的。我如何才能以有效的方式重写它，而无需与列表进行转换？

此外，如果您碰巧知道如何在 python 中编写数学函数的好指南，它可以同样处理数字和 np.arrays，请分享 :)

import numpy as np
from matplotlib import pyplot as plt

def L2normSquared(X, Y, x, y):
    return sum((X-x)**2 + (Y-y)**2)

X = np.random.normal(0, 1, 10)
Y = np.random.normal(0, 1, 10)

PX = np.arange(-1, 1, 0.1)
PY = np.arange(-1, 1, 0.1)
PXm, PYm = np.meshgrid(PX, PY)

rez = np.array([[L2normSquared(X, Y, x, y) for y in PY] for x in PX])

print(rez.shape)

plt.imshow(rez, cmap='jet', interpolation='nearest', origin='lower', extent=[-1, 1, -1, 1])
plt.show()

编辑:让我解释一下我想做什么

我在 2D 中生成 10 个随机点。然后我为 2D 中的任意点 (x,y) 定义了一个损失函数。该函数的结果是所有 10 个固定点到该任意点的欧几里德距离之和。最后，我想使用 2d imshow 方法绘制这个损失函数

编辑 2:根据 Nils Werner 的回答，可以使用 3D 阵列和广播，生成以下代码

import numpy as np
from matplotlib import pyplot as plt

def L2normSquared(X, Y, x, y):
    return np.sum((X-x)**2 + (Y-y)**2, axis=0)

X = np.random.normal(0, 1, 10)
Y = np.random.normal(0, 1, 10)

PX = np.arange(-1, 1, 0.1)
PY = np.arange(-1, 1, 0.1)
PXm, PYm = np.meshgrid(PX, PY)


rez = L2normSquared(X[:, None, None], Y[:, None, None], PXm, PYm)

print(rez.shape)

plt.imshow(rez, cmap='jet', interpolation='nearest', origin='lower', extent=[-1, 1, -1, 1])
plt.show()

然而，这段代码实际上比列表理解慢(对于 10000 个随机坐标，步长 0.01，大约慢 2-3 倍)。对于更大的输入，会出现内存崩溃，这让我相信这种方法在内部会导致 3D 数组动态规划，这在内存分配方面不能很好地扩展。

编辑 3: 非常抱歉，我的最小示例太小了。在我面临的原始问题中，坐标 X 和 Y 不解耦以允许单独计算它们。原始功能之一是

def gaussian(X, Y, x, y):
  return sum(np.exp(-(X-x)**2 -(Y-y)**2))

Answer 1

meshgrid 背后的想法是您可以获得两个或更多数组，您可以简单地将它们传递给操作本身。所以理想情况下，我们根本不需要 for 循环。

但由于您在 X 和 PX 之间进行“外部差异”，之后您对 X 轴求和，您也需要使用broadcasting先做外积，最后在正确的轴上求和:

import numpy as np
from matplotlib import pyplot as plt

def L2normSquared(X, Y, x, y):
    return np.sum((X-x)**2 + (Y-y)**2, axis=0)

X = np.random.normal(0, 1, 10)
Y = np.random.normal(0, 1, 10)

PX = np.arange(-1, 1, 0.1)
PY = np.arange(-1, 1, 0.1)
PXm, PYm = np.meshgrid(PX, PY)

rez = L2normSquared(X[:, None, None], Y[:, None, None], PXm, PYm)