Make sure parameter beam sigmas are always strictly positive

desy-ml · Jan 7, 2024 · a8c5e57 · a8c5e57
1 parent 5b638f9
commit a8c5e57
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/cheetah/particles.py b/cheetah/particles.py
@@ -683,39 +683,39 @@ def mu_xp(self) -> torch.Tensor:
 
     @property
     def sigma_xp(self) -> torch.Tensor:
-        return torch.sqrt(self._cov[:, 1, 1])
+        return torch.sqrt(self._cov[:, 1, 1]).clamp_min(1e-20)
 
     @property
     def mu_y(self) -> torch.Tensor:
         return self._mu[:, 2]
 
     @property
     def sigma_y(self) -> torch.Tensor:
-        return torch.sqrt(self._cov[:, 2, 2])
+        return torch.sqrt(self._cov[:, 2, 2]).clamp_min(1e-20)
 
     @property
     def mu_yp(self) -> torch.Tensor:
         return self._mu[:, 3]
 
     @property
     def sigma_yp(self) -> torch.Tensor:
-        return torch.sqrt(self._cov[:, 3, 3])
+        return torch.sqrt(self._cov[:, 3, 3]).clamp_min(1e-20)
 
     @property
     def mu_s(self) -> torch.Tensor:
         return self._mu[:, 4]
 
     @property
     def sigma_s(self) -> torch.Tensor:
-        return torch.sqrt(self._cov[:, 4, 4])
+        return torch.sqrt(self._cov[:, 4, 4]).clamp_min(1e-20)
 
     @property
     def mu_p(self) -> torch.Tensor:
         return self._mu[:, 5]
 
     @property
     def sigma_p(self) -> torch.Tensor:
-        return torch.sqrt(self._cov[:, 5, 5])
+        return torch.sqrt(self._cov[:, 5, 5]).clamp_min(1e-20)
 
     @property
     def sigma_xxp(self) -> torch.Tensor: