From 16b5948a36e709a64330c251de7ae0a7807e7d16 Mon Sep 17 00:00:00 2001
From: sakul-45 <81963567+sakul-45@users.noreply.github.com>
Date: Thu, 3 Mar 2022 23:32:02 +0100
Subject: [PATCH] Sigma and output type are optional arguments now. Also added
 an option to store calculated spectra as csv files.

---
 .idea/AMM.iml         |  5 ++++-
 exspectrum_plotter.py | 46 ++++++++++++++++++++++++++++++++-----------
 requirements.txt      |  1 -
 3 files changed, 39 insertions(+), 13 deletions(-)

diff --git a/.idea/AMM.iml b/.idea/AMM.iml
index 3357973..49f4c24 100644
--- a/.idea/AMM.iml
+++ b/.idea/AMM.iml
@@ -4,7 +4,10 @@
     <content url="file://$MODULE_DIR$">
       <excludeFolder url="file://$MODULE_DIR$/venv" />
     </content>
-    <orderEntry type="jdk" jdkName="Python 3.8 (AMM)" jdkType="Python SDK" />
+    <orderEntry type="inheritedJdk" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
+  <component name="PackageRequirementsSettings">
+    <option name="versionSpecifier" value="Greater or equal (&gt;=x.y.z)" />
+  </component>
 </module>
\ No newline at end of file
diff --git a/exspectrum_plotter.py b/exspectrum_plotter.py
index 9ef369f..008bad2 100644
--- a/exspectrum_plotter.py
+++ b/exspectrum_plotter.py
@@ -3,16 +3,31 @@ 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 sys
+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):
+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)
@@ -43,7 +58,6 @@ def exspectrum_plotter(data_folder: str, output_folder: str):
             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
@@ -57,7 +71,7 @@ def exspectrum_plotter(data_folder: str, output_folder: str):
         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)
+        ax.set_xlim(0, 6)
         plt.tight_layout()
 
         # save plot
@@ -65,16 +79,26 @@ def exspectrum_plotter(data_folder: str, output_folder: str):
         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
+        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__":
-    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"
+    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))
diff --git a/requirements.txt b/requirements.txt
index 1dacff3..1538cc4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,3 @@
 numpy~=1.22.1
 pandas~=1.4.0
 matplotlib~=3.5.1
-pathlib~=1.0.1