Update SynchrotronRadiation and add SkewQuadrupole

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

@@ -125,44 +125,18 @@ 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))
-            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
-        
-        #--------------------------------------
-        # 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
-        #---------------------------------------- 
+            rand = np.random.normal(size=len(bunch))
+            bunch["yp"] = ((1 - 2*self.ring.T0/self.ring.tau[1])*bunch["yp"] +
+                 2*self.ring.sigma()[1]*(self.ring.T0/self.ring.tau[1])**0.5*rand)
+        
 class TransverseMap(Element):
     """
     Transverse map for a single turn in the synchrotron.
@@ -226,3 +200,42 @@ class TransverseMap(Element):
         bunch["xp"] = xp
         bunch["y"] = y
         bunch["yp"] = yp
+        
+class SkewQuadrupole:
+    """
+    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
+        """
+        matrix = np.zeros((4,4,len(bunch)))
+        matrix[0,0,:] = 1
+        matrix[1,1,:] = 1
+        matrix[2,2,:] = 1
+        matrix[3,3,:] = 1
+        matrix[1,2,:] = -1*self.strength
+        matrix[3,0,:] = -1*self.strength
+        
+        xp = bunch['xp'] + matrix[1,2,:]*bunch['y']
+        yp = bunch['yp'] + matrix[3,0,:]*bunch['x']
+    
+        bunch['xp'] = xp
+        bunch['yp'] = yp