def Centroids(tfsfile, title='', outputfilename=None, machine=True):
"""
Plot the centroid (mean) x and y from the a Tfs file or pymadx.Tfs instance.
tfsfile - can be either a string or a pymadx.Tfs instance.
title - optional title for plot
outputfilename - optional name to save file to (extension determines format)
machine - if True (default) add machine diagram to top of plot
"""
import pymadx.Data as _Data
madx = _Data.CheckItsTfs(tfsfile)
d = _GetOpticalDataFromTfs(madx)
smax = madx.smax
f = _plt.figure(figsize=(9, 5))
axoptics = f.add_subplot(111)
#optics plots
axoptics.plot(d['s'], d['x'], 'b-', label=r'$\mu_{x}$')
axoptics.plot(d['s'], d['y'], 'g-', label=r'$\mu_{y}$')
axoptics.set_xlabel('S (m)')
axoptics.set_ylabel(r'$\mu_{(x,y)}$ (m)')
axoptics.legend(loc=0, fontsize='small') #best position
#add lattice to plot
if machine:
AddMachineLatticeToFigure(f, madx)
_plt.suptitle(title, size='x-large')
if outputfilename != None:
if '.' in outputfilename:
outputfilename = outputfilename.split('.')[0]
_plt.savefig(outputfilename + '.pdf')
_plt.savefig(outputfilename + '.png')
def ApertureN1(aperture, machine=None, outputfilename=None):
"""
Plot the N1 aperture value from MADX.
Requires N1 and S column.
Optional "machine" argument is string to or pymadx.Data.Tfs instance
for twiss description to provide a machine diagram on top.
"""
import pymadx.Data as _Data
aper = _Data.CheckItsTfsAperture(aperture)
f = _plt.figure(figsize=(9, 5))
s = aper.GetColumn('S')
n = aper.GetColumn('N1')
_plt.plot(s, n)
_plt.xlabel('S (m)')
_plt.ylabel(r'N1 ($\sigma$)')
if machine != None:
AddMachineLatticeToFigure(_plt.gcf(), machine)
if outputfilename != None:
if '.' in outputfilename:
outputfilename = outputfilename.split('.')[0]
_plt.savefig(outputfilename + '.pdf')
_plt.savefig(outputfilename + '.png')
def Aperture(aperture,
machine=None,
outputfilename=None,
plot="xy",
plotapertype=True):
"""
Plots the aperture extents vs. S from a pymadx.Data.Aperture instance.
Inputs:
aperture (pymadx.Data.Aperture) - the aperture model to plot from
machine (str or pymadx.Data.Tfs) - TFS file or TFS istance to plot a machine lattice from (default: None)
outputfilename (str) - Name without extension of the output file if desired (default: None)
plot (str) - Indicates whcih aperture to plot - 'x' for X, 'y' for Y and 'xy' for both (default: 'xy')
plotapertype (bool) - If enabled plots the aperture type at every definted aperture point as a color-coded dot (default: False)
"""
import pymadx.Data as _Data
aper = _Data.CheckItsTfsAperture(aperture)
allowed = ["x", "y", "xy", "X", "Y", "XY"]
if plot not in allowed:
raise ValueError("Invalid option plot: " + plot +
". Use 'x', 'y' or 'xy'")
f = _plt.figure(figsize=(9, 5))
s = aper.GetColumn('S')
x, y = aper.GetExtentAll()
if plotapertype:
t = aper.GetColumn('APERTYPE')
c = map(_ApertypeToColor, t)
if "x" in plot.lower():
line1, = _plt.plot(s, x, 'b-', label='X', alpha=0.6)
if plotapertype:
_plt.scatter(s, x, c=c, s=6)
if "y" in plot.lower():
line2, = _plt.plot(s, y, 'g-', label='Y', alpha=0.6)
if plotapertype:
_plt.scatter(s, y, c=c, s=6)
_plt.xlabel('S (m)')
_plt.ylabel('Aperture (m)')
if plotapertype:
_AddColorLegend(c)
_plt.legend(loc='best', numpoints=1, scatterpoints=1, fontsize='small')
if machine != None:
AddMachineLatticeToFigure(_plt.gcf(), machine)
if outputfilename != None:
if '.' in outputfilename:
outputfilename = outputfilename.split('.')[0]
_plt.savefig(outputfilename + '.pdf')
_plt.savefig(outputfilename + '.png')
def PlotBeta(tfsfile,
title='',
outputfilename=None,
machine=True,
dispersion=False,
squareroot=True):
"""
Plot sqrt(beta x,y) as a function of S. By default, a machine diagram is shown at
the top of the plot.
Optionally set dispersion=True to plot x dispersion as second axis.
Optionally turn off machine overlay at top with machine=False
Specify outputfilename (without extension) to save the plot as both pdf and png.
"""
import pymadx.Data as _Data
madx = _Data.CheckItsTfs(tfsfile)
d = _GetOpticalDataFromTfs(madx)
smax = madx.smax
f = _plt.figure(figsize=(11, 5))
axoptics = f.add_subplot(111)
#optics plots
if squareroot:
yx = _np.sqrt(d['betx'])
yy = _np.sqrt(d['bety'])
else:
yx = d['betx']
yy = d['bety']
axoptics.plot(d['s'], yx, 'b-', label='x')
axoptics.plot(d['s'], yy, 'g-', label='y')
if dispersion:
axoptics.plot([], [], 'r--',
label=r'$\mathrm{D}_{x} (S)$') #fake plot for legend
axoptics.set_xlabel('S (m)')
if squareroot:
axoptics.set_ylabel(r'$\sqrt{\beta}$ ($\sqrt{\mathrm{m}}$)')
else:
axoptics.set_ylabel(r'$\beta$ (m)')
axoptics.legend(loc=0, fontsize='small') #best position
#plot dispersion - only in horizontal
if dispersion:
ax2 = axoptics.twinx()
ax2.plot(d['s'], d['dispx'], 'r--')
ax2.set_ylabel('Dispersion (m)')
#add lattice to plot
if machine:
AddMachineLatticeToFigure(f, madx)
_plt.suptitle(title, size='x-large')
_plt.xlim((0 - 0.05 * smax, 1.05 * smax))
if outputfilename != None:
if '.' in outputfilename:
outputfilename = outputfilename.split('.')[0]
_plt.savefig(outputfilename + '.pdf')
_plt.savefig(outputfilename + '.png')
def _parse_tfs_input(tfs_in, name):
"""Return tfs_in as a Tfs instance, which should either be a path
to a TFS file or a Tfs instance, and in either case, generate a
name if None is provided, and return that as well."""
if isinstance(tfs_in, basestring):
if not _path.isfile(tfs_in):
raise IOError("file \"{}\" not found!".format(tfs_in))
name = (_path.splitext(_path.basename(tfs_in))[0]
if name is None else name)
return _Data.Tfs(tfs_in), name
try:
name = tfs_in.filename if name is None else name
return tfs_in, name
except AttributeError:
raise TypeError("Expected Tfs input is neither a "
"file path nor a Tfs instance: {}".format(tfs_in))
def Survey(tfsfile, title='', outputfilename=None):
"""
Plot the x and z coordinates from a tfs file.
"""
import pymadx.Data as _Data
madx = _Data.CheckItsTfs(tfsfile)
x = madx.GetColumn('X')
z = madx.GetColumn('Z')
f = _plt.figure()
ax = f.add_subplot(111)
ax.set_aspect('equal')
ax.plot(x, z, marker='.')
_plt.suptitle(title, size='x-large')
_plt.xlabel('X (m)')
_plt.ylabel('Z (m)')
def AddMachineLatticeToFigure(figure, tfsfile, tightLayout=True):
"""
Add a diagram above the current graph in the figure that represents the
accelerator based on a madx twiss file in tfs format.
Note you can use matplotlib's gcf() 'get current figure' as an argument.
>>> pymadx.Plot.AddMachineLatticeToFigure(gcf(), 'afile.tfs')
A pymadx.Tfs class instance or a string specifying a tfs file can be
supplied as the second argument interchangeably.
"""
import pymadx.Data as _Data
tfs = _Data.CheckItsTfs(tfsfile) #load the machine description
#check required keys
requiredKeys = ['KEYWORD', 'S', 'L', 'K1L']
okToProceed = all([key in tfs.columns for key in requiredKeys])
if not okToProceed:
print("The required columns aren't present in this tfs file")
print("The required columns are: ", requiredKeys)
raise IOError
axs = figure.get_axes() # get the existing graph
axoptics = figure.get_axes()[0]
_AdjustExistingAxes(figure, tightLayout=tightLayout)
axmachine = _PrepareMachineAxes(figure)
_DrawMachineLattice(axmachine, tfs)
#put callbacks for linked scrolling
def MachineXlim(ax):
axmachine.set_autoscale_on(False)
axoptics.set_xlim(axmachine.get_xlim())
def Click(a):
if a.button == 3:
print('Closest element: ', tfs.NameFromNearestS(a.xdata))
MachineXlim(axmachine)
axmachine.callbacks.connect('xlim_changed', MachineXlim)
figure.canvas.mpl_connect('button_press_event', Click)
def AddMachineLatticeToFigure(figure,
tfsfile,
tightLayout=True,
reverse=False,
offset=None):
"""
Add a diagram above the current graph in the figure that represents the
accelerator based on a madx twiss file in tfs format.
Note you can use matplotlib's gcf() 'get current figure' as an argument.
>>> pymadx.Plot.AddMachineLatticeToFigure(gcf(), 'afile.tfs')
A pymadx.Tfs class instance or a string specifying a tfs file can be
supplied as the second argument interchangeably.
If the reverse flag is used, the lattice is plotted in reverse only. The tfs
instance doesn't change.
Offset can optionally be the name of an object in the lattice (exact name match).
If both offset and reverse are used, reverse happens first. The right click searching
works with the reverse and offset similarly.
"""
import pymadx.Data as _Data
tfs = _Data.CheckItsTfs(tfsfile) #load the machine description
#check required keys
requiredKeys = ['KEYWORD', 'S', 'L', 'K1L']
okToProceed = all([key in tfs.columns for key in requiredKeys])
if not okToProceed:
print("The required columns aren't present in this tfs file")
print("The required columns are: ", requiredKeys)
raise IOError
axoptics = figure.get_axes()[0]
_AdjustExistingAxes(figure, tightLayout=tightLayout)
axmachine = _PrepareMachineAxes(figure)
axmachine.margins(x=0.02)
_DrawMachineLattice(axmachine, tfs, reverse, offset)
#put callbacks for linked scrolling
def MachineXlim(ax):
axmachine.set_autoscale_on(False)
axoptics.set_xlim(axmachine.get_xlim())
def Click(a):
if a.button == 3:
try:
x = a.xdata
if reverse:
x = tfs.smax - x
if offset:
ind = tfs.sequence.index(offset)
xoffset = tfs[ind]['S']
x += xoffset
if x > tfs.smax:
x -= tfs.smax
print('Closest element: ', tfs.NameFromNearestS(x))
except:
pass # don't complain if the S is out of bounds
MachineXlim(axmachine)
axmachine.callbacks.connect('xlim_changed', MachineXlim)
figure.canvas.mpl_connect('button_press_event', Click)
def _TfsToPtc(ptctype,inputfile,outputfilename,ptcfile,startname=None,
stopname=None,ignorezerolengthitems=True,samplers='all',
beampiperadius=0.2,beam=True,ptctrackaperture=[]):
"""
Prepare a madx model for PTC from a Tfs file containing the full
twiss output from MADX.
"""
madx = _Data.CheckItsTfs(inputfile)
ptcfilename = ptcfile
nitems = madx.nitems
opencollimatorsetting = beampiperadius
# data structures for checks
angtot = 0.0
lentot = 0.0
lldiff = []
dldiff = {}
itemsomitted = []
a = _Builder.Machine()
#prepare index for iteration:
if startname == None:
startindex = 0
elif type(startname) == int:
startindex = startname
else:
startindex = madx.IndexFromName(startname)
if stopname == None:
stopindex = nitems #this is 1 larger, but ok as range will stop at n-step -> step=1, range function issue
elif type(stopname) == int:
stopindex = stopname
else:
stopindex = madx.IndexFromName(stopname)
if stopindex <= startindex:
print('stopindex <= startindex'
stopindex = startindex + 1
requiredColumns = ['L', 'ANGLE', 'KSI', 'K0L', 'K0SL', 'K1L', 'K2L', 'K3L','K4L','K5L','K6L',
'K1SL', 'K2SL','K3SL','K4SL','K5SL','K6SL', 'TILT', 'KEYWORD',
'ALFX', 'ALFY', 'BETX', 'BETY', 'VKICK', 'HKICK', 'E1', 'E2', 'FINT', 'FINTX', 'HGAP']
missing = [column for column in requiredColumns if column not in madx.columns]
# raise a value error if a column is missing, otherwise the conversion will continue
if len(missing) > 0:
error = "Missing columns from tfs file - insufficient information to convert file\r\n"
error += "Columns missing: "
for column in missing:
error += column + " "
error += "\r\n"
error += "Given columns : "
for index,column in enumerate(madx.columns):
error += column
if index != len(madx.columns)-1:
error += ", "
raise ValueError(error)
# now assume all are columns present
lindex = madx.ColumnIndex('L')
angleindex = madx.ColumnIndex('ANGLE')
ksIindex = madx.ColumnIndex('KSI')
k0lindex = madx.ColumnIndex('K0L')
k0slindex = madx.ColumnIndex('K0SL')
k1lindex = madx.ColumnIndex('K1L')
k2lindex = madx.ColumnIndex('K2L')
k3lindex = madx.ColumnIndex('K3L')
k4lindex = madx.ColumnIndex('K4L')
k5lindex = madx.ColumnIndex('K5L')
k6lindex = madx.ColumnIndex('K6L')
k1slindex = madx.ColumnIndex('K1SL')
k2slindex = madx.ColumnIndex('K2SL')
k3slindex = madx.ColumnIndex('K3SL')
k4slindex = madx.ColumnIndex('K4SL')
k5slindex = madx.ColumnIndex('K5SL')
k6slindex = madx.ColumnIndex('K6SL')
tiltindex = madx.ColumnIndex('TILT')
tindex = madx.ColumnIndex('KEYWORD')
alphaxindex = madx.ColumnIndex('ALFX')
alphayindex = madx.ColumnIndex('ALFY')
betaxindex = madx.ColumnIndex('BETX')
betayindex = madx.ColumnIndex('BETY')
vkickangleindex = madx.ColumnIndex('VKICK')
hkickangleindex = madx.ColumnIndex('HKICK')
e1index = madx.ColumnIndex('E1')
e2index = madx.ColumnIndex('E2')
fintindex = madx.ColumnIndex('FINT')
fintxindex = madx.ColumnIndex('FINTX')
hgapindex = madx.ColumnIndex('HGAP')
h1index = madx.ColumnIndex('H1')
h2index = madx.ColumnIndex('H2')
# iterate through input file and construct machine
for i in range(startindex,stopindex):
name = madx.sequence[i]
#remove special characters like $, % etc 'reduced' name - rname
rname = _re.sub('[^a-zA-Z0-9_]+','',name) #only allow alphanumeric characters and '_'
t = madx.data[name][tindex]
l = madx.data[name][lindex]
ang = madx.data[name][angleindex]
if l <1e-9:
zerolength = True
else:
zerolength = False
if zerolength and ignorezerolengthitems:
itemsomitted.append(name)
continue #this skips the rest of the loop as we're ignoring this item
kws = {} # element-specific keywords
tilt = madx.data[name][tiltindex]
if tilt != 0:
kws['tilt'] = tilt
e1 = madx.data[name][e1index]
e2 = madx.data[name][e2index]
fint = madx.data[name][fintindex]
fintx = madx.data[name][fintxindex]
hgap = madx.data[name][hgapindex]
k1l = madx.data[name][k1lindex]
h1 = madx.data[name][h1index]
h2 = madx.data[name][h2index]
if k1l:
kws['k1'] = k1l / l
if t == 'DRIFT':
a.AddDrift(rname,l,**kws)
elif t == 'QUADRUPOLE':
k1 = madx.data[name][k1lindex] / l
a.AddQuadrupole(rname,l,**kws)
elif t == 'SEXTUPOLE':
k2 = madx.data[name][k2lindex] / l
a.AddSextupole(rname,l,k2=k2,**kws)
elif t == 'OCTUPOLE':
k3 = madx.data[name][k3lindex] / l
a.AddOctupole(rname,l,k3=k3,**kws)
elif t == 'SOLENOID':
ks = madx.data[name][ksIindex] / l
a.AddSolenoid(rname,l,ks=ks,**kws)
elif t == 'SBEND':
kws['e1'] = e1
kws['fint'] = fint
kws['e2'] = e2
# in madx, -1 means fintx was allowed to default to fint and we should do the same
# so if set to 0, this means we want it to be 0
if fintx != -1:
kws['fintx'] = fintx
if h1 != 0:
kws['h1'] = h1
if h2 != 0:
kws['h2'] = h2
if hgap != 0:
kws['hgap'] = hgap
angle = madx.data[name][angleindex]
a.AddDipole(rname,category='sbend',length=l,angle=angle,**kws)
elif t == 'RBEND':
angle = madx.data[name][angleindex]
# set element length to be the chord length - tfs output rbend
# length is arc length
chordLength = l
if angle != 0:
chordLength = 2 * (l / angle) * _np.sin(angle / 2.) # protect against 0 angle rbends
# subtract dipole angle/2 added on to poleface angles internally by
# madx
poleInAngle = e1 - 0.5 * angle
poleOutAngle = e2 - 0.5 * angle
if poleInAngle != 0:
kws['e1'] = poleInAngle
if poleOutAngle != 0:
kws['e2'] = poleOutAngle
if fint != 0:
kws['fint'] = fint
# in madx, -1 means fintx was allowed to default to fint and we should do the same
# so if set to 0, this means we want it to be 0
if fintx != -1:
kws['fintx'] = fintx
if h1 != 0:
kws['h1'] = h1
if h2 != 0:
kws['h2'] = h2
if hgap != 0:
kws['hgap'] = hgap
a.AddDipole(rname,category='rbend',length=chordLength,angle=angle,**kws)
elif t == 'MARKER':
angle = madx.data[name][angleindex]
a.AddMarker(rname, **kws)
elif t == 'MULTIPOLE':
kn0l = madx.data[name][k0lindex]
kn1l = madx.data[name][k1lindex]
kn2l = madx.data[name][k2lindex]
kn3l = madx.data[name][k3lindex]
kn4l = madx.data[name][k4lindex]
kn5l = madx.data[name][k5lindex]
kn6l = madx.data[name][k6lindex]
kn0sl = madx.data[name][k0slindex]
kn1sl = madx.data[name][k1slindex]
kn2sl = madx.data[name][k2slindex]
kn3sl = madx.data[name][k3slindex]
kn4sl = madx.data[name][k4slindex]
kn5sl = madx.data[name][k5slindex]
kn6sl = madx.data[name][k6slindex]
a.AddMultipole(
rname,
knl=(kn0l, kn1l, kn2l, kn3l, kn4l, kn5l, kn6l),
ksl=(kn0sl, kn1sl, kn2sl, kn3sl, kn4sl, kn5sl, kn6sl),
**kws)
elif t == 'HKICKER':
hkick = madx.data[name][hkickangleindex]
a.AddHKicker(rname, hkick=hkick, length=l)
elif t == 'VKICKER':
vkick = madx.data[name][vkickangleindex]
a.AddVKicker(rname, vkick=vkick, length=l)
elif t == 'TKICKER':
vkick = madx.data[name][vkickangleindex]
hkick = madx.data[name][hkickangleindex]
a.AddTKicker(rname, vkick=vkick, hkick=hkick, length=l)
else:
print('MadxTfs2Ptc> unknown element type: ',t,' for element named: ',name
if not zerolength:
print('MadxTfs2Ptc> replacing with drift')
a.AddDrift(rname,l)
else:
pass
if ptctype == 'ptctrack':
a.AddSampler(samplers)
a.AddPTCTrackAperture(ptctrackaperture)
# Make beam file
if beam:
if ptctype == 'ptctrack':
b = MadxTfsToPtcBeam(madx, ptcfilename, startname)
a.AddBeam(b)
if ptctype == 'ptctwiss':
b = MadxTfsToPtcTwissBeam(madx, startname)
a.AddBeam(b)
a.Write(outputfilename)
return a
def MadxTfsToPtcBeam(tfs, ptcfilename, startname=None):
if startname == None:
startindex = 0
elif type(startname) == int:
startindex = startname
else:
startindex = tfs.IndexFromName(startname)
#MADX defines parameters at the end of elements so need to go 1 element
#back if we can.
if startindex > 0:
startindex -= 1
energy = float(tfs.header['ENERGY'])
gamma = float(tfs.header['GAMMA'])
particle = str.lower(tfs.header['PARTICLE'])
ex = tfs.header['EX']
ey = tfs.header['EY']
sigmae = float(tfs.header['SIGE'])
sigmat = float(tfs.header['SIGT'])
data = tfs.GetRowDict(tfs.sequence[startindex])
beam = _Beam.Beam(particle, energy,'ptc', emitx=ex, emity=ey, sigmaE=sigmae)
beam.SetDistribFileName(ptcfilename)
return beam
def MadxTfsToPtcTwissBeam(tfs, startname=None):
if startname == None:
startindex = 0
elif type(startname) == int:
startindex = startname
else:
startindex = tfs.IndexFromName(startname)
# MADX defines parameters at the end of elements so need to go 1 element
# back if we can.
if startindex > 0:
startindex -= 1
energy = float(tfs.header['ENERGY'])
gamma = float(tfs.header['GAMMA'])
particle = str.lower(tfs.header['PARTICLE'])
ex = tfs.header['EX']
ey = tfs.header['EY']
sigmae = float(tfs.header['SIGE'])
sigmat = float(tfs.header['SIGT'])
beam = _Beam.Beam(particle, energy, 'ptctwiss', emitx=ex, emity=ey, sigmaE=sigmae)
if (not 'BETX' in tfs.columns) or (not 'BETY' in tfs.columns):
raise AttributeError('BETX or BETY columns missing from TFS file. Beam cannot be converted.')
beam.SetBetaX(tfs.GetColumn('BETX')[0])
beam.SetBetaY(tfs.GetColumn('BETY')[0])
if (not 'ALFX' in tfs.columns):
print("ALFX column missing from TFS file. Setting ALFX to 0.")
beam.SetAlphaX(0)
else:
beam.SetAlphaX(tfs.GetColumn('ALFX')[0])
if (not 'ALFY' in tfs.columns):
print("ALFY column missing from TFS file. Setting ALFY to 0.")
beam.SetAlphaY(0)
else:
beam.SetAlphaY(tfs.GetColumn('ALFY')[0])
return beam