In [1]:
%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

Electron Electromagnetic 1-D PIC simulation

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$$

In [2]:
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

Linear waves

note: If you do not consider a beam ($v_b=0$), you have to set $n_b$ finite.

Parameters: $\omega_{ce}/\omega_{pe}=2$, $\theta=0^\circ$

In [3]:
%%time
data=em1ele(0.05,2,0,100,1024,4096,0.5,0) #(vt,b0,th,ppc,nx,nt,nb,vb)
locals().update(data)
initial condition
100%|██████████| 4096/4096 [01:21<00:00, 50.28it/s]
CPU times: user 1min 19s, sys: 953 ms, total: 1min 20s
Wall time: 1min 22s

Space time evolution of $\vec{E}, \vec{B}$ and the dispersion relation
In [4]:
%%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()
CPU times: user 10.4 s, sys: 5.42 s, total: 15.8 s
Wall time: 16 s

Parameters: $\omega_{ce}/\omega_{pe}=2$, $\theta=45^\circ$

In [5]:
%%time
data=em1ele(0.05,2,45,100,1024,4096,0.5,0) #(vt,b0,th,ppc,nx,nt,nb,vb)
locals().update(data)
initial condition
100%|██████████| 4096/4096 [01:18<00:00, 52.28it/s]
CPU times: user 1min 16s, sys: 844 ms, total: 1min 17s
Wall time: 1min 19s

Space time evolution of $\vec{E}, \vec{B}$ and the dispersion relation
In [6]:
%%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()
CPU times: user 12.3 s, sys: 6.03 s, total: 18.4 s
Wall time: 20.2 s

Parameters: $\omega_{ce}/\omega_{pe}=2$, $\theta=90^\circ$

In [7]:
%%time
data=em1ele(0.05,2,90,100,1024,4096,0.5,0) #(vt,b0,th,ppc,nx,nt,nb,vb)
locals().update(data)
initial condition
100%|██████████| 4096/4096 [01:18<00:00, 52.22it/s]
CPU times: user 1min 15s, sys: 922 ms, total: 1min 16s
Wall time: 1min 19s
Space time evolution of $\vec{E}, \vec{B}$ and the dispersion relation
In [ ]:
%%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()
CPU times: user 10.5 s, sys: 5.98 s, total: 16.4 s
Wall time: 16.6 s