python - 避免 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 ，否则缓存将在每次调用时重新初始化)