Estimators restructure cleanup (#2512)

* Fix plasma continuum test name * Fix zero division handling in normalize_trans_probs (which breaks to_hdf) * Use correct references_idx for levels * Convert probabilities to float64 * Fixes downbranch tests * Add bound-free estimator utility functions and classes * Work ongoing Co-authored-by: Christian Vogl <[email protected]> * ongoing work * ongoing work * Update code in various files * Fix photo ion calculation and radiation field intensity calculation Co-authored-by: Christian Vogl <[email protected]> * Co-authored-by: Christian Vogl <[email protected]> * Delete unused files and workspace configuration * Add Visual Studio Code workspace file to .gitignore * Add import statement for radfield_mc_estimators * Fix most tests * add new solvers * Reference fix and unit fix * Skip segfaulting tests (!) --------- Co-authored-by: Christian Vogl <[email protected]> Co-authored-by: Wolfgang Kerzendorf <[email protected]>
tardis-sn · Feb 6, 2024 · 234c586 · 234c586
1 parent 17c5da9
commit 234c586
Show file tree

Hide file tree

Showing 39 changed files with 1,036 additions and 184 deletions.
diff --git a/.gitignore b/.gitignore
@@ -71,3 +71,6 @@ pip-wheel-metadata/
 
 # Mac OSX
 .DS_Store
+
+# Visual Studio Code
+*.code-workspace
diff --git a/astropy_helpers b/astropy_helpers
diff --git a/docs/physics/update_and_conv/update_and_conv.ipynb b/docs/physics/update_and_conv/update_and_conv.ipynb
@@ -311,7 +311,7 @@
    },
    "outputs": [],
    "source": [
-    "j_estimator = transport.transport_state.estimators.j_estimator * (u.erg * u.cm) \n",
+    "j_estimator = transport.transport_state.j_estimator * (u.erg * u.cm) \n",
     "j_estimator"
    ]
   },
@@ -324,7 +324,7 @@
    },
    "outputs": [],
    "source": [
-    "nu_bar_estimator = transport.transport_state.estimators.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n",
+    "nu_bar_estimator = transport.transport_state.nu_bar_estimator * (u.erg * u.cm * u.Hz)\n",
     "nu_bar_estimator"
    ]
   },

diff --git a/tardis/energy_input/tests/test_energy_source.py b/tardis/energy_input/tests/test_energy_source.py
@@ -1,6 +1,6 @@
-import pytest
 import numpy as np
 import numpy.testing as npt
+import pytest
 
 from tardis.energy_input.samplers import (
     create_energy_cdf,

diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py
@@ -1,16 +1,17 @@
-# atomic model
-
-
 import logging
+from collections import OrderedDict
+
 import numpy as np
 import pandas as pd
-
-from scipy import interpolate
-from collections import OrderedDict
 from astropy import units as u
-from tardis import constants as const
 from astropy.units import Quantity
+from scipy import interpolate
+
+from tardis import constants as const
 from tardis.io.atom_data.util import resolve_atom_data_fname
+from tardis.plasma.properties.continuum_processes import (
+    get_ground_state_multi_index,
+)
 
 
 class AtomDataNotPreparedError(Exception):
@@ -24,7 +25,7 @@ class AtomDataMissingError(Exception):
 logger = logging.getLogger(__name__)
 
 
-class AtomData(object):
+class AtomData:
     """
     Class for storing atomic data
 
@@ -167,7 +168,6 @@ def from_hdf(cls, fname=None):
             Path to the HDFStore file or name of known atom data file
             (default: None)
         """
-
         dataframes = dict()
         nonavailable = list()
 
@@ -208,7 +208,7 @@ def from_hdf(cls, fname=None):
                                     "Atomic Data Collisions Not a Valid Chanti or CMFGEN Carsus Data File"
                                 )
                         except KeyError as e:
-                            logger.warn(
+                            logger.warning(
                                 "Atomic Data is not a Valid Carsus Atomic Data File"
                             )
                             raise
@@ -253,7 +253,7 @@ def from_hdf(cls, fname=None):
                 )
                 atom_data.version = None
 
