Made new classes for updated geometry

tardis/montecarlo/montecarlo_numba/numba_interface.py

@@ -21,6 +21,27 @@
     ("v_outer", float64[:]),
 ]
 
+geometry_grid_spherical_1d_spec = [
+    ("r_inner", float64[:]),
+    ("r_outer", float64[:]),
+    ("time_explosion", float64),
+]
+
+geometry_grid_cartesian_2d_spec = [
+    ("x_coord", float64[:]),
+    ("y_coord", float64[:]),
+    ("time_explosion", float64),
+    ("cell_width", float64),
+]
+
+geometry_grid_cartesian_3d_spec = [
+    ("x_coord", float64[:]),
+    ("y_coord", float64[:]),
+    ("z_coord", float64[:]),
+    ("time_explosion", float64),
+    ("cell_width", float64),
+]
+
 
 @jitclass(numba_model_spec)
 class NumbaModel(object):
@@ -43,6 +64,65 @@ def __init__(self, r_inner, r_outer, v_inner, v_outer, time_explosion):
         self.time_explosion = time_explosion
 
 
+@jitclass(geometry_grid_spherical_1d_spec)
+class GeometryGridSpherical1D(object):
+    def __init__(self, r_inner, r_outer):
+        """
+        1D spherical shell model
+
+        Parameters
+        ----------
+        r_inner : numpy.ndarray
+        r_outer : numpy.ndarray
+        """
+        self.r_inner = r_inner
+        self.r_outer = r_outer
+
+
+@jitclass(geometry_grid_cartesian_2d_spec)
+class GeometryGridCartesian2D(object):
+    def __init__(self, x_coord, y_coord, cell_width):
+        """
+        Model for the 2D Cartesian grid system
+
+        Parameters
+        ----------
+        x_coord : numpy.ndarray
+        y_coord : numpy.ndarray
+        cell_width : float
+
+        Notes
+        -----
+        coords specify upper left coordinate of each cell
+        """
+        self.x_coord = x_coord
+        self.y_coord = y_coord
+        self.cell_width = cell_width
+
+
+@jitclass(geometry_grid_cartesian_3d_spec)
+class GeometryGridCartesian3D(object):
+    def __init__(self, x_coord, y_coord, z_coord, cell_width):
+        """
+        Model for the 3D Cartesian grid system
+
+        Parameters
+        ----------
+        x_coord : numpy.ndarray
+        y_coord : numpy.ndarray
+        z_coord : numpy.ndarray
+        cell_width : float
+
+        Notes
+        -----
+        coords specify upper left, out of the page coordinate of each cell
+        """
+        self.x_coord = x_coord
+        self.y_coord = y_coord
+        self.z_coord = z_coord
+        self.cell_width = cell_width
+
+
 numba_plasma_spec = [
     ("electron_density", float64[:]),
     ("t_electrons", float64[:]),

tardis/montecarlo/montecarlo_numba/tests/conftest.py

@@ -6,7 +6,12 @@
 
 from tardis.simulation import Simulation
 from tardis.montecarlo.montecarlo_numba import RPacket, PacketCollection
-from tardis.montecarlo.montecarlo_numba.numba_interface import Estimators
+from tardis.montecarlo.montecarlo_numba.numba_interface import (
+    Estimators,
+    GeometryGridCartesian2D,
+    GeometryGridSpherical1D,
+    GeometryGridCartesian3D,
+)
 
 
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
@@ -58,6 +63,37 @@ def verysimple_numba_model(nb_simulation_verysimple):
     )
 
 
+@pytest.fixture(scope="package")
+def verysimple_geometry_grid_spherical_1d(nb_simulation_verysimple):
+    runner = nb_simulation_verysimple.runner
+    model = nb_simulation_verysimple.model
+    return GeometryGridSpherical1D(
+        runner.r_inner_cgs,
+        runner.r_outer_cgs,
+    )
+
+
+@pytest.fixture(scope="package")
+def verysimple_geometry_grid_cartesian_2d(two_d_grid_verysimple):
+
+    width = 1.0
+    x_list = np.arange(-10, 10, width, dtype=np.float64)
+    y_list = np.arange(-10, 10, width, dtype=np.float64)
+
+    return GeometryGridCartesian2D(x_list, y_list, width)
+
+
+@pytest.fixture(scope="package")
+def verysimple_geometry_grid_cartesian_3d(nb_simulation_verysimple):
+
+    width = 1.0
+    x_list = np.arange(-10, 10, width, dtype=np.float64)
+    y_list = np.arange(-10, 10, width, dtype=np.float64)
+    z_list = np.arange(-10, 10, width, dtype=np.float64)
+
+    return GeometryGridCartesian3D(x_list, y_list, z_list, width)
+
+
 @pytest.fixture(scope="package")
 def verysimple_estimators(nb_simulation_verysimple):
     runner = nb_simulation_verysimple.runner