Hello everyone. I am attempting to use NUMBA for speeding up an MD simulation code. The jitted functions themselves work fine. I have a function that basically updates the position and velocities (and ‘angles’ which is some kind of direction pointing at the ‘vision’ of each particle). Calling the function
iterate(pos,vel,ang,Nrun=1), fails for multiple runs. (i.e. Nrun = 100 or so). I am not sure what is going wrong as the functions themselves work fine but the kernel dies if I call the function multiple times. Any help would be appreciated.
Note:
calling verlet_step(…) without updating parameters runs fine for multiple runs
but calling verlet_step(…update_each_step). fails if i update pos and vel at each step… I think the error is somehow related to updating them, but I cannot tell what the issue is.
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Sep 5 11:56:11 2022
@author: priyankaiyer
This code will be written in python and acclerated using numba. A code capable of simulating ABP with vision is written
"""
import numpy as np
import numba
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
@numba.jit(nopython=True)
def calc_neigh_grid(index,pos):
'''
Parameters
----------
index, pos : float
tag and Positions of particles
Returns
-------
particles sorted into grids in the simulation box for neighbour calcuation (O(N))
box_indx returns box indices of 'real' particles
box_count counts number of particles in each 'box'
'''
index_b, pos_b = add_periodic(index,pos)
##re-define r_cut such that simulation box is 'equaly divided'
box_count = np.zeros((nbox,nbox),dtype=np.int_)
box_indx = np.zeros((pos.shape[0],2),dtype=np.int_)
'''
First sort boudary particles into the grid.
'''
for i in range(pos_b.shape[0]):
nx = int(np.floor(((pos_b[i,0]+rcut)/rcut_new)))
ny = int(np.floor(((pos_b[i,1]+rcut)/rcut_new)))
box_count[nx,ny]+= 1
for i in range(pos.shape[0]):
nx = int(np.floor(((pos[i,0]+rcut)/rcut_new)))
ny = int(np.floor(((pos[i,1]+rcut)/rcut_new)))
box_count[nx,ny]+= 1
box_indx[i,0] = nx
box_indx[i,1] = ny
return box_indx, box_count
@numba.jit(nopython=True)
def add_periodic(index,pos):
'''
Parameters
----------
index,pos : float
tags and Position of particles
Returns
-------
Returns position of 'boundary' particles. This vector doesnt have a consstant size (possible orror ?)
'''
left = pos[pos[:,0]>(L-rcut)] -np.array([L,0.0]);
right = pos[pos[:,0]<(rcut)] + np.array([L,0.0]) ;
up = pos[pos[:,1]<(rcut)] + np.array([0.0,L]) ;
down = pos[pos[:,1]>(L-rcut)] - np.array([0.0,L]) ;
corner1 = pos[(pos[:,0]>(L-rcut)) & (pos[:,1]>(L-rcut))] + np.array([-L,-L]);
corner2 = pos[(pos[:,0]>(L-rcut)) & (pos[:,1]<(rcut))] + np.array([-L,L]);
corner3 = pos[(pos[:,0]<(rcut)) & (pos[:,1]>(L-rcut))] +np.array( [L,-L]);
corner4 = pos[(pos[:,0]<(rcut)) & (pos[:,1]<(rcut))] + np.array([L,L]);
pos_new = np.vstack((left,right,up,down,corner1,corner2,corner3,corner4))
left1 = index[pos[:,0]>(L-rcut)];
right1 = index[pos[:,0]<(rcut)];
up1 = index[pos[:,1]<(rcut)];
down1 = index[pos[:,1]>(L-rcut)];
corner11 = index[(pos[:,0]>(L-rcut)) & (pos[:,1]>(L-rcut))];
corner22 = index[(pos[:,0]>(L-rcut)) & (pos[:,1]<(rcut))];
corner33 = index[(pos[:,0]<(rcut)) & (pos[:,1]>(L-rcut))];
corner44 = index[(pos[:,0]<(rcut)) & (pos[:,1]<(rcut))];
index_new = np.hstack((left1,right1,up1,down1,corner11,corner22,corner33,corner44))
return index_new,pos_new
def plot_grid(pos,pos_b):
'''
plots the simulation box,also showing image particles
'''
fig = plt.figure(figsize=(6,6))
ax = fig.add_subplot(1, 1, 1)
n = int((L+3*rcut)/rcut)
rcut_new = (L+3*rcut)/n
major_ticks = np.arange(-3*rcut/2, L+3*rcut/2, rcut_new)
ax.set_xticks(major_ticks)
ax.set_yticks(major_ticks)
ax.grid(which='both')
rect=mpatches.Rectangle((0,0), L,L, linewidth=1, edgecolor='r', facecolor='none')
plt.gca().add_patch(rect)
ax.plot(pos_b[:,0],pos_b[:,1],'o')
ax.plot(pos[:,0],pos[:,1],'x')
plt.show()
@numba.jit(nopython=True)
def LJ(dr,r):
f = 24*eps*(2*((sigma/r)**(12))-(sigma/r)**(6))*dr
return f
@numba.jit(nopython=True)
def vision_(phi,psi1,r):
return omega*np.exp(-r/R0)*np.sin(phi-psi1)
@numba.jit(nopython=True)
def calc_force(pos,ang,box_indx,box_count):
'''
Parameters
----------
pos : TYPE
Positions of particles
ang:
orientation angle of the particles
box_indx : TYPE
array containing the index of the bin the particle belongs to
Returns
-------
f : TYPE
the net force acting on each particle due to- lennard jones, vision interactoion and noise
phidot:
time deriative of the vision angles of each partticle
'''
f = np.zeros_like(pos)
index_b, pos_b = add_periodic(index,pos)
phidot = np.zeros_like(ang)
n_hood = np.array(([0,0],[0,1],[0,-1],[1,0],[1,1],[1,-1],[-1,0],[-1,1],[-1,-1]))
for i in numba.prange(pos.shape[0]):
nx = box_indx[i,0]
ny = box_indx[i,1]
N_vc = 0.0
for j in range(9):
nxx = nx+ n_hood[j,0]
nyy = ny+n_hood[j,1]
if (box_count[nxx,nyy]):
part_neigh = np.where(((box_indx[:,0] == nxx) & (box_indx[:,1] == nyy)))[0]
for count in range(0,box_count[nxx,nyy]):
k = part_neigh[count] ##loop over particle neighbourhood
if (i!=k-1):
if nxx == 0 or nyy == 0 or (nxx == box_count.shape[0]-1) or (nyy == box_count.shape[0]-1):
kk = np.where(index_b==k)[0][0]
dist = ((pos[i,0]-pos_b[kk,0])**2 + (pos[i,1]-pos_b[kk,1])**2)**0.5
dr = (pos_b[kk]-pos[i])/dist
else:
dist = ((pos[i,0]-pos[k-1,0])**2 + (pos[i,1]-pos[k-1,1])**2)**0.5
dr = (pos[k-1]-pos[i])/dist
if dist<4*R0:
phi = np.arctan2(dr[1],dr[0])
if np.cos(phi-ang[i])>=np.cos(vis_angle):
phidot[i] += vision_(phi,ang[i],dist)
N_vc += np.exp(-dist/R0)
if dist < 1.22:
f[i]+= LJ(dr,dist)
if (N_vc):
phidot[i] = phidot[i]/N_vc
f = f + np.random.normal(loc=0.0,scale=np.sqrt(2*gamma*dt),size=f.shape)
phidot = phidot + np.random.normal(loc=0.0,scale=np.sqrt(2*Dr*dt),size=phidot.shape)
return f,phidot
@numba.jit(nopython=True)
def iterate(pos,vel,ang,Nrun):
'''
Parameters
----------
pos : TYPE
particle position
vel : TYPE
particle velocity
ang : TYPE
vision angle
Nrun : TYPE
number of MD steps
Returns
-------
pos : TYPE
updated position after time t= dt*nrun
vel : TYPE
updated velcotiy after time t= dt*nrun
ang : TYPE
updated vision angle after time t= dt*nrun
'''
#r1 = pos
box_indx, box_count = calc_neigh_grid(index,pos)
for i in range(Nrun):
pos,vel,ang = verlet_step(pos,vel,ang, box_indx, box_count)
#dr = pos-r1
#max_disp = np.max(np.sqrt((dr*dr).sum(axis=1)))
if i%2==0:
#'''
#track displacements and recalc. neighbours
#'''
box_indx, box_count = calc_neigh_grid(index,pos)
#r1=pos
return pos,vel,ang
@numba.jit(nopython=True)
def verlet_step(pos,vel,ang, box_indx, box_count):
'''
Parameters
----------
pos : TYPE
particle position
vel : TYPE
particle velocity
ang : TYPE
vision angle
box_indx : TYPE
box indices of 'real' particles
box_count : TYPE
number of particles in each 'box'
Returns
-------
pos : TYPE
single step update of position
vel : TYPE
single step update of velocity
ang : TYPE
single step update of angle
'''
f,phi_dot = calc_force(pos,ang,box_indx,box_count)
f = f -gamma*vel + gamma*v0*np.hstack((np.cos(ang),np.sin(ang)))
pos = pos + dt*vel + 0.5*dt*dt*f
pos = pos%L
ang = ang + dt*phi_dot
fnew ,phi_ = calc_force(pos,ang,box_indx,box_count)
fnew = fnew - gamma*vel + gamma*v0*np.hstack((np.cos(ang),np.sin(ang)))
vel = vel + 0.5*dt*(f+fnew)
return pos,vel,ang
'''
System parameters:
L box length,
N is number of particles,
pos x,y coordiates initaited randomly.
R particle radius
R_cut cut-off vision interactions
skin skin for recalc.
V0 paritcle velocity
'''
L = 250.0
N = 5*5
R =0.5
rcut =5*1.5
n_max = int(rcut*rcut/(np.pi*R*R))
###
dt=0.0001
gamma=100.0
v0=2.0
Dr=0.08
Dt=0.01
sigma=2*R
R0=1.5*sigma
Pe = sigma*v0/Dt
vis_angle=np.pi/4
omega=5
#eps=(1+Pe)
eps = 0.1
nbox = int((L+3*rcut)/rcut)
rcut_new = (L+3*rcut)/nbox
index = np.arange(1,N+1,dtype='int')
#pos = np.random.uniform(0,L,size=(N,2))
dx = np.linspace(rcut,L-rcut,int(N**0.5))
x,y=np.meshgrid(dx, dx, indexing='ij')
pos = np.vstack((x.flatten(),y.flatten()))
pos = pos.transpose()
ang = np.random.uniform(0,np.pi,size=(N,1))
vel = np.random.uniform(0,1,size=(N,1))
index_b,pos_b = add_periodic(index,pos)
'''
Compile functions
'''
box_indx, box_count= calc_neigh_grid(index,pos);
calc_force(pos,ang,box_indx,box_count);
verlet_step(pos,vel,ang, box_indx, box_count);
iterate(pos,vel,ang,Nrun=1);
print('compiled')
pos,vel,ang=iterate(pos,vel,ang,Nrun=10000); ##this fails
# index = np.arange(1,N+1,dtype='int')
# #pos = np.random.uniform(0,L,size=(N,2))
# dx = np.linspace(rcut,L-rcut,int(N**0.5))
# x,y=np.meshgrid(dx, dx, indexing='ij')
# pos = np.vstack((x.flatten(),y.flatten()))
# pos = pos.transpose()
# ang = np.random.uniform(0,np.pi,size=(N,1))
# vel = np.random.uniform(0,1,size=(N,1))
# index_b,pos_b = add_periodic(index,pos)
# plot_grid(pos,pos_b)
#pos,vel,ang=iterate(pos,vel,ang,Nrun=1000)
#index_b,pos_b = add_periodic(index,pos)
#plot_grid(pos,pos_b)
#box_indx, box_count = calc_neigh_grid(index,pos)
#calc_force(pos,ang,box_indx,box_count);
#%timeit calc_neigh_grid(index,pos)
#%timeit calc_force(pos,ang,box_indx,box_count)