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

Question

我正在尝试加速可以最低限度表示的函数：

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，同时保持 'vectorized' 允许使用 np.meshgrid[= 产生的网格调用它的行为12=]

Answer 1

这似乎主要是让广播形状正确的问题。要理解这里发生的事情，请注意给定一个一维数组 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

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

Replacing a for loop related to a parameter inside function of two variables by numpy sum

optimization

numpy

function

vectorization