Merge pull request #2 from SixTrack/master

merge

.flake8

@@ -0,0 +1,5 @@
+[flake8]
+extend-ignore =
+    # See https://github.com/PyCQA/pycodestyle/issues/373
+    E203,
+

examples/aperture/000_test_aperture.py

@@ -3,7 +3,6 @@
 import matplotlib.pyplot as plt
 
 import pysixtrack
-from numpy.random import uniform
 
 p = pysixtrack.Particles()
 

examples/bbsimple/testbb6d.py

@@ -131,7 +131,8 @@ def compare(prun, pbench, pbench_prev):
         out.append(abs(diff))
         out_rel.append(diffrel)
         print(
-            f"{att:<5} {vrun:22.13e} {vbench:22.13e} {diff:22.13g} {diffrel:22.13g}"
+            f"{att:<5} {vrun:22.13e} {vbench:22.13e}"
+            "{diff:22.13g} {diffrel:22.13g}"
         )
         if diffrel > 1e-5 or np.isnan(diffrel):
             if np.abs(diff) > 9e-10:

examples/beambeam/testbb6d.py

@@ -131,7 +131,8 @@ def compare(prun, pbench, pbench_prev):
         out.append(abs(diff))
         out_rel.append(diffrel)
         print(
-            f"{att:<5} {vrun:22.13e} {vbench:22.13e} {diff:22.13g} {diffrel:22.13g}"
+            f"{att:<5} {vrun:22.13e} {vbench:22.13e}"
+            "{diff:22.13g} {diffrel:22.13g}"
         )
         if diffrel > 1e-8 or np.isnan(diffrel):
             if np.abs(diff) > 1e-11:

examples/petra4/track_p1.py

@@ -34,7 +34,9 @@ def __repr__(self):
 print(mad.table.summ.q1, mad.table.summ.q2)
 
 
-twissout = pysixtrack.Particles.from_mad_twiss(mad.twiss(betx=1, bety=1, x=0.001))
+twissout = pysixtrack.Particles.from_madx_twiss(
+    mad.twiss(betx=1, bety=1, x=0.001)
+)
 
 line = pysixtrack.Line.from_madx_sequence(mad.sequence.ring)
 part = pysixtrack.Particles()
@@ -55,7 +57,8 @@ def mkd(name, t1, t2, ii, jj):
     sm = twissout.s[ii]
     sp = pysixout.s[ii]
     print(
-        f"{ii:3} {twiss.name[ii]:20} {line.element_names[ii]:20} {line.elements[ii]}"
+        f"{ii:3} {twiss.name[ii]:20}"
+        "{line.element_names[ii]:20} {line.elements[ii]}"
     )
     res = mkd("s", twissout, pysixout, ii, ii)
     if abs(res) > 1e-6:

examples/rf_multipole/rf_multipole.py

@@ -2,13 +2,12 @@
 
 from cpymad.madx import Madx
 import pysixtrack
-from pysixtrack.particles import Particles
 
 # run MADX tests
 mad = Madx()
 mad.call("rf_multipole.madx")
 
-# same test in pysixtrack
+# create pysixtrack rfmultipole
 mad_sequence = mad.sequence["sequ_rfmultipole"]
 rf_mulitpole_mad = mad_sequence.elements[1]
 freq = rf_mulitpole_mad.freq * 1e6  # MAD units are MHz
@@ -20,12 +19,22 @@
     voltage=0, frequency=freq, lag=lag, knl=knl, ksl=[0], pn=pn, ps=[0]
 )
 
-mad_part=Particles.from_mad_track(mad)
-p1=mad_part.copy(0)
-p2=mad_part.copy(1)
-p3=p1.copy()
+
+# track pysixtrack
+mad_part = pysixtrack.Particles.from_madx_track(mad)
+p1 = mad_part.copy(0)
+p2 = mad_part.copy(1)
+p3 = p1.copy()
 rf_multipole.track(p3)
 
+# compare
 p2.compare(p3)
 
