adding fft propagation basics and an example

examples/fft/spatial_filter.py

+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import matplotlib
+BACKEND = 'Qt4Agg'
+matplotlib.use(BACKEND)
+import pylab as pl
+
+from pykat.utilities.optics.gaussian_beams import HG_beam, beam_param
+from pykat.fft.fft import *
+import numpy as np
+import scipy
+
+def main():
+	# wavelength
+	Lambda = 1064.0E-9 
+	# distance to propagate/focal length of lens
+	D = 10
+	# mix coefficients
+	c1 = 0.7
+	c2 = 0.3
+	# mode indices 
+	n1 = 2
+	m1 = 3
+	n2 = 1
+	m2 = 0
+
+	######## Generate Grid stucture required for FFT propagation ####
+	xpoints = 512
+	ypoints = 512
+	xsize = 0.05
+	ysize = 0.05
+	# Apply offset such that the center of the beam lies in the
+	# center of a grid tile
+	xoffset = -0.5*xsize/xpoints
+	yoffset = -0.5*ysize/ypoints
+
+	shape = grid(xpoints, ypoints, xsize, ysize, xoffset, yoffset)
+	x = shape.xaxis
+	y = shape.yaxis
+
+	######## Generates a mixture of fields ################
+	gx = beam_param(w0=2e-3, z=0)
+	gy = beam_param(w0=2e-3, z=0)
+	beam = HG_beam(gx,gy,n1,m1)
+	field1 = beam.Unm(x,y) 
+	beam2 = HG_beam(gx,gy,n2,m2)
+	field2 = beam2.Unm(x,y) 
+	global field, laser1, laser2
+	field = np.sqrt(c1)*field1 + np.sqrt(c2)*field2
+ 
+	####### Apply phase plate #######################################
+
+	laser1 = field*(np.conjugate(field1))
+	laser2 = field*(np.conjugate(field2))
+
+	####### Propagates the field by FFT ##############################
+	laser1 = FFT_propagate(laser1,shape,Lambda,D,1) 
+	laser2 = FFT_propagate(laser2,shape,Lambda,D,1) 
+
+	f=D
+	#laser1 = apply_lens(laser1, shape, Lambda, f) 
+	#laser2 = apply_lens(laser2, shape, Lambda, f) 
+	laser1 = apply_thin_lens(laser1, shape, Lambda, f) 
+	laser2 = apply_thin_lens(laser2, shape, Lambda, f) 
+
+	laser1 = FFT_propagate(laser1,shape,Lambda,D,1) 
+	laser2 = FFT_propagate(laser2,shape,Lambda,D,1)
+	
+	# midpoint computation for even number of points only!
+	midx=(xpoints)//2
+	midy=(ypoints)//2
+	coef1 = np.abs(laser1[midx,midy])  
+	coef2 = np.abs(laser2[midx,midy])
+
+	ratio = (coef1/coef2)**2
+	pc2 = 1/(1+ratio)
+	pc1 = pc2*ratio
+
+	print("c1 {0}, coef1 {1}, error {3} (raw output {2})".format(c1, pc1, coef1, np.abs(c1-pc1)))
+	print("c2 {0}, coef2 {1}, error {3} (raw output {2})".format(c2, pc2, coef2, np.abs(c2-pc2)))
+
+	# plot hand tuned for certain ranges and sizes, not automtically scaled
+	fig=pl.figure(110)
+	fig.clear()
+	off1=xpoints/10
+	off2=xpoints/6
+	pl.subplot(1, 3, 1)
+	pl.imshow(abs(field))
+	pl.xlim(midx-off1,midx+off1)
+	pl.ylim(midy-off1,midy+off1)
+	pl.draw()
+	pl.subplot(1, 3, 2)
+	pl.imshow(abs(laser1))
+	pl.xlim(midx-off2,midx+off2)
+	pl.ylim(midy-off2,midy+off2)
+	pl.draw()
+	pl.subplot(1, 3, 3)
+	pl.imshow(abs(laser2))
+	pl.xlim(midx-off2,midx+off2)
+	pl.ylim(midy-off2,midy+off2)
+	pl.draw()
+	if in_ipython():
+		pl.show(block=0)
+	else:
+		pl.show(block=1)
+
+
+# testing if the script is run from within ipython
+def in_ipython():
+    try:
+        __IPYTHON__
+    except NameError:
+        return False
+    else:
+        return True
+
+if __name__ == '__main__':
+	main()
+

