Как использовать все ядра для ускорения программы моделирования на основе трехмерного массива numpy?

Я использую python для создания модели физического моделирования. Теперь у меня есть два пустых трехмерных массива arr_A и arr_B размером 50*50*15 (может быть увеличен до 1000*1000*50 в будущем). И я хочу посмотреть, как эти два массива эволюционируют на основе определенных вычислений. Я пытался ускорить свою программу с помощью параллельных вычислений на моей 12-ядерной машине, но результат был не очень хорошим. Я наконец понял, что Python очень медленный в научных вычислениях.

Должен ли я переписывать свою программу на языке C? Это довольно тяжелая работа. Я слышал, что Cython может быть решением. Должен ли я использовать его? Мне очень нужен совет, как ускорить мою программу, так как я новичок в программировании. Я работаю на машине Win10 x64 с 12 ядрами.

Мои вычисления примерно такие:
Значение в arr_A равно 1 или 0. Для каждой «1» в arr_A мне нужно вычислить определенное значение в соответствии с 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.

После прочтения это и некоторых документов о многопоточности и многопроцессорности я попытался использовать многопроцессорность, но симуляция не намного быстрее.

Я использую слайс типа это для многопроцессорной обработки. Если быть точным, «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