Beam get item method

.gitignore

@@ -15,5 +15,5 @@ build
 distributions
 
 docs/_build
-
+*.idea
 dev*

cheetah/particles/particle_beam.py

@@ -902,6 +902,32 @@ def momenta(self) -> torch.Tensor:
         """Momenta of the individual particles."""
         return torch.sqrt(self.energies**2 - electron_mass_eV**2)
 
+    def __getitem__(self, item):
+        # check if energy is batched
+        if len(self.energy.shape) == 0:
+            energy = self.energy
+        else:
+            # if energy is batched then make sure the batch dimension matches the
+            # batch dimension of the particle distribution
+            if not self.energy.shape == self.particles.shape[:-2]:
+                raise RuntimeError(
+                    "particle energy batch shape does not match " "particle batch shape"
+                )
+
+            energy = self.energy[item]
+
+        # check if particle charges is batched
+        if len(self.particle_charges.shape) > 1:
+            particle_charges = self.particle_charges[item]
+        else:
+            particle_charges = self.particle_charges
+
+        return ParticleBeam(
+            self.particles[item],
+            energy,
+            particle_charges,
+        )
+
     def __repr__(self) -> str:
         return (
             f"{self.__class__.__name__}(n={repr(self.num_particles)},"

tests/test_particle_beam.py

@@ -1,6 +1,8 @@
 import numpy as np
+import pytest
 import torch
 
+import cheetah
 from cheetah import ParticleBeam
 
 
@@ -144,3 +146,33 @@ def test_generate_uniform_ellipsoid_vectorized():
     assert torch.allclose(beam.sigma_p, sigma_p)
     assert torch.allclose(beam.energy, energy)
     assert torch.allclose(beam.total_charge, total_charge)
+
+
+def test_indexing():
+    # test batching with beamline parameters
+    quadrupole = cheetah.Quadrupole(
+        length=torch.tensor(0.2).unsqueeze(0), k1=torch.rand((5, 2))
+    )
+    incoming = cheetah.ParticleBeam.from_parameters(sigma_x=torch.tensor(1e-5))
+
+    outgoing = quadrupole.track(incoming)
+
+    sub_beam = outgoing[:2]
+    assert sub_beam.beta_x.shape == torch.Size([2, 2])
+    assert torch.equal(sub_beam.particle_charges, incoming.particle_charges)
+    assert torch.equal(sub_beam.energy, incoming.energy)
+
+    # test batching with energy
+    incoming = cheetah.ParticleBeam.from_parameters(sigma_x=torch.tensor(1e-5))
+    incoming.energy.data = torch.rand((5, 2))
+
+    outgoing = quadrupole.track(incoming)
+    sub_beam = outgoing[:2]
+
+    assert sub_beam.beta_x.shape == torch.Size([2, 2])
+    assert torch.equal(sub_beam.particle_charges, incoming.particle_charges)
+    assert torch.equal(sub_beam.energy, incoming.energy[:2])
+
+    with pytest.raises(RuntimeError):
+        outgoing.energy.data = torch.rand((4, 2))
+        outgoing[:2]