Merge pull request #30 from giadarol/develop

Testing and some fixes from Hannes and myself

Author: rdemaria

examples/load_from_mad/002_compare_against_sixinput.py

@@ -9,7 +9,7 @@
     lmad = pysixtrack.Line.from_dict(pickle.load(fid))
 
 sixinput = sixtracktools.sixinput.SixInput("sixtrack/")
-lsix, other = pysixtrack.Line.from_sixinput(sixinput)
+lsix = pysixtrack.Line.from_sixinput(sixinput)
 
 original_length = lmad.get_length()
 assert (lsix.get_length() - original_length) < 1e-6

pysixtrack/be_beamfields/beambeam.py

@@ -1,3 +1,4 @@
+import numpy as np
 from scipy.constants import e as qe
 from scipy.constants import c as clight
 
@@ -17,9 +18,9 @@ class BeamBeam4D(Element):
             "Charge of the interacting distribution (strong beam)",
             0,
         ),
-        ("sigma_x", "m", "Horizontal size of the strong beam (r.m.s.)", 0),
-        ("sigma_y", "m", "Vertical size of the strong beam (r.m.s.)", 0),
-        ("beta_r", "", "Relativistic beta of the stron beam", 0),
+        ("sigma_x", "m", "Horizontal size of the strong beam (r.m.s.)", 1.0),
+        ("sigma_y", "m", "Vertical size of the strong beam (r.m.s.)", 1.0),
+        ("beta_r", "", "Relativistic beta of the stron beam", 1.0),
         (
             "x_bb",
             "m",
@@ -112,20 +113,20 @@ class BeamBeam6D(Element):
             "charge_slices",
             "qe",
             "Charge of the interacting slices (strong beam)",
-            (),
+            (0.0),
         ),
         (
             "zeta_slices",
             "m",
             "Longitudinal position of the interacting"
             "slices (>0 head of the strong).",
-            (),
+            (0.0),
         ),
         (
             "sigma_11",
             "m^2",
             "Sigma_11 element of the sigma matrix of the strong beam",
-            0,
+            1.0,
         ),
         (
             "sigma_12",
@@ -167,7 +168,7 @@ class BeamBeam6D(Element):
             "sigma_33",
             "m^2",
             "Sigma_33 element of the sigma matrix of the strong beam",
-            0,
+            1.0,
         ),
         (
             "sigma_34",
@@ -223,8 +224,8 @@ def track(self, p):
                 self.alpha,
                 self.x_bb_co,
                 self.y_bb_co,
-                self.charge_slices,
-                self.zeta_slices,
+                np.atleast_1d(self.charge_slices),
+                np.atleast_1d(self.zeta_slices),
                 self.sigma_11,
                 self.sigma_12,
                 self.sigma_13,

pysixtrack/be_beamfields/gaussian_fields.py

@@ -123,32 +123,36 @@ def _get_Ex_Ey_Gx_Gy_gauss(
         )
 
         if not skip_Gs:
-            Gx = (
-                1
-                / (2.0 * (x * x + y * y))
-                * (
-                    y * Ey
-                    - x * Ex
-                    + 1.0
-                    / (2 * pi * epsilon_0 * sigma * sigma)
-                    * x
-                    * x
-                    * exp(-(x * x + y * y) / (2.0 * sigma * sigma))
+            if abs(x) + abs(y) < min_sigma_diff:
+                Gx = 0.0
+                Gy = 0.0
+            else:
+                Gx = (
+                    1
+                    / (2.0 * (x * x + y * y))
+                    * (
+                        y * Ey
+                        - x * Ex
+                        + 1.0
+                        / (2 * pi * epsilon_0 * sigma * sigma)
+                        * x
+                        * x
+                        * exp(-(x * x + y * y) / (2.0 * sigma * sigma))
+                    )
                 )
-            )
-            Gy = (
-                1.0
-                / (2 * (x * x + y * y))
-                * (
-                    x * Ex
-                    - y * Ey
-                    + 1.0
-                    / (2 * pi * epsilon_0 * sigma * sigma)
-                    * y
-                    * y
-                    * exp(-(x * x + y * y) / (2.0 * sigma * sigma))
+                Gy = (
+                    1.0
+                    / (2 * (x * x + y * y))
+                    * (
+                        x * Ex
+                        - y * Ey
+                        + 1.0
+                        / (2 * pi * epsilon_0 * sigma * sigma)
+                        * y
+                        * y
+                        * exp(-(x * x + y * y) / (2.0 * sigma * sigma))
+                    )
                 )
-            )
     else:
 
         sigma_x = sigma_x

pysixtrack/be_beamfields/spacecharge.py

@@ -6,16 +6,16 @@ class SpaceChargeCoasting(Element):
     """Space charge for a coasting beam"""
 
     _description = [
-        ("line_density", "1/m", "Desc...", 0.0),
-        ("sigma_x", "m", "Desc...", 1.0),
-        ("sigma_y", "m", "Desc...", 1.0),
-        ("length", "m", "Desc...", 0.0),
-        ("x_co", "m", "Desc...", 0.0),
-        ("y_co", "m", "Desc...", 0),
+        ("line_density", "1/m", "Number of particles per unit length", 0.0),
+        ("sigma_x", "m", "Horizontal size of the beam (r.m.s.)", 1.0),
+        ("sigma_y", "m", "Vertical size of the beam (r.m.s.)", 1.0),
+        ("length", "m", "Integration length of space charge kick", 0.0),
+        ("x_co", "m", "Horizontal closed orbit offset", 0.0),
+        ("y_co", "m", "Vertical closed orbit offset", 0),
     ]
     _extra = [
-        ("min_sigma_diff", "m", "Desc...", 1e-30),
-        ("enabled", "", "Desc...", True),
+        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-30),
+        ("enabled", "", "Switch to disable space charge effect", True),
     ]
 
     def track(self, p):
@@ -64,17 +64,17 @@ def track(self, p):
 
 class SpaceChargeBunched(Element):
     _description = [
-        ("number_of_particles", "", "Desc...", 0.0),
-        ("bunchlength_rms", "m", "Desc...", 1.0),
-        ("sigma_x", "m", "Desc...", 1.0),
-        ("sigma_y", "m", "Desc...", 1.0),
-        ("length", "m", "Desc...", 0.0),
-        ("x_co", "m", "Desc...", 0.0),
-        ("y_co", "m", "Desc...", 0),
+        ("number_of_particles", "", "Number of particles in the bunch", 0.0),
+        ("bunchlength_rms", "m", "Length of the bunch (r.m.s.)", 1.0),
+        ("sigma_x", "m", "Horizontal size of the beam (r.m.s.)", 1.0),
+        ("sigma_y", "m", "Vertical size of the beam (r.m.s.)", 1.0),
+        ("length", "m", "Integration length of space charge kick", 0.0),
+        ("x_co", "m", "Horizontal closed orbit offset", 0.0),
+        ("y_co", "m", "Vertical closed orbit offset", 0),
     ]
     _extra = [
-        ("min_sigma_diff", "m", "Desc...", 1e-30),
-        ("enabled", "", "Desc...", True),
+        ("min_sigma_diff", "m", "Threshold to detect round beam", 1e-30),
+        ("enabled", "", "Switch to disable space charge effect", True),
     ]
 
     def track(self, p):

pysixtrack/elements.py

@@ -309,7 +309,7 @@ def track(self, particle):
                 and y <= self.max_y
             )
             if particle.state != 1:
-                return particle.state
+                return "Particle lost"
         else:
             particle.state = np.int_(
                 (x >= self.min_x)
@@ -319,13 +319,13 @@ def track(self, particle):
             )
             particle.remove_lost_particles()
             if len(particle.state) == 0:
-                return -1
+                return "All particles lost"
 
 
 class LimitEllipse(Element):
     _description = [
-        ("a", "m^2", "Horizontal semiaxis", 1.0),
-        ("b", "m^2", "Vertical semiaxis", 1.0),
+        ("a", "m", "Horizontal semiaxis", 1.0),
+        ("b", "m", "Vertical semiaxis", 1.0),
     ]
 
     def track(self, particle):
@@ -345,7 +345,7 @@ def track(self, particle):
             )
             particle.remove_lost_particles()
             if len(particle.state) == 0:
