-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathplot_nominal.py
152 lines (135 loc) · 6.4 KB
/
plot_nominal.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
# -*- coding: utf-8 -*-
"""
Script to make nominal plots
@author: Linda Stoel
"""
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import os
import sys
import python.dataprocessing as datproc
ap = datproc.nomap
#TODO python 2/3 compatibility
# Get command line args
folder = sys.argv[1]
try:
kind = sys.argv[2]
except IndexError:
kind = '01'
# Check for errors
failed, messages = datproc.errorcheck(folder+'/error')
print (str(len(failed))+" processes failed")
if not len(failed)==0:
print "Errors:"
for message in messages:
print message
print ""
# Get study name
name = folder.split("/")[-1]
if name=="":
name = folder.split("/")[-2]
# Make plot folder
plotfolder = folder+"/../plots/"+name
if not os.path.exists(plotfolder):
os.makedirs(plotfolder)
# Get settings
settings = datproc.getsettings(folder)
sliced = False if settings['slices']=='None' else True
pycoll = eval(settings['pycollimate'])
myloc = "AP.UP.TPST21760" if pycoll else "AP.UP.ZS21633"
if '0' in kind or '1' in kind:
# Get losses
losses = datproc.getlosses(folder+'/losses', settings=settings)
myloss = losses[(losses['ELEMENT']==myloc)]
if '0' in kind:
# Print efficiency and loss stats
if pycoll:
myap = {'aperturex': ap['tpstcirc']+ap['tpstblade']+np.array([0, ap['tpstex']]),
'aperturey': [-0.0100,0.0100]}
else:
myap = {'aperturex': ap['zsupmid']+ap['zsthick']/2+np.array([0, ap['zsex']]),
'aperturex2': ap['zsdomid']+ap['zsthick']/2+np.array([0, ap['zsex']]),
'aperturey': [-0.023,0.023],
'zs_len': 18.77,
'zs_an': 4.1635E-4}
stdout = sys.stdout
with open(plotfolder+"/stats.txt", 'w') as sys.stdout:
print 'Probabilities with 95% confidence intervals:'
alltags,_ = datproc.extr_tagger(data=losses, pycoll=pycoll, **myap)
datproc.efficiency(losses, alltags=alltags, silent=False)
print '\nLoss stats:'
datproc.lossstats(losses, merge=False, silent=False)
print 'Beam stats:'
emita =datproc.beamstats(losses[losses['tag']=='extracted'], silent=False)
print '\nBeam fit (99.99%):'
emitb = datproc.get_ellipse(losses[losses['tag']=='extracted'],
x0=emita['X0'], px0=emita['PX0'],
alpha0=emita['alpha'], beta0=emita['beta'], silent=False)
sys.stdout = stdout
if '1' in kind:
emita3 = emita.copy()
emita5 = emita.copy()
emita3['emittance'] = emita['emittance']*9
emita5['emittance'] = emita['emittance']*25
# Make loss plots
datproc.plotter(myloss, xax="TURN", yax="PT", cax='PT', kind='hexbin',
ylim=[-0.0025, 0.0020], clim=[-0.0025, 0.0020],
mainkwargs={'bins': 'log'},
save=plotfolder+"/sweep_"+name+'.png')
if pycoll:
for plotkind in ['scatter', 'hist2d']:
datproc.plotter(myloss, xax='X', yax='PX', cax='PT', kind=plotkind,
xlim=[ap['tpstcirc'], 0.095], ylim=[0.0003,0.0027],
clim=[-0.0025, 0.0020], xbin=0.0005, ybin=0.000025, log=True,
save=plotfolder+'/tpst_loss_'+plotkind+'_h_'+name+'.png')
datproc.plotter(myloss, xax='Y', yax='PY', cax='PT', kind=plotkind,
xlim=[-0.02, 0.02], ylim=[-0.0006,0.0006],
clim=[-0.0025, 0.0020], xbin=0.0003, ybin=0.000015, log=True,
save=plotfolder+'/tpst_loss_'+plotkind+'_v_'+name+'.png')
datproc.plotter(myloss, xax='X', yax='Y', cax='PT', kind=plotkind,
xlim=[ap['tpstcirc'], 0.095], ylim=[-0.02, 0.02],
clim=[-0.0025, 0.0020], xbin=0.0005, ybin=0.0003, log=True,
save=plotfolder+'/tpst_loss_'+plotkind+'_s_'+name+'.png')
myrange = [1668,1688]
datproc.plotter(losses[losses['S'].between(*myrange)], xax='S', yax='X',
cax="PT", xlim=myrange, clim=[-0.0025, 0.0020],
ylim=[ap['zsdomid']-ap['zsthick']/2-ap['zsex'],
ap['zsupmid']+ap['zsthick']/2+ap['zsex']],
log=True, save=plotfolder+"/zs_loss_"+name+".png")
else:
lim_pt = [-0.0025, 0.0020]
lim_x = ap['zsupmid']+np.array([-ap['zsthick']/2, ap['zsthick']+ap['zsex']])
lim_px = [-0.00195, -0.00135]
g = datproc.plotter(myloss, xax='X', yax='PX', cax='PT',
xlim=lim_x, ylim=lim_px, clim=lim_pt,
xbin=0.0002, ybin=0.000005, log=False)
plt.savefig(plotfolder+"/losshist_"+name+".png")
datproc.draw_ellipse(emita, g.ax_joint, **{'color':'darkred','marker':'','linestyle':'-'})
datproc.draw_ellipse(emita3, g.ax_joint, **{'color':'darkred','marker':'','linestyle':'--'})
datproc.draw_ellipse(emita5, g.ax_joint, **{'color':'darkred','marker':'','linestyle':':'})
datproc.draw_ellipse(emitb, g.ax_joint, **{'color':'darkorange','marker':'','linestyle':'-'})
plt.savefig(plotfolder+"/losshist_ellipse_"+name+".png")
plt.close()
datproc.plotter(myloss, xax='X', yax='PX', cax='PT', kind='hist2d',
xlim=lim_x, ylim=lim_px, clim=lim_pt,
xbin=0.0002, ybin=0.000005, log=False,
save=plotfolder+"/losshist2_"+name+".png")
datproc.plotter(myloss, xax='X', yax='PX', cax='PT', kind='hist2d',
xlim=[lim_x[0],0.07215], ylim=[-0.00155, -0.0014], clim=lim_pt,
xbin=0.00008, ybin=0.000004, log=False,
save=plotfolder+"/losshist2_zoom_"+name+".png")
if '2' in kind:
# Get track data and make track plots
tracklocs = eval(settings['elements'])
tpt = None if sliced else 9
for trackloc in tracklocs:
tracks = datproc.gettracks(folder+'/tracks', settings=settings, obsloc=trackloc)
for obsnum in [1,3,6,9]:
datproc.plotter(tracks[tracks['obsnum']>-obsnum], cax="PT",
clim=[-0.0025, 0.0020],
save=plotfolder+'/'+trackloc+'_'+str(obsnum)+'obs_'+name+'.png')
if tpt is None:
datproc.plotter(tracks, cax="PT", clim=[-0.0025, 0.0020],
save=plotfolder+'/'+trackloc+'_full_'+name+'.png')