Add OpenPMD ParticleGroup Converters

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

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(),

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",
+    ],
 )

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)