function - 用 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，同时保持“矢量化”行为，允许使用由 np.meshgrid 产生的网格来调用它

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，那么您可以使用相同的 reshape 技巧来避免 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)

如果您想要一个完全通用的功能，只需完全删除 reshape 命令即可。这适用于标量输入。

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