如何利用所有内核加速基于numpy 3D阵列的仿真程序?

问题描述 投票:1回答:1

我正在使用python构建物理仿真模型。现在我有两个numpy 3d数组arr_A和arr_B,大小为50 * 50 * 15(将来可能放大到1000 * 1000 * 50)。我想看看这两个阵列是如何根据某些特定的计算而发展的。我试图使用我的12核机器并行计算来加速我的程序,但结果并不是那么好。我终于意识到python在科学计算方面非常缓慢。

我是否必须用C语言重写我的程序?这是一项相当艰苦的工作。我听说Cython可能是一个解决方案。我应该用吗?因为我是编程的初学者,所以我真的需要一些关于如何加速我的程序的建议。我正在开发一款带有12个内核的win10 x64机器。

我的计算是这样的: arr_A中的值为1或0.对于arr_A中的每个“1”,我需要根据arr_B计算某个值。例如,如果arr_A [x,y,z] == 1,则C [x,y,z] = 1 /(arr_B [x-1,y,z] + arr_B [x,y-1,z] + arr_B [X,Y,Z-1] + arr_B [X + 1,Y,Z] + arr_B [X,Y + 1,Z] + arr_B [X,Y,Z + 1])。然后我使用数组C中的最小值作为函数的参数。该函数可以稍微改变arr_A和arr_B,以便它们可以进化。然后我们再次计算“结果”并且循环继续。

请注意,对于每个单独的C [x,y,z],都涉及arr_B中的许多值。否则我可以这样做:

C = arr_B[arr_A>0]**2

我希望解决方案可以像这样简单。但除了三嵌套'for'循环外,我找不到任何可行的索引方法。

在阅读了this和一些有关多线程和多处理的文档后,我尝试使用多处理,但模拟速度并不快。

我使用像this这样的切片进行多处理。具体而言,carrier_3d和potential_3d分别是上面提到的arr_A和arr_B。我将切片放入不同的子流程中。这里没有给出功能的细节,但你可以得到主要的想法。

chunk = np.shape(carrier_3d)[0] // cores
p = Pool(processes=cores)
for i in range(cores):
    slice_of_carrier_3d = slice(i*chunk, 
                                np.shape(carrier_3d)[0] if i == cores-1 else (i+1)*chunk+2)
    p.apply_async(hopping_x_section, args=(i, chunk,carrier_3d[slice_of_carrier_3d, :, :], 
                                               potential_3d[slice_of_carrier_3d, :, :]), 
                                    callback=paral_site_record)
p.close()
p.join() 

如果您想了解有关计算的更多信息,以下代码基本上是我的计算如何在没有多处理的情况下工作。但我已经解释了上面的过程。

def vab(carrier_3d, potential_3d, a, b):
    try:
        Ea = potential_3d[a[0], a[1], a[2]]
        Eb = potential_3d[b[0], b[1], b[2]]
        if carrier_3d[b[0], b[1], b[2]] > 0:
            return 0
        elif b[2] < t_ox:
            return 0
        elif b[0] < 0 or b[1] < 0:
            return 0
        elif Eb > Ea:
            return math.exp(-10*math.sqrt((b[0]-a[0])**2+
                                              (b[1]-a[1])**2+(b[2]-a[2])**2)-
                                              q*(Eb-Ea)/(kB*T))
        else:
            return math.exp(-10*math.sqrt((b[0]-a[0])**2+
                                              (b[1]-a[1])**2+(b[2]-a[2])**2))
    except IndexError:
        return 0
#Given a point, get the vij to all 26 directions at the point
def v_all_drt(carrier_3d, potential_3d, x, y, z):
    x_neighbor = [-1, 0, 1]
    y_neighbor = [-1, 0, 1]
    z_neighbor = [-1, 0, 1]  
    v = []#v is the hopping probability
    drtn = []#direction
    for i in x_neighbor:
        for j in y_neighbor:
            for k in z_neighbor:
                v.append(vab(carrier_3d, potential_3d, 
                             [x, y, z], [x+i, y+j, z+k]))
                drtn.append([x+i, y+j, z+k])
    return np.array(v), np.array(drtn)
    #v is a list of probability(v_ij) hopping to nearest sites.
    #drt is the corresponding dirction(site).
