Fix vectorization issue in rbend (#173)

* Fix rbend issue w.r.t. batch input; Update tests

* Update changelog

* Remove old code comment

---------

Co-authored-by: Jan Kaiser <jan.kaiser@desy.de>
Co-authored-by: Jan Kaiser <jan.kaiser.email@googlemail.com>

Author: cr-xu

CHANGELOG.md

@@ -4,7 +4,7 @@
 
 ### 🚨 Breaking Changes
 
-- Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172) (@jank324, @cr-xu)
+- Cheetah is now vectorised. This means that you can run multiple simulations in parallel by passing a batch of beams and settings, resulting a number of interfaces being changed. For Cheetah developers this means that you now have to account for an arbitrary-dimensional tensor of most of the properties of you element, rather than a single value, vector or whatever else a property was before. (see #116, #157, #170, #172, #173) (@jank324, @cr-xu)
 
 ### 🚀 Features
 

cheetah/accelerator/dipole.py

@@ -85,7 +85,7 @@ def __init__(
             if fringe_integral_exit is None
             else torch.as_tensor(fringe_integral_exit, **factory_kwargs)
         )
-        # Rectangular bend
+        # Sector bend if not specified
         self.e1 = (
             torch.as_tensor(e1, **factory_kwargs)
             if e1 is not None

cheetah/accelerator/rbend.py

@@ -51,18 +51,19 @@ def __init__(
         device=None,
         dtype=torch.float32,
     ):
-        angle = angle if angle is not None else torch.tensor(0.0)
-        e1 = e1 if e1 is not None else torch.tensor(0.0)
-        e2 = e2 if e2 is not None else torch.tensor(0.0)
-        tilt = tilt if tilt is not None else torch.tensor(0.0)
-        fringe_integral = (
-            fringe_integral if fringe_integral is not None else torch.tensor(0.0)
+        super().__init__(
+            length=length,
+            angle=angle,
+            e1=e1,
+            e2=e2,
+            tilt=tilt,
+            fringe_integral=fringe_integral,
+            fringe_integral_exit=fringe_integral_exit,
+            gap=gap,
+            name=name,
+            device=device,
+            dtype=dtype,
         )
-        # fringe_integral_exit is left out on purpose
-        gap = gap if gap is not None else torch.tensor(0.0)
-
-        e1 = e1 + angle / 2
-        e2 = e2 + angle / 2
 
         super().__init__(
             length=length,
@@ -77,3 +78,7 @@ def __init__(
             device=device,
             dtype=dtype,
         )
+
+        # Rectangular bend
+        self.e1 = self.e1 + self.angle / 2
+        self.e2 = self.e2 + self.angle / 2

tests/test_compare_ocelot.py

@@ -525,7 +525,10 @@ def test_rbend():
         "tests/resources/ACHIP_EA1_2021.1351.001"
     )
     cheetah_dipole = cheetah.RBend(
-        length=torch.tensor([0.1]), angle=torch.tensor([0.2])
+        length=torch.tensor([0.1]),
+        angle=torch.tensor([0.2]),
+        fringe_integral=torch.tensor([0.1]),
+        gap=torch.tensor([0.2]),
     )
     cheetah_segment = cheetah.Segment(
         [
@@ -540,7 +543,7 @@ def test_rbend():
     incoming_p_array = ocelot.astraBeam2particleArray(
         "tests/resources/ACHIP_EA1_2021.1351.001", print_params=False
     )
-    ocelot_rbend = ocelot.RBend(l=0.1, angle=0.2)
+    ocelot_rbend = ocelot.RBend(l=0.1, angle=0.2, fint=0.1, gap=0.2)
     lattice = ocelot.MagneticLattice(
         [ocelot.Drift(l=0.1), ocelot_rbend, ocelot.Drift(l=0.1)]
     )