spectrum checkpoint - probably incorrect

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/plasma/base.py

@@ -541,7 +541,7 @@ def create_stellar_plasma(
         temperature=stellar_model.temperatures,
         dilution_factor=np.ones_like(stellar_model.temperatures),
     )
-    if True:
+    if config.opacity.line.include_molecules:
         plasma_modules.append(MoleculeNumberDensities)
         plasma_modules.append(AlphaLineValdMolecules)
 

stardis/plasma/molecules.py

@@ -72,10 +72,14 @@ def calculate(self, ion_number_density, t_electrons, atomic_data):
                 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_electrons * u.K)
-            ).cgs.value
+                (
+                    10 ** (pressure_equilibirium_const_at_depth_point)
+                    * (u.N / u.m**2)
+                    / (const.k_B * t_electrons * u.K)
+                )
+                .to(u.cm**-3)
+                .value
+            )
 
             molecule_number_density = (
                 ion1_number_density * ion2_number_density

stardis/radiation_field/opacities/opacities_solvers/base.py

@@ -7,6 +7,7 @@
 
 from stardis.radiation_field.opacities.opacities_solvers.broadening import (
     calculate_broadening,
+    calculate_molecule_broadening,
 )
 from stardis.radiation_field.opacities.opacities_solvers.voigt import voigt_profile
 from stardis.radiation_field.opacities.opacities_solvers.util import (
@@ -466,6 +467,92 @@ def calc_alpha_line_at_nu(
     return alpha_line_at_nu, gammas, doppler_widths
 
 
+def calc_molecular_alpha_line_at_nu(
+    stellar_plasma,
+    stellar_model,
+    tracing_nus,
+    line_opacity_config,
+    n_threads=1,
+):
+    if line_opacity_config.disable:
+        return 0, 0, 0
+
+    line_range = line_opacity_config.broadening_range
+
+    lines = stellar_plasma.molecule_lines_from_linelist
+    lines_sorted = lines.sort_values("nu")
+    lines_sorted_in_range = lines_sorted[
+        lines_sorted.nu.between(tracing_nus.min(), tracing_nus.max())
+    ]
+    line_nus = lines_sorted_in_range.nu.to_numpy()
+
+    alphas_and_nu = stellar_plasma.molecule_alpha_line_from_linelist
+
+    alphas_array = (
+        alphas_and_nu[alphas_and_nu.nu.between(tracing_nus.min(), tracing_nus.max())]
+        .drop(labels="nu", axis=1)
+        .to_numpy()
+    )
+
+    gammas, doppler_widths = calculate_molecule_broadening(
+        lines_sorted_in_range,
+        stellar_model,
+        stellar_plasma,
+        line_opacity_config.broadening,
+    )
+
+    delta_nus = tracing_nus.value - line_nus[:, np.newaxis]
+
+    line_range_value = None
+
+    # If there is a broadening range, first make sure the range is in frequency units, and then iterate through each frequency to calculate the contribution of each line within the broadening range.
+    if (
+        line_range is not None
+    ):  # This if statement block appropriately handles if the broadening range is in frequency or wavelength units.
+        h_lines_indices = np.full(
+            len(lines_sorted_in_range), False
+        )  # This is wonky but necessary for the calc_alan_entries function
+        if line_range.unit.physical_type == "length":
+            lambdas = tracing_nus.to(u.AA, equivalencies=u.spectral())
+            lambdas_plus_broadening_range = lambdas + line_range.to(u.AA)
+            nus_plus_broadening_range = lambdas_plus_broadening_range.to(
+                u.Hz, equivalencies=u.spectral()
+            )
+            line_range_value = (tracing_nus - nus_plus_broadening_range).value
+        elif line_range.unit.physical_type == "frequency":
+            line_range_value = line_range.to(u.Hz).value
+        else:
+            raise ValueError(
+                "Broadening range must be in units of length or frequency."
+            )
+
+    if n_threads == 1:  # Single threaded
+        alpha_line_at_nu = calc_alan_entries(
+            stellar_model.no_of_depth_points,
+            tracing_nus.value,
+            delta_nus,
+            doppler_widths,
+            gammas,
+            alphas_array,
+            line_range_value,
+            h_lines_indices,
+        )
+
+    else:  # Parallel threaded
+        alpha_line_at_nu = calc_alan_entries_parallel(
+            stellar_model.no_of_depth_points,
+            tracing_nus.value,
+            delta_nus,
+            doppler_widths,
+            gammas,
+            alphas_array,
+            line_range_value,
+            h_lines_indices,
+        )
+
+    return alpha_line_at_nu, gammas, doppler_widths
+
+
 @numba.njit(parallel=True)
 def calc_alan_entries_parallel(
     no_of_depth_points,
@@ -728,13 +815,35 @@ def calc_alphas(
         opacity_config.line,
         n_threads,
     )
-    stellar_radiation_field.opacities.opacities_dict[
-        "alpha_line_at_nu"
-    ] = alpha_line_at_nu
+    stellar_radiation_field.opacities.opacities_dict["alpha_line_at_nu"] = (
+        alpha_line_at_nu
+    )
     stellar_radiation_field.opacities.opacities_dict["alpha_line_at_nu_gammas"] = gammas
     stellar_radiation_field.opacities.opacities_dict[
         "alpha_line_at_nu_doppler_widths"
     ] = doppler_widths
+
+    if opacity_config.line.include_molecules:
+        molecule_alpha_line_at_nu, molecule_gammas, molecule_doppler_widths = (
+            calc_molecular_alpha_line_at_nu(
+                stellar_plasma,
+                stellar_model,
+                stellar_radiation_field.frequencies,
+                opacity_config.line,
+                n_threads,
+            )
+        )
+
+        stellar_radiation_field.opacities.opacities_dict[
+            "molecule_alpha_line_at_nu"
+        ] = molecule_alpha_line_at_nu
+        stellar_radiation_field.opacities.opacities_dict[
+            "molecule_alpha_line_at_nu_gammas"
+        ] = molecule_gammas
+        stellar_radiation_field.opacities.opacities_dict[
+            "molecule_alpha_line_at_nu_doppler_widths"
+        ] = molecule_doppler_widths
+
     alphas = stellar_radiation_field.opacities.calc_total_alphas()
 
     return alphas

stardis/radiation_field/opacities/opacities_solvers/broadening.py

@@ -704,6 +704,36 @@ def calculate_broadening(
     return gammas, doppler_widths
 
 
+def calculate_molecule_broadening(
+    lines,
+    stellar_model,
+    stellar_plasma,
+    broadening_line_opacity_config,
+):
+    if "radiation" in broadening_line_opacity_config:
+        gammas = lines.A_ul.values[:, np.newaxis]
+    else:
+        gammas = np.zeros(
+            (len(lines), len(stellar_model.geometry.no_of_depth_points)), dtype=float
+        )
+
+    ions = stellar_plasma.molecule_number_densities[["ion1", "ion2"]].loc[
+        lines.molecule
+    ]
+
+    ion1_masses = stellar_model.composition.nuclide_masses.loc[ions.ion1].values
+    ion2_masses = stellar_model.composition.nuclide_masses.loc[ions.ion2].values
+
+    molecule_masses = (ion1_masses + ion2_masses)[:, np.newaxis]
+    doppler_widths = calc_doppler_width(
+        lines.nu.values[:, np.newaxis],
+        stellar_model.temperatures.value,
+        molecule_masses,
+    )
+
+    return gammas, doppler_widths
+
+
 def rotation_broadening(
     velocity_per_pix, wavelength, flux, v_rot=0 * u.km / u.s, limb_darkening=0.6
 ):