Restructure/estimators

tardis/model/radiation_field_state.py

@@ -3,6 +3,8 @@
 
 from tardis.util.base import intensity_black_body
 
+from typing import Union
+
 
 class DiluteBlackBodyRadiationFieldState:
     """
@@ -30,7 +32,7 @@ def __init__(
         self.t_radiative = t_radiative
         self.dilution_factor = dilution_factor
 
-    def calculate_mean_intensity(self, nu):
+    def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]):
         """
         Calculate the intensity of the radiation field at a given frequency.
 
@@ -44,4 +46,6 @@ def calculate_mean_intensity(self, nu):
         intensity : u.Quantity
             Intensity of the radiation field at the given frequency
         """
-        return self.dilution_factor * intensity_black_body(nu, self.t_radiative)
+        return self.dilution_factor * intensity_black_body(
+            nu[np.newaxis].T, self.t_radiative
+        )

tardis/montecarlo/estimators/estimated_radfield_properties.py

@@ -1,16 +1,23 @@
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 
 import numpy as np
 import pandas as pd
+from astropy import units as u
 
 import tardis.constants as const
+from tardis.io.atom_data import AtomData
 from tardis.model.radiation_field_state import (
     DiluteBlackBodyRadiationFieldState,
 )
 from tardis.montecarlo.estimators.util import (
     bound_free_estimator_array2frame,
     integrate_array_by_blocks,
 )
+from tardis.montecarlo.estimators.radfield_mc_estimators import (
+    RadiationFieldMCEstimators,
+)
+from tardis.plasma.properties.continuum_processes import PhotoIonBoltzmannFactor
 
 H = const.h.cgs.value
 
@@ -25,143 +32,136 @@ def calculate_stimulated_recomb_rate_coefficient(self):
         pass
 
 
-class EstimatedMCContinuumProperties(EstimatedContinuumPropertiesABC):
+class MCContinuumPropertiesSolver(EstimatedContinuumPropertiesABC):
     def __init__(
         self,
-        radfield_mc_estimators,
-        level2continuum_idx,
+        atom_data: AtomData,
+    ):
+        self.atom_data = atom_data
+
+    def solve(
+        self,
+        radfield_mc_estimators: RadiationFieldMCEstimators,
         time_simulation,
         volume,
     ):
-        self.radfield_mc_estimators = radfield_mc_estimators
-        self.level2continuum_idx = level2continuum_idx
+        photo_ion_norm_factor = (time_simulation * volume * H) ** -1
 
-        # initializing the lazy properties
-        self._gamma_estimator_df = None
-        self._alpha_stimulated_estimator_df = None
-        self.photo_ion_norm_factor = (time_simulation * volume * H) ** -1
+        photo_ionization_rate_coefficient = bound_free_estimator_array2frame(
+            radfield_mc_estimators.photo_ion_estimator,
+            self.atom_data.level2continuum_edge_idx,
+        )
+        photo_ionization_rate_coefficient *= photo_ion_norm_factor
 
-    def calculate_photo_ionization_rate_coefficient(self):
-        return self.gamma_estimator_df * self.photo_ion_norm_factor.value
+        stimulated_recomb_rate_factor = bound_free_estimator_array2frame(
+            radfield_mc_estimators.stim_recomb_estimator,,
+            self.atom_data.level2continuum_edge_idx,
+        )
+        stimulated_recomb_rate_factor *= photo_ion_norm_factor
 
-    def calculate_stimulated_recomb_rate_factor(self):
-        """
-        Calculates the stimulated recombination rate coefficient.
+        return ContinuumProperties(stimulated_recomb_rate_factor, photo_ionization_rate_coefficient)
 
-        Returns
-        -------
-        alpha_stim : ndarray
-            The calculated stimulated recombination rate coefficient.
 
