示例#1
0
文件: trapplot.py 项目: natl/bg
def aniphase(filename):
  '''
  Plot the simulation generated in the file filename, plotting only 
  the phase
  Filename is a binary hd5 file generated by twodsave.
  '''
  
  #Setup------------------------------------------------------------------------
  infile = ts.h5file( filename, erase = False, read = True )
  fig = plt.figure()
  matplotlib.rcParams.update({'font.size':16}) #set inital font size
  
  steps = infile.head[ 'steps' ]
  npt   = infile.head[ 'npt' ]
  
  xvals, yvals = infile.readxy() #read x-y array
  
  p = infile.readpsi('0.0') #read initial state
  
  #Plot phase-------------------------------------------------------------------
  
  ax = fig.add_subplot(111, aspect = 'equal')
  
  #plt.xlim( [ -5, 5 ] )
  #plt.ylim( [ -5, 5 ] )
  
  plt.xlim( [ -10, 10 ] )
  plt.ylim( [ -10, 10 ] )
  
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  pi = np.pi
  im = plt.imshow( np.angle( p ),
                    extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                    vmin = -np.pi, vmax = np.pi, cmap=cm.gist_rainbow ) 
  cbar = plt.colorbar(im,
                       ticks = [round(j,2) for j in np.linspace(-pi,pi,5)]
                      )
  
  
  #Update the plot--------------------------------------------------------------
  def update_fig(i):
    p = infile.readpsi( str(i) )
    
    im.set_array( np.angle( p ) )
        
    return im,

  ani = animation.FuncAnimation( fig, update_fig,
                                 np.arange( 1, steps + 1., dtype=float ),
                                 interval=20., repeat_delay = 3000., 
                                 blit = False )
  plt.show()
  return ani
  infile.f.close() 


  

  
示例#2
0
文件: trapplot.py 项目: natl/bg
def twodani(filename):
  '''
  Plot the simulation generated in the file filename.
  Filename is a binary hd5 file generated by twodsave.
  '''
  
  #Setup------------------------------------------------------------------------
  infile = ts.h5file( filename, erase = False, read = True )
  fig = plt.figure()
  
  steps = infile.head[ 'steps' ]
  npt   = infile.head[ 'npt' ]
  
  xvals, yvals = infile.readxy() #read x-y array
  
  p = infile.readpsi('0.0') #read initial state
  
  #Plot probability density-----------------------------------------------------
  
  ax = fig.add_subplot(221, aspect = 'equal')
  
  plt.xlim( [ -5, 5 ] )
  plt.ylim( [ -5, 5 ] )
  
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  im = plt.imshow( abs( p ) ** 2 / ( abs( p ) ** 2 ).max(),
                  extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                  vmin = 0., vmax = 1., cmap = nicecmap() ) #(abs(p)**2).max())
 
  cbar = plt.colorbar(im, ticks = np.linspace( 0, 1, 5 ) )
  
  #if not autoscaling density:
  #v = np.linspace(0, ( abs( p ) ** 2. ).max() , 10)
  #for jj in range(0,len(v)): v[jj] = round(v[jj],3)
  #cbar = plt.colorbar(ticks = v)
  
  #Phase------------------------------------------------------------------------
  
  ax2 = fig.add_subplot( 222, aspect = 'equal' )
  plt.xlim( [ -5, 5 ] )
  plt.ylim( [ -5, 5 ] )
 
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  pi = np.pi
  im2 = plt.imshow( np.angle( p ),
                    extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                    vmin = -np.pi, vmax = np.pi, cmap=cm.gist_rainbow ) 
  cbar2 = plt.colorbar(im2,
                       ticks = [round(j,2) for j in np.linspace(-pi,pi,5)]
                      )
  
  #--------------------------------------------------------
  
  #x = 0.
  ax3 = fig.add_subplot( 223 )
  plt.xlim( [-10,10] )
  plt.ylim( [0,1] )
 
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'$|\psi|^2$ (atoms $a_{0}^{-1}$)')
  
  line1, = ax3.plot( xvals, ( abs( p ) ** 2. )[ npt / 2., : ] / 
                          ( abs( p ) ** 2. )[ npt / 2., : ].max() )
    
  #--------------------------------------------------------
  
  #y = 0.
  
  ax4 = fig.add_subplot(224)
  plt.xlim( [-10,10] )
  plt.ylim( [0, 1] )
  
  plt.xlabel( r'y ($a_0$)' )
  plt.ylabel( r'$|\psi|^2$ (atoms $a_{0}^{-1}$)' )
  
  line2, = ax4.plot( yvals, ( abs( p ) **2. )[ : , npt / 2. ] / 
                          ( abs( p ) **2. )[ : , npt / 2.].max() )
  
  #Update the plot--------------------------------------------------------------
  def update_fig(i):
    p = infile.readpsi( str(i) )
    
    im.set_array( abs( p ) ** 2 / (abs(p)**2).max() )
    im2.set_array( np.angle( p ) )
    line1.set_ydata( ( abs( p ) ** 2. )[ npt / 2., : ] / 
                     ( abs( p ) ** 2. )[ npt / 2., : ].max() )
    line2.set_ydata( ( abs( p ) ** 2. )[ :, npt / 2. ] / 
                     ( abs( p ) ** 2. )[ :, npt / 2. ].max() )
    
    return im, im2, line1, line2

  ani = animation.FuncAnimation( fig, update_fig,
                                 np.arange( 1, steps + 1., dtype=float ),
                                 interval=20., repeat_delay = 3000., 
                                 blit = False )
  plt.show()
  return ani
  infile.f.close()
