Transferred simulation to use ModelState

Added shape property to Radial1DGeometry

Use ModelState in montecarlo
Use Modelstate in numba
Use ModelState in plasma
Create RadiationField to hold temp/dilution etc properties
Use ModelState in simulation

tardis/model/geometry/radial1d.py

@@ -32,6 +32,10 @@ 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)
 
+    @property
+    def shape(self):
+        return self.r_inner.shape
+
     def to_numba(self):
         """
         Returns a new NumbaRadial1DGeometry object

tardis/montecarlo/base.py

@@ -185,18 +185,19 @@ def _initialize_continuum_estimator_arrays(self, gamma_shape):
             gamma_shape, dtype=np.int64
         )
 
-    def _initialize_geometry_arrays(self, model):
+    def _initialize_geometry_arrays(self, model_state):
         """
         Generate the cgs like geometry arrays for the montecarlo part
 
         Parameters
         ----------
         model : model.Radial1DModel
         """
-        self.r_inner_cgs = model.r_inner.to("cm").value
-        self.r_outer_cgs = model.r_outer.to("cm").value
-        self.v_inner_cgs = model.v_inner.to("cm/s").value
-        self.v_outer_cgs = model.v_outer.to("cm/s").value
+        geometry = model_state.geometry
+        self.r_inner_cgs = geometry.r_inner.to("cm").value
+        self.r_outer_cgs = geometry.r_outer.to("cm").value
+        self.v_inner_cgs = geometry.v_inner.to("cm/s").value
+        self.v_outer_cgs = geometry.v_outer.to("cm/s").value
 
     def _initialize_packets(
         self, temperature, no_of_packets, iteration, radius, time_explosion
@@ -299,8 +300,9 @@ def integrator(self):
 
     def run(
         self,
-        model,
+        model_state,
         plasma,
+        radiation_field,
         no_of_packets,
         no_of_virtual_packets=0,
         iteration=0,
@@ -326,8 +328,10 @@ def run(
 
         set_num_threads(self.nthreads)
 
-        self.time_of_simulation = self.calculate_time_of_simulation(model)
-        self.volume = model.volume
+        self.time_of_simulation = self.calculate_time_of_simulation(
+            model_state, radiation_field
+        )
+        self.volume = model_state.geometry.volume
 
         # Initializing estimator array
         self._initialize_estimator_arrays(plasma.tau_sobolevs.shape)
@@ -339,27 +343,27 @@ def run(
 
         self._initialize_continuum_estimator_arrays(gamma_shape)
 
-        self._initialize_geometry_arrays(model)
+        self._initialize_geometry_arrays(model_state)
 
         self._initialize_packets(
-            model.t_inner.value,
+            radiation_field.t_inner.value,
             no_of_packets,
             iteration,
-            model.r_inner[0],
-            model.time_explosion,
+            model_state.geometry.r_inner[0],
+            model_state.time_explosion,
         )
 
-        configuration_initialize(self, no_of_virtual_packets)
+        configuration_initialize(self, no_of_virtual_packets)  # TODO
         montecarlo_radial1d(
-            model,
+            model_state,
             plasma,
             iteration,
             no_of_packets,
             total_iterations,
             show_progress_bars,
             self,
         )
-        self._integrator = FormalIntegrator(model, plasma, self)
+        self._integrator = FormalIntegrator(model_state, plasma, self)  # TODO
 
     def legacy_return(self):
         return (
@@ -564,13 +568,13 @@ def calculate_radiationfield_properties(self):
 
         return t_rad * u.K, w
 
-    def calculate_luminosity_inner(self, model):
+    def calculate_luminosity_inner(self, model_state, radiation_field):
         """
         Calculate inner luminosity.
 
         Parameters
         ----------
-        model : model.Radial1DModel
+        model_state : ModelState
 
         Returns
         -------
@@ -580,23 +584,27 @@ def calculate_luminosity_inner(self, model):
             4
             * np.pi
             * const.sigma_sb.cgs
-            * model.r_inner[0] ** 2
-            * model.t_inner**4
+            * model_state.geometry.r_inner[0] ** 2
+            * radiation_field.t_inner**4
         ).to("erg/s")
 
-    def calculate_time_of_simulation(self, model):
+    def calculate_time_of_simulation(self, model_state, radiation_field):
         """
         Calculate time of montecarlo simulation.
 
         Parameters
         ----------
-        model : model.Radial1DModel
+        model_state : ModelState
 
         Returns
         -------
         float
         """
-        return 1.0 * u.erg / self.calculate_luminosity_inner(model)
+        return (
+            1.0
+            * u.erg
+            / self.calculate_luminosity_inner(model_state, radiation_field)
+        )
 
     def calculate_f_nu(self, frequency):
         pass

tardis/montecarlo/montecarlo_numba/base.py

