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

Я пытаюсь ускорить функцию, которая может быть минимально представлена:

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

Всего 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)

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

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

Есть идеи?

10000