Format

Author: austin-hoover

examples/Diagnostics/Tunes/test_tune.py

@@ -4,6 +4,7 @@
 are estimated from the phase space coordinates before/after tracking using the
 `BunchTuneAnalysis` class.
 """
+
 import math
 import os
 import pathlib
@@ -64,10 +65,10 @@
 
 tune_node = TeapotTuneAnalysisNode()
 tune_node.assignTwiss(
-    betax=lattice_beta_x, 
-    alphax=lattice_alpha_x, 
-    etax=lattice_eta_x, 
-    etapx=lattice_etap_x, 
+    betax=lattice_beta_x,
+    alphax=lattice_alpha_x,
+    etax=lattice_eta_x,
+    etapx=lattice_etap_x,
     betay=lattice_beta_y,
     alphay=lattice_alpha_y,
 )
@@ -77,8 +78,8 @@
 # Bunch
 # ------------------------------------------------------------------------------------
 
-# Generate a matched transverse phase space distribution. The longitudinal 
-# distribution will be uniform in position (z) and a delta function in energy 
+# Generate a matched transverse phase space distribution. The longitudinal
+# distribution will be uniform in position (z) and a delta function in energy
 # deviation (dE).
 emittance_x = 25.0e-06
 emittance_y = 25.0e-06
@@ -113,7 +114,7 @@
 bunch.dumpBunch(filename)
 
 
-# Analysis    
+# Analysis
 # ------------------------------------------------------------------------------------
 
 # Collect phase data from bunch
@@ -124,21 +125,21 @@
 # Read phase data from file
 particles = np.loadtxt(filename, comments="%")
 particles = pd.DataFrame(
-    particles, 
+    particles,
     columns=[  # https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/78
-        "x", 
-        "xp", 
-        "y", 
-        "yp", 
+        "x",
+        "xp",
+        "y",
+        "yp",
         "z",
-        "dE",   
+        "dE",
         "phase_x",
         "phase_y",
         "tune_x",
         "tune_y",
         "action_x",
         "action_y",
-    ] 
+    ],
 )
 print(particles.iloc[:, 6:])
 
@@ -159,4 +160,4 @@
 print("tune_y_err", tune_y_err)
 
 assert np.abs(tune_x_err) < 1.00e-08
-assert np.abs(tune_y_err) < 1.00e-08
+assert np.abs(tune_y_err) < 1.00e-08

examples/Diagnostics/Tunes/test_tune_4d.py

@@ -1,4 +1,5 @@
 """Test one-turn tune estimation in coupled lattice."""
+
 import math
 import os
 import pathlib
@@ -59,7 +60,7 @@ def calc_eigtune(eigval: float) -> float:
 def unit_symplectic_matrix(ndim: int) -> np.ndarray:
     U = np.zeros((ndim, ndim))
     for i in range(0, ndim, 2):
-        U[i: i + 2, i: i + 2] = [[0.0, 1.0], [-1.0, 0.0]]
+        U[i : i + 2, i : i + 2] = [[0.0, 1.0], [-1.0, 0.0]]
     return U
 
 
@@ -68,15 +69,15 @@ def normalize_eigvec(v: np.ndarray) -> np.ndarray:
 
     def _norm(v):
         return np.linalg.multi_dot([np.conj(v), U, v])
-    
+
     if _norm(v) > 0.0:
         v = np.conj(v)
-    
+
     v *= np.sqrt(2.0 / np.abs(_norm(v)))
     assert np.isclose(np.imag(_norm(v)), -2.0)
     assert np.isclose(np.real(_norm(v)), +0.0)
     return v
- 
+
 
 def calc_norm_matrix_from_eigvecs(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
     V = np.zeros((4, 4))
@@ -118,7 +119,7 @@ def calc_norm_matrix_from_eigvecs(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
 
 # Add tune diagnostic node
 # ------------------------------------------------------------------------------------
-    
+
 tune_node = TeapotTuneAnalysisNode()
 tune_node.setNormMatrix(V_inv)
 lattice.getNodes()[0].addChildNode(tune_node, 0)
@@ -162,7 +163,7 @@ def calc_norm_matrix_from_eigvecs(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
     print("turn={} xrms={:0.3f} yrms={:0.3f}".format(turn + 1, xrms, yrms))
 
 
-# Analysis    
+# Analysis
 # ------------------------------------------------------------------------------------
 
 # Collect phase data from bunch
@@ -184,4 +185,4 @@ def calc_norm_matrix_from_eigvecs(v1: np.ndarray, v2: np.ndarray) -> np.ndarray:
 print("tune_2_err", tune_2_err)
 
 assert np.abs(tune_1_err) < 1.00e-08
-assert np.abs(tune_2_err) < 1.00e-08
+assert np.abs(tune_2_err) < 1.00e-08

examples/Diagnostics/Tunes/test_tune_4d_uncoupled.py

@@ -5,6 +5,7 @@
 in the lattice, the (eigen)tunes {nu1, nu2} should be the same as horizontal
 and vertical tunes {nux, nuy}.
 """
