shortlist handling bugfix

tardis-sn · Aug 29, 2024 · b34b77e · b34b77e
1 parent 35cd09e
commit b34b77e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/stardis/plasma/base.py b/stardis/plasma/base.py
@@ -23,6 +23,7 @@
     MoleculeIonNumberDensity,
     AlphaLineValdMolecule,
     MoleculePartitionFunction,
+    AlphaLineShortlistValdMolecule,
 )
 
 
@@ -548,7 +549,11 @@ def create_stellar_plasma(
     if config.opacity.line.include_molecules:
         plasma_modules.append(MoleculeIonNumberDensity)
         plasma_modules.append(MoleculePartitionFunction)
-        plasma_modules.append(AlphaLineValdMolecule)
+        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,

diff --git a/stardis/plasma/molecules.py b/stardis/plasma/molecules.py
@@ -167,7 +167,7 @@ class AlphaLineValdMolecule(ProcessingPlasmaProperty):
     """
 
     outputs = ("molecule_alpha_line_from_linelist", "molecule_lines_from_linelist")
-    latex_name = (r"\alpha_{\textrm{line, vald}}",)
+    latex_name = (r"\alpha_{\textrm{moleculeline, 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)",
@@ -201,7 +201,7 @@ def calculate(
                 "stark",
                 "waals",
             ]
-        ]
+        ].copy()
 
         # Calculate degeneracies
         linelist["g_lo"] = linelist.j_lo * 2 + 1
@@ -295,8 +295,8 @@ class AlphaLineShortlistValdMolecule(ProcessingPlasmaProperty):
             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}}",)
+    outputs = ("molecule_alpha_line_from_linelist", "molecule_lines_from_linelist")
+    latex_name = (r"\alpha_{\textrm{moleculeline, 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)",
@@ -317,20 +317,17 @@ def calculate(
         ###TODO: handle other broadening parameters
         points = len(t_electrons)
 
-        linelist = atomic_data.linelist_atoms.rename(
-            columns={"ion_charge": "ion_number"}
-        )[
+        linelist = atomic_data.linelist_molecules[
             [
-                "atomic_number",
-                "ion_number",
+                "molecule",
                 "wavelength",
                 "log_gf",
                 "e_low",
                 "rad",
                 "stark",
                 "waals",
             ]
-        ]
+        ].copy()
 
         linelist["e_up"] = (
             (