Merge branch 'master' into 231-add-particle-species-in-particlebeam

Author: jank324

CHANGELOG.md

@@ -8,7 +8,7 @@
 
 ### 🚀 Features
 
-- `ParticleBeam` now supports importing from and exporting to [openPMD-beamphysics](https://github.com/ChristopherMayes/openPMD-beamphysics) HDF5 files and `ParticleGroup` objects. This allows for easy conversion to and from other file formats supported by openPMD-beamphysics. (see #305) (@cr-xu)
+- `ParticleBeam` now supports importing from and exporting to [openPMD-beamphysics](https://github.com/ChristopherMayes/openPMD-beamphysics) HDF5 files and `ParticleGroup` objects. This allows for easy conversion to and from other file formats supported by openPMD-beamphysics. (see #305, #320) (@cr-xu, @Hespe)
 - Add the option to choose the particle species `Species` for `Beam` classes (see #276) (@cr-xu)
 
 ### 🐛 Bug fixes

cheetah/particles/particle_beam.py

@@ -2,7 +2,6 @@
 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 scipy import constants
@@ -692,6 +691,14 @@ def from_openpmd_file(
         cls, path: str, energy: torch.Tensor, device=None, dtype=None
     ) -> "ParticleBeam":
         """Load an openPMD particle group HDF5 file as a Cheetah `ParticleBeam`."""
+        try:
+            import pmd_beamphysics as openpmd
+        except ImportError:
+            raise ImportError(
+                """To use the openPMD beam import, openPMD-beamphysics must be
+                installed."""
+            )
+
         particle_group = openpmd.ParticleGroup(path)
         return cls.from_openpmd_particlegroup(
             particle_group, energy, device=device, dtype=dtype
@@ -700,7 +707,7 @@ def from_openpmd_file(
     @classmethod
     def from_openpmd_particlegroup(
         cls,
-        particle_group: openpmd.ParticleGroup,
+        particle_group: "openpmd.ParticleGroup",  # noqa: F821
         energy: torch.Tensor,
         device=None,
         dtype=None,
@@ -749,7 +756,7 @@ def save_as_openpmd_h5(self, path: str) -> None:
         particle_group = self.to_openpmd_particlegroup()
         particle_group.write(path)
 
-    def to_openpmd_particlegroup(self) -> openpmd.ParticleGroup:
+    def to_openpmd_particlegroup(self) -> "openpmd.ParticleGroup":  # noqa: F821
         """
         Convert the `ParticleBeam` to an openPMD `ParticleGroup` object.
 
@@ -761,6 +768,14 @@ def to_openpmd_particlegroup(self) -> openpmd.ParticleGroup:
 
         :return: openPMD `ParticleGroup` object with the `ParticleBeam`'s particles.
         """
+        try:
+            import pmd_beamphysics as openpmd
+        except ImportError:
+            raise ImportError(
+                """To use the openPMD beam export, openPMD-beamphysics must be
+                installed."""
+            )
+
         # For now only support non-batched particles
         if len(self.particles.shape) != 2:
             raise ValueError("Only non-batched particles are supported.")

test_requirements.txt

@@ -1,3 +1,4 @@
 git+https://github.com/cr-xu/ocelot@update-scipy-compatibility # Ocelot
+openpmd-beamphysics
 pytest
 pytest-cov