Merge pull request #18 from Rodot-/fast_formal_integral_subbranch

Removed unneccessary branching
Rodot- · Jul 22, 2024 · de3adfa · de3adfa
2 parents abfa976 + 59a69da
commit de3adfa
Showing 1 changed file with 29 additions and 20 deletions.
diff --git a/tardis/transport/montecarlo/formal_integral_cuda.py b/tardis/transport/montecarlo/formal_integral_cuda.py
@@ -167,9 +167,6 @@ def cuda_formal_integral(
     pJred_lu = int(offset + idx_nu_start)
     pJblue_lu = int(offset + idx_nu_start)
 
-    # flag for first contribution to integration on current p-ray
-    first = 1
-
     # loop over all interactions
     for i in range(size_z - 1):
         escat_op = electron_density[int(shell_id_thread[i])] * SIGMA_THOMSON
@@ -178,26 +175,38 @@ def cuda_formal_integral(
         )  # +1 is the offset as the original is from z[1:]
 
         nu_end_idx = line_search_cuda(line_list_nu, nu_end, len(line_list_nu))
+        escat_contrib += (
+            (zend - zstart)
+            * escat_op
+            * (Jblue_lu[pJblue_lu] - I_nu_thread[p_idx])
+        )
+        pJred_lu += 1
+        I_nu_thread[p_idx] += escat_contrib
+        # // Lucy 1999, Eq 26
+        I_nu_thread[p_idx] *= exp_tau[pexp_tau]
+        I_nu_thread[p_idx] += att_S_ul[patt_S_ul]
+
+        # // reset e-scattering opacity
+        escat_contrib = 0.0
+        zstart = zend
+
+        pline += 1
+        pexp_tau += 1
+        patt_S_ul += 1
+        pJblue_lu += 1
 
-        for _ in range(max(nu_end_idx - pline, 0)):
+        for _ in range(max(nu_end_idx - pline, 1) - 1):
             # calculate e-scattering optical depth to next resonance point
             zend = time_explosion / C_INV * (1.0 - line_list_nu[pline] / nu)
-            if first == 1:
-                escat_contrib += (
-                    (zend - zstart)
-                    * escat_op
-                    * (Jblue_lu[pJblue_lu] - I_nu_thread[p_idx])
-                )
-                first = 0
-            else:
-                # Account for e-scattering, c.f. Eqs 27, 28 in Lucy 1999
-                Jkkp = 0.5 * (Jred_lu[pJred_lu] + Jblue_lu[pJblue_lu])
-                escat_contrib += (
-                    (zend - zstart) * escat_op * (Jkkp - I_nu_thread[p_idx])
-                )
-                # this introduces the necessary offset of one element between
-                # pJblue_lu and pJred_lu
-                pJred_lu += 1
+
+            # Account for e-scattering, c.f. Eqs 27, 28 in Lucy 1999
+            Jkkp = 0.5 * (Jred_lu[pJred_lu] + Jblue_lu[pJblue_lu])
+            escat_contrib += (
+                (zend - zstart) * escat_op * (Jkkp - I_nu_thread[p_idx])
+            )
+            # this introduces the necessary offset of one element between
+            # pJblue_lu and pJred_lu
+            pJred_lu += 1
             I_nu_thread[p_idx] += escat_contrib
             # // Lucy 1999, Eq 26
             I_nu_thread[p_idx] *= exp_tau[pexp_tau]