In [1]:
%matplotlib inline
import numpy as np
import sympy
import pylab
import random
sympy.init_printing()
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib import animation
from JSAnimation.IPython_display import display_animation
In [2]:
def create_grid(L,d,x_i,n,nt):
""" Creates a discrete grid based on our desired length and step size,
then also initializes our location storage matrix and calculates the number
of grid points in the directional arrays.
Parameters
----------
L : array of floats
Desired distances from origin to the edge of the cube in each direction
d : float
Discrete grid step size
x_i : float
Initial position of particles in cartesian coordinates
n : int
Number of particles
nt : int
Number of time steps
Returns
-------
x,y,z : arrays of floats
Arrays containing all of the discrete grid locations and their
true cartesian value
loc : array of floats
Storage matrix for the (x,y,z) location (0th dimension) for each
particle (1st dimension)
nx,ny,nz : ints
Number of grid points in each direction
X,Y,Z : Coordinate Matrices of floats
3D matrices (size nx,ny,nz) containing all of the discrete grid
points and their respective x, y, or z value. Allows for conversion
from index number to grid value.
"""
x = np.arange(-L[0],L[0],d)
y = np.arange(-L[1],L[1],d)
z = np.arange(-L[2],L[2],d)
X,Y,Z = np.meshgrid(x,y,z)
loc = np.zeros((3,n,nt))
loc[:,:,0] = x_i[:,:]
nx = int(2*L[0]/d)
ny = int(2*L[1]/d)
nz = int(2*L[2]/d)
return x, y, z, loc, nx, ny, nz, X, Y, Z
In [3]:
def PICsimAll(v_i,dt,nt,m,Zi):
""" Takes in the intial conditions of the problem to provide the velocity and
location of every particle for each time step.
Parameters
----------
v_i : Matrix of floats
Matrix that stores the cartesian initial velocity components
(0th dimension) for each particle (1st dimension)
dt : Array of float
Time step for each particle type
nt : int
Number of time steps desired
m : Array of float
Mass of each particle species in the simulation
Zi : float
Average charge number for the ion species
Returns
-------
vn : Matrix of floats
Matrix that stores the cartesian velocity components (0th dimension) for
each particle (1st dimension) throughout all of time (2nd dimension)
loca : Matrix of floats
Matrix that stores the cartesian location (0th dimension) for each
particle (1st dimension) throughout all of time (2nd dimension)
"""
# Create our local holder copies
vn = np.zeros((3,n,nt))
loca = np.zeros_like(vn)
rho = np.zeros_like(vn) # This will use the number density to hold rho
vn[:,:,0] = v_i[:,:]
pot = np.zeros((nx, ny, nz))
N = np.zeros_like(X) # This will hold the number of particles in each index
pd = np.zeros_like(pot) # Our holder matrix for recalculating phi
Exi = np.ones_like(X)*1000 # These are the initial cartesian field holders
Eyi = np.zeros_like(Y)
Ezi = np.zeros_like(Z)
Bxi = np.zeros_like(X)
Byi = np.zeros_like(Y)
Bzi = np.ones_like(Z)*0.1
# To simplify the math we're about to do...
den = 2*(dy**2*dz**2+dx**2*dz**2+dx**2*dy**2)
for i in range(1,nt):
# Get the new velocities
vn[:,:,i] = intEOMAll(vn[:,:,i-1],m,i,dt,Exi,Eyi,Ezi,Bxi,Byi,Bzi,loca[:,:,i])
# Use the old velocities to get the current location
loca[:,:n/2-1,i] = vn[:,:n/2-1,i-1]*dt[0]
loca[:,n/2:,i] = vn[:,n/2:,i-1]*dt[1]
# Use the current locations, find their nearest index
locInd = loca[:,:,i].astype(int)
# Fancy slicing to increment by the effective charge
N[locInd[:,:n/2-1]] += -1.
N[locInd[:,n/2:]] += Zi*1.
# Find rho!
rho = N * e / d
# Calculate the potential for all interior points given the current rho,
# so that you can set the stage for the next time iteration
pot[1:nx-1,1:ny-1,1:nz-1] = ( dx**2*dy**2*dz**2/(eps*den)*\
(rho[1:nx-1,1:ny-1,1:nz-1]) +\
dy**2*dz**2/(den)*(pd[2:nx,1:ny-1,1:nz-1]+pd[0:nx-2,1:ny-1,1:nz-1])+\
dx**2*dz**2/(den)*(pd[1:nx-1,2:ny,1:nz-1]+pd[1:nx-1,0:ny-2,1:nz-1])+\
dx**2*dy**2/(den)*(pd[1:nx-1,1:ny-1,2:nz]+pd[1:nx-1,1:ny-1,0:nz-2]) )
# No boundary conditions specified, eh? Why not Neumann?
pot[0,:,:] = pot[1,:,:]
pot[nx-1,:,:] = pot[nx-2,:,:]
pot[:,0,:] = pot[:,1,:]
pot[:,ny-1,:] = pot[:,ny-2,:]
pot[:,:,0] = pot[0,:,1]
pot[:,:,nz-1] = pot[:,:,nz-2]
# Finally, recalculate the Electric Field for the next time step
Exi[1:nx-1,1:ny-1,1:nz-1] = -(pot[2:nx,1:ny-1,1:nz-1]-\
pot[0:nx-2,1:ny-1,1:nz-1])/(2*dx)
Eyi[1:nx-1,1:ny-1,1:nz-1] = -(pot[1:nx-1,2:ny,1:nz-1]-\
pot[1:nx-1,0:ny-2,1:nz-1])/(2*dy)
Ezi[1:nx-1,1:ny-1,1:nz-1] = -(pot[1:nx-1,1:ny-1,2:nz]-\
pot[1:nx-1,1:ny-1,0:nz-2])/(2*dz)
return vn, loca
def intEOMAll(v,m,t,dt,Exi,Eyi,Ezi,Bxi,Byi,Bzi,loca):
""" Takes in the intial conditions of the problem to provide the velocity at the
next time step
Parameters
----------
v : Matrix of floats
Matrix that stores the cartesian velocity components from the previous
step (0th dimension) for each particle (1st dimension)
m : Array of float
Mass of each particle species in the simulation
t : int
Current time step
dt : Array of float
Time step for each particle type
Exi,
Eyi,
Ezi : 3D Matrix of floats
Cartesian electric field components for each grid point at the current
time step
Bxi,
Byi,
Bzi : 3D Matrix of floats
Cartesian magnetic field components for each grid point at the current
time step
loca : Matrix of floats
Matrix that contains the cartesian location (0th dimension) for each
particle (1st dimension) at the current time step
Returns
-------
v2 : Matrix of floats
Matrix that contains the cartesian velocity components (0th dimension)
for each particle (1st dimension) at the next time step
"""
v2 = np.zeros_like(v)
for i in range(0,n):
# We have to grab the right m, q, and dt values for each species
if i < n/2:
q = -e
mass = m[0]
dthold = dt[0]
else:
q = e
mass = m[1]
dthold = dt[0]
# Find the current E and B field values according to the particle's location
# from the cartesian component matrices for each particle individually.
# Notice that there is no need to save these values between steps, so we
# need only use holder arrays!
Ehold = np.array([Exi[loca[0,i],loca[1,i],loca[2,i]],\
Eyi[loca[0,i],loca[1,i],loca[2,i]],\
Ezi[loca[0,i],loca[1,i],loca[2,i]]])
Bhold = np.array([Bxi[loca[0,i],loca[1,i],loca[2,i]],\
Byi[loca[0,i],loca[1,i],loca[2,i]],\
Bzi[loca[0,i],loca[1,i],loca[2,i]]])
# This is all to build up to the final line, equations 1 & 4 discretized
# and solved for the next velocity:
v2[:,i] = (Ehold + np.cross(v[:,i],Bhold))*dthold*q/mass + v[:,i]
return v2
In [14]:
def animateAll(frame):
x = (loc_all[0,:n/2,frame]-L[0])
y = (loc_all[1,:n/2,frame]-L[1])
z = (loc_all[2,:n/2,frame]-L[2])
ax.scatter(x,y,z,c='b',marker='o')
x = (loc_all[0,n/2:,frame]-L[0])
y = (loc_all[1,n/2:,frame]-L[1])
z = (loc_all[2,n/2:,frame]-L[2])
ax.scatter(x,y,z,c='r',marker='o')
ax.set_xlim((np.ndarray.min(view[0,:,frame]), np.ndarray.max(view[0,:,frame])))
ax.set_ylim((np.ndarray.min(view[0,:,frame]), np.ndarray.max(view[0,:,frame])))
ax.set_zlim((np.ndarray.min(view[0,:,frame]), np.ndarray.max(view[0,:,frame])))
def initAll():
ax.scatter([],[],[],c='b',marker='o')
ax.scatter([],[],[],c='r',marker='o')
In [5]:
def initial_v(n):
""" Creates randomized initial velocities for ions and electrons separately
based on the desired velocity ranges.
Parameters
----------
n : int
Number of particles
Returns
-------
v_i : Matrix of floats
Matrix that stores the cartesian velocity components (0th dimension) for
each particle (1st dimension)
"""
v_i = np.zeros((3,n))
for i in range(0,n/2-1):
v_i[0,i] = random.randrange(1e5,5e5)*random.choice([-1,1]) # [m/s]
v_i[1,i] = random.randrange(1e5,5e5)*random.choice([-1,1]) # [m/s]
v_i[2,i] = random.randrange(1e5,5e5)*random.choice([-1,1]) # [m/s]
for i in range(n/2,n-1):
v_i[0,i] = random.randrange(1e3,5e3)*random.choice([-1,1]) # [m/s]
v_i[1,i] = random.randrange(1e3,5e3)*random.choice([-1,1]) # [m/s]
v_i[2,i] = random.randrange(1e3,5e3)*random.choice([-1,1]) # [m/s]
return v_i
In [24]:
# Initial Conditions
n = 20 # [particles]
dt_e = 2.50e-11 # [s]
dt_i = 5.0e-5 # [s]
dt = np.array([dt_e,dt_i])
nt = 200 # [steps]
m_e = 9.1094e-31 # [kg]
m_He = 1.6605e-27 # [kg]
Zi = 1.0 # [Average ion charge number, He]
m_tot = np.array([m_e,m_He])
e = 1.609e-19 # [C]
eps = 8.854e-12 # [F/m]
# Grid Info
dx = dy = dz = d = 5.0e-5 # [m]
L = 75.*np.array([d, d, d]) # [m]
x_i = np.zeros((3,n)) # [m]
x, y, z, loc, nx, ny, nz, X, Y, Z = create_grid(L,d,x_i,n,nt)
# Begin Math
v_i_tot = initial_v(n)
vt_all,loc_all = PICsimAll(v_i_tot,dt,nt,m_tot,Zi)
In [21]:
# Animate All
fig = pylab.figure(figsize=(8,6));
ax = Axes3D(fig)
ax.view_init(45,45)
view = loc_all
ax.set_xlim((np.ndarray.min(view[0,:,0]), np.ndarray.max(view[0,:,0])))
ax.set_ylim((np.ndarray.min(view[1,:,0]), np.ndarray.max(view[1,:,0])))
ax.set_zlim((np.ndarray.min(view[2,:,0]), np.ndarray.max(view[2,:,0])))
ax.set_xlabel(r'X Position [m]')
ax.set_ylabel(r'Y Position [m]')
ax.set_zlabel(r'Z Position [m]')
ax.set_title(r'20 Particle Simulation for Problem 3')
anim = animation.FuncAnimation(fig, animateAll, nt,\
init_func = initAll, interval=50)
display_animation(anim, default_mode='loop')
Out[21]: