From 4af6a5dd7b17e5b165f923bbaca737e7e149c3b1 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Mon, 29 Jul 2024 16:42:31 -0400 Subject: [PATCH 1/6] Refactor attribute setup --- tardis/io/atom_data/base.py | 37 +++++++------------------------------ tardis/io/atom_data/util.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 30 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 17a5b58886e..fdb106d50fc 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -8,7 +8,10 @@ from scipy import interpolate from tardis import constants as const -from tardis.io.atom_data.util import resolve_atom_data_fname +from tardis.io.atom_data.util import ( + resolve_atom_data_fname, + set_atom_data_attributes, +) from tardis.plasma.properties.continuum_processes import ( get_ground_state_multi_index, ) @@ -223,35 +226,9 @@ def from_hdf(cls, fname=None): atom_data = cls(**dataframes) - try: - atom_data.uuid1 = store.root._v_attrs["uuid1"] - if hasattr(atom_data.uuid1, "decode"): - atom_data.uuid1 = store.root._v_attrs["uuid1"].decode( - "ascii" - ) - except KeyError: - logger.debug( - "UUID not available for Atom Data. Setting value to None" - ) - atom_data.uuid1 = None - - try: - atom_data.md5 = store.root._v_attrs["md5"] - if hasattr(atom_data.md5, "decode"): - atom_data.md5 = store.root._v_attrs["md5"].decode("ascii") - except KeyError: - logger.debug( - "MD5 not available for Atom Data. Setting value to None" - ) - atom_data.md5 = None - - try: - atom_data.version = store.root._v_attrs["database_version"] - except KeyError: - logger.debug( - "VERSION not available for Atom Data. Setting value to None" - ) - atom_data.version = None + set_atom_data_attributes(atom_data, store, "uuid1") + set_atom_data_attributes(atom_data, store, "md5") + set_atom_data_attributes(atom_data, store, "database_version") # TODO: strore data sources as attributes in carsus diff --git a/tardis/io/atom_data/util.py b/tardis/io/atom_data/util.py index ce4b2765bce..f5359cce919 100644 --- a/tardis/io/atom_data/util.py +++ b/tardis/io/atom_data/util.py @@ -51,3 +51,32 @@ def resolve_atom_data_fname(fname): f"Atom Data {fname} is not found in current path or in TARDIS data repo. {atom_data_name} " "is also not a standard known TARDIS atom dataset." ) + + +def set_atom_data_attributes(atom_data, store, attribute): + """Sets arbitrary atom data attributes, throws error and sets to None + if they are not available. + + Parameters + ---------- + atom_data : AtomData + The atom data to modify + store : pd.HDFStore + Data source + property : str + Property to modify + """ + try: + setattr(atom_data, attribute, store.root._v_attrs[attribute]) + new_attribute = getattr(atom_data, attribute) + if hasattr(new_attribute, "decode"): + setattr( + atom_data, + attribute, + store.root._v_attrs[attribute].decode("ascii"), + ) + except KeyError: + logger.debug( + f"{attribute} not available for Atom Data. Setting value to None" + ) + setattr(atom_data, attribute, None) From 606946f781fa123c93b2f6f0c8f04d58cf3199a3 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 30 Jul 2024 10:35:26 -0400 Subject: [PATCH 2/6] Remove duplicate NLTEData class from atom_data --- tardis/io/atom_data/base.py | 111 +----------------------------------- 1 file changed, 1 insertion(+), 110 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index fdb106d50fc..91dbc5c9dad 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -5,13 +5,13 @@ import pandas as pd from astropy import units as u from astropy.units import Quantity -from scipy import interpolate from tardis import constants as const from tardis.io.atom_data.util import ( resolve_atom_data_fname, set_atom_data_attributes, ) +from tardis.plasma.nlte.nlte_data import NLTEData from tardis.plasma.properties.continuum_processes import ( get_ground_state_multi_index, ) @@ -630,112 +630,3 @@ def _check_selected_atomic_numbers(self): def __repr__(self): return f"" - - -class NLTEData: - def __init__(self, atom_data, nlte_species): - self.atom_data = atom_data - self.lines = atom_data.lines.reset_index() - self.nlte_species = nlte_species - - if nlte_species: - logger.info("Preparing the NLTE data") - self._init_indices() - if atom_data.collision_data is not None: - self._create_collision_coefficient_matrix() - - def _init_indices(self): - self.lines_idx = {} - self.lines_level_number_lower = {} - self.lines_level_number_upper = {} - self.A_uls = {} - self.B_uls = {} - self.B_lus = {} - - for species in self.nlte_species: - lines_idx = np.where( - (self.lines.atomic_number == species[0]) - & (self.lines.ion_number == species[1]) - ) - self.lines_idx[species] = lines_idx - self.lines_level_number_lower[ - species - ] = self.lines.level_number_lower.values[lines_idx].astype(int) - self.lines_level_number_upper[ - species - ] = self.lines.level_number_upper.values[lines_idx].astype(int) - - self.A_uls[species] = self.atom_data.lines.A_ul.values[lines_idx] - self.B_uls[species] = self.atom_data.lines.B_ul.values[lines_idx] - self.B_lus[species] = self.atom_data.lines.B_lu.values[lines_idx] - - def _create_collision_coefficient_matrix(self): - self.C_ul_interpolator = {} - self.delta_E_matrices = {} - self.g_ratio_matrices = {} - collision_group = self.atom_data.collision_data.groupby( - level=["atomic_number", "ion_number"] - ) - for species in self.nlte_species: - no_of_levels = self.atom_data.levels.loc[species].energy.count() - C_ul_matrix = np.zeros( - ( - no_of_levels, - no_of_levels, - len(self.atom_data.collision_data_temperatures), - ) - ) - delta_E_matrix = np.zeros((no_of_levels, no_of_levels)) - g_ratio_matrix = np.zeros((no_of_levels, no_of_levels)) - - for ( - ( - atomic_number, - ion_number, - level_number_lower, - level_number_upper, - ), - line, - ) in collision_group.get_group(species).iterrows(): - # line.columns : delta_e, g_ratio, temperatures ... - C_ul_matrix[ - level_number_lower, level_number_upper, : - ] = line.values[2:] - delta_E_matrix[level_number_lower, level_number_upper] = line[ - "delta_e" - ] - # TODO TARDISATOMIC fix change the g_ratio to be the otherway round - I flip them now here. - g_ratio_matrix[level_number_lower, level_number_upper] = ( - 1 / line["g_ratio"] - ) - self.C_ul_interpolator[species] = interpolate.interp1d( - self.atom_data.collision_data_temperatures, C_ul_matrix - ) - self.delta_E_matrices[species] = delta_E_matrix - - self.g_ratio_matrices[species] = g_ratio_matrix - - def get_collision_matrix(self, species, t_electrons): - """ - Creat collision matrix by interpolating the C_ul values for - the desired temperatures. - """ - c_ul_matrix = self.C_ul_interpolator[species](t_electrons) - no_of_levels = c_ul_matrix.shape[0] - c_ul_matrix[np.isnan(c_ul_matrix)] = 0.0 - - # TODO in tardisatomic the g_ratio is the other way round - here I'll flip it in prepare_collision matrix - - c_lu_matrix = ( - c_ul_matrix - * np.exp( - -self.delta_E_matrices[species].reshape( - (no_of_levels, no_of_levels, 1) - ) - / t_electrons.reshape((1, 1, t_electrons.shape[0])) - ) - * self.g_ratio_matrices[species].reshape( - (no_of_levels, no_of_levels, 1) - ) - ) - return c_ul_matrix + c_lu_matrix.transpose(1, 0, 2) From ef00b6998ee54ced2f98122353b326b516f509f1 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 30 Jul 2024 10:56:34 -0400 Subject: [PATCH 3/6] Cleanup unused attributes --- tardis/io/atom_data/base.py | 24 ++++-------------------- 1 file changed, 4 insertions(+), 20 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 91dbc5c9dad..0ff050b672c 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -114,8 +114,6 @@ class AtomData: collision_data_temperatures : numpy.array zeta_data : pandas.DataFrame synpp_refs : pandas.DataFrame - symbol2atomic_number : OrderedDict - atomic_number2symbol : OrderedDict photoionization_data : pandas.DataFrame two_photon_data : pandas.DataFrame decay_radiation_data : pandas.DataFrame @@ -171,8 +169,8 @@ def from_hdf(cls, fname=None): Path to the HDFStore file or name of known atom data file (default: None) """ - dataframes = dict() - nonavailable = list() + dataframes = {} + nonavailable = [] fname = resolve_atom_data_fname(fname) @@ -271,14 +269,7 @@ def __init__( # different values for the unit u and the constant. # This is changed in later versions of astropy ( # the value of constants.u is used in all cases) - if u.u.cgs == const.u.cgs: - atom_data.loc[:, "mass"] = Quantity( - atom_data["mass"].values, "u" - ).cgs.value - else: - atom_data.loc[:, "mass"] = ( - atom_data["mass"].values * const.u.cgs.value - ) + atom_data.loc[:, "mass"] = atom_data["mass"].values * const.u.cgs.value # Convert ionization energies to CGS ionization_data = ionization_data.squeeze() @@ -299,7 +290,7 @@ def __init__( self.atom_data = atom_data self.ionization_data = ionization_data self.levels = levels - # Not sure why this is need - WEK 17 Sep 2023 + # Cast to float so that Numba can use the values in numpy functions self.levels.energy = self.levels.energy.astype(np.float64) self.lines = lines @@ -327,13 +318,6 @@ def __init__( self.decay_radiation_data = decay_radiation_data self._check_related() - self.symbol2atomic_number = OrderedDict( - zip(self.atom_data["symbol"].values, self.atom_data.index) - ) - self.atomic_number2symbol = OrderedDict( - zip(self.atom_data.index, self.atom_data["symbol"]) - ) - def _check_related(self): """ Check that either all or none of the related dataframes are given. From 6b5144e6f86965013d1342435649f599cede55ff Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 30 Jul 2024 15:50:58 -0400 Subject: [PATCH 4/6] Move NLTE data back --- tardis/io/atom_data/base.py | 57 ++++++++++--------- .../nlte => io/atom_data}/nlte_data.py | 0 tardis/plasma/nlte/__init__.py | 0 3 files changed, 31 insertions(+), 26 deletions(-) rename tardis/{plasma/nlte => io/atom_data}/nlte_data.py (100%) delete mode 100644 tardis/plasma/nlte/__init__.py diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 0ff050b672c..936f78bb626 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -1,5 +1,4 @@ import logging -from collections import OrderedDict import numpy as np import pandas as pd @@ -11,7 +10,7 @@ resolve_atom_data_fname, set_atom_data_attributes, ) -from tardis.plasma.nlte.nlte_data import NLTEData +from tardis.io.atom_data.nlte_data import NLTEData from tardis.plasma.properties.continuum_processes import ( get_ground_state_multi_index, ) @@ -331,6 +330,10 @@ def _check_related(self): f'were not given: {", ".join(check_list)}' ) + @property + def atom_data_active(self): + return + def prepare_atom_data( self, selected_atomic_numbers, @@ -464,9 +467,9 @@ def prepare_continuum_interaction_data(self, continuum_interaction_species): ) level_idxs2continuum_idx = self.photo_ion_levels_idx.copy() - level_idxs2continuum_idx[ - "continuum_idx" - ] = self.level2continuum_edge_idx + level_idxs2continuum_idx["continuum_idx"] = ( + self.level2continuum_edge_idx + ) self.level_idxs2continuum_idx = level_idxs2continuum_idx.set_index( ["source_level_idx", "destination_level_idx"] ) @@ -500,31 +503,33 @@ def prepare_macro_atom_data( self.macro_atom_references = self.macro_atom_references.loc[ self.macro_atom_references["count_down"] > 0 ] - self.macro_atom_references.loc[ - :, "count_total" - ] = self.macro_atom_references["count_down"] - self.macro_atom_references.loc[ - :, "block_references" - ] = np.hstack( - ( - 0, - np.cumsum( - self.macro_atom_references["count_down"].values[:-1] - ), + self.macro_atom_references.loc[:, "count_total"] = ( + self.macro_atom_references["count_down"] + ) + self.macro_atom_references.loc[:, "block_references"] = ( + np.hstack( + ( + 0, + np.cumsum( + self.macro_atom_references["count_down"].values[ + :-1 + ] + ), + ) ) ) elif line_interaction_type == "macroatom": - self.macro_atom_references.loc[ - :, "block_references" - ] = np.hstack( - ( - 0, - np.cumsum( - self.macro_atom_references["count_total"].values[ - :-1 - ] - ), + self.macro_atom_references.loc[:, "block_references"] = ( + np.hstack( + ( + 0, + np.cumsum( + self.macro_atom_references[ + "count_total" + ].values[:-1] + ), + ) ) ) diff --git a/tardis/plasma/nlte/nlte_data.py b/tardis/io/atom_data/nlte_data.py similarity index 100% rename from tardis/plasma/nlte/nlte_data.py rename to tardis/io/atom_data/nlte_data.py diff --git a/tardis/plasma/nlte/__init__.py b/tardis/plasma/nlte/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 From 5f83fa4b46659635298d7224fab38a1d7aae1f17 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Tue, 30 Jul 2024 17:11:07 -0400 Subject: [PATCH 5/6] Picking apart atomdata --- tardis/io/atom_data/base.py | 207 ++++++-------------- tardis/io/atom_data/collision_data.py | 13 ++ tardis/io/atom_data/decay_radiation_data.py | 3 + tardis/io/atom_data/ionization_data.py | 28 +++ tardis/io/atom_data/levels_data.py | 16 ++ tardis/io/atom_data/lines_data.py | 19 ++ tardis/io/atom_data/macro_atom_data.py | 128 ++++++++++++ tardis/io/atom_data/nlte_data.py | 14 +- 8 files changed, 270 insertions(+), 158 deletions(-) create mode 100644 tardis/io/atom_data/collision_data.py create mode 100644 tardis/io/atom_data/decay_radiation_data.py create mode 100644 tardis/io/atom_data/ionization_data.py create mode 100644 tardis/io/atom_data/levels_data.py create mode 100644 tardis/io/atom_data/lines_data.py create mode 100644 tardis/io/atom_data/macro_atom_data.py diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 936f78bb626..addcfac27cf 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -2,15 +2,23 @@ import numpy as np import pandas as pd -from astropy import units as u -from astropy.units import Quantity from tardis import constants as const +from tardis.io.atom_data.collision_data import CollisionData +from tardis.io.atom_data.decay_radiation_data import DecayRadiationData +from tardis.io.atom_data.ionization_data import ( + IonizationData, + PhotoIonizationData, + ZetaData, +) +from tardis.io.atom_data.levels_data import LevelsData +from tardis.io.atom_data.lines_data import LineList, LinesData +from tardis.io.atom_data.macro_atom_data import MacroAtomData +from tardis.io.atom_data.nlte_data import NLTEData from tardis.io.atom_data.util import ( resolve_atom_data_fname, set_atom_data_attributes, ) -from tardis.io.atom_data.nlte_data import NLTEData from tardis.plasma.properties.continuum_processes import ( get_ground_state_multi_index, ) @@ -27,6 +35,11 @@ class AtomDataMissingError(Exception): logger = logging.getLogger(__name__) +class EnergyData: + def active_data(self): + pass + + class AtomData: """ Class for storing atomic data @@ -270,51 +283,48 @@ def __init__( # the value of constants.u is used in all cases) atom_data.loc[:, "mass"] = atom_data["mass"].values * const.u.cgs.value - # Convert ionization energies to CGS - ionization_data = ionization_data.squeeze() - ionization_data[:] = Quantity(ionization_data[:], "eV").cgs.value - - # Convert energy to CGS - levels.loc[:, "energy"] = Quantity( - levels["energy"].values, "eV" - ).cgs.value - - # Create a new columns with wavelengths in the CGS units - lines["wavelength_cm"] = Quantity( - lines["wavelength"], "angstrom" - ).cgs.value - # SET ATTRIBUTES self.atom_data = atom_data - self.ionization_data = ionization_data - self.levels = levels - # Cast to float so that Numba can use the values in numpy functions - self.levels.energy = self.levels.energy.astype(np.float64) - self.lines = lines - - # Rename these (drop "_all") when `prepare_atom_data` is removed! - self.macro_atom_data_all = macro_atom_data - self.macro_atom_references_all = macro_atom_references + self.ionization_data = IonizationData(ionization_data).data + self.levels = LevelsData(levels).data + self.lines = LinesData(lines).data + self.macro_atom_data_class = MacroAtomData( + macro_atom_data, macro_atom_references + ) + # replace this along with _check_related refactor + self.macro_atom_data_all = ( + self.macro_atom_data_class.transition_probability_data + ) + self.macro_atom_references_all = ( + self.macro_atom_data_class.block_reference_data + ) - self.zeta_data = zeta_data + self.zeta_data = ZetaData(zeta_data).data - self.collision_data = collision_data - self.collision_data_temperatures = collision_data_temperatures + collision_data_class = CollisionData( + collision_data, collision_data_temperatures, yg_data + ) + self.collision_data = collision_data_class.data + self.collision_data_temperatures = collision_data_class.temperatures self.synpp_refs = synpp_refs - self.photoionization_data = photoionization_data + self.photoionization_data = PhotoIonizationData( + photoionization_data + ).data - self.yg_data = yg_data + self.yg_data = collision_data_class.yg self.two_photon_data = two_photon_data if linelist is not None: - self.linelist = linelist + self.linelist = LineList(linelist).data if decay_radiation_data is not None: - self.decay_radiation_data = decay_radiation_data + self.decay_radiation_data = DecayRadiationData( + decay_radiation_data + ).data self._check_related() def _check_related(self): @@ -330,10 +340,6 @@ def _check_related(self): f'were not given: {", ".join(check_list)}' ) - @property - def atom_data_active(self): - return - def prepare_atom_data( self, selected_atomic_numbers, @@ -467,9 +473,9 @@ def prepare_continuum_interaction_data(self, continuum_interaction_species): ) level_idxs2continuum_idx = self.photo_ion_levels_idx.copy() - level_idxs2continuum_idx["continuum_idx"] = ( - self.level2continuum_edge_idx - ) + level_idxs2continuum_idx[ + "continuum_idx" + ] = self.level2continuum_edge_idx self.level_idxs2continuum_idx = level_idxs2continuum_idx.set_index( ["source_level_idx", "destination_level_idx"] ) @@ -481,120 +487,27 @@ def prepare_macro_atom_data( tmp_lines_upper2level_idx, ): if ( - self.macro_atom_data_all is not None + self.macro_atom_data_class.transition_probability_data is not None and not line_interaction_type == "scatter" ): - self.macro_atom_data = self.macro_atom_data_all.loc[ - self.macro_atom_data_all["atomic_number"].isin( - self.selected_atomic_numbers - ) - ].copy() - - self.macro_atom_references = self.macro_atom_references_all[ - self.macro_atom_references_all.index.isin( - self.selected_atomic_numbers, level="atomic_number" - ) - ].copy() - - if line_interaction_type == "downbranch": - self.macro_atom_data = self.macro_atom_data.loc[ - self.macro_atom_data["transition_type"] == -1 - ] - self.macro_atom_references = self.macro_atom_references.loc[ - self.macro_atom_references["count_down"] > 0 - ] - self.macro_atom_references.loc[:, "count_total"] = ( - self.macro_atom_references["count_down"] - ) - self.macro_atom_references.loc[:, "block_references"] = ( - np.hstack( - ( - 0, - np.cumsum( - self.macro_atom_references["count_down"].values[ - :-1 - ] - ), - ) - ) - ) - - elif line_interaction_type == "macroatom": - self.macro_atom_references.loc[:, "block_references"] = ( - np.hstack( - ( - 0, - np.cumsum( - self.macro_atom_references[ - "count_total" - ].values[:-1] - ), - ) - ) - ) - - self.macro_atom_references.loc[:, "references_idx"] = np.arange( - len(self.macro_atom_references) + new_data = self.macro_atom_data_class.active_data( + self.selected_atomic_numbers, + line_interaction_type, + self.lines_index, + tmp_lines_lower2level_idx, + tmp_lines_upper2level_idx, ) - self.macro_atom_data.loc[:, "lines_idx"] = self.lines_index.loc[ - self.macro_atom_data["transition_line_id"] - ].values - + self.macro_atom_data = new_data.transition_probability_data + self.macro_atom_references = new_data.block_reference_data + self.lines_lower2macro_reference_idx = ( + new_data.lines_lower2macro_reference_idx + ) self.lines_upper2macro_reference_idx = ( - self.macro_atom_references.loc[ - tmp_lines_upper2level_idx, "references_idx" - ] - .astype(np.int64) - .values + new_data.lines_upper2macro_reference_idx ) - if line_interaction_type == "macroatom": - self.lines_lower2macro_reference_idx = ( - self.macro_atom_references.loc[ - tmp_lines_lower2level_idx, "references_idx" - ] - .astype(np.int64) - .values - ) - # Sets all - tmp_macro_destination_level_idx = pd.MultiIndex.from_arrays( - [ - self.macro_atom_data["atomic_number"], - self.macro_atom_data["ion_number"], - self.macro_atom_data["destination_level_number"], - ] - ) - - tmp_macro_source_level_idx = pd.MultiIndex.from_arrays( - [ - self.macro_atom_data["atomic_number"], - self.macro_atom_data["ion_number"], - self.macro_atom_data["source_level_number"], - ] - ) - - self.macro_atom_data.loc[:, "destination_level_idx"] = ( - self.macro_atom_references.loc[ - tmp_macro_destination_level_idx, "references_idx" - ] - .astype(np.int64) - .values - ) - - self.macro_atom_data.loc[:, "source_level_idx"] = ( - self.macro_atom_references.loc[ - tmp_macro_source_level_idx, "references_idx" - ] - .astype(np.int64) - .values - ) - - elif line_interaction_type == "downbranch": - # Sets all the destination levels to -1 to indicate that they - # are not used in downbranch calculations - self.macro_atom_data.loc[:, "destination_level_idx"] = -1 - + # TODO: where does this go?! if self.yg_data is not None: self.yg_data = self.yg_data.reindex( self.selected_atomic_numbers, level=0 diff --git a/tardis/io/atom_data/collision_data.py b/tardis/io/atom_data/collision_data.py new file mode 100644 index 00000000000..493d2ab3fa0 --- /dev/null +++ b/tardis/io/atom_data/collision_data.py @@ -0,0 +1,13 @@ +class CollisionData: + def __init__( + self, + collision_energies=None, + temperatures=None, + thermally_averaged_strengths=None, + ): + self.data = collision_energies + self.temperatures = temperatures + self.yg = thermally_averaged_strengths + + def active_data(self): + pass diff --git a/tardis/io/atom_data/decay_radiation_data.py b/tardis/io/atom_data/decay_radiation_data.py new file mode 100644 index 00000000000..e2bda5c26ef --- /dev/null +++ b/tardis/io/atom_data/decay_radiation_data.py @@ -0,0 +1,3 @@ +class DecayRadiationData: + def __init__(self, decay_radiation_data): + self.data = decay_radiation_data diff --git a/tardis/io/atom_data/ionization_data.py b/tardis/io/atom_data/ionization_data.py new file mode 100644 index 00000000000..71e487247b8 --- /dev/null +++ b/tardis/io/atom_data/ionization_data.py @@ -0,0 +1,28 @@ +from astropy.units import Quantity + + +class IonizationData: + def __init__(self, ionization_data): + # Convert ionization energies to CGS + ionization_data = ionization_data.squeeze() + ionization_data[:] = Quantity(ionization_data[:], "eV").cgs.value + self.data = ionization_data + + def active_data(self): + pass + + +class ZetaData: + def __init__(self, zeta_data): + self.data = zeta_data + + def active_data(self): + pass + + +class PhotoIonizationData: + def __init__(self, photoionization_data): + self.data = photoionization_data + + def active_data(self): + pass diff --git a/tardis/io/atom_data/levels_data.py b/tardis/io/atom_data/levels_data.py new file mode 100644 index 00000000000..9aae9611e8d --- /dev/null +++ b/tardis/io/atom_data/levels_data.py @@ -0,0 +1,16 @@ +import numpy as np +from astropy.units import Quantity + + +class LevelsData: + def __init__(self, levels): + # Convert energy to CGS + levels.loc[:, "energy"] = Quantity( + levels["energy"].values, "eV" + ).cgs.value + self.data = levels + # Cast to float so that Numba can use the values in numpy functions + self.data.energy = self.data.energy.astype(np.float64) + + def active_data(self): + pass diff --git a/tardis/io/atom_data/lines_data.py b/tardis/io/atom_data/lines_data.py new file mode 100644 index 00000000000..7e64ea2ea72 --- /dev/null +++ b/tardis/io/atom_data/lines_data.py @@ -0,0 +1,19 @@ +import numpy as np +from astropy.units import Quantity + + +class LinesData: + def __init__(self, lines): + # Convert wavelengths to CGS + lines["wavelength_cm"] = Quantity( + lines["wavelength"], "angstrom" + ).cgs.value + self.data = lines + + def active_data(self): + pass + + +class LineList: + def __init__(self, linelist): + self.data = linelist diff --git a/tardis/io/atom_data/macro_atom_data.py b/tardis/io/atom_data/macro_atom_data.py new file mode 100644 index 00000000000..aabda0868a3 --- /dev/null +++ b/tardis/io/atom_data/macro_atom_data.py @@ -0,0 +1,128 @@ +import numpy as np +import pandas as pd + + +class MacroAtomData: + def __init__( + self, + transition_probabilities=None, + block_references=None, + lines_lower2macro_reference_idx=None, + lines_upper2macro_reference_idx=None, + ): + self.transition_probability_data = transition_probabilities + self.block_reference_data = block_references + self.lines_lower2macro_reference_idx = lines_lower2macro_reference_idx + self.lines_upper2macro_reference_idx = lines_upper2macro_reference_idx + + def active_data( + self, + selected_atomic_numbers, + line_interaction_type, + lines_index, + lines_upper2level_idx, + lines_lower2level_idx, + ): + new_transition_probabilities = self.transition_probability_data.loc[ + self.transition_probability_data["atomic_number"].isin( + selected_atomic_numbers + ) + ].copy() + + new_block_references = self.block_reference_data[ + self.block_reference_data.index.isin( + selected_atomic_numbers, level="atomic_number" + ) + ].copy() + + if line_interaction_type == "downbranch": + new_transition_probabilities = new_transition_probabilities.loc[ + new_transition_probabilities["transition_type"] == -1 + ] + new_block_references = new_block_references.loc[ + new_block_references["count_down"] > 0 + ] + new_block_references.loc[:, "count_total"] = new_block_references[ + "count_down" + ] + new_block_references.loc[:, "block_references"] = np.hstack( + ( + 0, + np.cumsum(new_block_references["count_down"].values[:-1]), + ) + ) + + elif line_interaction_type == "macroatom": + new_block_references.loc[:, "block_references"] = np.hstack( + ( + 0, + np.cumsum(new_block_references["count_total"].values[:-1]), + ) + ) + + new_block_references.loc[:, "references_idx"] = np.arange( + len(new_block_references) + ) + + new_transition_probabilities.loc[:, "lines_idx"] = lines_index.loc[ + new_transition_probabilities["transition_line_id"] + ].values + + lines_upper2macro_reference_idx = ( + new_block_references.loc[lines_upper2level_idx, "references_idx"] + .astype(np.int64) + .values + ) + + if line_interaction_type == "macroatom": + lines_lower2macro_reference_idx = ( + new_block_references.loc[ + lines_lower2level_idx, "references_idx" + ] + .astype(np.int64) + .values + ) + # Sets all + tmp_macro_destination_level_idx = pd.MultiIndex.from_arrays( + [ + new_transition_probabilities["atomic_number"], + new_transition_probabilities["ion_number"], + new_transition_probabilities["destination_level_number"], + ] + ) + + tmp_macro_source_level_idx = pd.MultiIndex.from_arrays( + [ + new_transition_probabilities["atomic_number"], + new_transition_probabilities["ion_number"], + new_transition_probabilities["source_level_number"], + ] + ) + + new_transition_probabilities.loc[:, "destination_level_idx"] = ( + new_block_references.loc[ + tmp_macro_destination_level_idx, "references_idx" + ] + .astype(np.int64) + .values + ) + + new_transition_probabilities.loc[:, "source_level_idx"] = ( + new_block_references.loc[ + tmp_macro_source_level_idx, "references_idx" + ] + .astype(np.int64) + .values + ) + + elif line_interaction_type == "downbranch": + # Sets all the destination levels to -1 to indicate that they + # are not used in downbranch calculations + new_transition_probabilities.loc[:, "destination_level_idx"] = -1 + + return MacroAtomData( + new_transition_probabilities, + new_block_references, + lines_lower2macro_reference_idx, + lines_upper2macro_reference_idx, + ) diff --git a/tardis/io/atom_data/nlte_data.py b/tardis/io/atom_data/nlte_data.py index b5381100d4a..86bbb341f7b 100644 --- a/tardis/io/atom_data/nlte_data.py +++ b/tardis/io/atom_data/nlte_data.py @@ -39,9 +39,9 @@ def _init_indices(self): species ] = self.lines.level_number_upper.values[lines_idx].astype(int) - self.A_uls[species] = self.atom_data.lines.A_ul.values[lines_idx] - self.B_uls[species] = self.atom_data.lines.B_ul.values[lines_idx] - self.B_lus[species] = self.atom_data.lines.B_lu.values[lines_idx] + self.A_uls[species] = self.lines.A_ul.values[lines_idx] + self.B_uls[species] = self.lines.B_ul.values[lines_idx] + self.B_lus[species] = self.lines.B_lu.values[lines_idx] def _create_collision_coefficient_matrix(self): self.C_ul_interpolator = {} @@ -113,11 +113,3 @@ def get_collision_matrix(self, species, t_electrons): ) ) return c_ul_matrix + c_lu_matrix.transpose(1, 0, 2) - - -class NLTEMatrixFactory: - def __init__(self, nlte_species): - pass - - def generate_indices(self): - pass From 51cdf4277ae0e1eb8b76f693674425a34eb64519 Mon Sep 17 00:00:00 2001 From: Andrew Fullard Date: Wed, 31 Jul 2024 16:32:23 -0400 Subject: [PATCH 6/6] Further split --- tardis/io/atom_data/atomic_mass_data.py | 14 ++++++++++++++ tardis/io/atom_data/base.py | 22 +++++++--------------- tardis/io/atom_data/lines_data.py | 1 - tardis/io/atom_data/two_photon_data.py | 3 +++ 4 files changed, 24 insertions(+), 16 deletions(-) create mode 100644 tardis/io/atom_data/atomic_mass_data.py create mode 100644 tardis/io/atom_data/two_photon_data.py diff --git a/tardis/io/atom_data/atomic_mass_data.py b/tardis/io/atom_data/atomic_mass_data.py new file mode 100644 index 00000000000..9587863a80e --- /dev/null +++ b/tardis/io/atom_data/atomic_mass_data.py @@ -0,0 +1,14 @@ +from tardis import constants as const + + +class AtomicMassData: + def __init__(self, atomic_mass_data): + # Convert atomic masses to CGS + # We have to use constants.u because astropy uses + # different values for the unit u and the constant. + # This is changed in later versions of astropy ( + # the value of constants.u is used in all cases) + atomic_mass_data.loc[:, "mass"] = ( + atomic_mass_data["mass"].values * const.u.cgs.value + ) + self.data = atomic_mass_data diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index addcfac27cf..8cd6e81274d 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -3,7 +3,7 @@ import numpy as np import pandas as pd -from tardis import constants as const +from tardis.io.atom_data.atomic_mass_data import AtomicMassData from tardis.io.atom_data.collision_data import CollisionData from tardis.io.atom_data.decay_radiation_data import DecayRadiationData from tardis.io.atom_data.ionization_data import ( @@ -15,6 +15,7 @@ from tardis.io.atom_data.lines_data import LineList, LinesData from tardis.io.atom_data.macro_atom_data import MacroAtomData from tardis.io.atom_data.nlte_data import NLTEData +from tardis.io.atom_data.two_photon_data import TwoPhotonData from tardis.io.atom_data.util import ( resolve_atom_data_fname, set_atom_data_attributes, @@ -274,18 +275,9 @@ def __init__( ): self.prepared = False - # CONVERT VALUES TO CGS UNITS - - # Convert atomic masses to CGS - # We have to use constants.u because astropy uses - # different values for the unit u and the constant. - # This is changed in later versions of astropy ( - # the value of constants.u is used in all cases) - atom_data.loc[:, "mass"] = atom_data["mass"].values * const.u.cgs.value - # SET ATTRIBUTES - self.atom_data = atom_data + self.atom_data = AtomicMassData(atom_data).data self.ionization_data = IonizationData(ionization_data).data self.levels = LevelsData(levels).data self.lines = LinesData(lines).data @@ -316,7 +308,7 @@ def __init__( self.yg_data = collision_data_class.yg - self.two_photon_data = two_photon_data + self.two_photon_data = TwoPhotonData(two_photon_data).data if linelist is not None: self.linelist = LineList(linelist).data @@ -473,9 +465,9 @@ def prepare_continuum_interaction_data(self, continuum_interaction_species): ) level_idxs2continuum_idx = self.photo_ion_levels_idx.copy() - level_idxs2continuum_idx[ - "continuum_idx" - ] = self.level2continuum_edge_idx + level_idxs2continuum_idx["continuum_idx"] = ( + self.level2continuum_edge_idx + ) self.level_idxs2continuum_idx = level_idxs2continuum_idx.set_index( ["source_level_idx", "destination_level_idx"] ) diff --git a/tardis/io/atom_data/lines_data.py b/tardis/io/atom_data/lines_data.py index 7e64ea2ea72..8395142640a 100644 --- a/tardis/io/atom_data/lines_data.py +++ b/tardis/io/atom_data/lines_data.py @@ -1,4 +1,3 @@ -import numpy as np from astropy.units import Quantity diff --git a/tardis/io/atom_data/two_photon_data.py b/tardis/io/atom_data/two_photon_data.py new file mode 100644 index 00000000000..6e04bda028e --- /dev/null +++ b/tardis/io/atom_data/two_photon_data.py @@ -0,0 +1,3 @@ +class TwoPhotonData: + def __init__(self, two_photon_data): + self.data = two_photon_data