Page 2 - Optimizing PIC (Particle in Cell) Code

import numpy
from numpy import zeros

def calc_density(position, ncells, L):
  
    density = zeros([ncells])
    nparticles = len(position)
	
    inversedx = ncells / L
    p = position*(inversedx)
    plower = numpy.array(p, dtype = numpy.integer)
    offset = p - plower 
    density = zeros([ncells + 1])
    rem = 1 - offset
    for (low, rem, ud) in zip(plower, rem, offset): 
        density[low] += rem
        density[low + 1] += ud
    density[0] += density[-1]
    density = density[:-1]
    
    density *= float(ncells) / float(nparticles) 
    return density

from distutils.core import setup
from distutils.extension import Extension
from Cython.Distutils import build_ext

setup(
    cmdclass = {'build_ext': build_ext},
    ext_modules = [Extension("density_calc", ["density_calc.pyx"])]
)

#!/usr/bin/env python
#
# Electrostatic PIC code in a 1D cyclic domain

import numpy 

from numpy import arange, concatenate, zeros, linspace, floor, array, pi
from numpy import diff, sign, sin, cos, sqrt, random, histogram
from numpy import *

from numpy import histogram

from scipy.optimize import curve_fit

import matplotlib.pyplot as plt # Matplotlib plotting library

try:
    import matplotlib.gridspec as gridspec  # For plot layout grid
    got_gridspec = True
except:
    got_gridspec = False

# Need an FFT routine, either from SciPy or NumPy
try:
    from scipy.fftpack import fft, ifft, integrate #added the integrate function
except:
    # No SciPy FFT routine. Import NumPy routine instead
    from numpy.fft import fft, ifft

def rk4step(f, y0, dt, args=()):
    """ Takes a single step using RK4 method """
    k1 = f(y0, *args)
    k2 = f(y0 + 0.5*dt*k1, *args)
    k3 = f(y0 + 0.5*dt*k2, *args)
    k4 = f(y0 + dt*k3, *args)

    return y0 + (k1 + 2.*k2 + 2.*k3 + k4)*dt / 6.

def calc_density(position, ncells, L):
    """ Calculate charge density given particle positions
    
    Input
      position  - Array of positions, one for each particle
                  assumed to be between 0 and L
      ncells    - Number of cells
      L         - Length of the domain

    Output
      density   - contains 1 if evenly distributed
    """
    # This is a crude method and could be made more efficient
    
    density_calc = zeros([ncells])
    nparticles = len(position)
	
    #Introducing a more efficient way of calculating the charge density of many particles using array calculations

    #dx = L / ncells       # Uniform cell spacing, L and ncells are values
	inversedx = ncells / L
    p = position*(inversedx)
    plower = numpy.array(p, dtype = numpy.integer)
    offset = p - plower #offset or remainder from the left of the cell
    density = zeros([ncells + 1])
    rem = 1 - offset
    for (low, rem, ud) in zip(plower, rem, offset): #tuples
        density[low] += rem
        density[low + 1] += ud
    density[0] += density[-1]
    density = density[:-1]
    
    # nparticles now distributed amongst ncells
    density *= float(ncells) / float(nparticles)  # Make average density equal to 1
    return density

def periodic_interp(y, x):
    """
    Linear interpolation of a periodic array y at index x
    
    Input

    y - Array of values to be interpolated
    x - Index where result required. Can be an array of values
    
    Output
    
    y[x] with non-integer x
    """
    ny = len(y)
    if len(x) > 1:
        y = array(y) # Make sure it's a NumPy array for array indexing
    xl = floor(x).astype(int) # Left index
    dx = x - xl
    xl = ((xl % ny) + ny) % ny  # Ensures between 0 and ny-1 inclusive
    return y[xl]*(1. - dx) + y[(xl+1)%ny]*dx

