From ed93addae751bb8eb2fd9fd0c14ca0bd101ad8dc Mon Sep 17 00:00:00 2001
From: Yngve Inntjore Levinsen <Yngve.Levinsen@esss.se>
Date: Fri, 22 Jan 2016 11:36:28 +0100
Subject: [PATCH] added get_elem_idx() which gets the element with tracewin
index given added update_adj() which updates the lattice based on
Adjusted_Values.txt
---
ess/lib_tw.py | 95 +++++++++++++++++++++++++++++++++++++++++++++++++--
1 file changed, 93 insertions(+), 2 deletions(-)
diff --git a/ess/lib_tw.py b/ess/lib_tw.py
index 7e760be..24cfd77 100644
--- a/ess/lib_tw.py
+++ b/ess/lib_tw.py
@@ -36,7 +36,7 @@ class LATTICE:
def __init__(self,file_name_lat,file_name_fmap,freq=352.21,gamma=1.0):
# In case file_name_fmap is str
- if isinstance(file_name_fmap,basestring)==True: file_name_fmap=[file_name_fmap]
+ if isinstance(file_name_fmap,basestring): file_name_fmap=[file_name_fmap]
# Elem/comm class dict
dic_cls={'DRIFT':DRIFT,
@@ -84,6 +84,28 @@ class LATTICE:
# Assign idx, idx_elem, s, freq, apt
self.update_idx()
+ def get_elem_idx(self,i):
+ '''
+ Get a TraceWin index number
+
+ Note: We start counting from 0, TW starts from 1
+ '''
+ for element in self.lst:
+ if element.idx_elem==i-1:
+ return element
+
+ def get_steerer_for(self,idx_elem):
+ '''
+ Returns the steerer object for an element (e.g. quad)
+ '''
+ previous=None
+ for element in self.lst:
+ if element.idx_elem+1==idx_elem:
+ if previous.typ=='STEERER':
+ return previous
+ return None
+ previous=element
+
def update_idx(self):
# Assign/update idx, idx_elem, s, freq
@@ -136,7 +158,7 @@ class LATTICE:
# Assign/update steerers
for i in range(len(self.lst)):
- if self.lst[i].idx_elem in BLx.keys():
+ if self.lst[i].idx_elem in BLx:
idx_elem=self.lst[i].idx_elem
if self.lst[i].typ=='THIN_STEERING':
self.lst[i].BLx=BLx[idx_elem]
@@ -148,6 +170,75 @@ class LATTICE:
self.lst[k].By=BLy[idx_elem]/self.lst[i].L
break
+ def update_adj(self,file_name="Adjusted_Values.txt"):
+ '''
+ Assign/update correction values from "Adjusted_Values.txt".
+
+ WARNING: many corner cases, what happens if multiple adjust
+ commands have corrected same element for example?
+ Use with care!
+ '''
+ with open(file_name,'r') as file:
+ # First we read in all corrections into dictionaries
+ values={}
+ counts={}
+ for lin in file:
+ lin=lin.split()
+ i=int(lin[1][1:-1])
+ if i=='ERROR' and lin[0]=='BEAM':
+ continue
+ settings={}
+ for j in xrange(len(lin[2:])/3):
+ k=int(lin[2+3*j])
+ val=float(lin[4+3*j])
+ if k in settings:
+ settings[k].append(val)
+ else:
+ settings[k]=[val]
+ values[i]=settings
+ for j in settings:
+ counts[j]=0
+
+ # now we will do all the ADJUST_STEERER ones
+ corr_next=False
+ for el in self.lst:
+ if isinstance(el,COMM) and el.typ=='ADJUST_STEERER':
+ i=int(el.para[0])
+ if i in values:
+ corr_next=values[i]
+ elif corr_next:
+ if isinstance(el,STEERER):
+ # Bx and By are corrected..
+ vals=corr_next.values()[0]
+ el.Bx=vals[0]
+ el.By=vals[1]
+ corr_next=False
+
+ # now we will do all the ADJUST ones
+ corr_next=False
+ # The TraceWin index number of the current element:
+ current=-1
+ for el in self.lst:
+ if el.idx_elem!=-1:
+ current=el.idx_elem+1
+ if isinstance(el,COMM) and el.typ=='ADJUST':
+ # The index of the corrector scheme
+ i = int(el.para[0])
+ # the parameter column to vary:
+ j = int(el.para[1])-1
+ # The TraceWin element index of the element we will vary:
+ k = current+1
+ # This corrector might not be used:
+ if i in values:
+ # We will correct the next active element in lattice:
+ value = values[i][k][counts[k]]
+ # We have now used one value of the total in the current corrector:
+ counts[k]+=1
+ # the element to vary:
+ vary=self.get_elem_idx(current+1)
+ vary.para[j]=value
+ vary.update()
+
def get_tw(self,file_name):
with open(file_name,'w') as file:
--
GitLab