Export raw Chianti collision and A_ul data instead of pre-processing

carsus/io/output/base.py

@@ -196,8 +196,8 @@ def to_hdf(self, fname):
             "/molecules/ionization_energies": (self.barklem_2016_data, "ionization_energies"),
             "/molecules/dissociation_energies": (self.barklem_2016_data, "dissociation_energies"),
             "/molecules/partition_functions": (self.barklem_2016_data, "partition_functions"),
-            "/collisions_data": (self.collisions_preparer, "collisions_prepared"),
-            "/collisions_metadata": (self.collisions_preparer, "collisions_metadata"),
+            "/collisions_data": (self, "collisions_prepared"),
+            "/collisions_metadata": (self, "collisions_metadata"),
             "/photoionization_data": (self.cross_sections_preparer, "cross_sections_prepared"),
         }
 

carsus/io/output/collisions.py

@@ -10,6 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 class CollisionsPreparer:
     def __init__(self, reader):
         collisions = reader.collisions.copy()
@@ -24,7 +25,7 @@ def __init__(self, reader):
 
         self.collisions = collisions
         self.collisions_metadata = reader.collisional_metadata
-    
+
     def prepare_collisions(self):
         """
         Prepare the DataFrame with electron collisions for TARDIS.
@@ -43,13 +44,22 @@ def prepare_collisions(self):
         ]
 
         if "chianti" in self.collisions_metadata.dataset:
-            collisions_columns = (
-                collisions_index
-                + ["g_ratio", "delta_e"]
-                + sorted(
-                    [col for col in self.collisions.columns if re.match(r"^t\d+$", col)]
-                )
-            )
+            collisions_columns = collisions_index + [
+                "e_col_id",
+                "lower_level_id",
+                "upper_level_id",
+                "btemp",
+                "bscups",
+                "ttype",
+                "cups",
+                "gf",
+                "g_l",
+                "energy_lower",
+                "g_u",
+                "energy_upper",
+                "delta_e",
+                "g_ratio",
+            ]
 
         elif "cmfgen" in self.collisions_metadata.dataset:
             collisions_columns = collisions_index + list(self.collisions.columns)
@@ -61,18 +71,18 @@ def prepare_collisions(self):
             self.collisions.reset_index().loc[:, collisions_columns].copy()
         )
         self.collisions_prepared = collisions_prepared.set_index(collisions_index)
-    
+
 
 class ChiantiCollisionsPreparer(CollisionsPreparer):
     def __init__(
-            self, 
-            chianti_reader, 
-            levels, 
-            levels_all, 
-            lines_all, 
-            chianti_ions, 
-            collisions_param = {"temperatures": np.arange(2000, 50000, 2000)}
-            ):
+        self,
+        chianti_reader,
+        levels,
+        levels_all,
+        lines_all,
+        chianti_ions,
+        collisions_param={"temperatures": np.arange(2000, 50000, 2000)},
+    ):
         self.chianti_reader = chianti_reader
         self.levels = levels
         self.levels_all = levels_all
@@ -87,7 +97,7 @@ def __init__(
                 "info": None,
             }
         )
-        
+
     def create_chianti_collisions(self, temperatures=np.arange(2000, 50000, 2000)):
         """
         Generates the definitive `collisions` DataFrame by adding new columns
@@ -169,9 +179,6 @@ def create_chianti_collisions(self, temperatures=np.arange(2000, 50000, 2000)):
         # Calculate g_ratio
         collisions["g_ratio"] = collisions["g_l"] / collisions["g_u"]
 
-        # Derive columns for collisional strengths
-        c_ul_temperature_cols = ["t{:06d}".format(t) for t in temperatures]
-
         collisions = collisions.rename(
             columns={"temperatures": "btemp", "collision_strengths": "bscups"}
         )
@@ -200,62 +207,6 @@ def create_chianti_collisions(self, temperatures=np.arange(2000, 50000, 2000)):
             ]
         ]
 
-        collisional_ul_factors = collisions.apply(
-            calculate_collisional_strength,
-            axis=1,
-            args=(temperatures, kb_ev, c_ul_temperature_cols),
-        )
-
-        collisions = pd.concat([collisions, collisional_ul_factors], axis=1)
         collisions = collisions.set_index("e_col_id")
 
         return collisions
-
-def calculate_collisional_strength(
-        row, temperatures, kb_ev, c_ul_temperature_cols
-    ):
-    """
-    Function to calculation upsilon from Burgess & Tully 1992 (TType 1 - 4; Eq. 23 - 38).
-
-    """
-
-    c = row["cups"]
-    x_knots = np.linspace(0, 1, len(row["btemp"]))
-    y_knots = row["bscups"]
-    delta_e = row["delta_e"]
-    g_u = row["g_u"]
-
-    ttype = row["ttype"]
-    if ttype > 5:
-        ttype -= 5
-
-    kt = kb_ev * temperatures
-
-    spline_tck = interpolate.splrep(x_knots, y_knots)
-
-    if ttype == 1:
-        x = 1 - np.log(c) / np.log(kt / delta_e + c)
-        y_func = interpolate.splev(x, spline_tck)
-        upsilon = y_func * np.log(kt / delta_e + np.exp(1))
-
-    elif ttype == 2:
-        x = (kt / delta_e) / (kt / delta_e + c)
-        y_func = interpolate.splev(x, spline_tck)
-        upsilon = y_func
-
-    elif ttype == 3:
-        x = (kt / delta_e) / (kt / delta_e + c)
-        y_func = interpolate.splev(x, spline_tck)
-        upsilon = y_func / (kt / delta_e + 1)
-
-    elif ttype == 4:
-        x = 1 - np.log(c) / np.log(kt / delta_e + c)
-        y_func = interpolate.splev(x, spline_tck)
-        upsilon = y_func * np.log(kt / delta_e + c)
-
-    elif ttype == 5:
-        raise ValueError("Not sure what to do with ttype=5")
-
-    #### 1992A&A...254..436B Equation 20 & 22 #####
-    collisional_ul_factor = 8.63e-6 * upsilon / (g_u * temperatures**0.5)
-    return pd.Series(data=collisional_ul_factor, index=c_ul_temperature_cols)