Intensity p rays (#1933)

* formal integral base functions now return I_nu per p-ray.  Now attaches this value to the runner.  Tests updated accordingly such that both the CPU and GPU implimentations should pass (this still needs to be tested though)

* Fixed a few small logical errors for testing the intensity p-rays.  Added the p-ray values themselves to the runner.  All tests pass now

* Added the p-rays to the runner

* Added docstrings specifiying new outputs of the formal integral base functions, removed old commented out code

* Reformatted files with black for PEP 8 consistency

* Restyled the intensity array assert

tardis/montecarlo/montecarlo_numba/formal_integral.py

@@ -18,7 +18,9 @@
     NumbaModel,
     NumbaPlasma,
 )
-from tardis.montecarlo.montecarlo_numba.formal_integral_cuda import CudaFormalIntegrator
+from tardis.montecarlo.montecarlo_numba.formal_integral_cuda import (
+    CudaFormalIntegrator,
+)
 
 from tardis.montecarlo.spectrum import TARDISSpectrum
 
@@ -48,6 +50,13 @@ def numba_formal_integral(
 ):
     """
     model, plasma, and estimator are the numba variants
+    Returns
+    -------
+    L : float64 array
+        integrated luminosities
+    I_nu_p : float64 2D array
+        intensities at each p-ray multiplied by p
+        frequency x p-ray grid
     """
     # todo: add all the original todos
     # Initialize the output which is shared among threads
