From 00233cc901631491400b672536fee87c3dffcaef Mon Sep 17 00:00:00 2001
From: Watanyu Foosang <watanyu.f@gmail.com>
Date: Tue, 2 Mar 2021 12:13:45 +0100
Subject: [PATCH] Fix cs_invariant method and sigma method

---
 tracking/particles.py   | 18 +++++++++++-------
 tracking/synchrotron.py |  8 ++++----
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/tracking/particles.py b/tracking/particles.py
index 87988e8..c272e98 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -233,13 +233,17 @@ class Bunch:
     
     @property
     def cs_invariant(self):
-        Jx = (self.ring.optics.local_gamma * np.mean(self['x']**2)) + \
-             (2*self.ring.optics.local_alpha * np.mean(self['x']*self['xp'])) + \
-             (self.ring.optics.local_beta * np.mean(self['xp']**2))
-        Jy = (self.ring.optics.local_gamma * np.mean(self['y']**2)) + \
-             (2*self.ring.optics.local_alpha * np.mean(self['y']*self['yp'])) + \
-             (self.ring.optics.local_beta * np.mean(self['yp']**2))
-        return np.array((Jx,Jy))
+        """
+        Return the average Couran-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)
+        return np.array((np.mean(Jx),np.mean(Jy)))
         
     @property
     def energy_change(self):
diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py
index 1c45d8c..42e2d9f 100644
--- a/tracking/synchrotron.py
+++ b/tracking/synchrotron.py
@@ -244,11 +244,11 @@ class Synchrotron:
             sigma = np.zeros((4,))
             sigma[0] = (self.emit[0]*self.optics.local_beta[0] +
                         self.optics.local_dispersion[0]**2*self.sigma_delta)**0.5
-            sigma[1] = (self.emit[0]*self.optics.local_alpha[0] +
+            sigma[1] = (self.emit[0]*self.optics.local_gamma[0] +
                         self.optics.local_dispersion[1]**2*self.sigma_delta)**0.5
             sigma[2] = (self.emit[1]*self.optics.local_beta[1] +
                         self.optics.local_dispersion[2]**2*self.sigma_delta)**0.5
-            sigma[3] = (self.emit[1]*self.optics.local_alpha[1] +
+            sigma[3] = (self.emit[1]*self.optics.local_gamma[1] +
                         self.optics.local_dispersion[3]**2*self.sigma_delta)**0.5
         else:
             if isinstance(position, (float, int)):
@@ -258,11 +258,11 @@ class Synchrotron:
             sigma = np.zeros((4, n))
             sigma[0,:] = (self.emit[0]*self.optics.beta(position)[0] +
                         self.optics.dispersion(position)[0]**2*self.sigma_delta)**0.5
-            sigma[1,:] = (self.emit[0]*self.optics.alpha(position)[0] +
+            sigma[1,:] = (self.emit[0]*self.optics.gamma(position)[0] +
                         self.optics.dispersion(position)[1]**2*self.sigma_delta)**0.5
             sigma[2,:] = (self.emit[1]*self.optics.beta(position)[1] +
                         self.optics.dispersion(position)[2]**2*self.sigma_delta)**0.5
-            sigma[3,:] = (self.emit[1]*self.optics.alpha(position)[1] +
+            sigma[3,:] = (self.emit[1]*self.optics.gamma(position)[1] +
                         self.optics.dispersion(position)[3]**2*self.sigma_delta)**0.5
         return sigma
     #------------------------------------------
-- 
GitLab