Electromagnetic 1 dimensional PIC simulation

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.

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
import math as mt
from scipy.fftpack import fft, ifft,fft2,fftshift
import time
from tqdm import tqdm
In [2]:
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

Linear waves

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

In [3]:
%%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)
initial condition
  0%|          | 0/4096 [00:00<?, ?it/s]/home/mnakanot/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:65: RuntimeWarning: invalid value encountered in true_divide
100%|██████████| 4096/4096 [01:17<00:00, 53.09it/s]
CPU times: user 1min 15s, sys: 906 ms, total: 1min 16s
Wall time: 1min 17s

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.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()
CPU times: user 6.02 s, sys: 4.12 s, total: 10.1 s
Wall time: 10.3 s

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

In [5]:
%%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)
initial condition
  0%|          | 0/4096 [00:00<?, ?it/s]/home/mnakanot/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:65: RuntimeWarning: invalid value encountered in true_divide
100%|██████████| 4096/4096 [01:38<00:00, 41.58it/s]
CPU times: user 1min 20s, sys: 1.38 s, total: 1min 21s
Wall time: 1min 38s

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.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()
CPU times: user 5.62 s, sys: 3.97 s, total: 9.59 s
Wall time: 10.8 s

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

In [7]:
%%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)
initial condition
  0%|          | 0/4096 [00:00<?, ?it/s]/home/mnakanot/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:65: RuntimeWarning: invalid value encountered in true_divide
100%|██████████| 4096/4096 [01:25<00:00, 47.74it/s]
CPU times: user 1min 17s, sys: 1.2 s, total: 1min 18s
Wall time: 1min 26s
In [8]:
%%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()
CPU times: user 5.41 s, sys: 3.53 s, total: 8.94 s
Wall time: 9.01 s