pykat/fft/fft.py

+"""
+------------------------------------------------------
+Functions related to FFT propogation of beams.
+Work in progress, currently these functions are
+untested!
+
+Andreas 30.11.2014
+http://www.gwoptics.org/pykat/
+------------------------------------------------------
+"""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+import numpy as np
+
+def apply_lens(field, grid, Lambda, f):
+	# apply a phase factor representing a lens
+	k= 2.0*np.pi/Lambda
+	return field*(np.exp(2.0 * 1j * k * (2*f - np.sign(f)*np.sqrt((2.0*f)**2-grid.r_squared))))
+
+def apply_thin_lens(field, grid, Lambda, f):
+	# apply a phase factor representing a thin lens
+	k= 2.0*np.pi/Lambda
+	return field*(np.exp(1.0 * 1j * k * grid.r_squared/(2.0*f)))
+
+def FFT_propagate(field, grid, Lambda, distance, nr):
+	# FFT propagation code in a fixed grid
+
+	k = 2.0*np.pi/Lambda*nr
+	plD = np.pi*Lambda*distance/nr
+
+	field = np.fft.fft2(field)
+	field = field * np.exp(-1j*k*distance) * np.exp(1j*plD*grid.fft_ir_squared)
+	field = np.fft.ifft2(field)
+
+	return field
+
+def FFT_scale_propagate(field, grid0, grid1, Lambda, distance, w0, w1, nr):
+	# FFT propagation code with an adaptive grid size.
+	# Propagates to a scaled coordinate system, see Virgo Book of
+	# Physics pages 179-184, the scaling factor is given
+	# as w1/w0 with w0 the beam size at the start of propagation
+	# and w1 the expected beam size at the end of propatation.
+	# NOT YET TESTED
+
+	k = 2.0*np.pi/Lambda*nr
+	plD = np.pi*Lambda*distance*w0/w1/nr
+	z0 = distance/(w1/w0-1.0)
+	
+	# initial scaling
+	field = field * np.exp(-1j*k*grid0.r_squared/(2.0*z0))
+	field = np.fft.fft2(field)
+	# scaled propagator
+	field = field * np.exp(-1j*k*distance) * np.exp(1j*plD*grid0.fft_ir_squared)
+	field = np.fft.ifft2(field)
+	# final scaling
+	field = field *w0/w1 * np.exp(1j* grid1.r_squared*(z0+L)/(2.0*z0*z0))
+	
+	return field
+
+class grid():
+	# Data structure to describe the size and axes for a (x,y) data array
+	# of complex beam amplitudes. Also contain also data structures for
+	# FFT propagation
+	
+	def __init__ (self,xpoints, ypoints, xsize, ysize, xoffset, yoffset):
+
+		self.xpoints=xpoints
+		self.ypoints=ypoints
+		self.xsize=xsize
+		self.ysize=ysize
+		self.xoffset=xoffset
+		self.yoffset=yoffset
+
+		# compute x and y axis
+		self.xstep=self.xsize/self.xpoints
+		self.ystep=self.ysize/self.ypoints
+		xvector= np.arange(self.xpoints)
+		yvector= np.arange(self.ypoints)
+		self.xaxis=-self.xsize/2.0 + self.xstep/2.0 + xvector*self.xstep + self.xoffset
+		self.yaxis=-self.ysize/2.0 + self.ystep/2.0 + yvector*self.ystep + self.yoffset
+
+		# and some useful variables based on the axis
+		self.X,self.Y = np.meshgrid(self.xaxis,self.yaxis)
+		self.r_squared = (self.X)**2 + (self.Y)**2
+		self.r = np.sqrt(self.r_squared)
+		self.angle = np.arctan2(self.Y,self.X)
+
+		# compute frequency axis
+		self.xaxis_fft = np.fft.fftshift(np.fft.fftfreq(self.xpoints))/self.xstep
+		self.yaxis_fft = np.fft.fftshift(np.fft.fftfreq(self.ypoints))/self.ystep
+
+		# some useful variables based on the frequency axis
+		self.fft_X,self.fft_Y = np.meshgrid(self.xaxis_fft, self.yaxis_fft)
+		self.fft_ir_squared= np.fft.ifftshift((self.fft_X)**2+(self.fft_Y)**2)