Example of scaled propagation

examples/fft/scale_propagation.py

@@ -3,50 +3,80 @@ from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import pylab as pl
+"""
+Demonstration of oscar.py's scaled propagation method.
+
+Sean Leavey
+sean.leavey@ligo.org
+May 2015
+"""
 
 import pykat.oscar as oscar
+import pylab as pl
 import numpy as np
+from mpl_toolkits.axes_grid1 import ImageGrid
 
 def main():
-	grid1 = oscar.grid(256, 256, 1e-2, 1e-2, xoffset=0, yoffset=0)
-	grid2 = oscar.grid(256, 256, 0.5e-2, 0.5e-2)
+	### parameters
 
-	field1 = oscar.field(grid1, w=1e-3, power=1, mode='HG 0 0')
-	field2 = oscar.field(grid1, w=1e-3, power=1, mode='HG 0 0')
+	# waist size at laser [m]
+	waist = 1e-3
 
-	field1.propagate(10)
-	field2.scalePropagate(10, 1, grid2)
+	# power at laser [W]
+	power = 1
 
-	Z1 = np.abs(field1.amplitude) ** 2
-	Z2 = np.abs(field2.amplitude) ** 2
+	# light mode
+	mode = 'HG 0 0'
+
+	# grid scale factor
+	scale = 2
 
-	fig = pl.figure()
+	# propagation distance [m]
+	distance = 1
 
-	pl.subplot(2, 1, 1)
-	ax = pl.gca()
+	### propagation
 
-	xLim = (field1.grid.xaxis.min(), field1.grid.xaxis.max())
-	yLim = (field1.grid.yaxis.min(), field1.grid.yaxis.max())
+	# create different grids to demonstrate scaled propagation
+	grid1 = oscar.grid(512, 512, 10e-3, 10e-3)
+	grid2 = oscar.grid(512, 512, 5e-3, 5e-3)
 
-	ax.set_xlim(*xLim)
-	ax.set_ylim(*yLim)
+	# create two identical fields
+	field1 = oscar.field(grid1, w=waist, power=power, mode=mode)
+	field2 = field1.copy()
+
+	# propagate without scaling
+	field1.propagate(distance)
+
+	# propagate with scaling
+	field2.scalePropagate(distance, scale, grid2)
+
+	# get magnitudes of the fields
+	Z1 = np.abs(field1.amplitude) ** 2
+	Z2 = np.abs(field2.amplitude) ** 2
 
-	im = ax.imshow(Z1, extent=[xLim[0], xLim[1], yLim[0], yLim[1]])
-	cb = fig.colorbar(im, ax=ax)
+	### plot
 
-	pl.subplot(2, 1, 2)
-	ax = pl.gca()
+	# create figure and image grid for two heatmaps
+	fig = pl.figure(figsize=(12, 8))
+	imgrid = ImageGrid(fig, 111, nrows_ncols=(1, 2), axes_pad=0.1)
 
-	xLim = (field2.grid.xaxis.min(), field2.grid.xaxis.max())
-	yLim = (field2.grid.yaxis.min(), field2.grid.yaxis.max())
+	# figure out global lowest and highest limits of the grids (so we can compare the sizes easily)
+	xLim = (min([field1.grid.xaxis.min(), field2.grid.xaxis.min()]), max([field1.grid.xaxis.max(), field2.grid.xaxis.max()]))
+	yLim = (min([field1.grid.yaxis.min(), field2.grid.yaxis.min()]), max([field1.grid.yaxis.max(), field2.grid.yaxis.max()]))
 
-	ax.set_xlim(*xLim)
-	ax.set_ylim(*yLim)
+	# plot first propagation
+	imgrid[0].set_xlim(*xLim)
+	imgrid[0].set_ylim(*yLim)
+	imgrid[0].set_title('Unscaled Propagation')
+	imgrid[0].imshow(Z1, extent=xLim+yLim)
 
-	im = ax.imshow(Z2, extent=[xLim[0], xLim[1], yLim[0], yLim[1]])
-	cb = fig.colorbar(im, ax=ax)
+	# plot second propagation
+	imgrid[1].set_xlim(*xLim)
+	imgrid[1].set_ylim(*yLim)
+	imgrid[1].set_title('Scaled Propagation')
+	imgrid[1].imshow(Z2, extent=xLim+yLim)
 
+	# show on screen
 	pl.show()
 if __name__ == '__main__':
 	main()

pykat/oscar.py

@@ -428,12 +428,20 @@ class field(object):
 	def scalePropagate(self, distance, scale, newGrid):
 		# 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.
+		# 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
 
-		# scale = w0 / w1
+		# For now, throw an error if distance == 0.
+		# If distance is zero, the function should return a simple
+		# scaled, unpropagated field.
+		# Currently, if distance == 0, a division by zero occurs due to
+		# the use of an inverse z0 function instead of simple z0 (which
+		# is itself used to avoid scale == 1 causing a division by zero).
+		assert(distance != 0)
 
 		if scale <= 0:
 			raise Exception('Specified scale factor must be > 0')
@@ -441,8 +449,10 @@ class field(object):
 		if newGrid.xpoints != self.grid.xpoints or newGrid.ypoints != self.grid.ypoints:
 			raise Exception('New grid must have same number of points as existing grid')
 
-		plD = np.pi * self.wavelength * distance * scale / self.ref_index
-		invz0 = (1.0 / scale - 1.0) / distance
+		plD = np.pi * self.wavelength * distance / (scale * self.ref_index)
+
+		# invz0 is used to avoid division by zero if scale == 1
+		invz0 = (scale - 1.0) / distance
 	
 		# initial scaling
 		field = self.amplitude * np.exp(-1j * self.k_prop * self.grid.r_squared * invz0 / 2.0)
@@ -451,7 +461,7 @@ class field(object):
 		field = field * np.exp(-1j * self.k_prop * distance) * np.exp(1j * plD * self.grid.fft_ir_squared)
 		field = np.fft.ifft2(field)
 		# final scaling
-		self.amplitude = field * scale * np.exp(1j * newGrid.r_squared * (invz0 + distance * invz0 * invz0) / 2.0)
+		self.amplitude = field * np.exp(1j * newGrid.r_squared * (invz0 + distance * invz0 * invz0) / (2.0 * scale))
 
 		# update grid
 		self.grid = newGrid