-        See Equation 15 in Lucy 2003
-        """
-        alpha_factor = (
-            self.alpha_stimulated_estimator_df * self.photo_ion_norm_factor
+class DiluteBlackBodyContinuumPropertiesSolver(EstimatedContinuumPropertiesABC):
+    def __init__(self, atom_data: AtomData) -> None:
+        self.atom_data = atom_data
+
+    def solve(
+        self,
+        dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState,
+        t_electrons: u.Quantity,
+    ):
+        photo_ion_boltzmann_factor = PhotoIonBoltzmannFactor.calculate(
+            self.atom_data.photoionization_data, t_electrons
         )
-        return alpha_factor
-
-    @property
-    def alpha_stimulated_estimator_df(self):
-        if self._alpha_stimulated_estimator_df is None:
-            self._alpha_stimulated_estimator_df = (
-                bound_free_estimator_array2frame(
-                    self.radfield_mc_estimators.stim_recomb_estimator,
-                    self.level2continuum_idx,
-                )
+        mean_intensity_photo_ion_df = (
+            self.calculate_mean_intensity_photo_ion_table(
+                dilute_blackbody_radiationfield_state
             )
-        return self._alpha_stimulated_estimator_df
-
-    @property
-    def gamma_estimator_df(self):
-        if self._gamma_estimator_df is None:
-            self._gamma_estimator_df = bound_free_estimator_array2frame(
-                self.radfield_mc_estimators.photo_ion_estimator,
-                self.level2continuum_idx,
-            )
-        return self._gamma_estimator_df
+        )
 
+        photo_ion_rate_coeff = self.calculate_photo_ionization_rate_coefficient(
+            mean_intensity_photo_ion_df
+        )
+        stimulated_recomb_rate_coeff = (
+            self.calculate_stimulated_recomb_rate_coefficient(
+                mean_intensity_photo_ion_df,
+                photo_ion_boltzmann_factor,
+            )
+        )
 
-class EstimatedDiluteBlackBodyContinuumProperties(
-    EstimatedContinuumPropertiesABC
-):
-    def __init__(
-        self,
-        dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState,
-    ) -> None:
-        self.dilute_blackbody_radiationfield_state = (
-            dilute_blackbody_radiationfield_state
+        return ContinuumProperties(
+            stimulated_recomb_rate_coeff, photo_ion_rate_coeff
         )
 
     def calculate_photo_ionization_rate_coefficient(
         self,
-        photo_ion_cross_sections_df,
-        photo_ion_block_references,
-        photo_ion_index,
+        mean_intensity_photo_ion_df: pd.DataFrame,
     ):
-        mean_intensity_df = self.calculate_mean_intensity_for_photo_ion_table(
-            photo_ion_cross_sections_df
-        )
-
-        gamma = mean_intensity_df.multiply(
+        gamma = mean_intensity_photo_ion_df.multiply(
             4.0
             * np.pi
-            * photo_ion_cross_sections_df.x_sect
-            / (photo_ion_cross_sections_df.nu * H),
+            * self.atom_data.photoionization_data.x_sect
+            / (self.atom_data.photoionization_data.nu * H),
             axis=0,
         )
         gamma = integrate_array_by_blocks(
             gamma.values,
-            photo_ion_cross_sections_df.nu.values,
-            photo_ion_block_references,
+            self.atom_data.photoionization_data.nu.values,
+            self.atom_data.photo_ion_block_references,
         )
-        gamma = pd.DataFrame(gamma, index=photo_ion_index)
+        gamma = pd.DataFrame(gamma, index=self.atom_data.photo_ion_unique_index)
         return gamma
 
     def calculate_stimulated_recomb_rate_coefficient(
         self,
-        photo_ion_cross_sections_df,
-        photo_ion_block_references,
-        photo_ion_index,
-        boltzmann_factor_photo_ion,
+        mean_intensity_photo_ion_df: pd.DataFrame,
+        photo_ion_boltzmann_factor: np.ndarray,
     ):
-        nu = photo_ion_cross_sections["nu"]
-        x_sect = photo_ion_cross_sections["x_sect"]
-        # TODO: duplicated; also used in calculate_photo_ionization_rate_coefficient
-        mean_intensity_df = self.calculate_mean_intensity_for_photo_ion_table(
-            photo_ion_cross_sections_df
+        alpha_stim_factor = (
+            mean_intensity_photo_ion_df * photo_ion_boltzmann_factor
         )
-
-        mean_intensity_df *= boltzmann_factor_photo_ion
-        alpha_stim = mean_intensity_df.multiply(
-            4.0 * np.pi * x_sect / nu / H, axis=0
+        alpha_stim_factor = alpha_stim_factor.multiply(
+            4.0
+            * np.pi
+            * self.atom_data.photoionization_data.x_sect
+            / (self.atom_data.photoionization_data.nu * H),
+            axis=0,
         )
-        alpha_stim = integrate_array_by_blocks(
-            alpha_stim.values, nu.values, photo_ion_block_references
+        alpha_stim_factor = integrate_array_by_blocks(
+            alpha_stim_factor.values,
+            self.atom_data.photoionization_data.nu.values,
+            self.atom_data.photo_ion_block_references,
         )
-        alpha_stim = pd.DataFrame(alpha_stim, index=photo_ion_index)
-        return alpha_stim
 
-    def calculate_mean_intensity_for_photo_ion_table(
-        self, photo_ion_cross_sections_df
+        alpha_stim_factor = pd.DataFrame(
+            alpha_stim_factor, index=self.atom_data.photo_ion_unique_index
+        )
+
+        return alpha_stim_factor
+
+    def calculate_mean_intensity_photo_ion_table(
+        self,
+        dilute_blackbody_radiationfield_state: DiluteBlackBodyRadiationFieldState,
     ):
         mean_intensity = (
-            self.dilute_blackbody_radiationfield_state.calculate_mean_intensity(
-                photo_ion_cross_sections_df.nu
+            dilute_blackbody_radiationfield_state.calculate_mean_intensity(
+                self.atom_data.photoionization_data.nu.values
             )
         )
         mean_intensity_df = pd.DataFrame(
             mean_intensity,
-            index=photo_ion_cross_sections_df.index,
+            index=self.atom_data.photoionization_data.index,
             columns=np.arange(
-                len(self.dilute_blackbody_radiationfield_state.t_radiative)
+                len(dilute_blackbody_radiationfield_state.t_radiative)
             ),
         )
         return mean_intensity_df
 
 
 """
 class PhotoIonBoltzmannFactor(ProcessingPlasmaProperty):
