%config InlineBackend.figure_format = 'retina'
from numba import jit,f8,u8
import numpy as Np
import math as mt
import matplotlib.pyplot as plt
from IPython import display
from scipy.fftpack import fft, ifft,fft2,fftshift
import time
from tqdm import tqdm
Only electrons are treated as super-particles. Ions are fixed. 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$$
def em1ele(vt,b0,th,ppc,nx,nt,nb,vb):
Np.seterr(divide='ignore', invalid='ignore')
#-----
#---wpe=1, c=1, me=1
#-----
np=ppc*nx
dx=vt; dt=dx;
lx=dx*nx
xx=dx*Np.arange(nx)
tt=dt*Np.arange(nt)
xp=Np.zeros((2,np))
vx,vy,vz=Np.zeros((3,2,np))
xp[0,:]=Np.linspace(0,lx-lx/np,np)
xp[1,:]=Np.linspace(0,lx-lx/np,np)
vx[0,:]=Np.random.normal(0,vt,np)-vb*nb/(1-nb)
vx[1,:]=Np.random.normal(0,vt,np)+vb
vy[0,:]=Np.random.normal(0,vt,np)
vy[1,:]=Np.random.normal(0,vt,np)
vz[0,:]=Np.random.normal(0,vt,np)
vz[1,:]=Np.random.normal(0,vt,np)
gm=Np.sqrt(1+vx**2+vy**2+vz**2)
vx=vx/gm; vy=vy/gm; vz=vz/gm
jj=Np.arange(nx); 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
qe=Np.zeros(2);qomedt=Np.zeros(2);qedt=Np.zeros(2)
qe[0]=-nx/np*(1-nb);qomedt[0]=-0.5*dt; qedt[0]=0.25*dt*qe[0]
qe[1]=-nx/np*nb ;qomedt[1]=-0.5*dt; qedt[1]=0.25*dt*qe[1]
ex,ey,ez=Np.zeros((3,nx))
by,bz=Np.zeros((2,nx))
bx0=b0*mt.cos(th*mt.pi/180); bz0=b0*mt.sin(th*mt.pi/180)
eym,eyp =Np.zeros((2,nx))
ezm,ezp =Np.zeros((2,nx))
exsave,eysave,ezsave =Np.zeros((3,nt,nx))
bzsave,bysave,enesave=Np.zeros((3,nt,nx))
print('initial condition')
plt.subplot(121);plt.plot(xp[0,:],vx[0,:],',');plt.plot(xp[1,:],vx[1,:],',')
plt.subplot(122);plt.hist(vx[0,:],bins=500,histtype='step',weights=(1-nb)*Np.ones(np));plt.hist(vx[1,:],bins=500,histtype='step',weights=nb*Np.ones(np))
plt.show()
for it in tqdm(range(nt)):
#-----
#calculate charge density and current on each grid
#-----
jym=Np.zeros((2,nx)); jzm=Np.zeros((2,nx))
jym,jzm=currnt(np,nx,qedt,dx,xp,vy,vz,jym,jzm)
#---calculate new particle posotions
xp=xp+dt*vx
xp[xp>lx]=xp[xp>lx]-lx
xp[xp<0 ]=xp[xp<0 ]+lx
#-----
#calculate charge density and current on each grid
#-----
ds =Np.zeros((2,nx))
jyp=Np.zeros((2,nx));jzp=Np.zeros((2,nx))
ds,jyp,jzp=currntden(np,nx,qe,qedt,dx,xp,vy,vz,jyp,jzp,ds)
#-----
#---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
ex=Np.real(ifft(exfft))
exsave[it,:]=ex
#---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[jj]=eym[jp]
ezm[jj]=ezm[jp]
eyp[jj]=eyp[jm]
ezp[jj]=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)
ey=eyp+eym
bz=eyp-eym
ez=ezp+ezm
by=ezm-ezp
eysave[it,:]=ey
ezsave[it,:]=ez
bysave[it,:]=by
bzsave[it,:]=bz
#-----
#solve eqs. of motion for each particle
#-----
vx,vy,vz,gm=acc(np,nx,dx,qomedt,bx0,bz0,xp,vx,vy,vz,gm,ex,ey,ez,by,bz)
#display.clear_output(True)
#plt.plot(xx,bz)
#plt.plot(xx,ex)
##plt.plot(xp[1,:],vx[1,:],',')
#plt.show()
kmin=2*mt.pi/(dx*nx)*(-nx/2); kmax=2*mt.pi/(dx*nx)*(nx/2)
wmin=2*mt.pi/(dt*nt)*(-nt/2); wmax=2*mt.pi/(dt*nt)*(nt/2)
kaxis=Np.linspace(kmin,kmax,nx); waxis=Np.linspace(wmin,wmax,nt)
return locals()
@jit('f8[:,:](u8,u8,f8,f8[:],f8,f8,f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:],f8[:],f8[:],f8[:],f8[:])')
def acc(np,nx,dx,qomedt,bx0,bz0,xp,vx,vy,vz,gm,ex,ey,ez,by,bz):
for isp in range(2):
for ip in range(np):
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)
exap=qomedt[isp]*(wxm*ex[ixm]+wxp*ex[ixp])
eyap=qomedt[isp]*(wxm*ey[ixm]+wxp*ey[ixp])
ezap=qomedt[isp]*(wxm*ez[ixm]+wxp*ez[ixp])
bxap=qomedt[isp]*bx0
byap=qomedt[isp]*(wxm*by[ixm]+wxp*by[ixp])
bzap=qomedt[isp]*(wxm*bz[ixm]+wxp*bz[ixp]+bz0)
#---half acceleration
gvxs=vx[isp,ip]+exap
gvys=vy[isp,ip]+eyap
gvzs=vz[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
f=2/(1+bxap**2+byap**2+bzap**2)
bxap=f*bxap
byap=f*byap
bzap=f*bzap
gvxs=gvxs+vvy*bzap-vvz*byap
gvys=gvys+vvz*bxap-vvx*bzap
gvzs=gvzs+vvx*byap-vvy*bxap
#---half acceleration
gm[isp,ip]=mt.sqrt(1+gvxs**2+gvys**2+gvzs**2)
vx[isp,ip]=gvxs+exap
vy[isp,ip]=gvys+eyap
vz[isp,ip]=gvzs+ezap
return vx,vy,vz,gm
@jit('f8[:,:](u8,u8,f8[:],f8,f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:])')
def currnt(np,nx,qedt,dx,xp,vy,vz,jym,jzm):
for isp in range(2):
for ip in range(np):
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)
jym[isp,ixm]=jym[isp,ixm]+wxm*vy[isp,ip]
jym[isp,ixp]=jym[isp,ixp]+wxp*vy[isp,ip]
jzm[isp,ixm]=jzm[isp,ixm]+wxm*vz[isp,ip]
jzm[isp,ixp]=jzm[isp,ixp]+wxp*vz[isp,ip]
jym[isp,:]=jym[isp,:]*qedt[isp]; jzm[isp,:]=jzm[isp,:]*qedt[isp]
return jym,jzm
@jit('f8[:,](u8,u8,f8[:],f8[:],f8,f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:],f8[:,:])')
def currntden(np,nx,qe,qedt,dx,xp,vy,vz,jyp,jzp,ds):
for isp in range(2):
for ip in range(np):
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
jyp[isp,ixm]=jyp[isp,ixm]+wxm*vy[isp,ip]
jyp[isp,ixp]=jyp[isp,ixp]+wxp*vy[isp,ip]
jzp[isp,ixm]=jzp[isp,ixm]+wxm*vz[isp,ip]
jzp[isp,ixp]=jzp[isp,ixp]+wxp*vz[isp,ip]
ds[isp,:]=ds[isp,:]*qe[isp]; jyp[isp,:]=jyp[isp,:]*qedt[isp]; jzp[isp,:] = jzp[isp,:]*qedt[isp]
return ds,jyp,jzp
%%time
data=em1ele(0.05,2,0,100,1024,4096,0.5,0) #(vt,b0,th,ppc,nx,nt,nb,vb)
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.001,0.001))
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
data=em1ele(0.05,2,45,100,1024,4096,0.5,0) #(vt,b0,th,ppc,nx,nt,nb,vb)
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.001,0.001))
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
data=em1ele(0.05,2,90,100,1024,4096,0.5,0) #(vt,b0,th,ppc,nx,nt,nb,vb)
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.001,0.001))
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()