避免在 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。

鉴于这些要求，我正在考虑使用带有数组作为属性的可调用对象。你怎么看待这件事？有没有更pythonic的方法？

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 cache 可以解决问题：

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

（不要在foo 中定义get_matrix，否则每次调用都会重新初始化缓存）