def test_generate_latfile_update1(self):
idx = 80
self.m.reconfigure(idx, {'theta_x': 0.1})
fout2_file = generate_latfile(self.m, latfile=self.latfile2)
with open(fout2_file) as f:
self.assertEqual(f.read().strip(), self.fout2_str)
with open(fout2_file, 'rb') as f:
m = Machine(f)
self.assertEqual(m.conf(idx)['theta_x'], 0.1)
def test_generate_latfile_update2(self):
idx = 0
s = self.m.allocState({})
self.m.propagate(s, 0, 1)
s.moment0 = [[0.1], [0.1], [0.1], [0.1], [0.1], [0.1], [1.0]]
fout2_file = generate_latfile(self.m, latfile=self.latfile2, state=s)
with open(fout2_file, 'rb') as f:
m = Machine(f)
assertAEqual(m.conf()['P0'], [0.1, 0.1, 0.1, 0.1, 0.1, 0.1, 1.0])
def test_reconfigure(self):
latfile = self.testfile
with open(latfile, 'rb') as f:
m0 = Machine(f)
s0 = m0.allocState({})
e_cor_idx = 10
e_name = m0.conf(e_cor_idx)['name']
m0.reconfigure(10, {'theta_x': 0.005})
r0 = m0.propagate(s0, 0, len(m0), range(len(m0)))
fm = ModelFlame(latfile)
fm.reconfigure(e_name, {'theta_x': 0.005})
r, s = fm.run(monitor=range(len(m0)))
rs0 = [ts for (ti,ts) in r0]
rs = [ts for (ti,ts) in r]
for (is1, is2) in zip(rs0, rs):
compare_mstates(self, is1, is2)
class PlotLat:
"""Lattice picture class from FLAME lattice file or FLAME Machine object.
Parameters
----------
source : str or callable
File path of the lattic file (str) or FLAME Machine object (callable)
output : str (None), optional
Output file name. If defined, the lattice plot is generated automatically.
auto_scaling : bool (True), optional
Flag for y-axis scaling by strength of the optical elements
starting_offset : float (0.0), optional
Position offset of starting point in the lattice file
Attributes
----------
types : dict
Element type list of the lattice. Each element type contains
on-off 'flag', plotting 'color', and y-axis 'scale'.
"""
def __init__(self, source, output=None, auto_scaling=False , starting_offset=0.0, **kws):
self._source = source
self._auto_scaling = auto_scaling
self._starting_offset = starting_offset
if type(self._source) == str:
with open(self._source, 'rb') as lat:
self.M = Machine(lat)
elif type(self._source) == Machine:
self.M = self._source
else:
raise ValueError('source must be a file path of .lat or flame.Machine object')
self.types = {'rfcavity': {'flag':True, 'name':'rfcavity', 'color':'orange', 'scale':0.0},
'solenoid': {'flag':True, 'name':'solenoid', 'color':'red', 'scale':0.0},
'quadrupole': {'flag':True, 'name':'quad', 'color':'purple', 'scale':0.0},
'sextupole': {'flag':True, 'name':'sext', 'color':'navy', 'scale':0.0},
'sbend': {'flag':True, 'name':'bend', 'color':'green', 'scale':0.0},
'equad': {'flag':True, 'name':'e-quad', 'color':'blue', 'scale':0.0},
'edipole': {'flag':True, 'name':'e-dipole', 'color':'lime', 'scale':0.0},
'bpm': {'flag':True, 'name':'bpm', 'color':'m', 'scale':0.0},
'orbtrim': {'flag':True, 'name':'corr', 'color':'black', 'scale':0.0},
'stripper': {'flag':True, 'name':'stripper', 'color':'y', 'scale':0.0},
'marker': {'flag':True, 'name':'pm', 'color':'c', 'scale':0.0}
}
if self._auto_scaling:
for i in range(len(self.M)):
elem = self.M.conf(i)
if elem['type'] in self.types.keys():
prv_scl = self.types[elem['type']]['scale']
tmp_scl = np.abs(self._get_scl(elem))
self.types[elem['type']]['scale'] = max(prv_scl,tmp_scl)
if isinstance(output, str):
self.generate(legend=True, option=True)
self.output(window=True, fname=output, **kws)
def _get_scl(self, elem):
"""Get arbital strength of the optical element.
"""
scl = 0.0
if elem['type'] == 'rfcavity':
scl = elem['scl_fac']*np.cos(2.0*np.pi*elem['phi']/360.0)
elif elem['type'] == 'solenoid':
scl = elem['B']
elif elem['type'] == 'quadrupole':
scl = elem['B2'] if 'B2' in elem else 1.0
elif elem['type'] == 'sextupole':
scl = elem['B3']
elif elem['type'] == 'sbend':
scl = elem['phi']
elif elem['type'] == 'equad':
scl = elem['V']/elem['radius']**2.0
elif elem['type'] == 'edipole':
scl = elem['phi']
return scl
def generate(self, start=None, end=None, xlim=None, ycen=0.0, yscl=1.0, aspect=5.0,
legend=True, option=True, axes = None):
"""Generate matplotlib Axes class object from lattice file.
Parameters
----------
start : int
Index of the lattice start.
end : int
Index of the lattice end.
xlim : list[2], optinal
Plot range of the lattice.
ycen : float (0.0)
Vertical offset of the plot.
yscl : float (1.0)
Plot scale of the element size
aspect : float (5.0), optional
Aspect ratio of the picture.
legend : bool
Add legend for the elements.
option : bool
Add optional setting for the plot.
Attributes
----------
axes : callable
Axes class object of matplotlib.
total_length : float
Total length of the lattice.
"""
if axes is None:
self._fig = plt.figure()
self.axes = self._fig.add_subplot(111)
else:
self.axes = axes
pos = self._starting_offset
bp = ycen
indexes = range(len(self.M))[start:end]
foundelm = []
for i in indexes:
elem = self.M.conf(i)
try:
dL = elem['L']
except:
dL = 0.0
if elem['type'] in self.types.keys():
info = self.types[elem['type']]
if foundelm.count(elem['type']) == 0:
foundelm.append(elem['type'])
if legend and info['flag']:
self.axes.fill_between([0,0],[0,0],[0,0], color=info['color'], label=info['name'])
if info['flag']:
if dL != 0.0:
bpp = bp
if info['scale'] != 0.0:
ht = yscl*self._get_scl(elem)/info['scale'] + 0.05
else:
ht = yscl*np.sign(self._get_scl(elem))
if elem['type'] == 'rfcavity' or elem['type'] == 'solenoid':
bpp = bp-yscl*0.7
ht = yscl*2.0*0.7
self.axes.add_patch(ptc.Rectangle((pos, bpp), dL, ht,
edgecolor='none',facecolor=info['color']))
else:
self.axes.add_line(lin.Line2D([pos,pos],[-yscl*0.3+bp, yscl*0.3+bp],color=info['color']))
pos += dL
self.total_length = pos
self.axes.add_line(lin.Line2D([0.0, pos], [bp,bp], color='gray', zorder=-5))
if len(foundelm) <= 4 :
ncol = 4
elif len(foundelm) <= 6:
ncol = 3
elif len(foundelm) <= 8 :
ncol = 4
elif len(foundelm) <= 9:
ncol = 3
else :
ncol = 4
if option:
if xlim != None:
self.axes.set_xlim(xlim)
ancscl = xlim[1]-xlim[0]
else :
self.axes.set_xlim((0.0,pos))
ancscl = pos
if legend:
self.axes.legend(ncol=ncol, loc=10, bbox_to_anchor=(0.5, -0.2*ancscl))
self.axes.set_aspect(aspect)
self.axes.set_ylim((-1.0,1.0))
self.axes.set_yticks(())
self.axes.grid()
self.axes.spines['right'].set_visible(False)
self.axes.spines['left'].set_visible(False)
self.axes.spines['top'].set_visible(False)
self.axes.spines['bottom'].set_visible(False)
plt.tight_layout()
def output(self,window=True,fname=None,**kws):
"""Output the lattice picture to window and/or file.
Parameters
----------
window : bool (True)
Output flag to the window.
fname : str, optional
File path of the output picutre.
**kwargs :
The same kwargs as pyplot.savefig() are available.
"""
if type(fname) == str :
plt.savefig(fname, **kws)
if window: plt.show()
class VAF:
"""
Contains:
getelem(int)
getvalue(int,str)
getindex(str)
getindexu(str)
setelem(int,str,float)
prt_lat_prms()
prt_lat()
tcs()
tcs(int,int)
looptcs()
SaveBeam(S)
"""
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S=self.M.allocState({})
self.M.propagate(S,0,1)
self.refIonZ = S.ref_IonZ
self.IonZ_io = S.IonZ# list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = numpy.zeros((len(self.M),9))
# store element position data
self.bpmls=[]
self.corls=[]
self.cavls=[]
self.solls=[]
self.cavpara=[]
self.solpara=[]
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm' : self.bpmls.append(i)
elif elem['type'] == 'orbtrim' : self.corls.append(i)
elif elem['type'] == 'rfcavity' :
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'],elem['phi']])
elif elem['type'] == 'solenoid' :
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info','w') as f:
iteminfo=[len(self.M),self.bpmls,self.corls,self.cavls,self.solls]
cPickle.dump(iteminfo,f)
def getelem(self,num):
"""
Get parameter of lattice element
getelem(index of lattice element)
"""
print self.lat[num]
def getvalue(self,num,name):
"""
Get parameter of lattice element
getelem(index of lattice element, parameter name)
"""
print self.M.conf(num)[name]
def getindex(self,name,searchby='name'):
"""
Get index list of lattice elements by python style regular expression
getindex(name or type,
searchby = 'name')
"""
name = name.replace(':','_').lower()
pat = re.compile(name)
result = []
for (i,elem) in enumerate(self.lat):
if pat.search(elem[searchby]):
result.append(i)
return result
def getindexu(self,name,searchby='name'):
"""
Get index list of lattice elements by unix style regular expression
getindex(name or type,
searchby = 'name')
"""
name = name.replace(':','_').lower()
result = []
for (i,elem) in enumerate(self.lat):
if fnmatch.fnmatch(elem[searchby],name):
result.append(i)
return result
def setelem(self,num,name,val):
"""
Set parameter of lattice element
setelem(position number of lattice element,
name of parameter
value of parameter)
"""
#D = self.M.conf()['elements'][num]
#D = self.lat[num]
#D[name] = float(val)
self.M.reconfigure(num,{name:float(val)})
"""
# Developing
def genRandomNoise(self):
for (i,cv) in enumerate(self.cavls):
D = self.M.conf()['elements'][cv]
for (j,st) in enumerate(['scl_fac','phi']):
defv = self.cavpara[i][j]
D[st] = defv + numpy.random.normal(0.0,abs(defv)*1e-3)
self.M.reconfigure(cv, D)
D.clear()
for (i,cv) in enumerate(self.solls):
D = self.M.conf()['elements'][cv]
for (j,st) in enumerate(['B']):
defv = self.solpara[i][j]
D[st] = defv + numpy.random.normal(0.0,abs(defv)*1e-3)
self.M.reconfigure(cv, D)
D.clear()
def loopRN(self):
while self.itr < 100:
self.genRandomNoise()
self.itr += 1
"""
def prt_lat_prms(self):
"""Print initial beam parameters in lattice file"""
print '\nIon Charge States = ', self.M.conf()['IonChargeStates']
print 'IonEs [MeV] = ', self.M.conf()['IonEs']/1e6
print 'IonEk [MeV] = ', self.M.conf()['IonEk']/1e6
print '\nBaryCenter 0:\n', self.M.conf()['BaryCenter0']
print '\nBaryCenter 1:\n', self.M.conf()['BaryCenter1']
print '\nBeam Envelope 0:\n', self.M.conf()['S0']
print '\nBeam Envelope 1:\n', self.M.conf()['S1']
def prt_lat(self):
"""Print all lattice elements"""
for (i,elem) in enumerate(self.lat):
print(i, elem['name'], elem['type'])
def tcs(self,lpf=0, opf=1):
"""
Main function of one through calculation
Calculation data is stored at LD.
LD[i] contains:
[pos, xcen, ycen, zrad, xrms, yrms, ref_phis, ref_IonEk]
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ_io
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
#S.clng = self.clng
fin = len(self.M)
# store initial beam data
self.LD[0][0] = S.pos
#Mean data
self.LD[0][1] = S.moment0_env[0]
self.LD[0][2] = S.moment0_env[2]
self.LD[0][3] = S.moment0_env[4]
self.LD[0][4] = S.moment0_rms[0]
self.LD[0][5] = S.moment0_rms[2]
self.LD[0][6] = S.moment0_rms[4]
self.LD[0][7] = S.ref_phis
self.LD[0][8] = S.ref_IonEk
# propagate step by step and store beam data
for i in range(1,len(self.M)):
self.M.propagate(S, i, 1)
self.LD[i][0] = S.pos
#Mean data
self.LD[i][1] = S.moment0_env[0]
self.LD[i][2] = S.moment0_env[2]
self.LD[i][3] = S.moment0_env[4]
self.LD[i][4] = S.moment0_rms[0]
self.LD[i][5] = S.moment0_rms[2]
self.LD[i][6] = S.moment0_rms[4]
self.LD[i][7] = S.ref_phis
self.LD[i][8] = S.ref_IonEk
#output data for plotting
if opf: numpy.savetxt('ldata.txt',self.LD)
if not lpf: return S
def tcs_seg(self,start,end,SD=None,allcs=0,opf=1):
"""
Main function of segmented section calculation
(S, RD) = tcs_seg(start_index, end_index, SD=beam_data)
beam_data can be generated by SaveBeam(S)
RD[i] contains:
[pos, xcen, ycen, zrad, xrms, yrms, ref_phis, ref_IonEk]
i: index of the element
For all charge outputs: allcs=1,
(S, RD, AD) = tcs_seg(start_index, end_index, SD=beam_data, allcs=1)
AD[k][i] contains:
[pos, xcen, ycen, zrad, xrms, yrms]
k: index of the charge state
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
if SD == None :
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ_io
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
else :
# set initial beam data
S.ref_IonZ = SD.refIonZ
S.IonZ = SD.IonZ_io
S.moment0 = SD.BC0
S.moment1 = SD.ENV0
S.ref_phis = SD.refphis
S.phis = SD.phis_io
S.ref_IonEk = SD.refIonEk
S.IonEk = SD.BC0[PS_PS,:]*MeVtoeV + SD.refIonEk
S.pos = SD.pos
phis_ini = S.ref_phis
#S.clng = self.clng
fin = end - start + 1
RD = numpy.zeros((fin,8))
#if allcs: AD = [[[0.0]*6 for i in range(fin)] for j in range(len(S.IonZ))]
if allcs: AD = numpy.zeros((len(S.IonZ),fin,8))
# store initial beam data
RD[0][0] = S.pos
RD[0][1] = S.moment0_env[0]
RD[0][2] = S.moment0_env[2]
RD[0][3] = S.moment0_env[4]
RD[0][4] = S.moment0_rms[0]
RD[0][5] = S.moment0_rms[2]
RD[0][6] = S.ref_phis - phis_ini
RD[0][7] = S.ref_IonEk
if allcs:
for k in range(len(S.IonZ)):
AD[k][0][0] = S.pos
AD[k][0][1] = S.moment0[0,k]
AD[k][0][2] = S.moment0[2,k]
AD[k][0][3] = S.moment0[4,k]
AD[k][0][4] = numpy.sqrt(S.moment1[0,0,k])
AD[k][0][5] = numpy.sqrt(S.moment1[2,2,k])
# propagate step by step and store beam data
for (j,i) in enumerate(range(start,end)):
self.M.propagate(S, i+1, 1)
RD[j+1][0] = S.pos
RD[j+1][1] = S.moment0_env[0]
RD[j+1][2] = S.moment0_env[2]
RD[j+1][3] = S.moment0_env[4]
RD[j+1][4] = S.moment0_rms[0]
RD[j+1][5] = S.moment0_rms[2]
RD[j+1][6] = S.ref_phis - phis_ini
RD[j+1][7] = S.ref_IonEk
if allcs:
for k in range(len(S.IonZ)):
AD[k][j+1][0] = S.pos
AD[k][j+1][1] = S.moment0[0,k]
AD[k][j+1][2] = S.moment0[2,k]
AD[k][j+1][3] = S.moment0[4,k]
AD[k][j+1][4] = numpy.sqrt(S.moment1[0,0,k])
AD[k][j+1][5] = numpy.sqrt(S.moment1[2,2,k])
if opf: numpy.savetxt('ldata.txt',RD)
if allcs: return (S,RD,AD)
else : return (S,RD)
def looptcs(self):
"""
Main loop for Virtual Accelerator
This function should be called as background job.
Loop is stopped by setting itr = 1.
"""
while self.itr < 1:
#self.genRandomNoise() #developing
self.tcs(lpf=1)
#self.itr +=1
class SaveBeam:
def __init__(self,S_in):
self.refIonZ = S_in.ref_IonZ
self.IonZ_io = S_in.IonZ
self.refphis = S_in.ref_phis
self.phis_io = S_in.phis
self.BC0 = S_in.moment0
self.ENV0 = S_in.moment1
self.refIonEk = S_in.ref_IonEk
self.pos = S_in.pos
#!/usr/bin/python
from flame import Machine
from numpy import ndarray
f = open('test.lat', 'r')
m = Machine(f)
fout = open('out.lat', 'w')
mconf = m.conf()
mconf_ks = mconf.keys()
[mconf_ks.remove(i) for i in ['elements', 'name'] if i in mconf_ks]
lines = []
for k in mconf_ks:
v = mconf[k]
if isinstance(v, ndarray):
v = v.tolist()
if isinstance(v, str):
v = '"{0}"'.format(v)
line = '{k} = {v};'.format(k=k, v=v)
lines.append(line)
fout.writelines('\n'.join(lines))
mconfe = mconf['elements']
eu = []
class VAF:
"""
VAF
===
Tool box for FLAME python interface
You can start VAF by
>>> va = vaf.VAF(fname=${lattice_file_path})
Initial beam parameters
(default is the same as lattice file)
-------------------------------------
BC0: 7*n array (n: number of charge state)
Barycenter vectors for each charge state.
contains:[x, px, y, py, z, pz, 1]
unit: [mm, rad, mm, rad, rad, MeV/u, 1]
ENV0: 7*7*n array (n: number of charge state)
Envelope matrixes for each charge state.
IonZ: n array (n: number of charge state)
Charge to mass ratio for each charge state.
RefIonEk: float
Reference kinetic energy at starting point [eV/u]
This tool box contains
(Check each documentation for detail.)
--------------------------------------
prt_beam_prms()
prt_lat_list()
getelem(int)
getvalue(int,str)
getindex(str)
getindexu(str)
setvalue(int,str,float)
tcs()
tcs_seg()
SaveBeam(S)
On developping
--------------
genRN()
"""
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
self.refIonZ = S.ref_IonZ
self.IonZ = S.IonZ # list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = np.zeros((len(self.M), 9))
self.LD2 = np.zeros((len(self.M), 7))
# store element position data
self.bpmls = []
self.corls = []
self.cavls = []
self.solls = []
self.cavpara = []
self.solpara = []
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm': self.bpmls.append(i)
elif elem['type'] == 'orbtrim': self.corls.append(i)
elif elem['type'] == 'rfcavity':
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'], elem['phi']])
elif elem['type'] == 'solenoid':
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info', 'w') as f:
iteminfo = [
len(self.M), self.bpmls, self.corls, self.cavls, self.solls
]
cPickle.dump(iteminfo, f)
def getelem(self, num):
"""
getelem(a)
Get base parameter of lattice element.
Parameters
----------
a : int
Input index of the element.
Return
------
out: dict
Python dict style lattice element.
Examples
--------
>>> print va.getelem(8)
OrderedDict([('B', 5.34), ('L', 0.1), ('aper', 0.02), ('name', 'ls1_ca01_sol1_d1131_1'), ('type', 'solenoid')])
"""
#return self.M.conf()['elements'][num]
return self.lat[num]
def getvalue(self, num, name):
"""
getvalue(a, b)
Get latest parameter of lattice element.
Parameters
----------
a: int
Input index of the element
b: str
Input name of the parameter
Return
------
out: depends on parameter name
Value of the parameter.
Example
-------
>>> print va.getvalue(8,'B')
5.34
"""
return self.M.conf(num)[name]
def getindex(self, name, searchfrom='name'):
"""
getindex(a, searchby='name')
Get index list of lattice elements by python style regular expression.
Parameters
----------
a: str
Keyword for search.
searchfrom: str
Search from 'name' or 'type'.
Return
------
out: list
Search result.
"""
name = name.replace(':', '_').lower()
pat = re.compile(name)
result = []
for (i, elem) in enumerate(self.lat):
if pat.search(elem[searchby]):
result.append(i)
return result
def getindexu(self, name, searchby='name'):
"""
getindex(a, searchby='name')
Get index list of lattice elements by unix style regular expression.
Parameters
----------
a: str
Keyword for search.
searchfrom: str
Search from 'name' or 'type'.
Return
------
out: list
Search result.
"""
name = name.replace(':', '_').lower()
result = []
for (i, elem) in enumerate(self.lat):
if fnmatch.fnmatch(elem[searchby], name):
result.append(i)
return result
def setvalue(self, num, name, val):
"""
setvalue(a, b, c)
Set parameter of lattice element.
Parameters
----------
a: int
Input index of the element.
b: str
Input name of the parameter.
c: float
Set value to the parameter.
"""
self.M.reconfigure(num, {name: float(val)})
# Developing
def genRN(self, scl=0.0005, seed=None):
np.random.seed(seed)
for (i, elem) in enumerate(self.lat):
rnds = np.random.randn(5) * scl #rms 0.5 mm and rms 0.5 mrad
etype = elem['type']
if etype == 'rfcavity' or etype == 'solenoid' or etype == 'quadrupole':
self.M.reconfigure(i, {'dx': float(rnds[0])})
self.M.reconfigure(i, {'dy': float(rnds[1])})
#self.M.reconfigure(i,{'pitch':float(rnds[2])})
#self.M.reconfigure(i,{'yaw':float(rnds[3])})
#self.M.reconfigure(i,{'tilt':float(rnds[4])})
def prt_beam_prms(self):
"""
Print initial beam parameters in lattice file
"""
print '\nIon Charge States = ', self.M.conf()['IonChargeStates']
print 'IonEs [MeV] = ', self.M.conf()['IonEs'] / 1e6
print 'IonEk [MeV] = ', self.M.conf()['IonEk'] / 1e6
print '\nBaryCenter 0:\n', self.M.conf()['BaryCenter0']
print '\nBaryCenter 1:\n', self.M.conf()['BaryCenter1']
print '\nBeam Envelope 0:\n', self.M.conf()['S0']
print '\nBeam Envelope 1:\n', self.M.conf()['S1']
def prt_lat_list(self):
"""
Print all lattice elements with index.
returns (index, name, type)
"""
for (i, elem) in enumerate(self.lat):
print(i, elem['name'], elem['type'])
def tcs(self, lpf=0, opf=1):
"""
tcs(lpf=0, opf=1)
Main function for one through calculation.
Calculation data is stored as VAF.LD.
VAF.LD[i]:
----------
i: int
index of the element
contains:
[pos, xcen, ycen, zcen, xrms, yrms, zrms, ref_phis, ref_IonEk]
units:
[m, mm, mm, mm, mm, mm, mm, rad, eV/u]
Parameters
----------
lpf: bool
flag for loop run.
opf: bool
flag for file type output.
Returns
-------
S: object
Beam parameters at the ending point.
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S, :]
S.IonEk = S.moment0[PS_PS, :] * MeVtoeV + S.ref_IonEk
#S.clng = self.clng
fin = len(self.M)
# store initial beam data
self.LD[0][0] = S.pos
#Mean data
self.LD[0][1] = S.moment0_env[0]
self.LD[0][2] = S.moment0_env[2]
self.LD[0][3] = S.moment0_env[4]
self.LD[0][4] = S.moment0_rms[0]
self.LD[0][5] = S.moment0_rms[2]
self.LD[0][6] = S.moment0_rms[4]
self.LD[0][7] = S.ref_phis
self.LD[0][8] = S.ref_IonEk
# store initial beam data
self.LD2[0][0] = S.pos
#Mean data
self.LD2[0][1] = S.moment0_env[1]
self.LD2[0][2] = S.moment0_env[3]
self.LD2[0][3] = S.moment0_env[5]
self.LD2[0][4] = S.moment0_rms[1]
self.LD2[0][5] = S.moment0_rms[3]
self.LD2[0][6] = S.moment0_rms[5]
# propagate step by step and store beam data
for i in range(1, len(self.M)):
self.M.propagate(S, i, 1)
self.LD[i][0] = S.pos
#Mean data
self.LD[i][1] = S.moment0_env[0]
self.LD[i][2] = S.moment0_env[2]
self.LD[i][3] = S.moment0_env[4]
self.LD[i][4] = S.moment0_rms[0]
self.LD[i][5] = S.moment0_rms[2]
self.LD[i][6] = S.moment0_rms[4]
self.LD[i][7] = S.ref_phis
self.LD[i][8] = S.ref_IonEk
self.LD2[i][0] = S.pos
#Mean data
self.LD2[i][1] = S.moment0_env[1]
self.LD2[i][2] = S.moment0_env[3]
self.LD2[i][3] = S.moment0_env[5]
self.LD2[i][4] = S.moment0_rms[1]
self.LD2[i][5] = S.moment0_rms[3]
self.LD2[i][6] = S.moment0_rms[5]
#output data for plotting
if opf: np.savetxt('ldata.txt', self.LD)
if not lpf: return S
def tcs2(self):
"""
tcs(lpf=0, opf=1)
Main function for one through calculation.
Calculation data is stored as VAF.LD.
VAF.LD[i]:
----------
i: int
index of the element
contains:
[pos, xcen, ycen, zcen, xrms, yrms, zrms, ref_phis, ref_IonEk]
units:
[m, mm, mm, mm, mm, mm, mm, rad, eV/u]
Parameters
----------
lpf: bool
flag for loop run.
opf: bool
flag for file type output.
Returns
-------
S: object
Beam parameters at the ending point.
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S, :]
S.IonEk = S.moment0[PS_PS, :] * MeVtoeV + S.ref_IonEk
### Reconstract data
U = collections.OrderedDict()
T = collections.OrderedDict()
U, T = self.save(U, T, S)
#S.clng = self.clng
fin = len(self.M)
H = self.M.propagate(S, 1, fin, observe=range(fin))
### Reconstract data
U = collections.OrderedDict()
T = collections.OrderedDict()
for i in range(fin - 1):
U, T = self.save(U, T, H[i][1])
return U, T
def tcss(self, start=None, end=None, S_in=None, obs=None):
if start == None: start = 1
if end == None: end = len(self.M) - 1
if S_in == None:
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S, :]
S.IonEk = S.moment0[PS_PS, :] * MeVtoeV + S.ref_IonEk
else:
# set initial beam data
S = S_in.clone()
if obs == None:
obs = range(len(self.M))
#S.clng = self.clng
H_ini = (0, S.clone())
fin = end - start + 1
#H = self.M.propagate(S, start, fin, observe=range(len(self.M)))
H = self.M.propagate(S, start, fin, observe=obs)
if not obs[0]: return [H_ini] + H
else: return H
def looptcs(self):
"""
Main loop for Virtual Accelerator
This function should be called as background job.
Loop is stopped by setting itr = 1.
"""
while self.itr < 1:
#self.genRandomNoise() #developing
self.tcs(lpf=1)
#self.itr +=1
def save(self, u, t, S, label='total'):
for (k, j) in enumerate(S.IonZ):
cs = str(int(round(j * 238.0)))
try:
u[cs].pos.append(S.pos)
except:
u[cs] = self.SPI()
u[cs].pos.append(S.pos)
for i, (n, m) in enumerate(zip(cen, rms)):
getattr(u[cs], n).append(S.moment0[i, k])
getattr(u[cs], m).append(np.sqrt(S.moment1[i, i, k]))
try:
t[label].pos.append(S.pos)
except:
t[label] = self.SPI()
t[label].pos.append(S.pos)
for i, (n, m) in enumerate(zip(cen, rms)):
getattr(t[label], n).append(S.moment0_env[i])
getattr(t[label], m).append(S.moment0_rms[i])
t[label].ek.append(S.ref_IonEk)
t[label].phis.append(S.ref_phis)
return u, t
class SaveBeam:
"""
SaveBeam(S_in)
Store beam parameters for VAF.tcs_seg().
Parameters
----------
S_in: object
Beam parameters made by VAF.tcs() or VAF.tcs_seg().
Returns
-------
out: object
Beam parameters object for VAF.tcs_seg().
"""
def __init__(self, S_in):
self.refIonZ = S_in.ref_IonZ
self.IonZ_io = S_in.IonZ
self.refphis = S_in.ref_phis
self.phis_io = S_in.phis
self.BC0 = S_in.moment0
self.ENV0 = S_in.moment1
self.refIonEk = S_in.ref_IonEk
self.pos = S_in.pos
class SPI:
def __init__(self):
self.pos = []
self.x = []
self.xp = []
self.y = []
self.yp = []
self.z = []
self.zp = []
self.xrms = []
self.xprms = []
self.yrms = []
self.yprms = []
self.zrms = []
self.zprms = []
self.ek = []
self.phis = []
def n(self, name):
self.x_n = np.array(self.x)
#---------------------------------------------------------
def rms2(self, xy, S, cs=0):
if xy == 'x': i = 0
elif xy == 'px': i = 1
elif xy == 'y': i = 2
elif xy == 'py': i = 3
elif xy == 'z': i = 4
elif xy == 'pz': i = 5
return S.moment1[i, i, cs]
def rms(self, xy, S, cs=0):
return np.sqrt(self.rms2(xy, S, cs))
def emit(self, xy, S, cs=0):
if xy == 'x': i = 0
elif xy == 'y': i = 2
else: raise NameError('First argument must be \'x\' or \'y\'')
xx = S.moment1[i, i, cs] * 1e-6
xxp = S.moment1[i, i + 1, cs] * 1e-3
xpxp = S.moment1[i + 1, i + 1, cs]
return np.sqrt(np.abs(xx * xpxp - xxp * xxp))
def nemit(self, xy, S, cs=0):
bg = S.beta[cs] * S.gamma[cs]
return bg * self.emit(xy, S, cs)
def bfunc(self, xy, S, cs=0):
return self.rms2(xy, S, cs) * 1e-6 / self.emit(xy, S, cs)
def bfuncp(self, xy, S, cs=0):
if xy == 'x': i = 0
elif xy == 'y': i = 2
return S.moment1[i, i + 1, cs] * 1e-3 / self.emit(xy, S, cs)
def hbfunc(self, xy, f, cs=0):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = self.bfunc(xy, f[i][1], cs)
return ans
def hbfuncp(self, xy, f, cs=0):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = self.bfuncp(xy, f[i][1], cs)
return ans
def hpos(self, f):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = f[i][1].pos
return ans
def hcen(self, xy, f, cs=0):
if xy == 'x': i = 0
elif xy == 'px': i = 1
elif xy == 'y': i = 2
elif xy == 'py': i = 3
elif xy == 'z': i = 4
elif xy == 'pz': i = 5
ans = np.zeros(len(f))
if cs == -1:
for k in range(len(f)):
ans[k] = f[k][1].moment0_env[i]
else:
for k in range(len(f)):
ans[k] = f[k][1].moment0[i, cs]
return ans
def hrms(self, xy, f, cs=0):
if xy == 'x': i = 0
elif xy == 'px': i = 1
elif xy == 'y': i = 2
elif xy == 'py': i = 3
elif xy == 'z': i = 4
elif xy == 'pz': i = 5
ans = np.zeros(len(f))
if cs == -1:
for k in range(len(f)):
ans[k] = f[k][1].moment0_rms[i]
else:
for k in range(len(f)):
ans[k] = self.rms(xy, f[k][1], cs)
return ans
def dpp(self, S, cs=0):
gm = S.gamma[cs]
iw = S.IonEk[cs]
dw = S.moment0[5][cs] * 1e6
return gm / (1 + gm) * dw / iw
#---------------------------------------------------------
def hdpp(self, f, cs=0):
ans = np.zeros(len(f))
for k in range(len(f)):
ans[k] = self.dpp(f[k][1], cs=cs)
return ans
def hdsp(self, xy, f, cs=0):
if xy == 'x': i = 0
elif xy == 'y': i = 2
ans = np.zeros(len(f))
for k in range(len(f)):
gm = f[k][1].gamma[cs]
iw = f[k][1].IonEk[cs]
xw = f[k][1].moment1[i, 5, cs] * 1e-3
ww = f[k][1].moment1[5, 5, cs] * 1e6 * gm / (1 + gm) / iw
ans[k] = xw / ww
return ans
def getek(self, f, ls):
ans = np.zeros(len(ls))
for (i, j) in enumerate(ls):
ans[i] = f[j][1].ref_IonEk
return ans
def getid0(self, id, f, ls):
ans = np.zeros(len(ls))
for (i, j) in enumerate(ls):
ans[i] = f[j][1].moment0_env[id]
return ans
def getid00(self, cs, id, f, ls):
ans = np.zeros(len(ls))
for (i, j) in enumerate(ls):
ans[i] = f[j][1].moment0[id, cs]
return ans
def getid1(self, id, f, ls):
ans = np.zeros(len(ls))
for (i, j) in enumerate(ls):
ans[i] = f[j][1].moment0_rms[id]
return ans
def getid01(self, cs, id, f, ls):
ans = np.zeros(len(ls))
for (i, j) in enumerate(ls):
ans[i] = np.sqrt(f[j][1].moment1[id, id, cs])
return ans
class FlameLatticeFile(LatticeFile):
element_type = [
'SOURCE', 'DRIFT', 'SBEND', 'QUADRUPOLE', 'SOLENOID', 'RFCAVITY',
'STRIPPER', 'EDIPOLE', 'EQUAD', 'BPM', 'MARKER', 'GENERIC'
]
def __init__(self, read_from=''):
LatticeFile.__init__(self)
self.lattice_format = 'Flame'
self.machine = None
self.linename = None
if read_from != '':
self.parseFrom(read_from)
def isDrift(self, ele_name):
parent_type = self.getParentElements(ele_name)[-1]
if 'DRIFT' in parent_type:
temp = self.getElementProperties(ele_name)
return temp
else:
return False
def isDipole(self, ele_name):
parent_type = self.getParentElements(ele_name)[-1]
if 'BEND' in parent_type or 'DIPOLE' in parent_type:
temp = self.getElementProperties(ele_name)
temp.update({'TYPE': parent_type})
if 'K1' not in temp:
temp['K1'] = 0
return temp
else:
return False
def isQuadrupole(self, ele_name):
parent_type = self.getParentElements(ele_name)[-1]
if 'QUAD' in parent_type:
temp = self.getElementProperties(ele_name)
if 'B2' in temp:
temp['K1'] = temp['B2']
elif 'V' in temp:
temp['K1'] = temp['V']
else:
temp['K1'] = 0
return temp
else:
return False
def isSolenoid(self, ele_name):
parent_type = self.getParentElements(ele_name)[-1]
if 'SOLENOID' in parent_type:
temp = self.getElementProperties(ele_name)
return temp
else:
return False
def isCavity(self, ele_name):
parent_type = self.getParentElements(ele_name)[-1]
if 'RFCAVITY' in parent_type:
temp = self.getElementProperties(ele_name)
return temp
else:
return False
def parseFrom(self, lattice_file_name):
with open(lattice_file_name, 'rb') as f:
self.machine = Machine(f)
lattice_dict = self.machine.conf()
self.lattice = copy.deepcopy(lattice_dict['elements'])
self.linename = lattice_dict['name']
'''['AMU',
'Eng_Data_Dir',
'HdipoleFitMode',
'IonChargeStates',
'IonEk',
'IonEs',
'IonW',
'IonZ',
'MpoleLevel',
'NCharge',
'Stripper_IonChargeStates',
'Stripper_NCharge',
'sim_type']
'''
self.AMU = lattice_dict['AMU']
self.Eng_Data_Dir = lattice_dict['Eng_Data_Dir']
self.HdipoleFitMode = float(lattice_dict.get("HdipoleFitMode", 1.0))
self.IonChargeStates = lattice_dict['IonChargeStates']
self.IonEk = lattice_dict['IonEk']
self.IonEs = lattice_dict['IonEs']
self.IonW = lattice_dict.get('IonW', -1)
self.IonZ = lattice_dict.get('IonZ', -1)
self.MpoleLevel = float(lattice_dict.get("MpoleLevel", 2.0))
self.NCharge = lattice_dict['NCharge']
self.Stripper_IonChargeStates = lattice_dict[
'Stripper_IonChargeStates']
self.Stripper_NCharge = lattice_dict['Stripper_NCharge']
self.sim_type = lattice_dict['sim_type']
n_charge_state = len(self.IonChargeStates)
self.BaryCenter = []
self.SMatrix = []
for i in range(n_charge_state):
temp = 'BaryCenter{}'.format(i)
self.BaryCenter.append(lattice_dict[temp])
temp = 'S{}'.format(i)
self.SMatrix.append(lattice_dict[temp])
for ele in self.lattice:
self.addElement(ele['name'], ele['type'], **ele)
self.appendToBeamline(self.linename, ele['name'].upper())
self.setUseLine(self.linename)
def checkType(self, typename, parameterName=None):
#Should be checked by the flame parser
return True
def run(self):
if self.machine is None:
print('No lattice is defined')
return
s = self.machine.allocState({})
self.machine.propagate(s, 0, 1)
startn = 1
stopn = len(self.machine) - 1
self.flame_result = self.machine.propagate(s,
startn,
stopn - startn + 1,
observe=range(
len(self.machine)))
self._collect_data()
def _collect_data(self):
arg_kws = [
'pos', 'ref_beta', 'ref_bg', 'ref_gamma', 'ref_IonEk', 'ref_IonEs',
'ref_IonQ', 'ref_IonW', 'ref_IonZ', 'ref_phis', 'ref_SampleIonK',
'beta', 'bg', 'gamma', 'IonEk', 'IonEs', 'IonQ', 'IonW', 'IonZ',
'phis', 'SampleIonK', 'moment0', 'moment0_rms', 'moment0_env',
'moment1'
]
self.result_ref = np.zeros((len(self.flame_result), 10))
self.result_moments = np.zeros((len(self.flame_result), 14))
for ind in range(len(self.flame_result)):
itm = self.flame_result[ind]
ind_ele = itm[0]
self.result_ref[ind, 0] = ind_ele
self.result_ref[ind, 1] = itm[1].pos
self.result_ref[ind, 2] = itm[1].ref_beta
self.result_ref[ind, 3] = itm[1].ref_bg
self.result_ref[ind, 4] = itm[1].ref_gamma
self.result_ref[ind, 5] = itm[1].ref_IonEk
self.result_ref[ind, 6] = itm[1].ref_phis
self.result_ref[ind, 7] = itm[1].ref_SampleIonK
self.result_ref[ind, 8] = itm[1].ref_IonQ
self.result_ref[ind, 9] = itm[1].ref_IonZ
self.result_moments[ind, 0] = ind_ele
self.result_moments[ind, 1] = itm[1].pos
self.result_moments[ind, 2:8] = itm[1].moment0_env[0:6]
self.result_moments[ind, 8:] = itm[1].moment0_rms[0:6]
def _output(self, f):
f.write()
def write(self, out_file_name):
with open(out_file_name, 'w') as f:
self._output(f)
def plot_references(self, axis, axis_layout, ploting='ENERGY'):
if isinstance(ploting, str):
plist = [ploting.upper(), None]
elif isinstance(ploting, list) and len(ploting) == 2:
plist = [ploting[0].upper(), ploting[1].upper()]
colors = ['b', 'g']
for i in range(2):
if i == 0:
ax = axis
elif plist[i] is None:
break
else:
axis_m = axis.twinx()
ax = axis_m
if plist[i] == 'E' or plist[i] == 'ENERGY':
ax.plot(self.result_ref[:, 1],
self.result_ref[:, 5] / 1e6,
color=colors[i],
label='Kinetic Energy')
ax.set_ylabel("Energy (MeV)")
elif plist[i] == 'PH' or plist[i] == 'PHASE':
ax.plot(self.result_ref[:, 1],
self.result_ref[:, 6],
color=colors[i],
label='Ref Phase')
ax.set_ylabel("Phase (rad)")
elif plist[i] == 'BG' or plist[i] == 'MOMENTUM':
ax.plot(self.result_ref[:, 1],
self.result_ref[:, 3],
color=colors[i],
label='Norm. Momentum')
ax.set_ylabel(r"$\beta\gamma$")
axis.legend(loc='best')
axis_layout.set_xlabel("s [m]")
self.plotBeamline(axis_layout)
def plot_transverse(self,
axis,
axis_layout,
plot_orbit='xy',
plot_size=None):
colors = ['b', 'g', 'r', 'y']
ci = 0
if plot_orbit is not None:
axis.set_ylabel("Centroid (mm)")
if 'X' in plot_orbit.upper():
axis.plot(self.result_moments[:, 1],
self.result_moments[:, 2],
color=colors[ci],
label='X centroid')
ci += 1
if 'Y' in plot_orbit.upper():
axis.plot(self.result_moments[:, 1],
self.result_moments[:, 4],
color=colors[ci],
label='Y centroid')
ci += 1
axis_m = axis.twinx()
else:
axis_m = axis
if plot_size is not None:
axis_m.set_ylabel("rms beamsize (mm)")
if 'X' in plot_size.upper():
axis_m.plot(self.result_moments[:, 1],
self.result_moments[:, 8],
color=colors[ci],
label='X rms size')
ci += 1
if 'Y' in plot_size.upper():
axis_m.plot(self.result_moments[:, 1],
self.result_moments[:, 10],
color=colors[ci],
label='Y rms size')
ci += 1
axis.legend(loc='best')
axis_layout.set_xlabel("s [m]")
self.plotBeamline(axis_layout)
class VAF:
"""
VAF
===
Tool box for FLAME python interface
You can start VAF by
>>> va = vaf.VAF(fname=${lattice_file_path})
Initial beam parameters
(default is the same as lattice file)
-------------------------------------
BC0: 7*n array (n: number of charge state)
Barycenter vectors for each charge state.
contains:[x, px, y, py, z, pz, 1]
unit: [mm, rad, mm, rad, rad, MeV/u, 1]
ENV0: 7*7*n array (n: number of charge state)
Envelope matrixes for each charge state.
IonZ: n array (n: number of charge state)
Charge to mass ratio for each charge state.
RefIonEk: float
Reference kinetic energy at starting point [eV/u]
This tool box contains
(Check each documentation for detail.)
--------------------------------------
prt_beam_prms()
prt_lat_list()
getelem(int)
getvalue(int,str)
getindex(str)
getindexu(str)
setvalue(int,str,float)
tcs()
tcs_seg()
SaveBeam(S)
On developping
--------------
genRN()
"""
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S=self.M.allocState({})
self.M.propagate(S,0,1)
self.refIonZ = S.ref_IonZ
self.IonZ = S.IonZ # list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = np.zeros((len(self.M),9))
self.LD2 = np.zeros((len(self.M),7))
# store element position data
self.bpmls=[]
self.corls=[]
self.cavls=[]
self.solls=[]
self.cavpara=[]
self.solpara=[]
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm' : self.bpmls.append(i)
elif elem['type'] == 'orbtrim' : self.corls.append(i)
elif elem['type'] == 'rfcavity' :
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'],elem['phi']])
elif elem['type'] == 'solenoid' :
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info','w') as f:
iteminfo=[len(self.M),self.bpmls,self.corls,self.cavls,self.solls]
cPickle.dump(iteminfo,f)
def getelem(self,num):
"""
getelem(a)
Get base parameter of lattice element.
Parameters
----------
a : int
Input index of the element.
Return
------
out: dict
Python dict style lattice element.
Examples
--------
>>> print va.getelem(8)
OrderedDict([('B', 5.34), ('L', 0.1), ('aper', 0.02), ('name', 'ls1_ca01_sol1_d1131_1'), ('type', 'solenoid')])
"""
#return self.M.conf()['elements'][num]
return self.lat[num]
def getvalue(self,num,name):
"""
getvalue(a, b)
Get latest parameter of lattice element.
Parameters
----------
a: int
Input index of the element
b: str
Input name of the parameter
Return
------
out: depends on parameter name
Value of the parameter.
Example
-------
>>> print va.getvalue(8,'B')
5.34
"""
return self.M.conf(num)[name]
def getindex(self,name,searchfrom='name'):
"""
getindex(a, searchby='name')
Get index list of lattice elements by python style regular expression.
Parameters
----------
a: str
Keyword for search.
searchfrom: str
Search from 'name' or 'type'.
Return
------
out: list
Search result.
"""
name = name.replace(':','_').lower()
pat = re.compile(name)
result = []
for (i,elem) in enumerate(self.lat):
if pat.search(elem[searchby]):
result.append(i)
return result
def getindexu(self,name,searchby='name'):
"""
getindex(a, searchby='name')
Get index list of lattice elements by unix style regular expression.
Parameters
----------
a: str
Keyword for search.
searchfrom: str
Search from 'name' or 'type'.
Return
------
out: list
Search result.
"""
name = name.replace(':','_').lower()
result = []
for (i,elem) in enumerate(self.lat):
if fnmatch.fnmatch(elem[searchby],name):
result.append(i)
return result
def setvalue(self,num,name,val):
"""
setvalue(a, b, c)
Set parameter of lattice element.
Parameters
----------
a: int
Input index of the element.
b: str
Input name of the parameter.
c: float
Set value to the parameter.
"""
self.M.reconfigure(num,{name:float(val)})
# Developing
def genRN(self,scl=0.0005,seed=None):
np.random.seed(seed)
for (i,elem) in enumerate(self.lat):
rnds = np.random.randn(5)*scl #rms 0.5 mm and rms 0.5 mrad
etype = elem['type']
if etype == 'rfcavity' or etype == 'solenoid' or etype == 'quadrupole':
self.M.reconfigure(i,{'dx':float(rnds[0])})
self.M.reconfigure(i,{'dy':float(rnds[1])})
#self.M.reconfigure(i,{'pitch':float(rnds[2])})
#self.M.reconfigure(i,{'yaw':float(rnds[3])})
#self.M.reconfigure(i,{'tilt':float(rnds[4])})
def prt_beam_prms(self):
"""
Print initial beam parameters in lattice file
"""
print '\nIon Charge States = ', self.M.conf()['IonChargeStates']
print 'IonEs [MeV] = ', self.M.conf()['IonEs']/1e6
print 'IonEk [MeV] = ', self.M.conf()['IonEk']/1e6
print '\nBaryCenter 0:\n', self.M.conf()['BaryCenter0']
print '\nBaryCenter 1:\n', self.M.conf()['BaryCenter1']
print '\nBeam Envelope 0:\n', self.M.conf()['S0']
print '\nBeam Envelope 1:\n', self.M.conf()['S1']
def prt_lat_list(self):
"""
Print all lattice elements with index.
returns (index, name, type)
"""
for (i,elem) in enumerate(self.lat):
print(i, elem['name'], elem['type'])
def tcs(self,lpf=0, opf=1):
"""
tcs(lpf=0, opf=1)
Main function for one through calculation.
Calculation data is stored as VAF.LD.
VAF.LD[i]:
----------
i: int
index of the element
contains:
[pos, xcen, ycen, zcen, xrms, yrms, zrms, ref_phis, ref_IonEk]
units:
[m, mm, mm, mm, mm, mm, mm, rad, eV/u]
Parameters
----------
lpf: bool
flag for loop run.
opf: bool
flag for file type output.
Returns
-------
S: object
Beam parameters at the ending point.
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
#S.clng = self.clng
fin = len(self.M)
# store initial beam data
self.LD[0][0] = S.pos
#Mean data
self.LD[0][1] = S.moment0_env[0]
self.LD[0][2] = S.moment0_env[2]
self.LD[0][3] = S.moment0_env[4]
self.LD[0][4] = S.moment0_rms[0]
self.LD[0][5] = S.moment0_rms[2]
self.LD[0][6] = S.moment0_rms[4]
self.LD[0][7] = S.ref_phis
self.LD[0][8] = S.ref_IonEk
# store initial beam data
self.LD2[0][0] = S.pos
#Mean data
self.LD2[0][1] = S.moment0_env[1]
self.LD2[0][2] = S.moment0_env[3]
self.LD2[0][3] = S.moment0_env[5]
self.LD2[0][4] = S.moment0_rms[1]
self.LD2[0][5] = S.moment0_rms[3]
self.LD2[0][6] = S.moment0_rms[5]
# propagate step by step and store beam data
for i in range(1,len(self.M)):
self.M.propagate(S, i, 1)
self.LD[i][0] = S.pos
#Mean data
self.LD[i][1] = S.moment0_env[0]
self.LD[i][2] = S.moment0_env[2]
self.LD[i][3] = S.moment0_env[4]
self.LD[i][4] = S.moment0_rms[0]
self.LD[i][5] = S.moment0_rms[2]
self.LD[i][6] = S.moment0_rms[4]
self.LD[i][7] = S.ref_phis
self.LD[i][8] = S.ref_IonEk
self.LD2[i][0] = S.pos
#Mean data
self.LD2[i][1] = S.moment0_env[1]
self.LD2[i][2] = S.moment0_env[3]
self.LD2[i][3] = S.moment0_env[5]
self.LD2[i][4] = S.moment0_rms[1]
self.LD2[i][5] = S.moment0_rms[3]
self.LD2[i][6] = S.moment0_rms[5]
#output data for plotting
if opf: np.savetxt('ldata.txt',self.LD)
if not lpf: return S
def tcs2(self):
"""
tcs(lpf=0, opf=1)
Main function for one through calculation.
Calculation data is stored as VAF.LD.
VAF.LD[i]:
----------
i: int
index of the element
contains:
[pos, xcen, ycen, zcen, xrms, yrms, zrms, ref_phis, ref_IonEk]
units:
[m, mm, mm, mm, mm, mm, mm, rad, eV/u]
Parameters
----------
lpf: bool
flag for loop run.
opf: bool
flag for file type output.
Returns
-------
S: object
Beam parameters at the ending point.
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
### Reconstract data
U = collections.OrderedDict()
T = collections.OrderedDict()
U,T = self.save(U, T, S)
#S.clng = self.clng
fin = len(self.M)
H = self.M.propagate(S, 1, fin, observe=range(fin))
### Reconstract data
U = collections.OrderedDict()
T = collections.OrderedDict()
for i in range(fin-1):
U,T = self.save(U, T, H[i][1])
return U,T
def tcss(self,start=None,end=None,S_in=None,obs=None):
if start == None : start = 1
if end == None : end = len(self.M) -1
if S_in == None :
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S,:]
S.IonEk = S.moment0[PS_PS,:]*MeVtoeV + S.ref_IonEk
else :
# set initial beam data
S = S_in.clone()
if obs == None :
obs = range(len(self.M))
#S.clng = self.clng
H_ini = (0, S.clone())
fin = end - start + 1
#H = self.M.propagate(S, start, fin, observe=range(len(self.M)))
H = self.M.propagate(S, start, fin, observe=obs)
if not obs[0]: return [H_ini] + H
else: return H
def looptcs(self):
"""
Main loop for Virtual Accelerator
This function should be called as background job.
Loop is stopped by setting itr = 1.
"""
while self.itr < 1:
#self.genRandomNoise() #developing
self.tcs(lpf=1)
#self.itr +=1
def save(self,u, t, S, label = 'total'):
for (k,j) in enumerate(S.IonZ):
cs = str(int(round(j*238.0)))
try:
u[cs].pos.append(S.pos)
except:
u[cs] = self.SPI()
u[cs].pos.append(S.pos)
for i,(n,m) in enumerate(zip(cen,rms)):
getattr(u[cs], n).append(S.moment0[i,k])
getattr(u[cs], m).append(np.sqrt(S.moment1[i,i,k]))
try:
t[label].pos.append(S.pos)
except:
t[label] = self.SPI()
t[label].pos.append(S.pos)
for i,(n,m) in enumerate(zip(cen,rms)):
getattr(t[label], n).append(S.moment0_env[i])
getattr(t[label], m).append(S.moment0_rms[i])
t[label].ek.append(S.ref_IonEk)
t[label].phis.append(S.ref_phis)
return u,t
class SaveBeam:
"""
SaveBeam(S_in)
Store beam parameters for VAF.tcs_seg().
Parameters
----------
S_in: object
Beam parameters made by VAF.tcs() or VAF.tcs_seg().
Returns
-------
out: object
Beam parameters object for VAF.tcs_seg().
"""
def __init__(self,S_in):
self.refIonZ = S_in.ref_IonZ
self.IonZ_io = S_in.IonZ
self.refphis = S_in.ref_phis
self.phis_io = S_in.phis
self.BC0 = S_in.moment0
self.ENV0 = S_in.moment1
self.refIonEk = S_in.ref_IonEk
self.pos = S_in.pos
class SPI:
def __init__(self):
self.pos = []
self.x = []
self.xp = []
self.y = []
self.yp = []
self.z = []
self.zp = []
self.xrms = []
self.xprms = []
self.yrms = []
self.yprms = []
self.zrms = []
self.zprms = []
self.ek = []
self.phis = []
def n(self, name):
self.x_n = np.array(self.x)
#---------------------------------------------------------
def rms2(self,xy,S,cs=0):
if xy == 'x': i=0
elif xy == 'px': i=1
elif xy == 'y': i=2
elif xy == 'py': i=3
elif xy == 'z': i=4
elif xy == 'pz': i=5
return S.moment1[i,i,cs]
def rms(self,xy,S,cs=0):
return np.sqrt(self.rms2(xy,S,cs))
def emit(self,xy,S,cs=0):
if xy == 'x': i=0
elif xy == 'y': i=2
else : raise NameError('First argument must be \'x\' or \'y\'')
xx = S.moment1[i,i,cs]*1e-6
xxp = S.moment1[i,i+1,cs]*1e-3
xpxp = S.moment1[i+1,i+1,cs]
return np.sqrt(np.abs(xx*xpxp - xxp*xxp))
def nemit(self,xy,S,cs=0):
bg = S.beta[cs]*S.gamma[cs]
return bg*self.emit(xy,S,cs)
def bfunc(self,xy,S,cs=0):
return self.rms2(xy,S,cs)*1e-6/self.emit(xy,S,cs)
def bfuncp(self,xy,S,cs=0):
if xy == 'x': i = 0
elif xy == 'y': i = 2
return S.moment1[i,i+1,cs]*1e-3/self.emit(xy,S,cs)
def hbfunc(self,xy,f,cs=0):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = self.bfunc(xy,f[i][1],cs)
return ans
def hbfuncp(self,xy,f,cs=0):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = self.bfuncp(xy,f[i][1],cs)
return ans
def hpos(self,f):
ans = np.zeros(len(f))
for i in range(len(f)):
ans[i] = f[i][1].pos
return ans
def hcen(self,xy,f,cs=0):
if xy == 'x': i=0
elif xy == 'px': i=1
elif xy == 'y': i=2
elif xy == 'py': i=3
elif xy == 'z': i=4
elif xy == 'pz': i=5
ans = np.zeros(len(f))
if cs == -1:
for k in range(len(f)):
ans[k] = f[k][1].moment0_env[i]
else:
for k in range(len(f)):
ans[k] = f[k][1].moment0[i,cs]
return ans
def hrms(self,xy,f,cs=0):
if xy == 'x': i=0
elif xy == 'px': i=1
elif xy == 'y': i=2
elif xy == 'py': i=3
elif xy == 'z': i=4
elif xy == 'pz': i=5
ans = np.zeros(len(f))
if cs == -1:
for k in range(len(f)):
ans[k] = f[k][1].moment0_rms[i]
else:
for k in range(len(f)):
ans[k] = self.rms(xy,f[k][1],cs)
return ans
def dpp(self,S,cs=0):
gm = S.gamma[cs]
iw = S.IonEk[cs]
dw = S.moment0[5][cs]*1e6
return gm/(1+gm)*dw/iw
#---------------------------------------------------------
def hdpp(self,f,cs=0):
ans = np.zeros(len(f))
for k in range(len(f)):
ans[k] = self.dpp(f[k][1],cs=cs)
return ans
def hdsp(self,xy,f,cs=0):
if xy == 'x' : i = 0
elif xy == 'y' : i = 2
ans = np.zeros(len(f))
for k in range(len(f)):
gm = f[k][1].gamma[cs]
iw = f[k][1].IonEk[cs]
xw = f[k][1].moment1[i,5,cs]*1e-3
ww = f[k][1].moment1[5,5,cs]*1e6*gm/(1+gm)/iw
ans[k] = xw/ww
return ans
def getek(self,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].ref_IonEk
return ans
def getid0(self,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].moment0_env[id]
return ans
def getid00(self,cs,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].moment0[id,cs]
return ans
def getid1(self,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = f[j][1].moment0_rms[id]
return ans
def getid01(self,cs,id,f,ls):
ans = np.zeros(len(ls))
for (i,j) in enumerate(ls):
ans[i] = np.sqrt(f[j][1].moment1[id,id,cs])
return ans
class DakotaOC(DakotaBase):
""" Dakota optimization class with orbit correction driver
:param lat_file: lattice file
:param elem_bpm: list of element indice of BPMs
:param elem_cor: list of element indice of correctors, always folders of 2
:param elem_hcor: list of element indice of horizontal correctors
:param elem_vcor: list of element indice of vertical correctors
:param ref_x0: reference orbit in x, list of BPM readings
:param ref_y0: reference orbit in y, list of BPM readings
:param ref_flag: string flag for objective functions:
1. "x": :math:`\sum \Delta x^2`, :math:`\Delta x = x-x_0`;
2. "y": :math:`\sum \Delta y^2`, :math:`\Delta y = y-y_0`;
3. "xy": :math:`\sum \Delta x^2 + \sum \Delta y^2`.
:param model: simulation model, 'flame' or 'impact'
:param optdriver: analysis driver for optimization, 'flamedriver_oc' by default
:param kws: keywords parameters for additional usage, defined in ``DakotaBase`` class
valid keys:
* *workdir*: root dir for dakota input/output files,
the defualt one should be created in /tmp, or define some dir path
* *dakexec*: full path of dakota executable,
the default one should be *dakota*, or define the full path
* *dakhead*: prefixed name for input/output files of *dakota*,
the default one is *dakota*
* *keep*: if keep the working directory (i.e. defined by *workdir*),
default is False
.. note:: ``elem_bpm`` should be treated as index of elemnt with type name of 'BPM',
however, for the simulation convenience, any element is acceptable, see :func:`set_bpms()`.
"""
def __init__(self,
lat_file=None,
elem_bpm=None,
elem_cor=None,
elem_hcor=None,
elem_vcor=None,
ref_x0=None,
ref_y0=None,
ref_flag=None,
model=None,
optdriver=None,
**kws):
super(self.__class__, self).__init__(**kws)
if lat_file is not None:
self._lat_file = os.path.realpath(os.path.expanduser(lat_file))
else: # use example lattice file
pass
self._elem_bpm = elem_bpm
self._elem_hcor, self._elem_vcor = elem_hcor, elem_vcor
if elem_cor is not None:
self._elem_hcor, self._elem_vcor = elem_cor[0::2], elem_cor[1::2]
elif elem_hcor is not None:
self._elem_hcor = elem_hcor
elif elem_vcor is not None:
self._elem_vcor = elem_vcor
self._ref_x0 = ref_x0
self._ref_y0 = ref_y0
self._ref_flag = "xy" if ref_flag is None else ref_flag
if model is None:
self._model = 'FLAME'
else:
self._model = model.upper()
if optdriver is None:
self._opt_driver = 'flamedriver_oc'
self.set_model()
self.create_machine(self._lat_file)
self.set_bpms(self._elem_bpm)
self.set_cors(self._elem_hcor, self._elem_vcor)
self.set_ref_x0(self._ref_x0)
self.set_ref_y0(self._ref_y0)
@property
def hcors(self):
return self._elem_hcor
@property
def vcors(self):
return self._elem_vcor
@property
def latfile(self):
return self._lat_file
@property
def ref_x0(self):
return self._ref_x0
@property
def ref_y0(self):
return self._ref_y0
@property
def ref_flag(sef):
return self._ref_flag
@ref_flag.setter
def ref_flag(self, s):
self._ref_flag = s
@property
def bpms(self):
return self._elem_bpm
@latfile.setter
def latfile(self, latfile):
self._lat_file = latfile
@property
def optdriver(self):
return self._opt_driver
@optdriver.setter
def optdriver(self, driver):
self._opt_driver = driver
def get_machine(self):
""" get flame machine object for potential usage
:return: flame machine object or None
"""
try:
return self._machine
except:
return None
def create_machine(self, lat_file):
""" create machine instance with model configuration
* setup _machine
* setup _elem_bpm, _elem_cor or (_elem_hcor and _elem_vcor)
"""
if self._model == "FLAME":
self._create_flame_machine(lat_file)
elif self._model == "IMPACT":
self._create_impact_machine(lat_file)
def _create_flame_machine(self, lat_file):
try:
self._machine = Machine(open(lat_file, 'r'))
except IOError as e:
print("Failed to open {fn}: {info}".format(fn=e.filename,
info=e.args[-1]))
sys.exit(1)
except (RuntimeError, KeyError) as e:
print("Cannot parse lattice, " + e.args[-1])
sys.exit(1)
except:
print("Failed to create machine")
sys.exit(1)
def _create_impact_machine(self, lat_file):
pass
def set_ref_x0(self, ref_arr=None):
""" set reference orbit in x, if not set, use 0s
:param ref_arr: array of reference orbit values
size should be the same number as selected BPMs
"""
if ref_arr is None:
self._ref_x0 = [0]*len(self._elem_bpm)
else:
self._ref_x0 = ref_arr
def set_ref_y0(self, ref_arr=None):
""" set reference orbit in y, if not set, use 0s
:param ref_arr: array of reference orbit values
size should be the same number as selected BPMs
"""
if ref_arr is None:
self._ref_y0 = [0]*len(self._elem_bpm)
else:
self._ref_y0 = ref_arr
def set_bpms(self, bpm=None, pseudo_all=False):
""" set BPMs, and trying to set reference orbit ``(x,y)`` if ``x`` and ``y``
is of one unique value.
:param bpm: list of bpm indices, if None, use all BPMs
:param pseudo_all: if set True, will use all elements, ignore ``bpm`` parameter
"""
if bpm is None:
self._elem_bpm = self.get_all_bpms()
else:
self._elem_bpm = bpm
bpm_count = len(bpm)
if pseudo_all:
self._elem_bpm = "all"
bpm_count = len(self._machine)
try:
if test_one_element(self._ref_x0):
self.set_ref_x0([self._ref_x0[0]]*bpm_count)
else:
pass
if test_one_element(self._ref_y0):
self.set_ref_y0([self._ref_y0[0]]*bpm_count)
else:
pass
except:
pass
#print("Warning, not set reference orbit, requires 'set_ref_x0()' and 'set_ref_y0()'")
def set_cors(self, cor=None, hcor=None, vcor=None):
""" set correctors, if cor, hcor and vcor are None, use all correctors
if cor is not None, use cor, ignore hcor and vcor
:param cor: list of corrector indices, hcor, vcor,...
:param hcor: list of horizontal corrector indices
:param vcor: list of vertical corrector indices
"""
if cor is not None:
self._elem_hcor = cor[0::2]
self._elem_vcor = cor[1::2]
else:
if hcor is None and vcor is None:
self._elem_hcor = self.get_all_cors(type='h')
self._elem_vcor = self.get_all_cors(type='v')
else:
if hcor is not None:
self._elem_hcor = hcor
if vcor is not None:
self._elem_vcor = vcor
def set_model(self, **kws):
""" configure model
:param kws: only for impact, available keys:
"execpath": path of impact executable
"""
if self._model == 'flame':
pass # nothing more needs to do if model is 'flame'
else: # self._model is 'impact'
execpath = kws.get('execpath')
if execpath is not None:
self._impexec = os.path.real(os.path.expanduser(execpath))
else: # just use impact as the default name
self._impexec = "impact"
def get_all_bpms(self):
""" get list of all valid bpms indices
:return: a list of bpm indices
:Example:
>>> dakoc = DakotaOC('test/test.lat')
>>> print(dakoc.get_all_bpms())
"""
return self.get_elem_by_type(type='bpm')
def get_all_cors(self, type=None):
""" get list of all valid correctors indices
:param type: define corrector type, 'h': horizontal, 'v': vertical,
if not defined, return all correctors
:return: a list of corrector indices
:Example:
>>> dakoc = DakotaOC('test/test.lat')
>>> print(dakoc.get_all_cors())
"""
all_cors = self.get_elem_by_type(type='orbtrim')
if type is None:
return all_cors
elif type == 'h':
return all_cors[0::2]
elif type == 'v':
return all_cors[1::2]
else:
print("warning: unrecongnized corrector type.")
return all_cors
def get_elem_by_name(self, name):
""" get list of element(s) by name(s)
:param name: tuple or list of name(s)
:return: list of element indices
:Example:
>>> dakoc = DakotaOC('test/test.lat')
>>> names = ('LS1_CA01:BPM_D1144', 'LS1_WA01:BPM_D1155')
>>> idx = dakoc.get_elem_by_name(names)
>>> print(idx)
[18, 31]
"""
if isinstance(name, str):
name = (name, )
retval = [self._machine.find(name=n)[0] for n in name]
return retval
def get_elem_by_type(self, type):
""" get list of element(s) by type
:param type: string name of element type
:return: list of element indices
:Example:
>>> dakoc = DakotaOC('test/test.lat')
>>> type = 'bpm'
>>> idx = dakoc.get_elem_by_type(type)
>>> print(idx)
"""
retval = self._machine.find(type=type)
return retval
def get_all_elem(self):
""" get all elements from ``Machine`` object
:return: list of element indices
"""
return range(len(self._machine))
def gen_dakota_input(self, infile=None, debug=False):
""" generate dakota input file
:param infile: dakota input filename
:param debug: if True, generate a simple test input file
"""
if not debug:
dakinp = dakutils.DakotaInput()
#dakinp.set_template(name='oc')
dakinp.interface = self._oc_interface
dakinp.variables = self._oc_variables
dakinp.model = self._oc_model
dakinp.responses = self._oc_responses
dakinp.method = self._oc_method
dakinp.environment = self._oc_environ
else: # debug is True
dakinp = dakutils.DakotaInput()
if infile is None:
infile = self._dakhead + '.in'
inputfile = os.path.join(self._workdir, infile)
outputfile = inputfile.replace('.in', '.out')
self._dakin = inputfile
self._dakout = outputfile
dakinp.write(inputfile)
def set_variables(self, plist=None, initial=1e-4, lower=-0.01, upper=0.01):
""" setup variables block, that is setup ``oc_variables``
should be ready to invoke after ``set_cors()``
:param plist: list of defined parameters (``DakotaParam`` object),
automatically setup if not defined
:param initial: initial values for all variables, only valid when plist is None
:param lower: lower bound for all variables, only valid when plist is None
:param upper: upper bound for all variables, only valid when plist is None
"""
if plist is None:
if self._elem_hcor is None and self._elem_vcor is None:
print("No corrector is selected, set_cors() first.")
sys.exit(1)
else:
x_len = len(
self._elem_hcor) if self._elem_hcor is not None else 0
y_len = len(
self._elem_vcor) if self._elem_vcor is not None else 0
n = x_len + y_len
oc_variables = []
oc_variables.append('continuous_design = {0}'.format(n))
oc_variables.append(' initial_point' + "{0:>14e}".format(
initial) * n)
oc_variables.append(' lower_bounds ' + "{0:>14e}".format(
lower) * n)
oc_variables.append(' upper_bounds ' + "{0:>14e}".format(
upper) * n)
xlbls = ["'x{0:03d}'".format(i) for i in range(1, x_len + 1)]
ylbls = ["'y{0:03d}'".format(i) for i in range(1, y_len + 1)]
oc_variables.append(' descriptors ' + ''.join(
["{0:>14s}".format(lbl) for lbl in xlbls + ylbls]))
self._oc_variables = oc_variables
else: # plist = [p1, p2, ...]
n = len(plist)
initial_point_string = ' '.join(["{0:>14e}".format(p.initial) for p in plist])
lower_bounds_string = ' '.join(["{0:>14e}".format(p.lower) for p in plist])
upper_bounds_string = ' '.join(["{0:>14e}".format(p.upper) for p in plist])
descriptors_string = ' '.join(["{0:>14s}".format(p.label) for p in plist])
oc_variables = []
oc_variables.append('continuous_design = {0}'.format(n))
oc_variables.append(' initial_point' + initial_point_string)
oc_variables.append(' lower_bounds ' + lower_bounds_string)
oc_variables.append(' upper_bounds ' + upper_bounds_string)
oc_variables.append(' descriptors ' + descriptors_string)
self._oc_variables = oc_variables
def set_interface(self, interface=None, **kws):
""" setup interface block, that is setup ``oc_interface``
should be ready to invoke after ``set_cors`` and ``set_bpms``
:param interface: ``DakotaInterface`` object, automatically setup if not defined
"""
if interface is None:
oc_interface = dakutils.DakotaInterface(mode='fork', latfile=self._lat_file,
driver='flamedriver_oc',
bpms=self._elem_bpm,
hcors=self._elem_hcor,
vcors=self._elem_vcor,
ref_x0=self._ref_x0,
ref_y0=self._ref_y0,
ref_flag=self._ref_flag,
deactivate='active_set_vector')
else:
oc_interface = interface
self._oc_interface = oc_interface.get_config()
def set_model(self, model=None, **kws):
""" setup model block, that is setup ``oc_model``
:param model: ``DakotaModel`` object, automatically setup if not defined
"""
if model is None:
oc_model = dakutils.DakotaModel()
else:
oc_model = model
self._oc_model = oc_model.get_config()
def set_responses(self, responses=None, **kws):
""" setup responses block, that is setup ``oc_responses``
:param responses: ``DakotaResponses`` object, automatically setup if not defined
"""
if responses is None:
oc_responses = dakutils.DakotaResponses(gradient='numerical')
else:
oc_responses = responses
self._oc_responses = oc_responses.get_config()
def set_environ(self, environ=None):
""" setup environment block, that is setup ``oc_environ``
:param environ: ``DakotaEnviron`` object, automatically setup if not defined
"""
if environ is None:
oc_environ = dakutils.DakotaEnviron(tabfile='dakota.dat')
else:
oc_environ = environ
self._oc_environ= oc_environ.get_config()
def set_method(self, method=None):
""" setup method block, that is setup ``oc_method``
:param method: ``DakotaMethod`` object, automatically setup if not defined
"""
if method is None:
oc_method = dakutils.DakotaMethod(method='cg')
else:
oc_method = method
self._oc_method = oc_method.get_config()
def run(self, mpi=False, np=None, echo=True):
""" run optimization
:param mpi: if True, run DAKOTA in parallel mode, False by default
:param np: number of processes to use, only valid when ``mpi`` is True
:param echo: suppress output if False, True by default
"""
if mpi:
max_core_num = multiprocessing.cpu_count()
if np is None or int(np) > max_core_num:
np = max_core_num
run_command = "mpirun -np {np} {dakexec} -i {dakin} -o {dakout}".format(
np=np,
dakexec=self._dakexec,
dakin=self._dakin,
dakout=self._dakout)
else: # mpi is False
run_command = "{dakexec} -i {dakin} -o {dakout}".format(
dakexec=self._dakexec,
dakin=self._dakin,
dakout=self._dakout)
if echo:
subprocess.call(run_command.split())
else:
devnull = open(os.devnull, 'w')
subprocess.call(run_command.split(), stdout=devnull)
def get_opt_results(self, outfile=None, rtype='dict', label='plain'):
""" extract optimized results from dakota output
:param outfile: file name of dakota output file,
'dakota.out' by default
:param rtype: type of returned results, 'dict' or 'list',
'dict' by default
:param label: label types for returned variables, only valid when rtype 'dict',
'plain' by default:
* *'plain'*: variable labels with format of ``x1``, ``x2``, ``y1``, ``y2``, etc.
e.g. ``{'x1': v1, 'y1': v2}``
* *'fancy'*: variable labels with the name defined in lattice file,
e.g. ``'LS1_CA01:DCH_D1131'``, dict returned sample:
``{'LS1_CA01:DCH_D1131': {'id':9, 'config':{'theta_x':v1}}}``
.. note:: The ``fancy`` option will make re-configuring flame machine in a more
convenient way, such as:
>>> opt_cors = get_opt_results(label='fancy')
>>> for k,v in opt_cors.items():
>>> m.reconfigure(v['id'], v['config'])
>>> # here m is an instance of flame.Machine class
>>>
:return: by default return a dict of optimized results with each item
of the format like "x1":0.1 or more fancy format by set label with 'fancy', etc.,
if rtype='list', return a list of values, when the keys are ascend sorted.
:Example:
>>> opt_vars = get_optresults(outfile='flame_oc.out', rtype='dict'):
>>> print(opt_vars)
{'x2': 0.0020782814353, 'x1': -0.0017913264033}
>>> opt_vars = get_optresults(outfile='flame_oc.out', rtype='list'):
>>> print(opt_vars)
[-0.0017913264033, 0.0020782814353]
"""
if outfile is None:
outfile=self._dakout
if rtype == 'list':
return dakutils.get_opt_results(outfile=outfile, rtype=rtype)
else:
rdict = dakutils.get_opt_results(outfile=outfile, rtype=rtype)
if label == 'plain':
return rdict
else: # label = 'fancy'
val_x = [v for (k,v) in sorted(rdict.items()) if k.startswith('x')]
val_y = [v for (k,v) in sorted(rdict.items()) if k.startswith('y')]
vx = [{'id': i, 'config':{'theta_x': v}} for i,v in zip(self._elem_hcor, val_x)]
vy = [{'id': i, 'config':{'theta_y': v}} for i,v in zip(self._elem_vcor, val_y)]
kx = [self._machine.conf(i)['name'] for i in self._elem_hcor]
ky = [self._machine.conf(i)['name'] for i in self._elem_vcor]
return dict(zip(kx+ky, vx+vy))
def plot(self, outfile=None, figsize=(10, 8), dpi=120, **kws):
""" show orbit
:param outfile: output file of dakota
:param figsize: figure size, (h, w)
:param dpi: figure dpi
"""
if outfile is None:
opt_vars = self.get_opt_results()
else:
opt_vars = self.get_opt_results(outfile=outfile)
idx_h, idx_v = self._elem_hcor, self._elem_vcor
zpos, x, y, mtmp = self.get_orbit((idx_h, idx_v), opt_vars)
fig = plt.figure(figsize=figsize, dpi=dpi, **kws)
ax = fig.add_subplot(111)
linex, = ax.plot(zpos, x, 'r-', label='$x$', lw=2)
liney, = ax.plot(zpos, y, 'b-', label='$y$', lw=2)
ax.set_xlabel('$z\,\mathrm{[m]}$', fontsize=20)
ax.set_ylabel('$\mathrm{Orbit\;[mm]}$', fontsize=20)
ax.legend(loc=3)
plt.show()
def get_orbit(self, idx=None, val=None, outfile=None):
""" calculate the orbit with given configurations
:param idx: (idx_hcor, idx_vcor), tuple of list of indices of h/v cors
:param val: values for each correctos, h/v
:param outfile: filename to save the data
:return: tuple of zpos, env_x, env_y, machine
"""
if idx is None:
idx_x, idx_y = self._elem_hcor, self._elem_vcor
else:
idx_x, idx_y = idx
if val is None:
val = self.get_opt_results()
else:
val = val
m = self._machine
val_x = [v for (k, v) in sorted(val.items()) if k.startswith('x')]
val_y = [v for (k, v) in sorted(val.items()) if k.startswith('y')]
[m.reconfigure(eid, {'theta_x': eval})
for (eid, eval) in zip(idx_x, val_x)]
[m.reconfigure(eid, {'theta_y': eval})
for (eid, eval) in zip(idx_y, val_y)]
s = m.allocState({})
r = m.propagate(s, 0, len(m), observe=range(len(m)))
ob_arr = range(len(m)) if self._elem_bpm == 'all' else self._elem_bpm
zpos = np.array([r[i][1].pos for i in ob_arr])
x, y = np.array(
[[r[i][1].moment0_env[j] for i in ob_arr] for j in [0, 2]])
if outfile is not None:
np.savetxt(outfile,
np.vstack((zpos, x, y)).T,
fmt="%22.14e",
comments='# orbit data saved at ' + time.ctime() + '\n',
header="#{0:^22s} {1:^22s} {2:^22s}".format(
"zpos [m]", "x [mm]", "y [mm]"),
delimiter=' ')
return zpos, x, y, m
def simple_run(self, method='cg', mpi=None, np=None, echo=True, **kws):
""" run optimization after ``set_bpms()`` and ``set_cors()``,
by using default configuration and make full use of computing resources.
:param method: optimization method, 'cg', 'ps', 'cg' by default
:param mpi: if True, run DAKOTA in parallel mode, False by default
:param np: number of processes to use, only valid when ``mpi`` is True
:param echo: suppress output if False, True by default
:param kws: keyword parameters
valid keys:
* step: gradient step, 1e-6 by default
* iternum: max iteration number, 20 by default
* evalnum: max function evaulation number, 1000 by default
"""
if method == 'cg':
max_iter_num = kws.get('iternum', 20)
step = kws.get('step', 1e-6)
md = dakutils.DakotaMethod(method='cg',
max_iterations=max_iter_num)
self.set_method(method=md)
re = dakutils.DakotaResponses(gradient='numerical', step=step)
self.set_responses(responses=re)
else: # 'ps'
max_eval_num = kws.get('evalnum', 1000)
md = dakutils.DakotaMethod(method='ps',
max_function_evaluations=max_eval_num)
self.set_method(method=md)
re = dakutils.DakotaResponses()
self.set_responses(responses=re)
self.set_environ()
self.set_model()
self.set_variables()
self.set_interface()
self.gen_dakota_input()
if mpi:
max_core_num = multiprocessing.cpu_count()
if np is None or int(np) > max_core_num:
np = max_core_num
self.run(mpi=mpi, np=np, echo=echo)
else:
self.run(echo=echo)
def get_opt_latfile(self, outfile='out.lat'):
""" get optimized lattice file for potential next usage,
``run()`` or ``simple_run()`` should be evoked first to get the
optimized results.
:param outfile: file name for generated lattice file
:return: lattice file name or None if failed
"""
try:
z, x, y, m = self.get_orbit()
rfile = generate_latfile(m, latfile=outfile)
except:
print("Failed to generate latfile.")
rfile = None
finally:
return rfile
class VAF:
"""
Contains:
getelem(int)
getvalue(int,str)
getindex(str)
getindexu(str)
setelem(int,str,float)
prt_lat_prms()
prt_lat()
tcs()
tcs(int,int)
looptcs()
SaveBeam(S)
"""
def __init__(self, fname=file_name):
self.name = fname
self.itr = 0
with open(self.name, 'rb') as inf:
self.M = Machine(inf)
self.lat = self.M.conf()['elements']
#Flag for limited longitudinal run
self.clng = 0
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
self.refIonZ = S.ref_IonZ
self.IonZ_io = S.IonZ # list of real IonZ
self.refIonEk = S.ref_IonEk
self.BC0 = S.moment0
self.ENV0 = S.moment1
# memory space for all beam data
# self.LD = [[0.0]*8 for i in range(len(self.M))]
self.LD = numpy.zeros((len(self.M), 9))
# store element position data
self.bpmls = []
self.corls = []
self.cavls = []
self.solls = []
self.cavpara = []
self.solpara = []
for i in range(len(self.M)):
#elem = self.M.conf()['elements'][i]
elem = self.lat[i]
if elem['type'] == 'bpm': self.bpmls.append(i)
elif elem['type'] == 'orbtrim': self.corls.append(i)
elif elem['type'] == 'rfcavity':
self.cavls.append(i)
self.cavpara.append([elem['scl_fac'], elem['phi']])
elif elem['type'] == 'solenoid':
self.solls.append(i)
self.solpara.append([elem['B']])
# output data for plotting
with open('plot.info', 'w') as f:
iteminfo = [
len(self.M), self.bpmls, self.corls, self.cavls, self.solls
]
cPickle.dump(iteminfo, f)
def getelem(self, num):
"""
Get parameter of lattice element
getelem(index of lattice element)
"""
print self.lat[num]
def getvalue(self, num, name):
"""
Get parameter of lattice element
getelem(index of lattice element, parameter name)
"""
print self.M.conf(num)[name]
def getindex(self, name, searchby='name'):
"""
Get index list of lattice elements by python style regular expression
getindex(name or type,
searchby = 'name')
"""
name = name.replace(':', '_').lower()
pat = re.compile(name)
result = []
for (i, elem) in enumerate(self.lat):
if pat.search(elem[searchby]):
result.append(i)
return result
def getindexu(self, name, searchby='name'):
"""
Get index list of lattice elements by unix style regular expression
getindex(name or type,
searchby = 'name')
"""
name = name.replace(':', '_').lower()
result = []
for (i, elem) in enumerate(self.lat):
if fnmatch.fnmatch(elem[searchby], name):
result.append(i)
return result
def setelem(self, num, name, val):
"""
Set parameter of lattice element
setelem(position number of lattice element,
name of parameter
value of parameter)
"""
#D = self.M.conf()['elements'][num]
#D = self.lat[num]
#D[name] = float(val)
self.M.reconfigure(num, {name: float(val)})
"""
# Developing
def genRandomNoise(self):
for (i,cv) in enumerate(self.cavls):
D = self.M.conf()['elements'][cv]
for (j,st) in enumerate(['scl_fac','phi']):
defv = self.cavpara[i][j]
D[st] = defv + numpy.random.normal(0.0,abs(defv)*1e-3)
self.M.reconfigure(cv, D)
D.clear()
for (i,cv) in enumerate(self.solls):
D = self.M.conf()['elements'][cv]
for (j,st) in enumerate(['B']):
defv = self.solpara[i][j]
D[st] = defv + numpy.random.normal(0.0,abs(defv)*1e-3)
self.M.reconfigure(cv, D)
D.clear()
def loopRN(self):
while self.itr < 100:
self.genRandomNoise()
self.itr += 1
"""
def prt_lat_prms(self):
"""Print initial beam parameters in lattice file"""
print '\nIon Charge States = ', self.M.conf()['IonChargeStates']
print 'IonEs [MeV] = ', self.M.conf()['IonEs'] / 1e6
print 'IonEk [MeV] = ', self.M.conf()['IonEk'] / 1e6
print '\nBaryCenter 0:\n', self.M.conf()['BaryCenter0']
print '\nBaryCenter 1:\n', self.M.conf()['BaryCenter1']
print '\nBeam Envelope 0:\n', self.M.conf()['S0']
print '\nBeam Envelope 1:\n', self.M.conf()['S1']
def prt_lat(self):
"""Print all lattice elements"""
for (i, elem) in enumerate(self.lat):
print(i, elem['name'], elem['type'])
def tcs(self, lpf=0, opf=1):
"""
Main function of one through calculation
Calculation data is stored at LD.
LD[i] contains:
[pos, xcen, ycen, zrad, xrms, yrms, ref_phis, ref_IonEk]
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ_io
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S, :]
S.IonEk = S.moment0[PS_PS, :] * MeVtoeV + S.ref_IonEk
#S.clng = self.clng
fin = len(self.M)
# store initial beam data
self.LD[0][0] = S.pos
#Mean data
self.LD[0][1] = S.moment0_env[0]
self.LD[0][2] = S.moment0_env[2]
self.LD[0][3] = S.moment0_env[4]
self.LD[0][4] = S.moment0_rms[0]
self.LD[0][5] = S.moment0_rms[2]
self.LD[0][6] = S.moment0_rms[4]
self.LD[0][7] = S.ref_phis
self.LD[0][8] = S.ref_IonEk
# propagate step by step and store beam data
for i in range(1, len(self.M)):
self.M.propagate(S, i, 1)
self.LD[i][0] = S.pos
#Mean data
self.LD[i][1] = S.moment0_env[0]
self.LD[i][2] = S.moment0_env[2]
self.LD[i][3] = S.moment0_env[4]
self.LD[i][4] = S.moment0_rms[0]
self.LD[i][5] = S.moment0_rms[2]
self.LD[i][6] = S.moment0_rms[4]
self.LD[i][7] = S.ref_phis
self.LD[i][8] = S.ref_IonEk
#output data for plotting
if opf: numpy.savetxt('ldata.txt', self.LD)
if not lpf: return S
def tcs_seg(self, start, end, SD=None, allcs=0, opf=1):
"""
Main function of segmented section calculation
(S, RD) = tcs_seg(start_index, end_index, SD=beam_data)
beam_data can be generated by SaveBeam(S)
RD[i] contains:
[pos, xcen, ycen, zrad, xrms, yrms, ref_phis, ref_IonEk]
i: index of the element
For all charge outputs: allcs=1,
(S, RD, AD) = tcs_seg(start_index, end_index, SD=beam_data, allcs=1)
AD[k][i] contains:
[pos, xcen, ycen, zrad, xrms, yrms]
k: index of the charge state
"""
S = self.M.allocState({})
self.M.propagate(S, 0, 1)
if SD == None:
# set initial beam data
S.ref_IonZ = self.refIonZ
S.IonZ = self.IonZ_io
S.moment0 = self.BC0
S.moment1 = self.ENV0
S.ref_IonEk = self.refIonEk
S.phis = S.moment0[PS_S, :]
S.IonEk = S.moment0[PS_PS, :] * MeVtoeV + S.ref_IonEk
else:
# set initial beam data
S.ref_IonZ = SD.refIonZ
S.IonZ = SD.IonZ_io
S.moment0 = SD.BC0
S.moment1 = SD.ENV0
S.ref_phis = SD.refphis
S.phis = SD.phis_io
S.ref_IonEk = SD.refIonEk
S.IonEk = SD.BC0[PS_PS, :] * MeVtoeV + SD.refIonEk
S.pos = SD.pos
phis_ini = S.ref_phis
#S.clng = self.clng
fin = end - start + 1
RD = numpy.zeros((fin, 8))
#if allcs: AD = [[[0.0]*6 for i in range(fin)] for j in range(len(S.IonZ))]
if allcs: AD = numpy.zeros((len(S.IonZ), fin, 8))
# store initial beam data
RD[0][0] = S.pos
RD[0][1] = S.moment0_env[0]
RD[0][2] = S.moment0_env[2]
RD[0][3] = S.moment0_env[4]
RD[0][4] = S.moment0_rms[0]
RD[0][5] = S.moment0_rms[2]
RD[0][6] = S.ref_phis - phis_ini
RD[0][7] = S.ref_IonEk
if allcs:
for k in range(len(S.IonZ)):
AD[k][0][0] = S.pos
AD[k][0][1] = S.moment0[0, k]
AD[k][0][2] = S.moment0[2, k]
AD[k][0][3] = S.moment0[4, k]
AD[k][0][4] = numpy.sqrt(S.moment1[0, 0, k])
AD[k][0][5] = numpy.sqrt(S.moment1[2, 2, k])
# propagate step by step and store beam data
for (j, i) in enumerate(range(start, end)):
self.M.propagate(S, i + 1, 1)
RD[j + 1][0] = S.pos
RD[j + 1][1] = S.moment0_env[0]
RD[j + 1][2] = S.moment0_env[2]
RD[j + 1][3] = S.moment0_env[4]
RD[j + 1][4] = S.moment0_rms[0]
RD[j + 1][5] = S.moment0_rms[2]
RD[j + 1][6] = S.ref_phis - phis_ini
RD[j + 1][7] = S.ref_IonEk
if allcs:
for k in range(len(S.IonZ)):
AD[k][j + 1][0] = S.pos
AD[k][j + 1][1] = S.moment0[0, k]
AD[k][j + 1][2] = S.moment0[2, k]
AD[k][j + 1][3] = S.moment0[4, k]
AD[k][j + 1][4] = numpy.sqrt(S.moment1[0, 0, k])
AD[k][j + 1][5] = numpy.sqrt(S.moment1[2, 2, k])
if opf: numpy.savetxt('ldata.txt', RD)
if allcs: return (S, RD, AD)
else: return (S, RD)
def looptcs(self):
"""
Main loop for Virtual Accelerator
This function should be called as background job.
Loop is stopped by setting itr = 1.
"""
while self.itr < 1:
#self.genRandomNoise() #developing
self.tcs(lpf=1)
#self.itr +=1
class SaveBeam:
def __init__(self, S_in):
self.refIonZ = S_in.ref_IonZ
self.IonZ_io = S_in.IonZ
self.refphis = S_in.ref_phis
self.phis_io = S_in.phis
self.BC0 = S_in.moment0
self.ENV0 = S_in.moment1
self.refIonEk = S_in.ref_IonEk
self.pos = S_in.pos
