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,16 +282,37 @@ 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 (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)
+        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
@@ -300,15 +321,19 @@ class Bunch:
         Return the average Courant-Snyder invariant of each plane.
 
         """
-        Jx = ((self.ring.optics.local_gamma[0] * self["x"]**2) +
-              (2 * self.ring.optics.local_alpha[0] * self["x"]) * self["xp"] +
-              (self.ring.optics.local_beta[0] * self["xp"]**2))
-        Jy = ((self.ring.optics.local_gamma[1] * self["y"]**2) +
-              (2 * self.ring.optics.local_alpha[1] * self["y"] * self["yp"]) +
-              (self.ring.optics.local_beta[1] * self["yp"]**2))
-        Js = ((self.ring.long_gamma * self["tau"]**2) +
-              (2 * self.ring.long_alpha * self["tau"] * self["delta"]) +
-              (self.ring.long_beta * self["delta"]**2))
+        xb = self['x'] - self['delta'] * self.ring.optics.local_dispersion[0]
+        yb = self['y'] - self['delta'] * self.ring.optics.local_dispersion[2]
+        xpb = self['xp'] - self['delta'] * self.ring.optics.local_dispersion[1]
+        ypb = self['yp'] - self['delta'] * self.ring.optics.local_dispersion[3]
+        Jx = (self.ring.optics.local_gamma[0] * xb**2) + \
+              (2*self.ring.optics.local_alpha[0] *xb*xpb) + \
+              (self.ring.optics.local_beta[0] * xpb**2)
+        Jy = (self.ring.optics.local_gamma[1] * yb**2) + \
+              (2*self.ring.optics.local_alpha[1] * yb*ypb) + \
+              (self.ring.optics.local_beta[1] * ypb**2)
+        Js = (self.ring.long_gamma * self['tau']**2) + \
+              (2*self.ring.long_alpha * self['tau']*self['delta']) + \
+              (self.ring.long_beta * self['delta']**2)
         return np.array((np.mean(Jx), np.mean(Jy), np.mean(Js)))
 
     def init_gaussian(self, cov=None, mean=None, **kwargs):