Merge pull request #9 from amanmdesai/minor-updates

Minor updates

pymcabc/cross_section.py

@@ -23,7 +23,11 @@ def __init__(self):
 
     def s_channel(self):
         """definition for s channel"""
-        return (self.g**2) / (self.Ecm**2 - self.mx**2)
+        deno = self.Ecm**2 - self.mx**2
+        if abs(deno) <= 0.01:
+            return (self.g**2) / (deno + 100)
+        else:
+            return (self.g**2) / deno
 
     def t_channel(self, costh, pf):
         """definition for t channel"""
@@ -134,7 +138,7 @@ def calc_xsection(self, N: int = 40000):
         w_square = library["w_square"][0]
         w_max = library["w_max"][0]
         sigma_x = w_sum * pymcabc.constants.convert / (N * 1e12)
-        variance = math.sqrt(w_square / N - (w_sum / N) ** 2)  # barn unit
+        variance = math.sqrt(abs(w_square / N - (w_sum / N) ** 2))  # barn unit
         error = (
             variance * pymcabc.constants.convert / (math.sqrt(N) * 1e12)
         )  # barn unit

pymcabc/decay_particle.py

@@ -24,27 +24,32 @@ def __init__(self):
 
     def rotate(self, pdecay: Particle):
         """rotate particle"""
-        costh = (2*random.random()) - 1
+        costh = (2 * random.random()) - 1
         sinth = math.sqrt(1 - costh**2)
         phi = 2 * math.pi * random.random()
         sinPhi = math.sin(phi)
         cosPhi = math.cos(phi)
-        pdecay_out = Particle(0,0,0,0)
-        pdecay_out.px = pdecay.pz * sinth *cosPhi
-        pdecay_out.py = pdecay.pz * sinth *sinPhi
+        pdecay_out = Particle(0, 0, 0, 0)
+        pdecay_out.px = pdecay.pz * sinth * cosPhi
+        pdecay_out.py = pdecay.pz * sinth * sinPhi
         pdecay_out.pz = pdecay.pz * costh
-        pdecay_out.E = pdecay.E 
-        
+        pdecay_out.E = pdecay.E
+
         return pdecay_out
 
     def decay(self, top: Particle):
         """decay particle"""
 
-        E1 = (top.mass()**2 - self.decay2_mass**2 + self.decay1_mass**2)/(2*top.mass())
-
-        E2 =  top.mass() - E1
+        E1 = (top.mass() ** 2 - self.decay2_mass**2 + self.decay1_mass**2) / (
+            2 * top.mass()
+        )
+
+        E2 = top.mass() - E1
 
-        self.decay_p = math.sqrt((top.mass()**2  - (self.decay1_mass + self.decay2_mass)**2) * (top.mass()**2  - (self.decay1_mass - self.decay2_mass)**2)) / (2*top.mass())
+        self.decay_p = math.sqrt(
+            (top.mass() ** 2 - (self.decay1_mass + self.decay2_mass) ** 2)
+            * (top.mass() ** 2 - (self.decay1_mass - self.decay2_mass) ** 2)
+        ) / (2 * top.mass())
 
         """ 
         self.decay_p = (
@@ -71,15 +76,15 @@ def decay(self, top: Particle):
         ) 
         """
 
-        decay1 = Particle(0,0,0,0)
-        decay2 = Particle(0,0,0,0)
+        decay1 = Particle(0, 0, 0, 0)
+        decay2 = Particle(0, 0, 0, 0)
 
         decay1.set4momenta(E1, 0, 0, self.decay_p)
         decay2.set4momenta(E2, 0, 0, -self.decay_p)
 
         decay1 = self.rotate(decay1)
         decay2.set4momenta(E2, -decay1.px, -decay1.py, -decay1.pz)
-        #decay2 = self.rotate(decay2)
+        # decay2 = self.rotate(decay2)
 
         return decay1, decay2
 
@@ -94,26 +99,25 @@ def prepare_decay(self, top: Particle):
             # decay_process = decay_process.replace(" < "," ")
             # decay_process = decay_process.split(" ")
             # print(top.mass()[0], self.massive)
-            decay_array1_px = [0]*len(top.px)
-            decay_array1_py = [0]*len(top.px)
-            decay_array1_pz = [0]*len(top.px)
-            decay_array1_E = [0]*len(top.px)
+            decay_array1_px = [0] * len(top.px)
+            decay_array1_py = [0] * len(top.px)
+            decay_array1_pz = [0] * len(top.px)
+            decay_array1_E = [0] * len(top.px)
 
