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()
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()
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()