From 4edc9492281c95e3d9e08062f9068e6bdce28fc2 Mon Sep 17 00:00:00 2001 From: Jan Kaiser Date: Fri, 10 Jan 2025 16:34:29 +0100 Subject: [PATCH] A little cleanup on OpenPMD conversion implementations --- cheetah/accelerator/aperture.py | 2 +- cheetah/accelerator/drift.py | 2 +- cheetah/accelerator/horizontal_corrector.py | 4 +- cheetah/particles/particle_beam.py | 124 ++++++++++---------- 4 files changed, 67 insertions(+), 65 deletions(-) diff --git a/cheetah/accelerator/aperture.py b/cheetah/accelerator/aperture.py index 592fb672..3c7e6ad4 100644 --- a/cheetah/accelerator/aperture.py +++ b/cheetah/accelerator/aperture.py @@ -16,7 +16,7 @@ class Aperture(Element): Physical aperture. NOTE: The aperture currently only affects beams of type `ParticleBeam` and only has - an effect when the aperture is active. + an effect when the aperture is active. :param x_max: half size horizontal offset in [m]. :param y_max: half size vertical offset in [m]. diff --git a/cheetah/accelerator/drift.py b/cheetah/accelerator/drift.py index cb29570f..38a8cb12 100644 --- a/cheetah/accelerator/drift.py +++ b/cheetah/accelerator/drift.py @@ -17,7 +17,7 @@ class Drift(Element): """ Drift section in a particle accelerator. - Note: the transfer map now uses the linear approximation. + NOTE: The transfer map now uses the linear approximation. Including the R_56 = L / (beta**2 * gamma **2) :param length: Length in meters. diff --git a/cheetah/accelerator/horizontal_corrector.py b/cheetah/accelerator/horizontal_corrector.py index 1a74342c..4fe4fc58 100644 --- a/cheetah/accelerator/horizontal_corrector.py +++ b/cheetah/accelerator/horizontal_corrector.py @@ -17,8 +17,8 @@ class HorizontalCorrector(Element): """ Horizontal corrector magnet in a particle accelerator. - Note: This is modeled as a drift section with - a thin-kick in the horizontal plane. + + NOTE: This is modeled as a drift section with a thin-kick in the horizontal plane. :param length: Length in meters. :param angle: Particle deflection angle in the horizontal plane in rad. diff --git a/cheetah/particles/particle_beam.py b/cheetah/particles/particle_beam.py index 47969510..19565269 100644 --- a/cheetah/particles/particle_beam.py +++ b/cheetah/particles/particle_beam.py @@ -2,9 +2,9 @@ from typing import List, Literal, Optional, Tuple, Union import numpy as np +import pmd_beamphysics as openpmd 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 @@ -672,11 +672,10 @@ def from_astra(cls, path: str, device=None, dtype=torch.float32) -> "ParticleBea @classmethod def from_openpmd_file( - cls, path: str, energy: torch.Tensor, device=None, dtype=torch.float32 + cls, path: str, energy: torch.Tensor, device=None, dtype=None ) -> "ParticleBeam": - """Load an OpenPMD particle distribution file as a Cheetah Beam.""" - - particle_group = ParticleGroup(path) + """Load an OpenPMD particle group HDF5 file as a Cheetah `ParticleBeam`.""" + particle_group = openpmd.ParticleGroup(path) return cls.from_openpmd_particlegroup( particle_group, energy, device=device, dtype=dtype ) @@ -684,21 +683,21 @@ def from_openpmd_file( @classmethod def from_openpmd_particlegroup( cls, - particle_group: ParticleGroup, + particle_group: openpmd.ParticleGroup, energy: torch.Tensor, device=None, - dtype=torch.float32, + dtype=None, ) -> "ParticleBeam": - """Convert an OpenPMD ParticleGroup to a Cheetah Beam. + """ + Create a Cheetah `ParticleBeam` from an OpenPMD `ParticleGroup` object. - :param particle_group: OpenPMD ParticleGroup object. + :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 + # For now, assume an electron beam p0c = torch.sqrt(energy**2 - electron_mass_eV**2) x = torch.from_numpy(particle_group.x) @@ -721,6 +720,59 @@ def from_openpmd_particlegroup( dtype=dtype, ) + def save_as_openpmd_h5(self, path: str) -> None: + """ + Save the `ParticleBeam` as an OpenPMD particle group HDF5 file. + + :param path: Path to the file where the beam should be saved. + """ + particle_group = self.to_openpmd_particlegroup() + particle_group.write(path) + + def to_openpmd_particlegroup(self) -> openpmd.ParticleGroup: + """ + Convert the `ParticleBeam` to an OpenPMD `ParticleGroup` object. + + NOTE: OpenPMD uses boolean particle status flags, i.e. alive or dead. Cheetah's + survival probabilities are converted to status flags by thresholding at 0.5. + + NOTE: At the moment this method only supports non-batched particles + distributions. + + :return: OpenPMD `ParticleGroup` object with the `ParticleBeam`'s particles. + """ + # For now only support non-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, + # TODO: Modify when support for other species was added + "species": "electron", + } + particle_group = openpmd.ParticleGroup(data=data) + + return particle_group + def transformed_to( self, mu_x: Optional[torch.Tensor] = None, @@ -1462,56 +1514,6 @@ 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(),