Fix emittance and CS invariant calculation when dispersion is included.

mbtrack2/instability/ions.py

@@ -31,7 +31,7 @@ def ion_cross_section(ring, ion):
     by a relativistic electron using the relativistic Bethe asymptotic formula
     [1].
 
-    Values of M02 and C02 from [2] and [3] (values of constants are independent of beam energy).
+    Values of M02 and C02 from [2-4] (values of constants are independent of beam energy).
 
     Parameters
     ----------
@@ -52,6 +52,9 @@ def ion_cross_section(ring, ion):
     [2] : P. F. Tavares, "Bremsstrahlung detection of ions trapped in the EPA
     electron beam", Part. Accel. 43 (1993).
     [3] : A. G. Mathewson, S. Zhang, "Beam-gas ionisation cross sections at 7.0 TEV", CERN Tech. rep. Vacuum-Technical-Note-96-01. https://cds.cern.ch/record/1489148/
+    [4] : F. Rieke and W. Prepejchal, "Ionization Cross Sections of Gaseous 
+    Atoms and Molecules for High-Energy Electrons and Positrons", 
+    Phys. Rev. A 6, 1507 (1972).
 
     """
     if ion == "CO":
@@ -66,6 +69,9 @@ def ion_cross_section(ring, ion):
     elif ion == "CH4":
         M02 = 4.23
         C0 = 42.85
+    elif ion == "H20":
+        M02 = 3.24
+        C0 = 32.26
     else:
         raise NotImplementedError
 
@@ -235,6 +241,10 @@ def fast_beam_ion(
         A = 28
     elif ion == "H2":
         A = 2
+    elif ion == "H2O":
+        A = 16
+    elif ion == "CO2":
+        A = 44
 
     sigma_i = ion_cross_section(ring, ion)
 

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
@@ -303,15 +324,19 @@ class Bunch:
         Twiss parameters need to be computed using the ring.get_longitudinal_twiss
         method.
         """
-        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):

tests/test_bunch.py

@@ -70,4 +70,37 @@ def test_bunch_plots(mybunch):
     mybunch.plot_phasespace()
     mybunch.plot_profile()
     assert True
+def test_bunch_emittance(demo_ring):
+    mp_number = 1_000_000
+    current = 1.2e-3
+    mybunch = Bunch(demo_ring, mp_number=mp_number, current=current, 
+                    track_alive=False)    
+    mybunch.init_gaussian()
+    np.testing.assert_allclose(mybunch.emit[0], demo_ring.emit[0], rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated')
+    np.testing.assert_allclose(mybunch.emit[1], demo_ring.emit[1], rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated')
+
+    np.testing.assert_allclose(mybunch.emit[0], mybunch.cs_invariant[0]/2, rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only')
+    np.testing.assert_allclose(mybunch.emit[1], mybunch.cs_invariant[1]/2, rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only')
+    
+
+def test_bunch_emittance_with_dispersion(demo_ring):
+    mp_number = 1_000_000
+    current = 1.2e-3
+    demo_ring.optics.local_dispersion = np.array([1e-2, 1e-3, 1e-2, 1e-3])
+    mybunch = Bunch(demo_ring, mp_number=mp_number, current=current, 
+                    track_alive=False)    
+    mybunch.init_gaussian()
+    np.testing.assert_allclose(mybunch.emit[0], demo_ring.emit[0], rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {demo_ring.emit[0]} initialised, {mybunch.emit[0]:} calculated')
+    np.testing.assert_allclose(mybunch.emit[1], demo_ring.emit[1], rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {demo_ring.emit[1]} initialised, {mybunch.emit[1]:} calculated')
+
+    np.testing.assert_allclose(mybunch.emit[0], mybunch.cs_invariant[0]/2, rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {mybunch.cs_invariant[0]/2} calculated with optics functions, {mybunch.emit[0]:} calculated with coordinates only')
+    np.testing.assert_allclose(mybunch.emit[1], mybunch.cs_invariant[1]/2, rtol=0, atol=1e-10,
+     err_msg=f'Emittances do not match. {mybunch.cs_invariant[1]/2} calculated with optics functions, {mybunch.emit[1]:} calculated with coordinates only')
     
\ No newline at end of file