add shortlist handling

stardis/plasma/base.py

@@ -308,8 +308,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
@@ -547,8 +547,8 @@ def create_stellar_plasma(
     )
     if config.opacity.line.include_molecules:
         plasma_modules.append(MoleculeIonNumberDensities)
-        plasma_modules.append(AlphaLineValdMolecules)
         plasma_modules.append(MoleculePartitionFunctions)
+        plasma_modules.append(AlphaLineValdMolecules)
 
     return BasePlasma(
         plasma_properties=plasma_modules,

stardis/plasma/molecules.py

@@ -274,8 +274,135 @@ 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
+
+        return alphas, linelist
+
+
+class AlphaLineShortlistValdMolecules(ProcessingPlasmaProperty):
+    """
+    Attributes
+    ----------
+    alpha_line_from_linelist : DataFrame
+            A pandas DataFrame with dtype float. This represents the alpha calculation
+            for each line from Vald at each depth point. This is adapted from the AlphaLineVald calculation
+            because shortlists do not contain js or an upper energy level. This works because the degeneracies
+            cancel between the n_lower recalculation and in calculating f_lu from g*f given by the linelist.
+    lines_from_linelist: DataFrame
+            A pandas dataframe containing the lines and information about the lines in
+            the same form and shape as alpha_line_from_linelist.
+    """
+
+    outputs = ("alpha_line_from_linelist", "lines_from_linelist")
+    latex_name = (r"\alpha_{\textrm{line, vald}}",)
+    latex_formula = (
+        r"\dfrac{\pi e^{2} n_{lower} f_{lu}}{m_{e} c}\
+        \Big(1-exp(-h \nu / k T) \phi(\nu)\Big)",
+    )
+
+    def calculate(
+        self,
+        atomic_data,
+        molecule_number_densities,
+        t_electrons,
+        molecule_partition_functions,
+    ):
+        # solve n_lower : n_i = N * g_i / U * e ^ (-E_i/kT)
+        # get f_lu : loggf -> use g = 2j+1
+        # emission_correction = (1-e^(-h*nu / kT))
+        # alphas = ALPHA_COEFFICIENT * n_lower * f_lu * emission_correction
+
+        ###TODO: handle other broadening parameters
+        points = len(t_electrons)
+
+        linelist = atomic_data.linelist_atoms.rename(
+            columns={"ion_charge": "ion_number"}
+        )[
+            [
+                "atomic_number",
+                "ion_number",
+                "wavelength",
+                "log_gf",
+                "e_low",
+                "rad",
+                "stark",
+                "waals",
+            ]
+        ]
+
+        linelist["e_up"] = (
+            (
+                linelist.e_low.values * u.eV
+                + (const.h * const.c / (linelist.wavelength.values * u.AA))
+            )
+            .to(u.eV)
+            .value
+        )
+
+        exponent_by_point = np.exp(
+            np.outer(
+                -linelist.e_low.values * u.eV, 1 / (t_electrons * u.K * const.k_B)
+            ).to(1)
+        )
+
+        molecule_densities_div_partition_function = (
+            molecule_number_densities.copy().div(molecule_partition_functions)
+        )
+        molecule_densities_div_partition_function.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_density_div_partition_function = linelist.merge(
+            molecule_densities_div_partition_function,
+            how="left",
+            on=["molecule"],
+        )
+
+        prefactor = (
+            exponent_by_point
+            * linelist_with_density_div_partition_function[np.arange(points)]
+        ).values.T  # arange mask of the dataframe returns the set of densities of the appropriate ion for the line at each point
+
+        line_nus = (linelist.wavelength.values * u.AA).to(
+            u.Hz, equivalencies=u.spectral()
+        )
+
+        emission_correction = 1 - np.exp(
+            (
+                -const.h
+                / const.k_B
+                * np.outer(
+                    line_nus,
+                    1 / (t_electrons * u.K),
+                )
+            ).to(1)
+        )
+
+        alphas = pd.DataFrame(
+            (
+                ALPHA_COEFFICIENT
+                * prefactor
+                * 10**linelist.log_gf.values
+                * emission_correction.T
+            ).T
+        )
+
+        if np.any(np.isnan(alphas)) or np.any(np.isinf(np.abs(alphas))):
+            raise ValueError(
+                "Some alpha_line from vald are nan, inf, -inf " " Something went wrong!"
+            )
+
+        alphas["nu"] = line_nus.value
+        linelist["nu"] = line_nus.value
+
+        # Linelist preparation below is taken from opacities_solvers/base/calc_alpha_line_at_nu
+        linelist["level_energy_lower"] = ((linelist["e_low"].values * u.eV).cgs).value
+        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"]) / (
+            4 * np.pi
         )  # see 1995A&AS..112..525P for appropriate units - may be off by a factor of 4pi
 
         return alphas, linelist