Fix dipole tracking issue; Add tests accordingly

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) (@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) (@jank324, @cr-xu)
 
 ### 🚀 Features
 

cheetah/accelerator/dipole.py

@@ -120,7 +120,7 @@ def transfer_map(self, energy: torch.Tensor) -> torch.Tensor:
         R_enter = self._transfer_map_enter()
         R_exit = self._transfer_map_exit()
 
-        if self.length != 0.0:  # Bending magnet with finite length
+        if any(self.length != 0.0):  # Bending magnet with finite length
             R = base_rmatrix(
                 length=self.length,
                 k1=torch.zeros_like(self.length),

tests/test_dipole.py

@@ -0,0 +1,50 @@
+import torch
+
+from cheetah import Dipole, Drift, ParameterBeam, ParticleBeam, Segment
+
+
+def test_dipole_off():
+    """
+    Test that a dipole with angle=0 behaves still like a drift.
+    """
+    dipole = Dipole(length=torch.tensor([1.0]), angle=torch.tensor([0.0]))
+    drift = Drift(length=torch.tensor([1.0]))
+    incoming_beam = ParameterBeam.from_parameters(
+        sigma_xp=torch.tensor([2e-7]), sigma_yp=torch.tensor([2e-7])
+    )
+    outbeam_dipole_off = dipole(incoming_beam)
+    outbeam_drift = drift(incoming_beam)
+
+    dipole.angle = torch.tensor([1.0], device=dipole.angle.device)
+    outbeam_dipole_on = dipole(incoming_beam)
+
+    assert torch.allclose(outbeam_dipole_off.sigma_x, outbeam_drift.sigma_x)
+    assert not torch.allclose(outbeam_dipole_on.sigma_x, outbeam_drift.sigma_x)
+
+
+def test_dipole_batched_execution():
+    """
+    Test that a dipole with batch dimensions behaves as expected.
+    """
+    batch_shape = torch.Size([3])
+    incoming = ParticleBeam.from_parameters(
+        num_particles=torch.tensor(1000000),
+        energy=torch.tensor([1e9]),
+        mu_x=torch.tensor([1e-5]),
+    ).broadcast(batch_shape)
+    segment = Segment(
+        [
+            Dipole(
+                length=torch.tensor([0.5, 0.5, 0.5]),
+                angle=torch.tensor([0.1, 0.2, 0.1]),
+            ),
+            Drift(length=torch.tensor([0.5])).broadcast(batch_shape),
+        ]
+    )
+    outgoing = segment(incoming)
+
+    # Check that dipole with same bend angle produce same output
+    assert torch.allclose(outgoing.particles[0], outgoing.particles[2])
+
+    # Check different angles do make a difference
+    assert not torch.allclose(outgoing.particles[0], outgoing.particles[1])