From 000a349589456904967d54d338320968403c60fe Mon Sep 17 00:00:00 2001
From: Yngve Levinsen <yngve.levinsen@ess.eu>
Date: Mon, 27 Jan 2020 10:48:27 +0100
Subject: [PATCH] new function lib_tw.get_diag_idx() some minor cleaning as
well
---
ess/lib_tw.py | 18 ++++++----
ess/lib_tw_elem.py | 90 +++++++++++++++-------------------------------
2 files changed, 40 insertions(+), 68 deletions(-)
diff --git a/ess/lib_tw.py b/ess/lib_tw.py
index 9744939..a271fcd 100644
--- a/ess/lib_tw.py
+++ b/ess/lib_tw.py
@@ -160,6 +160,17 @@ class LATTICE:
if element.idx_elem == i - 1:
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):
"""
Returns the steerer object for an element (e.g. quad)
@@ -201,12 +212,7 @@ class LATTICE:
except IndexError:
pass
else:
- for k in range(i)[::-1]:
- try:
- self.lst[i].apt = self.lst[k].apt
- break
- except IndexError:
- pass
+ self.lst[i].apt = self.lst[i - 1].apt
except AttributeError:
pass
diff --git a/ess/lib_tw_elem.py b/ess/lib_tw_elem.py
index 38feaa9..8e4bd9c 100644
--- a/ess/lib_tw_elem.py
+++ b/ess/lib_tw_elem.py
@@ -149,24 +149,30 @@ class ELEM:
# Option instances
self.L = 0.0
self.apt = None
+ # idx is essentially similar to "line number"
self.idx = -1
+ # idx_elem is the "tracewin index count" - 1 (we start at 0)
self.idx_elem = -1
+ # For diagnostic elements, the index given in the lattice file:
+ self.idx_diag = -1
self.s = 0.0
self.gamma = 1.0
self.freq = 352.21
+ self._idx_madx = "0"
def update_idx(self):
# Update idx, idx_elem, s
self.idx += 1
self.idx_elem += 1
+ self._idx_madx = f"{self.idx_elem+1:04d}"
self.s += self.L
def get_tw(self):
- if self.name != "":
+ if self.name:
lin = self.name + ": " + self.typ + " " + " ".join(self.para)
- if self.name == "":
+ else:
lin = self.typ + " " + " ".join(self.para)
return lin
@@ -297,12 +303,10 @@ class FIELD_MAP(ELEM):
def get_madx(self):
- i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
-
if self.name != "":
lin = self.name + ": RFCAVITY, "
if self.name == "":
- lin = "ELEM" + i + "_FIELDMAP" + ": RFCAVITY, "
+ lin = "ELEM" + self._idx_madx + "_FIELDMAP" + ": RFCAVITY, "
lin += "L=" + str(self.L) + ", "
lin += "FREQ=" + str(self.freq) + ", "
lin += "VOLT=" + str(self.E0TL) + ", "
@@ -312,12 +316,10 @@ class FIELD_MAP(ELEM):
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 + " RFCAVITY "
if self.name == "":
- lin = "ELEM" + i + "_FIELDMAP" + " RFCAVITY "
+ lin = "ELEM" + self._idx_madx + "_FIELDMAP" + " RFCAVITY "
lin += str(self.L) + " " + str(self.s) + " "
lin += "0.0 0.0 0.0 "
lin += "CIRCLE " + str(self.apt * 1e-3) + " 0.0"
@@ -381,23 +383,19 @@ class DRIFT(ELEM):
def get_madx(self):
- i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
-
if self.name != "":
lin = self.name + ": DRIFT, L=" + str(self.L) + ";"
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
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 + " DRIFT "
if self.name == "":
- lin = "ELEM" + i + "_DRIFT" + " DRIFT "
+ lin = "ELEM" + self._idx_madx + "_DRIFT" + " DRIFT "
lin += str(self.L) + " " + str(self.s) + " "
lin += "0.0 0.0 0.0 "
if self.apty == 0.0:
@@ -1048,86 +1046,54 @@ class APERTURE(ELEM):
return lin
-class DIAG_POSITION(ELEM):
+class DIAG(ELEM):
"""
"""
def __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):
- i = "0" * (3 - int(numpy.log10(self.idx_elem + 1))) + str(self.idx_elem + 1)
-
if self.name != "":
- lin = self.name + ": MONITOR, L=0;"
+ lin = f"{self.name}: {self.__diag_type}, L=0;"
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
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 + " MONITOR "
+ lin = f"{self.name} {self.__diag_type} "
if self.name == "":
- lin = "ELEM" + i + "_BPM" + " MONITOR "
- lin += str(self.L) + " " + str(self.s) + " "
- lin += "0.0 0.0 0.0 "
- lin += "CIRCLE " + str(self.apt * 1e-3) + " 0.0"
+ print(self._idx_madx)
+ lin = f"ELEM{self._idx_madx}_{self.__diag_name} {self.__diag_type} "
+ lin += f"{self.L} {self.s} 0.0 0.0 0.0 "
+ lin += f"CIRCLE {self.apt * 1e-3} 0.0"
return lin
def get_mars(self):
- lin = '"MONITOR" "" "" ' + str(self.s) + " " + str(self.L) + " "
- lin += "0 0 0 0 0"
+ lin = f'"{self.__diag_type}" "" "" {self.s} {self.L} 0 0 0 0 0'
return lin
-class DIAG(ELEM):
+class DIAG_POSITION(DIAG):
"""
"""
def __init__(self, name, typ, para):
- ELEM.__init__(self, name, typ, para)
-
- def get_madx(self):
-
- 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
+ DIAG.__init__(self, name, typ, para)
+ self.__diag_type = "MONITOR"
+ self.__diag_name = "BPM"
# ---- Active commands
--
GitLab