Merge pull request #40 from roman-martin/feature_arex

Add arex/arey to available alignment errors

Author: rdemaria

pysixtrack/line.py

@@ -10,6 +10,7 @@
 # missing access to particles._m:
 deg2rad = np.pi / 180.
 
+
 class Line(Element):
     _description = [
         ("elements", "", "List of elements", ()),
@@ -357,22 +358,37 @@ def find_element_ids(self, element_name):
         """
         # will raise error if element not present:
         idx_el = self.element_names.index(element_name)
-        idx_after_el = idx_el + 1
-        if self.element_names[idx_after_el] == element_name + "_aperture":
-            idx_after_el += 1
+        try:
+            # if aperture marker is present
+            idx_after_el = self.element_names.index(element_name + "_aperture") + 1
+        except ValueError:
+            # if aperture marker is not present
+            idx_after_el = idx_el + 1
         return idx_el, idx_after_el
 
     def add_offset_error_to(self, element_name, dx=0, dy=0):
         idx_el, idx_after_el = self.find_element_ids(element_name)
-        if not dx and not dy:
-            return
         xyshift = elements.XYShift(dx=dx, dy=dy)
         inv_xyshift = elements.XYShift(dx=-dx, dy=-dy)
         self.insert_element(idx_el, xyshift, element_name + "_offset_in")
         self.insert_element(
             idx_after_el + 1, inv_xyshift, element_name + "_offset_out"
         )
 
+    def add_aperture_offset_error_to(self, element_name, arex=0, arey=0):
+        idx_el, idx_after_el = self.find_element_ids(element_name)
+        idx_el_aper = idx_after_el - 1
+        if not self.element_names[idx_el_aper] == element_name + "_aperture":
+            # it is allowed to provide arex/arey without providing an aperture
+            print('Info: Element', element_name, ': arex/y provided without aperture -> arex/y ignored')
+            return
+        xyshift = elements.XYShift(dx=arex, dy=arey)
+        inv_xyshift = elements.XYShift(dx=-arex, dy=-arey)
+        self.insert_element(idx_el_aper, xyshift, element_name + "_aperture_offset_in")
+        self.insert_element(
+            idx_after_el + 1, inv_xyshift, element_name + "_aperture_offset_out"
+        )
+
     def add_tilt_error_to(self, element_name, angle):
         '''Alignment error of transverse rotation around s-axis.
         The element corresponding to the given `element_name`
@@ -384,8 +400,6 @@ def add_tilt_error_to(self, element_name, angle):
         by `angle` as well.
         '''
         idx_el, idx_after_el = self.find_element_ids(element_name)
-        if not angle:
-            return
         element = self.elements[self.element_names.index(element_name)]
         if isinstance(element, elements.Multipole) and (
                 element.hxl or element.hyl):
@@ -450,7 +464,7 @@ def apply_madx_errors(self, error_table):
             if error_type == "name":
                 continue
             if any(error_table[error_type]):
-                if error_type in ["dx", "dy", "dpsi"]:
+                if error_type in ["dx", "dy", "dpsi", "arex", "arey"]:
                     # available alignment error
                     continue
                 elif error_type[:1] == "k" and error_type[-1:] == "l":
@@ -478,14 +492,28 @@ def apply_madx_errors(self, error_table):
                 dy = error_table["dy"][i_line]
             except KeyError:
                 dy = 0
-            self.add_offset_error_to(element_name, dx, dy)
+            if dx or dy:
+                self.add_offset_error_to(element_name, dx, dy)
 
             # add tilt
             try:
                 dpsi = error_table["dpsi"][i_line]
+            except KeyError:
+                dpsi = 0
+            if dpsi:
                 self.add_tilt_error_to(element_name, angle=dpsi / deg2rad)
+
+            # add aperture-only offset
+            try:
+                arex = error_table["arex"][i_line]
+            except KeyError:
+                arex = 0
+            try:
+                arey = error_table["arey"][i_line]
             except KeyError:
-                pass
+                arey = 0
+            if arex or arey:
+                self.add_aperture_offset_error_to(element_name, arex, arey)
 
             # add multipole error
             knl = [

tests/test_line.py

@@ -35,13 +35,15 @@ def test_line():
     assert len(line) == n_elements
 
     line.add_offset_error_to(multipole_name, dx=0, dy=0)
+    n_elements += 2
     assert len(line) == n_elements
 
     line.add_offset_error_to(multipole_name, dx=0.2, dy=-0.003)
     n_elements += 2
     assert len(line) == n_elements
 
     line.add_tilt_error_to(multipole_name, angle=0)
+    n_elements += 2
     assert len(line) == n_elements
 
     line.add_tilt_error_to(multipole_name, angle=0.1)

tests/test_madx_import.py

@@ -107,3 +107,192 @@ def test_madx_import():
             )
         else:
             raise ValueError("mode not understood")
+
+
+def test_error_import():
+    cpymad_spec = util.find_spec("cpymad")
+    if cpymad_spec is None:
+        print("cpymad is not available - abort test")
+        sys.exit(0)
+
+    from cpymad.madx import Madx
+
+    madx = Madx()
+
+    madx.input('''
+        MQ1: Quadrupole, K1:=KQ1, L=1.0, apertype=CIRCLE, aperture={0.04};
+        MQ2: Quadrupole, K1:=KQ2, L=1.0, apertype=CIRCLE, aperture={0.04};
+        MQ3: Quadrupole, K1:=0.0, L=1.0, apertype=CIRCLE, aperture={0.04};
+
+        KQ1 = 0.02;
+        KQ2 = -0.02;
+
+        testseq: SEQUENCE, l = 20.0;
+            MQ1, at =  5;
+            MQ2, at = 12;
+            MQ3, at = 18;
+        ENDSEQUENCE;
+
+        !---the usual stuff
+        BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
+        USE, SEQUENCE=testseq;
+
+
+        Select, flag=makethin, pattern="MQ1", slice=2;
+        makethin, sequence=testseq;
+
+        use, sequence=testseq;
+
+        !---assign misalignments and field errors
+        select, flag = error, clear;
+        select, flag = error, pattern = "MQ1";
+        ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
+        select, flag = error, clear;
+        select, flag = error, pattern = "MQ2";
+        ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
+        select, flag = error, clear;
+        select, flag = error, pattern = "MQ3";
+        ealign, dx = 0.00, dy = 0.00, arex = 0.00, arey = 0.00, dpsi = 0.00;
+        efcomp, DKN = {0.0, 0.0, 0.001, 0.002}, DKS = {0.0, 0.0, 0.003, 0.004};
+        select, flag = error, full;
+    ''')
+    seq = madx.sequence.testseq
+    # store already applied errors:
+    madx.command.esave(file='lattice_errors.err')
+    madx.command.readtable(file='lattice_errors.err', table="errors")
+    os.remove('lattice_errors.err')
+    errors = madx.table.errors
+
+    pysixtrack_line = pysixtrack.Line.from_madx_sequence(seq, install_apertures=True)
+    pysixtrack_line.apply_madx_errors(errors)
+    madx.input('stop;')
+
+    expected_element_num = (
+            2           # start and end marker
+            + 6          # drifts (including drift between MQ1 slices)
+            + 3 + 2        # quadrupoles + MQ1 slices
+            + 3 + 2        # corresponding aperture elements
+            + 2*(3+1)    # dx/y in/out for MQ1 slices and MQ2
+            + 2          # tilt in/out for MQ2
+            + 2*3        # arex/y in/out for MQ1 slices
+            )
+    assert len(pysixtrack_line) == expected_element_num
+
+    expected_element_order = [
+            pysixtrack.elements.Drift,          # start marker
+            pysixtrack.elements.Drift,
+            pysixtrack.elements.XYShift,        # dx/y in of MQ1 1st slice
+            pysixtrack.elements.Multipole,      # MQ1 1st slice
+            pysixtrack.elements.XYShift,        # arex/y in for MQ1 1st slice
+            pysixtrack.elements.LimitEllipse,   # MQ1 1st slice aperture
+            pysixtrack.elements.XYShift,        # arex/y out for MQ1 1st slice
+            pysixtrack.elements.XYShift,        # dx/y out for MQ1 1st slice
+            pysixtrack.elements.Drift,
+            pysixtrack.elements.XYShift,        # dx/y in for MQ1 marker
+            pysixtrack.elements.Drift,          # MQ1 marker
+            pysixtrack.elements.XYShift,        # arex/y in for MQ1 marker
+            pysixtrack.elements.LimitEllipse,   # MQ1 marker aperture
+            pysixtrack.elements.XYShift,        # arex/y out for MQ1 marker
+            pysixtrack.elements.XYShift,        # dx/y out for MQ1 marker
+            pysixtrack.elements.Drift,
+            pysixtrack.elements.XYShift,        # dx/y in for MQ1 2nd slice
+            pysixtrack.elements.Multipole,      # MQ1 2nd slice
+            pysixtrack.elements.XYShift,        # arex/y in for MQ1 2nd slice
+            pysixtrack.elements.LimitEllipse,   # MQ1 2nd slice aperture
+            pysixtrack.elements.XYShift,        # arex/y out for MQ1 2nd slice
+            pysixtrack.elements.XYShift,        # dx/y out for MQ1 2nd slice
+            pysixtrack.elements.Drift,
+            pysixtrack.elements.XYShift,        # dx/y in for MQ2
+            pysixtrack.elements.SRotation,      # tilt in for MQ2
+            pysixtrack.elements.Multipole,      # MQ2
+            pysixtrack.elements.LimitEllipse,   # MQ2 aperture
+            pysixtrack.elements.SRotation,      # tilt out for MQ2
+            pysixtrack.elements.XYShift,        # dx/y out for MQ2
+            pysixtrack.elements.Drift,
+            pysixtrack.elements.Multipole,      # MQ3
+            pysixtrack.elements.LimitEllipse,   # MQ3 aperture
+            pysixtrack.elements.Drift,
+            pysixtrack.elements.Drift           # end marker
+            ]
+    for element, expected_element in zip(pysixtrack_line.elements,
+                                         expected_element_order):
+        assert isinstance(element, expected_element)
+
+    idx_MQ3 = pysixtrack_line.element_names.index('mq3')
+    MQ3 = pysixtrack_line.elements[idx_MQ3]
+    assert abs(MQ3.knl[2] - 0.001) < 1e-14
+    assert abs(MQ3.knl[3] - 0.002) < 1e-14
+    assert abs(MQ3.ksl[2] - 0.003) < 1e-14
+    assert abs(MQ3.ksl[3] - 0.004) < 1e-14
+
+
+def test_neutral_errors():
+    # make sure that some misaligned drifts do not influence particle
+    cpymad_spec = util.find_spec("cpymad")
+    if cpymad_spec is None:
+        print("cpymad is not available - abort test")
+        sys.exit(0)
+
+    from cpymad.madx import Madx
+
+    madx = Madx()
+
+    madx.input('''
+        T1: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
+        T2: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
+        T3: Collimator, L=1.0, apertype=CIRCLE, aperture={0.5};
+
+        KQ1 = 0.02;
+        KQ2 = -0.02;
+
+        testseq: SEQUENCE, l = 20.0;
+            T1, at =  5;
+            T2, at = 12;
+            T3, at = 18;
+        ENDSEQUENCE;
+
+        !---the usual stuff
+        BEAM, PARTICLE=PROTON, ENERGY=7000.0, EXN=2.2e-6, EYN=2.2e-6;
+        USE, SEQUENCE=testseq;
+
+
+        Select, flag=makethin, pattern="MQ1", slice=2;
+        makethin, sequence=testseq;
+
+        use, sequence=testseq;
+
+        !---misalign collimators
+        select, flag = error, clear;
+        select, flag = error, pattern = "T1";
+        ealign, dx = 0.01, dy = 0.01, arex = 0.02, arey = 0.02;
+        select, flag = error, clear;
+        select, flag = error, pattern = "T2";
+        ealign, dx = 0.04, dy = 0.04, dpsi = 0.1;
+        select, flag = error, clear;
+        select, flag = error, pattern = "T3";
+        ealign, dx = 0.02, dy = 0.01, arex = 0.03, arey = 0.02, dpsi = 0.1;
+        select, flag = error, full;
+    ''')
+    seq = madx.sequence.testseq
+    # store already applied errors:
+    madx.command.esave(file='lattice_errors.err')
+    madx.command.readtable(file='lattice_errors.err', table="errors")
+    os.remove('lattice_errors.err')
+    errors = madx.table.errors
+
+    pysixtrack_line = pysixtrack.Line.from_madx_sequence(seq, install_apertures=True)
+    pysixtrack_line.apply_madx_errors(errors)
+    madx.input('stop;')
+
+    initial_x = 0.025
+    initial_y = -0.015
+
+    particle = pysixtrack.Particles()
+    particle.x = initial_x
+    particle.y = initial_y
+    particle.state = 1
+
+    pysixtrack_line.track(particle)
+
+    assert abs(particle.x-initial_x) < 1e-14
+    assert abs(particle.y-initial_y) < 1e-14