Electrons and ions are treated as super-particles. Solving the following equations, $$\frac{dx_j}{dt}=v_{xj}$$ $$\frac{d{\bf v}_j}{dt}=\frac{q_j}{m_j}({\bf E}+{\bf v}_j\times{\bf B})$$ $$\nabla\times{\bf E}=-\frac{\partial {\bf B}}{\partial t}$$ $$\nabla\times{\bf B}=\mu_0{\bf J}+\epsilon_0\mu_0\frac{\partial {\bf E}}{\partial t}$$ $$\nabla\cdot{\bf E}=\frac{\rho}{\epsilon_0}$$ $$\nabla\cdot{\bf B}=0$$
Optimized by the jit compiler of Numba.
%config InlineBackend.figure_format = 'retina'
from numba import jit,f8,u8
import numpy as np
import math as mt
import matplotlib.pyplot as plt
import math as mt
from scipy.fftpack import fft, ifft,fft2,fftshift
import time
from tqdm import tqdm
def em1(nsp,rm,b0,th,vt,ppc,nx,nt,iphs):
#-----
#---wpe=1, c=1, me=1
#-----
nop=ppc*nx
dx=vt; dt=dx;lx=dx*nx
bx0=b0*mt.cos(th*mt.pi/180); bz0=b0*mt.sin(th*mt.pi/180)
xx=dx*np.arange(nx); tt=dt*np.arange(nt)
xp=np.zeros((nsp,nop))
vp=np.zeros((nsp,3,nop))
for ip in range(nsp):
xp[ip,:]=np.linspace(0,lx-lx/nop,nop)
if ip>0:
vp[ip,:,:]=np.random.randn(3,nop)*vt/mt.sqrt(rm)
else:
vp[ip,:,:]=np.random.randn(3,nop)*vt
gm=np.sqrt(1+vp[:,0,:]**2+vp[:,1,:]**2+vp[:,2,:]**2)
vp[:,0,:]=vp[:,0,:]/gm; vp[:,1,:]=vp[:,1,:]/gm; vp[:,2,:]=vp[:,2,:]/gm
jp=np.r_[np.arange(nx-1)+1, 0]; jm=np.r_[nx-1, np.arange(nx-1)]
kk=2*mt.pi/lx*np.r_[np.arange(nx/2),np.arange(-nx/2,0)]
Kl=(np.sin(0.5*kk*dx)/(0.5*dx))**2
kl=np.sin(kk*dx)/dx
q,qomdt,qdt=np.zeros((3,nsp))
q[0]=-nx/nop; qomdt[0]=-0.5*dt; qdt[0]=0.25*dt*q[0]
q[1]=-q[0]; qomdt[1]= 0.5*dt/rm; qdt[1]=0.25*dt*q[1]
ef=np.zeros((3,nx))
bf=np.zeros((3,nx))
eym,eyp =np.zeros((2,nx))
ezm,ezp =np.zeros((2,nx))
bf[0,:]=bx0; bf[2,:]=bz0
exsave,eysave,ezsave =np.zeros((3,nt//iphs,nx))
bzsave,bysave,enesave=np.zeros((3,nt//iphs,nx))
print('initial condition')
plt.subplot(121);plt.plot(xp[0,:],vp[0,0,:],',');plt.plot(xp[1,:],vp[1,0,:],',')
plt.subplot(122);plt.hist(vp[0,0,:],bins=500,histtype='step');plt.hist(vp[1,0,:],bins=500,histtype='step')
plt.show()
for it in tqdm(range(nt)):
#---solve eqs. of motion for each particle
vp,gm=acc(nsp,nop,nx,dx,qomdt,xp,vp,gm,bf,ef)
#enesave[it,5+isp]=0.5*enesave[it,5+isp]/np
#---calculate current on each grid
jym,jzm=currnt(nsp,nop,nx,dx,qdt,xp,vp)
#---calculate new particle posotions
xp=xp+dt*vp[:,0,:]
xp[xp>lx]=xp[xp>lx]-lx
xp[xp<0 ]=xp[xp<0 ]+lx
##---calculate charge density and current---#
jyp,jzp=currnt(nsp,nop,nx,dx,qdt,xp,vp)
ds =density(nsp,nop,nx,dx,q,xp)
#---solve Maxwell's equations on each grid
#---longitudinal
# exfft=-1j/kk*fft(Np.sum(ds,0))
exfft=-1j*kl*fft(np.sum(ds,0))/Kl
exfft[0]=0#; exfft[Np.int(nx/2)]=0#exfft[nx-1]=0
ef[0,:]=np.real(ifft(exfft))
#---transverse
eym=eym-np.sum(jym,0)
eyp=eyp-np.sum(jym,0)
ezm=ezm-np.sum(jzm,0)
ezp=ezp-np.sum(jzm,0)
eym=eym[jp]
ezm=ezm[jp]
eyp=eyp[jm]
ezp=ezp[jm]
eym=eym-np.sum(jyp,0)
eyp=eyp-np.sum(jyp,0)
ezm=ezm-np.sum(jzp,0)
ezp=ezp-np.sum(jzp,0)
ef[1,:]=eyp+eym
bf[2,:]=eyp-eym
ef[2,:]=ezp+ezm
bf[1,:]=ezm-ezp
if it%iphs==0:
exsave[it//iphs,:]=ef[0,:]
eysave[it//iphs,:]=ef[1,:]
ezsave[it//iphs,:]=ef[2,:]
bysave[it//iphs,:]=bf[1,:]
bzsave[it//iphs,:]=bf[2,:]
kmin=2*mt.pi/(dx*nx)*(-nx/2); kmax=2*mt.pi/(dx*nx)*(nx/2)
wmin=2*mt.pi/(dt*iphs*nt)*(-nt/2); wmax=2*mt.pi/(dt*iphs*nt)*(nt/2)
kaxis=np.linspace(kmin,kmax,nx); waxis=np.linspace(wmin,wmax,nt//iphs)
return locals()
@jit('(u8,u8,u8,f8,f8[:],f8[:,:],f8[:,:,:],f8[:,:],f8[:,:],f8[:,:])',nopython=True)
def acc(nsp,nop,nx,dx,qomdt,xp,vp,gm,bf,ef):
for isp in range(nsp):
for ip in range(nop):
ixm=int(xp[isp,ip]/dx); ixp=ixm+1
wxp=xp[isp,ip]/dx-ixm; wxm=1-wxp
if ixp>nx-1: ixp=ixp-nx
exap=qomdt[isp]*(wxm*ef[0,ixm]+wxp*ef[0,ixp])
eyap=qomdt[isp]*(wxm*ef[1,ixm]+wxp*ef[1,ixp])
ezap=qomdt[isp]*(wxm*ef[2,ixm]+wxp*ef[2,ixp])
bxap=qomdt[isp]*(wxm*bf[0,ixm]+wxp*bf[0,ixp])
byap=qomdt[isp]*(wxm*bf[1,ixm]+wxp*bf[1,ixp])
bzap=qomdt[isp]*(wxm*bf[2,ixm]+wxp*bf[2,ixp])
#---half acceleration
gvxs=vp[isp,0,ip]*gm[isp,ip]+exap
gvys=vp[isp,1,ip]*gm[isp,ip]+eyap
gvzs=vp[isp,2,ip]*gm[isp,ip]+ezap
#---first half rotation
gmm=mt.sqrt(1+gvxs**2+gvys**2+gvzs**2)
bxap=bxap/gmm
byap=byap/gmm
bzap=bzap/gmm
vvx=gvxs+gvys*bzap-gvzs*byap
vvy=gvys+gvzs*bxap-gvxs*bzap
vvz=gvzs+gvxs*byap-gvys*bxap
#---second half rotation + second half acceleration
f=2/(1+bxap**2+byap**2+bzap**2)
bxap=f*bxap
byap=f*byap
bzap=f*bzap
gvxs=gvxs+vvy*bzap-vvz*byap+exap
gvys=gvys+vvz*bxap-vvx*bzap+eyap
gvzs=gvzs+vvx*byap-vvy*bxap+ezap
#---half acceleration
gm[isp,ip]=mt.sqrt(1+gvxs**2+gvys**2+gvzs**2)
vp[isp,0,ip]=gvxs/gm[isp,ip]
vp[isp,1,ip]=gvys/gm[isp,ip]
vp[isp,2,ip]=gvzs/gm[isp,ip]
return vp,gm
@jit('(u8,u8,u8,f8,f8[:],f8[:,:],f8[:,:,:])',nopython=True)
def currnt(nsp,nop,nx,dx,qdt,xp,vp):
jy=np.zeros((nsp,nx)); jz=np.zeros((nsp,nx))
for isp in range(nsp):
for ip in range(nop):
ixm=mt.floor(xp[isp,ip]/dx); ixp=ixm+1
wxp=xp[isp,ip]/dx-ixm; wxm=1-wxp
if ixp>nx-1: ixp=ixp-nx
jy[isp,ixm]=jy[isp,ixm]+wxm*vp[isp,1,ip]
jy[isp,ixp]=jy[isp,ixp]+wxp*vp[isp,1,ip]
jz[isp,ixm]=jz[isp,ixm]+wxm*vp[isp,2,ip]
jz[isp,ixp]=jz[isp,ixp]+wxp*vp[isp,2,ip]
jy[isp,:]=jy[isp,:]*qdt[isp]
jz[isp,:]=jz[isp,:]*qdt[isp]
return jy,jz
@jit('f8[:,:](u8,u8,u8,f8,f8[:],f8[:,:])',nopython=True)
def density(nsp,nop,nx,dx,q,xp):
ds=np.zeros((nsp,nx))
for isp in range(nsp):
for ip in range(nop):
#print(xp[ip])
ixm=mt.floor(xp[isp,ip]/dx); ixp=ixm+1
wxp=xp[isp,ip]/dx-ixm; wxm=1-wxp
if ixp>nx-1: ixp=ixp-nx#; print(ixp)
ds[isp,ixm]=ds[isp,ixm]+wxm
ds[isp,ixp]=ds[isp,ixp]+wxp
ds[isp,:] =ds[isp,:]*q[isp]
return ds
%%time
nsp=2; rm=4; b0=2; th=0; vt=0.05; ppc=100; nx=1024; nt=4096; iphs=1
data=em1(nsp,rm,b0,th,vt,ppc,nx,nt,iphs)
locals().update(data)
%%time
fftex=fftshift(fft2(exsave[:,::-1]))
fftey=fftshift(fft2(eysave[:,::-1]))
fftez=fftshift(fft2(ezsave[:,::-1]))
fftby=fftshift(fft2(bysave[:,::-1]))
fftbz=fftshift(fft2(bzsave[:,::-1]))
plt.figure(figsize=(30,20))
plt.subplot(441);plt.imshow(exsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.01,0.01))
plt.subplot(445);plt.imshow(eysave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(449);plt.imshow(ezsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(447);plt.imshow(bysave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(4,4,11);plt.imshow(bzsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(442);plt.imshow(abs(fftex),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(446);plt.imshow(abs(fftey),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,10);plt.imshow(abs(fftez),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,8);plt.imshow(abs(fftby),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,12);plt.imshow(abs(fftbz),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.show()
%%time
nsp=2; rm=4; b0=2; th=45; vt=0.05; ppc=100; nx=1024; nt=4096; iphs=1
data=em1(nsp,rm,b0,th,vt,ppc,nx,nt,iphs)
locals().update(data)
%%time
fftex=fftshift(fft2(exsave[:,::-1]))
fftey=fftshift(fft2(eysave[:,::-1]))
fftez=fftshift(fft2(ezsave[:,::-1]))
fftby=fftshift(fft2(bysave[:,::-1]))
fftbz=fftshift(fft2(bzsave[:,::-1]))
plt.figure(figsize=(30,20))
plt.subplot(441);plt.imshow(exsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.01,0.01))
plt.subplot(445);plt.imshow(eysave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(449);plt.imshow(ezsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(447);plt.imshow(bysave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(4,4,11);plt.imshow(bzsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(442);plt.imshow(abs(fftex),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(446);plt.imshow(abs(fftey),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,10);plt.imshow(abs(fftez),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,8);plt.imshow(abs(fftby),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,12);plt.imshow(abs(fftbz),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.show()
%%time
nsp=2; rm=4; b0=2; th=90; vt=0.05; ppc=100; nx=1024; nt=4096; iphs=1
data=em1(nsp,rm,b0,th,vt,ppc,nx,nt,iphs)
locals().update(data)
%%time
fftex=fftshift(fft2(exsave[:,::-1]))
fftey=fftshift(fft2(eysave[:,::-1]))
fftez=fftshift(fft2(ezsave[:,::-1]))
fftby=fftshift(fft2(bysave[:,::-1]))
fftbz=fftshift(fft2(bzsave[:,::-1]))
plt.figure(figsize=(30,20))
plt.subplot(441);plt.imshow(exsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.01,0.01))
plt.subplot(445);plt.imshow(eysave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(449);plt.imshow(ezsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(447);plt.imshow(bysave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(4,4,11);plt.imshow(bzsave,origin='lower',cmap='bwr',aspect='auto',clim=(-0.001,0.001))
plt.subplot(442);plt.imshow(abs(fftex),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(446);plt.imshow(abs(fftey),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,10);plt.imshow(abs(fftez),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,8);plt.imshow(abs(fftby),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.subplot(4,4,12);plt.imshow(abs(fftbz),origin='lower',aspect='auto',extent=([kmin,kmax,wmin,wmax]))
plt.axis([0,kmax/16,0,wmax/16])
plt.clim(0,50)
plt.show()