Fast Fourier Transform (FFT) examples

The Discrete Fourier Transform(DFT) is defined by X_k = \sum_{n=0}^{N-1} x_n \cdot e^{-i 2 \pi \frac{k}{N} n} latex and the inverse fourier transform is defined as x_n = \frac{1}{N} \sum_{k=0}^{N-1} X_k \cdot e^{+i 2 \pi \frac{k}{N} n} latex

In [2]:
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
import time

'''
	Define two functions to compute fourier transforms using two different methods
	1.	dft(x,inverse=False,timer=False): for INVERSE: inverse=true
						  for timed: timer=true
	2.	timedfft(x,inverse=False,timer):for INVERSE: inverse=true
						  for timed: timer=true
	3.	func(x): defines function at x
'''

def func(x):
    f=10.0 
    val=np.sin(2*np.pi*f*x)
    return val

'''
   	Computes the discrete fourier transform /inverse fourier transform
	using basic summation convention.
	OPTIONS:
		inverse: computes inverse fourier transform
		timer:	 prints time required for computation
'''
def dft(x, inverse = False, timer = False) :

    t = time.clock()							#get current time to calculate the computation time
    N = len(x)
    if inverse:
        inv=1
    else:
        inv=-1

    X =[0] * N								#make an array of size N

    for k in xrange(N) :						#compute the required sum
        for n in xrange(N) :
            X[k] += x[n] * np.e**(inv * 2j * np.pi * k * n / N)
        if inverse :
            X[k] /= N

    t = time.clock() - t						#get delta time
    if timer :
        print "Computed DFT of size",N,"in",t,"sec."
    return X

'''
	fourier transform using built in fast fourier method
'''

def timedfft(x,inverse=False,timer=False):

	t=time.clock()
	if not inverse:X=np.fft.fft(x)
	else:	X=np.fft.ifft(x)
	t=time.clock()-t
	if timer: print "Computed DFT using fft of size",len(x),"in",t,"sec."
	return X

'''
	Computing part: 1. Define a function over a interval
			2. Compute fourier transform using the two methods
			3. Compare time
'''

#fix sample size, this gives the the rate of sampling of the function
sample_size=2**8
L=5.0
x=np.linspace(0,L,sample_size)

#there are four plots 1.original signal 2.fourier transformed signal 3.fast fourier transformed signal 4.Reconstructed signal

plt.subplot(411)
plt.xlabel('t')
plt.ylabel('f($t$)')
plt.xlim([0,L/10])
signal=func(x)
plt.plot(x,signal,label='original signal')

'''
				In [20]: run fourier_basic.py
				Computed DFT using fft of size 256 in 0.0 sec.
				Computed DFT of size 256 in 0.48 sec.
				Computed DFT of size 256 in 0.16 sec.

				In [21]: run fourier_basic.py
				Computed DFT using fft of size 4096 in 0.0 sec.
				Computed DFT of size 4096 in 123.94 sec.
				Computed DFT of size 4096 in 43.18 sec.
'''
#compute timed fft
fourierfft=timedfft(signal,timer=True)

#caluclate discrete fourier transform using normal method
fourier =dft(signal,timer=True)

#get the frequencies
n=x.size
d=x[1]-x[0]
freqsfft=np.fft.fftfreq(n,d)

#get positive index
indices=np.where(freqsfft>=0)

#plot
plt.subplot(412)
plt.xlabel('freq')
plt.ylabel('f($\omega$)')
plt.plot(freqsfft,fourierfft,label='FFT')

plt.subplot(413)
plt.xlabel('freq')
plt.ylabel('f($\omega$)')
plt.plot(freqsfft,fourier,color='green',label='Without FFT')

#inverse fourier transform using dft
invfourier=(dft(fourier,inverse=True,timer=True))

plt.subplot(414)
plt.xlabel('t')
plt.ylabel('f($t$)')
plt.xlim([0,L/10])
plt.plot(x,invfourier,'+-',color='red')
display(gcf())
Computed DFT using fft of size 256 in 0.0 sec.
Computed DFT of size 256 in 0.48 sec.
Computed DFT of size 256 in 0.16 sec.

/numpy/core/numeric.py:235: ComplexWarning: Casting complex values to real discards the imaginary part
  return array(a, dtype, copy=False, order=order)

N.B.: The fourier transform of the function is symmertric about 0.
Thus it suffices to display only the positive half,
as the negetive half gives us no more information

In [3]:
clf()
plt.subplot(111)
pi=np.pi
sampled=5*np.sin(2*pi*x)+np.sin(2*pi*10*x)
fourier=np.fft.fft(sampled)
#if you want to print only the positive frequencies
index=np.where(freqsfft>=0)
plt.plot(freqsfft[index],fourier[index])
display(gcf())

In [39]:
clf()
sampled=np.exp(-x**2)
fourier=np.fft.fft(sampled)
plt.plot(freqsfft,fourier)
display(gcf())

In [40]:
clf()
from scipy.special import jn
sampled=jn(2,x)
fourier=np.fft.fft(sampled)
plt.plot(freqsfft,fourier)
display(gcf())

