use mapping for ions instead of holding in densities df

tardis-sn · Aug 29, 2024 · 6d6b06b · 6d6b06b
1 parent b34b77e
commit 6d6b06b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 7 deletions.
diff --git a/stardis/plasma/molecules.py b/stardis/plasma/molecules.py
@@ -17,7 +17,7 @@ class MoleculeIonNumberDensity(ProcessingPlasmaProperty):
     # Need to think about negative ions - ignoring for now
     # applicable for equilibrium constants given by Barklem and Collet 2016, which are given in SI units
 
-    outputs = ("molecule_number_density",)
+    outputs = ("molecule_number_density", "molecule_ion_map")
 
     def calculate(self, ion_number_density, t_electrons, atomic_data):
         # Preprocessing - split ions into symbol, charge, and number
@@ -117,10 +117,13 @@ def calculate(self, ion_number_density, t_electrons, atomic_data):
             columns=ion_number_density.columns,
         )
         # Keep track of the individual ions - useful to calculate molecular masses later for doppler broadening
-        molecule_densities_df["ion1"] = molecules_df["Ion1"]
-        molecule_densities_df["ion2"] = molecules_df["Ion2"]
+        molecule_ion_map = pd.DataFrame(
+            molecules_df[["Ion1", "Ion2"]],
+        )
+        # molecule_densities_df["ion1"] = molecules_df["Ion1"]
+        # molecule_densities_df["ion2"] = molecules_df["Ion2"]
 
-        return molecule_densities_df
+        return molecule_densities_df, molecule_ion_map
 
 
 class MoleculePartitionFunction(ProcessingPlasmaProperty):

diff --git a/stardis/radiation_field/opacities/opacities_solvers/broadening.py b/stardis/radiation_field/opacities/opacities_solvers/broadening.py
@@ -717,10 +717,10 @@ def calculate_molecule_broadening(
             (len(lines), len(stellar_model.geometry.no_of_depth_points)), dtype=float
         )
 
-    ions = stellar_plasma.molecule_number_density[["ion1", "ion2"]].loc[lines.molecule]
+    ions = stellar_plasma.molecule_ion_map.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
+    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(