diff --git a/docs/physics/plasma/construction_simple_plasma.ipynb b/docs/physics/plasma/construction_simple_plasma.ipynb index 44849be37d3..10b53802885 100644 --- a/docs/physics/plasma/construction_simple_plasma.ipynb +++ b/docs/physics/plasma/construction_simple_plasma.ipynb @@ -4,7 +4,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "# Constructing a simple plasma" + "# Constructing a simple thermal plasma" ] }, { @@ -23,7 +23,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "bbbd27367e48465696aa4e40f25b8496", + "model_id": "adc8142e04de4ac4ad9fb6d9eba83286", "version_major": 2, "version_minor": 0 }, @@ -37,7 +37,7 @@ { "data": { "application/vnd.jupyter.widget-view+json": { - "model_id": "6830a2c86f754397999e6801ccca1001", + "model_id": "36b88ed7168e46c8864e17b0107f8f71", "version_major": 2, "version_minor": 0 }, @@ -51,9 +51,20 @@ ], "source": [ "import numpy as np\n", + "import pandas as pd\n", "from astropy import units as u\n", - "from tardis.radiation_field.planck_rad_field import DilutePlanckianRadiationField\n", - "#from tardis.plasma.properties.plasma_input import T" + "\n", + "from tardis.io.atom_data import AtomData\n", + "from tardis.plasma.assembly.base import PlasmaSolverFactory\n", + "from tardis.plasma.radiation_field import DilutePlanckianRadiationField\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In the default configuration the plasma solver assumes the conditions for a Local Thermodynamic Equilibrium meaning a planckian radiation field and a maxwellian electron velocity distribution at the same templerature" ] }, { @@ -62,8 +73,9 @@ "metadata": {}, "outputs": [], "source": [ - "temperatures_rad = np.ones(10) * 10000 * u.K\n", - "dilution_factor = np.ones(10) * 0.1" + "NO_OF_CELLS = 1000\n", + "temperatures_rad = np.ones(NO_OF_CELLS) * 10000 * u.K\n", + "dilution_factor = np.ones(NO_OF_CELLS) * 0.1" ] }, { @@ -77,12 +89,105 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "base_density = 1e-15 * u.Unit('g/cm^3')\n", + "number_densities = pd.DataFrame(columns=range(NO_OF_CELLS), index=[1,26,28])\n", + "number_densities.values[:] = 0.5 * (base_density / u.u).to(1/u.cm**3).value\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, "metadata": {}, "outputs": [], "source": [ - "test = DilutePlanckianRadFieldInput()\n", - "test.set_value(d_radfield)" + "atom_data = AtomData.from_hdf('kurucz_cd23_chianti_H_He.h5')\n", + "plasma_solver_factory = PlasmaSolverFactory(atom_data, selected_atomic_numbers=[1,26,28])\n", + "plasma_solver_factory.setup_factory()\n" + ] + }, + { + "cell_type": "code", + "execution_count": 36, + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
0
atomic_numberion_number
102.245918e-04
13.011070e+08
201.405189e-15
12.206015e-03
23.011070e+08
\n", + "
" + ], + "text/plain": [ + " 0\n", + "atomic_number ion_number \n", + "1 0 2.245918e-04\n", + " 1 3.011070e+08\n", + "2 0 1.405189e-15\n", + " 1 2.206015e-03\n", + " 2 3.011070e+08" + ] + }, + "execution_count": 36, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "plasma_solver.ion_number_density" ] }, { @@ -107,7 +212,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.11.5" + "version": "3.12.4" } }, "nbformat": 4, diff --git a/tardis/io/configuration/tests/test_config_reader.py b/tardis/io/configuration/tests/test_config_reader.py index f525557ca63..784797a2f03 100644 --- a/tardis/io/configuration/tests/test_config_reader.py +++ b/tardis/io/configuration/tests/test_config_reader.py @@ -10,7 +10,7 @@ from astropy.units import Quantity from tardis.io.configuration.config_reader import Configuration from tardis.plasma.exceptions import PlasmaConfigError -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma def test_convergence_section_parser(): diff --git a/tardis/model/base.py b/tardis/model/base.py index d3e1ce6cee4..27e11d2f20d 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -95,6 +95,7 @@ class SimulationState(HDFWriterMixin): "density", "r_inner", "time_explosion", + "abundance", ] hdf_name = "simulation_state" diff --git a/tardis/opacities/opacity_solver.py b/tardis/opacities/opacity_solver.py index 9b3b63f48ad..d5e821b99a7 100644 --- a/tardis/opacities/opacity_solver.py +++ b/tardis/opacities/opacity_solver.py @@ -46,7 +46,9 @@ def solve(self, legacy_plasma) -> OpacityState: legacy_plasma.atomic_data.lines.shape[ 0 ], # number of lines - legacy_plasma.abundance.shape[1], # number of shells + legacy_plasma.number_density.shape[ + 1 + ], # number of shells ), dtype=np.float64, ), diff --git a/tardis/plasma/assembly/__init__.py b/tardis/plasma/assembly/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tardis/plasma/assembly/base.py b/tardis/plasma/assembly/base.py new file mode 100644 index 00000000000..7470ee08ea9 --- /dev/null +++ b/tardis/plasma/assembly/base.py @@ -0,0 +1,656 @@ +import logging + +import numpy as np +import pandas as pd +from astropy import units as u + +from tardis.plasma import BasePlasma +from tardis.plasma.base import PlasmaSolverSettings +from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + LevelBoltzmannFactorNLTE, + MarkovChainTransProbsCollector, + RadiationFieldCorrection, + StimulatedEmissionFactor, +) +from tardis.plasma.properties.base import TransitionProbabilitiesProperty +from tardis.plasma.properties.level_population import LevelNumberDensity +from tardis.plasma.properties.nlte_rate_equation_solver import ( + NLTEPopulationSolverLU, + NLTEPopulationSolverRoot, +) +from tardis.plasma.properties.property_collections import ( + adiabatic_cooling_properties, + basic_inputs, + basic_properties, + continuum_interaction_inputs, + continuum_interaction_properties, + dilute_lte_excitation_properties, + helium_lte_properties, + helium_nlte_properties, + helium_numerical_nlte_properties, + lte_excitation_properties, + lte_ionization_properties, + macro_atom_properties, + nebular_ionization_properties, + nlte_lu_solver_properties, + nlte_properties, + nlte_root_solver_properties, + non_nlte_properties, + two_photon_properties, +) +from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( + DiluteBlackBodyContinuumPropertiesSolver, +) +from tardis.util.base import species_string_to_tuple + +logger = logging.getLogger(__name__) + +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. + """ + + ## 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 + + link_t_rad_t_electron: float = 1.0 + + radiative_rates_type: str = "dilute-blackbody" + + delta_treatment: float | None = None + + ## Statistical Balance Solver + legacy_nlte_species: list = [] + + nlte_excitation_species: list = [] + nlte_ionization_species: list = [] + nlte_solver: str = "lu" + + ## Helium Treatment options + helium_treatment: str = "none" + heating_rate_data_file: str = "none" + + ## Continuum Interaction + continuum_interaction_species: list = [] + enable_adiabatic_cooling: bool = False + enable_two_photon_decay: bool = False + + ## Opacities + line_interaction_type: str = "scatter" + + ## Assembly properties + plasma_modules: list = [] + kwargs: dict = {} + property_kwargs: dict = {} + + def __init__(self, atom_data, selected_atomic_numbers, config=None) -> None: + self.plasma_modules = [] + self.kwargs = {} + self.property_kwargs = {} + + 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"], + ) + + def parse_plasma_config(self, plasma_config): + """ + Parse the plasma configuration. + + Parameters + ---------- + plasma_config : PlasmaConfig + The plasma configuration object containing the plasma parameters. + + Returns + ------- + None + """ + self.continuum_interaction_species = ( + plasma_config.continuum_interaction.species + ) + self.set_nlte_species_from_string(plasma_config.nlte.species) + self.line_interaction_type = plasma_config.line_interaction_type + self.link_t_rad_t_electron = plasma_config.link_t_rad_t_electron + + self.excitation_analytical_approximation = plasma_config.excitation + self.ionization_analytical_approximation = plasma_config.ionization + self.delta_treatment = plasma_config.get("delta_treatment", None) + + self.helium_treatment = plasma_config.helium_treatment + self.heating_rate_data_file = plasma_config.heating_rate_data_file + + self.nlte_ionization_species = plasma_config.nlte_ionization_species + self.nlte_excitation_species = plasma_config.nlte_excitation_species + + self.nlte_solver = plasma_config.nlte_solver + + self.radiative_rates_type = plasma_config.radiative_rates_type + + self.enable_adiabatic_cooling = ( + plasma_config.continuum_interaction.enable_adiabatic_cooling + ) + self.enable_two_photon_decay = ( + plasma_config.continuum_interaction.enable_two_photon_decay + ) + + def setup_factory(self, config=None): + """ + Set up the plasma factory. + + Parameters + ---------- + config : object, optional + Configuration object containing plasma settings (default: None). + """ + self.check_continuum_interaction_species() + + self.plasma_modules = basic_inputs + basic_properties + + self.setup_analytical_approximations() + self.property_kwargs[RadiationFieldCorrection] = dict( + delta_treatment=self.delta_treatment + ) + if (config is not None) and len(self.legacy_nlte_species) > 0: + self.setup_legacy_nlte(config.plasma.nlte) + else: + self.plasma_modules += non_nlte_properties + + if self.line_interaction_type in ("downbranch", "macroatom") and ( + len(self.continuum_interaction_species) == 0 + ): + self.plasma_modules += macro_atom_properties + + self.setup_helium_treatment() + + if len(self.continuum_interaction_species) > 0: + self.setup_continuum_interactions() + + def setup_helium_treatment(self): + """ + Set up the helium treatment for the plasma assembly. + + Parameters + ---------- + helium_treatment : str + The type of helium treatment to be used. Possible values are: + - "recomb-nlte": Use recombination NLTE treatment for helium. + - "numerical-nlte": Use numerical NLTE treatment for helium. + + heating_rate_data_file : str or None + The path to the heating rate data file. Required when using + "numerical-nlte" helium treatment. + + Raises + ------ + PlasmaConfigError + If the helium NLTE treatment is incompatible with the NLTE ionization + and excitation treatment. + + If the heating rate data file is not specified when using + "numerical-nlte" helium treatment. + """ + if ( + self.helium_treatment == "recomb-nlte" + or self.helium_treatment == "numerical-nlte" + ) and ( + len(self.nlte_ionization_species + self.nlte_excitation_species) > 0 + ): + # Prevent the user from using helium NLTE treatment with + # NLTE ionization and excitation treatment. This is because + # the helium_nlte_properties could overwrite the NLTE ionization + # and excitation ion number and electron densities. + # helium_numerical_nlte_properties is also included here because + # it is currently in the same if else block, and thus may block + # the addition of the components from the else block. + raise PlasmaConfigError( + "Helium NLTE treatment is incompatible with the NLTE ionization and excitation treatment." + ) + + # TODO: Disentangle these if else block such that compatible components + # can be added independently. + if self.helium_treatment == "recomb-nlte": + self.plasma_modules += helium_nlte_properties + elif self.helium_treatment == "numerical-nlte": + self.plasma_modules += helium_numerical_nlte_properties + if self.heating_rate_data_file in ["none", None]: + raise PlasmaConfigError("Heating rate data file not specified") + self.property_kwargs[HeliumNumericalNLTE] = dict( + heating_rate_data_file=self.heating_rate_data_file + ) + else: + # If nlte ionization species are present, we don't want to add the + # IonNumberDensity from helium_lte_properties, since we want + # to use the IonNumberDensity provided by the NLTE solver. + if ( + len(self.nlte_ionization_species + self.nlte_excitation_species) + > 0 + ): + self.plasma_modules.append(LevelNumberDensity) + else: + self.plasma_modules += helium_lte_properties + + def setup_legacy_nlte(self, nlte_config): + """ + Set up the non-LTE (NLTE) properties for the legacy species. + + Parameters + ---------- + nlte_config : dict + A dictionary containing the NLTE configuration. + """ + self.plasma_modules += nlte_properties + self.plasma_modules.append( + LevelBoltzmannFactorNLTE.from_config(nlte_config) + ) + self.property_kwargs[StimulatedEmissionFactor] = dict( + nlte_species=self.legacy_nlte_species + ) + + def setup_analytical_approximations(self): + """ + Setup the analytical approximations for excitation and ionization. + + Returns + ------- + None + """ + if self.excitation_analytical_approximation == "lte": + self.plasma_modules += lte_excitation_properties + elif self.excitation_analytical_approximation == "dilute-lte": + self.plasma_modules += dilute_lte_excitation_properties + else: + raise PlasmaConfigError( + f'Invalid excitation analytical approximation. Configured as {self.excitation_analytical_approximation} but needs to be either "lte" or "dilute-lte"' + ) + + if self.ionization_analytical_approximation == "lte": + self.plasma_modules += lte_ionization_properties + elif self.ionization_analytical_approximation == "nebular": + self.plasma_modules += nebular_ionization_properties + else: + raise PlasmaConfigError( + f'Invalid excitation analytical approximation. Configured as {self.ionization_analytical_approximation} but needs to be either "lte" or "nebular"' + ) + + def initialize_j_blues(self, dilute_planckian_radiation_field, lines_df): + """ + Initialize the j_blues DataFrame based on the radiative_rates_type and the dilute_planckian_radiation_field. + + Parameters + ---------- + dilute_planckian_radiation_field : object + The dilute Planckian radiation field object. + lines_df : pandas.DataFrame + The DataFrame containing lines information. + + Returns + ------- + pandas.DataFrame + The DataFrame with calculated mean intensity values. + + Raises + ------ + ValueError + If the radiative_rates_type is unknown. + """ + if (self.radiative_rates_type == "dilute-blackbody") or ( + self.radiative_rates_type == "detailed" + ): + j_blues = pd.DataFrame( + dilute_planckian_radiation_field.calculate_mean_intensity( + lines_df.nu.values + ), + index=lines_df.index, + ) + + elif self.radiative_rates_type == "blackbody": + planckian_rad_field = ( + dilute_planckian_radiation_field.to_planckian_radiation_field() + ) + j_blues = pd.DataFrame( + planckian_rad_field.calculate_mean_intensity( + lines_df.nu.values + ), + index=lines_df.index, + ) + + else: + raise ValueError( + f"radiative_rates_type type unknown - {self.radiative_rates_type}" + ) + + return j_blues + + def set_continuum_interaction_species_from_string( + self, continuum_interaction_species + ): + """ + Set the continuum interaction species from a list of species strings. + + Parameters + ---------- + continuum_interaction_species : list of str + List of species strings representing the continuum interaction species. + + Returns + ------- + None + """ + self.continuum_interaction_species = [ + species_string_to_tuple(species) + for species in continuum_interaction_species + ] + + def check_continuum_interaction_species(self): + """ + Check if all continuum interaction species belong to atoms that have been specified in the configuration. + + Raises + ------ + PlasmaConfigError: If not all continuum interaction species belong to specified atoms. + """ + continuum_atoms = ( + self.continuum_interaction_species_multi_index.get_level_values( + "atomic_number" + ) + ) + + continuum_atoms_in_selected_atoms = np.all( + continuum_atoms.isin(self.atom_data.selected_atomic_numbers) + ) + + if not continuum_atoms_in_selected_atoms: + raise PlasmaConfigError( + "Not all continuum interaction species " + "belong to atoms that have been specified " + "in the configuration." + ) + + def set_nlte_species_from_string(self, nlte_species): + """ + Sets the non-LTE species from a string representation. + + Parameters + ---------- + nlte_species : str + A string representation of the non-LTE species. + + Returns + ------- + None + This method does not return anything. + """ + self.legacy_nlte_species = map_species_from_string(nlte_species) + + def setup_continuum_interactions(self): + """ + Set up continuum interactions for the plasma assembly. + + Raises + ------ + PlasmaConfigError: If the line_interaction_type is not "macroatom". + PlasmaConfigError: If an NLTE ionization species is not in the continuum species. + PlasmaConfigError: If an NLTE excitation species is not in the continuum species. + PlasmaConfigError: If the NLTE solver type is unknown. + """ + if self.line_interaction_type != "macroatom": + raise PlasmaConfigError( + "Continuum interactions require line_interaction_type " + f"macroatom (instead of {self.line_interaction_type})." + ) + + self.plasma_modules += continuum_interaction_properties + self.plasma_modules += continuum_interaction_inputs + + if self.enable_adiabatic_cooling: + self.plasma_modules += adiabatic_cooling_properties + + if self.enable_two_photon_decay: + self.plasma_modules += two_photon_properties + + transition_probabilities_outputs = [ + plasma_property.transition_probabilities_outputs + for plasma_property in self.plasma_modules + if issubclass(plasma_property, TransitionProbabilitiesProperty) + ] + transition_probabilities_outputs = [ + item + for sublist in transition_probabilities_outputs + for item in sublist + ] + + self.property_kwargs[MarkovChainTransProbsCollector] = { + "inputs": transition_probabilities_outputs + } + if len(self.nlte_ionization_species + self.nlte_excitation_species) > 0: + if self.nlte_ionization_species: + nlte_ionization_species = self.nlte_ionization_species + for species in nlte_ionization_species: + if species not in self.continuum_interaction_species: + raise PlasmaConfigError( + f"NLTE ionization species {species} not in continuum species." + ) + if self.nlte_excitation_species: + nlte_excitation_species = self.nlte_excitation_species + for species in nlte_excitation_species: + if species not in self.continuum_interaction_species: + raise PlasmaConfigError( + f"NLTE excitation species {species} not in continuum species." + ) + self.property_kwargs[NLTEIndexHelper] = { + "nlte_ionization_species": self.nlte_ionization_species, + "nlte_excitation_species": self.nlte_excitation_species, + } + if self.nlte_solver == "lu": + self.plasma_modules += nlte_lu_solver_properties + logger.warning( + "LU solver will be inaccurate for NLTE excitation, proceed with caution." + ) + elif self.nlte_solver == "root": + self.plasma_modules += nlte_root_solver_properties + else: + raise PlasmaConfigError( + f"NLTE solver type unknown - {self.nlte_solver}" + ) + + def setup_electron_densities(self, electron_densities): + if self.helium_treatment == "numerical-nlte": + self.property_kwargs[IonNumberDensityHeNLTE] = dict( + electron_densities=electron_densities + ) + elif ( + len(self.nlte_ionization_species + self.nlte_excitation_species) > 0 + ) and self.nlte_solver == "root": + self.property_kwargs[NLTEPopulationSolverRoot] = dict( + electron_densities=electron_densities + ) + elif ( + len(self.nlte_ionization_species + self.nlte_excitation_species) > 0 + ) and self.nlte_solver == "lu": + self.property_kwargs[NLTEPopulationSolverLU] = dict( + electron_densities=electron_densities + ) + else: + self.property_kwargs[IonNumberDensity] = dict( + electron_densities=electron_densities + ) + + def initialize_continuum_properties(self, dilute_planckian_radiation_field): + """ + Initialize the continuum properties of the plasma. + + Parameters + ---------- + dilute_planckian_radiation_field : DilutePlanckianRadiationField + The dilute Planckian radiation field. + + Returns + ------- + initial_continuum_properties : `~tardis.plasma.properties.ContinuumProperties` + The initial continuum properties of the plasma. + """ + t_electrons = dilute_planckian_radiation_field.temperature.to(u.K).value + + initial_continuum_solver = DiluteBlackBodyContinuumPropertiesSolver( + self.atom_data + ) + initial_continuum_properties = initial_continuum_solver.solve( + dilute_planckian_radiation_field, t_electrons + ) + return initial_continuum_properties + + def assemble( + self, + number_densities, + dilute_planckian_radiation_field, + time_explosion, + electron_densities=None, + ): + """ + Assemble the plasma based on the provided parameters and settings. + + Parameters + ---------- + number_densities : dict + Dictionary of number densities for different species. + dilute_planckian_radiation_field : object + The dilute Planckian radiation field object. + time_explosion : float + The time of explosion. + electron_densities : array-like, optional + Optional electron densities. + + Returns + ------- + BasePlasma + The assembled plasma object. + + Raises + ------ + ValueError + If an error occurs during assembly. + """ + j_blues = self.initialize_j_blues( + dilute_planckian_radiation_field, self.atom_data.lines + ) + plasma_solver_settings = PlasmaSolverSettings( + RADIATIVE_RATES_TYPE=self.radiative_rates_type + ) + + kwargs = dict( + time_explosion=time_explosion, + dilute_planckian_radiation_field=dilute_planckian_radiation_field, + number_density=number_densities, + link_t_rad_t_electron=self.link_t_rad_t_electron, + atomic_data=self.atom_data, + j_blues=j_blues, + continuum_interaction_species=self.continuum_interaction_species_multi_index, + 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( + gamma=initial_continuum_properties.photo_ionization_rate_coefficient, + bf_heating_coeff_estimator=None, + stim_recomb_cooling_coeff_estimator=None, + alpha_stim_factor=initial_continuum_properties.stimulated_recombination_rate_factor, + ) + + 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 + return BasePlasma( + plasma_properties=self.plasma_modules, + property_kwargs=self.property_kwargs, + plasma_solver_settings=plasma_solver_settings, + **kwargs, + ) diff --git a/tardis/plasma/assembly/legacy_assembly.py b/tardis/plasma/assembly/legacy_assembly.py new file mode 100644 index 00000000000..6f12fde618b --- /dev/null +++ b/tardis/plasma/assembly/legacy_assembly.py @@ -0,0 +1,39 @@ +from tardis.plasma.assembly.base import PlasmaSolverFactory +from tardis.plasma.radiation_field import DilutePlanckianRadiationField + + +def assemble_plasma(config, simulation_state, atom_data=None): + """ + Create a BasePlasma instance from a Configuration object + and a SimulationState. + + Parameters + ---------- + config : io.config_reader.Configuration + simulation_state : model.SimulationState + atom_data : atomic.AtomData + If None, an attempt will be made to read the atomic data + from config. + + Returns + ------- + : plasma.BasePlasma + + """ + atomic_numbers = simulation_state.abundance.index + plasma_solver_factory = PlasmaSolverFactory( + atom_data, + atomic_numbers, + config, + ) + plasma_solver_factory.setup_factory(config) + dilute_planckian_radiation_field = DilutePlanckianRadiationField( + simulation_state.t_radiative, simulation_state.dilution_factor + ) + + return plasma_solver_factory.assemble( + simulation_state.elemental_number_density, + dilute_planckian_radiation_field, + simulation_state.time_explosion, + simulation_state._electron_densities, + ) diff --git a/tardis/plasma/base.py b/tardis/plasma/base.py index 87af2cbd479..011cb17b50a 100644 --- a/tardis/plasma/base.py +++ b/tardis/plasma/base.py @@ -26,16 +26,16 @@ class BasePlasma(PlasmaWriterMixin): def __init__( self, plasma_properties, - plasma_solver_settings, property_kwargs=None, + plasma_solver_settings=None, **kwargs, ): - self.plasma_solver_settings = plasma_solver_settings self.outputs_dict = {} self.input_properties = [] self.plasma_properties = self._init_properties( plasma_properties, property_kwargs, **kwargs ) + self.plasma_solver_settings = plasma_solver_settings self._build_graph() self.update(**kwargs) diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index cb6e46ef650..7a66d16cc05 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -84,8 +84,8 @@ class SelectedAtoms(ProcessingPlasmaProperty): outputs = ("selected_atoms",) - def calculate(self, abundance): - return abundance.index + def calculate(self, number_density): + return number_density.index class ElectronTemperature(ProcessingPlasmaProperty): diff --git a/tardis/plasma/properties/ion_population.py b/tardis/plasma/properties/ion_population.py index 10f440b95c7..1d342163e21 100644 --- a/tardis/plasma/properties/ion_population.py +++ b/tardis/plasma/properties/ion_population.py @@ -221,7 +221,9 @@ def calculate( self._set_chi_0(ionization_data) if self.delta_treatment is None: if self.departure_coefficient is None: - departure_coefficient = 1.0 / w + departure_coefficient = ( + 1.0 / w + ) # see Equation 13 and explanations on page 451 lower right in ML 93 else: departure_coefficient = self.departure_coefficient radiation_field_correction = -np.ones( diff --git a/tardis/plasma/properties/nlte.py b/tardis/plasma/properties/nlte.py index f5a9a424046..a6e4d7ab230 100644 --- a/tardis/plasma/properties/nlte.py +++ b/tardis/plasma/properties/nlte.py @@ -29,7 +29,7 @@ class PreviousElectronDensities(PreviousIterationProperty): def set_initial_value(self, kwargs): initial_value = pd.Series( 1000000.0, - index=kwargs["abundance"].columns, + index=kwargs["number_density"].columns, ) self._set_initial_value(initial_value) @@ -47,6 +47,6 @@ def set_initial_value(self, kwargs): initial_value = pd.DataFrame( 1.0, index=kwargs["atomic_data"].lines.index, - columns=kwargs["abundance"].columns, + columns=kwargs["number_density"].columns, ) self._set_initial_value(initial_value) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index ea4c8ec37fb..0a71bb984de 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -17,7 +17,6 @@ class PlasmaPropertyCollection(list): basic_inputs = PlasmaPropertyCollection( [ DilutePlanckianRadField, - Abundance, NumberDensity, TimeExplosion, AtomicData, diff --git a/tardis/plasma/properties/radiative_properties.py b/tardis/plasma/properties/radiative_properties.py index a1b8d0a5d8e..2490a53683f 100644 --- a/tardis/plasma/properties/radiative_properties.py +++ b/tardis/plasma/properties/radiative_properties.py @@ -5,6 +5,7 @@ from astropy import units as u from numba import jit, prange + from tardis import constants as const from tardis.opacities.macro_atom.base import TransitionProbabilities from tardis.plasma.properties.base import ( @@ -12,6 +13,7 @@ TransitionProbabilitiesProperty, ) + logger = logging.getLogger(__name__) __all__ = [ diff --git a/tardis/plasma/tests/test_nlte_solver.py b/tardis/plasma/tests/test_nlte_solver.py index 48045d3cf09..27d5bafc617 100644 --- a/tardis/plasma/tests/test_nlte_solver.py +++ b/tardis/plasma/tests/test_nlte_solver.py @@ -14,7 +14,7 @@ calculate_jacobian_matrix, calculate_rate_matrix, ) -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma @pytest.fixture diff --git a/tardis/plasma/tests/test_tardis_model_density_config.py b/tardis/plasma/tests/test_tardis_model_density_config.py index 20c9ec46cda..5551e7c4eb4 100644 --- a/tardis/plasma/tests/test_tardis_model_density_config.py +++ b/tardis/plasma/tests/test_tardis_model_density_config.py @@ -4,7 +4,7 @@ from tardis.io.configuration.config_reader import Configuration from tardis.model import SimulationState -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma @pytest.fixture diff --git a/tardis/simulation/base.py b/tardis/simulation/base.py index c4c6b3c00e6..3f55bc0829f 100644 --- a/tardis/simulation/base.py +++ b/tardis/simulation/base.py @@ -16,7 +16,7 @@ ) from tardis.io.util import HDFWriterMixin from tardis.plasma.radiation_field import DilutePlanckianRadiationField -from tardis.plasma.standard_plasmas import assemble_plasma +from tardis.plasma.assembly.legacy_assembly import assemble_plasma from tardis.simulation.convergence import ConvergenceSolver from tardis.spectrum.base import SpectrumSolver from tardis.spectrum.formal_integral import FormalIntegrator diff --git a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py index 2751861736b..ef561073a21 100644 --- a/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py +++ b/tardis/transport/montecarlo/estimators/tests/test_continuum_property_solver.py @@ -1,14 +1,13 @@ from copy import deepcopy -import numpy.testing as npt import pandas.testing as pdt import pytest +from tardis.simulation import Simulation from tardis.transport.montecarlo.estimators.continuum_radfield_properties import ( DiluteBlackBodyContinuumPropertiesSolver, MCContinuumPropertiesSolver, ) -from tardis.simulation import Simulation @pytest.mark.continuum diff --git a/tardis/visualization/widgets/shell_info.py b/tardis/visualization/widgets/shell_info.py index bb3c2f46a86..ad67fbe2064 100644 --- a/tardis/visualization/widgets/shell_info.py +++ b/tardis/visualization/widgets/shell_info.py @@ -190,7 +190,7 @@ def __init__(self, sim_model): super().__init__( sim_model.simulation_state.t_radiative, sim_model.simulation_state.dilution_factor, - sim_model.plasma.abundance, + sim_model.simulation_state.abundance, sim_model.plasma.number_density, sim_model.plasma.ion_number_density, sim_model.plasma.level_number_density, @@ -216,7 +216,7 @@ def __init__(self, hdf_fpath): super().__init__( sim_data["/simulation/simulation_state/t_radiative"], sim_data["/simulation/simulation_state/dilution_factor"], - sim_data["/simulation/plasma/abundance"], + sim_data["/simulation/simulation_state/abundance"], sim_data["/simulation/plasma/number_density"], sim_data["/simulation/plasma/ion_number_density"], sim_data["/simulation/plasma/level_number_density"], diff --git a/tardis/visualization/widgets/tests/test_shell_info.py b/tardis/visualization/widgets/tests/test_shell_info.py index 7a99374dbc5..df7f4f168d9 100644 --- a/tardis/visualization/widgets/tests/test_shell_info.py +++ b/tardis/visualization/widgets/tests/test_shell_info.py @@ -16,7 +16,7 @@ def base_shell_info(simulation_verysimple): return BaseShellInfo( simulation_verysimple.simulation_state.t_radiative, simulation_verysimple.simulation_state.dilution_factor, - simulation_verysimple.plasma.abundance, + simulation_verysimple.simulation_state.abundance, simulation_verysimple.plasma.number_density, simulation_verysimple.plasma.ion_number_density, simulation_verysimple.plasma.level_number_density, @@ -58,12 +58,14 @@ def test_element_count_data( ): element_count_data = base_shell_info.element_count(1) assert element_count_data.shape == ( - len(simulation_verysimple.plasma.abundance[shell_num - 1]), + len( + simulation_verysimple.simulation_state.abundance[shell_num - 1] + ), 2, ) assert np.allclose( element_count_data.iloc[:, -1].map(np.float64), - simulation_verysimple.plasma.abundance[shell_num - 1], + simulation_verysimple.simulation_state.abundance[shell_num - 1], ) @pytest.mark.parametrize(("atomic_num", "shell_num"), [(12, 1), (20, 20)])