Замена для l oop, относящегося к параметру внутри функции двух переменных, на numpy sum

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)

Значение для l oop в основном является накопительной суммой, хотя я не могу понять, как использовать numpy здесь, сохраняя векторизованное поведение, позволяющее вызывать его с использованием сетки, полученной из np.meshgrid

Answer 1

Похоже, это в основном вопрос правильной формы вещания. Чтобы понять, что здесь происходит, обратите внимание, что с учетом 1-го массива a, a[None, :] создает 2-й массив с первым измерением длины 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)

Если вы хотите полностью обобщенную функцию c, просто полностью удалите команды изменения формы. Это будет работать со скалярными входами.

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