CUDA Version of the Formal Integral (#1837)

* Current work on the formal integral

* Work on cuda

* Work on the formal integral

* Work on cuda

* Removed units on inu when inputted to the formal integral

* Modified the trapz to take the correct input, and removed all comments and atomic add

* Changes to formal_integral to try and find the differences

* Uncommented a part of code regarding Jkkp, reduced error from 24% to 12%

* Removed random ;

* Initial unit tests for formal integral

* Added more tests, found error

* Fixed calculate_z to be the correct equation

* Made most functions test the number version against cuda

* First stage of adding full formal integral test

* Fixed test parameters as well as condensed them

* Remove unnecessary imports, improve tests

* Fixed call to points

* Slight changes trying to debug memory error

* First working test

* Implemented working CUDA formal_integral

* Clean up formal_integral and tests

* First idea for testing for GPUs

* Fixed name of dependency

* First method for automatically applying GPU, add dependency

* Add check for GPU for pytest tests

* Cleaning up the formal integral

* Remove outdated file

* Removed dependency, changed method for GPU detection

* Refactor to remove extra code

* Removed unnecessary imports

* Updates to test to match refactor

* Clean up imports

* Re-add pdb

* Removed wrapper for non-cuda test

* Black formatting

* Make docstrings match

* Updates to docstrings

* Add docstrings and update tests

* added docstring

* Fixed tests

* Add docstring to bound error

Co-authored-by: Kevin Cawley <[email protected]>

tardis/montecarlo/montecarlo_numba/formal_integral.py

@@ -18,6 +18,8 @@
     NumbaModel,
     NumbaPlasma,
 )
+from numba import cuda
+from tardis.montecarlo.montecarlo_numba.formal_integral_cuda import CudaFormalIntegrator
 
 from tardis.montecarlo.spectrum import TARDISSpectrum
 
@@ -225,7 +227,9 @@ def formal_integral(
         electron_density,
         N,
     ):
-        """simple wrapper for the numba implementation of the formal integral"""
+        """
+        Simple wrapper for the numba implementation of the formal integral
+        """
         return numba_formal_integral(
             self.model,
             self.plasma,
@@ -269,10 +273,14 @@ def generate_numba_objects(self):
         self.numba_plasma = numba_plasma_initialize(
             self.original_plasma, self.runner.line_interaction_type
         )
-
-        self.numba_integrator = NumbaFormalIntegrator(
-            self.numba_model, self.numba_plasma, self.points
-        )
+        if cuda.is_available():
+            self.integrator = CudaFormalIntegrator(
+                self.numba_model, self.numba_plasma, self.points
+            )
+        else:
+            self.integrator = NumbaFormalIntegrator(
+                self.numba_model, self.numba_plasma, self.points
+            )
 
     def check(self, raises=True):
         """
@@ -531,7 +539,7 @@ def formal_integral(self, nu, N):
         Jblue_lu = res[2].flatten(order="F")
 
         self.generate_numba_objects()
-        L = self.numba_integrator.formal_integral(
+        L = self.integrator.formal_integral(
             self.model.t_inner,
             nu,
             nu.shape[0],
@@ -612,7 +620,7 @@ def calculate_z(r, p, inv_t):
         return 0
 
 
-class BoundsError(ValueError):
+class BoundsError(IndexError):
     pass