避免在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缓存可以解决问题:

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否则每次调用都会重新初始化缓存)