105 lines
4.4 KiB
Python
105 lines
4.4 KiB
Python
"""
|
|
Turbomole-Spectrum-Plotter
|
|
(c) 2022 Lukas Schank
|
|
|
|
This script will run through all subfolders of the given folder and plot excitation spectra of TDDFT calculations.
|
|
|
|
Call on command line like this:
|
|
python exspectrum_plotter.py --data_folder "C:\Path\to\Calculations" --output_folder "C:\path\to\results" --calc_sigma 0.15 --output_type "pdf" --save_csv False
|
|
|
|
calc_sigma output_type and save_csv are optional.
|
|
"""
|
|
|
|
import argparse
|
|
import numpy as np
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
from pathlib import Path
|
|
|
|
|
|
def exspectrum_plotter(data_folder: str, output_folder: str, calc_sigma=0.15, output_type="pdf", save_csv=False):
|
|
|
|
"""
|
|
Iterates through all subfolders of data_folder and plots data of exspectrum files.
|
|
:param data_folder: path to input folder
|
|
:param output_folder: path to output folder, gets created if necessary
|
|
:param calc_sigma: standard deviation of gauss broadening for calculated transitions
|
|
:param output_type: graphs can be saved as "pdf" or "png"
|
|
:param save_csv: write interpolated spectrum as csv
|
|
:return: files to the folder specified in output_folder
|
|
"""
|
|
|
|
# convert input to path objects
|
|
data_folder = Path(data_folder)
|
|
output_folder = Path(output_folder)
|
|
if not output_folder.exists():
|
|
output_folder.mkdir()
|
|
|
|
# convert data to plot
|
|
for data_file in data_folder.glob('**/exspectrum'): # path.glob iterates through given folder for pattern
|
|
# read data
|
|
rawdata = pd.read_csv(
|
|
data_file,
|
|
sep='\s+', # seperator is one or more spaces
|
|
skiprows=2, # skip header rows
|
|
header=None # prevent getting first numbers as header
|
|
)
|
|
data_energy = np.array(rawdata.iloc[:, 3])
|
|
data_osc_strength = np.array(rawdata.iloc[:, 6])
|
|
|
|
# calculate gauss broadening for spectrum
|
|
def gauss_spectrum(energy, osc_strength, sigma, gauss_energy_range):
|
|
gauss_osc_strength = []
|
|
for E_i in gauss_energy_range:
|
|
tot = 0
|
|
for E_j, osc in zip(energy, osc_strength):
|
|
tot += osc * np.exp(-(((E_j - E_i) / sigma) ** 2))
|
|
gauss_osc_strength.append(tot)
|
|
return gauss_osc_strength
|
|
|
|
calc_energy_range = np.linspace(0, 6, num=500, endpoint=True) # x values for calculated spectrum
|
|
calc_osc_strength = gauss_spectrum(data_energy, data_osc_strength, calc_sigma, calc_energy_range)
|
|
|
|
# create plot
|
|
fig, ax = plt.subplots(dpi=300, figsize=(6, 4))
|
|
ax.plot(calc_energy_range, calc_osc_strength, "-k") # plot calculated spectrum
|
|
for plt_energy, plt_osc_strength in zip(data_energy, data_osc_strength):
|
|
ax.plot((plt_energy, plt_energy), (0, plt_osc_strength), c="k") # plot lines from data
|
|
ax.set_xlabel("Energy / eV", fontsize=16)
|
|
ax.set_ylabel("Osc. Strength / a.u.", fontsize=16)
|
|
ax.xaxis.set_tick_params(labelsize=14, width=1.5)
|
|
ax.yaxis.set_tick_params(labelsize=14, width=1.5)
|
|
for axis in ['top', 'bottom', 'left', 'right']:
|
|
ax.spines[axis].set_linewidth(1.5)
|
|
ax.set_xlim(0, 6)
|
|
plt.tight_layout()
|
|
|
|
# save plot
|
|
data_name = Path(data_file)
|
|
output_name = data_name.parts[data_folder.parts.__len__()]
|
|
for name in data_name.parts[(data_folder.parts.__len__() + 1):]:
|
|
output_name += "_" + name
|
|
output_name_plot = output_name + "." + output_type
|
|
output_path = output_folder / output_name_plot
|
|
plt.savefig(output_path)
|
|
plt.close(fig) # close plot
|
|
|
|
# save calculated data
|
|
if save_csv:
|
|
csv_output = pd.DataFrame({'Energy': calc_energy_range.tolist(), 'Osc. Strength': calc_osc_strength})
|
|
output_name_csv = output_name + ".csv"
|
|
output_path_2 = output_folder / output_name_csv
|
|
csv_output.to_csv(output_path_2)
|
|
|
|
|
|
# getting commandline input and pass it to function
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--data_folder', type=str, required=True)
|
|
parser.add_argument('--output_folder', type=str, required=True)
|
|
parser.add_argument('--calc_sigma', type=float, required=False)
|
|
parser.add_argument('--output_type', type=str, required=False)
|
|
parser.add_argument('--save_csv', type=bool, required=False)
|
|
args = parser.parse_args()
|
|
exspectrum_plotter(**vars(args))
|