def fft_integrate(y):
    """ Integrate a periodic function using FFTs
    """
    n = len(y) # Get the length of y
    
    f = fft(y) # Take FFT
    # Result is in standard layout with positive frequencies first then negative
    # n even: [ f(0), f(1), ... f(n/2), f(1-n/2) ... f(-1) ]
    # n odd:  [ f(0), f(1), ... f((n-1)/2), f(-(n-1)/2) ... f(-1) ]
    
    if n % 2 == 0: # If an even number of points
        k = concatenate( (arange(0, n/2+1), arange(1-n/2, 0)) )
    else:
        k = concatenate( (arange(0, (n-1)/2+1), arange( -(n-1)/2, 0)) )
    k = 2.*pi*k/n
    
    # Modify frequencies by dividing by ik
    f[1:] /= (1j * k[1:]) 
    f[0] = 0. # Set the arbitrary zero-frequency term to zero
    
    return ifft(f).real # Reverse Fourier Transform
   

def pic(f, ncells, L):
    """ f contains the position and velocity of all particles
    """
    nparticles = len(f)/2     # Two values for each particle
    pos = f[0:nparticles] # Position of each particle
    vel = f[nparticles:]      # Velocity of each particle

    dx = L / float(ncells)    # Cell spacing

    # Ensure that pos is between 0 and L
    pos = ((pos % L) + L) % L
    
    # Calculate number density, normalised so 1 when uniform
    density = calc_density(pos, ncells, L)
    
    # Subtract ion density to get total charge density
    rho = density_calc - 1.
    
    # Calculate electric field
    E = -fft_integrate(rho)*dx
    
    # Interpolate E field at particle locations
    accel = -periodic_interp(E, pos/dx)

    # Put back into a single array
    return concatenate( (vel, accel) )

####################################################################

def run(pos, vel, L, ncells=None, out=[], output_times=linspace(0,20,100), cfl=0.5):
    
    if ncells == None:
        ncells = int(sqrt(len(pos))) # A sensible default

    dx = L / float(ncells)
    
    f = concatenate( (pos, vel) )   # Starting state
    nparticles = len(pos)
    
    time = 0.0
    for tnext in output_times:
        # Advance to tnext
        stepping = True
        while stepping:
            # Maximum distance a particle can move is one cell
            dt = cfl * dx / max(abs(vel))
            if time + dt >= tnext:
                # Next time will hit or exceed required output time
                stepping = False
                dt = tnext - time
            #print "Time: ", time, dt
            f = rk4step(pic, f, dt, args=(ncells, L))
            time += dt
            
        # Extract position and velocities
        pos = ((f[0:nparticles] % L) + L) % L
        vel = f[nparticles:]
        
        # Send to output functions
        for func in out:
            func(pos, vel, ncells, L, time)
        
    return pos, vel

####################################################################
# 
# Output functions and classes
#

class Plot:
    """
    Displays three plots: phase space, charge density, and velocity distribution
    """
    def __init__(self, pos, vel, ncells, L):
        
        d = calc_density(pos, ncells, L)
        vhist, bins  = histogram(vel, int(sqrt(len(vel))))
        vbins = 0.5*(bins[1:]+bins[:-1])
        
        # Plot initial positions
        if got_gridspec:
            self.fig = plt.figure()
            self.gs = gridspec.GridSpec(4, 4)
            ax = self.fig.add_subplot(self.gs[0:3,0:3])
            self.phase_plot = ax.plot(pos, vel, '.')[0]
            ax.set_title("Phase space")
            
            ax = self.fig.add_subplot(self.gs[3,0:3])
            self.density_plot = ax.plot(linspace(0, L, ncells), d)[0]
            
            ax = self.fig.add_subplot(self.gs[0:3,3])
            self.vel_plot = ax.plot(vhist, vbins)[0]
        else:
            self.fig = plt.figure()
            self.phase_plot = plt.plot(pos, vel, '.')[0]
            
            self.fig = plt.figure()
            self.density_plot = plt.plot(linspace(0, L, ncells), d)[0]
            
            self.fig = plt.figure()
            self.vel_plot = plt.plot(vhist, vbins)[0]
        plt.ion()
        plt.show()
        
    def __call__(self, pos, vel, ncells, L, t):
        d = calc_density(pos, ncells, L)
        vhist, bins  = histogram(vel, int(sqrt(len(vel))))
        vbins = 0.5*(bins[1:]+bins[:-1])
        
        self.phase_plot.set_data(pos, vel) # Update the plot
        self.density_plot.set_data(linspace(0, L, ncells), d)
        self.vel_plot.set_data(vhist, vbins)
        plt.draw()
        

