Fix RBend broadcasting issue

CHANGELOG.md

@@ -17,7 +17,7 @@
 - `Screen` now offers the option to use KDE for differentiable images (see #200) (@cr-xu, @roussel-ryan)
 - Moving `Element`s and `Beam`s to a different `device` and changing their `dtype` like with any `torch.nn.Module` is now possible (see #209) (@jank324)
 - `Quadrupole` now supports tracking with Cheetah's matrix-based method or with Bmad's more accurate method (see #153) (@jp-ga, @jank324)
-- `Dipole` now takes a focusing moment `k1` (see #235) (@hespe)
+- `Dipole` now takes a focusing moment `k1` (see #235, #247) (@hespe)
 
 ### 🐛 Bug fixes
 

cheetah/accelerator/rbend.py

@@ -16,6 +16,7 @@ class RBend(Dipole):
 
     :param length: Length in meters.
     :param angle: Deflection angle in rad.
+    :param k1: Focussing strength in 1/m^-2.
     :param e1: The angle of inclination of the entrance face [rad].
     :param e2: The angle of inclination of the exit face [rad].
     :param tilt: Tilt of the magnet in x-y plane [rad].
@@ -30,6 +31,7 @@ def __init__(
         self,
         length: Optional[Union[torch.Tensor, nn.Parameter]],
         angle: Optional[Union[torch.Tensor, nn.Parameter]] = None,
+        k1: Optional[Union[torch.Tensor, nn.Parameter]] = None,
         e1: Optional[Union[torch.Tensor, nn.Parameter]] = None,
         e2: Optional[Union[torch.Tensor, nn.Parameter]] = None,
         tilt: Optional[Union[torch.Tensor, nn.Parameter]] = None,
@@ -43,20 +45,7 @@ def __init__(
         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,
-        )
-
-        super().__init__(
-            length=length,
-            angle=angle,
+            k1=k1,
             e1=e1,
             e2=e2,
             tilt=tilt,

tests/test_dipole.py

@@ -1,6 +1,15 @@
+import pytest
 import torch
 
-from cheetah import Dipole, Drift, ParameterBeam, ParticleBeam, Quadrupole, Segment
+from cheetah import (
+    Dipole,
+    Drift,
+    ParameterBeam,
+    ParticleBeam,
+    Quadrupole,
+    RBend,
+    Segment,
+)
 
 
 def test_dipole_off():
@@ -43,22 +52,23 @@ def test_dipole_focussing():
     assert not torch.allclose(outbeam_dipole_off.sigma_x, outbeam_quadrupole.sigma_x)
 
 
-def test_dipole_batched_execution():
+@pytest.mark.parametrize("DipoleType", [Dipole, RBend])
+def test_dipole_batched_execution(DipoleType):
     """
     Test that a dipole with batch dimensions behaves as expected.
     """
-    batch_shape = torch.Size([3])
+    batch_shape = torch.Size([6])
     incoming = ParticleBeam.from_parameters(
-        num_particles=torch.tensor(1000000),
+        num_particles=torch.tensor(1_000_000),
         energy=torch.tensor([1e9]),
         mu_x=torch.tensor([1e-5]),
     ).broadcast(batch_shape)
     segment = Segment(
         [
-            Dipole(
+            DipoleType(
                 length=torch.tensor([0.5, 0.5, 0.5]),
                 angle=torch.tensor([0.1, 0.2, 0.1]),
-            ),
+            ).broadcast((2,)),
             Drift(length=torch.tensor([0.5])).broadcast(batch_shape),
         ]
     )