+
 import math
 import os
 import pathlib
@@ -50,6 +51,7 @@
 # Analyze transfer matrix
 # ------------------------------------------------------------------------------------
 
+
 def build_norm_matrix_from_twiss_2d(alpha: float, beta: float) -> np.ndarray:
     norm_matrix_inv = np.array([[beta, 0.0], [-alpha, 1.0]]) * np.sqrt(1.0 / beta)
     norm_matrix = np.linalg.inv(norm_matrix_inv)
@@ -75,7 +77,7 @@ def build_norm_matrix_from_twiss_2d(alpha: float, beta: float) -> np.ndarray:
 
 # Add tune diagnostic node
 # ------------------------------------------------------------------------------------
-    
+
 tune_node = TeapotTuneAnalysisNode()
 tune_node.setNormMatrix(norm_matrix)
 lattice.getNodes()[0].addChildNode(tune_node, 0)
@@ -116,7 +118,7 @@ def build_norm_matrix_from_twiss_2d(alpha: float, beta: float) -> np.ndarray:
 bunch.dumpBunch(filename)
 
 
-# Analysis    
+# Analysis
 # ------------------------------------------------------------------------------------
 
 # Collect phase data from bunch
@@ -127,21 +129,21 @@ def build_norm_matrix_from_twiss_2d(alpha: float, beta: float) -> np.ndarray:
 # Read phase data from file
 particles = np.loadtxt(filename, comments="%")
 particles = pd.DataFrame(
-    particles, 
+    particles,
     columns=[  # https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/78
-        "x", 
-        "xp", 
-        "y", 
-        "yp", 
+        "x",
+        "xp",
+        "y",
+        "yp",
         "z",
-        "dE",   
+        "dE",
         "phase_1",
         "phase_2",
         "tune_1",
         "tune_2",
         "action_1",
         "action_2",
-    ] 
+    ],
 )
 print(particles.iloc[:, 6:])
 
@@ -161,4 +163,4 @@ def build_norm_matrix_from_twiss_2d(alpha: float, beta: float) -> np.ndarray:
 print("tune_2_err", tune_2_err)
 
 assert np.abs(tune_1_err) < 1.00e-08
-assert np.abs(tune_2_err) < 1.00e-08
+assert np.abs(tune_2_err) < 1.00e-08

examples/Diagnostics/Tunes/utils.py

@@ -18,7 +18,7 @@ def make_lattice(
     start: str = "drift",
 ) -> AccLattice:
     """Create FODO lattice.
-    
+
     Args:
         length: Length of lattice [m].
         fill_frac: Fraction of lattice occupied by quadrupoles.
@@ -38,14 +38,14 @@ def make_lattice(
             QuadTEAPOT("qd"),
             QuadTEAPOT("qf2"),
         ]
-        
+
         drift_nodes[0].setLength(length_drift)
         drift_nodes[1].setLength(length_drift)
-        
+
         quad_nodes[0].setLength(length_quad * 0.5)
         quad_nodes[1].setLength(length_quad)
         quad_nodes[2].setLength(length_quad * 0.5)
-        
+
         quad_nodes[0].setParam("kq", +kq)
         quad_nodes[1].setParam("kq", -kq)
         quad_nodes[2].setParam("kq", +kq)
@@ -69,14 +69,14 @@ def make_lattice(
             QuadTEAPOT("qd"),
             QuadTEAPOT("qf2"),
         ]
-        
+
         drift_nodes[0].setLength(length_drift)
         drift_nodes[1].setLength(length_drift)
-        
+
         quad_nodes[0].setLength(length_quad * 0.5)
         quad_nodes[1].setLength(length_quad)
         quad_nodes[2].setLength(length_quad * 0.5)
-        
+
         quad_nodes[0].setParam("kq", +kq)
         quad_nodes[1].setParam("kq", -kq)
         quad_nodes[2].setParam("kq", +kq)
@@ -91,7 +91,7 @@ def make_lattice(
 
     else:
         raise ValueError
-    
+
     for node in lattice.getNodes():
         node.setUsageFringeFieldIN(False)
         node.setUsageFringeFieldOUT(False)