-    
+
     Attributes
     ----------
     boltzmann_factor_photo_ion : pandas.DataFrame, dtype float
@@ -175,3 +175,9 @@ def calculate(self, photo_ion_cross_sections, t_electrons):
         boltzmann_factor = np.exp(-nu[np.newaxis].T / t_electrons * (H / K_B))
         return boltzmann_factor
 """
+
+
+@dataclass
+class ContinuumProperties:
+    stimulated_recomb_rate_factor: pd.DataFrame
+    photo_ionization_rate_coefficient: pd.DataFrame

tardis/montecarlo/montecarlo_numba/tests/test_continuum.py

@@ -1,35 +1,8 @@
 from copy import deepcopy
 
-import pytest
-
-from tardis.io.configuration.config_reader import Configuration
-from tardis.simulation import Simulation
-
 import numpy.testing as npt
 
-
-@pytest.fixture(scope="function")
-def continuum_config(
-    tardis_config_verysimple_nlte,
-):
-    continuum_config = Configuration.from_config_dict(
-        tardis_config_verysimple_nlte
-    )
-    continuum_config.plasma.continuum_interaction.species = ["H I"]
-    continuum_config.plasma.nlte_ionization_species = []
-
-    """"
-    montecarlo_configuration.LEGACY_MODE_ENABLED = True
-    # Setup model config from verysimple
-
-    config_montecarlo_1e5_verysimple.montecarlo.last_no_of_packets = 1e5
-    config_montecarlo_1e5_verysimple.montecarlo.no_of_virtual_packets = 0
-    config_montecarlo_1e5_verysimple.montecarlo.iterations = 1
-    config_montecarlo_1e5_verysimple.plasma.line_interaction_type = "macroatom"
-
-    del config_montecarlo_1e5_verysimple["config_dirname"]
-    """
-    return continuum_config
+from tardis.simulation import Simulation
 
 
 def test_montecarlo_continuum(

tardis/plasma/properties/continuum_processes.py

@@ -1017,7 +1017,8 @@ class PhotoIonBoltzmannFactor(ProcessingPlasmaProperty):
 
     outputs = ("boltzmann_factor_photo_ion",)
 
-    def calculate(self, photo_ion_cross_sections, t_electrons):
+    @staticmethod
+    def calculate(photo_ion_cross_sections, t_electrons):
         nu = photo_ion_cross_sections["nu"].values
 
         boltzmann_factor = np.exp(-nu[np.newaxis].T / t_electrons * (H / K_B))