Bugfixes to FFT scaled propagation

examples/fft/scale_propagation.py

@@ -14,35 +14,46 @@ May 2015
 import pykat.oscar as oscar
 import pylab as pl
 import numpy as np
-from mpl_toolkits.axes_grid1 import ImageGrid
+#from mpl_toolkits.axes_grid1 import ImageGrid
+import matplotlib as mpl
 
 def main():
 	### parameters
 
-	# waist size at laser [m]
-	waist = 1e-3
-
 	# power at laser [W]
 	power = 1
 
 	# light mode
 	mode = 'HG 0 0'
 
-	# grid scale factor
-	scale = 2
+	# waist size of beam at start [m]
+	waist = 1e-3
+
+	# unscaled physical grid size [m]
+	w0 = 10e-3
+
+	# scaled physical grid size [m]
+	w1 = 5e-3
 
 	# propagation distance [m]
-	distance = 1
+	distance = 5
+
+	# grid scale factor
+	scale = 2
 
 	### propagation
 
 	# 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)
+	grid1 = oscar.grid(512, 512, w0, w0)
+	grid2 = oscar.grid(512, 512, w1, w1)
+
+	# create input field
+	laser = oscar.field(grid1, w=waist, power=power, mode=mode)
 
-	# create two identical fields
-	field1 = oscar.field(grid1, w=waist, power=power, mode=mode)
-	field2 = field1.copy()
+	# create three identical fields
+	field0 = laser.copy()
+	field1 = laser.copy()
+	field2 = laser.copy()
 
 	# propagate without scaling
 	field1.propagate(distance)
@@ -51,32 +62,47 @@ def main():
 	field2.scalePropagate(distance, scale, grid2)
 
 	# get magnitudes of the fields
+	Z0 = np.abs(field0.amplitude) ** 2
 	Z1 = np.abs(field1.amplitude) ** 2
 	Z2 = np.abs(field2.amplitude) ** 2
 
 	### plot
 
-	# 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)
-
-	# 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()]))
-
-	# 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)
-
-	# 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)
+	# initial and final physical sizes of grids
+	extentInit = [min(field0.grid.xaxis), max(field0.grid.xaxis), min(field0.grid.yaxis), max(field0.grid.yaxis)]
+	extentFinal1 = [min(field1.grid.xaxis), max(field1.grid.xaxis), min(field1.grid.yaxis), max(field1.grid.yaxis)]
+	extentFinal2 = [min(field2.grid.xaxis), max(field2.grid.xaxis), min(field2.grid.yaxis), max(field2.grid.yaxis)]
+
+	# minimum/maximum values across all signals (for colourmap)
+	globalMin = min([Z0.min(), Z1.min(), Z2.min()])
+	globalMax = min([Z0.max(), Z1.max(), Z2.max()])
+
+	fig, axes = pl.subplots(nrows=2, ncols=2, figsize=(8, 8))
+
+	# original beams
+	axes[0, 0].imshow(Z0, extent=extentInit)
+	axes[0, 0].set_title('Original Beam')
+	axes[0, 0].set_xlabel('Physical width [m]')
+	axes[0, 0].set_ylabel('Physical height [m]')
+	axes[0, 1].imshow(Z0, extent=extentInit)
+	axes[0, 1].set_title('Original Beam')
+	axes[0, 1].set_xlabel('Physical width [m]')
+	axes[0, 1].set_ylabel('Physical height [m]')
+
+	# unscaled propagated beam
+	axes[1, 0].imshow(Z1, extent=extentFinal1)
+	axes[1, 0].set_title('Unscaled Propagated Beam')
+	axes[1, 0].set_xlabel('Physical width [m]')
+	axes[1, 0].set_ylabel('Physical height [m]')
+
+	# scaled propagated beam
+	axes[1, 1].imshow(Z2, extent=extentFinal2)
+	axes[1, 1].set_title('Scaled Propagated Beam')
+	axes[1, 1].set_xlabel('Physical width [m]')
+	axes[1, 1].set_ylabel('Physical height [m]')
 
 	# show on screen
+	pl.tight_layout()
 	pl.show()
 if __name__ == '__main__':
 	main()

pykat/oscar.py

@@ -310,7 +310,7 @@ class field(object):
 		
 	def normalise(self, power = 1):
 		if power == 0:
-			self.amplitude = self.amplitde * 0
+			self.amplitude = self.amplitude * 0
 		else:
 			current_power = self.power()
 			if power != 0:
@@ -461,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 * np.exp(1j * newGrid.r_squared * (invz0 + distance * invz0 * invz0) / (2.0 * scale))
+		self.amplitude = field * np.exp(1j * self.k_prop * newGrid.r_squared * invz0 * (1 + distance * invz0) / 2) / scale
 
 		# update grid
 		self.grid = newGrid