fix logging

stardis/base.py

@@ -7,6 +7,8 @@
 
 import logging
 
+logger = logging.getLogger(__name__)
+
 
 def run_stardis(
     config_fname, tracing_lambdas_or_nus, add_config_keys=None, add_config_vals=None
@@ -73,11 +75,11 @@ def set_num_threads(n_threads):
 
     """
     if n_threads == 1:
-        logging.info("Running in serial mode")
+        logger.info("Running in serial mode")
     elif n_threads == -99:
-        logging.info("Running with max threads")
+        logger.info("Running with max threads")
     elif n_threads > 1:
-        logging.info(f"Running with {n_threads} threads")
+        logger.info(f"Running with {n_threads} threads")
         numba.set_num_threads(n_threads)
     else:
         raise ValueError(

stardis/io/base.py

@@ -14,6 +14,8 @@
 BASE_DIR = Path(__file__).parent.parent
 SCHEMA_PATH = BASE_DIR / "config_schema.yml"
 
+logger = logging.getLogger(__name__)
+
 
 def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=None):
     """
@@ -49,7 +51,7 @@ def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=No
     ):  # If a dictionary was passed, update the config with the dictionary
         pass
     else:
-        print("Updating config with additional keys and values")
+        logger.info("Updating config with additional keys and values")
         if isinstance(add_config_keys, str):
             # Directly set the config item if add_config_keys is a string
             config.set_config_item(add_config_keys, add_config_vals)
@@ -75,7 +77,7 @@ def parse_config_to_model(config_fname, add_config_keys=None, add_config_vals=No
     adata = AtomData.from_hdf(config.atom_data)
 
     # model
-    logging.info("Reading model")
+    logger.info("Reading model")
     if config.model.type == "marcs":
         raw_marcs_model = read_marcs_model(
             Path(config.model.fname), gzipped=config.model.gzipped

stardis/plasma/base.py

@@ -19,7 +19,7 @@
     non_nlte_properties,
     helium_lte_properties,
 )
-from stardis.plasma.molecules import MoleculeNumberDensities
+from stardis.plasma.molecules import MoleculeNumberDensities, AlphaLineValdMolecules
 
 
 import tardis.plasma
@@ -52,6 +52,8 @@
     25200,
 ]  # see directly above
 
+logger = logging.getLogger(__name__)
+
 
 class HMinusDensity(ProcessingPlasmaProperty):
     """
@@ -447,9 +449,6 @@ def calculate(
         return alphas[valid_indices], linelist[valid_indices]
 
 
-
-
-
 # Properties that haven't been used in creating stellar plasma yet,
 # might be useful in future ----------------------------------------------------
 
@@ -500,7 +499,7 @@ def create_stellar_plasma(
     tardis.plasma.base.BasePlasma
     """
 
-    logging.info("Creating plasma")
+    logger.info("Creating plasma")
 
     # basic_properties.remove(tardis.plasma.properties.general.NumberDensity)
     plasma_modules = []
@@ -542,8 +541,9 @@ def create_stellar_plasma(
         temperature=stellar_model.temperatures,
         dilution_factor=np.ones_like(stellar_model.temperatures),
     )
