diff --git a/tracking/__init__.py b/tracking/__init__.py
index 41f06b9ee215b768e0f5cd99e1b7308eaf246a22..05db3c85cc9817ec1f16279adae27c34bc6167ff 100644
--- a/tracking/__init__.py
+++ b/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
diff --git a/tracking/element.py b/tracking/element.py
index 00f75ed6d6fd33cca7949e4b09711618befcebfc..6f23f7d1a39dbbc447c39140310a7a1a810fb1c7 100644
--- a/tracking/element.py
+++ b/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