Prvent zero division in twiss parameter computations

desy-ml · Jan 7, 2024 · b649094 · b649094
1 parent 5283bd1
commit b649094
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/cheetah/particles.py b/cheetah/particles.py
@@ -278,11 +278,11 @@ def normalized_emittance_x(self) -> torch.Tensor:
     @property
     def beta_x(self) -> torch.Tensor:
         """Beta function in x direction in meters."""
-        return self.sigma_x**2 / self.emittance_x
+        return self.sigma_x**2 / (self.emittance_x + 1e-20)
 
     @property
     def alpha_x(self) -> torch.Tensor:
-        return -self.sigma_xxp / self.emittance_x
+        return -self.sigma_xxp / (self.emittance_x + 1e-20)
 
     @property
     def emittance_y(self) -> torch.Tensor:
@@ -301,11 +301,11 @@ def normalized_emittance_y(self) -> torch.Tensor:
     @property
     def beta_y(self) -> torch.Tensor:
         """Beta function in y direction in meters."""
-        return self.sigma_y**2 / self.emittance_y
+        return self.sigma_y**2 / (self.emittance_y + 1e-20)
 
     @property
     def alpha_y(self) -> torch.Tensor:
-        return -self.sigma_yyp / self.emittance_y
+        return -self.sigma_yyp / (self.emittance_y + 1e-20)
 
     def broadcast(self, shape: torch.Size) -> "Beam":
         """Broadcast beam to new shape."""