-                return "All particle lost"
+                return "All particles lost"
 
 
 class BeamMonitor(Element):

tests/test_beamfields.py

@@ -0,0 +1,69 @@
+import numpy as np
+
+import pysixtrack
+
+
+def test_track_spacecharge():
+    x_co = 0.1
+    y_co = -0.5
+    sigma_x = 0.5
+    sigma_y = 0.1
+    el1 = pysixtrack.elements.SpaceChargeBunched(
+        number_of_particles=1e11,
+        bunchlength_rms=0.22,
+        sigma_x=sigma_x,
+        sigma_y=sigma_y,
+        length=2.0,
+        x_co=x_co,
+        y_co=y_co,
+    )
+    line_density = el1.number_of_particles / (
+        el1.bunchlength_rms * np.sqrt(2 * np.pi)
+    )
+    el2 = pysixtrack.elements.SpaceChargeCoasting(
+        line_density=line_density,
+        sigma_x=el1.sigma_x,
+        sigma_y=el1.sigma_y,
+        length=el1.length,
+        x_co=el1.x_co,
+        y_co=el1.y_co,
+    )
+
+    # test absolute kick for sigma_x > sigma_y
+    x_offset = 0.2
+    y_offset = -0.5
+    p1 = pysixtrack.Particles()
+    p1.x = x_co + x_offset
+    p1.y = y_co + y_offset
+    p2 = p1.copy()
+    el1.track(p1)
+    el2.track(p2)
+    assert np.isclose(p1.px, 1.33671170365656e-07, atol=1e-15)
+    assert np.isclose(p1.py, -6.228154616877465e-07, atol=1e-15)
+    assert p1.compare(p2, abs_tol=1e-15)
+
+    # test absolute kick for sigma_y > sigma_x
+    el1.sigma_x = sigma_y
+    el1.sigma_y = sigma_x
+    el2.sigma_x = sigma_y
+    el2.sigma_y = sigma_x
+    p1 = pysixtrack.Particles()
+    p1.x = x_co + y_offset
+    p1.y = y_co + x_offset
+    p2 = p1.copy()
+    el1.track(p1)
+    el2.track(p2)
+    assert np.isclose(p1.px, -6.228154616877465e-07, atol=1e-15)
+    assert np.isclose(p1.py, 1.33671170365656e-07, atol=1e-15)
+    assert p1.compare(p2, abs_tol=1e-15)
+
+    # test no kick for particle on closed orbit
+    p1 = pysixtrack.Particles()
+    p1.x = el1.x_co
+    p1.y = el1.y_co
+    p2 = p1.copy()
+    el1.track(p1)
+    el2.track(p2)
+    assert np.isclose(p1.px, 0.0, atol=1e-15)
+    assert np.isclose(p1.py, 0.0, atol=1e-15)
+    assert p1.compare(p2, abs_tol=1e-15)

