Fix indentation

desy-ml · Sep 11, 2023 · 5d1782c · 5d1782c
1 parent eddbb03
commit 5d1782c
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 42 deletions.
diff --git a/cheetah/accelerator.py b/cheetah/accelerator.py
@@ -915,14 +915,14 @@ def _track_beam(self, incoming: ParticleBeam) -> ParticleBeam:
                     + T555 * incoming.particles[:, 4] ** 2
                 )
 
-            if isinstance(incoming, ParameterBeam):
-                outgoing = ParameterBeam(outgoing_mu, outgoing_cov, outgoing_energy)
-                return outgoing
-            else:  # ParticleBeam
-                outgoing = ParticleBeam(
-                    outgoing_particles, outgoing_energy, device=incoming.device
-                )
-                return outgoing
+        if isinstance(incoming, ParameterBeam):
+            outgoing = ParameterBeam(outgoing_mu, outgoing_cov, outgoing_energy)
+            return outgoing
+        else:  # ParticleBeam
+            outgoing = ParticleBeam(
+                outgoing_particles, outgoing_energy, device=incoming.device
+            )
+            return outgoing
 
     def _cavity_rmatrix(self, energy: float) -> torch.Tensor:
         """Produces an R-matrix for a cavity when it is on, i.e. voltage > 0.0."""

diff --git a/test_bmad.ipynb b/test_bmad.ipynb