- emittance calculation from covariance matrices including dispersion

mbtrack2/tracking/particles.py

@@ -276,15 +276,31 @@ class Bunch:
         Return the bunch emittance for each plane.
         """
         cor = np.squeeze([[self[name] - self[name].mean()] for name in self])
-        emitX = np.sqrt(
-            np.mean(cor[0]**2) * np.mean(cor[1]**2) -
-            np.mean(cor[0] * cor[1])**2)
-        emitY = np.sqrt(
-            np.mean(cor[2]**2) * np.mean(cor[3]**2) -
-            np.mean(cor[2] * cor[3])**2)
-        emitS = np.sqrt(
-            np.mean(cor[4]**2) * np.mean(cor[5]**2) -
-            np.mean(cor[4] * cor[5])**2)
+        
+        cov_x = np.cov(self['x'], self['xp'])
+        cov_y = np.cov(self['y'], self['yp'])        
+        cov_z = np.cov(self['tau'], self['delta'])
+        
+        if (self.ring.optics.local_dispersion[0, 1] != 0).all():
+            cov_xdelta = np.cov(self['x'], self['delta'])
+            cov_xpdelta = np.cov(self['xp'], self['delta'])
+            cov_ydelta = np.cov(self['y'], self['delta'])        
+            cov_ypdelta = np.cov(self['yp'], self['delta'])        
+
+            sig11 = cov_x[0, 0] - cov_xdelta[0, 1] * cov_xdelta[0, 1] / cov_z[1,1]
+            sig12 = cov_x[0, 1] - cov_xdelta[0, 1] * cov_xpdelta[0, 1] / cov_z[1, 1]
+            sig22 = cov_x[1, 1] - cov_xpdelta[0, 1] * cov_xpdelta[0, 1] / cov_z[1, 1]
+            emitX = np.sqrt(sig11*sig22-sig12*sig12)
+
+            sig11 = cov_x[0, 0] - cov_xdelta[0, 1] * cov_xdelta[0, 1] / cov_z[1,1]
+            sig12 = cov_x[0, 1] - cov_xdelta[0, 1] * cov_xpdelta[0, 1] / cov_z[1, 1]
+            sig22 = cov_x[1, 1] - cov_xpdelta[0, 1] * cov_xpdelta[0, 1] / cov_z[1, 1]
+            emitY = np.sqrt(sig11*sig22-sig12*sig12)
+        else:
+            emitX = np.sqrt(np.linalg.det(cov_x))
+            emitY = np.sqrt(np.linalg.det(cov_y))
+
+        emitS = np.sqrt(np.linalg.det(cov_z))
         return np.array([emitX, emitY, emitS])
 
     @property