First round of matching

model.py

@@ -1,5 +1,5 @@
 import json
-
+import copy
 import numpy as np
 from onlinemodel.core import Facility
 from onlinemodel.madx import CMadX
@@ -13,6 +13,10 @@ class Model:
         self.order = None
 
         self.madx = CMadX()
+        self.startTwiss = None
+        self.startEnergy = None
+        self.energyReference ='SINLH02.MBND100'
+
         # hook up events
         self.eventHandling()
 
@@ -20,10 +24,11 @@ class Model:
         return self.om.Version
 
     def updateEnergy(self,E0):
-        self.om.forceEnergyAt('SINLH02.MBND100', E0*1e6)
+        if isinstance(E0,list):
+            E0=E0[0]
+        self.om.forceEnergyAt(self.energyReference, E0*1e6)
 
     def updateElement(self,name,val):
-
         if 'MQUA' in name or 'MQSK' in name:
             L = self.om.getRegExpElement(name[0:7], name[8:15],'Length')[0]
             self.om.setRegExpElement(name[0:7], name[8:15], 'k1', float(val[0])/L)
@@ -70,8 +75,8 @@ class Model:
         rf={}
         undulators={}
         kicker={}
-        loc = 'SINLH02.MBND100'
-        energy={'location': loc, 'energy':self.om.EnergyAt(loc)[0]}
+        energy={'location': self.energyReference, 'energy' : 1e-6*self.om.EnergyAt(self.energyReference)[0]}
+
         for ele in elements:
             if 'MQUA' in ele.Name or 'MQSK' in ele.Name:
                 quadrupoles[ele.Name]={'k1':ele.k1,'k1L':ele.k1*ele.Length}
@@ -86,7 +91,8 @@ class Model:
                 rf[ele.Name]={'Gradient':ele.Gradient*ele.Length,'Phase':ele.Phase}
             elif 'MKAC' in ele.Name or 'MKDC' in ele.Name:
                 kicker[ele.Name] = {'cory': ele.cory}
-        return {'Quadrupole':quadrupoles,'Sextupole':sextupoles,'Dipole':dipoles,'RF':rf,'Undulator':undulators,'Kicker':kicker,'Energy':energy}
+        return {'Quadrupole':quadrupoles,'Sextupole':sextupoles,'Dipole':dipoles,'RF':rf,'Undulator':undulators,
+                'Kicker':kicker,'Energy':energy, 'InitialCondition':self.startTwiss}
 
     def loadSettingsGroup(self,group,fields):
         for key in group.keys():
@@ -101,7 +107,11 @@ class Model:
         self.loadSettingsGroup(settings['RF'], ['Gradient','Phase'])
         self.loadSettingsGroup(settings['Undulator'], ['K','kx','ky'])
         self.loadSettingsGroup(settings['Kicker'], ['cory'])
-        self.om.forceEnergyAt(settings['Energy']['location'],settings['Energy']['energy'][0])
+        self.startEnergy = settings['Energy']['energy']
+        self.energyReference = settings['Energy']['location']
+        self.startTwiss = settings['InitialCondition']
+        print('initial condition:',self.startTwiss)
+        self.updateEnergy(self.startEnergy)
 
 
 
@@ -146,6 +156,18 @@ class Model:
     def eventHandling(self):
         self.parent.UITrack.clicked.connect(self.track)
 
+    def match(self,sequence,destination, variables,conditions,periodic=False,plot=True):
+        self.setBranch(destination.upper())
+        twiss={'energy0':150.}
+        self.madx.updateVariables(twiss)
+        res,twiss,tar=self.madx.match(sequence,variables,conditions,periodic)
+        energy = self.calcEnergyProfile(twiss)
+        matchtwiss={'betax':twiss.betx[0],'betay':twiss.betay[0],'alphax':twiss.alfx[0],'alphay':twiss.alfy[0]}
+        if plot:
+            self.parent.plot.newData(twiss, energy)
+        return matchtwiss,tar
+
+
     def track(self):
         start = str(self.parent.UITrackStart.text()).upper()
         end = str(self.parent.UITrackEnd.text()).upper()
@@ -155,7 +177,9 @@ class Model:
             end = end[0:7]
         twiss0 = self.parent.reference.getReferencePoint()
         refloc = self.parent.reference.getReferenceLocation().upper()
-        twiss0['energy'] = 150.
+        if refloc == 'START':
+            refloc = start.upper()
+        twiss0['energy'] = self.startEnergy
         start, end = self.checkRange(start, end, refloc[0:7])
         if start is None:
             return
@@ -168,13 +192,17 @@ class Model:
         self.setBranch(end.upper())
         if not refloc == start:
             twiss0 = self.doBackTrack(refloc,start,twiss0)
+        self.startTwiss=copy.deepcopy(twiss0)
         self.madx.updateVariables(twiss0)
         twiss = self.madx.track(start+'$START',end+'$END')
+        energy = self.calcEnergyProfile(twiss)
+        if plot:
+            self.parent.plot.newData(twiss,energy)
+
+    def calcEnergyProfile(self,twiss):
 
         energy = np.array([0. for i in range(len(twiss.betx))])
-
         e0 = 0.
-        de = 0
         for i, name in enumerate(twiss.name):
             if len(name) > 15:
                 elename = name[0:15]
@@ -185,8 +213,8 @@ class Model:
                         energy[:i]+=erg[0]*1e-6
                     e0 = (erg[0]+erg[1])*1e-6
             energy[i]=e0
-        if plot:
-            self.parent.plot.newData(twiss,energy)
+        return energy
+
 
     def doBackTrack(self,start=None,end=None,twiss0=None):
         twiss0['alphax0'] = -twiss0['alphax0']  # revert particle trajectories