Merge remote-tracking branch 'hannes-bartosik/master' into develop

Author: giadarol

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,
@@ -15,6 +18,8 @@
     pysixtrack.elements.LimitEllipse,
     pysixtrack.elements.BeamBeam4D,
     pysixtrack.elements.BeamBeam6D,
+    pysixtrack.elements.SpaceChargeCoasting,
+    pysixtrack.elements.SpaceChargeBunched,
 ]
 
 
@@ -37,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"