-
-    plasma_modules.append(MoleculeNumberDensities)
+    if True:
+        plasma_modules.append(MoleculeNumberDensities)
+        plasma_modules.append(AlphaLineValdMolecules)
 
     return BasePlasma(
         plasma_properties=plasma_modules,

stardis/plasma/molecules.py

@@ -7,6 +7,7 @@
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
 
+ALPHA_COEFFICIENT = (np.pi * const.e.gauss**2) / (const.m_e.cgs * const.c.cgs)
 
 logger = logging.getLogger(__name__)
 
@@ -18,18 +19,22 @@ class MoleculeNumberDensities(ProcessingPlasmaProperty):
 
     outputs = ("molecule_number_densities",)
 
-    def calculate(self, ion_number_density, t_rad, atomic_data):
+    def calculate(self, ion_number_density, t_electrons, atomic_data):
         # This first implementation takes ~half a second. Much slower than is reasonable I think.
 
         number_densities_arr = np.zeros(
-            (len(atomic_data.molecule_data.equilibrium_constants), len(t_rad))
+            (len(atomic_data.molecule_data.equilibrium_constants), len(t_electrons))
         )
 
+        ###TODO - KEEP TRACK OF THE IONS SO WE CAN GET THEIR MASSES LATER FOR DOPPLER BROADENING PURPOSES
+        ions_arr = np.zeros((len(atomic_data.molecule_data.equilibrium_constants), 2))
+
         equilibrium_const_temps = (
             atomic_data.molecule_data.equilibrium_constants.columns.values
         )
         included_elements = ion_number_density.index.get_level_values(0).unique()
 
+        row_tracker = 0
         for row in atomic_data.molecule_data.dissociation_energies.iterrows():
             ionization_state_1 = 0
             ionization_state_2 = 0
@@ -43,6 +48,8 @@ def calculate(self, ion_number_density, t_rad, atomic_data):
 
                 ion1 = element_symbol2atomic_number(ion1_arr[0])
                 ion2 = element_symbol2atomic_number(ion2_arr[0])
+                ions_arr[row_tracker] = [ion1, ion2]
+                row_tracker += 1
                 if ion1 not in included_elements:
                     logger.warning(
                         f"{row[1].Ion1} not in included elements. Assuming no {row[0]}."
@@ -54,19 +61,20 @@ def calculate(self, ion_number_density, t_rad, atomic_data):
                     )
                     continue
             except:
+                row_tracker += 1
                 continue  # This will currently skip over negative ions
             ion1_number_density = ion_number_density.loc[ion1, ionization_state_1]
             ion2_number_density = ion_number_density.loc[ion2, ionization_state_2]
 
             pressure_equilibirium_const_at_depth_point = np.interp(
-                t_rad,
+                t_electrons,
                 equilibrium_const_temps,
                 atomic_data.molecule_data.equilibrium_constants.loc[row[0]].values,
             )
             equilibirium_const_at_depth_point = (
                 10 ** (pressure_equilibirium_const_at_depth_point)
                 * (u.N * const.N_A / u.m**2)
-                / (const.R * t_rad * u.K)
+                / (const.R * t_electrons * u.K)
             ).cgs.value
 
             molecule_number_density = (
@@ -77,13 +85,17 @@ def calculate(self, ion_number_density, t_rad, atomic_data):
                 atomic_data.molecule_data.equilibrium_constants.index.get_loc(row[0])
             ] = molecule_number_density
 
-        return pd.DataFrame(
+        densities_df = pd.DataFrame(
             number_densities_arr,
             index=atomic_data.molecule_data.equilibrium_constants.index,
             columns=ion_number_density.columns,
         )
+        densities_df["ion1"] = ions_arr[:, 0].astype(int)
+        densities_df["ion2"] = ions_arr[:, 1].astype(int)
+        return densities_df
+
 
-class AlphaLineMolecules(ProcessingPlasmaProperty):
+class AlphaLineValdMolecules(ProcessingPlasmaProperty):
     """
     Attributes
     ----------
@@ -118,10 +130,9 @@ def calculate(
         ###TODO: handle other broadening parameters
         points = len(t_electrons)
 
-        linelist = atomic_data.linelist.rename(columns={"ion_charge": "ion_number"})[
+        linelist = atomic_data.linelist_molecules[
             [
-                "atomic_number",
-                "ion_number",
+                "molecule",
                 "wavelength",
                 "log_gf",
                 "e_low",
@@ -134,11 +145,6 @@ def calculate(
             ]
         ]
 
-        # Truncate to final atomic number
-        linelist = linelist[
-            linelist.atomic_number <= (atomic_data.selected_atomic_numbers.max())
-        ]
-
         # Calculate degeneracies
         linelist["g_lo"] = linelist.j_lo * 2 + 1
         linelist["g_up"] = linelist.j_up * 2 + 1
@@ -149,11 +155,14 @@ def calculate(
             ).to(1)
         )
 
+        prepared_molecule_number_densities = molecule_number_densities.copy()
+        prepared_molecule_number_densities.index.name = "molecule"
+
         # grab densities for n_lower - need to use linelist as the index and normalize by dividing by the partition function
         linelist_with_densities = linelist.merge(
-            ion_number_density / partition_function,
+            prepared_molecule_number_densities,
             how="left",
-            on=["atomic_number", "ion_number"],
+            on=["molecule"],
         )
 
         n_lower = (
@@ -200,16 +209,6 @@ def calculate(
         linelist["nu"] = line_nus.value
 
         # Linelist preparation below is taken from opacities_solvers/base/calc_alpha_line_at_nu
-        # Necessary for correct handling of ion numbers using charge instead of astronomy convention (i.e., 0 is neutral, 1 is singly ionized, etc.)
-        ionization_energies = ionization_data.reset_index()
-        ionization_energies["ion_number"] -= 1
-        linelist = pd.merge(
-            linelist,
-            ionization_energies,
-            how="left",
-            on=["atomic_number", "ion_number"],
-        )
-
         linelist["level_energy_lower"] = ((linelist["e_low"].values * u.eV).cgs).value
         linelist["level_energy_upper"] = ((linelist["e_up"].values * u.eV).cgs).value
 
@@ -218,7 +217,4 @@ def calculate(
             linelist["rad"]
         )  # see 1995A&AS..112..525P for appropriate units - may be off by a factor of 4pi
 
-        # Need to remove autoionization lines - can't handle with current broadening treatment because can't calculate effective principal quantum number
-        valid_indices = linelist.level_energy_upper < linelist.ionization_energy
-
-        return alphas[valid_indices], linelist[valid_indices]
+        return alphas, linelist

stardis/radiation_field/base.py

@@ -5,6 +5,8 @@
 from stardis.radiation_field.radiation_field_solvers import raytrace
 from stardis.radiation_field.source_functions.blackbody import blackbody_flux_at_nu
 
+logger = logging.getLogger(__name__)
+
 
 class RadiationField:
     """
@@ -66,15 +68,15 @@ def create_stellar_radiation_field(tracing_nus, stellar_model, stellar_plasma, c
     stellar_radiation_field = RadiationField(
         tracing_nus, blackbody_flux_at_nu, stellar_model
     )
-    logging.info("Calculating alphas")
+    logger.info("Calculating alphas")
     calc_alphas(
         stellar_plasma=stellar_plasma,
         stellar_model=stellar_model,
         stellar_radiation_field=stellar_radiation_field,
         opacity_config=config.opacity,
         n_threads=config.n_threads,
     )
-    logging.info("Raytracing")
+    logger.info("Raytracing")
     raytrace(
         stellar_model,
         stellar_radiation_field,