[Restructure/plasma] estimator cleanup (#2675)

* change numberdensity to input

* fixed number density

* some fixes

* removing density

* remove atomic and isotope mass

* add isotopic_number_density

* add opacities package

* Update imports in property_collections.py, base.py, test_numba_interface.py, transport_montecarlo_numba_interface.py, conftest.py, formal_integral.py, base.py, and macro_atom.py

* Add calculate_transition_probabilities function to util.py in macro_atom package

* Add calculate_transition_probabilities function to util.py in macro_atom package

* Remove unused imports and update plasma properties

* add __init__ to macroatom

* blackify tardis

* blackified

* chore: Update imports and remove unused code

* chore: Add PlanckRadiationField and DilutePlanckRadiationField classes

* chore: Update imports and remove unused code

* removed density

* ruff output

* cleanup and adding object mode

* starting to make radiation_field a thing

* switched over to old tau_sobolev calculation

* renamed function to indicate numba use

* address comments

* added dilute planckian radiation field

* refactor: Convert species lists to proper format in assemble_plasma function

* moved radiation field into plasma. Resulting in some renames

* some fixes

* black montecarlo

* chore: Initialize atom data and simulation state in `initialization.py`

* updating the documentation

* feat: Add EstimatedRadiationFieldProperties class for Monte Carlo estimators

* trying to get rid of j_blues in plasma with MC restructure

* completely restructure j_blues. Estimators and all.

* cleanup for the restructure

* remove parse_input.py

* chore: Refactor radiation field configuration parsing and state creation

Refactor the `parse_radiation_field_configuration.py` module to improve code organization and readability. Update the import statements to reflect the changes in the module structure. Replace the deprecated `DiluteBlackBodyRadiationFieldState` class with the new `DilutePlanckianRadiationField` class from the `tardis.plasma.radiation_field` module. This change ensures consistency and compatibility with the latest codebase.

Also, update the `tardis/simulation/base.py` module to import the `DilutePlanckianRadiationField` class from the `tardis.plasma.radiation_field` module. This change ensures that the correct class is used for creating the radiation field in the simulation.

* revert astropy_helpers

* remove astropy_helpers

* removed test.txt

* blackified code

* cleanup simulation from merges

* fix

* remove unused Input

* added description

* blackiefied codebase

* cleanup of branch

* blackify code

* chore: Refactor atom data parsing and simulation state initialization

* restructure logger

* working on continuum radfield properties

* Refactor continuum processes module structure

* Refactor opacities module structure

* cleanup from restructure

* cleanup

* clean up

* more cleanup

* cleanup standard_plasmas.py

* working nlte ionizations

* some more cleanup

* fix benchmarks

* blackify

* reverse the import pygraphviz

* fix docstrings

* Refactor code to address comments

benchmarks/opacities_opacity.py

@@ -6,6 +6,7 @@
 
 import tardis.opacities.opacities as calculate_opacity
 from benchmarks.benchmark_base import BenchmarkBase
+from tardis.opacities.opacities import compton_opacity_calculation
 
 
 class BenchmarkMontecarloMontecarloNumbaOpacities(BenchmarkBase):
@@ -28,9 +29,7 @@ class BenchmarkMontecarloMontecarloNumbaOpacities(BenchmarkBase):
         }
     )
     def time_compton_opacity_calculation(self, electron_number_density, energy):