Advertisements

2D Fast Poisson Solver

2D Fast Poisson Solver:

Introduction:

Here we will construct poisson solver for a 2D square boundary using a variation
fast fourier transform , dicrete sine transform.There are 4 types of discrete sine transform,you can read more about them in
wikipedia http://en.wikipedia.org/wiki/Discrete_sine_transform

We will solve the 2D electrostatic equation for square boundary,in

\epsilon_0=1 units

The poisson equation solver in 2D basically uses the matrix representation of the discrete poisson equation.
The poisson equation \nabla^2 u= \rho can be solved using fourier transform technique.The algorithm is basically
from
http://www.math.nus.edu.sg/~bao/teach/ma5233/lect10.ppt

This uses dst function(Type -I) from either fft.fftpack or a seperate
dst package I downloaded from internet called transforms

CODE:

In [2]:
from __future__ import division
import matplotlib.pyplot as plt
import numpy as np
from numpy.fft import fft,fft2,ifft2,ifft,irfft2,rfft2
import random as random
from mpl_toolkits.mplot3d import Axes3D

Define a function for discrete sin transform using fft algorithm,by symmetrizing the function x

In [3]:
def dst(x,axis=-1):
    """Discrete Sine Transform (DST-I)

    Implemented using 2(N+1)-point FFT
    xsym = r_[0,x,0,-x[::-1]]
    DST = (-imag(fft(xsym))/2)[1:(N+1)]

    adjusted to work over an arbitrary axis for entire n-dim array
    """
    n = len(x.shape)
    N = x.shape[axis]
    slices = [None]*3
    for k in range(3):
        slices[k] = []
        for j in range(n):
            slices[k].append(slice(None))
    newshape = list(x.shape)
    newshape[axis] = 2*(N+1)
    xsym = np.zeros(newshape,np.float)
    slices[0][axis] = slice(1,N+1)
    slices[1][axis] = slice(N+2,None)
    slices[2][axis] = slice(None,None,-1)
    for k in range(3):
        slices[k] = tuple(slices[k])
    xsym[slices[0]] = x
    xsym[slices[1]] = -x[slices[2]]
    DST = fft(xsym,axis=axis)
    #print xtilde
    return (-(DST.imag)/2)[slices[0]]

Define 2 dimensional DST, the idst is same as DST because of symmetry

In [4]:
def dst2(x,axes=(-1,-2)):
    return dst(dst(x,axis=axes[0]),axis=axes[1])

def idst2(x,axes=(-1,-2)):
    return dst(dst(x,axis=axes[0]),axis=axes[1])

Now define the poisson solver in 2D square boundary using the algorithm described above,this method can be
generalised to arbitary retangular boundary

In [5]:
def fft_poisson(f,h):

	m,n=f.shape
	f_bar=np.zeros([n,n])
	u_bar = f_bar			# make of the same shape
	u = u_bar

	f_bar=idst2(f)			# f_bar= fourier transform of f

	f_bar = f_bar * (2/n+1)**2  #Normalize
	#u_bar =np.zeros([n,n])
	pi=np.pi
	lam = np.arange(1,n+1)
	lam = -4/h**2 * (np.sin((lam*pi) / (2*(n + 1))))**2 #$compute $\lambda_x$
	#for rectangular domain add $lambda_y$
	for i in xrange(0,n):
		for j in xrange(0,n):
			u_bar[i,j] = (f_bar[i,j]) / (lam[i] + lam[j])

	u=dst2(u_bar)				#sine transform back
	u= u * (2/(n+1))**2			#normalize ,change for rectangular domain
	return u

Fix the initial parameters,domain of the equation etc.

In [6]:
#set bounds a,b,parameters
a = 0; b = 1;	
alpha=10				#alpha is grid points=2^alpha
n=2**alpha
L=b-a					#length of system

xe =np.linspace(a,b,n); 
ye =np.linspace(a,b,n);
x, y = np.meshgrid(xe,ye)
In [7]:
h=L/(n);			#size 
h2 = h**2;			#h squared
hx2=h2; 			#started with a cube,hx2=hy2
hy2=h2;
f=np.zeros([n,n]);

We put charges in the field.

Putting a random negetive charge in the field alog with a posituve charge

In [8]:
#initial conditions

#f[round(n/2):n,round(n/2):n]=1; #Initial condition
f[round((n+1)/2),round((n+1)/2-10)]=20
#f[round((n+1)/2),round((n+1)/2+10)]=-20
f[random.randint(0,n),random.randint(0,n)]=-10	#put a random charge
In [9]:
nx,ny=np.array(f.shape)-1		 #last index
U=np.zeros([n,n])

set boundary conditions

In [10]:
# BOUNDARY CONDITIONS
#set U(x,y)=g(x,y)
U[0,:]=0
U[nx,:]=0#5e-5
U[:,0]=0
U[:,ny]=0
In [11]:
##homogenize boundary condition

f[1,:]=f[1,:]+U[0,:]/hx2;
f[nx-1,:]=f[nx-1,:]+U[n-1,:]/hx2;

