Избегайте повторного вычисления больших массивов внутри функции в Python

Question

Я повторяю - во время процесса минимизации - функцию, которая требует больших массивов.Вот фиктивный пример

def foo(N,a):
    big_array = np.mgrid[0:N,0:N]
    b = np.fft.fft2(big_array[0]**a) #some heavy computation
    return b

Во время процесса минимизации размер массива N не меняется, поэтому я хотел бы использовать тот же массив, чтобы избежать бесполезных вычислений и выделения памяти.

Кроме того, я бы хотел, чтобы функция foo была самосогласованной, что означает, что я не хочу, чтобы другая функция создала массив и присвоила ему foo в процессе минимизации.

УчитываяЭти требования я думал использовать вызываемый объект с массивом в качестве атрибута.Что Вы думаете об этом?Есть ли более питонский способ сделать?

Answer 1

автономный подход (без глобальной переменной) будет использовать изменяемый аргумент по умолчанию (с которым вы не должны вызывать вашу функцию), чтобы запоминать ранее выделенные массивы, учитывая их размер

, если размер массиванет в словаре, создайте его и добавьте.

def foo(N,a,dict_container={}):
    if N in dict_container:
        big_array = dict_container[N]
    else:
        big_array = np.mgrid[0:N,0:N]
        dict_container[N] = big_array

    b = np.fft.fft2(big_array[0]**a) #some heavy computation
    return b

Основная проблема этого подхода заключается в том, что он отключает сборщик мусора для этого массива, поэтому, если N слишком сильно изменяется, выможет иметь истощение памяти.Та же методика, но использование LRU-кэша может решить проблему:

from functools import lru_cache
@lru_cache(maxsize=32)  # max 32 elements in cache
def get_matrix(N):
    return np.mgrid[0:N,0:N]

def foo(N,a):
    big_array = get_matrix(N)
    b = np.fft.fft2(big_array[0]**a) #some heavy computation
    return b

(не определяйте get_matrix внутри foo, иначе кэш будет переинициализироваться на каждомзвоните)