Macroatom restructure

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 @@
             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,
+        )