From 2d5105b4071817c55407b4f0a569caa20524797f Mon Sep 17 00:00:00 2001
From: gubaidulinvadim <gubaidulinvadim@gmail.com>
Date: Thu, 13 Jun 2024 10:29:25 +0200
Subject: [PATCH] - emittance calculation from covariance matrices including
dispersion
---
mbtrack2/tracking/particles.py | 34 +++++++++++++++++++++++++---------
1 file changed, 25 insertions(+), 9 deletions(-)
diff --git a/mbtrack2/tracking/particles.py b/mbtrack2/tracking/particles.py
index 939f476..8a205e2 100644
--- a/mbtrack2/tracking/particles.py
+++ b/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
--
GitLab