numpy 如何让Python更快地运行4个嵌套循环?

e5nszbig  于 2023-03-08  发布在  Python
关注(0)|答案(2)|浏览(168)

我的Python代码中有4个嵌套循环,当网格很大时,要花很长时间才能完成所有循环。例如,下面是我的一段代码:

from itertools import product

T        =  50
beta     =  0.98 
alpha    =  0.3
delta    =  0.1

xss      =  ((1 / beta - (1 - delta)) / alpha)**(1 / (alpha - 1))
yss      =  xss**alpha

x        =  np.linspace(0.8 * xss, 1.4 * xss, T,dtype='float32')
y        =  np.linspace(yss, 1.2 * yss, T,dtype='float32')

for i,j,a,b in product(range(T),range(T),range(T),range(T)):
    z[i*T+a,j*T+b]    =  x[i]**alpha + (1 - delta) * x[i] + y[j] - x[a] - y[b]

有没有什么解决方案可以帮助Python更快地创建矩阵z

okxuctiv

okxuctiv1#

Python中的循环是很慢的,尽可能地将其保存在numpy中,一个好的开始应该是注意到您正在创建一个基于四个参数的新矩阵;把参数转换成一个4维数组,让numpy广播为你处理循环。这一个命令应该就是你所需要的,而且它会非常非常快:

z4 = (
    (x ** alpha + (1 - delta) * x)[:, None, None, None]
    + y[None, None, :, None]
    - x[None, :, None, None]
    - y[None, None, None, :]
)

您可能会注意到,代码中与i索引有关的所有内容都放在了第一维([:, None, None, None])中;与j有关的所有内容(按[None, None, :, None]),等等。这将以[T, T, T, T]形状结束,每个参数(在代码中分别为iajb)对应一个维度。之后,由于维度的顺序正确,您可以按以下方式将其折叠为[T * T, T * T]形状

z = np.reshape(z4, [T * T, T * T])

结果中存在细微差别,因为我的计算完全停留在float32区域,而您的计算则徘徊在float64区域,然后被迫进入z.dtype区域(我在你的代码中看不到z的定义)。这是因为x ** alpha(1 - delta) * x将与x.dtype保持一致,完全在numpy内;但是x[i] ** alpha(1 - delta) * x[i]就不会受到这样的限制。如果你把代码改成这样,你可以验证你得到了同样的结果:

x[i]**np.float32(alpha) + np.float32(1 - delta) * x[i] + y[j] - x[a] - y[b]

也就是说,这种加速之所以可能,是因为xynumpy阵列:numpy对数字的存储效率非常高,在计算过程中不需要分配或解除分配,所有这些都是在Pythonland之外的一个非常低的级别上完成的,假设您没有在标记中指定numpy,这可能有点欺骗;但是你确实在你的代码中使用了numpy,所以我想应该没问题。如果你对numpy的答案没问题,那么添加numpy标记就好了。
在纯Python中,你能做的最好的事情就是预先计算你能做的:

xi = [vx**alpha + (1 - delta) * vx for vx in x]

然后运行循环,将x[i]**alpha + (1 - delta) * x[i]替换为xi[i]。没有其他方法可以使此过程更快。

y0u0uwnf

y0u0uwnf2#

首先,product(range(T),range(T),range(T),range(T))在python中是50^4次迭代,其中for循环极其缓慢。

  • 优化

缓存相同的值以避免在for循环中重复计算。
例如:

  • (1 - delta)重复50^4
  • x[i]**alpha重复50^3
  • ...

代码如下:

_delta = 1 - delta
# other caches
for i,j,a,b in product(range(T), repeat=4):
    z[i*T+a,j*T+b]    =  x[i]**alpha + _delta * x[i] + y[j] - x[a] - y[b]

相关问题