Merge pull request #311 from gagan-aryan/Qgas-improvements

refactor: minor Qgas loading improvements
radis · Jul 7, 2021 · 3dabb6e · 3dabb6e
2 parents ea7b606 + dca1b16
commit 3dabb6e
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/radis/lbl/base.py b/radis/lbl/base.py
@@ -1693,6 +1693,9 @@ def calc_weighted_trans_moment(self):
             id = df.attrs["id"]
 
         iso_set = self._get_isotope_list(self.input.molecule)  # df1.iso.unique()
+
+        Qref_dict = {}
+
         if (("id" in df and len(id_set) == 1) or ("id" not in df and id)) and len(
             iso_set
         ) == 1:
@@ -1711,6 +1714,8 @@ def calc_weighted_trans_moment(self):
             )  # stored as attribute, not column
             assert "Qref" not in df.columns
 
+            Qref = df.Qref
+
         else:
 
             # normal method
@@ -1725,7 +1730,7 @@ def calc_weighted_trans_moment(self):
                 parsum = self.get_partition_function_calculator(
                     molecule, iso, state
                 )  # partition function
-                df.at[idx, "Qref"] = parsum.at(Tref, update_populations=False)
+                Qref_dict[iso] = parsum.at(Tref, update_populations=False)
                 # ... note: do not update the populations here, so populations in the
                 # ... energy level list correspond to the one calculated for T and not Tref
 
@@ -1734,6 +1739,8 @@ def calc_weighted_trans_moment(self):
                         assert (df.loc[idx, "id"] == id).all()
                     assert (df.loc[idx, "iso"] == iso).all()
 
+            Qref = df["iso"].map(Qref_dict)
+
         # Get moment
         gl = df.gl
         El = df.El
@@ -1742,7 +1749,6 @@ def calc_weighted_trans_moment(self):
         h = h_CGS  # erg.s
         c = c_CGS
         S = df.int  # reference linestrength
-        Qref = df.Qref
 
         weighted_trans_moment_sq = (
             (3 * h * c / 8 / pi ** 3)
@@ -2000,8 +2006,9 @@ def _calc_Q(molecule, iso, state):
         if "iso" in df1:
             line_strength = df1.int * (df1["iso"].map(Qref_Qgas_ratio))
         else:
-            Qref_Qgas_ratio = list(Qref_Qgas_ratio.values())
-            line_strength = df1.int * Qref_Qgas_ratio[0]
+            assert len(Qref_Qgas_ratio) == 1
+            Qref_Qgas_ratio = list(Qref_Qgas_ratio.values())[0]
+            line_strength = df1.int * Qref_Qgas_ratio
         # ratio of Boltzman populations
         line_strength *= exp(-hc_k * df1.El * (1 / Tgas - 1 / Tref))
         # effect of stimulated emission
@@ -3010,8 +3017,9 @@ def _calc_linestrength_noneq(self):
                 df.gl * exp(-hc_k * df.El / Tref) / df["iso"].map(Qref_dict)
             )
         else:
-            Qref_dict = list(Qref_dict.values())
-            line_strength /= df.gl * exp(-hc_k * df.El / Tref) / Qref_dict[0]
+            assert len(Qref_dict) == 1
+            Qref = list(Qref_dict.values())[0]
+            line_strength /= df.gl * exp(-hc_k * df.El / Tref) / Qref
         line_strength *= nl
 
         # ... correct effect of stimulated emission