-
+# test conversion
+line = pysixtrack.Line.from_madx_sequence(mad_sequence)
+tw = mad.twiss(betx=1, bety=1, x=0.1, t=0.5)
+p_mad = pysixtrack.Particles.from_madx_twiss(tw)
+p_six = mad_part.copy(0)
+p_out = pysixtrack.Particles.from_list(
+    line.track_elem_by_elem(p_six, start=False)
+)

examples/rf_multipole/rf_multipole2.py

@@ -0,0 +1,10 @@
+import pysixtrack
+
+
+rfmult = {"freq": 100, "knl": [0.1]}
+out = pysixtrack.mad_benchmark("rfmultipole", rfmult, x=1, t=0.1)
+mad, line, p_mad, p_six = out
+
+mult = {"knl": [0.1]}
+out = pysixtrack.mad_benchmark("multipole", mult, x=0, pt=0.1)
+mad, line, p_mad, p_six = out

examples/spacecharge/example_helpers.py

@@ -105,9 +105,9 @@ def track_particle_sixtrack(
     lines_f3[i_start_ini + 7 + 6] = "    0.\n"
 
     lines_f13 = []
-    
+
     prev_part = None
-    
+
     for i_part in range(n_part):
         temp_part = pysixtrack.Particles(**partCO)
         temp_part.x += Dx_wrt_CO_m[i_part]
@@ -328,7 +328,8 @@ def track_particle_sixtracklib(
     sigma_tbt = res.particles[0].sigma.reshape(n_turns, n_part)
     delta_tbt = res.particles[0].delta.reshape(n_turns, n_part)
 
-    # For now data are saved at the end of the turn by STlib and at the beginning by the others
+    # For now data are saved at the end of the turn by STlib
+    # and at the beginning by the others
     # x_tbt[1:, :] = x_tbt[:-1, :]
     # px_tbt[1:, :] = px_tbt[:-1, :]
     # y_tbt[1:, :] = y_tbt[:-1, :]

pysixtrack/__init__.py

@@ -8,4 +8,14 @@
 from .particles import Particles
 
 from .be_beamfields import BB6Ddata
-from .loader_mad import MadPoint
+from .loader_mad import MadPoint, mad_benchmark
+
+__all__ = [
+    "base_classes",
+    "elements",
+    "Line",
+    "Particles",
+    "BB6Ddata",
+    "MadPoint",
+    "mad_benchmark",
+]

pysixtrack/be_beamfields/BB6D.py

@@ -45,7 +45,8 @@ def BB6D_track(x, px, y, py, sigma, delta, q0, p0, bb6ddata, mathlib):
         delta_subCO,
         parboost,
     )
-    # ~ x_star, px_star, y_star, py_star, sigma_star, delta_star = (x, px, y, py, sigma, delta)
+    # ~ x_star, px_star, y_star, py_star, sigma_star, delta_star =
+    #      (x, px, y, py, sigma, delta)
     for i_slice in range(N_slices):
         sigma_slice_star = sigma_slices_star[i_slice]
         x_slice_star = x_slices_star[i_slice]
@@ -72,7 +73,8 @@ def BB6D_track(x, px, y, py, sigma, delta, q0, p0, bb6ddata, mathlib):
             Sigmas_0_star, S, threshold_singular=threshold_singular
         )
 
-        # Evaluate transverse coordinates of the weake baem w.r.t. the strong beam centroid
+        # Evaluate transverse coordinates of the weake baem w.r.t.
+        # the strong beam centroid
         x_bar_star = x_star + px_star * S - x_slice_star
         y_bar_star = y_star + py_star * S - y_slice_star
 
@@ -87,8 +89,10 @@ def BB6D_track(x, px, y, py, sigma, delta, q0, p0, bb6ddata, mathlib):
         )
 
         # Compute normalized field
-        # Ex, Ey, Gx, Gy = tef.get_Ex_Ey_Gx_Gy_gauss(x=x_bar_hat_star, y=y_bar_hat_star,
-        #                     sigma_x=np.sqrt(Sig_11_hat_star), sigma_y=np.sqrt(Sig_33_hat_star),
+        # Ex, Ey, Gx, Gy = tef.get_Ex_Ey_Gx_Gy_gauss(
+        #                     x=x_bar_hat_star, y=y_bar_hat_star,
+        #                     sigma_x=np.sqrt(Sig_11_hat_star),
+        #                     sigma_y=np.sqrt(Sig_33_hat_star),
         #                     min_sigma_diff = min_sigma_diff)
 
         Ex, Ey, Gx, Gy = gf.get_Ex_Ey_Gx_Gy_gauss(

pysixtrack/be_beamfields/BB6Ddata.py

@@ -251,7 +251,8 @@ def BB6D_init(
     z_slices = np.take(z_slices, ind_sorted)
     N_part_per_slice = np.take(N_part_per_slice, ind_sorted)
 
-    # By boosting the strong z and all zeros, I get the transverse coordinates of the strong beam in the ref system of the weak
+    # By boosting the strong z and all zeros, I get the transverse coordinates
+    # of the strong beam in the ref system of the weak
     boost_vect = np.vectorize(boost.boost, excluded="parboost")
     (
         x_slices_star,

pysixtrack/be_beamfields/beambeam.py

@@ -91,7 +91,8 @@ class BeamBeam6D(Element):
         (
             "phi",
             "rad",
-            "Crossing angle (>0 weak beam increases x in the direction motion)",
+            "Crossing angle (>0 weak beam increases"
+            "x in the direction motion)",
             0,
         ),
         ("alpha", "rad", "Crossing plane tilt angle (>0 x tends to y)", 0),
@@ -116,7 +117,8 @@ class BeamBeam6D(Element):
         (
             "zeta_slices",
             "m",
-            "Longitudinal position of the interacting slices (>0 head of the strong).",
+            "Longitudinal position of the interacting"
+            "slices (>0 head of the strong).",
             (),
         ),
         (
@@ -249,7 +251,14 @@ def track(self, p):
                 self.d_delta,
                 self.enabled,
             )
-            x_ret, px_ret, y_ret, py_ret, zeta_ret, delta_ret = BB6D.BB6D_track(
+            (
+                x_ret,
+                px_ret,
+                y_ret,
+                py_ret,
+                zeta_ret,
+                delta_ret,
+            ) = BB6D.BB6D_track(
                 p.x,
                 p.px,
                 p.y,

pysixtrack/be_beamfields/gaussian_fields.py

@@ -38,7 +38,8 @@ def _get_transv_field_gauss_ellip(
     sqrt_pi = sqrt(pi)
 
     # I always go to the first quadrant and then apply the signs a posteriori
-    # numerically more stable (see http://inspirehep.net/record/316705/files/slac-pub-5582.pdf)
+    # numerically more stable
+    # (see http://inspirehep.net/record/316705/files/slac-pub-5582.pdf)
 
     abx = abs(x - Delta_x)
     aby = abs(y - Delta_y)

pysixtrack/be_beamfields/slicing.py

@@ -30,7 +30,8 @@ def constant_charge_slicing_gaussian(N_part_tot, sigmaz, N_slices):
                 )
                 * float(N_slices)
             )
-            # the multiplication times n slices comes from the fact that we have to divide by the slice charge, i.e. 1./N
+            # the multiplication times n slices comes from the fact that
+            # we have to divide by the slice charge, i.e. 1./N
             z_centroids.append(this_centroid)
 
         last_centroid = (

pysixtrack/be_beamfields/tools.py

@@ -118,7 +118,12 @@ def find_alpha_and_phi(dpx, dpy):
 def get_bb_names_madpoints_sigmas(
     mad, seq_name, use_survey=True, use_twiss=True
 ):
-    _, element_names, points, twissdata = get_points_twissdata_for_element_type(
+    (
+        _,
+        element_names,
+        points,
+        twissdata,
+    ) = get_points_twissdata_for_element_type(
         mad,
         seq_name,
         ele_type="beambeam",