Update readtanakajpltdata.py

artis-mcrt · Dec 18, 2024 · 7dc707d · 7dc707d
1 parent b5c9d92
commit 7dc707d
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/artisatomic/readtanakajpltdata.py b/artisatomic/readtanakajpltdata.py
@@ -79,8 +79,10 @@ def read_levels_and_transitions(atomic_number, ion_stage, flog):
             # dflevels = pd.concat(reader, ignore_index=True)
 
             dflevels = artisatomic.add_dummy_zero_level(
-                pl.from_pandas(reader.get_chunk(levelcount)).with_columns(
-                    pl.col("g").cast(pl.Float64),
+                pl.from_pandas(reader.get_chunk(levelcount)).select(
+                    energyabovegsinpercm=pl.col("energy_ev").cast(pl.Float64) / hc_in_ev_cm,
+                    parity=pl.when(pl.col("parity").str.strip_chars() == "odd").then(1).otherwise(0),
+                    g=pl.col("g").cast(pl.Float64),
                     levelname=pl.format(
                         "{},{},{}", pl.col("num"), pl.col("parity"), pl.col("configuration").str.strip_chars()
                     ),
@@ -90,12 +92,12 @@ def read_levels_and_transitions(atomic_number, ion_stage, flog):
 
             energy_levels = [None]
             for row in dflevels[1:].iter_rows(named=True):
-                parity = 1 if row["parity"].strip() == "odd" else 0
-                energyabovegsinpercm = float(row["energy_ev"] / hc_in_ev_cm)
-
-                energy_levels.append(
+                energy_levels.append(  # noqa: PERF401
                     EnergyLevel(
-                        levelname=row["levelname"], parity=parity, g=row["g"], energyabovegsinpercm=energyabovegsinpercm
+                        levelname=row["levelname"],
+                        parity=row["parity"],
+                        g=row["g"],
+                        energyabovegsinpercm=row["energyabovegsinpercm"],
                     )
                 )