Molecules (#210)

* fix plasma

* use aug 25 release

* molecules checkpoint

* fix logging

* spectrum checkpoint - probably incorrect

* separate out partition functions

* update docstrings

* clearer variable names

* slow checkpoint

* optimization

* add comments

* add shortlist handling

* densities -> density

* shortlist handling bugfix

* use mapping for ions instead of holding in densities df

* give value error if no molecular data provided

* docstrings

* clearer row iteration

* explain alpha

* add preprocessing function

* make preprocess function return processed dataframe

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/config_schema.yml

@@ -126,6 +126,10 @@ properties:
                                 type: boolean
                                 default: false
                         description: Options for whether to use a vald linelist. Linelist must be included the atom_data and cannot be supplied separately. 
+                    include_molecules:
+                        type: boolean
+                        default: false
+                        description: Whether to include molecular lines in the simulation.
         #required:
         #- line
     no_of_thetas:

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,6 +19,13 @@
     non_nlte_properties,
     helium_lte_properties,
 )
+from stardis.plasma.molecules import (
+    MoleculeIonNumberDensity,
+    AlphaLineValdMolecule,
+    MoleculePartitionFunction,
+    AlphaLineShortlistValdMolecule,
+)
+
 
 import tardis.plasma
 from tardis.opacities.tau_sobolev import TauSobolev
@@ -50,6 +57,8 @@
     25200,
 ]  # see directly above
 
+logger = logging.getLogger(__name__)
+
 
 class HMinusDensity(ProcessingPlasmaProperty):
     """
@@ -122,7 +131,7 @@ class AlphaLine(ProcessingPlasmaProperty):
     Attributes
     ----------
     alpha_line : Pandas DataFrame, dtype float
-        Sobolev optical depth for each line. Indexed by line.
+        Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each line at each depth point. Indexed by line.
         Columns as zones.
     """
 
@@ -167,6 +176,7 @@ def calculate(
 
 class AlphaLineVald(ProcessingPlasmaProperty):
     """
+    Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each line at each depth point. Uses VALD linelists for lines. Indexed by line.
     Attributes
     ----------
     alpha_line_from_linelist : DataFrame
@@ -200,7 +210,6 @@ def calculate(
         # alphas = ALPHA_COEFFICIENT * n_lower * f_lu * emission_correction
 
         ###TODO: handle other broadening parameters
-
         points = len(t_electrons)
 
         linelist = atomic_data.linelist_atoms.rename(
@@ -301,8 +310,8 @@ def calculate(
         linelist["level_energy_upper"] = ((linelist["e_up"].values * u.eV).cgs).value
 
         # Radiation broadening parameter is approximated as the einstein A coefficient. Vald parameters are in log scale.
-        linelist["A_ul"] = 10 ** (
-            linelist["rad"]
+        linelist["A_ul"] = 10 ** (linelist["rad"]) / (
+            4 * np.pi
         )  # 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
@@ -313,6 +322,9 @@ def calculate(
 
 class AlphaLineShortlistVald(ProcessingPlasmaProperty):
     """
+    Calculates the alpha (line integrated absorption coefficient in cm^-1) values for each line at each depth point. Uses VALD shortform linelists for lines. Indexed by line.
+    Subtley different from the full list calculation in that it does not require the upper level degeneracy, and pure number density is not directly calculated as a result.
+
     Attributes
     ----------
     alpha_line_from_linelist : DataFrame
@@ -496,7 +508,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 = []
@@ -538,6 +550,14 @@ def create_stellar_plasma(
         temperature=stellar_model.temperatures,
         dilution_factor=np.ones_like(stellar_model.temperatures),
     )
+    if config.opacity.line.include_molecules:
+        plasma_modules.append(MoleculeIonNumberDensity)
+        plasma_modules.append(MoleculePartitionFunction)
+        if config.opacity.line.vald_linelist.use_linelist:
+            if config.opacity.line.vald_linelist.shortlist:
+                plasma_modules.append(AlphaLineShortlistValdMolecule)
+            else:
+                plasma_modules.append(AlphaLineValdMolecule)
 
     return BasePlasma(
         plasma_properties=plasma_modules,