In [ ]:
%matplotlib inline

In [ ]:
import numpy as np
import pyopencl as cl
import pyopencl.array
from scipy.stats import maxwell, norm
import matplotlib.pyplot as plt

In [ ]:
import sys
sys.path.append("/home/user/dev/pyes1/")
from pyes1.pic1d2v import poisson_solve

In [ ]:
ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

In [ ]:
with open("pic.cl") as f:
    code = f.read()

prg = cl.Program(ctx, code).build()

normalize = prg.normalize_particles
normalize.set_scalar_arg_dtypes([None, np.int32, np.double])

weight = prg.weight_CIC
weight.set_scalar_arg_dtypes([None, np.int32, np.double, None, None, np.int32])

calc_E = prg.calc_E
calc_E.set_scalar_arg_dtypes([None, np.int32, np.double, None])

move = prg.move
move.set_scalar_arg_dtypes([None, None, np.double, np.int32])


accel = prg.accel
accel.set_scalar_arg_dtypes([None, None, None, np.double, np.int32])

interp = prg.interp_CIC
interp.set_scalar_arg_dtypes([None, np.int32, np.double,
                             None, None, np.int32])

In [ ]:
nt = 400
nx = 512
L  = 2*np.pi*2
dx = L/nx
dt = 0.1

Ne = nx*1000
N = Ne+Ne

n0e  = np.linspace(0., L, Ne+1)[:-1]
np.random.seed(12345+2)

# vx1 =  maxwell.rvs(size=Ne)
# vx2 = -maxwell.rvs(size=Ne)

vx1 = norm.rvs(size=Ne, scale=1.0/np.sqrt(2))+2.0
vx2 = norm.rvs(size=Ne, scale=1.0/np.sqrt(2))-2.0

k = 1.0*2*np.pi/L
x01 = n0e-0.1*np.sin(k*n0e)
x02 = n0e-0.1*np.sin(k*n0e)

xp = np.zeros(N)
xp[:len(x01)] = x01
xp[len(x01):] = x02

vp = np.zeros(len(vx1)+len(vx2))
vp[:len(vx1)] = vx1
vp[len(vx2):] = vx2

wp  = 1.0
qom = 1.0

n0 = (Ne+Ne)/L
q  = wp*wp/(n0*qom)

In [ ]:
mf = cl.mem_flags

xp_d = cl.array.Array(queue, N, np.double)
normalize(queue, (4,), None, xp_d.data, N, L)
vp_d = cl.array.Array(queue, N, np.double)
Ep_d = cl.array.Array(queue, N, np.double)
q_d  = cl.array.Array(queue, N, np.double)

xp_d.set(xp)
vp_d.set(vp)
q_d[:] = q

rho_d = cl.array.Array(queue, nx, np.double)
phi_d = cl.array.Array(queue, nx, np.double)
E_d   = cl.array.Array(queue, nx, np.double)

In [ ]:
for i in range(nt):
    
    rho_d[:] = 0.0
    weight(queue, (1,), None, rho_d.data, nx, dx, xp_d.data, q_d.data, N)
    
    rho = rho_d.get(queue)/dx
    phi = poisson_solve(rho, dx)
    phi_d.set(np.ascontiguousarray(phi))

    calc_E(queue, (4,), None, phi_d.data, nx, dx, E_d.data)

    interp(queue, (4,), None, E_d.data, nx, dx, xp_d.data,
          Ep_d.data, N)

    accel(queue, (4,), None, Ep_d.data, q_d.data, vp_d.data, 
         dt, N)
    move(queue, (4,), None, xp_d.data, vp_d.data, dt, N)

    normalize(queue, (4,), None, xp_d.data, N, L)

In [ ]:
queue.finish()

In [ ]:
x = xp_d.get()
v = vp_d.get()

k = 50
s = 1
plt.scatter(x[:Ne:k], v[:Ne:k], color='r', s=s)
plt.scatter(x[Ne::k], v[Ne::k], color='b', s=s)

In [ ]: