new closed_orbit module: add **kwargs for particles...

... such that mass and charge can be changed to other
than proton default.

Author: aoeftiger

pysixtrack/closed_orbit.py

@@ -2,8 +2,10 @@
 from .particles import Particles
 
 
-def get_init_particles_for_linear_map(closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum):
-    part = Particles(p0c=p0c)
+def get_init_particles_for_linear_map(
+        closed_orbit, p0c, d, longitudinal_coordinate,
+        longitudinal_momentum, **kwargs):
+    part = Particles(p0c=p0c, **kwargs)
 #    part = Particles()
 #    part.p0c = p0c
     part.x  = closed_orbit[0] + np.array([0.0, 1.0 * d, 0.0, 0.0, 0.0, 0.0, 0.0])
@@ -32,20 +34,27 @@ def part_to_array(part, longitudinal_coordinate, longitudinal_momentum):
     return X_array
 
 
-def linearize_around_closed_orbit(line, closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum):
+def linearize_around_closed_orbit(
+        line, closed_orbit, p0c, d, longitudinal_coordinate,
+        longitudinal_momentum, **kwargs):
 
-    part = get_init_particles_for_linear_map(closed_orbit, p0c, d, longitudinal_coordinate, longitudinal_momentum)
-    X_init = part_to_array(part, longitudinal_coordinate, longitudinal_momentum)
+    part = get_init_particles_for_linear_map(
+        closed_orbit, p0c, d, longitudinal_coordinate,
+        longitudinal_momentum, **kwargs)
+    X_init = part_to_array(
+        part, longitudinal_coordinate, longitudinal_momentum)
 
     line.track(part)
-    X_fin = part_to_array(part, longitudinal_coordinate, longitudinal_momentum)
+    X_fin = part_to_array(
+        part, longitudinal_coordinate, longitudinal_momentum)
 
     m = X_fin[:, 0] - X_init[:, 0]
     M = np.empty([6, 6])
     for j in range(6):
         M[:, j] = (X_fin[:, j + 1] - X_fin[:, 0]) / d
 
-    X_CO = X_init[:, 0] + np.matmul(np.linalg.inv(np.identity(6) - M), m.T)
+    X_CO = X_init[:, 0] + np.matmul(
+        np.linalg.inv(np.identity(6) - M), m.T)
 
     return X_CO, M
 
@@ -71,7 +80,8 @@ def healy_symplectify(M):
     V = np.matmul(S, np.matmul(I - M, np.linalg.inv(I + M)))
     W = (V + V.T) / 2
     if np.linalg.det(I - np.matmul(S, W)) != 0:
-        M_new = np.matmul(I + np.matmul(S, W), np.linalg.inv(I - np.matmul(S, W)))
+        M_new = np.matmul(I + np.matmul(S, W),
+                          np.linalg.inv(I - np.matmul(S, W)))
     else:
         print("WARNING: det(I - SW) = 0!")
         V_else = np.matmul(S, np.matmul(I + M, np.linalg.inv(I - M)))