def hopping():  
    global sys_time
    global time_counter
    global hop_ini
    global hop_finl
    global carrier_3d
    global potential_3d
    rt_min = 1000#1000 is meaningless. Just a large enough name to start
    for x in range(np.shape(carrier_3d)[0]):
        for y in range(np.shape(carrier_3d)[1]):
            for z in range(t_ox, np.shape(carrier_3d)[2]):
                if carrier_3d[x, y, z] == 1:
                    v, drt = v_all_drt(carrier_3d, potential_3d, x, y, z)
                    if v.sum() > 0:
                        rt_i = -math.log(random.random())/v.sum()/v0
                        if rt_i < rt_min:
                            rt_min = rt_i
                            v_hop = v
                            drt_hop = drt
                            hop_ini = np.array([x, y, z], dtype = int)
    #Above loop finds the carrier that hops. 
    #Yet we still need the hopping direction.
    rdm2 = random.random()
    for i in range(len(v_hop)):
        if (rdm2 > v_hop[:i].sum()/v_hop.sum()) and\
            (rdm2 <= v_hop[:i+1].sum()/v_hop.sum()):
                hop_finl = np.array(drt_hop[i], dtype = int)
                break      
    carrier_3d[hop_ini[0], hop_ini[1], hop_ini[2]] = 0
    carrier_3d[hop_finl[0], hop_finl[1], hop_finl[2]] = 1 
def update_carrier():
    pass
def update_potential():
    pass
#------------------------------------- 
carrier_3d = np.random.randn(len_x, len_y, len_z)
carrier_3d[carrier_3d>.5] = 1
carrier_3d[carrier_3d<=.5] = 0
carrier_3d = carrier_3d.astype(int)
potential_3d = np.random.randn(len_x, len_y, len_z)
while time_counter <= set_time:# set the running time of the simulation
    hopping() 
    update_carrier()
    update_potential()
    time_counter += 1
python arrays numpy cython scientific-computing
1个回答
0
投票

您可以使用numba创建分析函数的jit编译版本。仅这一点对您的代码来说是最大的加速,并且当您的问题适合约束时,它往往会很好地工作。您必须在for循环中编写更复杂的分析,但我没有看到任何理由说明您所概述的内容不起作用。请参阅以下代码,该代码通过使用numba进行编译显示330倍的加速。您还可以指定某些numba函数并行执行。但是,与此相关的开销只会在问题变得足够大时增加一个加速,所以这是你必须考虑的事情。

from numpy import *
from numba import njit

def function(A, B):
    C = zeros(shape=B.shape)
    X, Y, Z = B.shape
    for x in range(X):
        for y in range(Y):
            for z in range(Z):
                if A[x, y, z] == 1:
                    C[x, y, z] = B[x, y, z]**2
    return C

cfunction = njit(function)
cfunction_parallel = njit(function, parallel=True)

X, Y, Z = 50, 50, 10
A = random.randint(0, 2, size=X*Y*Z).reshape(X, Y, Z)
B = random.random(size=X*Y*Z).reshape(X, Y, Z)

_ = cfunction(A, B)  # force compilation so as not to slow down timers
_ = cfunction_parallel(A, B)

print('uncompiled function')
%timeit function(A, B)

print('\nfor smaller computations, the parallel overhead makes it slower')
%timeit cfunction(A, B)
%timeit cfunction_parallel(A, B)

X, Y, Z = 1000, 1000, 50
A = random.randint(0, 2, size=X*Y*Z).reshape(X, Y, Z)
B = random.random(size=X*Y*Z).reshape(X, Y, Z)

print('\nfor larger computations, parallelization helps')
%timeit cfunction(A, B)
%timeit cfunction_parallel(A, B)

这打印:

uncompiled function
23.2 ms ± 147 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

for smaller computations, the parallel overhead makes it slower
77.5 µs ± 1.69 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
121 µs ± 2.42 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

for larger computations, parallelization helps
138 ms ± 1.84 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
40.1 ms ± 633 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
© www.soinside.com 2019 - 2024. All rights reserved.