Add particle generation for 3D uniform ellispoid distribution

[cr-xu]

Description

Add a preliminary 3d uniform ellispoidal generation routine.

Motivation and Context

This addresses #145

Right now the code is lying in utils but we can think of moving it later.

To get a waterbag bunch, right now one can generate the distribution first and input it into ParticleBeam

from cheetah.utils import generate_3d_uniform_ellipsoid
from cheetah.particles import ParticleBeam
import torch

r_x = torch.tensor(1e-3)
r_y = torch.tensor(1e-3)
r_z = torch.tensor(1e-5)
num_particles = torch.tensor(100_000)

particles = generate_3d_uniform_ellipsoid(num_particles, r_x, r_y, r_z)
parray = ParticleBeam(particles=particles, energy=torch.tensor([1e6], dtype=torch.float32))

• I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)

Checklist

• I have updated the changelog accordingly (required).
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.
• I have reformatted the code and checked that formatting passes (required).
• I have have fixed all issues found by flake8 (required).
• I have ensured that all pytest tests pass (required).
• I have run pytest on a machine with a CUDA GPU and made sure all tests pass (required).
• I have checked that the documentation builds (required).

Note: We are using a maximum length of 88 characters per line

[cr-xu]

@jank324 Any suggestions on where to put the future routines?

@RemiLehe Is that what is needed for the benchmarking? From the papers I assume it's simply a cold beam $p_x=p_y=p_z=0$?

[jank324]

I think this should be a @classmethod of ParameterBeam named 3d_uniform_elipsoid, in the same way as from_parameters and from_twiss. Along that line of thinking it probably makes sense to rename and merge those two at some point into a single method called something like 3d_gaussian_elipsoid that takes either Twiss parameters or like sigmas and such (?)

[jank324]

I think the arguments should be called radius_x and so on. It's not too long and this way I don't have to look at the documentation to figure out what r_x is ... i.e. it's much more readable.

[cr-xu]

I think this should be a @classmethod of ParameterBeam named 3d_uniform_elipsoid, in the same way as from_parameters and from_twiss. Along that line of thinking it probably makes sense to rename and merge those two at some point into a single method called something like 3d_gaussian_elipsoid that takes either Twiss parameters or like sigmas and such (?)

Yes, that could be a solution. I don't know if it will become very crowded in the future, as particle distributions can be quite arbitrary. E.g. transverse 2D + longitudinal 1D with different distributions;
Technically what we already have is also only a decoupled case. e.g. x-x', y-y', $\tau - \delta$, rather than a fully-correlated gaussian ellipsoid (with x-y, x-z, y-z,...)

[RemiLehe]

Thanks for adding this!
Regarding your question: yes, we would only need a cold beam for now (zero divergence, zero energy spread)

[jank324]

I think this should be a @classmethod of ParameterBeam named 3d_uniform_elipsoid, in the same way as from_parameters and from_twiss. Along that line of thinking it probably makes sense to rename and merge those two at some point into a single method called something like 3d_gaussian_elipsoid that takes either Twiss parameters or like sigmas and such (?)

Yes, that could be a solution. I don't know if it will become very crowded in the future, as particle distributions can be quite arbitrary. E.g. transverse 2D + longitudinal 1D with different distributions; Technically what we already have is also only a decoupled case. e.g. x-x', y-y', τ−δ, rather than a fully-correlated gaussian ellipsoid (with x-y, x-z, y-z,...)

I'm not too worried that it would get too crowded. I think

beam = ParticleBeam.3d_uniform_elipsoid()

makes for a very intuitive interface. If I were to look for the distribution code, I would also immediately look where the ParticleBeam source is. I think it's just important, that we keep it structured in the sense that we don't interleave other ParticleBeam methods with the class methods for distributions.

Also, I think that at some point Beam, ParameterBeam and ParticleBeam might end up in separate files.

[jank324]

@cr-xu I reviewed the code by breaking and refactoring it to be more readable. I think it can be merged as it is now.