Restructuring NumbaModel (#2136)

* Added NumbaRadial1DGeometry

* Added to_numba function

* Fixed incorrect naming of numba spec

* Updated docstring

* Updated to_numba docstring

* Updated docstring

* Updated numba model to contain only time explosion

* updated docstring formatting error

* Moved geometry to new geometry sub package

* Updated all functions where numba model was used

* Updated to specify units

* converting time_explosion to seconds specifically

* Added numba geometry fixture

* Updated tests to work with numba geometry

* Updated formatting

* Fixed issue with geometry initialization

* Fixed formatting

tardis/model/base.py

@@ -6,6 +6,7 @@
 from tardis import constants
 import radioactivedecay as rd
 from radioactivedecay.utils import Z_DICT
+from tardis.model.geometry.radial1d import Radial1DGeometry
 
 from tardis.util.base import quantity_linspace, is_valid_nuclide_or_elem
 from tardis.io.parsers.csvy import load_csvy
@@ -24,35 +25,6 @@
 logger = logging.getLogger(__name__)
 
 
-class Radial1DGeometry:
-    """
-    Holds information about model geometry for radial 1D models.
-
-    Parameters
-    ----------
-    r_inner : astropy.units.quantity.Quantity
-    r_outer : astropy.units.quantity.Quantity
-    v_inner : astropy.units.quantity.Quantity
-    v_outer : astropy.units.quantity.Quantity
-
-    Attributes
-    ----------
-    volume : astropy.units.quantity.Quantity
-        Volume in each shell
-    """
-
-    def __init__(self, r_inner, r_outer, v_inner, v_outer):
-        self.r_inner = r_inner
-        self.r_outer = r_outer
-        self.v_inner = v_inner
-        self.v_outer = v_outer
-
-    @property
-    def volume(self):
-        """Volume in shell computed from r_outer and r_inner"""
-        return (4.0 / 3) * np.pi * (self.r_outer**3 - self.r_inner**3)
-
-
 class Composition:
     """
     Holds information about model composition
@@ -113,7 +85,7 @@ class ModelState:
     Parameters
     ----------
     composition : tardis.model.Composition
-    geometry : tardis.model.Radial1DGeometry
+    geometry : tardis.model.geometry.radial1d.Radial1DGeometry
     time_explosion : astropy.units.quantity.Quantity
 
     Attributes

tardis/model/geometry/radial1d.py

@@ -0,0 +1,76 @@
+from numba import float64
+from numba.experimental import jitclass
+import numpy as np
+from astropy import units as u
+
+
+class Radial1DGeometry:
+    """
+    Holds information about model geometry for radial 1D models.
+
+    Parameters
+    ----------
+    r_inner : astropy.units.quantity.Quantity
+    r_outer : astropy.units.quantity.Quantity
+    v_inner : astropy.units.quantity.Quantity
+    v_outer : astropy.units.quantity.Quantity
+
+    Attributes
+    ----------
+    volume : astropy.units.quantity.Quantity
+        Volume in each shell
+    """
+
+    def __init__(self, r_inner, r_outer, v_inner, v_outer):
+        self.r_inner = r_inner
+        self.r_outer = r_outer
+        self.v_inner = v_inner
+        self.v_outer = v_outer
+
+    @property
+    def volume(self):
+        """Volume in shell computed from r_outer and r_inner"""
+        return (4.0 / 3) * np.pi * (self.r_outer**3 - self.r_inner**3)
+
+    def to_numba(self):
+        """
+        Returns a new NumbaRadial1DGeometry object
+
+        Returns
+        -------
+        NumbaRadial1DGeometry
+            Numba version of Radial1DGeometry with properties in cgs units
+        """
+        return NumbaRadial1DGeometry(
+            self.r_inner.to(u.cm).value,
+            self.r_outer.to(u.cm).value,
+            self.v_inner.to(u.cm / u.s).value,
+            self.v_outer.to(u.cm / u.s).value,
+        )
+
+
+numba_geometry_spec = [
+    ("r_inner", float64[:]),
+    ("r_outer", float64[:]),
+    ("v_inner", float64[:]),
+    ("v_outer", float64[:]),
+]
+
+
+@jitclass(numba_geometry_spec)
+class NumbaRadial1DGeometry(object):
+    def __init__(self, r_inner, r_outer, v_inner, v_outer):
+        """
+        Radial 1D Geometry for the Numba mode
+
+        Parameters
+        ----------
+        r_inner : numpy.ndarray
+        r_outer : numpy.ndarray
+        v_inner : numpy.ndarray
+        v_outer : numpy.ndarray
+        """
+        self.r_inner = r_inner
+        self.r_outer = r_outer
+        self.v_inner = v_inner
+        self.v_outer = v_outer

tardis/montecarlo/montecarlo_numba/base.py

@@ -49,13 +49,9 @@ def montecarlo_radial1d(
         runner._output_nu,
         runner._output_energy,
     )
-
+    numba_radial_1d_geometry = model.model_state.geometry.to_numba()
     numba_model = NumbaModel(
-        runner.r_inner_cgs,
-        runner.r_outer_cgs,
-        runner.v_inner_cgs,
-        runner.v_outer_cgs,
-        model.time_explosion.to("s").value,
+        model.model_state.time_explosion.to("s").value,
     )
     numba_plasma = numba_plasma_initialize(plasma, runner.line_interaction_type)
     estimators = Estimators(
@@ -89,6 +85,7 @@ def montecarlo_radial1d(
         rpacket_trackers,
     ) = montecarlo_main_loop(
         packet_collection,
+        numba_radial_1d_geometry,
         numba_model,
         numba_plasma,
         estimators,
@@ -139,6 +136,7 @@ def montecarlo_radial1d(
 @njit(**njit_dict)
 def montecarlo_main_loop(
     packet_collection,
+    numba_radial_1d_geometry,
     numba_model,
     numba_plasma,
     estimators,
@@ -158,8 +156,9 @@ def montecarlo_main_loop(
     Parameters
     ----------
     packet_collection : PacketCollection
+    numba_radial_1d_geometry : NumbaRadial1DGeometry
     numba_model : NumbaModel
-        numba_plasma : NumbaPlasma
+    numba_plasma : NumbaPlasma
     estimators : NumbaEstimators
     spectrum_frequency : astropy.units.Quantity
         frequency binspas
@@ -247,7 +246,7 @@ def montecarlo_main_loop(
         seed = packet_seeds[i]
         np.random.seed(seed)
         r_packet = RPacket(
-            numba_model.r_inner[0],
+            numba_radial_1d_geometry.r_inner[0],
             packet_collection.packets_input_mu[i],
             packet_collection.packets_input_nu[i],
             packet_collection.packets_input_energy[i],
@@ -260,6 +259,7 @@ def montecarlo_main_loop(
 
         loop = single_packet_loop(
             r_packet,
+            numba_radial_1d_geometry,
             numba_model,
             numba_plasma,
             estimators,

tardis/montecarlo/montecarlo_numba/formal_integral.py

@@ -36,6 +36,7 @@ class IntegrationError(Exception):
 
 @njit(**njit_dict)
 def numba_formal_integral(
+    geometry,
     model,
     plasma,
     iT,
@@ -66,8 +67,8 @@ def numba_formal_integral(
     # global read-only values
     size_line, size_shell = tau_sobolev.shape
     size_tau = size_line * size_shell
-    R_ph = model.r_inner[0]  # make sure these are cgs
-    R_max = model.r_outer[size_shell - 1]
+    R_ph = geometry.r_inner[0]  # make sure these are cgs
+    R_max = geometry.r_outer[size_shell - 1]
     pp = np.zeros(N, dtype=np.float64)  # check
     exp_tau = np.zeros(size_tau, dtype=np.float64)
     exp_tau = np.exp(-tau_sobolev.T.ravel())  # maybe make this 2D?
@@ -108,7 +109,9 @@ def numba_formal_integral(
             p = pp[p_idx]
 
             # initialize z intersections for p values
-            size_z = populate_z(model, p, z, shell_id)  # check returns
+            size_z = populate_z(
+                geometry, model, p, z, shell_id
+            )  # check returns
             # initialize I_nu
             if p <= R_ph:
                 I_nu[p_idx] = intensity_black_body(nu * z[0], iT)
@@ -220,8 +223,8 @@ class NumbaFormalIntegrator(object):
     with numba.
     """
 
-    def __init__(self, model, plasma, points=1000):
-
+    def __init__(self, geometry, model, plasma, points=1000):
+        self.geometry = geometry
         self.model = model
         self.plasma = plasma
         self.points = points
@@ -242,6 +245,7 @@ def formal_integral(
         Simple wrapper for the numba implementation of the formal integral
         """
         return numba_formal_integral(
+            self.geometry,
             self.model,
             self.plasma,
             iT,
@@ -291,23 +295,33 @@ def __init__(self, model, plasma, runner, points=1000):
     def generate_numba_objects(self):
         """instantiate the numba interface objects
         needed for computing the formal integral"""
-        self.numba_model = NumbaModel(
+        from tardis.model.geometry.radial1d import NumbaRadial1DGeometry
+
+        self.numba_radial_1d_geometry = NumbaRadial1DGeometry(
             self.runner.r_inner_i,
             self.runner.r_outer_i,
             self.runner.r_inner_i / self.model.time_explosion.to("s").value,
             self.runner.r_outer_i / self.model.time_explosion.to("s").value,
-            self.model.time_explosion.to("s").value,
+        )
+        self.numba_model = NumbaModel(
+            self.model.time_explosion.cgs.value,
         )
         self.numba_plasma = numba_plasma_initialize(
             self.original_plasma, self.runner.line_interaction_type
         )
         if self.runner.use_gpu:
             self.integrator = CudaFormalIntegrator(
-                self.numba_model, self.numba_plasma, self.points
+                self.numba_radial_1d_geometry,
+                self.numba_model,
+                self.numba_plasma,
+                self.points,
             )
         else:
             self.integrator = NumbaFormalIntegrator(
-                self.numba_model, self.numba_plasma, self.points
+                self.numba_radial_1d_geometry,
+                self.numba_model,
+                self.numba_plasma,
+                self.points,
             )
 
     def check(self, raises=True):
@@ -602,7 +616,7 @@ def formal_integral(self, nu, N):
 
 
 @njit(**njit_dict_no_parallel)
-def populate_z(model, p, oz, oshell_id):
+def populate_z(geometry, model, p, oz, oshell_id):
     """Calculate p line intersections
 
     This function calculates the intersection points of the p-line with
@@ -615,13 +629,13 @@ def populate_z(model, p, oz, oshell_id):
         :oshell_id: (int64) will be set with the corresponding shell_ids
     """
     # abbreviations
-    r = model.r_outer
-    N = len(model.r_inner)  # check
+    r = geometry.r_outer
+    N = len(geometry.r_inner)  # check
     inv_t = 1 / model.time_explosion
     z = 0
     offset = N
 
-    if p <= model.r_inner[0]:
+    if p <= geometry.r_inner[0]:
         # intersect the photosphere
         for i in range(N):
             oz[i] = 1 - calculate_z(r[i], p, inv_t)

tardis/montecarlo/montecarlo_numba/formal_integral_cuda.py

@@ -41,11 +41,11 @@ def cuda_formal_integral(
     Parameters
     ----------
     r_inner : array(float64, 1d, C)
-        self.model.r_inner
+        self.geometry.r_inner
     r_outer : array(float64, 1d, C)
-        self.model.r_outer
+        self.geometry.r_outer
     time_explosion: float64
-        self.model.time_explosion
+        self.geometry.time_explosion
     line_list_nu : array(float64, 1d, A)
         self.plasma.line_list_nu
     iT : np.float64
@@ -216,8 +216,8 @@ class CudaFormalIntegrator(object):
     with CUDA.
     """
 
-    def __init__(self, model, plasma, points=1000):
-
+    def __init__(self, geometry, model, plasma, points=1000):
+        self.geometry = geometry
         self.model = model
         self.plasma = plasma
         self.points = points
@@ -246,7 +246,7 @@ def formal_integral(
         exp_tau = np.zeros(size_tau, dtype=np.float64)  # array(float64, 1d, C)
         exp_tau = np.exp(-tau_sobolev.T.ravel())  # array(float64, 1d, C)
         pp[::] = calculate_p_values(
-            self.model.r_outer[size_shell - 1], N
+            self.geometry.r_outer[size_shell - 1], N
         )  # array(float64, 1d, C)
 
         I_nu = np.zeros(
@@ -262,8 +262,8 @@ def formal_integral(
         blocks_per_grid = (inu_size // THREADS_PER_BLOCK) + 1
 
         cuda_formal_integral[blocks_per_grid, THREADS_PER_BLOCK](
-            self.model.r_inner,
-            self.model.r_outer,
+            self.geometry.r_inner,
+            self.geometry.r_outer,
             self.model.time_explosion,
             self.plasma.line_list_nu,
             iT.value,

tardis/montecarlo/montecarlo_numba/numba_interface.py

@@ -13,33 +13,22 @@
 
 C_SPEED_OF_LIGHT = const.c.to("cm/s").value
 
+
 numba_model_spec = [
-    ("r_inner", float64[:]),
-    ("r_outer", float64[:]),
     ("time_explosion", float64),
-    ("v_inner", float64[:]),
-    ("v_outer", float64[:]),
 ]
 
 
 @jitclass(numba_model_spec)
 class NumbaModel(object):
-    def __init__(self, r_inner, r_outer, v_inner, v_outer, time_explosion):
+    def __init__(self, time_explosion):
         """
         Model for the Numba mode
 
         Parameters
         ----------
-        r_inner : numpy.ndarray
-        r_outer : numpy.ndarray
-        v_inner : numpy.ndarray
-        v_outer : numpy.ndarray
         time_explosion : float
         """
-        self.r_inner = r_inner
-        self.r_outer = r_outer
-        self.v_inner = v_inner
-        self.v_outer = v_outer
         self.time_explosion = time_explosion