[SciPy-user] Python evangelism question: procedural programming

Gael Varoquaux gael.varoquaux@normalesup....
Fri Aug 3 10:34:41 CDT 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