-        calculate_opacity.compton_opacity_calculation(
-            energy, electron_number_density
-        )
+        compton_opacity_calculation(energy, electron_number_density)
 
     @parameterize(
         {

tardis/io/atom_data/base.py

@@ -9,7 +9,7 @@
 
 from tardis import constants as const
 from tardis.io.atom_data.util import resolve_atom_data_fname
-from tardis.plasma.properties.continuum_processes import (
+from tardis.plasma.properties.continuum_processes.rates import (
     get_ground_state_multi_index,
 )
 

tardis/opacities/continuum/bound_free.py

@@ -0,0 +1,36 @@
+from tardis.plasma.exceptions import PlasmaException
+from tardis.plasma.properties.base import ProcessingPlasmaProperty
+
+
+class BoundFreeOpacity(ProcessingPlasmaProperty):
+    """
+    Attributes
+    ----------
+    chi_bf : pandas.DataFrame, dtype float
+        Bound-free opacity corrected for stimulated emission.
+    """
+
+    outputs = ("chi_bf",)
+    latex_name = (r"\chi^{\textrm{bf}}",)
+
+    def calculate(
+        self,
+        photo_ion_cross_sections,
+        t_electrons,
+        phi_ik,
+        level_number_density,
+        lte_level_number_density,
+        boltzmann_factor_photo_ion,
+    ):
+        cross_section = photo_ion_cross_sections["x_sect"].values
+
+        n_i = level_number_density.loc[photo_ion_cross_sections.index]
+        lte_n_i = lte_level_number_density.loc[photo_ion_cross_sections.index]
+        chi_bf = (n_i - lte_n_i * boltzmann_factor_photo_ion).multiply(
+            cross_section, axis=0
+        )
+
+        num_neg_elements = (chi_bf < 0).sum().sum()
+        if num_neg_elements:
+            raise PlasmaException("Negative values in bound-free opacity.")
+        return chi_bf

tardis/opacities/macro_atom/continuum_processes/collisional_ion_trans_prob.py

@@ -0,0 +1,94 @@
+import pandas as pd
+
+from tardis import constants as const
+from tardis.plasma.properties.base import (
+    TransitionProbabilitiesProperty,
+)
+from tardis.plasma.properties.continuum_processes.rates import (
+    IndexSetterMixin,
+)
+
+H = const.h.cgs.value
+
+
+class RawCollIonTransProbs(TransitionProbabilitiesProperty, IndexSetterMixin):
+    """
+    Attributes
+    ----------
+    p_coll_ion : pandas.DataFrame, dtype float
+        The unnormalized transition probabilities for
+        collisional ionization.
+    p_coll_recomb : pandas.DataFrame, dtype float
+        The unnormalized transition probabilities for
+        collisional recombination.
+    cool_rate_coll_ion : pandas.DataFrame, dtype float
+        The collisional ionization cooling rates of the electron gas.
+    """
+
+    outputs = ("p_coll_ion", "p_coll_recomb", "cool_rate_coll_ion")
+    transition_probabilities_outputs = (
+        "p_coll_ion",
+        "p_coll_recomb",
+        "cool_rate_coll_ion",
+    )
+    latex_name = (
+        r"p^{\textrm{coll ion}}",
+        r"p^{\textrm{coll recomb}}",
+        r"C^{\textrm{ion}}",
+    )
+
+    def calculate(
+        self,
+        coll_ion_coeff,
+        coll_recomb_coeff,
+        nu_i,
+        photo_ion_idx,
+        electron_densities,
+        energy_i,
+        level_number_density,
+    ):
+        p_coll_ion = coll_ion_coeff.multiply(energy_i, axis=0)
+        p_coll_ion = p_coll_ion.multiply(electron_densities, axis=1)
+        p_coll_ion = self.set_index(p_coll_ion, photo_ion_idx, reverse=False)
+
+        coll_recomb_rate = coll_recomb_coeff.multiply(
+            electron_densities, axis=1
+        )  # The full rate is obtained from this by multiplying by the
+        # electron density and ion number density.
+        p_recomb_deactivation = coll_recomb_rate.multiply(nu_i, axis=0) * H
+        p_recomb_deactivation = self.set_index(
+            p_recomb_deactivation, photo_ion_idx, transition_type=-1
+        )
+        p_recomb_deactivation = p_recomb_deactivation.groupby(level=[0]).sum()
+        index_dd = pd.MultiIndex.from_product(
+            [p_recomb_deactivation.index.values, ["k"], [0]],
+            names=list(photo_ion_idx.columns) + ["transition_type"],
+        )
+        p_recomb_deactivation = p_recomb_deactivation.set_index(index_dd)
+
+        p_recomb_internal = coll_recomb_rate.multiply(energy_i, axis=0)
+        p_recomb_internal = self.set_index(
+            p_recomb_internal, photo_ion_idx, transition_type=0
+        )
+        p_coll_recomb = pd.concat([p_recomb_deactivation, p_recomb_internal])
+
+        cool_rate_coll_ion = (coll_ion_coeff * electron_densities).multiply(
+            nu_i * H, axis=0
+        )
+        level_lower_index = coll_ion_coeff.index
+        cool_rate_coll_ion = (
+            cool_rate_coll_ion
+            * level_number_density.loc[level_lower_index].values
+        )
+        cool_rate_coll_ion = self.set_index(
+            cool_rate_coll_ion, photo_ion_idx, reverse=False
+        )
+        cool_rate_coll_ion = cool_rate_coll_ion.groupby(
+            level="destination_level_idx"
+        ).sum()
+        ion_cool_index = pd.MultiIndex.from_product(
+            [["k"], cool_rate_coll_ion.index.values, [0]],
+            names=list(photo_ion_idx.columns) + ["transition_type"],
+        )
+        cool_rate_coll_ion = cool_rate_coll_ion.set_index(ion_cool_index)
+        return p_coll_ion, p_coll_recomb, cool_rate_coll_ion

tardis/opacities/macro_atom/transition_probabilities.py

@@ -5,7 +5,7 @@
 from scipy import sparse as sp
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
-from tardis.plasma.properties.continuum_processes import (
+from tardis.plasma.properties.continuum_processes.rates import (
     get_ground_state_multi_index,
 )
 from tardis.transport.montecarlo.macro_atom import (

tardis/opacities/opacities.py

@@ -278,7 +278,8 @@ def compton_opacity_partial(energy, fraction):
 
 @njit(**njit_dict_no_parallel)
 def compton_opacity_calculation(energy, electron_density):
-    """Calculate the Compton scattering opacity for a given energy
+    """
+    Calculate the Compton scattering opacity for a given energy
     (Rybicki & Lightman, 1979)
 
     $
@@ -364,7 +365,8 @@ def photoabsorption_opacity_calculation(
 def photoabsorption_opacity_calculation_kasen(
     energy, number_density, proton_count
 ):
-    """Calculates photoabsorption opacity for a given energy
+    """
+    Calculates photoabsorption opacity for a given energy
     Approximate treatment from Kasen et al. (2006)
 
     Parameters

tardis/plasma/base.py

@@ -1,24 +1,36 @@
-import os
-import re
+import dataclasses
+import fileinput
 import logging
+import re
 import tempfile
-import fileinput
 
 import networkx as nx
 
-from tardis.plasma.exceptions import PlasmaMissingModule, NotInitializedModule
-from tardis.plasma.properties.base import *
 from tardis.io.util import PlasmaWriterMixin
+from tardis.plasma.exceptions import NotInitializedModule, PlasmaMissingModule
+from tardis.plasma.properties.base import *
 
 logger = logging.getLogger(__name__)
 
 
+@dataclasses.dataclass(frozen=True)
+class PlasmaSolverSettings:
+    RADIATIVE_RATES_TYPE: str = "blackbody"
+
+
 class BasePlasma(PlasmaWriterMixin):
 
     outputs_dict = {}
     hdf_name = "plasma"
 
-    def __init__(self, plasma_properties, property_kwargs=None, **kwargs):
+    def __init__(
+        self,
+        plasma_properties,
+        plasma_solver_settings,
+        property_kwargs=None,
+        **kwargs,
+    ):
+        self.plasma_solver_settings = plasma_solver_settings
         self.outputs_dict = {}
         self.input_properties = []
         self.plasma_properties = self._init_properties(
@@ -63,7 +75,6 @@ def _build_graph(self):
         :param plasma_modules:
         :return:
         """
-
         self.graph = nx.DiGraph()
         # Adding all nodes
         self.graph.add_nodes_from(
@@ -200,8 +211,7 @@ def freeze(self, *args):
         for key in args:
             if key not in self.outputs_dict:
                 raise PlasmaMissingModule(
-                    "Trying to freeze property {0}"
-                    " that is unavailable".format(key)
+                    f"Trying to freeze property {key}" " that is unavailable"
                 )
             self.outputs_dict[key].frozen = True
 
@@ -224,8 +234,7 @@ def thaw(self, *args):
         for key in args:
             if key not in self.outputs_dict:
                 raise PlasmaMissingModule(
-                    "Trying to thaw property {0}"
-                    " that is unavailable".format(key)
+                    f"Trying to thaw property {key}" " that is unavailable"
                 )
             self.outputs_dict[key].frozen = False
 
@@ -249,7 +258,6 @@ def _resolve_update_list(self, changed_properties):
             : list
             all affected modules.
         """
-
         descendants_ob = []
 
         for plasma_property in changed_properties:
@@ -284,11 +292,10 @@ def write_to_dot(self, fname, args=None, latex_label=True):
             enables/disables writing LaTeX equations and
             edge labels into the file.
         """
-
         try:
             import pygraphviz
         except:
-            logger.warn(
+            logger.warning(
                 "pygraphviz missing. Plasma graph will not be " "generated."
             )
             return
@@ -303,7 +310,7 @@ def write_to_dot(self, fname, args=None, latex_label=True):
                     ] = f"\\\\textrm{{{node}: }}"
                     node_list = self.plasma_properties_dict[node]
                     formulae = node_list.latex_formula
-                    for output in range(0, len(formulae)):
+                    for output in range(len(formulae)):
                         formula = formulae[output]
                         label = formula.replace("\\", "\\\\")
                         print_graph.nodes[str(node)]["label"] += label
@@ -341,7 +348,7 @@ def write_to_dot(self, fname, args=None, latex_label=True):
                 )
 
         if args is not None:
-            with open(fname, "r") as file:
+            with open(fname) as file:
                 lines = file.readlines()
 
             for newline in args:
@@ -374,7 +381,7 @@ def write_to_tex(self, fname_graph, scale=0.5, args=None, latex_label=True):
         try:
             import dot2tex
         except:
-            logger.warn(
+            logger.warning(
                 "dot2tex missing. Plasma graph will not be " "generated."
             )
             return
@@ -383,7 +390,7 @@ def write_to_tex(self, fname_graph, scale=0.5, args=None, latex_label=True):
 
         self.write_to_dot(temp_fname, args=args, latex_label=latex_label)
 
-        with open(temp_fname, "r") as file:
+        with open(temp_fname) as file:
             dot_string = file.read().replace("\\\\", "\\")
 
         texcode = dot2tex.dot2tex(

tardis/plasma/properties/__init__.py

@@ -11,7 +11,6 @@
 from tardis.plasma.properties.general import *
 from tardis.plasma.properties.helium_nlte import *
 from tardis.plasma.properties.ion_population import *
-from tardis.plasma.properties.j_blues import *
 from tardis.plasma.properties.level_population import *
 from tardis.plasma.properties.nlte import *
 from tardis.plasma.properties.nlte_rate_equation_solver import *

tardis/plasma/properties/atomic.py

@@ -15,7 +15,7 @@
     HiddenPlasmaProperty,
     ProcessingPlasmaProperty,
 )
-from tardis.plasma.properties.continuum_processes import (
+from tardis.plasma.properties.continuum_processes.rates import (
     A0,
     BETA_COLL,
     K_B,

tardis/plasma/properties/continuum_processes/__init__.py

@@ -0,0 +1,9 @@
+from tardis.plasma.properties.continuum_processes.photo_ion_rate_coeff import (
+    PhotoIonRateCoeff,
+)
+from tardis.plasma.properties.continuum_processes.rates import *
+from tardis.plasma.properties.continuum_processes.recomb_rate_coeff import (
+    StimRecombRateFactor,
+    SpontRecombRateCoeff,
+    StimRecombRateCoeff,
+)