[SciPy-user] Python evangelism question: procedural programming
Gael Varoquaux
gael.varoquaux at normalesup.org
Fri Aug 3 11:34:41 EDT 2007
On Fri, Aug 03, 2007 at 09:55:17AM -0500, Ryan Krauss wrote:
> So, he wanted to know if Python could be used like Mathematica in
> terms of writing rule and procedure based programs. I don't know how
> difficult this would be.
I think I understand what he means. Python can almost be used like this,
but there is the UI missing. Ipython is currently working on a "notebook"
UI that would address this problem. I have developped my own workflow
using a home made script ("pyreport"
http://gael-varoquaux.info/computers/pyreport/ ) to approach this
workflow. Ultimately I want to merge this with ipython's effort in order
to get nice notebook functionnality with html and pdf output, but
xurrently I don't have time for this.
I am attaching a file that shows how I use Python with Pyreport (Python
file attached, pdf file on http://gael-varoquaux.info/FORT_MOT.pdf ). The file
also shows the limits of the approach, and for larger project I use more
modular programming. This example is a bit frightening because it has
really grown longer than a Mathematica script and is a bit too complex to
have in such a monolythic code, but it is the only thing I have around.
There are also some very simple examples on pyreport's page.
So I think the answer is that it is possible with Python, but the tools
are not very developped and only slow progress is made.
My 2 cents,
Gaël
-------------- next part --------------
##!#####################################################################
##! Description of the FORT
##!#####################################################################
#Global imports.
from __future__ import division # So that 2/3 = 0.6666 and not 0 !
from pylab import *
from scipy import *
main = ( __name__ == '__main__' )
#! Constants
#!#########################################################################
# ----------------------------------------------------------------------
# FORT constants:
P_FORT = 50
w_0 = 100e-6
lambda_L = 1565e-9
print("FORT power: " + str(P_FORT) + " W, \t waist: " + str(w_0) +
" m \t wavelength: " + str(lambda_L) + " m")
# Position of the center relative to the center of the MOT.
r_FORT = array([0.0, 0.0, 0.0])
# Transformation to go from a gaussian beam propagating along z to the trap
rot_FORT = array([ [0, 0, 1, ], [ 0, 1, 0], [ 1, 0, 0, ], ])
print("center of the FORT relativ to the MOT : %.4e %.4e %.4e (x,y,z, m)" %(
r_FORT[0], r_FORT[1], r_FORT[2]))
# ----------------------------------------------------------------------
# Rubidium 87 D2 transition:
Gamma_Rb = 2*pi*5.98e6 # Linewidth
lambda_Rb = 780e-9 # Wavelength
# Rubidium excited state transitions:
# 5P(3/2) -> 4D(5/2) is the most important transiton of the excited
# level around 1560nm
# For the wavelength, see NIST database, for the Einstein A-Coefficient,
# see Clark PRA 69 022509
lambda_Rb_e = 1528.948e-9 # Wavelength
# 4*alpha * (2*pi)**3 * c_0/(3*e**2 * lamda_Rb_e**3) * (e*a_0)**2 *
# (10.634) ** 2 * 1/(2*3/2+1)
Gamma_Rb_e = 16e6 # Linewidth
# ----------------------------------------------------------------------
# Tables of rubidium transition for the MOT ground state:
# Wavelength
lambda_g = array([ 794.979,
780.241,
421.671,
420.298,
359.259,
358.807,
335.177,
334.966 ]) * 1e-9
# Einstein A coefficients (ie Gamma)
A_g = array([ 36.1,
38.1,
1.50,
1.77,
0.289,
0.396,
0.0891,
0.137 ]) * 1e6
# And for the excited state:
# Wavelength
lambda_e = array([ 1366.875,
741.021,
616.133,
1529.261,
1529.366,
776.157,
775.978,
630.0966,
630.0067 ]) * 1e-9
# Einstein A coefficients (ie Gamma)
A_e = array([ 9.5618,
3.0339,
1.4555,
1.774,
10.675,
0.67097,
3.93703,
0.63045,
3.71235 ]) * 1e6
# ----------------------------------------------------------------------
# MOT constants:
delta_0 = -15e6 # Detuning
P_mot = 35e-3 # Power in a beam
d_mot = 25e-3 # Diameter of a beam
print("MOT power: %.3e W" % P_mot + "\t beam diameter: %.3e m" % d_mot +
"\n detuning %.3e Hz" % delta_0 )
I_sat_Rb = 16 # Rubidium saturation intensity
mu_Rb = 0.71e10 # Rubidium magnetic moment (T*s)
# FIXME: There should be a 2*pi here, but it doesn't give the right
# results, moving along for now.
#mu_Rb = 0.71e10/(2*pi) # Rubidium magnetic moment (T*s)
Bx = 15e-2 # magnetic field gradient (T/m and not G/cm !)
m_Rb = 1.45e-25 # Rubidium 87 mass
# ----------------------------------------------------------------------
# Second MOT constants:
delta_1 = -12e6 # Detuning to the bottom of the FORT
P_mot_2 = 15e-3 # Power in a beam
#P_mot_2 = 0 # Power in a beam
print("Second MOT frequency: \t power: %.3e W" % P_mot_2 +
"\n\t\t detuning to the bottom of the trap :%.3e Hz" % delta_1 )
# ----------------------------------------------------------------------
# Fondamental constants:
c_0 = 3e8 # Speed of light (m/s)
h_bar = 1.054e-34 # h_bar !
k_J = 1.38e-23 # Blotzmann constant, in J/K
# ----------------------------------------------------------------------
# mesh grid of the region of space we are interested in :
def spaced_grid(xrange, yrange, nx=50, ny=50):
dx, dy = xrange/float(nx), yrange/float(ny)
return mgrid[-xrange:(xrange+dx):dx, -yrange:(yrange+dy):dy]
def make_grids(xrange, yrange, vrange):
global extentxy, Xgrid, Ygrid
extentxy = (-xrange, xrange, -yrange, yrange)
Xgrid, Ygrid = spaced_grid(xrange, yrange)
global extentyv, Ygrid2, Vgrid, Ygrid2sparse, Vgridsparse
yrange2 = 0.5*yrange
extentyv = (-yrange2, yrange2, -vrange, vrange)
Ygrid2, Vgrid = spaced_grid(yrange2, vrange)
Ygrid2sparse, Vgridsparse = spaced_grid(yrange2, vrange, nx=8, ny=8)
xrange, yrange, vrange = 0.04, 0.002, 4
[dx, dy, dv] = [xrange/50.0, yrange/50.0, vrange/50.0]
make_grids(xrange, yrange, vrange)
# The origin :
O=array([[0,0,0],])
# ----------------------------------------------------------------------
# Useful subroutines:
def pathcolor(x,y,t,colormap=cm.copper,linewidth=1, linestyle='solid'):
""" Plots a path indexed by a parameter t, using a colormap
to represent the parameter """
# A plot with line color changing
from matplotlib.collections import LineCollection
points=zip(x,y)
segments = zip(points[:-1], points[1:])
t = t + t.min()
t = t/t.max()
colors=cm.copper(t)
ax=gca()
LC = LineCollection(segments, colors = colors)
LC.set_linewidth(linewidth)
LC.set_linestyle(linestyle)
ax.add_collection(LC)
axis()
def cmap_map(function,cmap):
""" Applies function (which should operate on vectors of shape 3:
[r, g, b]), on the color vectors of colormap cmap. This routine will
break any discontinuous points in a colormap. Beware, function should map
the [0, 1] segment to itself, or you are in for surprises.
See also cmap_map.
"""
cdict = cmap._segmentdata
step_dict = {}
# Firt get the list of points where the segments start or end
for key in ('red','green','blue'):
step_dict[key] = map(lambda x: x[0], cdict[key])
step_list = reduce(lambda x, y: x+y, step_dict.values())
step_list = array(list(set(step_list)))
# Then compute the LUT, and apply the function to the LUT
reduced_cmap = lambda step : array(cmap(step)[0:3])
old_LUT = array(map( reduced_cmap, step_list))
new_LUT = array(map( function, old_LUT))
# Now try to make a minimal segment definition of the new LUT
cdict = {}
for i,key in enumerate(('red','green','blue')):
this_cdict = {}
for j,step in enumerate(step_list):
if step in step_dict[key]:
this_cdict[step] = new_LUT[j,i]
elif new_LUT[j,i]!=old_LUT[j,i]:
this_cdict[step] = new_LUT[j,i]
colorvector= map(lambda x: x + (x[1], ), this_cdict.items())
colorvector.sort()
cdict[key] = colorvector
return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)
def cmap_xmap(function,cmap):
""" Applies function, on the indices of colormap cmap. Beware, function
should map the [0, 1] segment to itself, or you are in for surprises.
See also cmap_xmap.
"""
cdict = cmap._segmentdata
function_to_map = lambda x : (function(x[0]), x[1], x[2])
for key in ('red','green','blue'):
cdict[key] = map(function_to_map, cdict[key])
cdict[key].sort()
assert (cdict[key][0]<0 or cdict[key][-1]>1), "Resulting indices extend out of the [0, 1] segment."
return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024)
light_jet = cmap_map(lambda x: x*0.7+0.3, cm.jet)
def trajplot(y,time, delta,
extent = ( -xrange, xrange, -yrange, yrange, -yrange, yrange)):
xmin = extent[0]
xmax = extent[1]
dx = (xmax-xmin)/100.0
ymin = extent[2]
ymax = extent[3]
dy = (ymax - ymin)/100.0
zmin = extent[4]
zmax = extent[5]
dz = (zmax - zmin)/100.0
# YZ plane
subplot(2,3,1)
[Ygrid, Zgrid] = mgrid[ ymin:ymax:dy, zmin:zmax:dz ]
positions=column_stack((zeros(Zgrid.size),Ygrid.ravel(),Zgrid.ravel()))
deltaMap=delta(positions)
deltaMat=reshape(deltaMap,(Ygrid.shape))
imshow(rot90(deltaMat),origin='lower',extent=(xmin, xmax, ymin,
ymax),aspect="auto", cmap=cm.bone)
cset = contour(rot90(deltaMat), array([0,Gamma_Rb/2,-Gamma_Rb/2]),
cmap=cm.bone, origin='lower', linewidths=1, extent=extentxy,
aspect="auto")
pathcolor(y[:,2],y[:,1],time)
yticks(yticks()[0][0:-1:2]) # Keep one tick out of two
xticks(xticks()[0][0:-1:2])
xticks(xticks()[0], len(xticks()[1])*['',]) # Hide the labels
ax = gca()
text(0.05,0.95,'Y',verticalalignment='top',transform = ax.transAxes)
text(0.95,0.05,'Z',horizontalalignment='right',transform = ax.transAxes)
title(' Projected trajectories', horizontalalignment = 'left')
# XY plane
[Xgrid, Ygrid] = mgrid[ xmin:xmax:dx, ymin:ymax:dy ]
positions=column_stack((Xgrid.ravel(),Ygrid.ravel(),zeros(Ygrid.size)))
subplot(2,3,2)
deltaMap=delta(positions)
deltaMat=reshape(deltaMap,(Xgrid.shape))
imshow(rot90(deltaMat),origin='lower',extent=(xmin, xmax, ymin,
ymax),aspect="auto", cmap=cm.bone)
cset = contour(rot90(deltaMat), array([0,Gamma_Rb/2,-Gamma_Rb/2]),
cmap=cm.bone, origin='lower', linewidths=1, extent=extentxy,
aspect="auto")
pathcolor(y[:,0],y[:,1],time)
yticks(yticks()[0][0:-1:2]) # Keep one tick out of two
xticks(xticks()[0][0:-1:2])
yticks(yticks()[0], len(yticks()[1])*['',])
ax = gca()
text(0.05, 0.95, 'Y', verticalalignment='top', transform=ax.transAxes)
text(0.95, 0.05, 'X', horizontalalignment='right', transform=ax.transAxes)
# XZ plane
[Xgrid, Zgrid] = mgrid[ xmin:xmax:dx, zmin:zmax:dz ]
positions = column_stack((Xgrid.ravel(),zeros(Zgrid.size),Zgrid.ravel()))
subplot(2, 3, 4)
deltaMap = delta(positions)
deltaMat = reshape(deltaMap,(Xgrid.shape))
imshow(deltaMat, origin='lower', extent=(xmin, xmax, zmin,
zmax), aspect="auto", cmap=cm.bone)
cset = contour(deltaMat, array([0, Gamma_Rb/2, -Gamma_Rb/2]),
cmap=cm.bone, origin='lower', linewidths=1,
extent=(xmin, xmax, zmin, zmax), aspect="auto")
pathcolor(y[:,0], y[:,2], time)
yticks(yticks()[0][0:-1:2]) # Keep one tick out of two
xticks(xticks()[0][0:-1:2])
ax = gca()
text(0.05,0.95,'X',verticalalignment='top',transform = ax.transAxes)
text(0.95,0.05,'Z',horizontalalignment='right',transform = ax.transAxes)
# X timeserie
subplot(3,3,3)
title('Time series')
plot(time,y[:,0])
ax = gca()
text(0.05,0.95,'X',verticalalignment='top',transform = ax.transAxes)
text(0.95,0.05,'t',horizontalalignment='right',transform = ax.transAxes)
xticks(xticks()[0], len(xticks()[1])*['',])
# Y timeserie
subplot(3,3,6)
plot(time,y[:,1])
ax = gca()
text(0.05,0.95,'Y',verticalalignment='top',transform = ax.transAxes)
text(0.95,0.05,'t',horizontalalignment='right',transform = ax.transAxes)
xticks(xticks()[0], len(xticks()[1])*['',])
# Z timeserie
subplot(3,3,9)
ax = gca()
plot(time,y[:,2])
text(0.05,0.95,'Z',verticalalignment='top',transform = ax.transAxes)
text(0.95,0.05,'t',horizontalalignment='right',transform = ax.transAxes)
#!
#! FORT
#!############################################################################
#!
#! Intensity distribution
#!----------------------------------------------------------------------------
Z_R = pi*w_0**2/lambda_L
print "Raleigh Length: " + str(Z_R) + " m"
w = lambda z : w_0 *sqrt( 1 + (z/Z_R) **2 )
Gaussian_beam = lambda r : P_FORT/(pi*w(r[:,2])**2)*exp(-(r[:,0]**2+r[:,1]**2)/(2*w(r[:,2])**2))
NablaGaussian = lambda r : transpose(Gaussian_beam(r) * array([
r[:,0]/w(r[:,2])**2,
r[:,1]/w(r[:,2])**2,
r[:,2]/(r[:,2]**2+Z_R**2)*(-2 + (r[:,0]**2 + r[:,1]**2)/w(r[:,2])**2)
]))
Intensity = lambda r : Gaussian_beam(dot((r-r_FORT),transpose(rot_FORT)) )
NablaIntensity = lambda r : dot(NablaGaussian(dot((r-r_FORT),transpose(rot_FORT)) ),transpose(rot_FORT))
positions=column_stack((Xgrid.ravel(),Ygrid.ravel(),zeros(Xgrid.size)))
IntensityMap=Intensity(positions)
IntensityMat=reshape(IntensityMap,(Xgrid.shape))
rcParams.update({'figure.figsize': [10.5,5], 'text.usetex': True})
if main and True:
figure(1)
clf()
imshow(rot90(-IntensityMat),origin='lower',extent=extentxy,aspect="auto")
hot()
plot(Xgrid[:,1],w(Xgrid[:,1]),'k--')
plot(Xgrid[:,1],-w(Xgrid[:,1]),'k--')
title('FORT beam intensity')
xlabel("x{\small [m]}")
ylabel("y{\small [m]}")
show()
figure(2)
clf()
print("Intensity at center of trap: %.3e W/m^2" % Intensity(O) )
#!
#! Depth of the potential
#!----------------------------------------------------------------------------
# The pulsations:
omega = lambda l: 2*pi*c_0/l
omega_L = omega(lambda_L)
omega_Rb = omega(lambda_Rb)
omega_Rb_e = omega(lambda_Rb_e)
omega_g = omega(lambda_g)
omega_e = omega(lambda_e)
# ----------------------------------------------------------------------
# The potential for Rb:
""" Dipolar potential depth as a function of laser wavelengths, einstein
coefficients, and transition wavelengths """
u_dip = lambda l, A, l_0: sum(
-3*pi* c_0**2/(2*omega(l_0)**3) * (
A/(omega(l_0)-omega(l)) + A/(omega(l_0)+omega(l))
)) / (2*pi*h_bar)
gamma_sc = lambda l, A, l_0: sum(
3*pi* c_0**2/(2*h_bar*omega(l_0)**3) * (l_0/l)**3 * (
A/(omega(l_0)-omega(l)) + A/(omega(l_0)+omega(l))
)**2 )
u_g = sum(-3*pi* c_0**2/(2*omega_g**3) * ( A_g/(omega_g-omega_L) + A_g/(omega_g+omega_L) ))
U_g = lambda r : u_g * Intensity( r)
print("Ground state: Trap depth: %0.3e K" % (U_g(O)/k_J) +
"\t Detuning: %.3e Hz" % (U_g(O)/(2*pi*h_bar)) )
u_e = sum(-3*pi* c_0**2/(2*omega_e**3) * ( A_e/(omega_e-omega_L) + A_e/(omega_e+omega_L) ))
U_e = lambda r : u_e * Intensity( r)
print("Excited state: Trap depth: %.3e K" % (U_e(O)/k_J) +
"\t Detuning: %.3e Hz" % (U_e(O)/(2*pi*h_bar)) )
# Now transform delta_1 in a detuning relative to the max detuning:
delta_1=delta_1+(U_e(O)-U_g(O))/(2*pi*h_bar)
print("Detuning of the 2nd laser to the resonance: %.3e MHz" % (delta_1/1.0e6))
#!
#! MOT
#!############################################################################
# Wave-vectors of the different lasers:
kx = array( [ 1, 0, 0 ])
ky = array( [ 0, 1, 0 ])
kz = array( [ 0, 0, 1 ])
Kx = -kx
Ky = -ky
Kz = -kz
lasers = ( kx, Kx, ky, Ky, kz, Kz )
#!
#! Detuning
#!----------------------------------------------------------------------------
# Hack
delta_0_outside = delta_0
def delta(r,v,k,delta_0=None):
""" Detuning """
if not delta_0:
delta_0 = delta_0_outside
return(2*pi*(delta_0 + mu_Rb*Bx*dot(r, transpose(k))
+ (U_g(r) - U_e(r))/(2*pi*h_bar)
+ dot(v, transpose(k))/lambda_Rb)/(2*Gamma_Rb))
if main and True:
detuningMap=-(delta(positions,zeros(positions.shape),kx) +
delta(positions,zeros(positions.shape),Kx))/2
detuningMat=reshape(detuningMap,(Xgrid.shape))
imshow(rot90(-2*detuningMat), origin='lower', extent=extentxy,
aspect="auto", cmap=cmap_xmap(lambda x: 1-x, cm.bone))
co=colorbar()
co.set_label('Detuning, in units of $\Gamma_{Rb}$')
cset = contour(rot90(detuningMat), array([0,Gamma_Rb/2,-Gamma_Rb/2]),
cmap=cm.bone, origin='lower', linewidths=1, extent=extentxy,
aspect="auto")
#clabel(cset, inline=1, fmt='%1.1f', fontsize=10, colors='black')
title(r'Detuning, in units of $\Gamma_{Rb}$')
xlabel("x{\small [m]}")
ylabel("y{\small [m]}")
show()
#! Forces on the atoms
#!----------------------------------------------------------------------------
# saturation parameter :
s_0 = P_mot / (pi* (d_mot/2)**2 ) * 1 / I_sat_Rb
print("Saturation parameter at resonnance: %g " % s_0)
s_1 = P_mot_2 / (pi* (d_mot/2)**2 ) * 1 / I_sat_Rb
print("Saturation parameter at resonnance for second frequency: %g " % s_1)
# ----------------------------------------------------------------------
# capture velocity :
#v_capt = ( -2*pi*delta_0 + mu_Rb*Bx*d_mot) * lambda_Rb/(2*pi)
v_capt = mu_Rb*Bx*5e-3 * lambda_Rb/(2*pi)
print "Capture velocity: %.3e m/s" % v_capt
v_capt = h_bar*Gamma_Rb*pi / (lambda_Rb**2*m_Rb*mu_Rb*Bx)* s_0/(1+s_0)
print "Capture velocity of the MOT: %.3e m/s" % v_capt
# ----------------------------------------------------------------------
# excited level population :
def Pop_e(r, v, k):
global delta_1
return 0.5*( s_0/(1 + s_0 + delta(r, v, k)**2 )
+ s_1/(1 + s_1 + delta(r, v, k, delta_0=delta_1)**2 ) )
# ----------------------------------------------------------------------
# radiation pressure force
F_mot = lambda r, v: - h_bar * Gamma_Rb * 2*pi/lambda_Rb *(
sum(
dstack([ dot(reshape(Pop_e(r, v, k),(-1,1)),reshape(k,(1,-1)))
for k in lasers])
, axis=-1 ))
# ----------------------------------------------------------------------
# dipolar force of the FORT
F_fort = lambda r, v: (
reshape((u_g + (u_e-u_g)
* sum(dstack([ Pop_e(r, v, k) for k in lasers ]), axis=-1
)), (-1,1))*NablaIntensity(r)
)
# ----------------------------------------------------------------------
# dipolar force of the MOT beams
F_dipMOT = lambda r, v: (
(u_e-u_g)*1/pi*NablaIntensity(r) * reshape(
sum(dstack([
1/(1+s_0+delta(r,v,k)**2)*(1-1/(1+delta(r,v,k)**2))
+ 1/(1+s_1+delta(r,v,k,delta_0=delta_1)**2)*(1-1/(1+delta(r,v,k,delta_0=delta_1)**2))
for k in lasers]), axis=-1), (-1,1))
)
# FIXME: I think I found a mistake !!
# !!!!!!!!!!!!!!!!! Why isn't delta_O in the prefix of this equation !!!!!!!!
#This need to be checked with my notes
F_dipMOT = lambda r, v: (
(u_e-u_g)*1/pi*NablaIntensity(r) * reshape(
sum(dstack([
1/(1+s_0+delta(r,v,k)**2)*(1-1/(1+delta(r,v,k)**2))
+ 1/(1+s_1+delta(r,v,k,delta_0=delta_1)**2)*(1-1/(1+delta(r,v,k,delta_0=delta_1)**2))
for k in lasers]), axis=-1), (-1,1))
)
# ----------------------------------------------------------------------
# total force
F_tot = lambda r, v: F_mot(r,v) + F_dipMOT(r,v) + F_fort(r,v)
# Initial conditions to study MOT capture
y0_capture = map(lambda z : array([0., -yrange, 0., 0., z, 0.]), arange(2,48,5))
# ----------------------------------------------------------------------
def plot_phase_space(traj=False,
y0list=[ array([0., 0.0005, 0., 0., 0.2, 0.]),
array([0., 0.0001, 0., 0., 0. , 0.]),
array([0., 0.0007, 0., 0., 2. , 0.]),
],
plot_traj=False):
# ----------------------------------------------------------------------
# color plot of the intensity of the force
positions=column_stack((zeros(Ygrid2.size), Ygrid2.ravel(),
zeros(Ygrid2.size)))
velocities=column_stack((zeros(Vgrid.size), Vgrid.ravel(),
zeros(Vgrid.size)))
RVforceMap = F_tot(positions, velocities) / m_Rb
YVforceMat = reshape(RVforceMap[:,1],(Ygrid2.shape))
figure(3)
clf()
imshow(rot90(YVforceMat), cmap=light_jet, extent=extentyv, aspect="auto")
null_format = FuncFormatter(lambda x,y: '')
co = colorbar()
co.set_label('Force intensity, {\small [arb. units]}.')
# ----------------------------------------------------------------------
# plot of the resonance line for the 4 lasers
detuningMap = delta(positions, velocities, ky)
detuningMat = reshape(detuningMap,(Ygrid2.shape))
cset = contour(transpose(detuningMat), array([0]),
colors=((1,1,1),), origin='lower', linewidths=1,
extent=extentyv, aspect="auto")
detuningMap = delta(positions, velocities, Ky)
detuningMat = reshape(detuningMap,(Ygrid2.shape))
cset = contour(transpose(detuningMat), array([0]),
colors=((1,1,1),), origin='lower', linewidths=1,
extent=extentyv, aspect="auto")
detuningMap = delta(positions, velocities, ky, delta_0=delta_1)
detuningMat = reshape(detuningMap,(Ygrid2.shape))
cset = contour(transpose(detuningMat), array([0]),
colors=((1,1,1),), origin='lower', linewidths=1,
extent=extentyv, aspect="auto")
detuningMap = delta(positions, velocities, Ky, delta_0=delta_1)
detuningMat = reshape(detuningMap,(Ygrid2.shape))
cset = contour(transpose(detuningMat), array([0]),
colors=((1,1,1),), origin='lower', linewidths=1,
extent=extentyv, aspect="auto")
# ----------------------------------------------------------------------
# vector plot of the dynamical flow
positionssparse=column_stack((zeros(Ygrid2sparse.size),
Ygrid2sparse.ravel(), zeros(Ygrid2sparse.size)))
velocitiessparse=column_stack((zeros(Vgridsparse.size),
Vgridsparse.ravel(), zeros(Vgridsparse.size)))
RVforceMapsparse = F_tot(positionssparse, velocitiessparse)/m_Rb
YVforceMatsparse = reshape(RVforceMapsparse[:,1], (Ygrid2sparse.shape))
forceScale=YVforceMatsparse.max()
# Autoscaling fails, here. I have to do it myself
quiver2(Ygrid2sparse,Vgridsparse+0.001, (Vgridsparse+0.001)/(2*vrange),
YVforceMatsparse/(2*forceScale),width=0.0015, color=(0.3,0.3,0.3))
title('Amplitude of the MOT force in the y:v plane')
xlabel('y{\small [m]}')
ylabel('v{\small [m s$^{-1}$]}')
# ----------------------------------------------------------------------
# integration of trajectories
if traj:
def deriv(y,t):
""" Computes the derivative of y=hstack((r,v)) at time t """
return hstack(( y[3:],
F_tot(reshape(y[:3],(1,-1)), reshape(y[3:],(1,-1)))[0,:]/m_Rb ))
# Integration parameters
start=0
end=0.009
numsteps=4000
time=linspace(start,end,numsteps)
from scipy import integrate
for y0 in y0list:
y=integrate.odeint(deriv,y0.copy(),time)
if y[:, 1].max() > yrange:
# The trajectory leaves MOT region:
pathcolor(y[:,1], y[:,4], time, linewidth=1,
linestyle='dotted')
else:
pathcolor(y[:,1], y[:,4], time, linewidth=2)
show()
y0 = array([ 0.002, 0.0008, 0.004, 0., 7., -3.])
time=linspace(start,2*end,numsteps)
if plot_traj:
figure(4)
y=integrate.odeint(deriv,y0,time)
delta_mean = lambda x : -(delta(x,zeros_like(x),kx)
+ delta(x,zeros_like(x),Kx))/2
trajplot(y,time,delta_mean)
show()
else:
show()
# ----------------------------------------------------------------------
# photon scattering rate (# photons/sec):
ph_rate = lambda r, v: Gamma_Rb * reshape(sum(dstack(
[ Pop_e(r, v, k) for k in lasers ]
), axis=-1), (-1, 1))
#! Heating, cooling and temperature
#!----------------------------------------------------------------------------
# ----------------------------------------------------------------------
# damping coefficient in the y direction:
damping_coef_y = lambda r, v: -h_bar*Gamma_Rb*2*pi/(m_Rb*lambda_Rb)*(
s_0*delta(r, v, ky)/(1 + s_0 + delta(r, v, ky)**2 )* pi/(lambda_Rb*Gamma_Rb)
+ s_1*delta(r, v, ky, delta_0=delta_1)/(1 + s_0 +
delta(r, v, ky, delta_0=delta_1)**2 )
* pi/(lambda_Rb*Gamma_Rb)
+ s_0*delta(r, v, Ky)/(1 + s_0 + delta(r, v, Ky)**2 )* pi/(lambda_Rb*Gamma_Rb)
+ s_1*delta(r, v, Ky, delta_0=delta_1)/(1 + s_0 +
delta(r, v, Ky, delta_0=delta_1)**2 )
* pi/(lambda_Rb*Gamma_Rb)
)
# ----------------------------------------------------------------------
# inverse equilibrium (?) temperature (compute inverse to avoid overflows):
inv_temperature = lambda r, v: ( k_J*2/(m_Rb*(h_bar*2*pi/(lambda_Rb*m_Rb))**2 )
* damping_coef_y(r, v)/ph_rate(r, v) )
if main and True:
yrange2 = yrange
# ----------------------------------------------------------------------
# color plot of the scattering rate, damping coefficient, and co.
[Ygrid3,Vgrid3] = mgrid[-yrange2:(yrange2+dy/8.):dy/8.,
-vrange:(vrange+dv):dv]
positions=column_stack((zeros(Ygrid3.size), Ygrid3.ravel(),
zeros(Ygrid3.size)))
velocities=column_stack((zeros(Vgrid3.size), Vgrid3.ravel(),
zeros(Vgrid3.size)))
figure(5)
clf()
subplot(3,1,1)
ph_scat_map = ph_rate(positions, velocities)
ph_scat_mat = reshape(ph_scat_map, (Ygrid3.shape))
imshow(rot90(ph_scat_map), cmap=cm.hot, extent=extentyv,
aspect="auto")
co=colorbar(fraction=0.05, aspect=7.)
#co.set_label('Photon scattering rate {\small [ph s$^{-1}$]}')
title('Photon scattering rate {\small [ph s$^{-1}$]}')
#xlabel('y{\small [m]}')
xticks(xticks()[0], len(xticks()[1])*['',]) # Hide the labels
ylabel('v{\small [m s$^{-1}$]}')
subplot(3,1,2)
damp_coef_map = damping_coef_y(positions, velocities)
damp_coef_mat = reshape(damp_coef_map, (Ygrid3.shape))
imshow(rot90(damp_coef_mat), cmap=cm.jet_r, extent=extentyv,
aspect="auto")
null_format = FuncFormatter(lambda x,y: '')
co=colorbar(fraction=0.05, aspect=7.)
#co.set_label('Damping coefficient, {\small [arb. units]}')
cset = contour(rot90(damp_coef_mat), array([0]),
colors=('black', ) , origin='lower', linewidths=1, extent=extentyv,
aspect="auto")
title('Damping coefficient')
# Store the ticks for later use
standard_xticks = xticks()
standard_yticks = yticks()
xticks(xticks()[0], len(xticks()[0])*['', ])
xlabel(r'\strut \hskip 9cm y{\small [m]}')
ylabel('v{\small [m s$^{-1}$]}')
# ----------------------------------------------------------------------
# color plot of the temperature resulting of balance between heating and
# damping.
# Calling the inv_temperature function does not work, I am not to sure why,
# but it raises a MemoryError, so we do this by hand.
inv_temperature_mat = ( k_J*2/(m_Rb*(h_bar*2*pi/(lambda_Rb*m_Rb))**2 ) *
damp_coef_mat / ph_scat_mat )
temperature_mat = log(abs(1/inv_temperature_mat))/log(10)
if any(inv_temperature_mat<0):
# Define masked arrays for positive and negative temperatures.
temperature_pos = ma.array(rot90(temperature_mat),
mask=rot90(inv_temperature_mat>0))
temperature_neg = -ma.array(rot90(temperature_mat),
mask=rot90(inv_temperature_mat<0))
last_ax=subplot(3,1,3)
pcolor(temperature_neg[:-1,:], shading='flat', cmap=cm.Blues_r)
# Store the image, to be able to give it a colorbar, later on
pos_im = gci()
pcolor(temperature_pos[:-1,:], shading='flat', cmap=cm.hot_r)
neg_im = gci()
xticks(xticks()[0][::2], map(str, standard_xticks[0]))
yticks(yticks()[0], map(str, standard_yticks[0]))
#xlabel('y{\small [m]}')
ylabel('v{\small [m s$^{-1}$]}')
title('Temperature limit')
### Double colorbar ###
# A phony colorbar just to shape these axis like to others:
co = colorbar(fraction=0.05, aspect=7.)
gcf().axes.pop()
axes(last_ax)
# Create a phony colorbar just for the sake of retrieving the axis info:
co_neg = colorbar(orientation='horizontal', pad= 0.2)
co_x, co_y, co_w, co_h = co_neg.ax.get_position()
gcf().axes.pop()
co_neg_axes = axes( [co_x, co_y, co_w/2., min(co_h, 0.04) ] )
neg_format = FuncFormatter(lambda x, y: '%.1f' % ((-10**(-x))*1000))
co_neg = colorbar(pos_im, co_neg_axes, orientation='horizontal',
format=neg_format)
co_neg.set_label('Temperature {\small [mK]}')
co_pos_axes = axes( [co_x + co_w/2., co_y, co_w/2., min(co_h, 0.04) ] )
pos_format = FuncFormatter(lambda x, y: '%.2f' % ((10**x)*1000))
co_pos = colorbar(neg_im, co_pos_axes, orientation='horizontal',
format=pos_format)
co_pos.set_label('Imaginary temperature {\small [mK]}')
else:
last_ax = subplot(3,1,3)
neg_im = imshow(rot90(temperature_mat), extent=extentyv, aspect="auto",
origin="lower", cmap=cm.Blues_r)
# A phony colorbar just to shape these axis like to others:
co = colorbar(fraction=0.05, aspect=7.)
gcf().axes.pop()
axes(last_ax)
neg_format = FuncFormatter(lambda x, y: '%d' % ((10**(-x))*1000))
co_neg = colorbar(neg_im, orientation='horizontal', pad= 0.2,
format=neg_format)
co_neg.set_label('Temperature {\small [mK]}')
show()
if main and True:
plot_phase_space()
More information about the SciPy-User
mailing list