-            decay_array2_px = [0]*len(top.px)
-            decay_array2_py = [0]*len(top.px)
-            decay_array2_pz = [0]*len(top.px)
-            decay_array2_E = [0]*len(top.px)
+            decay_array2_px = [0] * len(top.px)
+            decay_array2_py = [0] * len(top.px)
+            decay_array2_pz = [0] * len(top.px)
+            decay_array2_E = [0] * len(top.px)
 
             for i in range(len(top.px)):
-                heavy_state = Particle(top.E[i],top.px[i], top.py[i],top.pz[i])
-                if self.nearlyequal(heavy_state.mass(), self.massive) and heavy_state.mass() > (
-                    self.decay1_mass + self.decay2_mass
-                ):
-
+                heavy_state = Particle(top.E[i], top.px[i], top.py[i], top.pz[i])
+                if self.nearlyequal(
+                    heavy_state.mass(), self.massive
+                ) and heavy_state.mass() > (self.decay1_mass + self.decay2_mass):
                     decay1, decay2 = self.decay(heavy_state)
                     decay1 = decay1.boost(heavy_state)
                     decay2 = decay2.boost(heavy_state)
-                    
+
                     decay_array1_px[i] = decay1.px
                     decay_array1_py[i] = decay1.py
                     decay_array1_pz[i] = decay1.pz
@@ -134,10 +138,11 @@ def prepare_decay(self, top: Particle):
                     decay_array2_py[i] = output
                     decay_array2_pz[i] = output
                     decay_array2_E[i] = output
-
-            decay1 = Particle(decay_array1_E,decay_array1_px, decay_array1_py, decay_array1_pz)
-            decay2 = Particle(decay_array2_E,decay_array2_px, decay_array2_py, decay_array2_pz)
-            return decay1, decay2
-
-
 
+            decay1 = Particle(
+                decay_array1_E, decay_array1_px, decay_array1_py, decay_array1_pz
+            )
+            decay2 = Particle(
+                decay_array2_E, decay_array2_px, decay_array2_py, decay_array2_pz
+            )
+            return decay1, decay2

pymcabc/detector.py

@@ -5,7 +5,8 @@
 
 class Detector:
     """Applies gaussian smearing on E and momenta"""
-    def __init__(self, sigma: float =1., factor: float =1.):
+
+    def __init__(self, sigma: float = 1.0, factor: float = 1.0):
         self.sigma = sigma
         self.factor = factor
 
@@ -20,28 +21,35 @@ def gauss_smear(self, particle: Particle):
         if particle.px[0] == -9 and particle.py[0] == -9:
             return particle
         else:
-            output_px = [0]*len(particle.px)
-            output_py = [0]*len(particle.px)
-            output_pz = [0]*len(particle.px)
-            output_E = [0]*len(particle.px)
+            output_px = [0] * len(particle.px)
+            output_py = [0] * len(particle.px)
+            output_pz = [0] * len(particle.px)
+            output_E = [0] * len(particle.px)
             for i in range(len(particle.px)):
-
-                momentum =  math.sqrt(particle.px[i]**2  + particle.py[i]**2  + particle.pz[i]**2 )
-                random_measure_momentum = random.gauss(momentum, self.factor*self.sigma) / momentum
-                random_measure_energy = random.gauss(particle.E[i], self.sigma) / particle.E[i]
-
-                output_px[i] = random_measure_momentum*particle.px[i]
-                output_py[i] = random_measure_momentum*particle.py[i]
-                output_pz[i] = random_measure_momentum*particle.pz[i]
+                momentum = math.sqrt(
+                    particle.px[i] ** 2 + particle.py[i] ** 2 + particle.pz[i] ** 2
+                )
+                random_measure_momentum = (
+                    random.gauss(momentum, self.factor * self.sigma) / momentum
+                )
+                random_measure_energy = (
+                    random.gauss(particle.E[i], self.sigma) / particle.E[i]
+                )
+
+                output_px[i] = random_measure_momentum * particle.px[i]
+                output_py[i] = random_measure_momentum * particle.py[i]
+                output_pz[i] = random_measure_momentum * particle.pz[i]
 
-                output_E[i] = random_measure_energy*particle.E[i]
+                output_E[i] = random_measure_energy * particle.E[i]
 
                 """output_px[i] = random.gauss(particle.px[i], self.sigma)
                 output_py[i] = random.gauss(particle.py[i], self.sigma)
                 output_pz[i] = random.gauss(particle.pz[i], self.sigma)
                 output_E[i] = random.gauss(particle.E[i], self.sigma)
                 """
