Fix "break" in voltage scan

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

Hide file tree

Showing 2 changed files with 895 additions and 49 deletions.
diff --git a/cheetah/accelerator.py b/cheetah/accelerator.py
@@ -939,10 +939,11 @@ def _cavity_rmatrix(self, energy: float) -> torch.Tensor:
 
         r11 = torch.cos(alpha) - torch.sqrt(2 / eta) * torch.cos(phi) * torch.sin(alpha)
 
-        if torch.abs(Ep) > 10:
-            r12 = torch.sqrt(8 / eta) * Ei / Ep * torch.cos(phi) * torch.sin(alpha)
-        else:
-            r12 = self.length
+        # In Ocelot r12 is defined as below only if abs(Ep) > 10, and self.length
+        # otherwise. This is implemented differently here in order to achieve results
+        # closer to Bmad.
+        r12 = torch.sqrt(8 / eta) * Ei / Ep * torch.cos(phi) * torch.sin(alpha)
+
         r21 = (
             -Ep
             / Ef
@@ -989,6 +990,8 @@ def _cavity_rmatrix(self, energy: float) -> torch.Tensor:
         r66 = Ei / Ef * beta0 / beta1
         r65 = k * torch.sin(phi) * self.voltage / (Ef * beta1 * electron_mass_eV)
 
+        print(f"{r12 = }")
+
         R = torch.tensor(
             [
                 [r11, r12, 0, 0, 0, 0, 0],