Restructure/detailed balance foundation (tardis-sn#2770)

* 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

* start of assemble plasma cleanup

* cleanup standard_plasmas.py

* updated tests

* some more cleanup

* fix benchmarks

* cleanup assembly

* cleanup assembly

* working on the restructure

* slowly fixing the assembly module

* blackify

* reverse the import pygraphviz

* fix docstrings

* fixed all plasma

* slow fixes

* Refactor recomb_rate_coeff.py and test_continuum_property_solver.py

* fixing assembly

* restructure the read in

* fixup plasma assemble to be clean

* fix shell info widget

* fix the widgets

* fixing hopefully last bugs

* slowly fix up assembly

* working on getting the notebook running

* Refactor code to address comments

* feat: Add RadiativeRatesSolver class for detailed balance rates calculation

* feat: Add solve method to RadiativeRatesSolver class

* add rates ipynb

* remove abundance

* feat: Refactor code to improve readability and maintainability

* first commit of the rates solver working

* fix up detailed balance foundation

* further work on including the rates

* including the new rates structure

* getting the rates ready

* feat: Refactor collision_rates module to improve code organization and readability

* further work on collisional strengths

* further work on the collisional cross sections

* further work o nthe collisional rates

* feat: Add ThermalCollisionalRateSolver class for calculating thermal collisional rates

* further fixes

* further testing

* refactor: Refactor collisional rates calculation in ThermalCollisionalRateSolver

* finished collisional_rates

* Fixes collisional rates test so it can run (still fails)

* Correct de-excitation rate calculation

* Add Chianti upsilon solver

* Docstrings and better variable names

* Fix rename error

* Comparison notebook between Chianti, CMFGEN et al.

* Add reference data and config, save notebook with plots

* Chianti test beginnings

* Working tests

* Up to date comparison notebook

* Remove irrelevant change

* Corrected changes to the plasma solver factory

* Partially fixes tests

* Slightly dirty fix for electron densities not refreshing properly

* Fix NLTE tests by passing electron densities as an option input to the population solvers

* Ruff formatting

---------

Co-authored-by: Andrew Fullard <andrewgfullard@gmail.com>

docs/physics/plasma/detailed_balance/test_continuum_template_wkerzen_rate_coeffs.yml

@@ -0,0 +1,64 @@
+# Example YAML configuration for TARDIS
+tardis_config_version: v1.0
+
+supernova:
+  time_explosion: 16 day
+
+atom_data: TestNLTE_He_Ti.h5
+
+model:
+  structure:
+    type: specific
+    velocity:
+      start: 5700 km/s
+      stop: 12500 km/s
+      num: 10
+    density:
+      type : power_law
+      time_0: 16.0 day
+      rho_0: 1.3636e-14 g/cm^3 #1.948e-14 g/cm^3
+      v_0: 8000 km/s
+      exponent: -10
+
+  abundances:
+    type: uniform
+    H: 1.0
+
+plasma:
+  initial_t_inner: 9000 K
+  ionization: nebular
+  excitation: dilute-lte
+  radiative_rates_type: dilute-blackbody
+  line_interaction_type: macroatom
+  #nlte:
+   #   species:
+   #   - H I
+  continuum_interaction:
+      species:
+      - H I
+  nlte_ionization_species:
+    - H I
+  nlte_excitation_species:
+    - H I
+
+montecarlo:
+  seed: 23111963
+  no_of_packets: 500000
+  iterations: 1
+  nthreads: 1
+
+  last_no_of_packets: 100000
+  no_of_virtual_packets: 0
+
+  convergence_strategy:
+    type: damped
+    damping_constant: 0.5
+    threshold: 0.05
+    fraction: 0.8
+    hold_iterations: 3
+
+
+spectrum:
+  start: 800 angstrom
+  stop: 10000 angstrom
+  num: 4000

tardis/plasma/assembly/base.py

@@ -55,81 +55,24 @@ def map_species_from_string(species):
     return [species_string_to_tuple(spec) for spec in species]
 
 
-class PlasmaSolverFactory:
-    """Factory class for creating plasma solvers.
-
-    atom_data : object
-        Object containing atomic data.
-    selected_atomic_numbers : list
-        List of selected atomic numbers.
-
-    Attributes
-    ----------
-    excitation_analytical_approximation : str
-        Analytical approximation for excitation (default: "lte").
-    ionization_analytical_approximation : str
-        Analytical approximation for ionization (default: "lte").
-    nebular_ionization_delta_treatment : tuple
-        Species to use for the delta_treatment in nebular ionization ML93 (default: ()).
-    link_t_rad_t_electron : float
-        Link between t_rad and t_electron (default: 1.0).
-    radiative_rates_type : str
-        Type of radiative rates (default: "dilute-blackbody").
-    delta_treatment : float or None
-        Delta treatment (default: None).
-    legacy_nlte_species : list
-        List of legacy non-LTE species (default: []).
-    nlte_excitation_species : list
-        List of non-LTE excitation species (default: []).
-    nlte_ionization_species : list
-        List of non-LTE ionization species (default: []).
-    nlte_solver : str
-        Non-LTE solver (default: "lu").
-        Helium treatment options (default: "none").
-    heating_rate_data_file : str
-        Heating rate data file (default: "none").
-    continuum_interaction_species : list
-        List of continuum interaction species (default: []).
-    enable_adiabatic_cooling : bool
-        Flag for enabling adiabatic cooling (default: False).
-    enable_two_photon_decay : bool
-        Flag for enabling two-photon decay (default: False).
-    line_interaction_type : str
-        Type of line interaction (default: "scatter").
-    plasma_modules : list
-        List of plasma modules (default: []).
-    kwargs : dict
-        Additional keyword arguments (default: {}).
-    property_kwargs : dict
-        Additional keyword arguments for properties (default: {}).
-
-    Methods
-    -------
-    parse_plasma_config(plasma_config)
-    continuum_interaction_species_multi_index()
-        Get the continuum interaction species as a multi-index.
-    setup_factory(config)
-    setup_helium_treatment()
-    setup_legacy_nlte(nlte_config)
-        Set up the non-LTE properties for the legacy species.
-    setup_analytical_approximations()
-        Set up the analytical approximations for excitation and ionization.
-    initialize_j_blues(dilute_planckian_radiation_field, lines_df)
-        Initialize j_blues.
-    """
+def convert_species_to_multi_index(species_strs):
+    return pd.MultiIndex.from_tuples(
+        map_species_from_string(species_strs),
+        names=["atomic_number", "ion_number"],
+    )
+
 
+class PlasmaSolverFactory:
     ## Analytical Approximations
     excitation_analytical_approximation: str = "lte"
     ionization_analytical_approximation: str = "lte"
-    nebular_ionization_delta_treatment: (
-        tuple
-    ) = ()  # species to use for the delta_treatment in nebular ionization ML93
+    nebular_ionization_delta_treatment: tuple  # species to use for the delta_treatment in nebular ionization ML93
 
     link_t_rad_t_electron: float = 1.0
 
     radiative_rates_type: str = "dilute-blackbody"
 
-    delta_treatment: float | None = None
+    delta_treatment = None
 
     ## Statistical Balance Solver
     legacy_nlte_species: list = []
@@ -155,26 +98,19 @@ class PlasmaSolverFactory:
     kwargs: dict = {}
     property_kwargs: dict = {}
 
-    def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None:
-        self.plasma_modules = []
-        self.kwargs = {}
-        self.property_kwargs = {}
-
+    def __init__(
+        self,
+        atom_data,
+        config=None,
+    ) -> None:
         if config is not None:
             self.parse_plasma_config(config.plasma)
         self.atom_data = atom_data
-        self.atom_data.prepare_atom_data(
-            selected_atomic_numbers,
-            line_interaction_type=self.line_interaction_type,
-            continuum_interaction_species=self.continuum_interaction_species_multi_index,
-            nlte_species=self.legacy_nlte_species,
-        )
 
     @property
     def continuum_interaction_species_multi_index(self):
-        return pd.MultiIndex.from_tuples(
-            map_species_from_string(self.continuum_interaction_species),
-            names=["atomic_number", "ion_number"],
+        return convert_species_to_multi_index(
+            self.continuum_interaction_species
         )
 
     def parse_plasma_config(self, plasma_config):
@@ -218,7 +154,7 @@ def parse_plasma_config(self, plasma_config):
             plasma_config.continuum_interaction.enable_two_photon_decay
         )
 
-    def setup_factory(self, config=None):
+    def prepare_factory(self, selected_atomic_numbers, config=None):
         """
         Set up the plasma factory.
 
@@ -227,6 +163,13 @@ def setup_factory(self, config=None):
         config : object, optional
             Configuration object containing plasma settings (default: None).
         """
+        self.atom_data.prepare_atom_data(
+            selected_atomic_numbers,
+            line_interaction_type=self.line_interaction_type,
+            continuum_interaction_species=self.continuum_interaction_species_multi_index,
+            nlte_species=self.legacy_nlte_species,
+        )
+
         self.check_continuum_interaction_species()
 
         self.plasma_modules = basic_inputs + basic_properties
@@ -590,6 +533,7 @@ def assemble(
         dilute_planckian_radiation_field,
         time_explosion,
         electron_densities=None,
+        **kwargs,
     ):
         """
         Assemble the plasma based on the provided parameters and settings.
@@ -622,7 +566,7 @@ def assemble(
             RADIATIVE_RATES_TYPE=self.radiative_rates_type
         )
 
-        kwargs = dict(
+        plasma_assemble_kwargs = dict(
             time_explosion=time_explosion,
             dilute_planckian_radiation_field=dilute_planckian_radiation_field,
             number_density=number_densities,
@@ -633,12 +577,11 @@ def assemble(
             nlte_ionization_species=self.nlte_ionization_species,
             nlte_excitation_species=self.nlte_excitation_species,
         )
-
         if len(self.continuum_interaction_species) > 0:
             initial_continuum_properties = self.initialize_continuum_properties(
                 dilute_planckian_radiation_field
             )
-            kwargs.update(
+            plasma_assemble_kwargs.update(
                 gamma=initial_continuum_properties.photo_ionization_rate_coefficient,
                 bf_heating_coeff_estimator=None,
                 stim_recomb_cooling_coeff_estimator=None,
@@ -647,11 +590,13 @@ def assemble(
 
         if electron_densities is not None:
             electron_densities = pd.Series(electron_densities.cgs.value)
-            self.setup_electron_densities(electron_densities)
-        kwargs["helium_treatment"] = self.helium_treatment
+
+        self.setup_electron_densities(electron_densities)
+        plasma_assemble_kwargs["helium_treatment"] = self.helium_treatment
+        plasma_assemble_kwargs.update(kwargs)
         return BasePlasma(
             plasma_properties=self.plasma_modules,
             property_kwargs=self.property_kwargs,
             plasma_solver_settings=plasma_solver_settings,
-            **kwargs,
+            **plasma_assemble_kwargs,
         )

tardis/plasma/assembly/legacy_assembly.py

@@ -23,10 +23,9 @@ def assemble_plasma(config, simulation_state, atom_data=None):
     atomic_numbers = simulation_state.abundance.index
     plasma_solver_factory = PlasmaSolverFactory(
         atom_data,
-        atomic_numbers,
         config,
     )
-    plasma_solver_factory.setup_factory(config)
+    plasma_solver_factory.prepare_factory(atomic_numbers, config)
     dilute_planckian_radiation_field = DilutePlanckianRadiationField(
         simulation_state.t_radiative, simulation_state.dilution_factor
     )

tardis/plasma/detailed_balance/rates/__init__.py

@@ -0,0 +1,10 @@
+from tardis.plasma.detailed_balance.rates.collision_strengths import (
+    UpsilonCMFGENSolver,
+    UpsilonRegemorterSolver,
+)
+from tardis.plasma.detailed_balance.rates.collisional_rates import (
+    ThermalCollisionalRateSolver,
+)
+from tardis.plasma.detailed_balance.rates.radiative_rates import (
+    RadiativeRatesSolver,
+)