Merge pull request #105 from SNEWS2/vec-nu-decay

Vectorize NeutrinoDecay
SNEWS2 · Sep 13, 2021 · 7bd54a4 · 7bd54a4
2 parents 74ee03c + 27162b2
commit 7bd54a4
Showing 1 changed file with 8 additions and 26 deletions.
diff --git a/python/snewpy/flavor_transformation.py b/python/snewpy/flavor_transformation.py
@@ -1144,23 +1144,14 @@ def prob_ee(self, t, E):
         prob : float or ndarray
             Transition probability.
         """
-        pe_array = []
-
         # NMO case.
         if self.mass_order == MassHierarchy.NORMAL:
-            for energy in E:
-                pe_array.append(
-                    self.De1*(1-np.exp(-self.gamma(energy)*self.d)) +
-                    self.De3*np.exp(-self.gamma(energy)*self.d))
-            pe_array = np.array(pe_array)
+            pe_array = self.De1*(1-np.exp(-self.gamma(E)*self.d)) + \
+                       self.De3*np.exp(-self.gamma(E)*self.d)
         # IMO case.
         else:
-            for energy in E:
-                pe_array.append(
-                    self.De2*np.exp(-self.gamma(energy)*self.d) +
-                    self.De3*(1-np.exp(-self.gamma(energy)*self.d)))
-            pe_array = np.array(pe_array)
-
+            pe_array = self.De2*np.exp(-self.gamma(E)*self.d) + \
+                       self.De3*(1-np.exp(-self.gamma(E)*self.d))
         return pe_array
 
     def prob_ex(self, t, E):
@@ -1251,23 +1242,14 @@ def prob_exbar(self, t, E):
         prob : float or ndarray
             Transition probability.
         """
-        pxbar_array = []
-
         # NMO case.
         if self.mass_order == MassHierarchy.NORMAL:
-            for energy in E:
-                pxbar_array.append(
-                    self.De1*(1-np.exp(-self.gamma(energy)*self.d)) +
-                    self.De2 + self.De3*np.exp(-self.gamma(energy)*self.d))
-            pxbar_array = np.array(pxbar_array)
+            pxbar_array = self.De1*(1-np.exp(-self.gamma(E)*self.d)) + \
+                          self.De2 + self.De3*np.exp(-self.gamma(E)*self.d)
         # IMO case.
         else:
-            for energy in E:
-                pxbar_array.append(
-                    self.De1 + self.De2*np.exp(-self.gamma(energy)*self.d) +
-                    self.De3*(1-np.exp(-self.gamma(energy)*self.d)))
-            pxbar_array = np.array(pxbar_array)
-
+            pxbar_array = self.De1 + self.De2*np.exp(-self.gamma(E)*self.d) + \
+                          self.De3*(1-np.exp(-self.gamma(E)*self.d))
         return pxbar_array
 
     def prob_xxbar(self, t, E):