f[:,1]=f[:,1]+U[:,0]/hy2;
f[:,ny-1]=f[:,ny-1]+U[:,n-1]/hy2;

Solve the equation:

In [12]:
U=fft_poisson(f,h)
In [13]:
plt.figure()
#plt.imshow((U),cmap='hot')
plt.contour(U,50)
#plt.contour(f)

plt.show()
In [14]:
display(gcf())

In [15]:
imshow(U)
Out[15]:
<matplotlib.image.AxesImage at 0xb2b8a8c>
In [16]:
display(gcf())

 

In [ ]:

//

Sieve of Eratosthenes-primes

SyntaxHighlighter.config.bloggerMode=true; SyntaxHighlighter.config.clipboardSwf=’http://alexgorbatchev.com/pub/sh/current/scripts/clipboard.swf’; SyntaxHighlighter.all(); My python diaries… hello all, Recently I was experimenting with sieve method for finding primes in python… This code is really fast and memory saving than the ones commonly avaliable have a look…
It produced 100000000 primes on my 2.4Ghz i7 machine in just few seconds…



# Aritra Kundu
# arsenous.blogspot.com
'''
create prime numbers list using sieve method
'''
import numpy as np
import time as ti

def primes_up_to(n):
rootn=int(np.sqrt(n))
sieve=[True]*(n+1) #create a n+1 arrray
#set 0,1, false +ve
sieve[0]=sieve[1]=False

for i in xrange(2,rootn+1):
if sieve[i]:
no_of_multiples=n/i-i
sieve[i*i:n+1:i]=[False]*(no_of_multiples+1)

return [i for i in xrange(n+1) if sieve[i]]

#a second commonly avaliable method..the above is almost 30% faster as it uses binary and memory saving...
def primes_up_to2(n):
rootn=int(np.sqrt(n))
sieve=xrange(n+1) #create a n+1 arrray
#set 0,1, false +ve
sieve[0]=sieve[1]=0

for i in xrange(2,rootn+1):
if sieve[i] !=0:
no_of_multiples=n/i-i
sieve[i*i:n+1:i]=[0]*(no_of_multiples+1)

sieve=[x for x in sieve if x !=0]
#print sieve
return sieve

def count_primes_up_to(n):
return len(primes_up_to(n))

N=1000000

def save_primes():
t=ti.time()
sieve_file=open('sieve_primes_%s'%N,'w')
p=primes_up_to(N)
print len(p)
sieve_file.writelines("%s\n"%len(p))
sieve_file.writelines(["%s\n" %item for item in p])
t=ti.time()-t
print 'TF modified',t
sieve_file.close()

def is_prime(n):
return n in primes_up_to(n)


Quantum Harmonic Oscillator -Simulations

Heya…For the past couple of weeks I have been doing simulations for quantum systems, like the harmonic oscillator(mainly) and square well..as a part of the course..We produced some really nice graphics (using python) for multiple harmonic wells.
In this post I will be sharing some of the simulations..It can give a good insight to the physics of quantum mechanics.
Plots of wavefunctions for  multiple harmonic wells

Some useful points to be noted
  1. The degeneracy in each levels increase as the number of wells is increased.(remember the lattice structure?)
  2. For large number of wells, it begins to show band structure.
  3. The separation between the wells determine the magnitude of their interaction

formation of band structure

Random Walker – Python Simulation

SyntaxHighlighter.config.bloggerMode=true; SyntaxHighlighter.config.clipboardSwf=’http://alexgorbatchev.com/pub/sh/current/scripts/clipboard.swf’; SyntaxHighlighter.all(); I havent posted for long ..have been quite busy..here is some very simple code to animate random walk I just love seeing how the random walk simulates…experiment with it..and freely redistribute..;) Recently Im learing python simulations…so this was what i did while reading statistical mechanics by richel

# Random_walker.py
# Uses turtle class to simulate random walk through pseudo random generator
# external interrupt is enabled ,such that one can give external impacts
# to the random walker and see if that small external interrupts change random walk pattern
#1.give brisk shocks with up /right/left keys
#2.try combinations
# Aritra Kundu
# arsenous.blogspot.com


import turtle
import numpy as np
turtle.setup(400,500) # determine the window size
wn = turtle.Screen() # get a reference to the window
wn.title("Handling keypresses!") # Change the window title
tess = turtle.Turtle() # Create our favorite turtle

# The next four functions are our "event handlers".
def h1():
tess.forward(30)

def h2():
tess.left(45)

def h3():
tess.right(45)

def h4():
wn.bye() # Close down the turtle window

#These lines "wire up" keypresses to the handlers we've defined.
wn.onkey(h1, "Up")
wn.onkey(h2, "Left")
wn.onkey(h3, "Right")
wn.onkey(h4, "q")

wn.listen()

#listen responses from keyboard
N=1000 # no of steps
for i in xrange(1,N):
x=np.random.randint(-15,15) #step size for random walker,generate integer
theta=np.random.randint(0,360) # rotation angle
tess.forward(x)
tess.right(theta)

turtle.clear()
#wn.mainloop()