black

Author: austin-hoover

orbit_tools/bunch.py

@@ -54,9 +54,7 @@ def get_bunch_coords(bunch: Bunch, axis: tuple[int, ...] = None) -> np.ndarray:
     return X[:, axis]
 
 
-def set_bunch_coords(
-    bunch: Bunch, X: np.ndarray, axis: tuple[int, ...] = None
-) -> Bunch:
+def set_bunch_coords(bunch: Bunch, X: np.ndarray, axis: tuple[int, ...] = None) -> Bunch:
     if axis is None:
         axis = tuple(range(6))
 
@@ -115,9 +113,7 @@ def shift_bunch_centroid(bunch: Bunch, offset: np.ndarray) -> Bunch:
     return bunch
 
 
-def set_bunch_cov(
-    bunch: Bunch, covariance_matrix: np.ndarray, block_diag: bool = True
-) -> Bunch:
+def set_bunch_cov(bunch: Bunch, covariance_matrix: np.ndarray, block_diag: bool = True) -> Bunch:
     X_old = get_bunch_coords(bunch)
     S_old = np.cov(X_old.T)
 
@@ -136,9 +132,7 @@ def set_bunch_cov(
     return bunch
 
 
-def transform_bunch(
-    bunch: Bunch, transform: Callable, axis: tuple[int, ...] = None
-) -> Bunch:
+def transform_bunch(bunch: Bunch, transform: Callable, axis: tuple[int, ...] = None) -> Bunch:
     if axis is None:
         axis = tuple(range(6))
 
@@ -147,9 +141,7 @@ def transform_bunch(
     return set_bunch_coords(bunch, X)
 
 
-def transform_bunch_linear(
-    bunch: Bunch, matrix: np.ndarray, axis: tuple[int, ...] = None
-) -> Bunch:
+def transform_bunch_linear(bunch: Bunch, matrix: np.ndarray, axis: tuple[int, ...] = None) -> Bunch:
     return transform_bunch(bunch, lambda x: np.matmul(x, matrix.T), axis=axis)
 
 
@@ -170,9 +162,7 @@ def set_bunch_current(bunch: Bunch, current: float, frequency: float) -> Bunch:
 
     Assumes bunch charge is already set.
     """
-    intensity = current_to_intensity(
-        current=current, frequency=frequency, charge=bunch.charge()
-    )
+    intensity = current_to_intensity(current=current, frequency=frequency, charge=bunch.charge())
     bunch_size_global = bunch.getSizeGlobal()
     if bunch_size_global > 0:
         macro_size = intensity / bunch_size_global
@@ -315,9 +305,7 @@ def get_bunch_cov(
     return S
 
 
-def generate_bunch(
-    sample: Callable, size: int, bunch: Bunch = None, verbose: bool = True
-) -> Bunch:
+def generate_bunch(sample: Callable, size: int, bunch: Bunch = None, verbose: bool = True) -> Bunch:
     """Generate bunch from particle sampler..
 
     Parameters

orbit_tools/cov.py

@@ -99,9 +99,7 @@ def normalization_matrix_from_twiss_2d(
     return np.linalg.inv(V)
 
 
-def normalization_matrix_from_twiss(
-    twiss_params: list[tuple[float, float, float]]
-) -> np.ndarray:
+def normalization_matrix_from_twiss(twiss_params: list[tuple[float, float, float]]) -> np.ndarray:
     """2N x 2N block-diagonal normalization matrix from Twiss parameters.
 
     Parameters
@@ -117,9 +115,7 @@ def normalization_matrix_from_twiss(
     ndim = len(twiss_params) // 2
     V = np.zeros((ndim, ndim))
     for i in range(0, ndim, 2):
-        V[i : i + 2, i : i + 2] = normalization_matrix_from_twiss_2d(
-            *twiss_params[i : i + 2]
-        )
+        V[i : i + 2, i : i + 2] = normalization_matrix_from_twiss_2d(*twiss_params[i : i + 2])
     return np.linalg.inv(V)
 
 
@@ -175,9 +171,7 @@ def cov_to_corr(S: np.ndarray) -> np.ndarray:
     return np.linalg.multi_dot([Dinv, S, Dinv])
 
 
-def rms_ellipse_params(
-    S: np.ndarray, axis: tuple[int, ...] = None
-) -> tuple[float, ...]:
+def rms_ellipse_params(S: np.ndarray, axis: tuple[int, ...] = None) -> tuple[float, ...]:
     """Return projected rms ellipse dimensions and orientation.
 
     Parameters

orbit_tools/diag.py

@@ -16,17 +16,16 @@
 from orbit.lattice import AccNode
 
 
-def get_grid_points(coords: list[np.ndarray]) -> np.ndarray:        
+def get_grid_points(coords: list[np.ndarray]) -> np.ndarray:
     return np.vstack([C.ravel() for C in np.meshgrid(*coords, indexing="ij")]).T
 
 
 def make_grid(
-    shape: tuple[int, ...], 
-    limits: list[tuple[float, float]]
+    shape: tuple[int, ...], limits: list[tuple[float, float]]
 ) -> Union[Grid1D, Grid2D, Grid3D]:
-    
+
     ndim = len(shape)
-    
+
     grid = None
     if ndim == 1:
         grid = Grid1D(shape[0] + 1, limits[0][0], limits[0][1])
@@ -52,7 +51,7 @@ def make_grid(
         raise ValueError
 
     return grid
-    
+
 
 class Diagnostic:
     def __init__(self, output_dir: str = None, verbose: bool = True) -> None:
@@ -96,7 +95,9 @@ def __init__(
         self.ndim = len(axis)
 
         if self.ndim > 2:
-            raise NotImplementedError("BunchHistogram does not yet support 3D grids. See https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/46 and https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/47.")
+            raise NotImplementedError(
+                "BunchHistogram does not yet support 3D grids. See https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/46 and https://github.com/PyORBIT-Collaboration/PyORBIT3/issues/47."
+            )
 
         self.dims = ["x", "xp", "y", "yp", "z", "dE"]
         self.dims = [self.dims[i] for i in self.axis]
@@ -107,28 +108,28 @@ def __init__(
             np.linspace(self.limits[i][0], self.limits[i][1], self.shape[i] + 1)
             for i in range(self.ndim)
         ]
-        self.coords = [0.5 * (e[:-1] + e[1:]) for e in self.edges]        
+        self.coords = [0.5 * (e[:-1] + e[1:]) for e in self.edges]
         self.values = np.zeros(shape)
-        
+
         self.points = get_grid_points(self.coords)
         self.cell_volume = np.prod([e[1] - e[0] for e in self.edges])
 
         self.grid = make_grid(self.shape, self.limits)
         self.method = method
         self.transform = transform
         self.normalize = normalize
-    
+
     def reset(self) -> None:
         self.grid.setZero()
 
     def sync_mpi(self) -> None:
         self.grid.synchronizeMPI(self.mpi_comm)
 
-    def bin_bunch(self, bunch: Bunch) -> None:     
+    def bin_bunch(self, bunch: Bunch) -> None:
         macrosize = bunch.macroSize()
         if macrosize == 0:
             bunch.macroSize(1.0)
-            
+
         if self.method == "bilinear":
             self.grid.binBunchBilinear(bunch, *self.axis)
         else:
@@ -139,7 +140,7 @@ def bin_bunch(self, bunch: Bunch) -> None:
     def compute_histogram(self, bunch: Bunch) -> np.ndarray:
         self.bin_bunch(bunch)
         self.sync_mpi()
-        
+
         values = np.zeros(self.points.shape[0])
         if self.method == "bilinear":
             for i, point in enumerate(self.points):
@@ -158,7 +159,7 @@ def compute_histogram(self, bunch: Bunch) -> np.ndarray:
                 values /= values_sum
             values /= self.cell_volume
         return values
-    
+
     def track(self, params_dict: dict) -> None:
         bunch_copy = Bunch()
         bunch = params_dict["bunch"]
@@ -186,7 +187,7 @@ def get_filename(self) -> str:
 class BunchHistogram1D(BunchHistogram):
     def __init__(self, **kwargs) -> None:
         super().__init__(**kwargs)
-        
+
 
 class BunchHistogram2D(BunchHistogram):
     def __init__(self, **kwargs) -> None:
@@ -195,4 +196,4 @@ def __init__(self, **kwargs) -> None:
 
 class BunchHistogram3D(BunchHistogram):
     def __init__(self, **kwargs) -> None:
-        super().__init__(**kwargs)
+        super().__init__(**kwargs)

orbit_tools/hydra.py

@@ -12,9 +12,7 @@ def on_job_start(self, config: DictConfig, **kwargs: Any) -> None:
         repo = git.Repo(search_parent_directories=True)
         sha = repo.head.object.hexsha
 
-        output_dir = os.path.join(
-            config.hydra.runtime.output_dir, config.hydra.output_subdir
-        )
+        output_dir = os.path.join(config.hydra.runtime.output_dir, config.hydra.output_subdir)
         filename = os.path.join(output_dir, "git_sha.txt")
 
         with open(filename, "w") as file:

orbit_tools/linac/core.py

@@ -64,9 +64,7 @@ def unnormalize_beta_z(mass: float, kin_energy: float, beta_z: float) -> float:
     return beta_z
 
 
-def get_node_info(
-    name: str = None, position: float = None, lattice: AccLattice = None
-) -> dict:
+def get_node_info(name: str = None, position: float = None, lattice: AccLattice = None) -> dict:
     """Return node, node index, start and stop position from node name or center position.
 
     Returns dict:
@@ -228,9 +226,7 @@ def make_phase_aperture_node(
     return aperture_node
 
 
-def make_energy_aperture_node(
-    energy_min: float, energy_max: float
-) -> LinacEnergyApertureNode:
+def make_energy_aperture_node(energy_min: float, energy_max: float) -> LinacEnergyApertureNode:
     aperture_node = LinacEnergyApertureNode()
     aperture_node.setMinMaxEnergy(energy_min, energy_max)
     return aperture_node
@@ -256,9 +252,7 @@ def check_sync_time(
     sync_time_design = 0.0
 
     if start_node_info["index"] > 0:
-        design_bunch = lattice.trackDesignBunch(
-            bunch, index_start=0, index_stop=start["index"]
-        )
+        design_bunch = lattice.trackDesignBunch(bunch, index_start=0, index_stop=start["index"])
         sync_time_design = design_bunch.getSyncParticle().time()
 
     if _mpi_rank == 0 and verbose:
@@ -271,11 +265,7 @@ def check_sync_time(
             print("    Setting to design value.")
         bunch.getSyncParticle().time(sync_time_design)
         if _mpi_rank == 0 and verbose:
-            print(
-                "bunch.getSyncParticle().time() = {}".format(
-                    bunch.getSyncParticle().time()
-                )
-            )
+            print("bunch.getSyncParticle().time() = {}".format(bunch.getSyncParticle().time()))
 
 
 def estimate_transfer_matrix(
@@ -309,9 +299,7 @@ def estimate_transfer_matrix(
             (-0.000, 0.000),
         ]
     test_bunch_lb, test_bunch_ub = list(zip(*test_bunch_limits))
-    test_bunch_coords = rng.uniform(
-        test_bunch_lb, test_bunch_ub, size=(test_bunch_size, 6)
-    )
+    test_bunch_coords = rng.uniform(test_bunch_lb, test_bunch_ub, size=(test_bunch_size, 6))
 
     test_bunch = Bunch()
     bunch.copyEmptyBunchTo(test_bunch)
@@ -327,21 +315,15 @@ def estimate_transfer_matrix(
     return matrix
 
 
-def save_node_positions(
-    lattice: AccLattice, filename: str = "lattice_nodes.txt"
-) -> None:
+def save_node_positions(lattice: AccLattice, filename: str = "lattice_nodes.txt") -> None:
     file = open(filename, "w")
     file.write("node position length\n")
     for node in lattice.getNodes():
-        file.write(
-            "{} {} {}\n".format(node.getName(), node.getPosition(), node.getLength())
-        )
+        file.write("{} {} {}\n".format(node.getName(), node.getPosition(), node.getLength()))
     file.close()
 
 
-def save_lattice_structure(
-    lattice: AccLattice, filename: str = "lattice_structure.txt"
-) -> None:
+def save_lattice_structure(lattice: AccLattice, filename: str = "lattice_structure.txt") -> None:
     file = open(filename, "w")
     file.write(lattice.structureToText())
     file.close()
@@ -583,9 +565,7 @@ def __call__(self, params_dict: dict, force_update: bool = False) -> None:
         # Measure mean and covariance.
         twiss_analysis = BunchTwissAnalysis()
         order = 2
-        twiss_analysis.computeBunchMoments(
-            bunch, order, self.dispersion_flag, self.emit_norm_flag
-        )
+        twiss_analysis.computeBunchMoments(bunch, order, self.dispersion_flag, self.emit_norm_flag)
 
         mean = np.zeros(6)
         for i in range(6):
@@ -650,9 +630,7 @@ def __call__(self, params_dict: dict, force_update: bool = False) -> None:
 
         # Measure maximum phase space coordinates.
         extrema_calculator = BunchExtremaCalculator()
-        (x_min, x_max, y_min, y_max, z_min, z_max) = extrema_calculator.extremaXYZ(
-            bunch
-        )
+        (x_min, x_max, y_min, y_max, z_min, z_max) = extrema_calculator.extremaXYZ(bunch)
         if _mpi_rank == 0:
             self.history["x_min"] = x_min
             self.history["x_max"] = x_max
@@ -691,17 +669,13 @@ def __call__(self, params_dict: dict, force_update: bool = False) -> None:
 
         # Write phase space coordinates to file.
         if self.write is not None:
-            if force_update or (
-                self.stride_write <= (position - self.last_write_position)
-            ):
+            if force_update or (self.stride_write <= (position - self.last_write_position)):
                 self.write(bunch, tag=node.getName())
                 self.last_write_position = position
 
         # Call plotting routines.
         if self.plot is not None and _mpi_rank == 0:
-            if force_update or (
-                self.stride_plot <= (position - self.last_plot_position)
-            ):
+            if force_update or (self.stride_plot <= (position - self.last_plot_position)):
                 info = dict()
                 for key in self.history:
                     if self.history[key]:
@@ -749,9 +723,7 @@ def action(self, params_dict: dict) -> None:
         order = 2
         dispersion_flag = 0
         emitt_norm_flag = 0
-        twiss_analysis.computeBunchMoments(
-            bunch, order, dispersion_flag, emitt_norm_flag
-        )
+        twiss_analysis.computeBunchMoments(bunch, order, dispersion_flag, emitt_norm_flag)
         x_rms = np.sqrt(twiss_analysis.getCorrelation(0, 0))
         y_rms = np.sqrt(twiss_analysis.getCorrelation(2, 2))
 
@@ -817,9 +789,7 @@ def forward(self, x: np.ndarray) -> np.ndarray:
 
         bunch = self.get_new_bunch()
         bunch = set_bunch_coords(bunch, x_new, verbose=False)
-        self.lattice.trackBunch(
-            bunch, index_start=self.index_start, index_stop=self.index_stop
-        )
+        self.lattice.trackBunch(bunch, index_start=self.index_start, index_stop=self.index_stop)
         U = bunch.get_bunch_coords(bunch)
         U = U[:, self.axis]
         return U
@@ -835,9 +805,7 @@ def inverse(self, u: np.ndarray) -> np.ndarray:
         bunch = oset_bunch_coords(bunch, u_new, verbose=False)
         bunch = reverse_bunch(bunch)
         self.lattice.reverseOrder()
-        self.lattice.trackBunch(
-            bunch, index_start=self.index_stop, index_stop=self.index_start
-        )
+        self.lattice.trackBunch(bunch, index_start=self.index_stop, index_stop=self.index_start)
         self.lattice.reverseOrder()
         bunch = reverse_bunch(bunch)
         x = get_bunch_coords(bunch)

orbit_tools/linac/diag.py

@@ -122,9 +122,7 @@ def __init__(self, rf_frequency: float = 402.5e06, **kwargs) -> None:
                 self.history[key] = None
 
             if self.output_dir is not None:
-                self.history_file = open(
-                    os.path.join(self.output_dir, "history.dat"), "w"
-                )
+                self.history_file = open(os.path.join(self.output_dir, "history.dat"), "w")
                 line = ",".join(list(self.history))
                 line = line[:-1] + "\n"
                 self.history_file.write(line)
@@ -178,9 +176,7 @@ def measure_stats(self, params_dict: dict) -> None:
 
     def measure_extrema(self) -> None:
         extrema_calculator = BunchExtremaCalculator()
-        (x_min, x_max, y_min, y_max, z_min, z_max) = extrema_calculator.extremaXYZ(
-            self.bunch
-        )
+        (x_min, x_max, y_min, y_max, z_min, z_max) = extrema_calculator.extremaXYZ(self.bunch)
         if self._mpi_rank == 0:
             self.history["x_min"] = x_min
             self.history["x_max"] = x_max