Filtering Levels Parameter within Plasma Calculations (#1899)

* Moved Filtering from Atom Data Module to Plasma Module

* Reformatted into black-friendly formatting

* Adding levels filter back to atom data base for time being

* Fixed Error when Performing Formal Integral

* Eliminating White Spaces

tardis/io/atom_data/base.py

@@ -330,12 +330,6 @@ def prepare_atom_data(
 
         self.nlte_species = nlte_species
 
-        self.levels = self.levels[
-            self.levels.index.isin(
-                self.selected_atomic_numbers, level="atomic_number"
-            )
-        ]
-
         self.levels_index = pd.Series(
             np.arange(len(self.levels), dtype=int), index=self.levels.index
         )

tardis/montecarlo/montecarlo_numba/formal_integral.py

@@ -274,6 +274,7 @@ def __init__(self, model, plasma, runner, points=1000):
             )
             self.atomic_data = plasma.atomic_data
             self.original_plasma = plasma
+            self.levels_index = plasma.levels
 
     def generate_numba_objects(self):
         """instantiate the numba interface objects
@@ -387,7 +388,7 @@ def make_source_function(self):
         #macro_data = self.atomic_data.macro_atom_data
         macro_data = self.original_plasma.macro_atom_data
 
-        no_lvls = len(self.atomic_data.levels)
+        no_lvls = len(self.levels_index)
         no_shells = len(model.w)
 
         if runner.line_interaction_type == "macroatom":

tardis/plasma/properties/atomic.py

@@ -72,8 +72,7 @@ class Levels(BaseAtomicDataProperty):
     )
 
     def _filter_atomic_property(self, levels, selected_atoms):
-        return levels
-        # return levels[levels.atomic_number.isin(selected_atoms)]
+        return levels[levels.index.isin(selected_atoms, level="atomic_number")]
 
     def _set_index(self, levels):
         # levels = levels.set_index(['atomic_number', 'ion_number',