##############################################################################
# Simulation of Optical Pumping
#
# Reinhard A. Schumacher (All rights reserved)
# Carnegie Mellon University
# March 2008
#
# Modifications:
# 3-13-2008 - branch off from NMR program
#
from visual import *
from random import random, randrange
from visual.graph import *
print """
##############################################################################
Simulation of Optical Pumping using a classical vector model.
------------------------------------------------------------------------------
A set of Rubidium-87 atoms precesses in a static  magnetic field. A perp-
endicular oscillating magnetic field causes the reorientation of the spins,
inducing a non-zero net magnetization in the plane perpendicular to the
static external field.

Quantum physics shows that the expectation value of the quantum spin operator
evolves in time in the same way as a classical spin vector.  Thus, this
simulation is a faithful representation of what the quantum system does.

The individual spins flash red for a moment when they have "interacted" with
a photon and been forced into one of the Z eigenstates according to the rule
that the light absorption increases m by one, and then light emmision changes
the m of the atom by 0, +1, or -1.

Controls:
    Left-clicking spheres on the main display screen activates given action.
    Left-clicking anywhere else pauses/resumes animation.
    Left-clicking on the "Net Polarization" panel clears the arrow "trail".

Adjustments:
    Inside the program, various parameters can be set:
numspin    = number of randomly placed, randomly oriented (in phi), spins
bfield0    = vector components of the strong external field
bfield1mag = magnitude of the rotating field
freq1      = oscillating field frequency
teeone     = interaction time scale on which optical pumping occurs
speed      = overall speed of the simulation, set depending on CPU speed

 R. A. Schumacher
 Carnegie Mellon University

Version 1.0
###########################################################################
"""
###########################################################################
# USER ADJUSTABLE INPUTS GO HERE:
numspin = 100           #100 number of randomly placed spins
bfield0 = vector(0,5,0) #0,5,0 components of the strong external field
bfield1mag = 2.0        #2.  magnitude of the 'rotating' field
freq1  = 10.            #10. oscillating field frequency
teeone = 5.             #5.  optical pumping relaxation time towards maximum spin
speed  = 20             #maximum update speed per second: depends on CPU power
###########################################################################

ampdata=[0.,0.,0.,0.,0.,0.]  #VPython starts enumerating lists at zero, hence 6 elements
scene1 = display(
    title = "Optical Pumping" ,
    width=600, height=600,
    autoscale=0, background=(.3,.4,.5),
    x=0, y=0, forward=(-0.0,-0.25,-1))
scene2 = display(title = "Net Polarization" , width=300, height=300,
    center=(0,0,0), autoscale=0, background=(.3,.4,.5),
    x=600, y=0, forward=(0,-.1,-1.), range=(4,4,4))
scene3 = gdisplay(title='Expectation', xtitle='Time', ytitle='Net Z Spin Projection',
    background=(.2,.2,.2),
    xmax=100.0, ymax=1.1, ymin=-1.1, x=600, y=300, width=300, height=225)
mt = gcurve(color=color.magenta)
#
# Hack in a "bar graph" of the magnetic sub-state amplitudes.  This
# awkward code was needed since VPython's built-in version of such a
# display does not work properly
#
scene4 = display(title="F=2 Sub-state Amplitudes",width=300,height=225,
                 autoscale = 0, background=(.3,.4,.5),forward=(0.,0.,1.),
                 x=600, y=525, range=(4,2.2,4))
scene4.lights=[vector(0.,0.2,-0.8)]
xaxis = curve(pos=[(-3,0,0), (3,0,0)], radius=0.01 ,color=color.blue)
line1 = curve(pos=[(-3, .5,0), (3, .5,0)], radius=0.006 ,color=color.white)
line2 = curve(pos=[(-3,-.5,0), (3,-.5,0)], radius=0.006 ,color=color.white)
line3 = curve(pos=[(-3,1.0,0), (3,1.0,0)], radius=0.006 ,color=color.white)
line4 = curve(pos=[(-3,-1.,0), (3,-1.,0)], radius=0.006 ,color=color.white)
yaxis = curve(pos=[( 0,-1,0),(0,1,0)], radius=0.02 ,color=color.blue)
label(pos=vector(0,0,0), text="-2         -1          0          1          2",
      xoffset=-3,yoffset=-5, height=12,opacity=0, box=0, line=0)
