model_CUDA.py

import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
import numpy as np
import matplotlib.pyplot as plt
from time import time
import math as m

sigma = 1/10
mu = 0.5
pc = 10**-5
delta = 1.758820149*10**11 #q/m
e0 = 8.854184817*10**-12
N_p = 2000
N_R = 1000
N_t = 1000
T = 2*10**-9
q = 1.6*10**-19
R_0 = 1
V_0 = 4*m.pi*R_0**3/3
Q = pc
rho_0 = Q/V_0
h_t = T/N_t
h_r = R_0/N_R #?
Velosity = [[0, 0] for i in np.arange(0, R_0+h_r, h_r)]
R = [x for x in np.arange(0, R_0+h_r, h_r)]

def Q(r) :
    C1 = 4*pc*m.sqrt(m.pi/2)
    C2 = m.sqrt(m.pi/2)*(sigma**2 + mu**2)
    C3 = m.erf( (r-mu)/(m.sqrt(2)*sigma) ) + m.erf( mu/(m.sqrt(2)*sigma) )
    C4 = sigma**2
    C5 = mu*m.exp(-mu**2/(2*sigma**2))
    C6 = (r+mu)*m.exp(-(r-mu)**2/(2*sigma**2))
    return C1*(C2*C3 + C4*(C5 - C6))

def gamma(r) :
    C1 = delta/(4*m.pi*e0)
    return C1*Q(r)

def lamda(r) :
    return m.sqrt(2*gamma(r))/r**(3/2)

def density_start(r):
    C1 = pc/(m.sqrt(2*m.pi)*sigma)
    return C1*m.exp(-(r-mu)**2/(2*sigma**2))

def density(R0, R, t):
    p0 = density_start(R0)
    C1 = R0**2 * p0 / R**2
    lamda0 = lamda(R0)
    C2 = R/R0 + t*m.sqrt(1 - R0/R)*(delta*p0/(e0*lamda0) - 3*lamda0/2)
    return C1 / C2
	
Density = [density_start(x) for x in np.arange(0, R_0+h_r, h_r)]	
D = [[Q(x)/(4*m.pi*x**2), Q(x)/(4*m.pi*x**2)] for x in np.arange(0, R_0+h_r, h_r)]
pc = Q(R_0)
D[0] = [0, 0]
Density[0] = 0
#print(Density[0], Density[1])
#D=e*e0*E = e0*(1/4pie0)Q/r^2 = Q / (4pi r**2)

def iter() :
	for _ in range(N_t) :
		for i in range(N_R+1 +2) :
			if i == 0 :
				pass
			elif i == 1 :
				dV = Velosity[i][0]*(Velosity[i+1][0]-Velosity[i][0])/(h_r)
				Velosity[i][1] = Velosity[i][0] + h_t*(delta*D[i][0]/e0 - dV)
				R[i] += h_t*Velosity[i][0] + h_t**2*(delta*D[i][0]/e0 - dV)/2
				Density[i] = (D[i+1][0] - D[i][0])/(h_r)
				D[i][1] = D[i][0] - h_t*Density[i]*Velosity[i][0]
			elif i == N_R :
				dV = Velosity[i][0]*(Velosity[i][0]-Velosity[i-1][0])/(h_r)
				Velosity[i][1] = Velosity[i][0] + h_t*(delta*D[i][0]/e0 - dV)
				R[i] += h_t*Velosity[i][0] + h_t**2*(delta*D[i][0]/e0 - dV)/2
				#D[i][1] = pc/(4*m.pi*R[i]**2)
				D[i][1] = D[i][0] - h_t*Density[i]*Velosity[i][0]
				Density[i] = (D[i][0] - D[i-1][0])/(h_r)
				#Density[i] = D[i][1]/R[i]
			elif i < N_R:
				dV = Velosity[i][0]*(Velosity[i][0]-Velosity[i-1][0])/(h_r)
				Velosity[i][1] = Velosity[i][0] + h_t*(delta*D[i][0]/e0 - dV)
				R[i] += h_t*Velosity[i][0] + h_t**2*(delta*D[i][0]/e0 - dV)/2
				Density[i] = (D[i][0] - D[i-1][0])/(h_r)
				D[i][1] = D[i][0] - h_t*Density[i]*Velosity[i][0]
			if i <= N_R and Density[i] < 0:
				Density[i] = 0
				
			if i > 1 :
				D[i-2][0], D[i-2][1] = D[i-2][1], D[i-2][0]
				Velosity[i-2][0], Velosity[i-2][1] = Velosity[i-2][1], Velosity[i-2][0]

	
