diff --git a/examples/fft/spatial_filter.py b/examples/fft/spatial_filter.py new file mode 100644 index 0000000000000000000000000000000000000000..c255507781881e019b01725ca4932259c8491974 --- /dev/null +++ b/examples/fft/spatial_filter.py @@ -0,0 +1,122 @@ +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() + diff --git a/pykat/fft/fft.py b/pykat/fft/fft.py new file mode 100644 index 0000000000000000000000000000000000000000..90a36d432275e76aad1e0ce38096246891006ef5 --- /dev/null +++ b/pykat/fft/fft.py @@ -0,0 +1,96 @@ +""" +------------------------------------------------------ +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)