label(pos=vector(-2.8,0,0), text="m_F",
      xoffset=5,yoffset=-5, height=12,opacity=0, box=0, line=0)
label(pos=vector(0,1.0,0), text="Amplitude",
      xoffset=3,yoffset=2, height=12,opacity=0, box=0, line=0)
mstates1 = cylinder(pos=(-2,0,0), axis=(0,1,0),  radius=.3,color=color.red)
mstates2 = cylinder(pos=(-1,0,0), axis=(0,.5,0), radius=.3,color=color.red)
mstates3 = cylinder(pos=( 0,0,0), axis=(0,1,0),  radius=.3,color=color.red)
mstates4 = cylinder(pos=( 1,0,0), axis=(0,.5,0), radius=.3,color=color.red)
mstates5 = cylinder(pos=( 2,0,0), axis=(0,1,0),  radius=.3,color=color.red)

def wirecube (s,x,y,z):
    c=curve (color=color.white, radius=.05)
    pts = [(-s+x, -s+y, -s+z),(-s+x, -s+y, s+z), (-s+x, s+y, s+z),
           (-s+x, s+y, -s+z), (-s+x, -s+y, -s+z), (s+x, -s+y, -s+z),
           (s+x, s+y, -s+z), (-s+x, s+y, -s+z), (s+x, s+y, -s+z),
           (s+x, s+y, s+z), (-s+x, s+y, s+z), (s+x, s+y, s+z),
           (s+x, -s+y, s+z), (-s+x, -s+y, s+z), (s+x, -s+y, s+z),(s+x, -s+y, -s+z)]
    for pt in pts:
        c.append(pos=pt)
#
# Set up the visual scene
#
scene1.select()  #makes first scene the current default for new objects
wirecube(5,0,0,0)
oring1 =ring(pos=(6,0,0),  axis=(1,0,0), radius=6, thickness=0.1, color=color.red)
oring2 =ring(pos=(-6,0,0), axis=(1,0,0), radius=6, thickness=0.1, color=color.red)
oring1.visible = 0
oring2.visible = 0
detector = cylinder(pos=(0,7,0), axis=(0,0.5,0), radius=3,opacity=.2,color=color.blue)
label(pos=vector(0,7,0), text="Detector",
      xoffset=0,yoffset=0,zoffset=4, height=12,opacity=0, box=0, line=0)

scene2.select()
wirecube(2,0,0,0)

scene1.select()
base = box(pos=vector(0,-5,0),height=0.1, length=10.0, width=10.0,
           color=color.yellow) #color=(0.0,0.5,1.0))
s_reset=sphere(pos=(-5,5,5), radius=0.25, color=(0.6, 1.0, 0.6))
s_reset_label=label(pos=s_reset.pos, text="Reset Spins", yoffset=+10,
                    height=12,opacity=0, box=0, line=0)
s_osc=sphere(pos=(5,5,5), radius=0.2, color=(1.0, 0.0, 0.0))
s_osc_label=label(pos=s_osc.pos, text="Oscillating Field ON/OFF", yoffset=+10,
                    height=12,opacity=0, box=0, line=0)
#s_dbdr=sphere(pos=(-5,-5,5), radius=0.25, color=(1.0, 0.0, 0.0))
#s_dbdr_label=label(pos=s_dbdr.pos, text="Gradients ON/OFF", yoffset=-10,
#                    height=12,opacity=0, box=0, line=0)
s_freq1dn=sphere(pos=(-5,0,5), radius=0.25, color=(0.6, 1.0, 0.6))
s_freq1dn_label=label(pos=s_freq1dn.pos, text="Lower Freq.", xoffset=-10,
                    height=12,opacity=0, box=0, line=0)
s_freq1up=sphere(pos=(5,0,5), radius=0.25, color=(0.6, 1.0, 0.6))
s_freq1dn_label=label(pos=s_freq1up.pos, text="Raise Freq.", xoffset=10,
                    height=12,opacity=0, box=0, line=0)
