Use polars for transitions write out

artisatomic/__init__.py

@@ -19,6 +19,7 @@
 import argcomplete
 import numpy as np
 import pandas as pd
+import polars as pl
 from astropy import constants as const
 from astropy import units as u
 from scipy import integrate
@@ -232,7 +233,7 @@ def process_files(ion_handlers: list[tuple[int, list[int | tuple[int, str]]]], a
         thresholds_ev_dict: list[dict] = [{} for _ in listions]
 
         # list of named tuples (hillier_transition_row)
-        transitions: list = [[] for x in listions]
+        transitions: list = [[] for _ in listions]
         transition_count_of_level_name: list[dict] = [{} for _ in listions]
         upsilondicts: list[dict] = [{} for x in listions]
 
@@ -1533,9 +1534,17 @@ def write_output_files(
                 flog,
             )
 
+        dftransitions_ion = pl.DataFrame(transitions[i])
         with open(os.path.join(args.output_folder, "transitiondata.txt"), "a") as ftransitiondata:
             write_transition_data(
-                ftransitiondata, atomic_number, ion_stage, energy_levels[i], transitions[i], upsilondicts[i], args, flog
+                ftransitiondata,
+                atomic_number,
+                ion_stage,
+                energy_levels[i],
+                dftransitions_ion,
+                upsilondicts[i],
+                args,
+                flog,
             )
 
         if i < len(listions) - 1 and not args.nophixs:  # ignore the top ion
@@ -1612,28 +1621,25 @@ def write_adata(
 
 
 def write_transition_data(
-    ftransitiondata, atomic_number, ion_stage, energy_levels, transitions, upsilondict, args, flog
+    ftransitiondata, atomic_number, ion_stage, energy_levels, dftransitions_ion, upsilondict, args, flog
 ):
-    log_and_print(flog, f"Writing {len(transitions)} transitions to 'transitiondata.txt'")
+    log_and_print(flog, f"Writing {len(dftransitions_ion)} transitions to 'transitiondata.txt'")
 
     num_forbidden_transitions = 0
     num_collision_strengths_applied = 0
-    ftransitiondata.write(f"{atomic_number:7d}{ion_stage:7d}{len(transitions):12d}\n")
+    ftransitiondata.write(f"{atomic_number:7d}{ion_stage:7d}{len(dftransitions_ion):12d}\n")
 
     level_ids_with_permitted_down_transitions = set()
-    for transition in transitions:
-        levelid_lower = transition.lowerlevel
-        levelid_upper = transition.upperlevel
+    for levelid_lower, levelid_upper in dftransitions_ion[["lowerlevel", "upperlevel"]].itertuples(index=False):
         forbidden = energy_levels[levelid_lower].parity == energy_levels[levelid_upper].parity
 
         if not forbidden:
             level_ids_with_permitted_down_transitions.add(levelid_upper)
 
-    for transition in transitions:
-        levelid_lower = transition.lowerlevel
-        levelid_upper = transition.upperlevel
+    for levelid_lower, levelid_upper, A, coll_str in dftransitions_ion[
+        ["lowerlevel", "upperlevel", "A", "coll_str"]
+    ].itertuples(index=False):
         assert levelid_lower < levelid_upper
-        coll_str = transition.coll_str
 
         if coll_str > 0:
             num_collision_strengths_applied += 1
@@ -1642,19 +1648,14 @@ def write_transition_data(
 
         if forbidden:
             num_forbidden_transitions += 1
-            # flog.write(f'Forbidden transition: lambda_angstrom= {float(transition.lambdaangstrom):7.1f}, {transition.namefrom:25s} to {transition.nameto:25s}\n')
 
-        # ftransitiondata.write('{0:4d} {1:4d} {2:16.10E} {3:9.2e} {4:d}\n'.format(
-        #     levelid_lower, levelid_upper, float(transition.A), coll_str, forbidden))
-        ftransitiondata.write(
-            f"{levelid_lower:4d} {levelid_upper:4d} {float(transition.A):11.5e} {coll_str:9.2e} {forbidden:d}\n"
-        )
+        ftransitiondata.write(f"{levelid_lower:4d} {levelid_upper:4d} {float(A):11.5e} {coll_str:9.2e} {forbidden:d}\n")
 
     ftransitiondata.write("\n")
 
     log_and_print(
         flog,
-        f"  output {len(transitions):d} transitions of which {num_forbidden_transitions:d} are forbidden and"
+        f"  output {len(dftransitions_ion):d} transitions of which {num_forbidden_transitions:d} are forbidden and"
         f" {num_collision_strengths_applied:d} have collision strengths",
     )