diff --git a/CHANGELOG.md b/CHANGELOG.md index 57f6a044..ca9d36d3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ ### 🚀 Features +- Add a conversion to [OpenPMD-beamphysics](https://github.com/ChristopherMayes/openPMD-beamphysics) with export to h5 files for the `ParticleBeam` class. This allows also exporting to other file formats via the beamphysics package. (#305 @cr-xu) + ### 🐛 Bug fixes ### 🐆 Other diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index 1d8d04be..47969510 100644 --- a/cheetah/particles/particle_beam.py +++ b/cheetah/particles/particle_beam.py @@ -4,6 +4,7 @@ import numpy as np import torch from matplotlib import pyplot as plt +from pmd_beamphysics import ParticleGroup from scipy import constants from scipy.constants import physical_constants from scipy.ndimage import gaussian_filter @@ -669,6 +670,57 @@ def from_astra(cls, path: str, device=None, dtype=torch.float32) -> "ParticleBea dtype=dtype, ) + @classmethod + def from_openpmd_file( + cls, path: str, energy: torch.Tensor, device=None, dtype=torch.float32 + ) -> "ParticleBeam": + """Load an OpenPMD particle distribution file as a Cheetah Beam.""" + + particle_group = ParticleGroup(path) + return cls.from_openpmd_particlegroup( + particle_group, energy, device=device, dtype=dtype + ) + + @classmethod + def from_openpmd_particlegroup( + cls, + particle_group: ParticleGroup, + energy: torch.Tensor, + device=None, + dtype=torch.float32, + ) -> "ParticleBeam": + """Convert an OpenPMD ParticleGroup to a Cheetah Beam. + + :param particle_group: OpenPMD ParticleGroup object. + :param energy: Reference energy of the beam in eV. + :param device: Device to move the beam's particle array to. If set to `"auto"` a + CUDA GPU is selected if available. The CPU is used otherwise. + :param dtype: Data type of the generated particles. + """ + + # Assume only electron now + p0c = torch.sqrt(energy**2 - electron_mass_eV**2) + + x = torch.from_numpy(particle_group.x) + y = torch.from_numpy(particle_group.y) + px = torch.from_numpy(particle_group.px) / p0c + py = torch.from_numpy(particle_group.py) / p0c + tau = torch.from_numpy(particle_group.t) * speed_of_light + delta = (torch.from_numpy(particle_group.energy) - energy) / p0c + + particles = torch.stack([x, px, y, py, tau, delta, torch.ones_like(x)], dim=-1) + particle_charges = torch.from_numpy(particle_group.weight) + survival_probabilities = torch.from_numpy(particle_group.status) + + return cls( + particles=particles, + energy=energy, + particle_charges=particle_charges, + survival_probabilities=survival_probabilities, + device=device, + dtype=dtype, + ) + def transformed_to( self, mu_x: Optional[torch.Tensor] = None, @@ -1410,6 +1462,56 @@ def momenta(self) -> torch.Tensor: """Momenta of the individual particles.""" return torch.sqrt(self.energies**2 - electron_mass_eV**2) + def save_as_openpmd_h5(self, filename: str) -> None: + """Save the ParticleBeam as an OpenPMD beamphysics HDF5 file. + + :param filename: Path to the file where the beam should be saved. + """ + particle_group = self.to_openpmd_particlegroup() + particle_group.write(filename) + + def to_openpmd_particlegroup(self) -> None: + """ + Convert the beam to an OpenPMD-beamphysics ParticleGroup object. + + Note: Currently OpenPMD-beamphysics only supports boolean particle status + flags, i.e. alive or dead. The survival_probabilities will be converted to + status flags by thresholding at 0.5. + + :return: OpenPMD-beamphysics ParticleGroup object with the beam's particles. + """ + # For now only support none-batched particles + if len(self.particles.shape) != 2: + raise ValueError("Only non-batched particles are supported.") + + n_particles = self.num_particles + weights = np.ones(n_particles) + px = self.px * self.p0c + py = self.py * self.p0c + p_total = torch.sqrt(self.energies**2 - electron_mass_eV**2) + pz = torch.sqrt(p_total**2 - px**2 - py**2) + t = self.tau / speed_of_light + weights = self.particle_charges + # To be discussed + status = self.survival_probabilities > 0.5 + + data = { + "x": self.x.numpy(), + "y": self.y.numpy(), + "z": self.tau.numpy(), + "px": px.numpy(), + "py": py.numpy(), + "pz": pz.numpy(), + "t": t.numpy(), + "weight": weights, + "status": status, + # To be modified after adding support for other species + "species": "electron", + } + particle_group = ParticleGroup(data=data) + + return particle_group + def clone(self) -> "ParticleBeam": return ParticleBeam( particles=self.particles.clone(), diff --git a/setup.py b/setup.py index 61fb4e46..74d1a384 100644 --- a/setup.py +++ b/setup.py @@ -19,5 +19,12 @@ long_description_content_type="text/markdown", packages=[package for package in find_packages() if package.startswith("cheetah")], python_requires=">=3.9", - install_requires=["matplotlib", "numpy", "scipy", "torch"], + install_requires=[ + "h5py", + "matplotlib", + "numpy", + "openpmd-beamphysics", + "scipy", + "torch", + ], ) diff --git a/tests/test_openpmd_conversion.py b/tests/test_openpmd_conversion.py new file mode 100644 index 00000000..73efc9b1 --- /dev/null +++ b/tests/test_openpmd_conversion.py @@ -0,0 +1,63 @@ +import torch + +from cheetah import ParticleBeam + + +def test_particlebeam_to_and_from_particlegroup(): + + ref_energy = torch.tensor(1e6) + beam = ParticleBeam.from_parameters( + num_particles=10000, + mu_x=torch.tensor(1e-4), + sigma_x=torch.tensor(2e-5), + mu_y=torch.tensor(1e-4), + sigma_y=torch.tensor(2e-5), + sigma_p=torch.tensor(1e-4), + energy=ref_energy, + total_charge=torch.tensor(1e-9), + dtype=torch.float64, + ) + + particle_group = beam.to_openpmd_particlegroup() + + # Test that the loaded particle group is the same as the original beam + beam_loaded = ParticleBeam.from_openpmd_particlegroup( + particle_group, energy=ref_energy, dtype=torch.float64 + ) + + assert beam.num_particles == beam_loaded.num_particles + assert torch.allclose(beam.particles, beam_loaded.particles) + assert torch.allclose(beam.particle_charges, beam_loaded.particle_charges) + + +def test_particlebeam_to_and_from_openpmd_h5(): + + ref_energy = torch.tensor(1e6) + beam = ParticleBeam.from_parameters( + num_particles=10000, + mu_x=torch.tensor(1e-4), + sigma_x=torch.tensor(2e-5), + mu_y=torch.tensor(1e-4), + sigma_y=torch.tensor(2e-5), + sigma_p=torch.tensor(1e-4), + energy=ref_energy, + total_charge=torch.tensor(1e-9), + dtype=torch.float64, + ) + + filename = "tests/resources/test_particlebeam_to_and_from_openpmd_h5.h5" + beam.save_as_openpmd_h5(filename) + + # Test that the loaded particle group is the same as the original beam + beam_loaded = ParticleBeam.from_openpmd_file( + filename, energy=ref_energy, dtype=torch.float64 + ) + + assert beam.num_particles == beam_loaded.num_particles + assert torch.allclose(beam.particles, beam_loaded.particles) + assert torch.allclose(beam.particle_charges, beam_loaded.particle_charges) + + # Clean up + import os + + os.remove(filename)