|
| 1 | +import typing as t |
| 2 | +from pathlib import Path |
| 3 | + |
| 4 | +import pandas as pd |
| 5 | +import polars as pl |
| 6 | + |
| 7 | +import artisatomic |
| 8 | + |
| 9 | +# from astropy import units as u |
| 10 | +# import os.path |
| 11 | +# import pandas as pd |
| 12 | +# from carsus.util import parse_selected_species |
| 13 | + |
| 14 | +BASEPATH = Path(Path.home() / "Google Drive/Shared drives/Atomic Data Group/FloersOpacityPaper/OutputFiles") |
| 15 | + |
| 16 | + |
| 17 | +def extend_ion_list(ion_handlers, calibrated=True): |
| 18 | + assert BASEPATH.is_dir() |
| 19 | + handlername = "floers25" + "calib" if calibrated else "uncalib" |
| 20 | + for s in BASEPATH.glob("*_levels_calib.txt"): |
| 21 | + ionstr = s.name.lstrip("0123456789").split("_")[0] |
| 22 | + elsym = ionstr.rstrip("IVX") |
| 23 | + ion_stage_roman = ionstr.removeprefix(elsym) |
| 24 | + atomic_number = artisatomic.elsymbols.index(elsym) |
| 25 | + |
| 26 | + ion_stage = artisatomic.roman_numerals.index(ion_stage_roman) |
| 27 | + |
| 28 | + found_element = False |
| 29 | + for tmp_atomic_number, list_ions in ion_handlers: |
| 30 | + if tmp_atomic_number == atomic_number: |
| 31 | + if ion_stage not in [x[0] if len(x) > 0 else x for x in list_ions]: |
| 32 | + list_ions.append((ion_stage, handlername)) |
| 33 | + list_ions.sort() |
| 34 | + found_element = True |
| 35 | + if not found_element: |
| 36 | + ion_handlers.append( |
| 37 | + ( |
| 38 | + atomic_number, |
| 39 | + [(ion_stage, handlername)], |
| 40 | + ) |
| 41 | + ) |
| 42 | + |
| 43 | + ion_handlers.sort(key=lambda x: x[0]) |
| 44 | + return ion_handlers |
| 45 | + |
| 46 | + |
| 47 | +class FloersEnergyLevel(t.NamedTuple): |
| 48 | + levelname: str |
| 49 | + energyabovegsinpercm: float |
| 50 | + g: float |
| 51 | + parity: int |
| 52 | + |
| 53 | + |
| 54 | +class FloersTransition(t.NamedTuple): |
| 55 | + lowerlevel: int |
| 56 | + upperlevel: int |
| 57 | + A: float |
| 58 | + coll_str: float |
| 59 | + |
| 60 | + |
| 61 | +def read_levels_and_transitions(atomic_number: int, ion_stage: int, flog, calibrated: bool): |
| 62 | + # ion_charge = ion_stage - 1 |
| 63 | + elsym = artisatomic.elsymbols[atomic_number] |
| 64 | + ion_stage_roman = artisatomic.roman_numerals[ion_stage] |
| 65 | + calibstr = "calib" if calibrated else "uncalib" |
| 66 | + |
| 67 | + ionstr = f"{atomic_number}{elsym}{ion_stage_roman}" |
| 68 | + levels_file = BASEPATH / f"{ionstr}_levels_{calibstr}.txt" |
| 69 | + lines_file = BASEPATH / f"{ionstr}_transitions_{calibstr}.txt" |
| 70 | + |
| 71 | + artisatomic.log_and_print( |
| 72 | + flog, |
| 73 | + f"Reading Floers+25 {calibstr}rated data for Z={atomic_number} ion_stage {ion_stage} ({elsym} {ion_stage_roman})", |
| 74 | + ) |
| 75 | + |
| 76 | + ionization_energy_in_ev = artisatomic.get_nist_ionization_energies_ev()[(atomic_number, ion_stage)] |
| 77 | + |
| 78 | + assert Path(levels_file).exists() |
| 79 | + dflevels = pl.from_pandas( |
| 80 | + pd.read_csv(levels_file, sep=r"\s+", skiprows=18, dtype_backend="pyarrow", dtype={"J": str}) |
| 81 | + ).with_columns( |
| 82 | + pl.when(pl.col("J").str.ends_with("/2")) |
| 83 | + .then(pl.col("J").str.strip_suffix("/2").cast(pl.Int32) + 1) |
| 84 | + .otherwise( |
| 85 | + pl.col("J").str.strip_suffix("/2").cast(pl.Int32) * 2 + 1 |
| 86 | + ) # the strip_suffix should not be needed (does not end in "/2" but prevents a polars error) |
| 87 | + .alias("g") |
| 88 | + ) |
| 89 | + |
| 90 | + dflevels = dflevels.with_columns(pl.col("J").str.strip_suffix("/2").cast(pl.Float32).alias("2J")) |
| 91 | + |
| 92 | + energy_levels_zerodindexed = [ |
| 93 | + FloersEnergyLevel( |
| 94 | + levelname=row["Configuration"], |
| 95 | + parity=row["Parity"], |
| 96 | + g=row["g"], |
| 97 | + energyabovegsinpercm=float(row["Energy"]), |
| 98 | + ) |
| 99 | + for row in dflevels.iter_rows(named=True) |
| 100 | + ] |
| 101 | + |
| 102 | + energy_levels = [None, *energy_levels_zerodindexed] |
| 103 | + |
| 104 | + artisatomic.log_and_print(flog, f"Read {len(energy_levels[1:]):d} levels") |
| 105 | + |
| 106 | + transitions = [] |
| 107 | + dftransitions = pl.from_pandas(pd.read_csv(lines_file, sep=r"\s+", skiprows=28, dtype_backend="pyarrow")) |
| 108 | + |
| 109 | + transitions = [ |
| 110 | + FloersTransition(lowerlevel=lowerindex + 1, upperlevel=upperindex + 1, A=A, coll_str=-1) |
| 111 | + for lowerindex, upperindex, A in dftransitions[["Lower", "Upper", "A"]].iter_rows() |
| 112 | + ] |
| 113 | + |
| 114 | + transition_count_of_level_name = { |
| 115 | + config: ( |
| 116 | + dftransitions.filter(pl.col("Config_Lower") == config).height |
| 117 | + + dftransitions.filter(pl.col("Config_Upper") == config).height |
| 118 | + ) |
| 119 | + for config in dflevels["Configuration"] |
| 120 | + } |
| 121 | + |
| 122 | + # this check is slow |
| 123 | + # assert sum(transition_count_of_level_name.values()) == len(transitions) * 2 |
| 124 | + |
| 125 | + artisatomic.log_and_print(flog, f"Read {len(transitions)} transitions") |
| 126 | + |
| 127 | + return ionization_energy_in_ev, energy_levels, transitions, transition_count_of_level_name |
| 128 | + |
| 129 | + |
| 130 | +def get_level_valence_n(levelname: str): |
| 131 | + part = levelname.split(".")[-1] |
| 132 | + if part[-1] not in "spdfg": |
| 133 | + # end of string is a number of electrons in the orbital, not a principal quantum number, so remove it |
| 134 | + assert part[-1].isdigit() |
| 135 | + part = part.rstrip("0123456789") |
| 136 | + return int(part.rstrip("spdfg")) |
0 commit comments