def test_configure(self): latfile = self.testfile with open(latfile, 'rb') as f: m0 = Machine(f) s0 = m0.allocState({}) e_cor_idx = 10 m0.reconfigure(10, {'theta_x': 0.005}) m0.propagate(s0, 0, 1) r0 = m0.propagate(s0, 1, -1, range(len(m0))) fm = ModelFlame(latfile) e = fm.get_element(index=10)[0] e['properties']['theta_x'] = 0.005 fm.configure(e) 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)
def test_run_5(self): """ test_run_5: using BeamState object """ latfile = self.testfile with open(latfile, 'rb') as f: m0 = Machine(f) ms = BeamState(machine=m0) fm = ModelFlame() fm.bmstate = ms fm.machine = m0 obs = fm.get_index_by_type(type='bpm')['bpm'] r,s = fm.run(monitor=obs) s0 = m0.allocState({}) m0.propagate(s0, 0, 1) r0 = m0.propagate(s0, 1, len(m0), observe=obs) 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) compare_mstates(self, s, s0)
def test_run_6(self): """ test_run_6: optional monitor setting 'all' """ latfile = self.testfile with open(latfile, 'rb') as f: m0 = Machine(f) s0 = m0.allocState({}) fm = ModelFlame(latfile) r0 = m0.propagate(s0, 0, len(m0), observe=range(len(m0))) r,s = fm.run(monitor='all') 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) compare_mstates(self, s, s0)
def test_run_2(self): """ test_run_2: propagate from the first to last, monitor all BPMs """ latfile = self.testfile with open(latfile, 'rb') as f: m0 = Machine(f) s0 = m0.allocState({}) fm = ModelFlame(latfile) obs = fm.get_index_by_type(type='bpm')['bpm'] r0 = m0.propagate(s0, 0, len(m0), observe=obs) r,s = fm.run(monitor=obs) 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) compare_mstates(self, s, s0)
def test_run_7(self): """ test_run_7: optional monitor setting 'type' """ latfile = self.testfile with open(latfile, 'rb') as f: m0 = Machine(f) s0 = m0.allocState({}) fm = ModelFlame(latfile) obs = fm.get_index_by_type(type='bpm')['bpm'] r0 = m0.propagate(s0, 0, len(m0), observe=obs) r,s = fm.run(monitor='bpm') 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) compare_mstates(self, s, s0)
def test_collect_data(self): """ test_collect_data: get pos, x0, IonEk """ latfile = self.testfile with open(latfile, 'rb') as f: m0 = Machine(f) s0 = m0.allocState({}) r0 = m0.propagate(s0, 0, 100, observe=range(100)) data0 = collect_data(r0, pos=True, x0=True, IonEk=True) fm = ModelFlame(latfile) r, s = fm.run(from_element=1, to_element=99, monitor=range(100)) data = fm.collect_data(r, pos=True, x0=True, IonEk=True) self.assertEqual(data0['pos'][1:].tolist(), data['pos'].tolist()) self.assertEqual(data0['x0'][1:].tolist(), data['x0'].tolist()) self.assertEqual(data0['IonEk'][1:].tolist(), data['IonEk'].tolist())
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)
from flame import Machine import matplotlib.pyplot as plt f = '../lattice/test_392.lat' m = Machine(open(f)) s = m.allocState({}) ms = len(m) r = m.propagate(s, 0, ms, range(ms)) z = [r[i][1].pos for i in range(ms)] x = [r[i][1].moment0_env[0] for i in range(ms)] opt_cor = { 'LS1_CB10:DCH_D1849': { 'config': { 'theta_x': 5.0134431486e-05 }, 'id': 619 }, 'LS1_CB02:DCV_D1339': { 'config': { 'theta_y': 0.00036972498014 }, 'id': 198 }, 'FS1_CSS:DCH_D2275': { 'config': { 'theta_x': 6.5212371248e-05 },
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
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)
def _co_execute(self): """Execute the virtual accelerator. This includes the following: 1. Creating a temporary working directory for execution of FLAME. 2. Set up the working directory by symlinking from the data directory. 3. Writing the EPICS DB to the working directory (va.db). 4. Starting the softIoc and channel initializing monitors. 5. Add noise to the settings for all input (CSET) channels. 6. Create or update the FLAME machine configuration. 7. Propagate the FLAME simulation and read the results. 8. Update the READ channels of all devives. 9. Update the REST channels of input devies. 10. Repeat from step #5. """ _LOGGER.debug("VA: Execute virtual accelerator") _LOGGER.info("VA: Running at " + self._ts_now) chanprefix = self._chanprefix if self._pv_suffix != '': suffix_str = "_" + self._pv_suffix else: suffix_str = "" # Add channel for VA rep-rate chanrate = f"{chanprefix}:SVR:RATE{suffix_str}" self._epicsdb.append(("ao", chanrate, OrderedDict([ ("DESC", "Rep-rate of Simulation Engine"), ("VAL", self._rate), ("PREC", 1), ]))) _LOGGER.info("VA: Reprate PV is " + chanrate) # Add channel for sample counting chansample_cnt = f"{chanprefix}:SVR:CNT{suffix_str}" self._epicsdb.append(("ao", chansample_cnt, OrderedDict([ ("DESC", "Sample counter for scan client"), ("VAL", 0) ]))) # Add channel for VA configuration and control channoise = f"{chanprefix}:SVR:NOISE{suffix_str}" self._epicsdb.append(("ao", channoise, OrderedDict([ ("DESC", "Noise level of Virtual Accelerator"), ("VAL", self._noise), ("PREC", 5) ]))) _LOGGER.info("VA: Noise PV is " + channoise) chanstat = f"{chanprefix}:SVR:STATUS{suffix_str}" self._epicsdb.append(("bi", chanstat, OrderedDict([ ("DESC", "Status of Virtual Accelerator"), ("VAL", 1), ("ZNAM", "ERR"), ("ONAM", "OK"), ("PINI", "1") ]))) _LOGGER.info("VA: Status PV is " + chanstat) # MPS status chan_mps_stat = f"{chanprefix}:SVR:MpsStatus" self._epicsdb.append(("mbbi", chan_mps_stat, OrderedDict([ ("DESC", "MPS Status of Virtual Accelerator"), ("VAL", 3), ("ZRST", "Fault"), ("ONST", "Disable"), ("TWST", "Monitor"), ("THST", "Enable"), ("PINI", "1") ]))) _LOGGER.info("VA: MPS PV is " + chan_mps_stat) # chancharge = f"{chanprefix}:SVR:CHARGE{suffix_str}" self._epicsdb.append(("ai", chancharge, OrderedDict([ ("DESC", "Q/M of Virtual Accelerator"), ("VAL", 0.0), ("PREC", 5) ]))) # initial beam condition chanbsrc = f"{chanprefix}:SVR:BEAM{suffix_str}" self._epicsdb.append(("waveform", chanbsrc, OrderedDict([ ("DESC", "Init beam of Virtual Accelerator"), ("NELM", 4096), ("FTVL", "UCHAR") ]))) _LOGGER.info("VA: Init beam condition PV is " + chanbsrc) # if self.work_dir is not None: os.makedirs(self.work_dir) self._rm_work_dir = False latticepath = os.path.join(self.work_dir, "test.lat") else: self.work_dir = tempfile.mkdtemp(_TEMP_DIRECTORY_SUFFIX) self._rm_work_dir = True latticepath = None _LOGGER.info("VA: Working directory: %s", self._work_dir) # input file paths epicsdbpath = os.path.join(self.work_dir, "va.db") #output file paths epicslogpath = os.path.join(self.work_dir, "softioc.log") if os.path.isabs(self.data_dir): abs_data_dir = self.data_dir self._latfactory.dataDir = self.data_dir else: abs_data_dir = os.path.abspath(self.data_dir) self._latfactory.dataDir = os.path.abspath(self.data_dir) with open(epicsdbpath, "w") as outfile: self._write_epicsdb(outfile) _LOGGER.info("VA: Write EPICS database to %s", epicsdbpath) #_LOGGER.debug("VA: Write EPICS database to %s", epicsdbpath) self._ioc_logfile = open(epicslogpath, "w") self._ioc_process = Popen(["softIoc", "-d", "va.db"], cwd=self.work_dir, stdout=self._ioc_logfile, stderr=subprocess.STDOUT) _LOGGER.debug("VA: Start EPICS soft IOC with log %s", epicslogpath) _LOGGER.debug("VA: Connecting to channels: {}".format(len(self._csetmap.keys()))) self._subscriptions = [] self._subscriptions.append(catools.camonitor(chanrate, self._handle_rate_monitor)) self._subscriptions.append(catools.camonitor(channoise, self._handle_noise_monitor)) self._subscriptions.append(catools.camonitor(chanbsrc, self._handle_bsrc_monitor)) self._subscriptions.extend(catools.camonitor(self._csetmap.keys(), self._handle_cset_monitor)) _LOGGER.debug("VA: Connecting to channels: Done") machine = None while self._continue: # update the RSET channels with new settings batch = catools.CABatch() for cset in self._csetmap.items(): name, field = self._fieldmap[cset[0]] batch[cset[1][0]] = self._settings[name][field] batch.caput() settings = self._copy_settings_with_noise() self._latfactory.settings = settings lattice = self._latfactory.build() start = time.time() if machine is None: _LOGGER.debug("VA: Create FLAME machine from configuration") machine = Machine(lattice.conf()) else: _LOGGER.debug("VA: Reconfigure FLAME machine from configuration") for idx, elem in enumerate(lattice.elements): machine.reconfigure(idx, elem[2]) if self._bsrc is not None: _LOGGER.info("VA: Reconfigure FLAME machine with init beam config") machine.reconfigure(self._bsrc['index'], self._bsrc['properties']) if latticepath is not None: _LOGGER.debug(f"VA: Write FLAME lattice file to {outfile}") generate_latfile(machine, latfile=latticepath) _LOGGER.debug("VA: Allocate FLAME state from configuration") S = machine.allocState({}) output_map = [] for elem in lattice.elements: if 'name' in elem[3]: output_map.append(elem[3]['name']) else: output_map.append(None) batch = catools.CABatch() for i in range(0, len(machine)): machine.propagate(S, i, 1) if output_map[i] in self._elemmap: elem = self._elemmap[output_map[i]] if isinstance(elem, BPMElement): x_centroid = S.moment0_env[0]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.x_phy], x_centroid) batch[self._readfieldmap[elem.name][elem.fields.x_phy]] = x_centroid y_centroid = S.moment0_env[2]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.y_phy], y_centroid) batch[self._readfieldmap[elem.name][elem.fields.y_phy]] = y_centroid # convert rad to deg and adjust for 161MHz sampling frequency phase = _normalize_phase(2.0 * S.ref_phis * (180.0 / math.pi)) _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.phase_phy], phase) batch[self._readfieldmap[elem.name][elem.fields.phase_phy]] = phase energy = S.ref_IonEk/1.0e6 # convert eV to MeV _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.energy_phy], energy) batch[self._readfieldmap[elem.name][elem.fields.energy_phy]] = energy elif isinstance(elem, PMElement): x_centroid = S.moment0_env[0]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.x], x_centroid) batch[self._readfieldmap[elem.name][elem.fields.x]] = x_centroid y_centroid = S.moment0_env[2]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.y], y_centroid) batch[self._readfieldmap[elem.name][elem.fields.y]] = y_centroid x_rms = S.moment0_rms[0]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.xrms], x_rms) batch[self._readfieldmap[elem.name][elem.fields.xrms]] = x_rms y_rms = S.moment0_rms[2]/1.0e3 _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.yrms], y_rms) batch[self._readfieldmap[elem.name][elem.fields.yrms]] = y_rms sign = elem.sign xy_centroid = (sign*x_centroid + y_centroid)/math.sqrt(2.0) # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.xy], xy_centroid) batch[self._readfieldmap[elem.name][elem.fields.xy]] = xy_centroid xy_rms = 1.0e-3*math.sqrt( (S.moment1_env[0, 0] + S.moment1_env[2, 2])*0.5 + sign*S.moment1_env[0, 2] ) _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.xyrms], xy_rms) batch[self._readfieldmap[elem.name][elem.fields.xyrms]] = xy_rms cxy = sign * S.moment1_env[0, 2] * 1e-6 / x_rms / y_rms _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.cxy], cxy) batch[self._readfieldmap[elem.name][elem.fields.cxy]] = cxy elif isinstance(elem, (FCElement, VDElement, TargetElement, DumpElement, WedgeElement)): x_centroid = S.moment0_env[0]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.x], x_centroid) batch[self._readfieldmap[elem.name][elem.fields.x]] = x_centroid y_centroid = S.moment0_env[2]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.y], y_centroid) batch[self._readfieldmap[elem.name][elem.fields.y]] = y_centroid x_rms = S.moment0_rms[0]/1.0e3 # convert mm to m _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.xrms], x_rms) batch[self._readfieldmap[elem.name][elem.fields.xrms]] = x_rms y_rms = S.moment0_rms[2]/1.0e3 _LOGGER.debug("VA: Update read: %s to %s", self._readfieldmap[elem.name][elem.fields.yrms], y_rms) batch[self._readfieldmap[elem.name][elem.fields.yrms]] = y_rms batch.caput() _LOGGER.info("VA: FLAME execution time: %f s", time.time()-start) # Allow the BPM, PM, etc. readbacks to update # before the device setting readbacks PVs. cothread.Yield() batch = catools.CABatch() for name, value in self._csetmap.items(): name, field = self._fieldmap[name] _LOGGER.debug("VA: Update read: %s to %s", value[1], settings[name][field]) batch[value[1]] = settings[name][field] batch.caput() # Sleep for a fraction (10%) of the total execution time # when one simulation costs more than 0.50 seconds. # Otherwise, sleep for the rest of 1 second. # If a scan is being done on this virtual accelerator, # then the scan server has a period of time to update # setpoints before the next run of IMPACT. delt = time.time() - start if delt > 0.50: cothread.Sleep(delt * 0.1) else: cothread.Sleep(1.0 / self._rate - delt)
#!/usr/bin/python from flame import Machine import matplotlib.pyplot as plt f1 = '../lattice/test_392.lat' #f2 = 'optout.lat' f1 = 'opt6_1.lat' f2 = 'opt6_2.lat' m1 = Machine(open(f1)) s1 = m1.allocState({}) r1 = m1.propagate(s1, 0, len(m1), range(len(m1))) z1 = [r1[i][1].pos for i in range(len(m1))] x1 = [r1[i][1].moment0_env[0] for i in range(len(m1))] y1 = [r1[i][1].moment0_env[2] for i in range(len(m1))] m2 = Machine(open(f2)) s2 = m2.allocState({}) r2 = m2.propagate(s2, 0, len(m2), range(len(m2))) z2 = [r2[i][1].pos for i in range(len(m2))] x2 = [r2[i][1].moment0_env[0] for i in range(len(m2))] y2 = [r2[i][1].moment0_env[2] for i in range(len(m2))] fig = plt.figure() ax = fig.add_subplot(111) lx1, = ax.plot(z1, x1, 'r--', lw=2, label='x0') lx2, = ax.plot(z2, x2, 'b-', lw=2, label='x') ly1, = ax.plot(z1, y1, 'm--', lw=2, label='y0') ly2, = ax.plot(z2, y2, 'g-', lw=2, label='y')
def main(): if len(sys.argv) < 2: parser.print_help() sys.exit(1) args = parser.parse_args(sys.argv[1:]) if args.latfile is None: parser.print_help() sys.exit(1) from flame import Machine import traceback import logging logging.getLogger('flame.machine').disabled = True try: with open(args.latfile, 'rb') as f: m = Machine(f) s = m.allocState({}) m_len = len(m) r = m.propagate(s, 0, m_len, range(m_len)) z = [r[i][1].pos for i in range(m_len)] x = [r[i][1].moment0_env[0] for i in range(m_len)] y = [r[i][1].moment0_env[2] for i in range(m_len)] except: traceback.print_exc() sys.exit(1) import matplotlib.pyplot as plt import matplotlib.style as mplstyle mplstyle.use('bmh') plt.rcParams['font.size'] = 18 plt.rcParams['font.family'] = 'serif' plt.rcParams['axes.linewidth'] = 1.0 plt.rcParams['lines.linewidth'] = 2.0 plt.rcParams['figure.figsize'] = (10, 8) plt.rcParams['figure.dpi'] = 90 fig = plt.figure() ax = fig.add_subplot(111) if args.x is not None: line_x, = ax.plot(z, x, 'r-', label='$x$') if args.y is not None: line_y, = ax.plot(z, y, 'b-', label='$y$') if args.x is None and args.y is None: line_x, = ax.plot(z, x, 'r-', label='$x$') line_y, = ax.plot(z, y, 'b-', label='$y$') ax.set_xlabel('$z\,\mathrm{[m]}$') ax.set_ylabel('$\mathrm{Centroid\,\mathrm{[mm]}}$') ax.set_title(args.latfile, family="monospace") ax.legend(loc='best') if args.output is not None: fig.savefig(args.output) print("Generate figure to file: {}.".format(args.output)) else: plt.show() return 0
# # just show env # # Tong Zhang <*****@*****.**> # 2016-10-16 20:23:35 PM EDT # import numpy as np import time from flame import Machine import matplotlib.pyplot as plt lat_fid = open('../lattice/test_392.lat', 'r') m = Machine(lat_fid) s = m.allocState({}) r = m.propagate(s, 0, len(m), observe=range(len(m))) bpms = m.find(type='bpm') x, y = np.array([[r[i][1].moment0_env[j] for i in bpms] for j in [0, 2]]) pos = np.array([r[i][1].pos for i in bpms]) np.savetxt('orbit0.dat', np.vstack((pos, 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=' ') fig = plt.figure() ax1 = fig.add_subplot(211) linex, = ax1.plot(pos, x, 'r-', label='$x$')
import numpy as np import matplotlib.pyplot as plt lat_fid = open('test.lat', 'r') m = Machine(lat_fid) ## all BPMs and Correctors (both horizontal and vertical) bpm_ids, cor_ids = m.find(type='bpm'), m.find(type='orbtrim') corh_ids = cor_ids[0::2] corv_ids = cor_ids[1::2] observe_ids = bpm_ids ## before distortion s = m.allocState({}) r = m.propagate(s, 0, len(m), observe=range(len(m))) x, y = np.array([[r[i][1].moment0_env[0] for i in range(len(m))] for j in [0, 2]]) pos = np.array([r[i][1].pos for i in range(len(m))]) fig1 = plt.figure(figsize=(10, 8), dpi=120) ax1 = fig1.add_subplot(111) linex, = ax1.plot(pos[observe_ids], x[observe_ids], 'r-', alpha=0.6, label='$\mathrm{ref\;orbit}$') linex.set_lw(2) ## apply random kicks N = 1
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
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
#!/usr/bin/env python import sys "[includes]" from flame import Machine "[includes]" "[Parse lattice file]" with open(sys.argv[1], "r") as F: mymachine = Machine(F) "[Parse lattice file]" "[Allocate State]" thestate = mymachine.allocState({}) "[Allocate State]" print "Initial state", thestate "[Run simulation]" mymachine.propagate(thestate) "[Run simulation]" print "Final state", thestate
class TestGenerateLatfile(unittest.TestCase): def setUp(self): testfile = os.path.join(curdir, 'lattice/test_0.lat') out1file = os.path.join(curdir, 'lattice/out1_0.lat') out2file = os.path.join(curdir, 'lattice/out2_0.lat') self.testfile = make_latfile(testfile) self.out1file = make_latfile(out1file) self.out2file = make_latfile(out2file) with open(self.testfile, 'rb') as f: self.m = Machine(f) s0 = self.m.allocState({}) self.r = self.m.propagate(s0, 0, len(self.m), range(len(self.m))) self.s = s0 k = [ 'beta', 'bg', 'gamma', 'IonEk', 'IonEs', 'IonQ', 'IonW', 'IonZ', 'phis', 'SampleIonK' ] self.keys = [i for i in k] self.ref_keys = ['ref_{}'.format(i) for i in k] + ['pos'] self.keys_ps = ['moment0', 'moment0_env', 'moment0_rms', 'moment1'] with open(self.out1file, 'rb') as f: self.fout1_str = f.read().strip().decode() with open(self.out2file, 'rb') as f: self.fout2_str = f.read().strip().decode() self.latfile0 = os.path.join(curdir, 'lattice/test_1.lat') self.latfile1 = os.path.join(curdir, 'lattice/out1_1.lat') self.latfile2 = os.path.join(curdir, 'lattice/out2_1.lat') self.testcfile = os.path.join(curdir, 'lattice/test_c.lat') self.out3file = os.path.join(curdir, 'lattice/out3.lat') with open(self.out3file, 'rb') as f: self.fout3_str = f.read().strip().decode() with open(self.testcfile, 'rb') as f: self.mc = Machine(f) self.foutc_str = f.read().strip().decode() out4file = os.path.join(curdir, 'lattice/out4_0.lat') self.out4file = make_latfile(out4file) with open(self.out4file, 'rb') as f: self.fout4_str = f.read().strip().decode() self.latfile4 = os.path.join(curdir, 'lattice/out4_1.lat') def tearDown(self): for f in [self.latfile0, self.latfile1, self.latfile2, self.latfile4]: if os.path.isfile(f): os.remove(f) def test_generate_latfile_original1(self): sio = StringIO() sioname = generate_latfile(self.m, out=sio) self.assertEqual(sioname, 'string') self.assertEqual(sio.getvalue().strip(), self.fout1_str) def test_generate_latfile_original2(self): # TEST LATFILE fout1_file = generate_latfile(self.m, latfile=self.latfile1) with open(fout1_file, 'rb') as f: f_str = f.read().strip().decode() self.assertEqual(f_str, self.fout1_str) with open(fout1_file, 'rb') as f: m = Machine(f) s = m.allocState({}) r = m.propagate(s, 0, len(m), range(len(m))) # TEST RESULTS for i in range(len(m)): s1, s0 = r[i][1], self.r[i][1] for k in self.ref_keys: self.assertEqual(getattr(s1, k), getattr(s0, k)) for k in self.keys: self.assertEqual( getattr(s1, k).tolist(), getattr(s0, k).tolist()) for k in self.keys_ps: self.assertEqual( getattr(s1, k).tolist(), getattr(s0, k).tolist()) @unittest.skip("output precision depends on the python version") def test_generate_latfile_original3(self): sio = StringIO() sioname = generate_latfile(self.mc, original=self.testcfile, out=sio) self.assertEqual(sioname, 'string') self.assertEqual(sio.getvalue().strip(), self.fout3_str) def test_generate_latfile_original4(self): sio = StringIO() sioname = generate_latfile(self.m, start=30, end=60, out=sio) self.assertEqual(sioname, 'string') self.assertEqual(sio.getvalue().strip(), self.fout4_str) 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])
"[Allocate State]" thestate = mymachine.allocState({}) "[Allocate State]" print "Initial state", thestate "[Run simulation]" bpms = mymachine.find(type='bpm') #sel_cor = mymachine.find(name='LS1_CA01:DCH_D1131')[-1] #mymachine.reconfigure(sel_cor, {'theta_x': 10}) mymachine.reconfigure(9, {'theta_x': 1.00017}) mymachine.reconfigure(10, {'theta_y': 2.00035}) r = mymachine.propagate(thestate, observe=bpms) "[Run simulation]" print "Final state", thestate r01 = r[0][1] print r01.pos print r01.moment0_env[0] print r01.moment0_env[2] import numpy as np import matplotlib.pyplot as plt xx = np.array([(i, r[i][1].moment0_env[0]) for i in range(len(bpms))]) print(np.sum(xx[:,1]**2))