diff --git a/Tracking/beam.py b/Tracking/beam.py
index b7041d21237da9fbc66a836b28ac64179cac98a1..75143ac592b88d95bd1187860927fd9431f9b90d 100644
--- a/Tracking/beam.py
+++ b/Tracking/beam.py
@@ -62,7 +62,9 @@ class Bunch:
self.alive = pd.Series(np.ones((self.mp_number,),dtype=bool))
self.current = current
if not alive:
- self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool))
+ self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool))
+
+ self._energy_change = pd.Series(np.zeros((self.mp_number,),dtype=float))
def __len__(self):
"""Return the number of alive particles"""
@@ -83,6 +85,18 @@ class Bunch:
def __repr__(self):
"""Return representation of alive particles"""
return f'{self[:]!r}'
+
+ @property
+ def energy_change(self):
+ """Store the particle relative energy change in the last turn. Used by
+ the SynchrotronRadiation class to compute radiation damping. To be
+ changed by Element objects which change the energy, for example RF
+ cavities."""
+ return self._energy_change.loc[self.alive]
+
+ @energy_change.setter
+ def energy_change(self, value):
+ self._energy_change.loc[self.alive] = value
@property
def mp_number(self):
diff --git a/Tracking/one_turn_matrix.py b/Tracking/one_turn_matrix.py
index 45452db027b36851dc9d0d8b2c8c51273498b39e..30e18071b7086b60c9380e54985164cf159ee702 100644
--- a/Tracking/one_turn_matrix.py
+++ b/Tracking/one_turn_matrix.py
@@ -72,15 +72,28 @@ class Long_one_turn(Element):
class SynchrotronRadiation(Element):
"""SyncRad"""
- def __init__(self, ring):
+ def __init__(self, ring, switch = np.ones((3,), dtype=bool)):
self.ring = ring
+ self.switch = switch
@Element.not_empty
def track(self, bunch):
"""track"""
- 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[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]
+ bunch.energy_change = 0
class RF_cavity(Element):
""" Perfect RF cavity class for main and harmonic RF cavities."""
@@ -93,8 +106,12 @@ class RF_cavity(Element):
@Element.not_empty
def track(self,bunch):
"""Track """
- bunch["delta"] += self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*bunch["tau"] + self.theta)
+ energy_change = self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*bunch["tau"] + self.theta)
+ bunch["delta"] += energy_change
+ bunch.energy_change += energy_change
+ def value(self, val):
+ return self.Vc/self.ring.E0*np.cos(self.m*self.ring.omega1*val + self.theta)
class Trans_one_turn(Element):
"""
@@ -118,7 +135,7 @@ class Trans_one_turn(Element):
phase_advance_y = self.phase_advance[1]*(1+self.ring.chro[1]*bunch["delta"])
matrix = np.zeros((6,6,len(bunch)))
matrix[0,0,:] = np.cos(phase_advance_x) + self.alpha[0]*np.sin(phase_advance_x)
- matrix[0,1,:] = self.beta[0]*np.cos(phase_advance_x)
+ matrix[0,1,:] = self.beta[0]*np.sin(phase_advance_x)
matrix[0,2,:] = self.disp[0]
matrix[1,0,:] = -1*self.gamma[0]*np.sin(phase_advance_x)
matrix[1,1,:] = np.cos(phase_advance_x) - self.alpha[0]*np.sin(phase_advance_x)
@@ -126,7 +143,7 @@ class Trans_one_turn(Element):
matrix[2,2,:] = 1
matrix[3,3,:] = np.cos(phase_advance_y) + self.alpha[1]*np.sin(phase_advance_y)
- matrix[3,4,:] = self.beta[1]*np.cos(phase_advance_y)
+ matrix[3,4,:] = self.beta[1]*np.sin(phase_advance_y)
matrix[3,5,:] = self.disp[1]
matrix[4,3,:] = -1*self.gamma[1]*np.sin(phase_advance_y)
matrix[4,4,:] = np.cos(phase_advance_y) - self.alpha[1]*np.sin(phase_advance_y)
diff --git a/Tracking/parallel.py b/Tracking/parallel.py
index c349c124d4beab96bc924d8ae3922c46b9f8c156..ab2eb479c358cde063784bf416f4da4f8842e03f 100644
--- a/Tracking/parallel.py
+++ b/Tracking/parallel.py
@@ -19,8 +19,7 @@ class Mpi:
from mpi4py import MPI
except(ModuleNotFoundError):
print("Please install mpi4py module.")
- else:
- print("MPI error.")
+
self.comm = MPI.COMM_WORLD
self.rank = self.comm.Get_rank()
self.size = self.comm.Get_size()
diff --git a/Tracking/synchrotron.py b/Tracking/synchrotron.py
index 38e1e23e4a4bcbd038af9f795a0242d8ded3ccd2..a75f98bf150756ec341de50465f0634a801e8794 100644
--- a/Tracking/synchrotron.py
+++ b/Tracking/synchrotron.py
@@ -181,4 +181,13 @@ class Synchrotron:
def eta(self):
"""Momentum compaction"""
return self.ac - 1/(self.gamma**2)
-
\ No newline at end of file
+
+ @property
+ def sigma(self):
+ """RMS beam size at equilibrium"""
+ sigma = np.zeros((4,))
+ sigma[0] = (self.emit[0]*self.mean_optics.beta[0] + self.mean_optics.disp[0]**2*self.sigma_delta)**0.5
+ sigma[1] = (self.emit[0]*self.mean_optics.alpha[0] + self.mean_optics.dispp[0]**2*self.sigma_delta)**0.5
+ sigma[2] = (self.emit[1]*self.mean_optics.beta[1] + self.mean_optics.disp[1]**2*self.sigma_delta)**0.5
+ sigma[3] = (self.emit[1]*self.mean_optics.alpha[1] + self.mean_optics.dispp[1]**2*self.sigma_delta)**0.5
+ return sigma
\ No newline at end of file