Function 用numpy sum替换与两个变量函数内的参数相关的for循环

我正在尝试加速一个函数，该函数可以通过以下方式最低限度地表示：

import numpy as np

def simple_function(x, y, a1, a2, a3):
    return a1 + a2*x**2/(1 + a3*y**2)


def to_optimnize(x, y, a1, a2, a3, N):
    Sigma = 0
    for i in range(len(N)):
        yn = N[i]*y
        Sigma = Sigma + N[i]*simple_function(x, yn, a1, a2, a3)
    return Sigma


x = np.linspace(0, 10, 1000)
y = np.linspace(0, 4, 200)

N = np.random.random((180,))
a1, a2, a3 = 1, 2, 3


X, Y = np.meshgrid(x, y)
test = simple_function(X, Y, a1, a2, a3)
result = to_optimnize(X, Y, a1, a2, a3, N)

---

for循环主要是一个累积和，尽管我不知道如何在这里使用numpy，同时保持“矢量化”行为，允许使用

np.meshgrid

生成的网格调用它。这似乎主要是为了获得正确的广播形状。要理解这里发生了什么，请注意，给定一个一维数组

a

，

a[None，：]

将创建一个第一维长度为

1

的二维数组

a[：，None]

创建一个二维数组，其第二个维度的长度为

1

def to_optimize_new(x, y, a1, a2, a3, n):
    n = n[None, None, :]
    y = y[:, :, None]
    x = x[:, :, None]
    yn = n * y
    series = n * simple_function(x, yn, a1, a2, a3)
    return series.sum(axis=2)

当使用

np进行测试时，这会给出正确的结果。allclose

：

>>> np.allclose(to_optimize_new(X, Y, a1, a2, a3, N), 
...             to_optimize(X, Y, a1, a2, a3, N))
True

这种方法可能会占用大量内存，因为所有操作的结果都会在最后存储和汇总。但是对于这个例子，它工作得很好

顺便说一下，如果除了启用广播之外，您没有理由使用

meshgrid

，那么您可以使用相同的整形技巧来避免调用

meshgrid

，如下所示：

def to_optimize_nomesh(x, y, a1, a2, a3, n):
    n = n[None, None, :]
    x = x[None, :, None]
    y = y[:, None, None]
    yn = n * y
    series = n * simple_function(x, yn, a1, a2, a3)
    return series.sum(axis=2)

如果您想要一个完全通用的函数，只需完全删除重塑命令。这将适用于标量输入

def to_optimize_generic(x, y, a1, a2, a3, n):
    yn = n * y
    series = n * simple_function(x, yn, a1, a2, a3)
    return series

但是，如果要使用向量输入，必须在函数外为它们指定正确的形状，并在返回后执行求和：

>>> np.allclose(
...     to_optimize_generic(x[None, :, None], 
...                         y[:, None, None], 
...                         a1, a2, a3, 
...                         N[None, None, :]).sum(axis=2), 
...     to_optimize(X, Y, a1, a2, a3, N))
True

您可以根据

N

中具有

：

的轴来判断要求和的轴

>>> np.allclose(
...     to_optimize_generic(x[None, None, :], 
...                         y[None, :, None], 
...                         a1, a2, a3, 
...                         N[:, None, None]).sum(axis=0), 
...     to_optimize(X, Y, a1, a2, a3, N))
True

---

这就行了。然而，我有点担心的是，

x

和

y

的维度失去了一般性。虽然我主要使用此函数绘制计算曲面，但我也喜欢计算特定点处的值（

x

和

y

为浮动），或沿路径（

x

和

y

为1D numpy数组）。“我怎样才能使你的回答适应这个问题呢？”Liris。通过仔细跟踪您的尺寸。我想我无法避免在开始时测试尺寸，并在更改形状时根据长度1调整尺寸数。@Liris，这取决于您的意思。您可以将该部分工作推到功能之外。请参阅我的第二次编辑。