From 0a4831f12045afb5d4e0363923e267532764c05e Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Mon, 19 Aug 2024 13:22:50 -0400 Subject: [PATCH 01/12] apply black --- tardis/io/atom_data/base.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 6596bcefd7b..1dc7e0506f2 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -221,6 +221,17 @@ def from_hdf(cls, fname=None): if "linelist" in store: dataframes["linelist"] = store["linelist"] + if "molecules" in store: + dataframes["molecular_equilibrium_constants"] = store[ + "molecules/equilibrium_constants" + ] + dataframes["molecular_partition_functions"] = store[ + "molecules/partition_functions" + ] + dataframes["molecular_dissociation_energies"] = store[ + "molecules/dissociation_energies" + ] + atom_data = cls(**dataframes) try: @@ -284,6 +295,9 @@ def __init__( two_photon_data=None, linelist=None, decay_radiation_data=None, + molecular_equilibrium_constants=None, + molecular_partition_functions=None, + molecular_dissociation_energies=None, ): self.prepared = False @@ -346,6 +360,17 @@ def __init__( if linelist is not None: self.linelist = linelist + if molecular_equilibrium_constants is not None: + self.molecular_equilibrium_constants = ( + molecular_equilibrium_constants + ) + if molecular_partition_functions is not None: + self.molecular_partition_functions = molecular_partition_functions + if molecular_dissociation_energies is not None: + self.molecular_dissociation_energies = ( + molecular_dissociation_energies + ) + if decay_radiation_data is not None: self.decay_radiation_data = decay_radiation_data self._check_related() From 65d744900ca514fd289018386a6ae3631aef449f Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Mon, 19 Aug 2024 13:33:23 -0400 Subject: [PATCH 02/12] add docstrings to class --- tardis/io/atom_data/base.py | 86 +++++++++++++++++++++++-------------- 1 file changed, 53 insertions(+), 33 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 1dc7e0506f2..c9a4883a9ba 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -100,6 +100,21 @@ class AtomData: columns: atomic_number, element, Rad energy, Rad intensity decay mode. Curated from nndc + molecular_equilibrium_constants : pandas.DataFrame + A DataFrame containing the *molecular equilibrium constants* with: + index: molecule + columns: temperatures + + molecular_partition_functions : pandas.DataFrame + A DataFrame containing the *molecular partition functions* with: + index: molecule + columns: temperatures + + molecular_dissociation_energies : pandas.DataFrame + A DataFrame containing the *molecular dissociation energies* with: + index: molecule + + Attributes ---------- prepared : bool @@ -116,6 +131,9 @@ class AtomData: photoionization_data : pandas.DataFrame two_photon_data : pandas.DataFrame decay_radiation_data : pandas.DataFrame + molecular_equilibrium_constants : pandas.DataFrame + molecular_partition_functions : pandas.DataFrame + molecular_dissociation_energies : pandas.DataFrame Methods ------- @@ -528,9 +546,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"] ) @@ -564,31 +582,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] + ), + ) ) ) @@ -706,12 +726,12 @@ def _init_indices(self): & (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.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] @@ -746,9 +766,9 @@ def _create_collision_coefficient_matrix(self): 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:] + C_ul_matrix[level_number_lower, level_number_upper, :] = ( + line.values[2:] + ) delta_E_matrix[level_number_lower, level_number_upper] = line[ "delta_e" ] From 3f8b9cf425596f08ccc1625ec640e6e20e43779a Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Mon, 19 Aug 2024 13:35:02 -0400 Subject: [PATCH 03/12] apply black --- tardis/io/atom_data/base.py | 68 ++++++++++++++++++------------------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index c9a4883a9ba..76fc424a3c4 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -546,9 +546,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"] ) @@ -582,33 +582,31 @@ 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 + ] + ), ) ) @@ -726,12 +724,12 @@ def _init_indices(self): & (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.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] @@ -766,9 +764,9 @@ def _create_collision_coefficient_matrix(self): 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:] - ) + C_ul_matrix[ + level_number_lower, level_number_upper, : + ] = line.values[2:] delta_E_matrix[level_number_lower, level_number_upper] = line[ "delta_e" ] From 358ee415c36026245b716f9412e45572fb25fa09 Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Tue, 20 Aug 2024 11:09:33 -0400 Subject: [PATCH 04/12] change to dataclass object --- tardis/io/atom_data/base.py | 79 +++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 39 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 76fc424a3c4..06d26df4639 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -6,6 +6,7 @@ from astropy import units as u from astropy.units import Quantity from scipy import interpolate +from dataclasses import dataclass from tardis import constants as const from tardis.io.atom_data.util import resolve_atom_data_fname @@ -100,19 +101,9 @@ class AtomData: columns: atomic_number, element, Rad energy, Rad intensity decay mode. Curated from nndc - molecular_equilibrium_constants : pandas.DataFrame - A DataFrame containing the *molecular equilibrium constants* with: - index: molecule - columns: temperatures - - molecular_partition_functions : pandas.DataFrame - A DataFrame containing the *molecular partition functions* with: - index: molecule - columns: temperatures - - molecular_dissociation_energies : pandas.DataFrame - A DataFrame containing the *molecular dissociation energies* with: - index: molecule + molecule_data : MolecularData + A class containing the *molecular data* with: + equilibrium_constants, partition_functions, dissociation_energies Attributes @@ -131,9 +122,7 @@ class AtomData: photoionization_data : pandas.DataFrame two_photon_data : pandas.DataFrame decay_radiation_data : pandas.DataFrame - molecular_equilibrium_constants : pandas.DataFrame - molecular_partition_functions : pandas.DataFrame - molecular_dissociation_energies : pandas.DataFrame + molecule_data : MolecularData Methods ------- @@ -240,17 +229,13 @@ def from_hdf(cls, fname=None): dataframes["linelist"] = store["linelist"] if "molecules" in store: - dataframes["molecular_equilibrium_constants"] = store[ - "molecules/equilibrium_constants" - ] - dataframes["molecular_partition_functions"] = store[ - "molecules/partition_functions" - ] - dataframes["molecular_dissociation_energies"] = store[ - "molecules/dissociation_energies" - ] + molecule_data = MoleculeData( + store["molecules/equilibrium_constants"], + store["molecules/partition_functions"], + store["molecules/dissociation_energies"], + ) - atom_data = cls(**dataframes) + atom_data = cls(**dataframes, molecule_data=molecule_data) try: atom_data.uuid1 = store.root._v_attrs["uuid1"] @@ -313,9 +298,7 @@ def __init__( two_photon_data=None, linelist=None, decay_radiation_data=None, - molecular_equilibrium_constants=None, - molecular_partition_functions=None, - molecular_dissociation_energies=None, + molecule_data=None, ): self.prepared = False @@ -378,16 +361,8 @@ def __init__( if linelist is not None: self.linelist = linelist - if molecular_equilibrium_constants is not None: - self.molecular_equilibrium_constants = ( - molecular_equilibrium_constants - ) - if molecular_partition_functions is not None: - self.molecular_partition_functions = molecular_partition_functions - if molecular_dissociation_energies is not None: - self.molecular_dissociation_energies = ( - molecular_dissociation_energies - ) + if molecule_data is not None: + self.molecule_data = molecule_data if decay_radiation_data is not None: self.decay_radiation_data = decay_radiation_data @@ -805,3 +780,29 @@ def get_collision_matrix(self, species, t_electrons): ) ) return c_ul_matrix + c_lu_matrix.transpose(1, 0, 2) + + +@dataclass +class MoleculeData(object): + """ + Class to hold molecular data. Held by the AtomData object. + + equilibrium_constants : pandas.DataFrame + A DataFrame containing the *molecular equilibrium constants* with: + index: molecule + columns: temperatures + + partition_functions : pandas.DataFrame + A DataFrame containing the *molecular partition functions* with: + index: molecule + columns: temperatures + + dissociation_energies : pandas.DataFrame + A DataFrame containing the *molecular dissociation energies* with: + index: molecule + + """ + + equilibrium_constants: pd.DataFrame + partition_functions: pd.DataFrame + dissociation_energies: pd.DataFrame From 3a3efb5b65439df1225d6e103dbabc55227a307b Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Tue, 20 Aug 2024 11:13:42 -0400 Subject: [PATCH 05/12] remove object inheritance --- tardis/io/atom_data/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 06d26df4639..e957cdfcbe3 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -783,7 +783,7 @@ def get_collision_matrix(self, species, t_electrons): @dataclass -class MoleculeData(object): +class MoleculeData(): """ Class to hold molecular data. Held by the AtomData object. From 5e21bf1994f7eafe11d12b2deaeda6c8cb50e971 Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Tue, 20 Aug 2024 11:16:01 -0400 Subject: [PATCH 06/12] apply black --- tardis/io/atom_data/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index e957cdfcbe3..fcd914f94e3 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -783,7 +783,7 @@ def get_collision_matrix(self, species, t_electrons): @dataclass -class MoleculeData(): +class MoleculeData: """ Class to hold molecular data. Held by the AtomData object. From f6c2019d061626f22109aa064a7493c4e0915573 Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Tue, 20 Aug 2024 11:26:09 -0400 Subject: [PATCH 07/12] fix benchmark bug --- tardis/io/atom_data/base.py | 70 ++++++++++++++++++++----------------- 1 file changed, 37 insertions(+), 33 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index fcd914f94e3..4e9a3d63ca3 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -234,6 +234,8 @@ def from_hdf(cls, fname=None): store["molecules/partition_functions"], store["molecules/dissociation_energies"], ) + else: + molecule_data = None atom_data = cls(**dataframes, molecule_data=molecule_data) @@ -521,9 +523,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"] ) @@ -557,31 +559,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] + ), + ) ) ) @@ -699,12 +703,12 @@ def _init_indices(self): & (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.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] @@ -739,9 +743,9 @@ def _create_collision_coefficient_matrix(self): 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:] + C_ul_matrix[level_number_lower, level_number_upper, :] = ( + line.values[2:] + ) delta_E_matrix[level_number_lower, level_number_upper] = line[ "delta_e" ] From 0cc4060552a7d073ba6a10176a58216434ba879a Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Tue, 20 Aug 2024 11:27:23 -0400 Subject: [PATCH 08/12] apply black --- tardis/io/atom_data/base.py | 68 ++++++++++++++++++------------------- 1 file changed, 33 insertions(+), 35 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 4e9a3d63ca3..d23707fc31e 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -523,9 +523,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"] ) @@ -559,33 +559,31 @@ 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 + ] + ), ) ) @@ -703,12 +701,12 @@ def _init_indices(self): & (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.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] @@ -743,9 +741,9 @@ def _create_collision_coefficient_matrix(self): 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:] - ) + C_ul_matrix[ + level_number_lower, level_number_upper, : + ] = line.values[2:] delta_E_matrix[level_number_lower, level_number_upper] = line[ "delta_e" ] From f46667bc78620f6354e4207905045d0a85a12743 Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Tue, 20 Aug 2024 11:38:19 -0400 Subject: [PATCH 09/12] apply black (I really need to fix my local formatting extension --- tardis/io/atom_data/base.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index 80eb4d963ed..ff824fce4fb 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -691,7 +691,7 @@ def get_attributes_from_store(store, store_key): ) attribute = None - return attribute + return attribute @dataclass @@ -718,4 +718,3 @@ class MoleculeData: equilibrium_constants: pd.DataFrame partition_functions: pd.DataFrame dissociation_energies: pd.DataFrame - \ No newline at end of file From 82a364ae3a783d974e5fd558b8818d0e2d28fb3d Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Thu, 22 Aug 2024 14:06:47 -0400 Subject: [PATCH 10/12] add molecule exposure for vald linelists --- tardis/io/atom_data/base.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index ff824fce4fb..db22b9105a9 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -154,7 +154,8 @@ class AtomData: "photoionization_data", "yg_data", "two_photon_data", - "linelist", + "linelist_atoms", + "linelist_molecules", "decay_radiation_data", ] @@ -226,8 +227,10 @@ def from_hdf(cls, fname=None): raise ValueError( f"Current carsus version, {carsus_version}, is not supported." ) - if "linelist" in store: - dataframes["linelist"] = store["linelist"] + if "linelist_atoms" in store: + dataframes["linelist_atoms"] = store["linelist_atoms"] + if "linelist_molecules" in store: + dataframes["linelist_molecules"] = store["linelist_molecules"] if "molecules" in store: molecule_data = MoleculeData( @@ -275,7 +278,8 @@ def __init__( photoionization_data=None, yg_data=None, two_photon_data=None, - linelist=None, + linelist_atoms=None, + linelist_molecules=None, decay_radiation_data=None, molecule_data=None, ): @@ -346,8 +350,10 @@ def __init__( self.two_photon_data = two_photon_data - if linelist is not None: - self.linelist = linelist + if linelist_atoms is not None: + self.linelist_atoms = linelist_atoms + if linelist_molecules is not None: + self.linelist_molecules = linelist_molecules if molecule_data is not None: self.molecule_data = molecule_data From c2a1aa04d4e08df0de54bbb734a7f8fae87a55aa Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Thu, 22 Aug 2024 14:27:34 -0400 Subject: [PATCH 11/12] update docstrings --- tardis/io/atom_data/base.py | 58 ++++++++++++++++++++++--------------- 1 file changed, 34 insertions(+), 24 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index db22b9105a9..fa6e1abace0 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -104,6 +104,12 @@ class AtomData: columns: atomic_number, element, Rad energy, Rad intensity decay mode. Curated from nndc + linelist_atoms : pandas.DataFrame + A DataFrame containing a linelist of input atoms + + linelist_molecules : pandas.DataFrame + A DataFrame containing a linelist of input molecules + molecule_data : MolecularData A class containing the *molecular data* with: equilibrium_constants, partition_functions, dissociation_energies @@ -123,6 +129,8 @@ class AtomData: photoionization_data : pandas.DataFrame two_photon_data : pandas.DataFrame decay_radiation_data : pandas.DataFrame + linelist_atoms : pandas.DataFrame + linelist_molecules : pandas.DataFrame molecule_data : MolecularData Methods @@ -520,9 +528,9 @@ def prepare_continuum_interaction_data(self, continuum_interaction_species): ) level_idxs2continuum_idx = 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"] ) @@ -556,31 +564,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] + ), + ) ) ) From ec4d62ad7b17332ad72001f315e642ade5bfd9ec Mon Sep 17 00:00:00 2001 From: Josh Shields Date: Thu, 22 Aug 2024 14:28:46 -0400 Subject: [PATCH 12/12] blackify --- tardis/io/atom_data/base.py | 50 ++++++++++++++++++------------------- 1 file changed, 24 insertions(+), 26 deletions(-) diff --git a/tardis/io/atom_data/base.py b/tardis/io/atom_data/base.py index fa6e1abace0..3f53408a210 100644 --- a/tardis/io/atom_data/base.py +++ b/tardis/io/atom_data/base.py @@ -528,9 +528,9 @@ def prepare_continuum_interaction_data(self, continuum_interaction_species): ) level_idxs2continuum_idx = 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"] ) @@ -564,33 +564,31 @@ 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 + ] + ), ) )