Opacity state restructure (tardis-sn#2773)

* created file for opacity solver

* Added opacity states for the continuum and macroatom

* Updated the opacity state and macroatom state to work with legacy plasma

* Added numba initialize from opacity state, added methods to solve opacity state with the solver

* moved line2macro_level_upper to the macroatom state

* renamed a macroatom variable, added solver to montecarlo solver

* ran black

* corrected an import

* corrected an import

* ran black

* Added line_interaction_type to the to_numba function, should be reqorked in the future

* fixed initialization of opacity state

* Updated handling of continuum vs noncontinuum macroblock references

* Fixed a typo

* refactored the state objects to put them in the correct folders

* another typo

* added basic docstrings to the continuum state

* Added docstrings to macroatom state

* Updated docstrings onithe opacity_state file

* Added docstrings to the solver

* Updated Docstring Style

* ran black

* changed Opacity solver to take kwargs

* ran black

* Added class attributes to Opacity Solver

* Changed names of the opacity state and numba equivalent

* Added test for the opacity solver

* added conftest to the opacities module

* fixed up tests for the opacity_solver

* ran black

* Added new test for the numba opacity state

* ran black, updated name of the test

* improved tau_sobelev disable_line_scattering initialization

* Fixed tau_sobolev initialization

* added some comments

* Changed scope of the simulation in conftest

tardis/opacities/continuum/continuum_state.py

@@ -0,0 +1,194 @@
+import numpy as np
+
+
+class ContinuumState:
+    def __init__(
+        self,
+        nu_i,
+        level2continuum_idx,
+        p_fb_deactivation,
+        photo_ion_cross_sections,
+        chi_bf,
+        ff_cooling_factor,
+        fb_emission_cdf,
+        photo_ion_idx,
+        k_packet_idx,
+    ):
+        """
+        Current State of the Continuum Required for Opacity Computation
+
+        Parameters
+        ----------
+        nu_i : pd.DataFrame
+            frequencies for the bound-free thresholds
+        level2continuum_idx : pd.DataFrame
+            mapping from levels to the continuum
+        p_fb_deactivation : pd.DataFrame
+            probabilities of free-bound deactivation channels
+        photo_ion_cross_sections : pd.DataFrame
+            Photoionization cross sections
+        chi_bf : pd.DataFrame
+            Bound-free opacities
+        ff_cooling_factor : np.ndarray
+            free-free cooling factor
+        fb_emission_cdf : pd.DataFrame
+            free-bound emission cumulative distribution function
+        photo_ion_idx : pd.DataFrame
+            photoionization indices
+        k_packet_idx : pd.DataFrame
+            k-packet indices
+        """
+        self.nu_i = nu_i
+        self.level2continuum_idx = level2continuum_idx
+        self.p_fb_deactivation = p_fb_deactivation
+        self.photo_ion_cross_sections = photo_ion_cross_sections
+        self._chi_bf = chi_bf
+        self.ff_cooling_factor = ff_cooling_factor
+        self.fb_emission_cdf = fb_emission_cdf
+        self.photo_ion_idx = photo_ion_idx
+        self.k_packet_idx = k_packet_idx
+
+    @classmethod
+    def from_legacy_plasma(cls, plasma):
+        """
+        Generates a ContinuumState object from a tardis BasePlasma
+
+        Parameters
+        ----------
+        plasma : tarids.plasma.BasePlasma
+            legacy base plasma
+
+        Returns
+        -------
+        ContinuumState
+        """
+        nu_i = plasma.nu_i
+        level2continuum_idx = plasma.level2continuum_idx
+        p_fb_deactivation = plasma.p_fb_deactivation
+        photo_ion_cross_sections = plasma.photo_ion_cross_sections
+        chi_bf = plasma.chi_bf
+        ff_cooling_factor = plasma.ff_cooling_factor
+        fb_emission_cdf = plasma.fb_emission_cdf
+        photo_ion_idx = plasma.photo_ion_idx
+        k_packet_idx = plasma.k_packet_idx
+
+        return cls(
+            nu_i,
+            level2continuum_idx,
+            p_fb_deactivation,
+            photo_ion_cross_sections,
+            chi_bf,
+            ff_cooling_factor,
+            fb_emission_cdf,
+            photo_ion_idx,
+            k_packet_idx,
+        )
+
+    @property
+    def bf_threshold_list_nu(self):
+        """
+        List of Bound-Free Threshold Frequencies
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.nu_i.loc[self.level2continuum_idx.index]
+
+    @property
+    def phot_nus(self):
+        """
+        Frequencies corresponding to Photoionization Cross Sections
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.photo_ion_cross_sections.nu.loc[
+            self.level2continuum_idx.index
+        ]
+
+    @property
+    def photo_ion_block_references(self):
+        """Photoionization Block References
+
+        Returns
+        -------
+        np.ndarray
+        """
+        return np.pad(
+            self.phot_nus.groupby(level=[0, 1, 2], sort=False)
+            .count()
+            .values.cumsum(),
+            [1, 0],
+        )
+
+    @property
+    def photo_ion_nu_threshold_mins(self):
+        """
+        Minimum Edges of the photoionization threshold frequencies
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.phot_nus.groupby(level=[0, 1, 2], sort=False).first()
+
+    @property
+    def photo_ion_nu_threshold_maxs(self):
+        """
+        Maximum Edges of the photoionization threshold frequencies
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.phot_nus.groupby(level=[0, 1, 2], sort=False).last()
+
+    @property
+    def x_sect(self):
+        """
+        Photoionization Cross Sections mapped to the continuum indices
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.photo_ion_cross_sections.x_sect.loc[
+            self.level2continuum_idx.index
+        ]
+
+    @property
+    def chi_bf(self):
+        """
+        Bound-Free Opacities indices corresponding to the continuum levels
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self._chi_bf.loc[self.level2continuum_idx.index]
+
+    @property
+    def emissivities(self):
+        """
+        Free-bound Emissivities corresponding to the continuum levels
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.fb_emission_cdf.loc[self.level2continuum_idx.index]
+
+    @property
+    def photo_ion_activation_idx(self):
+        """
+        Index corresponding to photoionization activation
+
+        Returns
+        -------
+        pd.DataFrame
+        """
+        return self.photo_ion_idx.loc[
+            self.level2continuum_idx.index, "destination_level_idx"
+        ]

tardis/opacities/macro_atom/macroatom_state.py

@@ -0,0 +1,90 @@
+from tardis.transport.montecarlo.configuration import montecarlo_globals
+from tardis.io.util import HDFWriterMixin
+
+
+class MacroAtomState(HDFWriterMixin):
+
+    hdf_name = "macro_atom_state"
+
+    hdf_properties = [
+        "transition_probabilities",
+        "transition_type",
+        "destination_level_id",
+        "transition_line_id",
+        "macro_block_references",
+        "line2macro_level_upper",
+    ]
+
+    def __init__(
+        self,
+        transition_probabilities,
+        transition_type,
+        destination_level_id,
+        transition_line_id,
+        macro_block_references,
+        line2macro_level_upper,
+    ):
+        """
+        Current State of the MacroAtom
+
+        Parameters
+        ----------
+        transition_probabilities : pd.DataFrame
+            Macro Atom Transition probabilities between levels
+        transition_type : pd.DataFrame)
+            Macro Atom transition types
+        destination_level_id : pd.DataFrame
+            ID of destination levels of the Macro Atom
+        transition_line_id : pd.DataFrame
+            ID of lines corresponding to Macro Atom transitions
+        macro_block_references : pd.DataFrame or np.ndarray
+            Index references to the Macro Atom blocks
+        line2macro_level_upper : pd.DataFrame
+            Mapping from lines to Macro Atom upper levels
+        """
+        self.transition_probabilities = transition_probabilities
+        self.transition_type = transition_type
+        self.destination_level_id = destination_level_id
+        self.transition_line_id = transition_line_id
+        self.macro_block_references = macro_block_references
+        self.line2macro_level_upper = line2macro_level_upper
+
+    @classmethod
+    def from_legacy_plasma(cls, plasma):
+        """
+        Generates a MacroAtomState object from a tardis BasePlasma
+
+        Parameters
+        ----------
+        plasma : tarids.plasma.BasePlasma
+            legacy base plasma
+
+        Returns
+        -------
+        MacroAtomState
+        """
+        transition_probabilities = plasma.transition_probabilities
+        transition_type = plasma.macro_atom_data["transition_type"]
+        destination_level_id = plasma.macro_atom_data["destination_level_idx"]
+        transition_line_id = plasma.macro_atom_data["lines_idx"]
+        line2macro_level_upper = (
+            plasma.atomic_data.lines_upper2macro_reference_idx
+        )
+
+        if (
+            montecarlo_globals.CONTINUUM_PROCESSES_ENABLED
+        ):  # TODO: Unify this in the plasma solver
+            macro_block_references = plasma.macro_block_references
+        else:
+            macro_block_references = plasma.atomic_data.macro_atom_references[
+                "block_references"
+            ]
+
+        return cls(
+            transition_probabilities,
+            transition_type,
+            destination_level_id,
+            transition_line_id,
+            macro_block_references,
+            line2macro_level_upper,
+        )

tardis/opacities/opacity_solver.py

@@ -0,0 +1,67 @@
+from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity
+from tardis.opacities.opacity_state import (
+    OpacityState,
+)
+import numpy as np
+import pandas as pd
+
+
+class OpacitySolver(object):
+
+    line_interaction_type: str = "scatter"
+    disable_line_scattering: bool = False
+
+    def __init__(
+        self, line_interaction_type="scatter", disable_line_scattering=False
+    ):
+        """Solver class for opacities
+
+        Parameters
+        ----------
+        line_interaction_type: str
+            "scatter", "downbranch", or "macroatom"
+        disable_line_scattering: bool
+        """
+
+        self.line_interaction_type = line_interaction_type
+        self.disable_line_scattering = disable_line_scattering
+
+    def solve(self, legacy_plasma) -> OpacityState:
+        """
+        Solves the opacity state
+
+        Parameters
+        ----------
+        plasma : tarids.plasma.BasePlasma
+            legacy base plasma
+
+        Returns
+        -------
+        OpacityState
+        """
+        if self.disable_line_scattering:
+            tau_sobolev = pd.DataFrame(
+                np.zeros(
+                    (
+                        legacy_plasma.atomic_data.lines.shape[
+                            0
+                        ],  # number of lines
+                        legacy_plasma.abundance.shape[1],  # number of shells
+                    ),
+                    dtype=np.float64,
+                ),
+                index=legacy_plasma.atomic_data.lines.index,
+            )
+        else:
+            tau_sobolev = calculate_sobolev_line_opacity(
+                legacy_plasma.atomic_data.lines,
+                legacy_plasma.level_number_density,
+                legacy_plasma.time_explosion,
+                legacy_plasma.stimulated_emission_factor,
+            )
+
+        opacity_state = OpacityState.from_legacy_plasma(
+            legacy_plasma, tau_sobolev
+        )
+
+        return opacity_state