Merge branch 'Beam_coupling'

# Conflicts:
#	tracking/element.py

tracking/__init__.py

@@ -12,7 +12,8 @@ from mbtrack2.tracking.rf import RFCavity, CavityResonator
 from mbtrack2.tracking.parallel import Mpi
 from mbtrack2.tracking.optics import Optics, PhyisicalModel
 from mbtrack2.tracking.element import (Element, LongitudinalMap, 
-                                       TransverseMap, SynchrotronRadiation)
+                                       TransverseMap, SynchrotronRadiation,
+                                       SkewQuadrupole)
 from mbtrack2.tracking.aperture import (CircularAperture, ElipticalAperture,
                                         RectangularAperture)
 from mbtrack2.tracking.wakepotential import WakePotential

tracking/element.py

@@ -128,19 +128,14 @@ class SynchrotronRadiation(Element):
                  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]
-        
+            rand = np.random.normal(size=len(bunch))
+            bunch["xp"] = ((1 - 2*self.ring.T0/self.ring.tau[0])*bunch["xp"] +
+                 2*self.ring.sigma()[1]*(self.ring.T0/self.ring.tau[0])**0.5*rand)
+       
         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
+            rand = np.random.normal(size=len(bunch))
+            bunch["yp"] = ((1 - 2*self.ring.T0/self.ring.tau[1])*bunch["yp"] +
+                 2*self.ring.sigma()[3]*(self.ring.T0/self.ring.tau[1])**0.5*rand)
         
 class TransverseMap(Element):
     """
@@ -205,3 +200,33 @@ class TransverseMap(Element):
         bunch["xp"] = xp
         bunch["y"] = y
         bunch["yp"] = yp
+        
+class SkewQuadrupole:
+    """
+    Thin skew quadrupole element used to introduce betatron coupling (the 
+    length of the quadrupole is neglected).
+    
+    Parameters
+    ----------
+    strength : float
+        Integrated strength of the skew quadrupole [m].
+        
+    """
+    def __init__(self, strength):
+        self.strength = strength
+        
+    @Element.parallel        
+    def track(self, bunch):
+        """
+        Tracking method for the element.
+        No bunch to bunch interaction, so written for Bunch objects and
+        @Element.parallel is used to handle Beam objects.
+        
+        Parameters
+        ----------
+        bunch : Bunch or Beam object
+        """
+        
+        bunch['xp'] = bunch['xp'] - self.strength * bunch['y']
+        bunch['yp'] = bunch['yp'] - self.strength * bunch['x']
+

tracking/monitors/monitors.py

@@ -230,7 +230,8 @@ class Monitor(Element, metaclass=ABCMeta):
             
 class BunchMonitor(Monitor):
     """
-    Monitor a single bunch and save attributes (mean, std, emit and current).
+    Monitor a single bunch and save attributes 
+    (mean, std, emit, current, and action).
     
     Parameters
     ----------
@@ -267,9 +268,11 @@ class BunchMonitor(Monitor):
         self.bunch_number = bunch_number
         group_name = "BunchData_" + str(self.bunch_number)
         dict_buffer = {"mean":(6, buffer_size), "std":(6, buffer_size),
-                     "emit":(3, buffer_size), "current":(buffer_size,)}
+                     "emit":(3, buffer_size), "current":(buffer_size,),
+                     "cs_invariant":(2, buffer_size)}
         dict_file = {"mean":(6, total_size), "std":(6, total_size),
-                     "emit":(3, total_size), "current":(total_size,)}
+                     "emit":(3, total_size), "current":(total_size,),
+                     "cs_invariant":(2, total_size)}
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode)
         

tracking/monitors/plotting.py

@@ -112,7 +112,7 @@ def plot_beamdata(filename, dataset, option=None, stat_var=None, x_var="time"):
     file.close()
     return fig    
               
-def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
+def plot_bunchdata(filename, bunch_number, dataset, dimension="x", x_var="time"):
     """
     Plot data recorded by BunchMonitor.
     
@@ -123,13 +123,14 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
     bunch_number : int
         Bunch to plot. This has to be identical to 'bunch_number' parameter in 
         'BunchMonitor' object.
-    detaset : {"current","emit","mean","std"}
+    dataset : {"current", "emit", "mean", "std", "cs_invariant"}
         HDF5 file's dataset to be plotted.
-    option : str, optional
-        If dataset is "emit", "mean", or "std", the variable name to be plotted
-        needs to be specified :
-            for "emit", option = {"x","y","s"}
-            for "mean" and "std", option = {"x","xp","y","yp","tau","delta"}    
+    dimension : str, optional
+        The dimension of the dataset to plot. Use "None" for "current",
+        otherwise use the following : 
+            for "emit", dimension = {"x","y","s"},
+            for "mean" and "std", dimension = {"x","xp","y","yp","tau","delta"},
+            for "action", dimension = {"x","y"}.
     x_var : {"time", "current"}, optional
         Variable to be plotted on the horizontal axis. The default is "time".
         
@@ -149,19 +150,19 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
         label = "current (mA)"
         
     elif dataset == "emit":
-        option_dict = {"x":0, "y":1, "s":2}
+        dimension_dict = {"x":0, "y":1, "s":2}
                          
-        y_var = file[group][dataset][option_dict[option]]*1e9
+        y_var = file[group][dataset][dimension_dict[dimension]]*1e9
         
-        if option == "x": label = "hor. emittance (nm.rad)"
-        elif option == "y": label = "ver. emittance (nm.rad)"
-        elif option == "s": label = "long. emittance (nm.rad)"
+        if dimension == "x": label = "hor. emittance (nm.rad)"
+        elif dimension == "y": label = "ver. emittance (nm.rad)"
+        elif dimension == "s": label = "long. emittance (nm.rad)"
         
         
-    elif dataset == "mean" or "std":                        
-        option_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} 
+    elif dataset == "mean" or dataset == "std":                        
+        dimension_dict = {"x":0, "xp":1, "y":2, "yp":3, "tau":4, "delta":5} 
         scale = [1e6, 1e6, 1e6, 1e6, 1e12, 1]        
-        axis_index = option_dict[option]
+        axis_index = dimension_dict[dimension]
         
         y_var = file[group][dataset][axis_index]*scale[axis_index]
         if dataset == "mean":
@@ -174,6 +175,13 @@ def plot_bunchdata(filename, bunch_number, dataset, option=None, x_var="time"):
         
         label = label_list[axis_index]
         
+    elif dataset == "cs_invariant":
+        dimension_dict = {"x":0, "y":1}
+        axis_index = dimension_dict[dimension]
+        y_var = file[group][dataset][axis_index]
+        label_list = ['$J_x$ (m)', '$J_y$ (m)']
+        label = label_list[axis_index]
+        
     if x_var == "current":
         x_axis = file[group]["current"][:] * 1e3
         xlabel = "current (mA)"
@@ -544,5 +552,4 @@ def plot_tunedata(filename, bunch_number):
     ax2.set_ylabel("Synchrotron tune")
     
     file.close()
-    return (fig1, fig2)
-    
\ No newline at end of file
+    return (fig1, fig2)  
\ No newline at end of file

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 Courant-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):

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