s_t1onoff=sphere(pos=(-5,-5,5), radius=0.25, color=(0.0, 1.0, 0.0))
s_t1onoff_label=label(pos=s_t1onoff.pos, text="Pumping ON/OFF", yoffset=-10,
                    height=12,opacity=0, box=0, line=0)
#s_t2onoff=sphere(pos=(0,-5,5), radius=0.25, color=(1.0, 0.0, 0.0))
#s_t2onoff_label=label(pos=s_t2onoff.pos, text="T2 Relax. ON/OFF", yoffset=-10,
#                    height=12,opacity=0, box=0, line=0)

Spins      = []
Dots       = []
dera       = pi/180.
polarangle = 2    #a hack...to make the arrows rock a bit
magsum     = 0
rescale    = 0.5  #rescale vector length such that they are unit length
for i in range(numspin):
    xpos = randrange(-500.,500.)/100.
    ypos = randrange(-500.,500.)/100.
    zpos = randrange(-500.,500.)/100.
    phiangle = randrange(0.,360.)
    spinmag = randrange(-2,3) #pick a random z eigenstate to start...
#    print i,spinmag
#    magsum = magsum + spinmag
    if i==1:
        xpos=0.
        ypos=0.
        zpos=0.
        phiangle = 0.
    # (VPython forces us to accept "up" as the +Y direction)
    xval = spinmag*sin(polarangle*dera)*cos(phiangle*dera)*rescale
    zval = spinmag*sin(polarangle*dera)*sin(phiangle*dera)*rescale
    yval = spinmag*cos(polarangle*dera)*rescale
    spinvector = vector (xval,yval,zval)
    posvector  = vector (xpos,ypos,zpos)
    Spins = Spins+[arrow(pos=posvector,
                axis=spinvector, shaftwidth=0.10, color=color.green)]
    Dots  = Dots+[sphere(pos=posvector,radius=0.10,color=color.blue)]
#
# Draw permanent arrows representing magnetic field in space.
#
bfield = bfield0 
btail  = -bfield.mag / 2.
bleft  = arrow(pos=(-5,btail,0), axis=bfield, shaftwidth=0.05, color=color.blue)
bfield = bfield0 
btail  = -bfield.mag / 2.
brght  = arrow(pos=( 5,btail,0), axis=bfield, shaftwidth=0.05, color=color.blue)
btail  = -bfield0.mag / 2.
bback  = arrow(pos=( 0,btail,-5),axis=bfield0, shaftwidth=0.05, color=color.blue)
bfrnt  = arrow(pos=( 0,btail,5), axis=bfield0, shaftwidth=0.05, color=color.blue)
#
brot   = arrow(pos=( 0,0,0), axis=(bfield1mag,0,0), shaftwidth=0.05, color=color.cyan,
               visible=0,display=scene1)
bosc   = arrow(pos=( 0,0,0), axis=(bfield1mag,0,0), shaftwidth=0.05,
               color=color.magenta, visible=0,display=scene1)
bosc2  = arrow(pos=( 0,0,0), axis=(bfield1mag,0,0), shaftwidth=0.05,
               color=color.magenta, display=scene2, visible=0)
netmag = arrow(pos=(0,0,0), axis=(0,0,0), shaftwidth=0.2,
               color=color.green, display=scene2)