class Summary:
    def __init__(self):
        self.t = []
        self.firstharmonic = []
        
    def __call__(self, pos, vel, ncells, L, t):
        # Calculate the charge density
        d = calc_density(pos, ncells, L)
        
        # Amplitude of the first harmonic
        fh = 2.*abs(fft(d)[1]) / float(ncells)
        
        print "Time:", t, "First:", fh
        
        self.t.append(t)
        self.firstharmonic.append(fh)

####################################################################
# 
# Functions to create the initial conditions
#

def landau(npart, L, alpha=0.2):
    """
    Creates the initial conditions for Landau damping
    
    """
    # Start with a uniform distribution of positions
    pos = random.uniform(0., L, npart) #random pos
    pos0 = pos.copy()
    k = 2.*pi / L
    for i in range(10): # Adjust distribution using Newton iterations
        pos -= ( pos + alpha*sin(k*pos)/k - pos0 ) / ( 1. + alpha*cos(k*pos) )
        
    # Normal velocity distribution
    vel = random.normal(0.0, 1.0, npart) #random vel
    
    return pos, vel

def twostream(npart, L, vbeam=2):
    # Start with a uniform distribution of positions
    pos = random.uniform(0., L, npart) #random pos
    # Normal velocity distribution
    vel = random.normal(0.0, 1.0, npart) #random vel
    
    np2 = int(npart / 2)
    vel[:np2] += vbeam  # Half the particles moving one way
    vel[np2:] -= vbeam  # and half the other
    
    return pos,vel

####################################################################