def iter_gpu() :
	dens = np.array(Density, dtype = np.float64)
	rad = np.array(R, dtype = np.float64)
	Velosityy = np.array([np.array(xi, dtype = np.float64) for xi in Velosity])
	Dd = np.array([np.array(xi, dtype = np.float64) for xi in D])
	
	mod = SourceModule(""" 
		__global__ void run(double* Density, double* R, double** D, double** Velosity, double h_r, double h_t, int N_R, double delta,  double e0)
		{
			int i = blockDim.x * blockIdx.x + threadIdx.x;
			double dV;
			if (i==1) 
				{
					dV = Velosity[i][0]*(Velosity[i+1][0]-Velosity[i][0])/(h_r);
					Velosity[i][1] = Velosity[i][0] + h_t*(delta*D[i][0]/e0 - dV);
					R[i] += h_t*Velosity[i][0] + h_t*h_t*(delta*D[i][0]/e0 - dV)/2;
					Density[i] = (D[i+1][0] - D[i][0])/(h_r);
					D[i][1] = D[i][0] - h_t*Density[i]*Velosity[i][0];
				}
			else if (i == N_R) 
				{
					dV = Velosity[i][0]*(Velosity[i][0]-Velosity[i-1][0])/(h_r);
					Velosity[i][1] = Velosity[i][0] + h_t*(delta*D[i][0]/e0 - dV);
					R[i] += h_t*Velosity[i][0] + h_t*h_t*(delta*D[i][0]/e0 - dV)/2;
					D[i][1] = D[i][0] - h_t*Density[i]*Velosity[i][0];
					Density[i] = (D[i][0] - D[i-1][0])/(h_r);
				}
			else if (i < N_R)
				{
					dV = Velosity[i][0]*(Velosity[i][0]-Velosity[i-1][0])/(h_r);
					Velosity[i][1] = Velosity[i][0] + h_t*(delta*D[i][0]/e0 - dV);
					R[i] += h_t*Velosity[i][0] + h_t*h_t*(delta*D[i][0]/e0 - dV)/2;
					Density[i] = (D[i][0] - D[i-1][0])/(h_r);
					D[i][1] = D[i][0] - h_t*Density[i]*Velosity[i][0];
				}
			
			if (i <= N_R && Density[i] < 0)
				{Density[i] = 0;}
				
			if (i > 1 && i <= N_R+2) 
				{
					D[i-2][0] = D[i-2][1];
					Velosity[i-2][0] = Velosity[i-2][1];
				}
		}
	""")
	
	st = time()
	dens_gpu = cuda.to_device(dens)
	rad_gpu = cuda.to_device(rad)
	D_gpu = cuda.to_device(Dd)
	Vel_gpu = cuda.to_device(Velosityy)
	h_r_gpu = cuda.to_device(np.float64(h_r))
	h_t_gpu = cuda.to_device(np.float64(h_t))
	N_R_gpu = cuda.to_device(np.int64(N_R))
	delta_gpu = cuda.to_device(np.float64(delta))
	e0_gpu = cuda.to_device(np.float64(e0))

	func = mod.get_function("run")
	for i in range(N_t) :
		func(dens_gpu, rad_gpu, D_gpu, Vel_gpu, h_r_gpu, h_t_gpu, N_R_gpu, delta_gpu, e0_gpu, block=(256, 4, 1))
		
	rad = cuda.from_device_like(rad_gpu, rad)
	dens = cuda.from_device_like(dens_gpu, dens)
	print("GPU: " + str(time()-st))
	return rad, dens
	
	
def main() :
	graph = plt.figure()
	ax = graph.add_subplot(111)
	ax.plot(R, Density)
	st = time()
	iter()
	print("CPU: " + str(time() - st))
	ax.plot(ar, ad)
	ax.plot(R[1:], Density[1:])
	r, d = iter_gpu()
	plt.grid(True)
	plt.show()
	
if __name__ == "__main__" :
	main()