Restructure/detailed balance foundation (#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 <[email protected]>
tardis-sn · Oct 14, 2024 · 8106ede · 8106ede
1 parent 47d64bf
commit 8106ede
Show file tree

Hide file tree

Showing 15 changed files with 4,427 additions and 95 deletions.
diff --git a/docs/physics/plasma/detailed_balance/collisional_rate_coefficients.hdf b/docs/physics/plasma/detailed_balance/collisional_rate_coefficients.hdf
diff --git a/docs/physics/plasma/detailed_balance/comparison.ipynb b/docs/physics/plasma/detailed_balance/comparison.ipynb
diff --git a/docs/physics/plasma/detailed_balance/rates.ipynb b/docs/physics/plasma/detailed_balance/rates.ipynb
diff --git a/docs/physics/plasma/detailed_balance/test_continuum_template_wkerzen_rate_coeffs.yml b/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
diff --git a/tardis/plasma/assembly/base.py b/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,
         )
diff --git a/tardis/plasma/assembly/legacy_assembly.py b/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
     )

diff --git a/tardis/plasma/detailed_balance/__init__.py b/tardis/plasma/detailed_balance/__init__.py
diff --git a/tardis/plasma/detailed_balance/rates/__init__.py b/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,
+)