Skip to content
Snippets Groups Projects
Commit 000a3495 authored by Yngve Levinsen's avatar Yngve Levinsen
Browse files

new function lib_tw.get_diag_idx() 

some minor cleaning as well
parent dfdfed46
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
ess/lib_tw.py
+ 12
− 6
View file @ 000a3495
...@@ -160,6 +160,17 @@ class LATTICE: ...@@ -160,6 +160,17 @@ class LATTICE:
if element.idx_elem == i - 1: if element.idx_elem == i - 1:
return element return element
def get_diag_idx(self, i):
"""
Get a list of DIAG that has the given index
"""
diags = []
for element in self.lst:
if isinstance(element, lib_tw_elem.DIAG):
if element.idx_diag == i:
diags.append(element)
return diags
def get_steerer_for(self, idx_elem): def get_steerer_for(self, idx_elem):
""" """
Returns the steerer object for an element (e.g. quad) Returns the steerer object for an element (e.g. quad)
...@@ -201,12 +212,7 @@ class LATTICE: ...@@ -201,12 +212,7 @@ class LATTICE:
except IndexError: except IndexError:
pass pass
else: else:
for k in range(i)[::-1]: self.lst[i].apt = self.lst[i - 1].apt
try:
self.lst[i].apt = self.lst[k].apt
break
except IndexError:
pass
except AttributeError: except AttributeError:
pass pass
......
ess/lib_tw_elem.py
+ 28
− 62
View file @ 000a3495
...@@ -149,24 +149,30 @@ class ELEM: ...@@ -149,24 +149,30 @@ class ELEM:
# Option instances # Option instances
self.L = 0.0 self.L = 0.0
self.apt = None self.apt = None
# idx is essentially similar to "line number"
self.idx = -1 self.idx = -1
# idx_elem is the "tracewin index count" - 1 (we start at 0)
self.idx_elem = -1 self.idx_elem = -1
# For diagnostic elements, the index given in the lattice file:
self.idx_diag = -1
self.s = 0.0 self.s = 0.0
self.gamma = 1.0 self.gamma = 1.0
self.freq = 352.21 self.freq = 352.21
self._idx_madx = "0"
def update_idx(self): def update_idx(self):
# Update idx, idx_elem, s # Update idx, idx_elem, s
self.idx += 1 self.idx += 1
self.idx_elem += 1 self.idx_elem += 1
self._idx_madx = f"{self.idx_elem+1:04d}"
self.s += self.L self.s += self.L
def get_tw(self): def get_tw(self):
if self.name != "": if self.name:
lin = self.name + ": " + self.typ + " " + " ".join(self.para) lin = self.name + ": " + self.typ + " " + " ".join(self.para)
if self.name == "": else:
lin = self.typ + " " + " ".join(self.para) lin = self.typ + " " + " ".join(self.para)
return lin return lin
...@@ -297,12 +303,10 @@ class FIELD_MAP(ELEM): ...@@ -297,12 +303,10 @@ class FIELD_MAP(ELEM):
def get_madx(self): def get_madx(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "": if self.name != "":
lin = self.name + ": RFCAVITY, " lin = self.name + ": RFCAVITY, "
if self.name == "": if self.name == "":
lin = "ELEM" + i + "_FIELDMAP" + ": RFCAVITY, " lin = "ELEM" + self._idx_madx + "_FIELDMAP" + ": RFCAVITY, "
lin += "L=" + str(self.L) + ", " lin += "L=" + str(self.L) + ", "
lin += "FREQ=" + str(self.freq) + ", " lin += "FREQ=" + str(self.freq) + ", "
lin += "VOLT=" + str(self.E0TL) + ", " lin += "VOLT=" + str(self.E0TL) + ", "
...@@ -312,12 +316,10 @@ class FIELD_MAP(ELEM): ...@@ -312,12 +316,10 @@ class FIELD_MAP(ELEM):
def get_fluka(self): def get_fluka(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "": if self.name != "":
lin = self.name + " RFCAVITY " lin = self.name + " RFCAVITY "
if self.name == "": if self.name == "":
lin = "ELEM" + i + "_FIELDMAP" + " RFCAVITY " lin = "ELEM" + self._idx_madx + "_FIELDMAP" + " RFCAVITY "
lin += str(self.L) + " " + str(self.s) + " " lin += str(self.L) + " " + str(self.s) + " "
lin += "0.0 0.0 0.0 " lin += "0.0 0.0 0.0 "
lin += "CIRCLE " + str(self.apt * 1e-3) + " 0.0" lin += "CIRCLE " + str(self.apt * 1e-3) + " 0.0"
...@@ -381,23 +383,19 @@ class DRIFT(ELEM): ...@@ -381,23 +383,19 @@ class DRIFT(ELEM):
def get_madx(self): def get_madx(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "": if self.name != "":
lin = self.name + ": DRIFT, L=" + str(self.L) + ";" lin = self.name + ": DRIFT, L=" + str(self.L) + ";"
if self.name == "": if self.name == "":
lin = "ELEM" + i + "_DRIFT" + ": DRIFT, L=" + str(self.L) + ";" lin = "ELEM" + self._idx_madx + "_DRIFT" + ": DRIFT, L=" + str(self.L) + ";"
return lin return lin
def get_fluka(self): def get_fluka(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "": if self.name != "":
lin = self.name + " DRIFT " lin = self.name + " DRIFT "
if self.name == "": if self.name == "":
lin = "ELEM" + i + "_DRIFT" + " DRIFT " lin = "ELEM" + self._idx_madx + "_DRIFT" + " DRIFT "
lin += str(self.L) + " " + str(self.s) + " " lin += str(self.L) + " " + str(self.s) + " "
lin += "0.0 0.0 0.0 " lin += "0.0 0.0 0.0 "
if self.apty == 0.0: if self.apty == 0.0:
...@@ -1048,86 +1046,54 @@ class APERTURE(ELEM): ...@@ -1048,86 +1046,54 @@ class APERTURE(ELEM):
return lin return lin
class DIAG_POSITION(ELEM): class DIAG(ELEM):
""" """
""" """
def __init__(self, name, typ, para): def __init__(self, name, typ, para):
ELEM.__init__(self, name, typ, para) ELEM.__init__(self, name, typ, para)
self.idx_diag = int(para[0])
self.__diag_type = "INSTRUMENT"
self.__diag_name = "DIAG"
def get_madx(self): def get_madx(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "": if self.name != "":
lin = self.name + ": MONITOR, L=0;" lin = f"{self.name}: {self.__diag_type}, L=0;"
if self.name == "": if self.name == "":
lin = "ELEM" + i + "_BPM" + ": MONITOR, L=0;" lin = f"ELEM{self._idx_madx}_{self.__diag_name}: {self.__diag_type}, L=0;"
return lin return lin
def get_fluka(self): def get_fluka(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "": if self.name != "":
lin = self.name + " MONITOR " lin = f"{self.name} {self.__diag_type} "
if self.name == "": if self.name == "":
lin = "ELEM" + i + "_BPM" + " MONITOR " print(self._idx_madx)
lin += str(self.L) + " " + str(self.s) + " " lin = f"ELEM{self._idx_madx}_{self.__diag_name} {self.__diag_type} "
lin += "0.0 0.0 0.0 " lin += f"{self.L} {self.s} 0.0 0.0 0.0 "
lin += "CIRCLE " + str(self.apt * 1e-3) + " 0.0" lin += f"CIRCLE {self.apt * 1e-3} 0.0"
return lin return lin
def get_mars(self): def get_mars(self):
lin = '"MONITOR" "" "" ' + str(self.s) + " " + str(self.L) + " " lin = f'"{self.__diag_type}" "" "" {self.s} {self.L} 0 0 0 0 0'
lin += "0 0 0 0 0"
return lin return lin
class DIAG(ELEM): class DIAG_POSITION(DIAG):
""" """
""" """
def __init__(self, name, typ, para): def __init__(self, name, typ, para):
ELEM.__init__(self, name, typ, para) DIAG.__init__(self, name, typ, para)
self.__diag_type = "MONITOR"
def get_madx(self): self.__diag_name = "BPM"
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "":
lin = self.name + ": INSTRUMENT, L=0;"
if self.name == "":
lin = "ELEM" + i + "_DIAG" + ": INSTRUMENT, L=0;"
return lin
def get_fluka(self):
i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
if self.name != "":
lin = self.name + " INSTRUMENT "
if self.name == "":
lin = "ELEM" + i + "_DIAG" + " INSTRUMENT "
lin += str(self.L) + " " + str(self.s) + " "
lin += "0.0 0.0 0.0 "
lin += "CIRCLE " + str(self.apt * 1e-3) + " 0.0"
return lin
def get_mars(self):
lin = '"INSTRUMENT" "" "" ' + str(self.s) + " " + str(self.L) + " "
lin += "0 0 0 0 0"
return lin
# ---- Active commands # ---- Active commands
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment