Remove tau sobolevs from plasma (#2874)

* Create legacy plasma graph

* Set up legacy plasma collections and factory

* Legacy macro atom transition probs

* Modifications to the formal integral, and simple workflow testbed

* Partial fix for the formal integral

* Drop macro atom data from plasma

* Fixes formal integral tests

* Resolve macro atom solver comments

* Address opacity state and solver comments

* Use restructured sobolev calculations for legacy plasma modules

* Fixes formal integral test failures

* Simplify workflow radiation field setup

* Fixes standard workflow

* Adds tests for the sobolev calculations

* Fixes CUDA formal integral tests

* Remove plasma atomic data variable assignment

* Clean up macro atom a bit

tardis/opacities/macro_atom/base.py

@@ -142,7 +142,6 @@ def calculate_transition_probability(
     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
@@ -287,7 +286,6 @@ def _calculate_transition_probability(
         transition_probabilities = np.empty(
             (self.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
@@ -304,14 +302,6 @@ def _calculate_transition_probability(
         )
         return transition_probabilities
 
-    def calculate_transition_probabilities(
-        self, macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor
-    ):
-        transition_probabilities = self.prepare_transition_probabilities(
-            macro_atom_data, beta_sobolev, j_blues, stimulated_emission_factor
-        )
-        return transition_probabilities
-
     def initialize_macro_atom_transition_type_filters(
         self, atomic_data, macro_atom_data
     ):

tardis/opacities/macro_atom/macroatom_solver.py

@@ -5,8 +5,7 @@
 from tardis.opacities.macro_atom.macroatom_state import MacroAtomState
 
 
-class MacroAtomSolver(object):
-
+class MacroAtomSolver:
     initialize: bool = True
     normalize: bool = True
 
@@ -20,7 +19,6 @@ def __init__(self, initialize=True, normalize=True):
         normalize: bool
             Whether or not to normalize the transition probabilities to unity. Default True
         """
-
         self.initialize = initialize
         self.normalize = normalize
 
@@ -32,15 +30,14 @@ def initialize_transition_probabilities(self, atomic_data):
         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(
+    def solve_legacy_transition_probabilities(
         self,
         atomic_data,
         legacy_plasma,
@@ -80,12 +77,57 @@ def solve_transition_probabilities(
 
         return transition_probabilities
 
+    def solve_transition_probabilities(
+        self,
+        atomic_data,
+        mean_intensities,
+        tau_sobolev,
+        beta_sobolev,
+        stimulated_emission_factor,
+    ):
+        """Solve the basic transition probabilities for the macroatom
+
+        Parameters
+        ----------
+        atomic_data : tardis.io.atom_data.AtomData
+            Atomic Data
+        mean_intensities : pd.DataFrame
+            Mean intensity of the radiation field for each shell
+        tau_sobolev : pd.DataFrame
+            Expansion Optical Depths
+        beta_sobolev : pd.DataFrame
+            Modified 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,
+            beta_sobolev,
+            mean_intensities,
+            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,
+        beta_sobolev=None,
+        legacy_mode=True,
     ):
         """Solved the Macro Atom State
 
@@ -104,13 +146,23 @@ def solve(
         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,
-        )
+        if legacy_mode:
+            transition_probabilities = (
+                self.solve_legacy_transition_probabilities(
+                    atomic_data,
+                    legacy_plasma,
+                    tau_sobolev,
+                    stimulated_emission_factor,
+                )
+            )
+        else:
+            transition_probabilities = self.solve_transition_probabilities(
+                atomic_data,
+                legacy_plasma.j_blues,
+                tau_sobolev,
+                beta_sobolev,
+                stimulated_emission_factor,
+            )
 
         macro_block_references = atomic_data.macro_atom_references[
             "block_references"

tardis/opacities/opacity_solver.py

@@ -1,18 +1,23 @@
-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
 
+from tardis.opacities.opacity_state import (
+    OpacityState,
+)
+from tardis.opacities.tau_sobolev import (
+    calculate_beta_sobolev,
+    calculate_sobolev_line_opacity,
+)
 
-class OpacitySolver(object):
 
+class OpacitySolver:
     line_interaction_type: str = "scatter"
     disable_line_scattering: bool = False
 
     def __init__(
-        self, line_interaction_type="scatter", disable_line_scattering=False
+        self,
+        line_interaction_type="scatter",
+        disable_line_scattering=False,
     ):
         """Solver class for opacities
 
@@ -22,11 +27,10 @@ def __init__(
             "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:
+    def legacy_solve(self, plasma) -> OpacityState:
         """
         Solves the opacity state
 
@@ -39,33 +43,64 @@ def solve(self, legacy_plasma) -> OpacityState:
         -------
         OpacityState
         """
-        atomic_data = legacy_plasma.atomic_data
+        if self.disable_line_scattering:
+            tau_sobolev = pd.DataFrame(
+                np.zeros(
+                    (
+                        plasma.atomic_data.lines.shape[0],  # number of lines
+                        plasma.number_density.shape[1],  # number of shells
+                    ),
+                    dtype=np.float64,
+                ),
+                index=plasma.atomic_data.lines.index,
+            )
+        else:
+            tau_sobolev = calculate_sobolev_line_opacity(
+                plasma.atomic_data.lines,
+                plasma.level_number_density,
+                plasma.time_explosion,
+                plasma.stimulated_emission_factor,
+            )
+
+        opacity_state = OpacityState.from_legacy_plasma(plasma, tau_sobolev)
+
+        return opacity_state
+
+    def solve(self, 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.number_density.shape[
-                            1
-                        ],  # number of shells
+                        plasma.atomic_data.lines.shape[0],  # number of lines
+                        plasma.number_density.shape[1],  # number of shells
                     ),
                     dtype=np.float64,
                 ),
-                index=atomic_data.lines.index,
+                index=plasma.atomic_data.lines.index,
             )
         else:
             tau_sobolev = calculate_sobolev_line_opacity(
-                atomic_data.lines,
-                legacy_plasma.level_number_density,
-                legacy_plasma.time_explosion,
-                legacy_plasma.stimulated_emission_factor,
+                plasma.atomic_data.lines,
+                plasma.level_number_density,
+                plasma.time_explosion,
+                plasma.stimulated_emission_factor,
             )
+            beta_sobolev = calculate_beta_sobolev(tau_sobolev)
 
-        opacity_state = OpacityState.from_legacy_plasma(
-            legacy_plasma, tau_sobolev
+        opacity_state = OpacityState.from_plasma(
+            plasma, tau_sobolev, beta_sobolev
         )
 
         return opacity_state

tardis/opacities/opacity_state.py

@@ -2,10 +2,10 @@
 from numba import float64, int64
 from numba.experimental import jitclass
 
-from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity
-from tardis.transport.montecarlo.configuration import montecarlo_globals
 from tardis.opacities.continuum.continuum_state import ContinuumState
 from tardis.opacities.macro_atom.macroatom_state import MacroAtomState
+from tardis.opacities.tau_sobolev import calculate_sobolev_line_opacity
+from tardis.transport.montecarlo.configuration import montecarlo_globals
 
 
 class OpacityState:
@@ -15,6 +15,7 @@ def __init__(
         t_electrons,
         line_list_nu,
         tau_sobolev,
+        beta_sobolev,
         continuum_state,
     ):
         """
@@ -26,6 +27,7 @@ def __init__(
         t_electrons : numpy.ndarray
         line_list_nu : pd.DataFrame
         tau_sobolev : pd.DataFrame
+        beta_sobolev : pd.DataFrame
         continuum_state: tardis.opacities.continuum.continuum_state.ContinuumState
         """
         self.electron_density = electron_density
@@ -34,6 +36,8 @@ def __init__(
 
         self.tau_sobolev = tau_sobolev
 
+        self.beta_sobolev = beta_sobolev
+
         # Continuum Opacity Data
         self.continuum_state = continuum_state
 
@@ -53,7 +57,38 @@ def from_legacy_plasma(cls, plasma, tau_sobolev):
         -------
         OpacityStatePython
         """
+        if hasattr(plasma, "photo_ion_cross_sections"):
+            continuum_state = ContinuumState.from_legacy_plasma(plasma)
+        else:
+            continuum_state = None
+
+        return cls(
+            plasma.electron_densities,
+            plasma.t_electrons,
+            plasma.atomic_data.lines.nu,
+            tau_sobolev,
+            plasma.beta_sobolev,
+            continuum_state,
+        )
 
+    @classmethod
+    def from_plasma(cls, plasma, tau_sobolev, beta_sobolev):
+        """
+        Generates an OpacityStatePython object from a tardis BasePlasma
+
+        Parameters
+        ----------
+        plasma : tarids.plasma.BasePlasma
+            legacy base plasma
+        tau_sobolev : pd.DataFrame
+            Expansion Optical Depths
+        beta_sobolev : pd.DataFrame
+            Modified expansion Optical Depths
+
+        Returns
+        -------
+        OpacityStatePython
+        """
         if hasattr(plasma, "photo_ion_cross_sections"):
             continuum_state = ContinuumState.from_legacy_plasma(plasma)
         else:
@@ -66,6 +101,7 @@ def from_legacy_plasma(cls, plasma, tau_sobolev):
             plasma.t_electrons,
             atomic_data.lines.nu,
             tau_sobolev,
+            beta_sobolev,
             continuum_state,
         )