#! /usr/bin/python import matplotlib.pyplot as plt from matplotlib.ticker import MaxNLocator import numpy as np from numpy import ma import re import sys ########################################################## fileList = [] if len(sys.argv) > 1: plotType = sys.argv[1] if plotType == 'bmad': fileList = ['ring_ma2.bmad.out'] if plotType == 'meas': fileList = ['ring_ma2.meas.out'] else: # if no input, default to showing both. filelist = ['ring_ma2.bmad.out', 'ring_ma2.meas.out'] ####################################################### # for plotting: plotBySeed = False plotHist = True dxtick = 10 #plotCorrections = [0] plotCorrections = [0,1,2,3] allMeas = [10.8,0.013,0.004] # emit.y, RMS eta.y, RMS cbar12 ####################################################### # Parameters for data plotting: markersize = 2 markertype = 'd' colorList = ['r','b','g','k'] dataType = ['Bmad', 'Meas', '(Meas-Bmad)'] nbins = 25 # for histograms yticks = 6 xtickFontSize = 15 ytickFontSize = 15 titlesize = 15 ######################################################### #titles = ['emit.x [nm]', 'emit.y [pm]', 'x [m]','y [m]','$\eta_x$ [m]','$\eta_y$ [m]','d$\phi_x$ [rad]','d$\phi_y$ [rad]','Cbar22','Cbar11','Cbar12'] titles = ['$\epsilon_y$ [pm]','$\eta_y$ [m]','Cbar12'] ################## # figure prep: xsize = 8 ysize = 8 thedpi = 100 left = 0.10 right = 0.95 top = 0.95 bottom = 0.08 hspace = .35 #xsize = 19.2 #ysize = 10.8 #thedpi = 100 #left = 0.06 #right = 0.95 #top = 0.95 #bottom = 0.05 #hspace = .8 # Prep figures for seed-by-seed plots: #if plotBySeed == True: # fig1 = plt.figure(figsize=(xsize,ysize),dpi=thedpi) # fig1.subplots_adjust(left=left,right=right,top=top,bottom=bottom, hspace=hspace) #fig2 = plt.figure(figsize=(xsize,ysize),dpi=thedpi) #fig2.subplots_adjust(left=left,right=right,top=top,bottom=bottom, hspace=hspace) #fig3 = plt.figure(figsize=(xsize,ysize),dpi=thedpi) #fig3.subplots_adjust(left=left,right=right,top=top,bottom=bottom, hspace=hspace) # Prep figures for histograms: if plotHist == True: fig4 = plt.figure(figsize=(xsize,ysize),dpi=thedpi) fig4.subplots_adjust(left=left,right=right,top=top,bottom=bottom, hspace=hspace) #fig5 = plt.figure(figsize=(xsize,ysize),dpi=thedpi) #fig5.subplots_adjust(left=left,right=right,top=top,bottom=bottom, hspace=hspace) #fig6 = plt.figure(figsize=(xsize,ysize),dpi=thedpi) #fig6.subplots_adjust(left=left,right=right,top=top,bottom=bottom, hspace=hspace) ################### # parse files: for kx in range(len(fileList)): # loop over all 3 files # prep data structures: x = [] y = [] phix = [] phiy = [] etax = [] etay = [] cbar22 = [] cbar12 = [] cbar11 = [] emitx = [] emity = [] seedIter = [] failedSeeds = [] jobFailed = False file = open(fileList[kx]) for line in file: fileline = line.split() if len(fileline) < 2: continue if fileline[1] == 'Summary': break if fileline[0] == '#': continue # a few checks for failed jobs: for jx in range(2,14): if float(fileline[jx]) < 0.: jobFailed = True # can't have negative RMS or emittance if len(emitx) > 0 and (len(emitx[-1]) < len(plotCorrections)) and (int(fileline[1]) == 0): jobFailed = True # didn't run through all corrections if jobFailed == True: if len(emitx) not in failedSeeds: #print "Correction ", len(emitx), "failed." failedSeeds.append(len(emitx)) while len(emitx[-1]) < len(plotCorrections): emitx[-1].append(-1.) emity[-1].append(-1.) x[-1].append(-1.) y[-1].append(-1.) etax[-1].append(-1.) etay[-1].append(-1.) cbar22[-1].append(-1.) cbar12[-1].append(-1.) cbar11[-1].append(-1.) phix[-1].append(-1.) phiy[-1].append(-1.) if int(fileline[1]) == 0: # new seed started jobFailed = False # reset x.append([]) y.append([]) phix.append([]) phiy.append([]) etax.append([]) etay.append([]) cbar22.append([]) cbar12.append([]) cbar11.append([]) emitx.append([]) emity.append([]) seedIter.append(int(fileline[0])) if jobFailed == False: emitx[-1].append(1.e9*float(fileline[2])) emity[-1].append(1.e12*float(fileline[3])) # skip emitz = float(fileline[4]) x[-1].append(float(fileline[5])) y[-1].append(float(fileline[6])) etax[-1].append(float(fileline[7])) etay[-1].append(float(fileline[8])) cbar22[-1].append(float(fileline[9])) cbar12[-1].append(float(fileline[10])) cbar11[-1].append(float(fileline[11])) phix[-1].append(float(fileline[12])) phiy[-1].append(float(fileline[13])) file.close() # check if last seed is completed (if job is still running) if len(emitx[-1]) < len(plotCorrections): print "Seed #", len(emitx), "incomplete-- masking." failedSeeds.append(len(emitx)) while len(emitx[-1]) < len(plotCorrections): emitx[-1].append(-1.) emity[-1].append(-1.) x[-1].append(-1.) y[-1].append(-1.) etax[-1].append(-1.) etay[-1].append(-1.) cbar22[-1].append(-1.) cbar12[-1].append(-1.) cbar11[-1].append(-1.) phix[-1].append(-1.) phiy[-1].append(-1.) if len(failedSeeds) > 0: #print "Masked seeds: ", failedSeeds print len(failedSeeds), "seeds failed." else: print "All seeds ok!" ######################################################### # Data post-processing: convert to np.arrays, and mask # any iterations NOT being plotted: Ax = ma.array(x) Ay = ma.array(y) Aphix = ma.array(phix) Aphiy = ma.array(phiy) Aetax = ma.array(etax) Aetay = ma.array(etay) Acbar22 = ma.array(cbar22) Acbar12 = ma.array(cbar12) Acbar11 = ma.array(cbar11) Aemitx = ma.array(emitx) Aemity = ma.array(emity) #allData = [Aemitx,Aemity,Ax,Ay,Aetax,Aetay,Aphix,Aphiy,Acbar22,Acbar11,Acbar12] allData = [Aemity, Aetay, Acbar12] # unmask only the data we want for data in allData: for ix in range(len(data[0,:])): if ix not in plotCorrections: data[:,ix] = ma.masked for ix in failedSeeds: data[ix-1] = ma.masked ######################## # Plot: # Seed-by-seed: # if plotBySeed == True: # jx = 0 # for data in allData: # iterate over plots on one page: # jx = jx + 1 # # determine which plot to put data on: # #if kx == 0: # ax1 = fig1.add_subplot(6,2,jx) # #elif kx == 1: # # ax1 = fig2.add_subplot(6,2,jx) # #elif kx == 2: # # ax1 = fig3.add_subplot(6,2,jx) # # ulim = 0 # for plot ymax # for ix in plotCorrections: # thisLabel = 'iter '+str(ix) # ax1.plot(seedIter, data[:,ix], color=colorList[ix], marker=markertype, ms=markersize, label=thisLabel) # thisulim = np.amax(data[:,ix]) # if thisulim > ulim: ulim = thisulim # # # ax1.grid() # ax1.set_ylabel(titles[jx-1]) # ax1.set_xlabel('Seed Index') # ax1.set_xticks(range(0,len(seedIter)+1,dxtick)) # ax1.set_title(dataType[kx]+': '+titles[jx-1]) # ax1.legend(prop={'size':8}) # ax1.yaxis.set_major_locator(MaxNLocator(yticks-1)) # # set tick fontsize: # for tick in ax1.xaxis.get_major_ticks(): # tick.label.set_fontsize(xtickFontSize) # for tick in ax1.yaxis.get_major_ticks(): # tick.label.set_fontsize(ytickFontSize) # Histograms: if plotHist == True: jx = 0 for data in allData: jx = jx + 1 # determine which plot to put data on: if kx == 0: ax1 = fig4.add_subplot(3,1,jx) elif kx == 1: ax1 = fig5.add_subplot(3,1,jx) elif kx == 2: ax1 = fig6.add_subplot(3,1,jx) ylim = 0 for ix in plotCorrections: thisLabel = 'iter '+str(ix) n, bins, patches = ax1.hist(data[:,ix],bins=np.logspace(start=np.log10(np.amin(data)),stop=np.log10(np.amax(data)),num=nbins),color=colorList[ix],histtype='step', lw=2, align='mid',label=thisLabel) ymax = np.amax(n) ylim = np.amax([ymax, ylim]) ax1.set_xscale('log') ylim = ylim * 1.05 ax1.axvline(x=allMeas[jx-1],ymin=0,ymax=ylim, lw=2,ls='dashed',c='k') #ax1.grid() ax1.set_ylabel('# Seeds', fontsize=ytickFontSize) ax1.set_xlabel(titles[jx-1], fontsize=xtickFontSize) ax1.xaxis.labelpad = -.1 if ylim > 0: ax1.set_ylim(0,ylim) else: ax1.set_ylim(0,100) ax1.legend(prop={'size':8}) ax1.yaxis.set_major_locator(MaxNLocator(yticks-1)) # set tick fontsize: for tick in ax1.xaxis.get_major_ticks(): tick.label.set_fontsize(xtickFontSize) for tick in ax1.yaxis.get_major_ticks(): tick.label.set_fontsize(ytickFontSize) if len(fileList) == 1: outFile1 = fileList[kx][0:-3]+'png' outFile2 = fileList[kx][0:-3]+'eps' plt.savefig(outFile1) plt.savefig(outFile2) print "Plots saved as "+ outFile1 + ', '+ outFile2 if len(fileList) > 1: plt.show()