integrate radiation field

tardis/model/base.py

@@ -23,6 +23,8 @@
 from tardis.model.density import HomologousDensity
 from tardis.montecarlo.packet_source import BlackBodySimpleSource
 
+from tardis.radiation_field.base import MonteCarloRadiationFieldState
+
 logger = logging.getLogger(__name__)
 
 
@@ -290,24 +292,37 @@ def __init__(
             lambda_wien_inner = (
                 constants.b_wien / self.blackbody_packet_source.temperature
             )
-            self._t_radiative = constants.b_wien / (
+            t_radiative = constants.b_wien / (
                 lambda_wien_inner
                 * (1 + (self.v_middle - self.v_boundary_inner) / constants.c)
             )
+        elif len(t_radiative) != self.no_of_shells:
+            t_radiative = t_radiative[
+                self.v_boundary_inner_index
+                + 1 : self.v_boundary_outer_index
+                + 1
+            ]
         else:
-            # self._t_radiative = t_radiative[self.v_boundary_inner_index + 1:self.v_boundary_outer_index]
-            self._t_radiative = t_radiative
+            assert len(t_radiative) == self.no_of_shells
 
         if dilution_factor is None:
-            self._dilution_factor = 0.5 * (
+            dilution_factor = 0.5 * (
                 1
                 - np.sqrt(
                     1 - (self.r_inner[0] ** 2 / self.r_middle**2).to(1).value
                 )
             )
-        else:
-            # self.dilution_factor = dilution_factor[self.v_boundary_inner_index + 1:self.v_boundary_outer_index]
-            self._dilution_factor = dilution_factor
+        elif len(dilution_factor) != self.no_of_shells:
+            dilution_factor = dilution_factor[
+                self.v_boundary_inner_index
+                + 1 : self.v_boundary_outer_index
+                + 1
+            ]
+            assert len(dilution_factor) == self.no_of_shells
+
+        self.radiation_field = MonteCarloRadiationFieldState(
+            t_radiative, dilution_factor, None, None
+        )
 
     @property
     def w(self):
@@ -335,41 +350,12 @@ def t_inner(self, value):
 
     @property
     def dilution_factor(self):
-        if len(self._dilution_factor) == self.no_of_shells:
-            return self._dilution_factor
-
-        #        if self.v_boundary_inner in self.raw_velocity:
-        #            v_inner_ind = np.argwhere(self.raw_velocity == self.v_boundary_inner)[0][0]
-        #        else:
-        #            v_inner_ind = np.searchsorted(self.raw_velocity, self.v_boundary_inner) - 1
-        #        if self.v_boundary_outer in self.raw_velocity:
-        #            v_outer_ind = np.argwhere(self.raw_velocity == self.v_boundary_outer)[0][0]
-        #        else:
-        #            v_outer_ind = np.searchsorted(self.raw_velocity, self.v_boundary_outer)
-
-        return self._dilution_factor[
-            self.v_boundary_inner_index + 1 : self.v_boundary_outer_index + 1
-        ]
+        return self.radiation_field.dilution_factor
 
     @dilution_factor.setter
     def dilution_factor(self, value):
-        if len(value) == len(self._dilution_factor):
-            self._dilution_factor = value
-        elif len(value) == self.no_of_shells:
-            #            if self.v_boundary_inner in self.raw_velocity:
-            #                v_inner_ind = np.argwhere(self.raw_velocity == self.v_boundary_inner)[0][0]
-            #            else:
-            #                v_inner_ind = np.searchsorted(self.raw_velocity, self.v_boundary_inner) - 1
-            #            if self.v_boundary_outer in self.raw_velocity:
-            #                v_outer_ind = np.argwhere(self.raw_velocity == self.v_boundary_outer)[0][0]
-            #            else:
-            #                v_outer_ind = np.searchsorted(self.raw_velocity, self.v_boundary_outer)
-            #            assert v_outer_ind - v_inner_ind == self.no_of_shells, "trad shape different from number of shells"
-            self._dilution_factor[
-                self.v_boundary_inner_index
-                + 1 : self.v_boundary_outer_index
-                + 1
-            ] = value
+        if len(value) == self.no_of_shells:
+            self.radiation_field.dilution_factor = value
         else:
             raise ValueError(
                 "Trying to set dilution_factor for unmatching number"
@@ -378,44 +364,15 @@ def dilution_factor(self, value):
 
     @property
     def t_radiative(self):
-        if len(self._t_radiative) == self.no_of_shells:
-            return self._t_radiative
-
-        #        if self.v_boundary_inner in self.raw_velocity:
-        #            v_inner_ind = np.argwhere(self.raw_velocity == self.v_boundary_inner)[0][0]
-        #        else:
-        #            v_inner_ind = np.searchsorted(self.raw_velocity, self.v_boundary_inner) - 1
-        #        if self.v_boundary_outer in self.raw_velocity:
-        #            v_outer_ind = np.argwhere(self.raw_velocity == self.v_boundary_outer)[0][0]
-        #        else:
-        #            v_outer_ind = np.searchsorted(self.raw_velocity, self.v_boundary_outer)
-
-        return self._t_radiative[
-            self.v_boundary_inner_index + 1 : self.v_boundary_outer_index + 1
-        ]
+        return self.radiation_field.t_radiative
 
     @t_radiative.setter
     def t_radiative(self, value):
-        if len(value) == len(self._t_radiative):
-            self._t_radiative = value
-        elif len(value) == self.no_of_shells:
-            #            if self.v_boundary_inner in self.raw_velocity:
-            #                v_inner_ind = np.argwhere(self.raw_velocity == self.v_boundary_inner)[0][0]
-            #            else:
-            #                v_inner_ind = np.searchsorted(self.raw_velocity, self.v_boundary_inner) - 1
-            #            if self.v_boundary_outer in self.raw_velocity:
-            #                v_outer_ind = np.argwhere(self.raw_velocity == self.v_boundary_outer)[0][0]
-            #            else:
-            #                v_outer_ind = np.searchsorted(self.raw_velocity, self.v_boundary_outer)
-            #            assert v_outer_ind - v_inner_ind == self.no_of_shells, "trad shape different from number of shells"
-            self._t_radiative[
-                self.v_boundary_inner_index
-                + 1 : self.v_boundary_outer_index
-                + 1
-            ] = value
+        if len(value) == self.no_of_shells:
+            self.radiation_field.t_radiative = value
         else:
             raise ValueError(
-                "Trying to set t_radiative for unmatching number" "of shells."
+                "Trying to set t_radiative for unmatching number of shells."
             )
 
     @property
@@ -567,26 +524,6 @@ def v_boundary_outer_index(self):
             )
         return v_outer_ind
 
-    #    @property
-    #    def v_boundary_inner_index(self):
-    #        if self.v_boundary_inner <= self.raw_velocity[0]:
-    #            return 0
-    #        else:
-    #            idx = max(0,
-    #                      self.raw_velocity.searchsorted(self.v_boundary_inner) - 1)
-    #            # check for zero volume of designated first cell
-    #            if np.isclose(self.v_boundary_inner, self.raw_velocity[idx + 1],
-    #                          atol=1e-8 * u.km / u.s) and (self.v_boundary_inner <=
-    #                                                           self.raw_velocity[idx + 1]):
-    #                idx += 1
-    #            return idx
-    #
-    #    @property
-    #    def v_boundary_outer_index(self):
-    #        if self.v_boundary_outer >= self.raw_velocity[-1]:
-    #            return None
-    #        return self.raw_velocity.searchsorted(self.v_boundary_outer) + 1
-
     @classmethod
     def from_config(cls, config, atom_data=None):
         """

tardis/radiation_field/base.py

@@ -0,0 +1,34 @@
+import numpy as np
+from astropy import units as u
+
+from tardis.montecarlo.packet_source import BasePacketSource
+from tardis.montecarlo.montecarlo_numba.numba_interface import OpacityState
+
+
+class MonteCarloRadiationFieldState:
+    """_summary_
+
+    Parameters
+    ----------
+    t_radiative : u.Quantity
+        Radiative temperature in each shell
+    dilution_factor : numpy.ndarray
+        Dilution Factors in each shell
+    opacities : OpacityState
+        Opacity container object
+    packet_source : SourceFunction
+        Source function for radiative transfer, for example a packet_source
+    """
+
+    def __init__(
+        self,
+        t_radiative: u.Quantity,
+        dilution_factor: np.ndarray,
+        opacities: OpacityState,
+        packet_source: BasePacketSource,
+    ):
+        self.t_radiative = t_radiative
+        self.dilution_factor = dilution_factor
+        self.t_rad = self.t_radiative
+        self.opacities = opacities
+        self.packet_source = packet_source

tardis/radiation_field/opacities/__init__.py

@@ -0,0 +1 @@
+from tardis.radiation_field.opacities.base import *