Add the contribution of dispersion when bunch/beam is initialized

tracking/particles.py

@@ -252,10 +252,9 @@ class Bunch:
     def init_gaussian(self, cov=None, mean=None, **kwargs):
         """
         Initialize bunch particles with 6D gaussian phase space.
-        Covariance matrix is taken from [1].
-        
-        !!! -> disp & dispp not included yet.
-
+        Covariance matrix is taken from [1] and dispersion is added following
+        the method explained in [2].
+                
         Parameters
         ----------
         cov : (6,6) array, optional
@@ -267,6 +266,7 @@ class Bunch:
         ----------
         [1] Wiedemann, H. (2015). Particle accelerator physics. 4th 
         edition. Springer, Eq.(8.38) of p223.
+        [2] http://www.pp.rhul.ac.uk/bdsim/manual-develop/dev_beamgeneration.html
 
         """
         if mean is None:
@@ -278,14 +278,22 @@ class Bunch:
             optics = kwargs.get("optics", self.ring.optics)
             
             cov = np.zeros((6,6))
-            cov[0,0] = self.ring.emit[0]*optics.local_beta[0]
-            cov[1,1] = self.ring.emit[0]*optics.local_gamma[0]
-            cov[0,1] = -1*self.ring.emit[0]*optics.local_alpha[0]
-            cov[1,0] = -1*self.ring.emit[0]*optics.local_alpha[0]
-            cov[2,2] = self.ring.emit[1]*optics.local_beta[1]
-            cov[3,3] = self.ring.emit[1]*optics.local_gamma[1]
-            cov[2,3] = -1*self.ring.emit[1]*optics.local_alpha[1]
-            cov[3,2] = -1*self.ring.emit[1]*optics.local_alpha[1]
+            cov[0,0] = self.ring.emit[0]*optics.local_beta[0] + (optics.local_dispersion[0]*self.ring.sigma_delta)**2
+            cov[1,1] = self.ring.emit[0]*optics.local_gamma[0] + (optics.local_dispersion[1]*self.ring.sigma_delta)**2
+            cov[0,1] = -1*self.ring.emit[0]*optics.local_alpha[0] + (optics.local_dispersion[0]*optics.local_dispersion[1]*self.ring.sigma_delta**2)
+            cov[1,0] = -1*self.ring.emit[0]*optics.local_alpha[0] + (optics.local_dispersion[0]*optics.local_dispersion[1]*self.ring.sigma_delta**2)
+            cov[0,5] = optics.local_dispersion[0]*self.ring.sigma_delta**2
+            cov[5,0] = optics.local_dispersion[0]*self.ring.sigma_delta**2
+            cov[1,5] = optics.local_dispersion[1]*self.ring.sigma_delta**2
+            cov[5,1] = optics.local_dispersion[1]*self.ring.sigma_delta**2
+            cov[2,2] = self.ring.emit[1]*optics.local_beta[1] + (optics.local_dispersion[2]*self.ring.sigma_delta)**2
+            cov[3,3] = self.ring.emit[1]*optics.local_gamma[1] + (optics.local_dispersion[3]*self.ring.sigma_delta)**2
+            cov[2,3] = -1*self.ring.emit[1]*optics.local_alpha[1] + (optics.local_dispersion[2]*optics.local_dispersion[3]*self.ring.sigma_delta**2)
+            cov[3,2] = -1*self.ring.emit[1]*optics.local_alpha[1] + (optics.local_dispersion[2]*optics.local_dispersion[3]*self.ring.sigma_delta**2)
+            cov[2,5] = optics.local_dispersion[2]*self.ring.sigma_delta**2
+            cov[5,2] = optics.local_dispersion[2]*self.ring.sigma_delta**2
+            cov[3,5] = optics.local_dispersion[3]*self.ring.sigma_delta**2
+            cov[5,3] = optics.local_dispersion[3]*self.ring.sigma_delta**2
             cov[4,4] = sigma_0**2
             cov[5,5] = sigma_delta**2