This is an example of how to implement the split step propagation algorithm using python and numpy. Click here to see the results http://www.alexfb.com/cgi-bin/twiki/view/PtPhysics/VortexLattice
Python, 191 lines
Download
Copy to clipboard1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | # Date : 29th July
# Author : Alexander Baker
import sys
from pylab import figure, show, clf, savefig, cm
from numpy.fft import fft, fft2, ifft, ifft2, fftshift
from optparse import OptionParser
from numpy import *
def generateResolution(n):
return 2**n, (2**n)-1
def store_value(option, opt_str, value, parser):
setattr(parser.values, option.dest, value)
def main():
parser = OptionParser()
parser.add_option("-r", "--resolution", action="callback", callback=store_value, type="int", nargs=1, dest="resolution", help="resolution of the grid parameter")
parser.add_option("-g", "--grid", action="callback", callback=store_value, type="int", nargs=1, dest="grid", help="grid size parameter")
parser.add_option("-s", "--steps", action="callback", callback=store_value, type="int", nargs=1, dest="steps", help="number of steps")
parser.add_option("-b", "--beam", action="callback", callback=store_value, type="int", nargs=1, dest="beam", help="beam size")
parser.add_option("-c", "--core", action="callback", callback=store_value, type="float", nargs=1, dest="core", help="beam size")
parser.add_option("-p", "--phase", action="callback", callback=store_value, type="int", nargs=1, dest="phase", help="phase size parameter")
parser.add_option("-i", "--image", action="callback", callback=store_value, type="string", nargs=1, dest="image", help="phase size parameter")
(options, args) = parser.parse_args()
grid_resolution = 6
grid_size = 8
steps = 35
beam = 8
core = 0.3
phase = 15
image = 'bw'
if options.resolution:
grid_resolution = options.resolution
if options.grid:
grid_size = options.grid
if options.steps:
steps = options.steps
if options.beam:
beam = options.beam
if options.core:
core = options.core
if options.phase:
phase = options.phase
if options.image:
image = options.image
print '\n######################################################'
print '# Numerical Lattice Model'
print '# Alexander Baker - August 2007'
print '######################################################\n'
grid_resolution, grid_resolution_1 = generateResolution(grid_resolution)
# The grid size effects the number of steps we apply to get to the upper limit
print "grid_size : [%d]\ngrid_resolution [%d]\ngrid_resolution-1 [%d]" %(grid_size, grid_resolution, grid_resolution_1)
x1 = arange(-1 * 1.0* grid_size,(-1 * grid_size) + grid_resolution_1 * 1.0 * (2 * grid_size)/grid_resolution, (grid_size * 1.0)/grid_resolution)
x2 = arange(-1 * 1.0 * grid_size,(-1 * grid_size) + grid_resolution_1 * 1.0 * (2 * grid_size)/grid_resolution, (grid_size * 1.0)/grid_resolution)
xmin, xmax, ymin, ymax = -1 * 1.0 * grid_size, (-1 * grid_size) + grid_resolution_1 * 1.0 * (2 * grid_size)/grid_resolution, -1 * 1.0 * grid_size, (-1 * grid_size) + grid_resolution_1 * 1.0 * (2 * grid_size)/grid_resolution
extent = xmin, xmax, ymin, ymax
print "\ngrid extent X-[%f][%f] Y-[%f][%f]" % (xmin, xmax, ymin, ymax)
print "shape x1", shape(x1)
print "shape x2", shape(x2)
print "step [%f]" % ((grid_size * 1.0)/grid_resolution)
print "start [%f], end[%f]" % (-1 * 1.0* grid_size, grid_resolution_1 * 1.0 * (2 * grid_size)/grid_resolution)
[x1,x2] = meshgrid(x1,x2)
rbeam = beam
wcore = core
vpos = 1
_del = 1
print "\nrbeam [%d]\nwcore [%f]\nvpos [%f]\n_del [%f]" % (rbeam, wcore, vpos, _del)
#
# Step 1. Take your initial beam profile (gaussian, a vortex etc etc) at z = 0
#
y = 1.0*exp(-2*(x1**2 + x2**2)/rbeam**2)*exp(1j* phase *(+arctan2(x2,x1))) * tanh(pow((x1**2 + x2**2), 0.5)/wcore)
fig1 = figure(1)
fig1.clf()
ax1a = fig1.add_subplot(121)
if image == 'bw':
ax1a.imshow(angle((y)), cmap=cm.gist_gray, alpha=.9, interpolation='bilinear', extent=extent)
else:
ax1a.imshow(angle((y)), cmap=cm.jet, alpha=.9, interpolation='bilinear', extent=extent)
ax1a.set_title(r'Angle')
ax1b = fig1.add_subplot(122)
if image == 'bw':
ax1b.imshow(abs((y)), cmap=cm.gist_gray, alpha=.9, interpolation='bilinear', extent=extent)
else:
ax1b.imshow(abs((y)), cmap=cm.jet, alpha=.9, interpolation='bilinear', extent=extent)
ax1b.set_title(r'Amplitude')
savefig('big_start_' + str(wcore)+'_' + str(vpos) +'.png')
u1 = arange(-1.0,-1+1.0*grid_resolution_1* ((2 * 1.0)/grid_resolution), 1.0/grid_resolution)
u2 = arange(-1.0,-1+1.0*grid_resolution_1* ((2 * 1.0)/grid_resolution), 1.0/grid_resolution)
u1min, u1max, u2min, u2max = -1.0, -1+1.0*grid_resolution_1* ((2 * 1.0)/grid_resolution), -1.0, -1+1.0*grid_resolution_1* ((2 * 1.0)/grid_resolution)
print "\npropagation grid X-[%f][%f] Y-[%f][%f]" % (u1min, u1max, u2min, u2max)
print "shape u1", shape(u1)
print "shape u2", shape(u2)
print "step [%f]" % (1.0/grid_resolution)
print "start [%f], end[%f]" % (-1.0, -1+1.0*grid_resolution_1* ((2 * 1.0)/grid_resolution))
print "\nbeam power (start) - [%f]" % (sum(sum(abs(y**2))))
[u1,u2] = meshgrid(u1,u2)
t = exp(2*pi*1j*(u1**2 + u2**2)*_del)
w = fftshift(t)
#
# Step 2. Split step progagation
#
for i in arange(100,100+steps, 1):
z = fft2(y)
zp = z * w
yp = ifft2(zp)
p = (exp(+0.01*pi*1j*(x1**2 + x2**2)*_del + 0.05*pi*1j*y*conj(y))*_del);
yp = yp * p
y = yp
zp = fft2(yp)
zp = zp * w
yp = ifft2(zp)
fig3 = figure(3)
fig3.clf()
ax3 = fig3.add_subplot(111)
if image == 'bw':
ax3.imshow(abs((yp)), cmap=cm.gist_gray, alpha=.9, interpolation='bilinear', extent=extent)
else:
ax3.imshow(abs((yp)), cmap=cm.jet, alpha=.9, interpolation='bilinear', extent=extent)
ax3 = fig3.add_subplot(111)
if image == 'bw':
ax3.imshow(angle((yp)), cmap=cm.gist_gray, alpha=.9, interpolation='bilinear', extent=extent)
else :
ax3.imshow(angle((yp)), cmap=cm.jet, alpha=.9, interpolation='bilinear', extent=extent)
print sum(sum(abs(yp**2))), i-100
print "beam power (end) - [%f]" % (sum(sum(abs(yp**2))))
fig2 = figure(2)
fig2.clf()
ax2a = fig2.add_subplot(121)
if image == 'bw':
ax2a.imshow(angle((yp)), cmap=cm.gist_gray, alpha=.9, interpolation='bilinear', extent=extent)
else:
ax2a.imshow(angle((yp)), cmap=cm.jet, alpha=.9, interpolation='bilinear', extent=extent)
ax2a.set_title(r'Angle')
ax2b = fig2.add_subplot(122)
if image == 'bw':
ax2b.imshow(abs((yp)), cmap=cm.gist_gray, alpha=.9, interpolation='bilinear', extent=extent)
else:
ax2b.imshow(abs((yp)), cmap=cm.jet, alpha=.9, interpolation='bilinear', extent=extent)
ax2b.set_title(r'Amplitude')
savefig('big_end_' + str(wcore)+'_' + str(vpos) +'.png')
print '\ndone. ok'
if __name__ == "__main__":
main()
|
This is a particluarly tricky algorithm to implement and is a very popular method for beam propagation in optics.
Tags: algorithms
