Update readkuruczdata.py

Author: lukeshingles

artisatomic/readkuruczdata.py

@@ -71,30 +71,11 @@ def __init__(self, ions: list[tuple[int, int]], fname: str) -> None:
         self.fname = fname
 
         self.ions = ions
-
-        self._gfall_raw = None
-        self._gfall = None
-        self._levels = None
-        self._lines = None
         self.unique_level_identifier = ["energy", "j"]
 
-    @property
-    def gfall(self):
-        if self._gfall is None:
-            self._gfall = self.parse_gfall()
-        return self._gfall
-
-    @property
-    def levels(self):
-        if self._levels is None:
-            self._levels = self.extract_levels()
-        return self._levels
-
-    @property
-    def lines(self):
-        if self._lines is None:
-            self._lines = self.extract_lines()
-        return self._lines
+        self.gfall = self.parse_gfall()
+        self.levels = self.extract_levels()
+        self.lines = self.extract_lines()
 
     def read_gfall_raw(self):
         print(f"Parsing GFALL from: {self.fname}")
@@ -192,36 +173,15 @@ def parse_gfall(self):
 
         return gfall
 
-    def extract_levels(self, gfall=None, selected_columns=None):
-        """
-        Extract levels from `gfall`. We first generate a concatenated DataFrame
-        of all lower and upper levels. Then we drop the duplicate levels
-
-        Parameters
-        ----------
-        gfall: pandas.DataFrame
-        selected_columns: list
-            list of which columns to select (optional - default=None which selects
-            a default set of columns)
-
-        Returns
-        -------
-            pandas.DataFrame
-                a level DataFrame
-        """
-
-        if gfall is None:
-            gfall = self.gfall
-
-        if selected_columns is None:
-            selected_columns = [
-                "atomic_number",
-                "ion_charge",
-                "energy",
-                "j",
-                "label",
-                "theoretical",
-            ]
+    def extract_levels(self):
+        selected_columns = [
+            "atomic_number",
+            "ion_charge",
+            "energy",
+            "j",
+            "label",
+            "theoretical",
+        ]
 
         column_renames = {
             "energy_{0}": "energy",
@@ -230,9 +190,13 @@ def extract_levels(self, gfall=None, selected_columns=None):
             "energy_{0}_predicted": "theoretical",
         }
 
-        e_lower_levels = gfall.rename(columns={key.format("lower"): value for key, value in column_renames.items()})
+        e_lower_levels = self.gfall.rename(
+            columns={key.format("lower"): value for key, value in column_renames.items()}
+        )
 
-        e_upper_levels = gfall.rename(columns={key.format("upper"): value for key, value in column_renames.items()})
+        e_upper_levels = self.gfall.rename(
+            columns={key.format("upper"): value for key, value in column_renames.items()}
+        )
 
         levels = pd.concat([e_lower_levels[selected_columns], e_upper_levels[selected_columns]])
         unique_level_id = ["atomic_number", "ion_charge", *self.unique_level_identifier]
@@ -263,7 +227,7 @@ def extract_levels(self, gfall=None, selected_columns=None):
 
         return levels
 
-    def extract_lines(self, gfall=None, levels=None, selected_columns=None):
+    def extract_lines(self):
         """
         Extract lines from `gfall`
 
@@ -279,17 +243,14 @@ def extract_lines(self, gfall=None, levels=None, selected_columns=None):
             pandas.DataFrame
                 a level DataFrame
         """
-        if gfall is None:
-            gfall = self.gfall
+        gfall = self.gfall
 
-        if levels is None:
-            levels = self.levels
+        levels = self.levels
 
-        if selected_columns is None:
-            selected_columns = ["atomic_number", "ion_charge"]
-            selected_columns += [item + "_lower" for item in self.unique_level_identifier]
-            selected_columns += [item + "_upper" for item in self.unique_level_identifier]
-            selected_columns += ["wavelength_nm", "loggf"]
+        selected_columns = ["atomic_number", "ion_charge"]
+        selected_columns += [item + "_lower" for item in self.unique_level_identifier]
+        selected_columns += [item + "_upper" for item in self.unique_level_identifier]
+        selected_columns += ["wavelength_nm", "loggf"]
 
         print(f"Extracting line data: {', '.join(selected_columns)}.")
         unique_level_id = ["atomic_number", "ion_charge", *self.unique_level_identifier]
@@ -408,13 +369,13 @@ def read_levels_and_transitions(atomic_number: int, ion_stage: int, flog) -> tup
         fname=str(path_gfall),
     )
 
-    dflevels = gfall_reader.extract_levels().loc[atomic_number, ion_charge]
+    dflevels = gfall_reader.levels.loc[atomic_number, ion_charge]
 
     energy_levels = read_levels_data(dflevels)
 
     artisatomic.log_and_print(flog, f"Read {len(energy_levels[1:]):d} levels")
 
-    dflines = gfall_reader.extract_lines().loc[atomic_number, ion_charge]
+    dflines = gfall_reader.lines.loc[atomic_number, ion_charge]
 
     # wavelengths are in nanometers, so multiply by 10 to get Angstroms
     dflines = dflines.assign(