示例#3
0
文件: trapplot.py 项目: natl/bg
def aniprob( filename,
             speed = 1.):
  '''
  Plot the simulation generated in the file filename, plotting only 
  the probability density
  def aniprob( filename,
               speed = 1.):
  Filename is a binary hd5 file generated by twodsave.
  speed adjusts the frame rate ( lit. interval = 20ms/speed )
  '''
  
  #Setup------------------------------------------------------------------------
  infile = ts.h5file( filename, erase = False, read = True )
  fig = plt.figure()
  matplotlib.rcParams.update({'font.size':16}) #set inital font size
  
  steps = infile.head[ 'steps' ]
  npt   = infile.head[ 'npt' ]
  
  xvals, yvals = infile.readxy() #read x-y array
  
  p = infile.readpsi('0.0') #read initial state
  
  #Plot probability density-----------------------------------------------------
  
  ax = fig.add_subplot(111, aspect = 'equal')
  
  #plt.xlim( [ -5, 5 ] )
  #plt.ylim( [ -5, 5 ] )
  
  plt.xlim( [ -10, 10 ] )
  plt.ylim( [ -10, 10 ] )
  
  plt.xlabel(r'x ($a_0$)')
  plt.ylabel(r'y ($a_0$)')
  
  im = plt.imshow( abs( p ) ** 2 / ( abs( p ) ** 2 ).max(),
                  extent = ( min(xvals), max(xvals), min(yvals), max(yvals) ),
                  vmin = 0., vmax = 1., cmap = nicecmap() ) #(abs(p)**2).max())
 
  cbar = plt.colorbar(im, ticks = np.linspace( 0, 1, 5 ) )
  
  #if not autoscaling density:
  #v = np.linspace(0, ( abs( p ) ** 2. ).max() , 10)
  #for jj in range(0,len(v)): v[jj] = round(v[jj],3)
  #cbar = plt.colorbar(ticks = v)
  
  
  #Update the plot--------------------------------------------------------------
  def update_fig(i):
    p = infile.readpsi( str(i) )
    
    im.set_array( abs( p ) ** 2 / (abs(p)**2).max() )
        
    return im,

  ani = animation.FuncAnimation( fig, update_fig,
                                 np.arange( 1, steps + 1., dtype=float ),
                                 interval=20./speed, repeat_delay = 3000., 
                                 blit = False )
  plt.show()
  return ani
  infile.f.close()