Adds a SynchrotronRadiation.qexcitation optional parameter to turn off quantum excitation.

mbtrack2/tracking/element.py

@@ -144,12 +144,22 @@ class SynchrotronRadiation(Element):
     ----------
     ring : Synchrotron object
     switch : bool array of shape (3,), optional
-        allow to choose on which plane the synchrotron radiation is active
+        If False in one plane (long, x, y), the synchrotron radiation is turned 
+        off.
+        The default is True, in all three planes.
+    qexcitation : bool, optional
+        If False, the quantum excitation is turned off.
+        The default is True.
+        
     """
 
-    def __init__(self, ring, switch=np.ones((3, ), dtype=bool)):
+    def __init__(self,
+                 ring,
+                 switch=np.ones((3, ), dtype=bool),
+                 qexcitation=True):
         self.ring = ring
         self.switch = switch
+        self.qexcitation = qexcitation
 
     @Element.parallel
     def track(self, bunch):
@@ -163,23 +173,37 @@ class SynchrotronRadiation(Element):
         bunch : Bunch or Beam object
         """
         N = len(bunch)
-        if self.switch[0] == True:
-            rand = np.random.standard_normal(size=N)
-            bunch["delta"] = (1 - 2 * self.ring.T0 / self.ring.tau[2]) * bunch[
-                "delta"] + 2 * self.ring.sigma_delta * (
+
+        excitation = 0
+        if self.switch[0]:
+
+            if self.qexcitation:
+                rand = np.random.standard_normal(size=N)
+                excitation = 2 * self.ring.sigma_delta * (
                     self.ring.T0 / self.ring.tau[2])**0.5 * rand
 
-        if self.switch[1] == True:
-            rand = np.random.standard_normal(size=N)
-            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.standard_normal(size=N)
-            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
+            bunch["delta"] = (1 - 2 * self.ring.T0 /
+                              self.ring.tau[2]) * bunch["delta"] + excitation
+
+        if self.switch[1]:
+
+            if self.qexcitation:
+                rand = np.random.standard_normal(size=N)
+                excitation = 2 * self.ring.sigma()[1] * (
+                    self.ring.T0 / self.ring.tau[0])**0.5 * rand
+
+            bunch["xp"] = (1 - 2 * self.ring.T0 /
+                           self.ring.tau[0]) * bunch["xp"] + excitation
+
+        if self.switch[2]:
+
+            if self.qexcitation:
+                rand = np.random.standard_normal(size=N)
+                excitation = 2 * self.ring.sigma()[3] * (
+                    self.ring.T0 / self.ring.tau[1])**0.5 * rand
+
+            bunch["yp"] = (1 - 2 * self.ring.T0 /
+                           self.ring.tau[1]) * bunch["yp"] + excitation
 
 
 class SkewQuadrupole: