Emittance calculation

Hello, There are a few changes that need to be verified. I've changed the way emittance is computed. For nonzero dispersion, it gives the correct results. When dispersion is included, I think the difference is minimal for realistic values. But I am not sure.

• Calculation of emittances and CS invariants for nonzero dispersion (based on PyHEADTAIL "cobra" functions)
• Different ion species data in impedance.ions

mbtrack2/tracking/particles.py

@@ -282,27 +282,32 @@ class Bunch:
         """
         Return the bunch emittance for each plane.
         """
-        cor = np.squeeze([[self[name] - self[name].mean()] for name in self])
-        
+
         cov_x = np.cov(self['x'], self['xp'])
-        cov_y = np.cov(self['y'], self['yp'])        
+        cov_y = np.cov(self['y'], self['yp'])
         cov_z = np.cov(self['tau'], self['delta'])
-        
-        if (self.ring.optics.local_dispersion != [0, 0, 0, 0]):
+
+        if (np.array(self.ring.optics.local_dispersion) != np.array([0, 0, 0, 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_y[0, 0] - cov_ydelta[0, 1] * cov_ydelta[0, 1] / cov_z[1,1]
-            sig12 = cov_y[0, 1] - cov_ydelta[0, 1] * cov_ypdelta[0, 1] / cov_z[1, 1]
-            sig22 = cov_y[1, 1] - cov_ypdelta[0, 1] * cov_ypdelta[0, 1] / cov_z[1, 1]
-            emitY = np.sqrt(sig11*sig22-sig12*sig12)
+            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_y[
+                0, 0] - cov_ydelta[0, 1] * cov_ydelta[0, 1] / cov_z[1, 1]
+            sig12 = cov_y[
+                0, 1] - cov_ydelta[0, 1] * cov_ypdelta[0, 1] / cov_z[1, 1]
+            sig22 = cov_y[
+                1, 1] - cov_ypdelta[0, 1] * cov_ypdelta[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))