@@ -37,7 +37,7 @@
 
 
 def montecarlo_radial1d(
-    model,
+    model_state,
     plasma,
     iteration,
     no_of_packets,
@@ -53,9 +53,9 @@ def montecarlo_radial1d(
         transport._output_nu,
         transport._output_energy,
     )
-    numba_radial_1d_geometry = model.model_state.geometry.to_numba()
+    numba_radial_1d_geometry = model_state.geometry.to_numba()
     numba_model = NumbaModel(
-        model.model_state.time_explosion.to("s").value,
+        model_state.time_explosion.to("s").value,
     )
     numba_plasma = numba_plasma_initialize(
         plasma, transport.line_interaction_type

tardis/plasma/standard_plasmas.py

@@ -49,15 +49,16 @@
 logger = logging.getLogger(__name__)
 
 
-def assemble_plasma(config, model, atom_data=None):
+def assemble_plasma(config, model_state, radiation_field atom_data=None):
     """
     Create a BasePlasma instance from a Configuration object
     and a Radial1DModel.
 
     Parameters
     ----------
     config : io.config_reader.Configuration
-    model : model.Radial1DModel
+    model_state : model.ModelState
+    radiation_field : radiationField
     atom_data : atomic.AtomData
         If None, an attempt will be made to read the atomic data
         from config.
@@ -106,7 +107,7 @@ def assemble_plasma(config, model, atom_data=None):
             raise
 
     atom_data.prepare_atom_data(
-        model.abundance.index,
+        model_state.composition.elemental_mass_fraction.index,
         line_interaction_type=config.plasma.line_interaction_type,
         nlte_species=nlte_species,
     )
@@ -135,12 +136,12 @@ def assemble_plasma(config, model, atom_data=None):
     ]
 
     kwargs = dict(
-        t_rad=model.t_radiative,
-        abundance=model.abundance,
-        density=model.density,
+        t_rad=radiation_field.t_radiative,
+        abundance=model_state.composition.elemental_mass_fraction,
+        density=model_state.composition.density,
         atomic_data=atom_data,
-        time_explosion=model.time_explosion,
-        w=model.dilution_factor,
+        time_explosion=model_state.time_explosion,
+        w=radiation_field.dilution_factor,
         link_t_rad_t_electron=config.plasma.link_t_rad_t_electron,
         continuum_interaction_species=continuum_interaction_species,
         nlte_ionization_species=nlte_ionization_species,
@@ -213,9 +214,9 @@ def assemble_plasma(config, model, atom_data=None):
             bf_heating_coeff_estimator=None,
             stim_recomb_cooling_coeff_estimator=None,
             alpha_stim_estimator=None,
-            volume=model.volume,
-            r_inner=model.r_inner,
-            t_inner=model.t_inner,
+            volume=model_state.geometry.volume,
+            r_inner=model_state.geometry.r_inner,
+            t_inner=radiation_field.t_inner,
         )
     if config.plasma.radiative_rates_type == "blackbody":
         plasma_modules.append(JBluesBlackBody)
@@ -224,9 +225,9 @@ def assemble_plasma(config, model, atom_data=None):
     elif config.plasma.radiative_rates_type == "detailed":
         plasma_modules += detailed_j_blues_properties + detailed_j_blues_inputs
         kwargs.update(
-            r_inner=model.r_inner,
-            t_inner=model.t_inner,
-            volume=model.volume,
+            r_inner=model_state.geometry.r_inner,
+            t_inner=radiation_field.t_inner,
+            volume=model_state.geometry.volume,
             j_blue_estimator=None,
         )
         property_kwargs[JBluesDetailed] = {"w_epsilon": config.plasma.w_epsilon}
@@ -278,6 +279,7 @@ def assemble_plasma(config, model, atom_data=None):
     else:
         plasma_modules += helium_lte_properties
 
+    # TODO: compute electron densities somewhere!
     if model._electron_densities:
         electron_densities = pd.Series(model._electron_densities.cgs.value)
         if config.plasma.helium_treatment == "numerical-nlte":
@@ -289,11 +291,9 @@ def assemble_plasma(config, model, atom_data=None):
                 electron_densities=electron_densities
             )
 
-    if not model.raw_isotope_abundance.empty:
+    if not model_state.composition.raw_isotope_abundance.empty:
         plasma_modules += isotope_properties
-        isotope_abundance = model.raw_isotope_abundance.loc[
-            :, model.v_boundary_inner_index : model.v_boundary_outer_index - 1
-        ]
+        isotope_abundance = model_state.composition.raw_isotope_abundance
         kwargs.update(isotope_abundance=isotope_abundance)
 
     kwargs["helium_treatment"] = config.plasma.helium_treatment

tardis/radiation_field/base.py

@@ -1,12 +1,15 @@
 import numpy as np
 
+from tardis.montecarlo.packet_source import BasePacketSource
+
+
 class RadiationField:
     """_summary_
 
     Parameters
     ----------
     t_rad : numpy.ndarray
-        Radiative temperature in each shell 
+        Radiative temperature in each shell
     w : numpy.ndarray
         Dilution Factors in each shell
     opacities : Opacites
@@ -20,4 +23,33 @@ def __init__(self, t_rad, w, opacities, source_function):
         self.t_rad = t_rad
         self.w = w
         self.opacities = opacities
-        self.source_function = source_function
+        self.source_function = source_function
+
+
+class MonteCarloRadiationFieldState:
+    """_summary_
+
+    Parameters
+    ----------
+    t_rad : numpy.ndarray
+        Radiative temperature in each shell
+    w : numpy.ndarray
+        Dilution Factors in each shell
+    opacities : Opacites
+        Opacity container object
+    packet_source : SourceFunction
+        Source function for radiative transfer, for example a packet_source
+    """
+
+    def __init__(
+        self,
+        t_rad: np.ndarray,
+        w: np.ndarray,
+        opacities,
+        packet_source: BasePacketSource,
+    ):
+
+        self.t_rad = t_rad
+        self.w = w
+        self.opacities = opacities
+        self.source_function = packet_source