Add particle generation for 3D uniform ellispoid distribution

CHANGELOG.md

@@ -9,6 +9,7 @@
 ### 🚀 Features
 
 - `CustomTransferMap` elements created by combining multiple other elements will now reflect that in their `name` attribute (see #100) (@jank324)
+- Add a new class method for `ParticleBeam` to generate a 3D uniformly distributed ellipsoidal beam (see #146) (@cr-xu, @jank324)
 
 ### 🐛 Bug fixes
 

cheetah/particles.py

@@ -1019,6 +1019,129 @@ def from_twiss(
             dtype=dtype,
         )
 
+    @classmethod
+    def uniform_3d_ellispoid(
+        cls,
+        num_particles: Optional[torch.Tensor] = None,
+        radius_x: Optional[torch.Tensor] = None,
+        radius_y: Optional[torch.Tensor] = None,
+        radius_s: Optional[torch.Tensor] = None,
+        sigma_xp: Optional[torch.Tensor] = None,
+        sigma_yp: Optional[torch.Tensor] = None,
+        sigma_p: Optional[torch.Tensor] = None,
+        energy: Optional[torch.Tensor] = None,
+        total_charge: Optional[torch.Tensor] = None,
+        device=None,
+        dtype=torch.float32,
+    ):
+        """
+        Generate a particle beam with spatially uniformly distributed particles inside
+        an ellipsoid, i.e. a waterbag distribution.
+
+        Note that:
+         - The generated particles do not have correlation in the momentum directions,
+           and by default a cold beam with no divergence is generated.
+         - For batched generation, parameters that are not `None` must have the same
+           shape.
+
+        :param num_particles: Number of particles to generate.
+        :param radius_x: Radius of the ellipsoid in x direction in meters.
+        :param radius_y: Radius of the ellipsoid in y direction in meters.
+        :param radius_s: Radius of the ellipsoid in s (longitudinal) direction
+        in meters.
+        :param sigma_xp: Sigma of the particle distribution in x' direction in rad,
+        default is 0.
+        :param sigma_yp: Sigma of the particle distribution in y' direction in rad,
+        default is 0.
+        :param sigma_p: Sigma of the particle distribution in p, dimensionless.
+        :param energy: Energy of the beam in eV.
+        :param total_charge: Total charge of the beam in C.
+        :param device: Device to move the beam's particle array to.
+        :param dtype: Data type of the generated particles.
+
+        :return: ParticleBeam with uniformly distributed particles inside an ellipsoid.
+        """
+
+        # Figure out if arguments were passed, figure out their shape
+        not_nones = [
+            argument
+            for argument in [
+                radius_x,
+                radius_y,
+                radius_s,
+                sigma_xp,
+                sigma_yp,
+                sigma_p,
+                energy,
+                total_charge,
+            ]
+            if argument is not None
+        ]
+        shape = not_nones[0].shape if len(not_nones) > 0 else torch.Size([1])
+        if len(not_nones) > 1:
+            assert all(
+                argument.shape == shape for argument in not_nones
+            ), "Arguments must have the same shape."
+
+        # Set default values without function call in function signature
+        # NOTE that this does not need to be done for values that are passed to the
+        # Gaussian beam generation.
+        num_particles = (
+            num_particles if num_particles is not None else torch.tensor(1_000_000)
+        )
+        radius_x = radius_x if radius_x is not None else torch.full(shape, 1e-3)
+        radius_y = radius_y if radius_y is not None else torch.full(shape, 1e-3)
+        radius_s = radius_s if radius_s is not None else torch.full(shape, 1e-3)
+
+        # Generate xs, ys and ss within the ellipsoid
+        flattened_xs = torch.empty(*shape, num_particles).flatten(end_dim=-2)
+        flattened_ys = torch.empty(*shape, num_particles).flatten(end_dim=-2)
+        flattened_ss = torch.empty(*shape, num_particles).flatten(end_dim=-2)
+        for i, (r_x, r_y, r_s) in enumerate(
+            zip(radius_x.flatten(), radius_y.flatten(), radius_s.flatten())
+        ):
+            num_successful = 0
+            while num_successful < num_particles:
+                xs = (torch.rand(num_particles) - 0.5) * 2 * r_x
+                ys = (torch.rand(num_particles) - 0.5) * 2 * r_y
+                ss = (torch.rand(num_particles) - 0.5) * 2 * r_s
+
+                is_in_ellipsoid = xs**2 / r_x**2 + ys**2 / r_y**2 + ss**2 / r_s**2 < 1
+                num_to_add = min(num_particles - num_successful, is_in_ellipsoid.sum())
+
+                flattened_xs[i, num_successful : num_successful + num_to_add] = xs[
+                    is_in_ellipsoid
+                ][:num_to_add]
+                flattened_ys[i, num_successful : num_successful + num_to_add] = ys[
+                    is_in_ellipsoid
+                ][:num_to_add]
+                flattened_ss[i, num_successful : num_successful + num_to_add] = ss[
+                    is_in_ellipsoid
+                ][:num_to_add]
+
+                num_successful += num_to_add
+
+        # Generate an uncorrelated Gaussian beam
+        beam = cls.from_parameters(
+            num_particles=num_particles,
+            mu_xp=torch.full(shape, 0.0),
+            mu_yp=torch.full(shape, 0.0),
+            sigma_xp=sigma_xp,
+            sigma_yp=sigma_yp,
+            sigma_p=sigma_p,
+            energy=energy,
+            total_charge=total_charge,
+            device=device,
+            dtype=dtype,
+        )
+
+        # Replace the spatial coordinates with the generated ones
+        beam.xs = flattened_xs.view(*shape, num_particles)
+        beam.ys = flattened_ys.view(*shape, num_particles)
+        beam.ss = flattened_ss.view(*shape, num_particles)
+
+        return beam
+
     @classmethod
     def make_linspaced(
         cls,

tests/test_particle_beam.py

@@ -103,3 +103,44 @@ def test_from_twiss_to_twiss():
     assert np.isclose(beam.alpha_y.cpu().numpy(), 1.0, rtol=1e-2)
     assert np.isclose(beam.emittance_y.cpu().numpy(), 3.497810737006068e-09, rtol=1e-2)
     assert np.isclose(beam.energy.cpu().numpy(), 6e6)
+
+
+def test_generate_uniform_ellipsoid_batched():
+    """
+    Test that a `ParticleBeam` generated from a uniform 3D ellipsoid has the correct
+    parameters, i.e. the all particles are within the ellipsoid, and that the other
+    beam parameters are as they would be for a Gaussian beam.
+    """
+    radius_x = torch.tensor([1e-3, 2e-3])
+    radius_y = torch.tensor([1e-4, 2e-4])
+    radius_s = torch.tensor([1e-5, 2e-5])
+
+    num_particles = torch.tensor(1_000_000)
+    sigma_xp = torch.tensor([2e-7, 1e-7])
+    sigma_yp = torch.tensor([3e-7, 2e-7])
+    sigma_p = torch.tensor([0.000001, 0.000002])
+    energy = torch.tensor([1e7, 2e7])
+    total_charge = torch.tensor([1e-9, 3e-9])
+
+    num_particles = torch.tensor(1_000_000)
+    beam = ParticleBeam.uniform_3d_ellispoid(
+        num_particles=num_particles,
+        radius_x=radius_x,
+        radius_y=radius_y,
+        radius_s=radius_s,
+        sigma_xp=sigma_xp,
+        sigma_yp=sigma_yp,
+        sigma_p=sigma_p,
+        energy=energy,
+        total_charge=total_charge,
+    )
+
+    assert beam.num_particles == num_particles
+    assert torch.all(beam.xs.abs().transpose(0, 1) <= radius_x)
+    assert torch.all(beam.ys.abs().transpose(0, 1) <= radius_y)
+    assert torch.all(beam.ss.abs().transpose(0, 1) <= radius_s)
+    assert torch.allclose(beam.sigma_xp, sigma_xp)
+    assert torch.allclose(beam.sigma_yp, sigma_yp)
+    assert torch.allclose(beam.sigma_p, sigma_p)
+    assert torch.allclose(beam.energy, energy)
+    assert torch.allclose(beam.total_charge, total_charge)