Merge pull request #15 from desy-ml/14-confusing-variable-naming-in-m…

…agnet-misalignment Fix variable naming for misalignment treatment
desy-ml · Jan 16, 2023 · a19f262 · a19f262
2 parents 1e6a4e4 + 7a91a06
commit a19f262
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/cheetah/accelerator.py b/cheetah/accelerator.py
@@ -267,7 +267,7 @@ def transfer_map(self, energy):
         if self.misalignment[0] == 0 and self.misalignment[1] == 0:
             return R
         else:
-            R_entry = torch.tensor(
+            R_exit = torch.tensor(
                 [
                     [1, 0, 0, 0, 0, 0, self.misalignment[0]],
                     [0, 1, 0, 0, 0, 0, 0],
@@ -280,7 +280,7 @@ def transfer_map(self, energy):
                 dtype=torch.float32,
                 device=self.device,
             )
-            R_exit = torch.tensor(
+            R_entry = torch.tensor(
                 [
                     [1, 0, 0, 0, 0, 0, -self.misalignment[0]],
                     [0, 1, 0, 0, 0, 0, 0],
@@ -293,8 +293,7 @@ def transfer_map(self, energy):
                 dtype=torch.float32,
                 device=self.device,
             )
-            R = torch.matmul(R_entry, R)
-            R = torch.matmul(R, R_exit)
+            R = torch.matmul(R_exit, torch.matmul(R, R_entry))
             return R
 
     @property