-            # ToDo: strore data sources as attributes in carsus
+            # TODO: strore data sources as attributes in carsus
 
             logger.info(
                 f"Reading Atom Data with: UUID = {atom_data.uuid1} MD5  = {atom_data.md5} "
@@ -373,8 +373,9 @@ def _check_related(self):
     def prepare_atom_data(
         self,
         selected_atomic_numbers,
-        line_interaction_type="scatter",
-        nlte_species=[],
+        line_interaction_type,
+        nlte_species,
+        continuum_interaction_species,
     ):
         """
         Prepares the atom data to set the lines, levels and if requested macro
@@ -437,6 +438,84 @@ def prepare_atom_data(
             .values
         )
 
+        self.prepare_macro_atom_data(
+            line_interaction_type,
+            tmp_lines_lower2level_idx,
+            tmp_lines_upper2level_idx,
+        )
+        if len(continuum_interaction_species) > 0:
+            self.prepare_continuum_interaction_data(
+                continuum_interaction_species
+            )
+
+        self.nlte_data = NLTEData(self, nlte_species)
+
+    def prepare_continuum_interaction_data(self, continuum_interaction_species):
+        """
+        Prepares the atom data for the continuum interaction
+
+        Parameters
+        ----------
+        continuum_interaction : ContinuumInteraction
+            The continuum interaction object
+        """
+        # photoionization_data = atomic_data.photoionization_data.set_index(
+        #    ["atomic_number", "ion_number", "level_number"]
+        # )
+        mask_selected_species = self.photoionization_data.index.droplevel(
+            "level_number"
+        ).isin(continuum_interaction_species)
+        self.photoionization_data = self.photoionization_data[
+            mask_selected_species
+        ]
+        self.photo_ion_block_references = np.pad(
+            self.photoionization_data.nu.groupby(level=[0, 1, 2])
+            .count()
+            .values.cumsum(),
+            [1, 0],
+        )
+        self.photo_ion_unique_index = self.photoionization_data.index.unique()
+        self.nu_ion_threshold = (
+            self.photoionization_data.groupby(level=[0, 1, 2]).first().nu
+        )
+        self.energy_photo_ion_lower_level = self.levels.loc[
+            self.photo_ion_unique_index
+        ].energy
+
+        source_idx = self.macro_atom_references.loc[
+            self.photo_ion_unique_index
+        ].references_idx
+        destination_idx = self.macro_atom_references.loc[
+            get_ground_state_multi_index(self.photo_ion_unique_index)
+        ].references_idx
+        self.photo_ion_levels_idx = pd.DataFrame(
+            {
+                "source_level_idx": source_idx.values,
+                "destination_level_idx": destination_idx.values,
+            },
+            index=self.photo_ion_unique_index,
+        )
+
+        self.level2continuum_edge_idx = pd.Series(
+            np.arange(len(self.nu_ion_threshold)),
+            self.nu_ion_threshold.sort_values(ascending=False).index,
+            name="continuum_idx",
+        )
+
+        level_idxs2continuum_idx = self.photo_ion_levels_idx.copy()
+        level_idxs2continuum_idx[
+            "continuum_idx"
+        ] = self.level2continuum_edge_idx
+        self.level_idxs2continuum_idx = level_idxs2continuum_idx.set_index(
+            ["source_level_idx", "destination_level_idx"]
+        )
+
+    def prepare_macro_atom_data(
+        self,
+        line_interaction_type,
+        tmp_lines_lower2level_idx,
+        tmp_lines_upper2level_idx,
+    ):
         if (
             self.macro_atom_data_all is not None
             and not line_interaction_type == "scatter"
@@ -505,6 +584,13 @@ def prepare_atom_data(
             )
 
             if line_interaction_type == "macroatom":
+                self.lines_lower2macro_reference_idx = (
+                    self.macro_atom_references.loc[
+                        tmp_lines_lower2level_idx, "references_idx"
+                    ]
+                    .astype(np.int64)
+                    .values
+                )
                 # Sets all
                 tmp_macro_destination_level_idx = pd.MultiIndex.from_arrays(
                     [
@@ -548,8 +634,6 @@ def prepare_atom_data(
                     self.selected_atomic_numbers, level=0
                 )
 
-        self.nlte_data = NLTEData(self, nlte_species)
-
     def _check_selected_atomic_numbers(self):
         selected_atomic_numbers = self.selected_atomic_numbers
         available_atomic_numbers = np.unique(
@@ -571,7 +655,7 @@ def __repr__(self):
         return f"<Atomic Data UUID={self.uuid1} MD5={self.md5} Lines={self.lines.line_id.count():d} Levels={self.levels.energy.count():d}>"
 
 
-class NLTEData(object):
+class NLTEData:
     def __init__(self, atom_data, nlte_species):
         self.atom_data = atom_data
         self.lines = atom_data.lines.reset_index()

diff --git a/tardis/io/tests/test_atomic.py b/tardis/io/tests/test_atomic.py
@@ -55,7 +55,12 @@ def test_atom_data_lines(lines):
 
 
 def test_atomic_reprepare(kurucz_atomic_data):
-    kurucz_atomic_data.prepare_atom_data([14, 20])
+    kurucz_atomic_data.prepare_atom_data(
+        [14, 20],
+        line_interaction_type="scatter",
+        nlte_species=[],
+        continuum_interaction_species=[],
+    )
     lines = kurucz_atomic_data.lines.reset_index()
     assert lines["atomic_number"].isin([14, 20]).all()
     assert len(lines.loc[lines["atomic_number"] == 14]) > 0

diff --git a/tardis/model/base.py b/tardis/model/base.py
@@ -23,7 +23,9 @@
     parse_packet_source,
 )
 from tardis.montecarlo.packet_source import BlackBodySimpleSource
-from tardis.model.radiation_field_state import DiluteThermalRadiationFieldState
+from tardis.model.radiation_field_state import (
+    DiluteBlackBodyRadiationFieldState,
+)
 from tardis.util.base import is_valid_nuclide_or_elem
 
 logger = logging.getLogger(__name__)

diff --git a/tardis/model/parse_input.py b/tardis/model/parse_input.py
@@ -17,7 +17,9 @@
 from tardis.model.geometry.radial1d import HomologousRadial1DGeometry
 from tardis.model.matter.composition import Composition
 from tardis.model.matter.decay import IsotopicMassFraction
-from tardis.model.radiation_field_state import DiluteThermalRadiationFieldState
+from tardis.model.radiation_field_state import (
+    DiluteBlackBodyRadiationFieldState,
+)
 from tardis.montecarlo.packet_source import (
     BlackBodySimpleSource,
     BlackBodySimpleSourceRelativistic,
@@ -572,7 +574,7 @@ def parse_radiation_field_state(
         ]
         assert len(dilution_factor) == geometry.no_of_shells
 
-    return DiluteThermalRadiationFieldState(t_radiative, dilution_factor)
+    return DiluteBlackBodyRadiationFieldState(t_radiative, dilution_factor)
 
 
 def initialize_packet_source(config, geometry, packet_source):
@@ -687,7 +689,7 @@ def parse_csvy_radiation_field_state(
     else:
         dilution_factor = calculate_geometric_dilution_factor(geometry)
 
-    return DiluteThermalRadiationFieldState(t_radiative, dilution_factor)
+    return DiluteBlackBodyRadiationFieldState(t_radiative, dilution_factor)
 
 
 def calculate_t_radiative_from_t_inner(geometry, packet_source):

diff --git a/tardis/model/radiation_field_state.py b/tardis/model/radiation_field_state.py
@@ -1,11 +1,12 @@
-from tardis.montecarlo.montecarlo_numba.numba_interface import OpacityState
-
-
 import numpy as np
 from astropy import units as u
 
+from tardis.util.base import intensity_black_body
+
+from typing import Union
+
 
-class DiluteThermalRadiationFieldState:
+class DiluteBlackBodyRadiationFieldState:
     """
     Represents the state of a dilute thermal radiation field.
 
@@ -30,3 +31,21 @@ def __init__(
         assert np.all(dilution_factor > 0)
         self.t_radiative = t_radiative
         self.dilution_factor = dilution_factor
+
+    def calculate_mean_intensity(self, nu: Union[u.Quantity, np.ndarray]):
+        """
+        Calculate the intensity of the radiation field at a given frequency.
+
+        Parameters
+        ----------
+        nu : u.Quantity
+            Frequency at which the intensity is to be calculated
+
+        Returns
+        -------
+        intensity : u.Quantity
+            Intensity of the radiation field at the given frequency
+        """
+        return self.dilution_factor * intensity_black_body(
+            nu[np.newaxis].T, self.t_radiative
+        )
diff --git a/tardis/montecarlo/base.py b/tardis/montecarlo/base.py
@@ -8,14 +8,16 @@
 from tardis.io.logger import montecarlo_tracking as mc_tracker
 from tardis.io.util import HDFWriterMixin
 from tardis.montecarlo import montecarlo_configuration
+from tardis.montecarlo.estimators.radfield_mc_estimators import (
+    initialize_estimator_statistics,
+)
 from tardis.montecarlo.montecarlo_configuration import (
     configuration_initialize,
 )
 from tardis.montecarlo.montecarlo_numba import (
     montecarlo_main_loop,
     numba_config,
 )
-from tardis.montecarlo.montecarlo_numba.estimators import initialize_estimators
 from tardis.montecarlo.montecarlo_numba.formal_integral import FormalIntegrator
 from tardis.montecarlo.montecarlo_numba.numba_interface import (
     NumbaModel,
@@ -97,8 +99,8 @@ def __init__(
         mc_tracker.DEBUG_MODE = debug_packets
         mc_tracker.BUFFER = logger_buffer
 
-        if self.spectrum_method == "integrated":
-            self.optional_hdf_properties.append("spectrum_integrated")
+        # if self.spectrum_method == "integrated":
+        #    self.optional_hdf_properties.append("spectrum_integrated")
 
     def initialize_transport_state(
         self,
@@ -116,7 +118,7 @@ def initialize_transport_state(
         packet_collection = self.packet_source.create_packets(
             no_of_packets, seed_offset=iteration
         )
-        estimators = initialize_estimators(
+        estimators = initialize_estimator_statistics(
             plasma.tau_sobolevs.shape, gamma_shape
         )
 
@@ -182,7 +184,7 @@ def run(
             transport_state.geometry_state,
             numba_model,
             transport_state.opacity_state,
-            transport_state.estimators,
+            transport_state.radfield_mc_estimators,
             transport_state.spectrum_frequency.value,
             number_of_vpackets,
             iteration=iteration,
@@ -231,8 +233,8 @@ def legacy_return(self):
         return (
             self.transport_state.packet_collection.output_nus,
             self.transport_state.packet_collection.output_energies,
-            self.transport_state.estimators.j_estimator,
-            self.transport_state.estimators.nu_bar_estimator,
+            self.transport_state.j_estimator,
+            self.transport_state.nu_bar_estimator,
             self.transport_state.last_line_interaction_in_id,
             self.transport_state.last_line_interaction_out_id,
             self.transport_state.last_interaction_type,

diff --git a/tardis/montecarlo/conftest.py b/tardis/montecarlo/conftest.py
@@ -0,0 +1,16 @@
+import pytest
+
+from tardis.io.configuration.config_reader import Configuration
+
+
+@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", "Ti II"]
+    continuum_config.plasma.nlte_ionization_species = []
+
+    return continuum_config
diff --git a/tardis/montecarlo/estimators/__init__.py b/tardis/montecarlo/estimators/__init__.py