diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral.py b/tardis/montecarlo/montecarlo_numba/formal_integral.py index 70f82ef626d..404daec5249 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral.py +++ b/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) diff --git a/tardis/montecarlo/montecarlo_numba/formal_integral_cuda.py b/tardis/montecarlo/montecarlo_numba/formal_integral_cuda.py index 5ecb1d38444..13a1a47ebaa 100644 --- a/tardis/montecarlo/montecarlo_numba/formal_integral_cuda.py +++ b/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 diff --git a/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py b/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py index 054a444d0a6..90a36010d91 100644 --- a/tardis/montecarlo/montecarlo_numba/tests/test_cuda_formal_integral.py +++ b/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)