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