trail  = curve(display=scene2)  #follows the net magentization in scene2
btrail = curve()                #follows the rotating field vector
tmark = sphere(pos = (-5,-5,5.2), radius = 0.05, color = color.red, visible=0 )
#
# Set up the physics
#
tflag     = 0
opfact    = 1       #start with the pumping turned on
gamma     = 1.0
magmoment = 2.0
ratio     = magmoment/gamma
freq0     = ratio*bfield0.mag       #Larmor frequency at center of space
t  = 0.     #global time
t1 = 0.     #time since rotating field was turned on
dt = 0.03   #time step
print "Larmor frequency=",freq0,"cyles/time, Rotating field frequency=",freq1,"cycles/time"
while 1:
    rate(speed)  #VPython's maximum update rate, in inverse seconds
    t=t+dt
    if tflag==1:            #this flag says the rotating field is "ON"
        t1 = t1 + dt
        angle1 = freq1*t1   #angle of rotating B-field in horizontal plane
        brot.axis.x =  bfield1mag*cos(angle1)  #effective rotating field
        brot.axis.z = -bfield1mag*sin(angle1)
        bosc.axis.x =  brot.axis.x             #actual applied oscillating field...
        bosc2.axis.x=  brot.axis.x             #...duplicate for 2nd scene
        bfield = bfield0 + 0.5*bosc.axis
        if angle1<2.5*pi:
            btrail.append(pos=brot.axis)
    elif tflag==0:
        bfield = bfield0
        
    sumx = 0
    sumy = 0
    sumz = 0
    for i in range(numspin):
        xxx   = Spins[i].pos.x
        yyy   = Spins[i].pos.y
        zzz   = Spins[i].pos.z
        btot  = bfield
        danglerad = ratio*(btot.mag)*dt #increment in precession angle
        Spins[i].rotate(angle=danglerad, axis=btot)
#
#       Simulate the effect of optical pumping wherein the "Z" eigestate
#       rachets upward in steps of 0, +1, or +2.  After the first interaction,
#       assume that the spin is in an eigenstate of the z-basis.
#
        sxaxis = Spins[i].axis.x         
        syaxis = Spins[i].axis.y         
        szaxis = Spins[i].axis.z
        if random() < dt/teeone*opfact:   #stochastic "hit" on this spin
            dm = randrange(0,3)           #random integer 0,1,2
            if syaxis >= -0.5:
                if syaxis >= 0.0:
                    if syaxis >= 0.5:
                        syaxis = 1.0
                    else:
                        syaxis = 0.0
                else:
                    syaxis = -0.5
            else:
                syaxis = -1.0              
            syaxis = syaxis + dm/2.  #divide by 2 since visual length is "1" not "2"
#            sxaxis = 0.01 #hack to make spins wobble a bit
#            szaxis = 0.01
            cflag = 1   #make this spin red (temporarily)
        else:
            cflag = 0   #make this spin green
        if syaxis>1:    #use "1" as the size scale regardless of total spin value
            syaxis = 1.
            cflag = 0
        if syaxis<-1:
            syaxis = -1.
            cflag = 0
        sxzmax = sqrt(1. - syaxis**2.)
        sxz    = sqrt(sxaxis**2. + szaxis**2.)
        if sxz > 0.:
            sizrat = sxzmax/sxz
        else:
            sizrat = 10.  #don't scale transverse projections
        if sizrat < 1 & opfact==1:
            sxaxis = sxaxis*sizrat
            szaxis = szaxis*sizrat
        Spins[i].axis.x = sxaxis
        Spins[i].axis.y = syaxis
        Spins[i].axis.z = szaxis

#        redfrac = 1. - (syaxis+1.)/2.
#        grnfrac = 1. - redfrac
#        blufrac = 0.
#        Spins[i].color=(redfrac,grnfrac,blufrac)
        if cflag==1:
            Spins[i].color = color.red
            Dots[i].color  = color.red
        else:
            Spins[i].color = color.green
            Dots[i].color  = color.blue

        sumx = sumx+sxaxis
        sumy = sumy+syaxis
        sumz = sumz+szaxis
#
#   Make the net magnetization arrow some blended color...
#
#    blufrac = 0.
#    redfrac = 1.-((sumy/numspin)+1.)/2.
#    grnfrac = 1.-redfrac
#    netmag.color =(redfrac,grnfrac,blufrac)
    netmag.axis  =(2*sumx/numspin,2*sumy/numspin,2*sumz/numspin)
    zzmag        = sumy/numspin   #recall that VPython's "y" is our "z"
    xymag        = sqrt(netmag.axis.x**2 + netmag.axis.z**2)/2.
    mt.plot(pos=(t,zzmag))
    trail.append(pos=netmag.axis)