if __name__ == "__main__":
    # Generate initial condition
    #
    if False:
        # 2-stream instability
        L = 100
        ncells = 20
        pos, vel = twostream(10000, L, 3.)
    else:
        # Landau damping
        L = 4.*pi
        ncells = 20
        npart = input("How many Particles? ") #Requires the users input in order to aid repeat runs.
        pos, vel = landau(npart, L)
		
		
    
    # Create some output classes
    #p = Plot(pos, vel, ncells, L) # This displays an animated figure
    s = Summary()                 # Calculates, stores and prints summary info
   
	
    # Run the simulation
    pos, vel = run(pos, vel, L, ncells, 
                   out=[s],    #Animation off
				   #out=[p, s],            #Animation On          # These are called each output
                   output_times=linspace(0.,40,100)) # The times to output
  
    # Summary stores an array of the first-harmonic amplitude
    # Make a semilog plot to see exponential damping
	
    x = s.t
    data = s.firstharmonic
	
	#Local minimum and local maximum
    b = (diff(sign(diff(data))) > 0).nonzero()[0] + 1 # local min
    c = (diff(sign(diff(data))) < 0).nonzero()[0] + 1 # local max
	
    timemin = []
    minh = []
	
    for j in b:
	    timemin.append(x[j])
	    minh.append(data[j])
		
    derivListmin = [] #Finding the change in gradients between the minimum peaks of the Harmonics
	
    for j in range(len(timemin)-1):
        derivmin = (minh[j]-minh[j+1])/(timemin[j]-timemin[j+1])
        derivListmin.append(derivmin)
		
    #print derivListmin
	
    for indexmin, valuemin in enumerate(derivListmin):
	    if valuemin > 0: #j refers to each element in derivlist
		   #print indexmin, valuemin
		   variablemin = indexmin 
		   break
		   
    #print "timemin =", timemin[variablemin]
    #print "minh =", minh[variablemin]
		   
    higherindexminh = []
    for j in range(len(minh)):
		if j > variablemin:
		   higherindexminh.append(minh[j])
	
    #print "higherindexminh = ", higherindexminh
	
    noiselevelmin = mean(higherindexminh)
	
    #print noiselevelmin
	
	
    timemax = []
    maxh = []
    for i in c:
        timemax.append(x[i])
        maxh.append(data[i])
		
	
    derivListmax = [] #Finding the change in gradients between the maximum peaks of the Harmonics
	
    for i in range(len(timemax)-1):
        derivmax = (maxh[i]-maxh[i+1])/(timemax[i]-timemax[i+1])
        derivListmax.append(derivmax) #This is the gradient
		
    #print derivListmax
	
	#Index is the position of the gradient in the list of derivList
    for indexmax, valuemax in enumerate(derivListmax):
	    if valuemax > 0: #For when the derivative between the peaks first becomes postive, print all postive values after it
		   #print indexmax, valuemax
		   variablemax = indexmax
		   break
		   
    #print "timemax =", timemax[variablemax]
    #print "maxh =", maxh[variablemax]
		   
    higherindexmaxh = [] #creating a new list that includes every peak that follows when the derivative first becomes positive
	
    for i in range(len(maxh)):
		if i > variablemax: #An index/position of higher value than when the derivative first becomes positive
		   higherindexmaxh.append(maxh[i]) #Fill the list with such values
	
    #print "higherindexmaxh = ", higherindexmaxh
	
    noiselevelmax = mean(higherindexmaxh) #Average all the values in the list to imitate a constant noise.
	
    #print noiselevelmax
	
    noiselevelavg = (noiselevelmax + noiselevelmin)/2

    file = open('c:/Users/Pinky/Programming/PIC.txt','a')
    file.write((str(noiselevelavg) + '   ' + str(npart)) + "\n")
    file.close()
	
	# file = open('c:/Users/Pinky/Programming/PIC.txt','r')
	#file.read()
	#file.close()
	
    noiselist = []
    particlelist = []
	
    for line in open('c:/Users/Pinky/Programming/PIC.txt', 'r'):
        words = line.split() #This will split the rows in the file into 2 list by spaces
        noiselist.append(words[0])
        particlelist.append(words[1])
		
    file.close()
		
    print noiselist
	
    """
	Investigating a way of calculating the accuracy of the noise level in conjuction
	with finding the period of the Laundau damping.  Will superimpose a sinusoidal
	function which will exponentially damp itself over time.  Subtracting this function away
	from the original Laundau function will give a residue - take this as noise
    """
    def func(x, b, k, a):
	    return (b*abs(cos(k*x))*exp(-x*a))
	
    x = linspace(0, 40, 100)
    popt, pcov = curve_fit(func, x, data, [0.2, 1.4, 0.1])
    print "params = ", popt
    fitdata = func(x, popt[0], popt[1], popt[2])
	
    residue = ((s.firstharmonic - fitdata))
	
    residueavg = mean(residue)
	
    plt.plot(particlelist, noiselist, 'or')
    plt.xlabel("Noise")
    plt.ylabel("Number of Particles")
	
    plt.figure()
    plt.plot(s.t, s.firstharmonic, label='Harmonics')
    plt.plot([0,40],[residueavg,residueavg],"-o", label='Residue Noise')
    plt.xlabel("Time [Normalised]")
    plt.ylabel("First harmonic amplitude [Normalised]")
    plt.yscale('log')
    plt.plot(timemax, maxh, 'or')
    plt.plot(timemin, minh, 'oy')
    plt.plot([0,40],[noiselevelavg, noiselevelavg], label='Average Noise')
    plt.plot(x, fitdata)
    plt.legend()
	
	#plt.plot([0,40],[noiselevelmax, noiselevelmax])
    #plt.plot([0,40],[noiselevelmin, noiselevelmin])
    
	
	#show()

    plt.ioff() # This so that the windows stay open
    plt.show()