@@ -64,8 +73,9 @@ def numba_formal_integral(
     line_list_nu = plasma.line_list_nu
     # done with instantiation
     # now loop over wavelength in spectrum
+    I_nu_p = np.zeros((inu_size, N), dtype=np.float64)
     for nu_idx in prange(inu_size):
-        I_nu = np.zeros(N, dtype=np.float64)
+        I_nu = I_nu_p[nu_idx]
         z = np.zeros(2 * size_shell, dtype=np.float64)
         shell_id = np.zeros(2 * size_shell, dtype=np.int64)
         offset = 0
@@ -191,17 +201,17 @@ def numba_formal_integral(
             I_nu[p_idx] *= p
         L[nu_idx] = 8 * M_PI * M_PI * trapezoid_integration(I_nu, R_max / N)
 
-    return L
+    return L, I_nu_p
 
 
-#integrator_spec = [
+# integrator_spec = [
 #    ("model", NumbaModel.class_type.instance_type),
 #    ("plasma", NumbaPlasma.class_type.instance_type),
 #    ("points", int64),
-#]
+# ]
 
 
-#@jitclass(integrator_spec)
+# @jitclass(integrator_spec)
 class NumbaFormalIntegrator(object):
     """
     Helper class for performing the formal integral
@@ -246,15 +256,15 @@ def formal_integral(
 
 class FormalIntegrator(object):
     """
-    Class containing the formal integrator. 
-    
-    If there is a NVIDIA CUDA GPU available, 
-    the formal integral will automatically run 
-    on it. If multiple GPUs are available, it will 
-    choose the first one that it sees. You can 
-    read more about selecting different GPUs on 
+    Class containing the formal integrator.
+
+    If there is a NVIDIA CUDA GPU available,
+    the formal integral will automatically run
+    on it. If multiple GPUs are available, it will
+    choose the first one that it sees. You can
+    read more about selecting different GPUs on
     Numba's CUDA documentation.
-    
+
     Parameters
     ----------
     model : tardis.model.Radial1DModel
@@ -383,9 +393,9 @@ def make_source_function(self):
         model = self.model
         runner = self.runner
 
-        #macro_ref = self.atomic_data.macro_atom_references
+        # macro_ref = self.atomic_data.macro_atom_references
         macro_ref = self.atomic_data.macro_atom_references
-        #macro_data = self.atomic_data.macro_atom_data
+        # macro_data = self.atomic_data.macro_atom_data
         macro_data = self.original_plasma.macro_atom_data
 
         no_lvls = len(self.levels_index)
@@ -555,7 +565,7 @@ def formal_integral(self, nu, N):
         Jblue_lu = res[2].flatten(order="F")
 
         self.generate_numba_objects()
-        L = self.integrator.formal_integral(
+        L, I_nu_p = self.integrator.formal_integral(
             self.model.t_inner,
             nu,
             nu.shape[0],
@@ -566,6 +576,24 @@ def formal_integral(self, nu, N):
             self.runner.electron_densities_integ,
             N,
         )
+        R_max = self.runner.r_outer_i[-1]
+        ps = calculate_p_values(R_max, N)[None, :]
+        I_nu_p[:, 1:] /= ps[:, 1:]
+        self.runner.I_nu_p = I_nu_p
+        self.runner.p_rays = ps
+
+        I_nu = self.runner.I_nu_p * ps
+        L_test = np.array(
+            [
+                8 * M_PI * M_PI * trapezoid_integration((I_nu)[i, :], R_max / N)
+                for i in range(nu.shape[0])
+            ]
+        )
+        error = np.max(np.abs((L_test - L) / L))
+        assert (
+            error < 1e-7
+        ), f"Incorrect I_nu_p values, max relative difference:{error}"
+
         return np.array(L, np.float64)
 
 

tardis/montecarlo/montecarlo_numba/formal_integral_cuda.py

@@ -283,7 +283,7 @@ def formal_integral(
             shell_id,
         )
 
-        return L
+        return L, I_nu
 
 
 @cuda.jit(device=True)
@@ -375,6 +375,7 @@ class BoundsError(IndexError):
     Used to check bounds in reverse
     binary search
     """
+
     pass
 
 

tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py

@@ -66,9 +66,7 @@ def black_body_caller(nu, T, actual):
 @pytest.mark.skipif(
     not GPUs_available, reason="No GPU is available to test CUDA function"
 )
-@pytest.mark.parametrize(
-    "N", (1e2, 1e3, 1e4, 1e5)
-)
+@pytest.mark.parametrize("N", (1e2, 1e3, 1e4, 1e5))
 def test_trapezoid_integration_cuda(N):
     """
     Initializes the test of the cuda version
@@ -85,10 +83,10 @@ def test_trapezoid_integration_cuda(N):
 
     expected = formal_integral_numba.trapezoid_integration(data, h)
     trapezoid_integration_caller[1, 1](data, h, actual)
-    
-    #This is 1e-13, as more points are added to the integration
-    #there will be more floating point error due to the difference
-    #in how the trapezoid integration is called. 
+
+    # This is 1e-13, as more points are added to the integration
+    # there will be more floating point error due to the difference
+    # in how the trapezoid integration is called.
     ntest.assert_allclose(actual[0], expected, rtol=1e-13)
 
 
@@ -102,7 +100,6 @@ def trapezoid_integration_caller(data, h, actual):
     actual[x] = formal_integral_cuda.trapezoid_integration_cuda(data, h)
 
 
-
 TESTDATA_model = [
     {
         "r": np.linspace(1, 2, 3, dtype=np.float64),
@@ -375,7 +372,7 @@ def test_full_formal_integral(
         formal_integrator_cuda.runner.tau_sobolevs_integ,
         formal_integrator_cuda.runner.electron_densities_integ,
         formal_integrator_cuda.points,
-    )
+    )[0]
 
     L_numba = formal_integrator_numba.integrator.formal_integral(
         formal_integrator_numba.model.t_inner,
@@ -387,6 +384,6 @@ def test_full_formal_integral(
         formal_integrator_numba.runner.tau_sobolevs_integ,
         formal_integrator_numba.runner.electron_densities_integ,
         formal_integrator_numba.points,
-    )
+    )[0]
 
     ntest.assert_allclose(L_cuda, L_numba, rtol=1e-14)