#
#   Compute the magnetic sub-state amplitudes of a pure quantum J=2
#   state positioned at the angle of the net magnetization vector.
#   Use the Wigner D-functions for this calculation.
#
    theta = atan2(xymag,zzmag)  #polar angle from 0 to pi radians
    ctheta = cos(theta)
    stheta = sin(theta)
    ampdata[1] = ((1. + ctheta)/2.)**2.
    ampdata[2] = -(1. + ctheta)/2. *stheta
    ampdata[3] = sqrt(6.)/4.*stheta**2.
    ampdata[4] = -(1. - ctheta)/2. *stheta
    ampdata[5] = ((1. - ctheta)/2.)**2.
    mstates1.axis=vector(0.,ampdata[1],0.)*zzmag #really, the net z magnitude is NOT this
    mstates2.axis=vector(0.,ampdata[2],0.)*zzmag #for a pure state
    mstates3.axis=vector(0.,ampdata[3],0.)*zzmag
    mstates4.axis=vector(0.,ampdata[4],0.)*zzmag
    mstates5.axis=vector(0.,ampdata[5],0.)*zzmag

#
#   Mousing around...
#
    if scene1.mouse.events:
        m = scene1.mouse.getevent()
#
#       Click on a labeled sphere to get the stated action.
#
        if m.click:
            if m.pick is s_reset:
                print "Resetting spins..."
                for i in range(numspin):
                    Spins[i].axis.x = (2.*random()-1.)*0.01
                    Spins[i].axis.y = 1.0
                    Spins[i].axis.z = (2.*random()-1.)*0.01
                    Spins[i].color = color.green
                    trail.visible = 0
                    trail = curve(display=scene2)
                    mt.visible = 0
                    mt = gcurve(display=scene3,color=color.cyan)
                    btrail.visible = 0
                    btrail = curve()
                    t  = 0
                    t1 = 0
            elif m.pick is s_freq1up:
                freq1 = freq1 + 2.
                print "Raising oscillation frequency..."
                print "Larmor frequency=",freq0,"cyles/time, Rotating field frequency=",freq1,"cycles/time"
            elif m.pick is s_freq1dn:
                freq1 = freq1 - 2.
                print "Lowering oscillation frequency..."
                print "Larmor frequency=",freq0,"cyles/time, Rotating field frequency=",freq1,"cycles/time"
            elif m.pick is s_t1onoff:
                if opfact == 0:
                    opfact = 1
                    s_t1onoff.color = color.green
                    base.color = color.yellow
                    print "Optical pumping light is turned ON"
                else:
                    opfact = 0
                    s_t1onoff.color = color.red
                    base.color = (0.0,0.5,1.0)
                    print "Optical pumping light is turned OFF"
            elif m.pick is s_osc:
                if tflag ==0:
                    tmark.visible = 1
                    btrail = curve()
                    btrail.visible = 1
                    bosc.visible   = 1
                    bosc2.visible  = 1
                    brot.visible   = 1
                    oring1.visible = 1
                    oring2.visible = 1
                    tflag = 1
                    t1 = 0.  #start the rotating field at phi=0
                    s_osc.color = color.green
                    print "Rotating field turned ON"
                    scene1.title = 'RF magnetic field turned ON'
                else:
                    tmark.visible  = 0
                    btrail.visible = 0
                    bosc.visible   = 0
                    bosc2.visible  = 0
                    brot.visible   = 0
                    oring1.visible = 0
                    oring2.visible = 0
                    tflag = 0
                    s_osc.color = color.red
                    print "Rotating field turned OFF"
                    scene1.title = 'RF magnetic field turned OFF'
            else:
                print "Pausing  animation"
                scene1.mouse.getclick() #inhibits animation until the mouse is clicked
                print "Resuming animation"
    if scene2.mouse.events:
        m = scene2.mouse.getevent()
        print "Clearing trail"
        trail.visible = 0
        trail  = curve(display=scene2)  #follows the net magentization in scene2
#            else:
#                print "Pausing  animation"
#                scene2.mouse.getclick() #inhibits the calculation until the mouse is clicked
#                print "Resuming animation"

#  END OF CODE