tests/test_track.py

@@ -1,10 +1,13 @@
+import numpy as np
+
 import pysixtrack
 
 element_list = [
     pysixtrack.elements.Drift,
     pysixtrack.elements.DriftExact,
     pysixtrack.elements.Multipole,
     pysixtrack.elements.Cavity,
+    pysixtrack.elements.SawtoothCavity,
     pysixtrack.elements.XYShift,
     pysixtrack.elements.SRotation,
     pysixtrack.elements.RFMultipole,
@@ -13,6 +16,10 @@
     pysixtrack.elements.Line,
     pysixtrack.elements.LimitRect,
     pysixtrack.elements.LimitEllipse,
+    pysixtrack.elements.BeamBeam4D,
+    pysixtrack.elements.BeamBeam6D,
+    pysixtrack.elements.SpaceChargeCoasting,
+    pysixtrack.elements.SpaceChargeBunched,
 ]
 
 
@@ -35,3 +42,51 @@ def test_track_rfmultipole():
     el2.track(p2)
 
     assert p1.compare(p2, abs_tol=1e-15)
+
+
+def test_track_LimitRect():
+    min_x = -0.1
+    max_x = 0.3
+    min_y = -0.5
+    max_y = 0.1
+    el = pysixtrack.elements.LimitRect(
+        min_x=min_x, max_x=max_x, min_y=min_y, max_y=max_y
+    )
+
+    p1 = pysixtrack.Particles()
+    p1.x = 1
+    p1.y = 1
+    ret = el.track(p1)
+    assert ret == "Particle lost"
+
+    arr = np.arange(0, 1, 0.001)
+    p2 = pysixtrack.Particles(x=arr, y=arr)
+    ret = el.track(p2)
+
+    p2.x += max_x + 1e-6
+    ret = el.track(p2)
+    assert ret == "All particles lost"
+
+
+def test_track_LimitEllipse():
+    limit_a = 0.1
+    limit_b = 0.2
+    el = pysixtrack.elements.LimitEllipse(a=limit_a, b=limit_b)
+
+    p1 = pysixtrack.Particles()
+    p1.x = 1
+    p1.y = 1
+    ret = el.track(p1)
+    assert ret == "Particle lost"
+
+    arr = np.arange(0, 1, 0.001)
+    p2 = pysixtrack.Particles(x=arr, y=arr)
+    ret = el.track(p2)
+    survived = np.where(
+        (p2.x ** 2 / limit_a ** 2 + p2.y ** 2 / limit_b ** 2 <= 1.0)
+    )
+    assert len(p2.state) == len(survived[0])
+
+    p2.x += limit_a + 1e-6
+    ret = el.track(p2)
+    assert ret == "All particles lost"