Недавно я написал на Python реализацию модели Беддингтона ДеАнджелиса, которая используется для моделирования популяций хищников и жертв.
Моя проблема в том, что код очень медленный. 10.000 итераций занимают 230 секунд, когда эта программа должна иметь возможность выполнять итерации 1 миллион раз за разумные сроки.
Я понимаю, что я мог бы просто переписать на C, так как это в основном математика, но я действительно хочу научиться правильно векторизовать такую программу в Python.
Ситуация для упрощения заключается в том, что у меня есть два массива фигур размером 200x200, и мне нужно выполнить итерацию каждого элемента каждого массива, используя один и тот же индексный элемент из другого массива плюс некоторые окружающие элементы этого же массива. Так, например, если я работаю над [1] [1], мне также понадобится:
б [1] [1]
а [0] [1]
а [0] [- 1]
а [1] [0]
а [-1] [0]
Вся операция должна быть полностью векторизованной, потому что я изменяю все элементы 200x200x2 за один временной шаг.
Так как бы я вызвал эту функцию, чтобы получить эти индексы?
Любой совет будет принят с благодарностью.
Полный код для контекста: (это выглядит пугающе, но на самом деле очень просто)
import numpy as np
import copy
import time
def get_cell_zero_flux(m,i,j,x,y,prey):
"""
Fetch an array element that could be outside of the border
"""
if i >= x or i < 0 or j >= y or j < 0:
if prey: return 0.43058
else: return 0.718555
return m[i][j]
def get_laplacian(n,i,j,x,y,neighbors,h,prey):
"""
Generate the laplacian value for the given element
"""
total = 0
for ng in neighbors:
cell = get_cell_zero_flux(n,i+ng[0],j+ng[1],x,y,prey)
total += cell
return (total - 4*n[i][j]) / (h**2)
def next_n(n,p,nl,pl,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b):
"""
Integrate prey population function
"""
return n + t * (r * ( 1 - n / k ) * n
- beta * n / ( b + n + w * p ) * p + d11 * nl + d12 * pl)
def next_p(n,p,nl,pl,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b):
"""
Integrate predator population function
"""
return p + t * (e * beta * n / ( b + n + w * p )
* p - ni * p + d21 * nl + d22 * pl)
def generate_preys(x,y,epsilon,n_start):
"""
Generate the initial population of preys
"""
n = np.random.rand(x, y)
n = np.interp(n,(n.min(),n.max()),(-epsilon/2,epsilon/2))
n = n + n_start
return n
def generate_predators(x,y,p_start):
"""
Generate the initial population of predators
"""
p = np.ones((x,y))
p.fill(p_start)
return p
def generate_n(n0,n,p,x,y,neighbors,h,d11,d12,t,r,e,beta,k,ni,w,b):
"""
Vectorized element iteration attempt for preys
"""
i,j = np.where(n==n0) # this wouldnt work, need the current element
n_laplacian = get_laplacian(n,i,j,x,y,neighbors,h,True)
p_laplacian = get_laplacian(p,i,j,x,y,neighbors,h,False)
p0 = p[i,j]
return next_n(n0,p0,laplacian,d11,d12,t,r,e,beta,k,ni,w,b)
def generate_p(p0,p,n,x,y,neighbors,h,d21,d22,t,r,e,beta,k,ni,w,b):
"""
Vectorized element iteration attempt for predators
"""
i,j = np.where(p==p0) # this wouldnt work, need the current element
n_laplacian = get_laplacian(n,i,j,x,y,neighbors,h,True)
p_laplacian = get_laplacian(p,i,j,x,y,neighbors,h,False)
n0 = n[i,j]
return next_p(n0,p0,n_laplacian,
p_laplacian,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b)
def generate_system(x,y,h,d11,d12,d21,d22,t,r,e,
beta,k,ni,w,b,ite,n_start,p_start,epsilon):
"""
System generation
"""
# Initial distribution
n = generate_preys(x,y,epsilon,n_start)
p = generate_predators(x,y,p_start)
#n = n.tolist()
#p = p.tolist()
ps = []
ns = []
# neighbor list for easy laplacian neighbor fetch
neighbors = [[-1,0],[1,0],[0,1],[0,-1]]
t1 = time.time()
for it in range(ite):
# record each iteration
old_n = copy.copy(n)
old_p = copy.copy(p)
ns.append(old_n)
ps.append(old_p)
# main array element iteration for prey and predator arrays
for i in range(x):
for j in range(y):
n_laplacian = get_laplacian(old_n,i,j,x,y,neighbors,h,True)
p_laplacian = get_laplacian(old_p,i,j,x,y,neighbors,h,False)
n0 = old_n[i][j]
p0 = old_p[i][j]
n[i][j] = next_n(n0,p0,n_laplacian,p_laplacian,
d11,d12,d21,d22,t,r,e,beta,k,ni,w,b)
p[i][j] = next_p(n0,p0,n_laplacian,p_laplacian,
d11,d12,d21,d22,t,r,e,beta,k,ni,w,b)
"""
n = generate_n(old_n,old_n,old_p,x,y,neighbors,
h,d11,d12,t,r,e,beta,k,ni,w,b)
p = generate_p(old_p,old_p,old_n,x,y,neighbors,
h,d21,d22,t,r,e,beta,k,ni,w,b)
"""
t2 = time.time()
print(t2-t1)
return ns,ps
ns,ps = generate_system(x=50,y=50,h=0.25,d11=0.01,d12=0.0115,d21=0.01,d22=1,
t=0.01,r=0.5,e=1,beta=0.6,k=2.6,ni=0.25,w=0.4,b=0.3154,
ite=10,n_start=0.43058,p_start=0.718555,epsilon=0.001)
Ожидаемый результат - 1 миллион итераций в течение нескольких минут в сетке 200x200, но 230 секунд только на 10.000 в сетке 40x40
EDIT
Мне удалось векторизовать всю программу. Увеличение производительности было в 400 раз. WOW
Вот новый код:
import numpy as np
import copy
import time
def next_n(n,p,nl,pl,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b):
"""
Integrate prey population function
"""
return n + t * (r * ( 1 - n / k ) * n
- beta * n / ( b + n + w * p ) * p + d11 * nl + d12 * pl)
def next_p(n,p,nl,pl,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b):
"""
Integrate predator population function
"""
return p + t * (e * beta * n / ( b + n + w * p )
* p - ni * p + d21 * nl + d22 * pl)
def generate_preys(x,y,epsilon,n_start):
"""
Generate the initial population of preys
"""
n = np.random.rand(x, y)
n = np.interp(n,(n.min(),n.max()),(-epsilon/2,epsilon/2))
n = n + n_start
n[0,:] = n_start
n[-1:,:] = n_start
n[:,0] = n_start
n[:,-1:] = n_start
return n
def generate_predators(x,y,p_start):
"""
Generate the initial population of predators
"""
p = np.ones((x,y))
p.fill(p_start)
return p
def get_laps(a,x,y,h):
center = a[1:-1,1:-1]
left = a[1:-1,0:-2]
right = a[1:-1,2:]
top = a[0:-2,1:-1]
bottom = a[2:,1:-1]
return (left+right+top+bottom - 4*center) / (h**2)
def generate_system(x,y,h,d11,d12,d21,d22,t,r,e,
beta,k,ni,w,b,ite,n_start,p_start,epsilon):
"""
System generation
"""
# Initial distribution
n = generate_preys(x+2,y+2,epsilon,n_start)
p = generate_predators(x+2,y+2,p_start)
ps = []
ns = []
t1 = time.time()
for it in range(ite):
if it % 10000 == 0:
print(f"iterations passed: {it}")
ns.append(copy.copy(n))
ps.append(copy.copy(p))
# record each iteration
nl = get_laps(n,x,y,h)
pl = get_laps(p,x,y,h)
nc = n[1:-1,1:-1]
pc = p[1:-1,1:-1]
n[1:-1,1:-1] = next_n(nc,pc,nl,pl,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b)
p[1:-1,1:-1] = next_p(nc,pc,nl,pl,d11,d12,d21,d22,t,r,e,beta,k,ni,w,b)
t2 = time.time()
print(f"Time taken: {t2-t1}")
return ns,ps
ns,ps = generate_system(x=200,y=200,h=0.25,d11=0.01,d12=0.0115,d21=0.01,d22=1,
t=0.01,r=0.5,e=1,beta=0.6,k=2.6,ni=0.25,w=0.4,b=0.3154,
ite=100,n_start=0.43058,p_start=0.718555,epsilon=0.001)