Plotting/DyLevelStructure/visualizeDyLevelStructureWithTransitions.py

import os, csv
from pathlib import Path
import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

sns.set_theme(style="white")

def parse_NIST_data(path, min_J, max_J, max_wavenumber):

    data_list = []
    with open(path) as f:
        reader = csv.reader(f, delimiter="\t")
        data_list = list(reader)
    
    #collect data
    Parity = np.zeros(len(data_list))
    J = np.zeros(len(data_list))
    wavenumber = np.zeros(len(data_list))
    for i in range(1, len(data_list)):
        try:
            tmp = data_list[:][i]
            if not tmp[0] == '' and tmp[0].find('*') != -1:
                Parity[i] = 1
            elif not tmp[0] == '':
                Parity[i] = 0
            J[i] = int(tmp[1])
            wavenumber[i] = float(tmp[3])
            if wavenumber[i] > max_wavenumber:
                J[i] = np.nan
                wavenumber[i] = np.nan    
        except ValueError:
            J[i] = np.nan
            wavenumber[i] = np.nan

    remove_idxs = []
    for i in range(1, len(data_list)):
        p = Parity[i]
        j = J[i]
        wn = wavenumber[i]
        if np.isnan(p) or np.isnan(j) or np.isnan(wn):
            remove_idxs.append(i)
    
    Parity = np.delete(Parity, remove_idxs)
    J = np.delete(J, remove_idxs)
    wavenumber = np.delete(wavenumber, remove_idxs)
    
    #sort data
    sorting_indices = np.argsort(J)
    Parity = Parity[sorting_indices]
    J = J[sorting_indices]
    wavenumber = wavenumber[sorting_indices]
    
    # splice data to within user-defined range
    splice_idx_start = np.where(J==min_J)[0][0]
    splice_idx_stop = len(J) - 1 - np.where(J[::-1]==max_J)[0][0]

    Parity = Parity[splice_idx_start:splice_idx_stop]
    J = J[splice_idx_start:splice_idx_stop]
    wavenumber = wavenumber[splice_idx_start:splice_idx_stop]

    # Create a Pandas data frame with the data
    dataset = pd.DataFrame(np.array(list(zip(Parity, J, wavenumber))), columns=['Parity', 'J', 'Wavenumber'])

    return dataset

def plot_level_structure_with_red_and_blue_transitions(*args, **kwargs):
    
    #start plotting
    f, ax = plt.subplots(figsize=(4, 8))
    named_colors = ['r', 'm']
    Red_Blue_colors = ['#ab162a', '#cf5246', '#eb9172', '#fac8af', '#faeae1', '#e6eff4', '#bbdaea', '#7bb6d6', '#3c8abe', '#1e61a5']

    #draw levels
    plot_handle = sns.scatterplot(x='J', y='Wavenumber', data = dataframe, s=500, hue = 'Parity', palette = sns.color_palette(named_colors), marker = '_', linewidth=1, legend=False)
    
    #write electronic configuration for GS
    ax.text(gs_J + 0.15, gs_wavenumber + 400, '$6s^2$') 
    #draw guide line for GS
    plt.axhline(y=gs_wavenumber, color='m', linestyle='--', linewidth=1, alpha=0.5)
    
    #write wavelength of red transition
    ax.text(red_J - 0.4, red_wavenumber * 0.5, '$626.082 ~ \mathrm{nm}$')
    #draw red transition arrow
    ax.annotate('', 
            xy=(red_J, red_wavenumber),
            xytext=(gs_J, gs_wavenumber), 
            arrowprops=dict(color='#db2929', alpha=0.8, width=1.5),
            horizontalalignment='right',
            verticalalignment='top')

    #write electronic configuration for triplet excited state
    ax.text(red_J + 0.18, red_wavenumber + 400, '$6s6p(^3P_1)$') 
    #draw guide line for triplet excited state
    plt.axhline(y=red_wavenumber, color='m', linestyle='--', linewidth=1, alpha=0.5)

    #write wavelength of red transition
    ax.text(blue_J - 1.5, blue_wavenumber * 0.55, '$421.291~ \mathrm{nm}$')
    #draw blue transition arrow
    ax.annotate('', 
            xy=(blue_J, blue_wavenumber),
            xytext=(gs_J, gs_wavenumber), 
            arrowprops=dict(color='#2630ea', alpha=0.8, width=3.5),
            horizontalalignment='right',
            verticalalignment='top')

    #write electronic configuration for singlet excited state
    ax.text(blue_J + 0.18, blue_wavenumber + 400, '$6s6p(^1P_1)$') 
    #draw guide line for singlet excited state
    plt.axhline(y=blue_wavenumber, color='m', linestyle='--', linewidth=1, alpha=0.5)

    #figure options
    plt.xlabel('$J$', fontsize=16)
    plt.ylabel('$\\tilde{v}~(cm^{-1})$', fontsize=16)
    #plt.title('Dysprosium I Energy Level Structure', fontsize=20)
    plot_handle.set_xticks(range(min_J-1, max_J+2))
    plt.tick_params(axis='both', which='major', labelsize=12)
    ax.get_xticklabels()[0].set_visible(False)
    ax.get_xticklabels()[-1].set_visible(False)
    #plt.show()
    
    f.savefig(Path(home_path + os.sep + 'result.pdf'), format='pdf', bbox_inches = "tight")
    
if __name__ == '__main__':
    
    min_J = 8
    max_J = 9
    max_wavenumber = 25500.0

    #Ground State: [Xe]4f^{10} 6s^2(^5I_8)
    gs_J = 8
    gs_wavenumber = 0.0

    #626 nm transition GS ---> [Xe]4f^{10}(^5I_8) 6s6p(^3P_1)(8,1)_9
    red_J = 9
    red_wavenumber = 15972.35

    #421 nm transition GS ---> [Xe]4f^{10}(^5I_8) 6s6p(^1P_1)(8,1)_9
    blue_J = 9
    blue_wavenumber = 23736.610

    NIST_Dy_level_data_filename= 'Dylevels.txt'
    home_path = str(Path(__file__).parent.resolve())
    
    dataframe = parse_NIST_data(home_path + os.sep + NIST_Dy_level_data_filename, min_J, max_J, max_wavenumber)
    plot_level_structure_with_red_and_blue_transitions(dataframe, gs_J, gs_wavenumber, red_J, red_wavenumber, blue_J, blue_wavenumber, home_path)