Prevent NaNs in emittance computation

desy-ml · Jan 7, 2024 · 5283bd1 · 5283bd1
1 parent df48d85
commit 5283bd1
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/cheetah/particles.py b/cheetah/particles.py
@@ -264,7 +264,11 @@ def sigma_yyp(self) -> torch.Tensor:
     @property
     def emittance_x(self) -> torch.Tensor:
         """Emittance of the beam in x direction in m*rad."""
-        return torch.sqrt(self.sigma_x**2 * self.sigma_xp**2 - self.sigma_xxp**2)
+        return torch.sqrt(
+            torch.clamp_min(
+                self.sigma_x**2 * self.sigma_xp**2 - self.sigma_xxp**2, 0.0
+            )
+        )
 
     @property
     def normalized_emittance_x(self) -> torch.Tensor:
@@ -283,7 +287,11 @@ def alpha_x(self) -> torch.Tensor:
     @property
     def emittance_y(self) -> torch.Tensor:
         """Emittance of the beam in y direction in m*rad."""
-        return torch.sqrt(self.sigma_y**2 * self.sigma_yp**2 - self.sigma_yyp**2)
+        return torch.sqrt(
+            torch.clamp_min(
+                self.sigma_y**2 * self.sigma_yp**2 - self.sigma_yyp**2, 0.0
+            )
+        )
 
     @property
     def normalized_emittance_y(self) -> torch.Tensor: