AMM/exspectrum_plotter.py
2022-02-20 18:39:04 +01:00

81 lines
3.1 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.
"""
import sys
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):
# 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_sigma = 0.15 # std for gauss broadening
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(1.5, 5)
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 += ".pdf" # change here for pdf or png output
output_path = output_folder / output_name
plt.savefig(output_path)
plt.close(fig) # close plot
# getting commandline input and pass it to function
if __name__ == "__main__":
a = str(sys.argv[1])
b = str(sys.argv[2])
exspectrum_plotter(a, b)
# python exspectrum_plotter.py "C:\Path\to\Calculations" "C:\path\to\results"