-                mass =  (output_E[i]**2  - (output_px[i]**2 + output_py[i]**2 +output_pz[i]**2 ))
+                mass = output_E[i] ** 2 - (
+                    output_px[i] ** 2 + output_py[i] ** 2 + output_pz[i] ** 2
+                )
                 print(mass)
             particle_output = Particle(output_E, output_px, output_py, output_pz)
         return particle_output

pymcabc/feynman_diagram.py

@@ -31,11 +31,11 @@ def __init__(self):
                     return 0
         else:
             if channel == "s":
-                    self.s_chan()
+                self.s_chan()
             if channel == "t":
-                    self.t_chan()
+                self.t_chan()
             if channel == "u":
-                    self.u_chan()
+                self.u_chan()
             else:
                 print("Possible channels: s, t, and u")
                 return

pymcabc/identify_process.py

@@ -41,8 +41,6 @@ class DefineProcess:
         mC (float): mass of particle C
         Ecm (float): center of mass energy
         channel (str): optional, use to study effect a particular channel
-
-
     """
 
     def __init__(
@@ -63,8 +61,6 @@ def __init__(
             mC (float): mass of particle C
             Ecm (float): center of mass energy
             channel (str): optional, use to study effect a particular channel
-
-
         """
 
         build_json()
@@ -75,11 +71,16 @@ def __init__(
         self.mB = mB
         self.mC = mC
         self.Ecm = Ecm
+        if  self.mA<0 or self.mB<0 or self.mC < 0:
+            raise Exception("Negative masses not accepted")
+        if self.Ecm < 0:
+            raise Exception("Negative center of mass energy not accepted")
         self.library["mA"].append(mA)
         self.library["mB"].append(mB)
         self.library["mC"].append(mC)
         self.library["channel"].append(channel)
         self.process()
+        self.channel()
         self.masses()
         self.ECM()
         self.identify_mediator()
@@ -105,6 +106,19 @@ def process(self):
             json.dump(self.library, f)
         return None
 
+    def channel(self):
+        process_type = self.library["process_type"][0]
+        channel = self.library["channel"][0]
+        if channel != "none":
+            if channel not in process_type:
+                raise Exception(
+                    "Channel "
+                    + channel
+                    + " not available for process type "
+                    + process_type
+                )
+        return None
+
     def masses(self):
         """assign masses to m1, m2, m3, m4 and mediator"""
         string = self.input_string.replace(" > ", " ")
@@ -202,5 +216,4 @@ def identify_decay(self):
             self.library["decay_process"].append("NaN")
         with open("library.json", "w") as f:
             json.dump(self.library, f)
-
-        return None
+        return None

pymcabc/particle.py

@@ -67,9 +67,11 @@ def mass(self):
             else:
                 x = math.sqrt(x)
         except:
-            x = [0]*len(self.px)
+            x = [0] * len(self.px)
             for i in range(len(x)):
-                x[i] = self.E[i]**2 - self.px[i]**2 - self.py[i]**2 - self.pz[i]**2
+                x[i] = (
+                    self.E[i] ** 2 - self.px[i] ** 2 - self.py[i] ** 2 - self.pz[i] ** 2
+                )
                 print(x[i])
                 if x[i] < 0:
                     x[i] = 0
@@ -91,26 +93,25 @@ def boost(self, other):
 
          other is used to boost
         """
-        new = Particle(0., 0., 0., 0.)
-        new_other = Particle(0., 0., 0., 0.)
-        
+        new = Particle(0.0, 0.0, 0.0, 0.0)
+        new_other = Particle(0.0, 0.0, 0.0, 0.0)
+
         new_other.set4momenta(
             other.E, other.px / other.E, other.py / other.E, other.pz / other.E
         )
         beta = new_other.p2()
         gamma = 1.0 / math.sqrt(1.0 - beta)
 
-        if beta>0:
+        if beta > 0:
             gamma_2 = (gamma - 1.0) / beta
         else:
             gamma_2 = 0.0
 
-
         dotproduct = (
             self.px * new_other.px + self.py * new_other.py + self.pz * new_other.pz
         )
         new.px = self.px + (gamma_2 * dotproduct + gamma * self.E) * new_other.px
         new.py = self.py + (gamma_2 * dotproduct + gamma * self.E) * new_other.py
         new.pz = self.pz + (gamma_2 * dotproduct + gamma * self.E) * new_other.pz
         new.E = gamma * (self.E + dotproduct)
-        return new
+        return new

pymcabc/plotting.py

@@ -16,7 +16,7 @@ def plot(data, key):
         label = key.replace("_", " ")
         plt.xlabel(label)
         plt.ylabel("Events")
-        #plt.show()
+        # plt.show()
         plt.savefig(key + ".png")
 
     def file(filename="ABC_events.root"):