import json import copy from PyQt5 import QtCore, QtGui, QtWidgets from fontTools.misc.cython import returns CBeige = QtGui.QColor(250, 240, 200) CGreen = QtGui.QColor(100, 255, 100) CYellow = QtGui.QColor(255, 255, 0) CRed = QtGui.QColor(255, 100, 100) class ReferenceManager: def __init__(self, parent): self.parent = parent self.settings={} self.reference={} self.save={} self.label=None self.twisswidget={'betax':self.parent.UIBetax,'betay':self.parent.UIBetay, 'alphax':self.parent.UIAlphax,'alphay':self.parent.UIAlphay, 'etax':self.parent.UIEtax,'etay':self.parent.UIEtay, 'etapx':self.parent.UIEtapx,'etapy':self.parent.UIEtapy, 'x':self.parent.UIX,'y':self.parent.UIY, 'px':self.parent.UIPX,'py':self.parent.UIPY} self.loadReference('Settings/Reference.json') self.parent.UITrackReference.currentIndexChanged.connect(self.updateReferenceWidgets) self.parent.UIMatchOrder.itemClicked.connect(self.displayMatchingPoint) def clearMatching(self): self.save.clear() self.updateMatchOrder() def updateMatchOrder(self): self.parent.UIMatchOrder.clear() for i,ele in enumerate(self.settings['Order']): self.parent.UIMatchOrder.addItem(ele['MatchID']) self.parent.UIMatchOrder.item(i).setBackground(CBeige) if self.parent.UIMatchOrder.count() > 0: self.parent.UIMatchOrder.setCurrentRow(0) self.displayMatchingPoint() def loadReference(self, filename): with open(filename) as f: self.settings = json.load(f) self.order={ele['MatchID']:i for i,ele in enumerate(self.settings['Order'])} self.reference=self.settings['Reference'] self.label=self.settings['Label'] self.parameter=self.settings['Parameter'] self.dependence=self.settings['Dependence'] self.parent.UIMatchSequenceLabel.setText(self.label) self.clearMatching() self.updateReferenceComboBox() self.updateReferenceWidgets() def saveTwiss(self, ID, twiss): self.save[ID]={'betax':twiss.betx[0],'betay':twiss.bety[0],'alphax':twiss.alfx[0],'alphay':twiss.alfy[0]} print('saving twissvalue for',ID,':',self.save[ID]) def getMatchingPoint(self): idx = self.parent.UIMatchOrder.currentRow() if idx < 0: return None,None,False status = True ID = str(self.parent.UIMatchOrder.item(idx).text()) idx = self.order[ID] # get index in the list order = copy.deepcopy(self.settings['Order'][idx]) # getting the specific entry from the order list if not order['Reference'] is None: # has a fixed referenc epoint instead of periodic solution twiss0 = self.getReferenceByTag(order['Reference']) # get the starting parameters by its reference order['Reference'] = twiss0 # put it back into the order variable order['Variable'] = self.parameter[ID]['Variable'] nvar = len(order['Variable'].keys()) target0= self.parameter[ID]['Target'] target = {} ncon = 0 for tar in target0: if 'Reference' in tar.keys(): par=self.getReferenceByTag(tar['Reference']) loc = par['Location'] con = {key:{'Val':par['Twiss'][key],'Condition':0} for key in par['Twiss'].keys()} target[loc]=con ncon += len(con.keys()) elif 'Script' in tar.keys(): loc = tar['Script']['Location'] file = tar['Script']['Script'] target[loc] = {'File':file} # all condition are in the script ncon += 100 elif 'Fixed' in tar.keys(): loc=tar['Fixed']['Location'] if 'Limit' in tar['Fixed'].keys(): condi = tar['Fixed']['Limit'] else: condi = 0 con = {key: {'Val': tar['Fixed']['Twiss'][key], 'Condition': condi} for key in tar['Fixed']['Twiss'].keys()} target[loc]=con ncon += len(con.keys()) elif 'Save' in tar.keys(): loc=tar['Save']['Location'] IDsave = tar['Save']['SaveID'] if not IDsave in self.save.keys(): print('Needs first the matching point of',IDsave) status = False con = {} ncon+=100 # make sure that this will abort any matching and not filled by dummy constraints else: con = {key:{'Val': self.save[IDsave][key],'Condition':0} for key in self.save[IDsave].keys()} ncon += len(con.keys()) target[loc]=con print('Saved Target',loc,target[loc]) if nvar>ncon: dummycon={} for key in order['Reference']['Twiss'].keys(): if ncon < nvar: dummycon[key]={'Val':order['Reference']['Twiss'][key],'Condition':0} ncon+=1 target['#s']=dummycon order['Target'] = target return order, ID, status def updateMatchPoint(self,ID,val): idx = self.order[ID] if val < 1e-6: color = CGreen elif val < 1: color = CYellow else: color = CRed self.parent.UIMatchOrder.item(idx).setBackground(color) if val<1: idx = self.parent.UIMatchOrder.currentRow()+1 if idx < self.parent.UIMatchOrder.count(): self.parent.UIMatchOrder.setCurrentRow(idx) self.displayMatchingPoint() def getDependence(self,ele): eletag=ele.upper().replace('.','-') if eletag in self.dependence.keys(): return [mag.lower().replace('-','.') for mag in self.dependence[eletag]] else: return None #-------------------------------- # match info def displayMatchingPoint(self): order,ID,status = self.getMatchingPoint() if status is None: return self.parent.UIMatchLabel.setText(ID) if order['Reference'] is None: self.parent.UIMatchStart.setText('Start') else: self.parent.UIMatchStart.setText(order['Reference']['Location']) self.parent.UIMatchKnobs.clear() for ele in order['Variable'].keys(): self.parent.UIMatchKnobs.addItem(ele) self.parent.UIMatchTargets.clear() self.parent.UIMatchTargets.setColumnCount(2) self.parent.UIMatchTargets.setRowCount(0) irow = 0 sym=['=','<','>'] for key0 in order['Target'].keys(): info = order['Target'][key0] for key in info.keys(): self.parent.UIMatchTargets.insertRow(irow) if key == 'File': self.parent.UIMatchTargets.setItem(irow, 0, QtWidgets.QTableWidgetItem('Script')) self.parent.UIMatchTargets.setItem(irow, 1, QtWidgets.QTableWidgetItem(info['File'])) else: self.parent.UIMatchTargets.setItem(irow, 0, QtWidgets.QTableWidgetItem(key0)) val = '%7.3f' % info[key]['Val'] txt = key + sym[info[key]['Condition']] + val.strip() self.parent.UIMatchTargets.setItem(irow, 1, QtWidgets.QTableWidgetItem(txt)) irow+=1 if irow == 0: self.parent.UIMatchTargets.insertRow(irow) self.parent.UIMatchTargets.setItem(0, 0, QtWidgets.QTableWidgetItem('Needs previous matching')) self.parent.UIMatchTargets.resizeColumnsToContents() self.parent.UIMatchTargets.verticalHeader().hide() self.parent.UIMatchTargets.horizontalHeader().hide() return #-------------------------------- # obtaining reference data def getReference(self): return self.getReferenceLocation(),self.getReferenceTwiss() def getReferenceByTag(self,tag): if not tag in self.reference.keys(): return None return self.reference[tag] def getReferenceTwiss(self): self.updateReferenceWidgets() # enforce that the data is consistent twiss = {key:float(str(self.twisswidget[key].text())) for key in self.twisswidget.keys()} return twiss def getReferenceLocation(self): return str(self.parent.UITrackLocation.text()) #################################################### ###### only relevant for tracking ##### interaction of the tracking reference point def swithToUserDefinedLocation(self): n = self.parent.UITrackReference.count() self.parent.UITrackReference.setCurrentIndex(n-1) def updateReferenceWidgets(self): key = str(self.parent.UITrackReference.currentText()) if key == 'User Defined': return twiss = self.reference[key]['Twiss'] name = self.reference[key]['Location'] self.parent.UITrackLocation.setText(name) for key in self.twisswidget.keys(): if key in twiss.keys(): value = twiss[key] else: value = 0 if 'beta' in key: value = 30 self.twisswidget[key].setText('%7.3f' % value) def updateReferenceComboBox(self): self.parent.UITrackReference.clear() for ref in self.reference.keys(): self.parent.UITrackReference.addItem(ref) self.parent.UITrackReference.addItem('User Defined')