Macroatom restructure (#2786)

* 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

* Make a function to calculate markov chain indices

* Extracted more functions

* Moved more calculations over to functions (maybe this should be a separate graph?) and implimented a basic MacroAtomSolver

* Forgot an extra arg to the macroatom state

* fixed up continuum vs non-continuum initialization and handling of continuum_interaction_species when no continuum_interaction

* broke about the base transition probabilities properties into functions

* Added the default transition probabilities to the solver

* See previous commit, added a comment

* Improved code flow for continuum vs non-continuum

* See previous commit

* See previous commit

* removed non-continuum interaction species from non-continuum solver path

* broke up the solver to have a common method for both continuum and non-continuum

* Made a separate solver for the macroatom when working with the continuum:

* cleaned up default solver a bit

* ditto

* ran black, added personal comments

* first attempt to work basic macroatom solver into the opacity solver

* Fixed calculate_deactivation_channel_probs missing argument 'B'

* Fixed calculate_deactivation_channel_probs missing argument 'B'

* Fixed accessign the macro_atom_data from the plasma

* Fixed typo

* Small fix to make tests pass, need to revisit later

* Small fix to make tests pass, need to revisit later

* Small fix to make tests pass, need to revisit later

* Small fix to make tests pass, need to revisit later

* Small fix to make tests pass, need to revisit later

* Small fix to make tests pass, need to revisit later

* Added switch to swap our the macroatom solver when using the continuum

* Made a passthrough when continuum processes are enabled in case we don't want to worry about that part of the restructure

* Removed switch to auto-build continuum macrostate using legacy mode

* Fixed initialization in default macroatom solver

* grab continuum interaction species from the legacy plasma

* Added in the rest of the transition probability properties to the continuum macro atom solver.  Manually hard-coded and should be changed

* Removed old method of doing this

* switch back to using legacy solve for continuum in case we don't want to work on that solver

* Moved the MacroAtomContinuumSolver to a new file (deleted from this PR) and restrctured code accordingly

* moved 'calculate_' functions derived from transition_probabilities.py to a new file (deleted in this PR, will be posted in another)

* Moved MacroAtomState Solving out of the plasma state from_legacy_plasma

* Added docstrings to MacroAtomSolver

* Added docstrings to macro_atom base.py

* Fixed up docstrings in macroatom_sovler

* Added check to verify macroatom_state is None when line_interaction_type is scatter

* removed some comments, fixed up OpacitySolver initialization

* Moved handling of the opacity solver over to the simulation base

* Fixed typo

* Grabbed the line interaction type from the right place (makes me think we should just to a to_numba method here or not require it)

* Changed 'non_markov_transition...' to 'transtiition...'

* decoupled macro_atom from the opacity solver

* Fixed typo

* Updated the to_numba tests

* UPdated docstrings/removed unused imports

tardis/opacities/macro_atom/base.py

@@ -9,7 +9,231 @@
 logger = logging.getLogger(__name__)
 
 
-class TransitionProbabilities(ProcessingPlasmaProperty):
+def initialize_transition_probabilities(
+    atomic_data,
+):
+    """Convienience Function for initializing the transition probabilities
+
+    Parameters
+    ----------
+    atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+
+    Returns
+    -------
+    dict
+        "transition_probability_coef" : np.ndarray
+            Reshaped macro atom transition probabilities
+        "block_references": np.ndarray
+            macro atom block references
+    """
+    macro_atom_data = get_macro_atom_data(atomic_data)
+    (
+        transition_up_filter,
+        transition_up_line_filter,
+        block_references,
+    ) = initialize_macro_atom_transition_type_filters(
+        atomic_data, macro_atom_data
+    )
+    transition_probability_coef = get_transition_probability_coefs(
+        macro_atom_data
+    )
+
+    return {
+        "transition_probability_coef": transition_probability_coef,
+        "block_references": block_references,
+    }
+
+
+def calculate_transition_probabilities(
+    atomic_data,
+    beta_sobolev,
+    j_blues,
+    stimulated_emission_factor,
+    tau_sobolevs,
+    transition_probability_coef,
+    block_references,
+    normalize=True,
+):
+    """Computes transition probabilities and provides them as a pd.DataFrame
+
+    Parameters
+    ----------
+    atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+    beta_sobolev : pd.DataFrame
+        Beta Sobolevs
+    j_blues : pd.DataFrame
+        mean intensity
+    stimulated_emission_factor : np.ndarray
+        Stimulated Emission Factors
+    tau_sobolev : pd.DataFrame
+            Expansion Optical Depths
+    transition_probability_coef : np.ndarray
+        Reshaped macro atom transition probabilities
+    block_references : np.ndarray
+        macro atom block references
+    normalize : bool
+        Whether or not to normalize the transition probabilities to unity
+
+    Returns
+    -------
+    pd.DataFrame
+        transition probabilities
+    """
+    # I wonder why?
+    # Not sure who wrote this but the answer is that when the plasma is
+    # first initialised (before the first iteration, without temperature
+    # values etc.) there are no j_blues values so this just prevents
+    # an error. Aoife.
+    if len(j_blues) == 0:
+        return None
+    macro_atom_data = get_macro_atom_data(atomic_data)
+
+    transition_probabilities = calculate_transition_probability(
+        macro_atom_data,
+        beta_sobolev,
+        j_blues,
+        stimulated_emission_factor,
+        transition_probability_coef,
+        block_references,
+        normalize,
+    )
+    transition_probabilities = pd.DataFrame(
+        transition_probabilities,
+        index=macro_atom_data.transition_line_id,
+        columns=tau_sobolevs.columns,
+    )
+    return transition_probabilities
+
+
+def calculate_transition_probability(
+    macro_atom_data,
+    beta_sobolev,
+    j_blues,
+    stimulated_emission_factor,
+    transition_probability_coef,
+    block_references,
+    normalize,
+):
+    """Calculate the transition probabilities using optimized functions
+    Parameters
+    ----------
+    macro_atom_data : pd.DataFrame
+        Macro Atom Data
+    beta_sobolev : pd.DataFrame
+        Beta Sobolevs
+    j_blues : pd.DataFrame
+        mean intensity
+    stimulated_emission_factor : np.ndarray
+        Stimulated Emission Factors
+    transition_probability_coef : np.ndarray
+        Reshaped macro atom transition probabilities
+    block_references : np.ndarray
+        macro atom block references
+    normalize : bool
+        Whether or not to normalize the transition probabilities to unity
+
+    Returns
+    -------
+    np.ndarray
+        transition probabilities
+    """
+    transition_probabilities = np.empty(
+        (transition_probability_coef.shape[0], beta_sobolev.shape[1])
+    )
+    # trans_old = self.calculate_transition_probabilities(macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor)
+    transition_type = macro_atom_data.transition_type.values
+    lines_idx = macro_atom_data.lines_idx.values
+    tpos = macro_atom_data.transition_probability.values
+    util.fast_calculate_transition_probabilities(
+        tpos,
+        beta_sobolev.values,
+        j_blues.values,
+        stimulated_emission_factor,
+        transition_type,
+        lines_idx,
+        block_references,
+        transition_probabilities,
+        normalize,
+    )
+    return transition_probabilities
+
+
+def initialize_macro_atom_transition_type_filters(atomic_data, macro_atom_data):
+    """Get the filters and block references from the macro atom
+
+    Parameters
+    ----------
+    atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+    macro_atom_data : pd.DataFrame
+        Macro Atom Data
+
+    Returns
+    -------
+    np.ndarray
+        Mask where the transition type is 1
+    np.ndarray
+        index of lines at these locations
+    pd.ndarray
+        macro atom block references
+    """
+    transition_up_filter = macro_atom_data.transition_type.values == 1
+    transition_up_line_filter = macro_atom_data.lines_idx.values[
+        transition_up_filter
+    ]
+    block_references = np.hstack(
+        (
+            atomic_data.macro_atom_references.block_references,
+            len(macro_atom_data),
+        )
+    )
+
+    return transition_up_filter, transition_up_line_filter, block_references
+
+
+def get_transition_probability_coefs(macro_atom_data):
+    """Coefficients of the transition probabilities
+
+    Parameters
+    ----------
+    macro_atom_data : pd.DataFrame
+        Macro Atom Data
+
+    Returns
+    -------
+    np.ndarray
+        Reshaped macro atom transition probabilities
+    """
+    return macro_atom_data.transition_probability.values[np.newaxis].T
+
+
+def get_macro_atom_data(atomic_data):
+    """Get the macro atom data from the atomic data
+
+    Parameters
+    ----------
+    atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+
+    Returns
+    -------
+    pd.DataFrame
+        The macro atom data in the plasma
+    """
+    try:
+        return atomic_data.macro_atom_data
+    except:
+        logger.debug(
+            "Macro Atom Data was not found. Instead returning All Macro Atom Data"
+        )
+        return atomic_data.macro_atom_data_all
+
+
+class TransitionProbabilities(
+    ProcessingPlasmaProperty
+):  # Base MacroAtom Property
     """
     Attributes
     ----------
@@ -137,5 +361,7 @@ def _get_macro_atom_data(atomic_data):
             return atomic_data.macro_atom_data_all
 
 
-class NonMarkovChainTransitionProbabilities(TransitionProbabilities):
+class NonMarkovChainTransitionProbabilities(
+    TransitionProbabilities
+):  # Continuum Only
     outputs = ("non_markov_transition_probabilities",)

tardis/opacities/macro_atom/macroatom_solver.py

@@ -0,0 +1,133 @@
+from tardis.opacities.macro_atom.base import (
+    calculate_transition_probabilities,
+    initialize_transition_probabilities,
+)
+from tardis.opacities.macro_atom.macroatom_state import MacroAtomState
+
+
+class MacroAtomSolver(object):
+
+    initialize: bool = True
+    normalize: bool = True
+
+    def __init__(self, initialize=True, normalize=True):
+        """Solver class for Macro Atom related opacities
+
+        Parameters
+        ----------
+        initialize: bool
+            Whether or not to initialize the transition probabilitiy coefficients and block references when solving the first time (default True)
+        normalize: bool
+            Whether or not to normalize the transition probabilities to unity. Default True
+        """
+
+        self.initialize = initialize
+        self.normalize = normalize
+
+    def initialize_transition_probabilities(self, atomic_data):
+        """initialize the transition probability coefficients and block references when solving the first time
+
+        Parameters
+        ----------
+        atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+        """
+
+        coef_and_block_ref = initialize_transition_probabilities(
+            atomic_data
+        )
+        self.transition_probability_coef = coef_and_block_ref[
+            "transition_probability_coef"
+        ]
+        self.block_references = coef_and_block_ref["block_references"]
+        self.initialize = False
+
+    def solve_transition_probabilities(
+        self,
+        atomic_data,
+        legacy_plasma,
+        tau_sobolev,
+        stimulated_emission_factor,
+    ):
+        """Solve the basic transition probabilities for the macroatom
+
+        Parameters
+        ----------
+        atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+        legacy_plasma : tarids.plasma.BasePlasma
+            legacy base plasma
+        tau_sobolev : pd.DataFrame
+            Expansion Optical Depths
+        stimulated_emission_factor : np.ndarray
+
+        Returns
+        -------
+        pd.DataFrame
+            Transition Probabilities
+        """
+        if self.initialize:
+            self.initialize_transition_probabilities(atomic_data)
+
+        transition_probabilities = (
+            calculate_transition_probabilities(
+                atomic_data,
+                legacy_plasma.beta_sobolev,
+                legacy_plasma.j_blues,
+                stimulated_emission_factor,
+                tau_sobolev,
+                self.transition_probability_coef,
+                self.block_references,
+                normalize=self.normalize,
+            )
+        )
+
+        return transition_probabilities
+
+    def solve(
+        self,
+        legacy_plasma,
+        atomic_data,
+        tau_sobolev,
+        stimulated_emission_factor,
+    ):
+        """Solved the Macro Atom State
+
+        Parameters
+        ----------
+        legacy_plasma : tarids.plasma.BasePlasma
+            legacy base plasma
+        atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+        tau_sobolev : pd.DataFrame
+            Expansion Optical Depths
+        stimulated_emission_factor : pd.DataFrame
+
+        Returns
+        -------
+        tardis.opacities.macroatom_state.MacroAtomState
+            State of the macro atom ready to be placed into the OpacityState
+        """
+
+        transition_probabilities = (
+            self.solve_transition_probabilities(
+                atomic_data,
+                legacy_plasma,
+                tau_sobolev,
+                stimulated_emission_factor,
+            )
+        )
+
+        macro_block_references = atomic_data.macro_atom_references[
+            "block_references"
+        ]
+        macro_atom_info = legacy_plasma.atomic_data.macro_atom_data
+
+        return MacroAtomState(
+            transition_probabilities,
+            macro_atom_info["transition_type"],
+            macro_atom_info["destination_level_idx"],
+            macro_atom_info["lines_idx"],
+            macro_block_references,
+            legacy_plasma.atomic_data.lines_upper2macro_reference_idx,
+        )