diff --git a/mbtrack2/instability/ions.py b/mbtrack2/instability/ions.py
index ee20a309c184c3af144ba4221554aecb02fd263e..8ddd8f7eac65c80ca4f28dc71f263719b9b41a3f 100644
--- a/mbtrack2/instability/ions.py
+++ b/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)
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index 395de32462c2d3a0692a8ec2ff88ab9ede3eac81..a96919fd71ebe6d2d2ebcc334ad4109535467728 100644
--- a/mbtrack2/tracking/particles.py
+++ b/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):
diff --git a/tests/test_bunch.py b/tests/test_bunch.py
index 7c1499a6af032b73f44a5f375cfb669f6c3861ce..cac173ee54f7f19aa3b74d17304e8dbd2ed3ab8e 100644
--- a/tests/test_bunch.py
+++ b/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