201 lines
7.8 KiB
Python
201 lines
7.8 KiB
Python
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)
|
|
|
|
|