cross-section-TD.py

#!/usr/bin/python3

import math
import sys
import os
import time

def read_inp(file_name):
    pass

def read_data(path_file, method=float) -> list:
    tmp = open(path_file, 'r').read().splitlines()
    data = list(map(method, tmp))
    data = list(map(abs, data))
    return data

def data_set_process(path_file) -> (dict, dict):
    E_dict = {}
    f_dict = {}
    for each_line in open(path_file, 'r'):
        index, E, f = each_line.split()
        E_dict[int(index)] = float(E)
        f_dict[int(index)] = float(f)
    return E_dict, f_dict
        

def gauss(E, dE_0n) -> float:
    # coeff = 1 / (c.delta * math.sqrt(pi / 2))
    # devided_num = - c.delta**2 / 2
    to_be_exp = ((E - dE_0n) ** 2) / devided_num
    return coeff * math.exp(to_be_exp)
    

class Config():
    spec_type = 1
    init_state = 1
    # final_state = [2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
    final_state = [2, 3, 4]
    prob_kind = 'F'
    screen = 0
    os_condon = -1
    norm = 'local'
    seed = 0
    function = 'gauss'
    delta = 0.05
    temp = 0
    nref = 1
    eps = 0.002
    kappa = 3
    run_IS = 0

if __name__ == "__main__":
    # start_time = time.process_time()
    ######  config read   #####
    c = Config()

    #####  TD init     #####
    TD_num = sys.argv[1]
    data_path = sys.argv[2]
    max_state = int(sys.argv[3])
    data_name_template = 'state%d.index.E.f'
    itrain_name = 'itrain.dat'
    train_index = read_data(itrain_name, method=int)


    ######  data init   #####
    if True:
        BK        = 0.3166813639E-5     # Boltzmann constant Hartree/kelvin
        bk_ev     = 8.617343E-5         # Boltzmann constant eV/kelvin
        proton    = 1822.888515         # (a.m.u.)/(electron mass)
        timeunit  = 24.188843265E-3     # femtoseconds
        au2ev     = 27.21138386         # au to eV
        pi        = 3.141592653589793   # pi
        au2ang    = 0.52917720859       # au to angstrom
        deg2rad   = 1.745329252E-2      # degree to radian (pi/180)
        au2cm     = 219474.625          # hartree to cm-1
        cm2au     = 4.55633539E-6       # cm-1 to hartree
        h_planck  = 1.5198298508E-16    # h hartree.s
        h_evs     = 4.13566733E-15      # h eV.s
        light     = 299792458           # speed of light m/s
        lightau   = 137.035999139       # speed of light au
        e_charge  = -1.602176487E-19    # electron charge C
        e_mass    = 9.10938215E-31      # electron mass kg
        eps0      = 8.854187817e-12     # vacuum permitivity C^2*s^2*kg^-1*m^-3
        # gauss function const
        coeff = 1 / (c.delta * math.sqrt(pi / 2))
        devided_num = - c.delta**2 / 2
        # cross section const
        hplanck = h_evs
        ev2nm = 1E9 * hplanck * light
        gamma = 1 # temperature = 0  ==>  gamma = 1
        nref = 1 # ratio
        cs_coeff_pre = pi * e_charge**2 / (2 * e_mass * light * eps0 * nref)  # m^2/s  unit
        cs_coeff = cs_coeff_pre * hplanck / (2 * pi) * 1E20                   # Angstrom^2*eV

    #####   read_data     #####
    E_data_list = []
    f_data_list = []
    for i in range(1, max_state + 1):
        path_file_name = data_path + '/' + data_name_template % i
        E_data, f_data = data_set_process(path_file_name)
        E_data = [v for k, v in E_data.items() if k in train_index]
        f_data = [v for k, v in f_data.items() if k in train_index]
        # E_data = list(filter(lambda x: x in train_index, E_data))
        # f_data = list(filter(lambda x: x in train_index, f_data))
        E_data_list.append(E_data)
        f_data_list.append(f_data)
    
    #####   data pre-process     #####
    E_max = -10000
    E_min = 10000
    whole_num_point = 0
    for E_data in E_data_list:
        tmp_max = max(E_data)
        tmp_min = min(E_data)
        E_max = tmp_max if tmp_max > E_max else E_max
        E_min = tmp_min if tmp_min < E_min else E_min
    whole_num_point = len(E_data)  # todo: check the E and f is has the same number 
    range_min = E_min - c.kappa * c.delta 
    range_max = E_max + c.kappa * c.delta
    # de = round(range_min, 4)
    de = range_min
    
    ##### calc cross section #####
    file_handle = open('cross-section.dat', 'w')
    file_handle.write('DE/eV    lambda/nm    sigma/A^2    +/-error/A^2\n')
    while de < range_max:
        spect = 0.0
        for state_index, E_data in enumerate(E_data_list):
            f_data = f_data_list[state_index]
            for target_index, E in enumerate(E_data):
                f = f_data[target_index]
                spect += E * f * gauss(de, E)
        cross_section = cs_coeff * spect * gamma / (whole_num_point * de)
        wavelength = ev2nm / de
        file_handle.write('%.4f   %.4f     %.8f    0.00000\n' % (de, wavelength, cross_section))
        de += c.eps
    file_handle.close()
    # end_time = time.process_time()
    # print(f'start time: {start_time} \nend time: {end_time}')
    # print(f'cost time: {end_time - start_time}')