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

carsus/__init__.py

@@ -28,4 +28,4 @@
 logging.getLogger('py.warnings').addHandler(console_handler)
 
 # Set atomic file format version
-FORMAT_VERSION = "1.0"
+FORMAT_VERSION = "2.0"

carsus/io/chianti_/chianti_.py

@@ -192,8 +192,9 @@ def read_levels(self):
         levels = pd.DataFrame(levels_dict)
 
         # Replace empty labels with NaN
-        levels.loc[:, "label"] = levels["label"].replace(
-            r'\s+', np.nan, regex=True)
+        with pd.option_context('future.no_silent_downcasting', True):
+            levels.loc[:, "label"] = levels["label"].replace(
+                r'\s+', np.nan, regex=True).infer_objects(copy=False)
 
         # Extract configuration and term from the "pretty" column
         levels[["term", "configuration"]] = levels["pretty"].str.rsplit(
@@ -355,7 +356,8 @@ def _get_levels_lines(self, get_collisions=False):
         lines = lines.reset_index()
         lines = lines.rename(columns={'lower_level_index': 'level_index_lower',
                                       'upper_level_index': 'level_index_upper',
-                                      'gf_value': 'gf'})
+                                      'gf_value': 'gf',
+                                      'a_value': 'A_ul'})
 
         # Kurucz levels starts from zero, Chianti from 1.
         lines['level_index_lower'] = lines['level_index_lower'] - 1
@@ -368,7 +370,7 @@ def _get_levels_lines(self, get_collisions=False):
         lines['j_upper'] = None
         lines['j_lower'] = None
         lines = lines[['energy_upper', 'j_upper', 'energy_lower', 'j_lower',
-                       'wavelength', 'gf']]
+                       'wavelength', 'gf', 'A_ul']]
 
         lines['wavelength'] = u.Quantity(lines['wavelength'], u.AA).to('nm').value
 

carsus/io/nist/weightscomp.py

@@ -125,8 +125,8 @@ def load(self, input_data):
 
     def _prepare_atomic_weights(self, atomic):
         grouped = atomic.groupby([AW_TYPE_COL])
-        interval_gr = grouped.get_group(INTERVAL).copy()
-        stable_mass_num_gr = grouped.get_group(STABLE_MASS_NUM).copy()
+        interval_gr = grouped.get_group((INTERVAL,)).copy()
+        stable_mass_num_gr = grouped.get_group((STABLE_MASS_NUM,)).copy()
 
         def atomic_weight_interval_to_nom_val_and_std(row):
             nom_val, std_dev = to_nom_val_and_std_dev(

carsus/io/output/base.py

@@ -49,7 +49,6 @@ def __init__(
             "levels_metastable_loggf_threshold": -3,
             "lines_loggf_threshold": -3,
         },
-        collisions_param={"temperatures": np.arange(2000, 50000, 2000)},
     ):
         self.atomic_weights = atomic_weights
 
@@ -61,7 +60,6 @@ def __init__(
         self.vald_reader = vald_reader
         self.barklem_2016_data = barklem_2016_data
         self.levels_lines_param = levels_lines_param
-        self.collisions_param = collisions_param
 
         self.ionization_energies_preparer = IonizationEnergiesPreparer(self.cmfgen_reader, ionization_energies)
 
@@ -86,7 +84,7 @@ def __init__(
         if cmfgen_reader is not None and hasattr(cmfgen_reader, "collisions"):
             self.collisions_preparer = CollisionsPreparer(self.cmfgen_reader)
         elif hasattr(chianti_reader, "collisions"):
-            self.collisions_preparer = ChiantiCollisionsPreparer(self.chianti_reader, self.levels, self.levels_all, self.lines_all, self.levels_lines_preparer.chianti_ions, self.collisions_param)
+            self.collisions_preparer = ChiantiCollisionsPreparer(self.chianti_reader, self.levels, self.levels_all, self.lines_all, self.levels_lines_preparer.chianti_ions)
         else:
             logger.warning("No source of collisions was selected.")
             self.collisions_preparer = None
@@ -196,8 +194,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

@@ -1,15 +1,13 @@
 import logging
-import re
 
 import astropy.constants as const
-import numpy as np
 import pandas as pd
-from scipy import interpolate
 
 from carsus.io.util import get_lvl_index2id, exclude_artificial_levels
 
 logger = logging.getLogger(__name__)
 
+
 class CollisionsPreparer:
     def __init__(self, reader):
         collisions = reader.collisions.copy()
@@ -24,7 +22,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 +41,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,34 +68,32 @@ 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,
+    ):
         self.chianti_reader = chianti_reader
         self.levels = levels
         self.levels_all = levels_all
         self.lines_all = lines_all
         self.chianti_ions = chianti_ions
 
-        self.collisions = self.create_chianti_collisions(**collisions_param)
+        self.collisions = self.create_chianti_collisions()
         self.collisions_metadata = pd.Series(
             {
-                "temperatures": collisions_param["temperatures"],
                 "dataset": ["chianti"],
                 "info": None,
             }
         )
-        
-    def create_chianti_collisions(self, temperatures=np.arange(2000, 50000, 2000)):
+
+    def create_chianti_collisions(self):
         """
         Generates the definitive `collisions` DataFrame by adding new columns
         and making some calculations.
@@ -169,9 +174,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 +202,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)

carsus/io/output/levels_lines.py

@@ -304,7 +304,7 @@ def all_lines_data(self):
         )
 
         lines = lines[
-            ["lower_level_id", "upper_level_id", "wavelength", "gf", "loggf", "ds_id"]
+            ["lower_level_id", "upper_level_id", "wavelength", "gf", "loggf", "A_ul", "ds_id"]
         ]
 
         return lines
@@ -424,8 +424,12 @@ def create_levels_lines(
             pd.DataFrame(index=levels.index), on="lower_level_id", how="inner"
         ).join(pd.DataFrame(index=levels.index), on="upper_level_id", how="inner")
 
-        # Culling lines with low gf values
-        lines = lines.loc[lines["loggf"] > lines_loggf_threshold]
+        # Culling lines with low gf values if needed
+        if lines_loggf_threshold > -99:
+            lines = lines.loc[lines["loggf"] < lines_loggf_threshold]
+
+        # get a mask of the lines with loggf above the threshold
+        high_gf_mask = lines["loggf"] > lines_loggf_threshold
 
         # Do not clean levels that don't exist in lines
 
@@ -480,7 +484,7 @@ def create_levels_lines(
         lines["nu"] = u.Quantity(lines["wavelength"], "AA").to("Hz", u.spectral()).value
 
         # Create Einstein coefficients
-        create_einstein_coeff(lines)
+        create_einstein_coeff(lines, high_gf_mask)
 
         # Reset indexes because `level_id` cannot be an index once we
         # add artificial levels for fully ionized ions that don't have ids (-1)
@@ -495,7 +499,7 @@ def create_levels_lines(
         self.levels = levels
         self.lines = lines
 
-def create_einstein_coeff(lines):
+def create_einstein_coeff(lines, high_gf_mask):
     """
     Create Einstein coefficients columns for the `lines` DataFrame.
 
@@ -517,7 +521,7 @@ def create_einstein_coeff(lines):
         einstein_coeff * lines["f_ul"] / (const.h.cgs.value * lines["nu"])
     )
 
-    lines["A_ul"] = (
+    lines.loc[high_gf_mask, "A_ul"] = (
         2
         * einstein_coeff
         * lines["nu"] ** 2