import os import re from onlinemodel.madx import CMadX import io import copy class MatchMaker: def __init__(self,signal = None,phase = 0): self.matchlist = self.getAllScripts(phase) self.referencePoints={} self.variables={} self.scriptdir = None self.signal=signal self.matchresult=[] self.filter={} def getAllScripts(self,phase): phases = ['Current','Planned','Final'] path = 'Scripts/%s' % phases[phase] print(path) res={} dirs = [f.path for f in os.scandir(path) if f.is_dir()] for ele in dirs: tags=ele.split('/')[-1] res[tags]=ele return res def initScripts(self,target): self.scriptdir = self.matchlist[target] if os.path.exists(self.scriptdir+'/initTwiss.madx'): self.parseReferencePoints() self.scriptInjector = os.path.exists(self.scriptdir+'/matchInjector.madx') self.scriptAramis = os.path.exists(self.scriptdir+'/matchAramis.madx') self.scriptAthos = os.path.exists(self.scriptdir + '/matchAthos.madx') self.scriptPorthos = os.path.exists(self.scriptdir + '/matchPorthos.madx') if os.path.exists(self.scriptdir + '/settings.json'): self.settings=self.scriptdir + '/settings.json' else: self.settings=None def parseReferencePoints(self): file = self.scriptdir+'/initTwiss.madx' self.referencePoints.clear() self.variables.clear() self.filter.clear() # default values self.filter['injector']='s[i1].*k[12]|s20cb.*k1|s20sy01.*k1' self.filter['athos']='sat.*mqua.*k1|sat.*msex.*k2|s20sy02.*m[kq]' self.filter['aramis'] = 's[3a][0r].*k[12]|s20sy03.*k1' with open(file,'r') as f: lines = f.readlines() for line in lines: if 'location:' in line: tags = line.split('location:') location = tags[1].strip() if 'label:' in tags[0]: label = tags[0].split('label:')[1].strip() else: label=None fields_all = tags[0].split(',') fields = [fld.split(';')[0].strip() for fld in fields_all if "=" in fld] self.referencePoints[location] = {'Twiss':{twiss.split('=')[0].strip().lower():float(twiss.split('=')[1].strip()) for twiss in fields},'Label':label} if 'Description:' in line: tags = line.split('Description:') description = tags[1].strip() flds=tags[0].split('=') name = flds[0].split(':')[0].strip() val = float(flds[1].split(';')[0].strip()) self.variables[name] = {'Value':val,'Description':description} print('Found Variable',name,'with value',val) if 'filter:' in line: tag = line.split('filter:') location = tag[1].strip() filterexo = tag[0].split('=')[1].split(';')[0].replace("'","").strip() self.filter[location.lower()]=filterexo print('Found Filter', filterexo, 'for location', location) print('##### Reference Twiss values parsed') def match(self, om, variables = None, Injector=True, Athos = True, Aramis = False, Porthos = False): self.matchresult.clear() cwd = os.getcwd() os.chdir(self.scriptdir) if Athos: target = 'SATBD02' else: target = 'SARBD02' f = io.StringIO() madx = CMadX(f) madx.E0 = 2000 om.setBranch(target, 'SINLH01') madx.updateLattice(om, target) # write lattice madx.commonHeader('SwissFEL', '#s/#e', None) # sets header madx.madx.call('initTwiss.madx') if not variables is None: madx.definePresets(variables) if Injector: print('Matching Injector ...') madx.madx.call('matchInjector.madx') self.updateOnlineModel(om,madx.madx, self.filter['injector']) self.parseMatchOutput(f.getvalue().split('\n'),'Injector') if Athos: f.truncate(0) print('Matching Athos ...') madx.madx.call('matchAthos.madx') self.updateOnlineModel(om, madx.madx, self.filter['athos']) self.parseMatchOutput(f.getvalue().split('\n'), 'Athos') if Aramis: f.truncate(0) print('Matching Aramis ...') if target=='SATBD02': target = 'SARUN20' om.setBranch(target, 'SINLH01') madx.updateLattice(om, target) # write lattice madx.commonHeader('SwissFEL', '#s/#e', None) # sets header madx.madx.call('initTwiss.madx') madx.madx.call('matchAramis.madx') self.updateOnlineModel(om, madx.madx, self.filter['aramis']) self.parseMatchOutput(f.getvalue().split('\n'), 'Aramis') os.chdir(cwd) return madx.madx.table.twiss def parseMatchOutput(self, result,prefix,full=False): loc = None locidx = 1 isMatch = False penalty = 1 report=[] for line in result: if full: print(line) if 'MATCH POINT' in line: isMatch = False print("--------------------------------------") print(line) loc = line.split(':')[1].strip() continue if 'MATCH SUMMARY' in line.upper(): isMatch = True continue if 'VARIABLE "TAR"' in line: isMatch = False print('MATCHING RESULTS: %s (error: %e)' % (prefix,penalty)) if loc is None: loc = 'Matchpoint %d' % locidx locidx +=1 self.matchresult.append({'Location':loc,'Error':penalty}) if penalty > 1e-5: for rep in report: print(rep) penalty = 1 loc = None report.clear() continue if isMatch and len(line) > 4: if 'Final Penal' in line: penalty = float(line.split('=')[1].strip()) else: report.append(line) return def updateOnlineModel(self,om,madx,filter): reg = re.compile(filter) for var in list(madx.globals): if reg.match(var): if '.k1' in var: element = var.split('.k1')[0] ele = om.getElement(element) if not ele is None: ele.k1 = madx.globals.get(var,0) print('... Updating', ele.Name, 'to k1=',ele.k1) else: print('### Unidentified element', element) elif '.k2' in var: element = var.split('.k2')[0] ele = om.getElement(element) if not ele is None: ele.k2 = madx.globals.get(var, 0) print('... Updating', ele.Name, 'to k2=', ele.k2) else: print('### Unidentified element', element) elif '.cory' in var: element = var.split('.cory')[0] ele = om.getElement(element) if not ele is None: ele.design_kick = madx.globals.get(var, 0) print('... Updating', ele.Name, 'to design_kick=', ele.design_kick) else: print('### Unidentified element', element)