change molecule calculation to use final densities

jvshields · Jan 13, 2025 · fc41227 · fc41227
1 parent 8d8984c
commit fc41227
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 13 deletions.
diff --git a/stardis/plasma/molecules.py b/stardis/plasma/molecules.py
@@ -3,6 +3,7 @@
 import logging
 
 from astropy import constants as const, units as u
+from scipy.interpolate import CubicSpline
 from tardis.util.base import element_symbol2atomic_number
 
 from tardis.plasma.properties.base import ProcessingPlasmaProperty
@@ -81,26 +82,49 @@ def calculate(self, ion_number_density, t_electrons, atomic_data):
                 molecule_row.Ion2, molecule_row.Ion2_charge
             ]
 
-            # Barklem and Collet 2016 equilibrium constants are pressure constants and are in SI units
-            pressure_equilibirium_const_at_depth_point = np.interp(
-                t_electrons,
+            pressure_equil_spline = CubicSpline(
                 equilibrium_const_temps,
                 atomic_data.molecule_data.equilibrium_constants.loc[molecule].values,
+                extrapolate=True,
+            )
+
+            pressure_equilibirium_const_at_depth_point = pressure_equil_spline(
+                t_electrons
             )
+
             # Convert from pressure constants to number density constants using ideal gas law
-            equilibirium_const_at_depth_point = (
+            # k is the equilibrium concentration constant, pressure is Pa
+            k = (
                 (
-                    10 ** (pressure_equilibirium_const_at_depth_point)
-                    * (u.N / u.m**2)
+                    (10**pressure_equilibirium_const_at_depth_point)
+                    * (u.Pa)
                     / (const.k_B * t_electrons * u.K)
+                ).to(u.cm**-3)
+            ).value
+
+            # Different formulae for homonuclear heteronuclear diatomic molecules
+            if (molecule_row.Ion1 == molecule_row.Ion2) and (
+                molecule_row.Ion1_charge == molecule_row.Ion2_charge
+            ):
+                molecule_number_density = (1 / 8) * (
+                    (-((k * (k + 8 * ion1_number_density)) ** 0.5))
+                    + k
+                    + 4 * ion1_number_density
                 )
-                .to(u.cm**-3)
-                .value
-            )
 
-            molecule_number_density = (ion1_number_density * ion2_number_density) / (
-                equilibirium_const_at_depth_point
-            )
+            else:
+                molecule_number_density = 0.5 * (
+                    -np.sqrt(
+                        k**2
+                        + 2 * k * (ion1_number_density + ion2_number_density)
+                        + (ion1_number_density - ion2_number_density) ** 2
+                    )
+                    + k
+                    + ion1_number_density
+                    + ion2_number_density
+                )
+
+            np.maximum(molecule_number_density, 0, out=molecule_number_density)
 
             number_densities_arr[
                 atomic_data.molecule_data.equilibrium_constants.index.get_loc(molecule)

diff --git a/stardis/radiation_field/base.py b/stardis/radiation_field/base.py
@@ -68,7 +68,7 @@ def __init__(
             )
 
         # This was our original theta sampling method
-        # dtheta = (np.pi / 2) / num_of_thetas  # Korg uses Gauss-Legendre quadrature here
+        # dtheta = (np.pi / 2) / num_of_thetas
         # start_theta = dtheta / 2
         # end_theta = (np.pi / 2) - (dtheta / 2)
         # self.thetas = np.linspace(start_theta, end_theta, num_of_thetas)