diff --git a/tardis/model/base.py b/tardis/model/base.py index 9b741381ac0..adb86b99bd6 100644 --- a/tardis/model/base.py +++ b/tardis/model/base.py @@ -161,6 +161,45 @@ def t_radiative(self, new_t_radiative): "Trying to set t_radiative for different number of shells." ) + @property + def elemental_number_density(self): + elemental_number_density = ( + ( + self.composition.elemental_mass_fraction + * self.composition.density + ) + .divide(self.composition.element_masses, axis=0) + .dropna() + ) + elemental_number_density = elemental_number_density.iloc[ + :, + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index, + ] + elemental_number_density.columns = range( + len(elemental_number_density.columns) + ) + return elemental_number_density + + @property + def isotopic_number_density(self): + isotopic_number_density = ( + self.composition.isotopic_mass_fraction * self.composition.density + ).divide( + self.composition.isotope_masses.loc[ + self.composition.isotopic_mass_fraction.index + ] + * u.u.to(u.g), + axis=0, + ) + isotopic_number_density = isotopic_number_density.iloc[ + :, + self.geometry.v_inner_boundary_index : self.geometry.v_outer_boundary_index, + ] + isotopic_number_density.columns = range( + len(isotopic_number_density.columns) + ) + return isotopic_number_density + @property def radius(self): return self.time_explosion * self.velocity diff --git a/tardis/plasma/properties/atomic.py b/tardis/plasma/properties/atomic.py index f206264d3e4..72fd4b6b472 100644 --- a/tardis/plasma/properties/atomic.py +++ b/tardis/plasma/properties/atomic.py @@ -31,8 +31,6 @@ "Lines", "LinesLowerLevelIndex", "LinesUpperLevelIndex", - "AtomicMass", - "IsotopeMass", "IonizationData", "ZetaData", "NLTEData", @@ -76,8 +74,6 @@ def _filter_atomic_property(self, levels, selected_atoms): return levels[levels.index.isin(selected_atoms, level="atomic_number")] def _set_index(self, levels): - # levels = levels.set_index(['atomic_number', 'ion_number', - # 'level_number']) return ( levels.index, levels["energy"], @@ -541,69 +537,6 @@ def calculate(self, level_idxs2line_idx, level_idxs2continuum_idx): return level_idxs2transition_idx -class AtomicMass(ProcessingPlasmaProperty): - """ - Attributes - ---------- - atomic_mass : pandas.Series - Atomic masses of the elements used. Indexed by atomic number. - """ - - outputs = ("atomic_mass",) - - def calculate(self, atomic_data, selected_atoms): - if getattr(self, self.outputs[0]) is not None: - return (getattr(self, self.outputs[0]),) - else: - return atomic_data.atom_data.loc[selected_atoms].mass - - -class IsotopeMass(ProcessingPlasmaProperty): - """ - Attributes - ---------- - isotope_mass : pandas.Series - Masses of the isotopes used. Indexed by isotope name e.g. 'Ni56'. - """ - - outputs = ("isotope_mass",) - - def calculate(self, isotope_abundance): - """ - Determine mass of each isotope. - - Parameters - ---------- - isotope_abundance : pandas.DataFrame - Fractional abundance of isotopes. - - Returns - ------- - pandas.DataFrame - Masses of the isotopes used. Indexed by isotope name e.g. 'Ni56'. - """ - if getattr(self, self.outputs[0]) is not None: - return (getattr(self, self.outputs[0]),) - else: - if isotope_abundance.empty: - return None - isotope_mass_dict = {} - for Z, A in isotope_abundance.index: - element_name = rd.utils.Z_to_elem(Z) - isotope_name = element_name + str(A) - - isotope_mass_dict[(Z, A)] = rd.Nuclide(isotope_name).atomic_mass - - isotope_mass_df = pd.DataFrame.from_dict( - isotope_mass_dict, orient="index", columns=["mass"] - ) - isotope_mass_df.index = pd.MultiIndex.from_tuples( - isotope_mass_df.index - ) - isotope_mass_df.index.names = ["atomic_number", "mass_number"] - return isotope_mass_df / const.N_A - - class IonizationData(BaseAtomicDataProperty): """ Attributes diff --git a/tardis/plasma/properties/general.py b/tardis/plasma/properties/general.py index 9988f656411..660991c1d32 100644 --- a/tardis/plasma/properties/general.py +++ b/tardis/plasma/properties/general.py @@ -12,8 +12,6 @@ __all__ = [ "BetaRadiation", "GElectron", - "NumberDensity", - "IsotopeNumberDensity", "SelectedAtoms", "ElectronTemperature", "BetaElectron", @@ -79,64 +77,6 @@ def calculate(self, beta_electron): return super(ThermalGElectron, self).calculate(beta_electron) -class NumberDensity(ProcessingPlasmaProperty): - """ - Attributes - ---------- - number_density : Pandas DataFrame, dtype float - Indexed by atomic number, columns corresponding to zones - """ - - outputs = ("number_density",) - latex_name = ("N_{i}",) - - @staticmethod - def calculate(atomic_mass, abundance, density): - number_densities = abundance * density - return number_densities.div(atomic_mass.loc[abundance.index], axis=0) - - -class IsotopeNumberDensity(ProcessingPlasmaProperty): - """ - Calculate the atom number density based on isotope mass. - - Attributes - ---------- - isotope_number_density : Pandas DataFrame, dtype float - Indexed by atomic number, columns corresponding to zones - """ - - outputs = ("isotope_number_density",) - latex_name = ("N_{i}",) - - @staticmethod - def calculate(isotope_mass, isotope_abundance, density): - """ - Calculate the atom number density based on isotope mass. - - Parameters - ---------- - isotope_mass : pandas.DataFrame - Masses of isotopes. - isotope_abundance : pandas.DataFrame - Fractional abundance of isotopes. - density : pandas.DataFrame - Density of each shell. - - Returns - ------- - pandas.DataFrame - Indexed by atomic number, columns corresponding to zones. - """ - number_densities = isotope_abundance * density - isotope_number_density_array = ( - number_densities.to_numpy() / isotope_mass.to_numpy() - ) - return pd.DataFrame( - isotope_number_density_array, index=isotope_abundance.index - ) - - class SelectedAtoms(ProcessingPlasmaProperty): """ Attributes diff --git a/tardis/plasma/properties/plasma_input.py b/tardis/plasma/properties/plasma_input.py index 2d7b67ca4d6..ae8b6ab4adc 100644 --- a/tardis/plasma/properties/plasma_input.py +++ b/tardis/plasma/properties/plasma_input.py @@ -1,12 +1,15 @@ -from tardis.plasma.properties.base import Input, ArrayInput, DataFrameInput +from tardis.plasma.properties.base import ( + ArrayInput, + Input, +) __all__ = [ "TRadiative", "DilutionFactor", "AtomicData", "Abundance", + "NumberDensity", "IsotopeAbundance", - "Density", "TimeExplosion", "JBlueEstimator", "LinkTRadTElectron", @@ -76,18 +79,6 @@ class IsotopeAbundance(Input): outputs = ("isotope_abundance",) -class Density(ArrayInput): - """ - Attributes - ---------- - density : Numpy array, dtype float - Total density values - """ - - outputs = ("density",) - latex_name = (r"\rho",) - - class TimeExplosion(Input): """ Attributes @@ -160,3 +151,15 @@ class NLTEIonizationSpecies(Input): class NLTEExcitationSpecies(Input): outputs = ("nlte_excitation_species",) + + +class NumberDensity(Input): + """ + Attributes + ---------- + number_density : Pandas DataFrame, dtype float + Indexed by atomic number, columns corresponding to zones + """ + + outputs = ("number_density",) + latex_name = ("N_{i}",) diff --git a/tardis/plasma/properties/property_collections.py b/tardis/plasma/properties/property_collections.py index deb2405347b..d93b6016442 100644 --- a/tardis/plasma/properties/property_collections.py +++ b/tardis/plasma/properties/property_collections.py @@ -9,7 +9,7 @@ class PlasmaPropertyCollection(list): [ TRadiative, Abundance, - Density, + NumberDensity, TimeExplosion, AtomicData, DilutionFactor, @@ -25,11 +25,9 @@ class PlasmaPropertyCollection(list): BetaRadiation, Levels, Lines, - AtomicMass, PartitionFunction, GElectron, IonizationData, - NumberDensity, LinesLowerLevelIndex, LinesUpperLevelIndex, TauSobolev, @@ -153,6 +151,3 @@ class PlasmaPropertyCollection(list): TwoPhotonFrequencySampler, ] ) -isotope_properties = PlasmaPropertyCollection( - [IsotopeAbundance, IsotopeMass, IsotopeNumberDensity] -) diff --git a/tardis/plasma/standard_plasmas.py b/tardis/plasma/standard_plasmas.py index 42d81af1e1b..da772508140 100644 --- a/tardis/plasma/standard_plasmas.py +++ b/tardis/plasma/standard_plasmas.py @@ -1,57 +1,53 @@ -from astropy import units as u -import os import logging import numpy as np import pandas as pd +from astropy import units as u -from tardis.io.atom_data import AtomData +from tardis.plasma import BasePlasma +from tardis.plasma.exceptions import PlasmaConfigError +from tardis.plasma.properties import ( + HeliumNumericalNLTE, + IonNumberDensity, + IonNumberDensityHeNLTE, + JBluesBlackBody, + JBluesDetailed, + JBluesDiluteBlackBody, + 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.rate_matrix_index import NLTEIndexHelper -from tardis.util.base import species_string_to_tuple -from tardis.plasma import BasePlasma -from tardis.plasma.properties.base import TransitionProbabilitiesProperty from tardis.plasma.properties.property_collections import ( + adiabatic_cooling_properties, basic_inputs, basic_properties, + continuum_interaction_inputs, + continuum_interaction_properties, + detailed_j_blues_inputs, + detailed_j_blues_properties, + dilute_lte_excitation_properties, + helium_lte_properties, + helium_nlte_properties, + helium_numerical_nlte_properties, lte_excitation_properties, lte_ionization_properties, macro_atom_properties, - dilute_lte_excitation_properties, nebular_ionization_properties, - non_nlte_properties, - nlte_properties, - helium_nlte_properties, - helium_numerical_nlte_properties, - helium_lte_properties, - detailed_j_blues_properties, - detailed_j_blues_inputs, - continuum_interaction_properties, - continuum_interaction_inputs, - adiabatic_cooling_properties, - two_photon_properties, - isotope_properties, nlte_lu_solver_properties, + nlte_properties, nlte_root_solver_properties, + non_nlte_properties, + two_photon_properties, ) -from tardis.plasma.exceptions import PlasmaConfigError - -from tardis.plasma.properties import ( - LevelBoltzmannFactorNLTE, - JBluesBlackBody, - JBluesDiluteBlackBody, - JBluesDetailed, - RadiationFieldCorrection, - StimulatedEmissionFactor, - HeliumNumericalNLTE, - IonNumberDensity, - IonNumberDensityHeNLTE, - MarkovChainTransProbsCollector, -) +from tardis.plasma.properties.rate_matrix_index import NLTEIndexHelper +from tardis.util.base import species_string_to_tuple logger = logging.getLogger(__name__) @@ -121,7 +117,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): kwargs = dict( t_rad=simulation_state.t_radiative, abundance=simulation_state.abundance, - density=simulation_state.density, + number_density=simulation_state.elemental_number_density, atomic_data=atom_data, time_explosion=simulation_state.time_explosion, w=simulation_state.dilution_factor, @@ -138,7 +134,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): if line_interaction_type != "macroatom": raise PlasmaConfigError( "Continuum interactions require line_interaction_type " - "macroatom (instead of {}).".format(line_interaction_type) + f"macroatom (instead of {line_interaction_type})." ) plasma_modules += continuum_interaction_properties @@ -171,8 +167,9 @@ def assemble_plasma(config, simulation_state, atom_data=None): if config.plasma.nlte_ionization_species: nlte_ionization_species = config.plasma.nlte_ionization_species for species in nlte_ionization_species: - if not ( - species in config.plasma.continuum_interaction.species + if ( + species + not in config.plasma.continuum_interaction.species ): raise PlasmaConfigError( f"NLTE ionization species {species} not in continuum species." @@ -180,8 +177,9 @@ def assemble_plasma(config, simulation_state, atom_data=None): if config.plasma.nlte_excitation_species: nlte_excitation_species = config.plasma.nlte_excitation_species for species in nlte_excitation_species: - if not ( - species in config.plasma.continuum_interaction.species + if ( + species + not in config.plasma.continuum_interaction.species ): raise PlasmaConfigError( f"NLTE excitation species {species} not in continuum species." @@ -199,9 +197,7 @@ def assemble_plasma(config, simulation_state, atom_data=None): plasma_modules += nlte_root_solver_properties else: raise PlasmaConfigError( - "NLTE solver type unknown - {}".format( - config.plasma.nlte_solver - ) + f"NLTE solver type unknown - {config.plasma.nlte_solver}" ) kwargs.update( diff --git a/tardis/plasma/tests/test_complete_plasmas.py b/tardis/plasma/tests/test_complete_plasmas.py index ed7d4428852..a28aeba63e8 100644 --- a/tardis/plasma/tests/test_complete_plasmas.py +++ b/tardis/plasma/tests/test_complete_plasmas.py @@ -193,7 +193,7 @@ def test_plasma_properties(self, plasma, attr): def test_levels(self, plasma): actual = pd.DataFrame(plasma.levels) - key = f"plasma/levels" + key = "plasma/levels" expected = pd.read_hdf(self.regression_data.fpath, key) pdt.assert_frame_equal(actual, expected) @@ -202,13 +202,13 @@ def test_scalars_properties(self, plasma, attr): actual = getattr(plasma, attr) if hasattr(actual, "cgs"): actual = actual.cgs.value - key = f"plasma/scalars" + key = "plasma/scalars" expected = pd.read_hdf(self.regression_data.fpath, key)[attr] npt.assert_equal(actual, expected) def test_helium_treatment(self, plasma): actual = plasma.helium_treatment - key = f"plasma/scalars" + key = "plasma/scalars" expected = pd.read_hdf(self.regression_data.fpath, key)[ "helium_treatment" ] @@ -217,6 +217,6 @@ def test_helium_treatment(self, plasma): def test_zeta_data(self, plasma): if hasattr(plasma, "zeta_data"): actual = plasma.zeta_data - key = f"plasma/zeta_data" + key = "plasma/zeta_data" expected = pd.read_hdf(self.regression_data.fpath, key) npt.assert_allclose(actual, expected.values) diff --git a/tardis/plasma/tests/test_hdf_plasma.py b/tardis/plasma/tests/test_hdf_plasma.py index d0c4f4c6433..13169bdfa4e 100644 --- a/tardis/plasma/tests/test_hdf_plasma.py +++ b/tardis/plasma/tests/test_hdf_plasma.py @@ -86,7 +86,7 @@ def test_atomic_data_uuid(simulation_verysimple, regression_data): assert actual == expected -COLLECTION_PROPERTIES = ["t_rad", "w", "density"] +COLLECTION_PROPERTIES = ["t_rad", "w"] @pytest.mark.parametrize("attr", COLLECTION_PROPERTIES) diff --git a/tardis/tests/test_tardis_full_formal_integral.py b/tardis/tests/test_tardis_full_formal_integral.py index f28203045f1..d3a0eeaecc5 100644 --- a/tardis/tests/test_tardis_full_formal_integral.py +++ b/tardis/tests/test_tardis_full_formal_integral.py @@ -1,11 +1,12 @@ from pathlib import Path -import pytest + import numpy.testing as npt +import pytest from astropy import units as u from astropy.tests.helper import assert_quantity_allclose -from tardis.simulation.base import Simulation from tardis.io.configuration.config_reader import Configuration +from tardis.simulation.base import Simulation config_line_modes = ["downbranch", "macroatom"] interpolate_shells = [-1, 30]