diff --git a/tracking/element.py b/tracking/element.py
index 0ef071c4bbde1c90c9be9c9727277017454af4da..00f75ed6d6fd33cca7949e4b09711618befcebfc 100644
--- a/tracking/element.py
+++ b/tracking/element.py
@@ -125,7 +125,7 @@ class SynchrotronRadiation(Element):
if (self.switch[0] == True):
rand = np.random.normal(size=len(bunch))
bunch["delta"] = ((1 - 2*self.ring.T0/self.ring.tau[2])*bunch["delta"] +
- 2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
+ 2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
if (self.switch[1] == True):
rand = np.random.normal(size=(len(bunch),2))
@@ -142,6 +142,27 @@ class SynchrotronRadiation(Element):
# Reset energy change to 0 for next turn
bunch.energy_change = 0
+ #--------------------------------------
+ # if (self.switch[0] == True):
+ # rand = np.random.normal(size=len(bunch))
+ # bunch["delta"] = ((1 - 2*self.ring.T0/self.ring.tau[2])*bunch["delta"] +
+ # 2*self.ring.sigma_delta*(self.ring.T0/self.ring.tau[2])**0.5*rand)
+
+ # if (self.switch[1] == True):
+ # rand = np.random.normal(size=(len(bunch),2))
+ # bunch["x"] += self.ring.sigma[0]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,0]
+ # bunch["xp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["xp"]
+ # bunch["xp"] += self.ring.sigma[1]*(2*self.ring.T0/self.ring.tau[0])**0.5*rand[:,1]
+
+ # if (self.switch[2] == True):
+ # rand = np.random.normal(size=(len(bunch),2))
+ # bunch["y"] += self.ring.sigma[2]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,0]
+ # bunch["yp"] = (1 + bunch["delta"])/(1 + bunch["delta"] + bunch.energy_change)*bunch["yp"]
+ # bunch["yp"] += self.ring.sigma[3]*(2*self.ring.T0/self.ring.tau[1])**0.5*rand[:,1]
+
+ # # Reset energy change to 0 for next turn
+ # bunch.energy_change = 0
+ #----------------------------------------
class TransverseMap(Element):
"""
Transverse map for a single turn in the synchrotron.
diff --git a/tracking/monitors/__init__.py b/tracking/monitors/__init__.py
index 1d186dae35b7e1759984464db4dacf838da6fc71..bd39dbdae9a30b8b1e9fe42f48e5e55f52d4c5a4 100644
--- a/tracking/monitors/__init__.py
+++ b/tracking/monitors/__init__.py
@@ -15,4 +15,5 @@ from mbtrack2.tracking.monitors.plotting import (plot_bunchdata,
plot_phasespacedata,
plot_profiledata,
plot_beamdata,
- plot_wakedata)
\ No newline at end of file
+ plot_wakedata,
+ plot_tunedata)
\ No newline at end of file
diff --git a/tracking/monitors/monitors.py b/tracking/monitors/monitors.py
index 6a46e38322048d2005375e45d55911925de4fae3..7a2b2095a60f152199d3c44b92aa5c8b26d38e9c 100644
--- a/tracking/monitors/monitors.py
+++ b/tracking/monitors/monitors.py
@@ -16,7 +16,7 @@ from mbtrack2.tracking.element import Element
from mbtrack2.tracking.particles import Bunch, Beam
from scipy.interpolate import interp1d
from abc import ABCMeta
-from mpi4py import MPI
+#from mpi4py import MPI
class Monitor(Element, metaclass=ABCMeta):
"""
diff --git a/tracking/optics.py b/tracking/optics.py
index 58049c53d30e36c3307eae3baf3d35c5a1a9d456..865a1366f78c7fd137a5431da56f4728b4f84244 100644
--- a/tracking/optics.py
+++ b/tracking/optics.py
@@ -160,7 +160,7 @@ class Optics:
kind='linear')
self.disppY = interp1d(self.position, self.dispersion_array[3,:],
kind='linear')
-
+
@property
def local_beta(self):
"""
@@ -215,6 +215,7 @@ class Optics:
"""
return self._local_gamma
+
def beta(self, position):
"""
Return beta functions at specific locations given by position. If no
diff --git a/tracking/particles.py b/tracking/particles.py
index 630e2323b75d9bd651e1997be16b9333b50a2cad..87988e833ef880848155d12dd1764e40cb654939 100644
--- a/tracking/particles.py
+++ b/tracking/particles.py
@@ -231,6 +231,16 @@ class Bunch:
np.mean(self['tau']*self['delta'])**2)**(0.5)
return np.array([emitX, emitY, emitS])
+ @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))
+
@property
def energy_change(self):
"""Store the particle relative energy change in the last turn. Used by
diff --git a/tracking/synchrotron.py b/tracking/synchrotron.py
index f5e2d69a6f5f65e42229ff9f41e5c604a6a24b68..1c45d8c702263724a5298ba3fb8b999a2138957c 100644
--- a/tracking/synchrotron.py
+++ b/tracking/synchrotron.py
@@ -265,6 +265,21 @@ class Synchrotron:
sigma[3,:] = (self.emit[1]*self.optics.alpha(position)[1] +
self.optics.dispersion(position)[3]**2*self.sigma_delta)**0.5
return sigma
+ #------------------------------------------
+ # @property
+ # def sigma(self):
+ # """RMS beam size at equilibrium"""
+ # sigma = np.zeros((4,))
+ # sigma[0] = (self.emit[0]*self.optics.beta(10)[0] +
+ # self.optics.dispersion(10)[0]**2*self.sigma_delta)**0.5
+ # sigma[1] = (self.emit[0]*self.optics.alpha(10)[0] +
+ # self.optics.dispersion(10)[1]**2*self.sigma_delta)**0.5
+ # sigma[2] = (self.emit[1]*self.optics.beta(10)[1] +
+ # self.optics.dispersion(10)[2]**2*self.sigma_delta)**0.5
+ # sigma[3] = (self.emit[1]*self.optics.alpha(10)[1] +
+ # self.optics.dispersion(10)[3]**2*self.sigma_delta)**0.5
+ # return sigma
+ #------------------------------------------
def synchrotron_tune(self, Vrf):
"""