plot_style_ScattRate.py

#   This file takes a directory of dat files, combines them and plots the output 
#   It is adapted from a script that Philip wrote
#   I just reconfigured them to use the GUI from the main script 


import os
import numpy as np
import matplotlib.pyplot as plt
import plot_settings as ps
from plot_settings import  *

def scattrate_plotter(real_folder, complex_folder):


    # Define file paths for combined data files in the base directory
    base_dir = os.path.dirname(real_folder)  # Use the parent directory
    write_real_Ek_scattrate_full = os.path.join(base_dir, 'realEk_scattering_rate_FullData.dat')
    write_complex_Ek_scattrate_full = os.path.join(base_dir, 'complexEk_scattering_rate_FullData.dat')

    # Collect all .dat files in the folders
    real_files = [os.path.join(real_folder, f) for f in os.listdir(real_folder) if f.endswith('.dat')]
    complex_files = [os.path.join(complex_folder, f) for f in os.listdir(complex_folder) if f.endswith('.dat')]

    # Combine files if needed
    ps.combine_file_data(real_files, write_real_Ek_scattrate_full)
    ps.combine_file_data(complex_files, write_complex_Ek_scattrate_full)

    # Load the combined data using genfromtxt
    data_real = np.genfromtxt(write_real_Ek_scattrate_full, skip_header=1)
    data_complex = np.genfromtxt(write_complex_Ek_scattrate_full, skip_header=1)

    # Plot the real and complex scattering rates
    plot_scattering_rate(data_real, 'Real')
    plot_scattering_rate(data_complex, 'Complex')

def plot_scattering_rate(data, flag):
    E = data[:, 0]
    kx = data[:, 1]
    ky = data[:, 2]
    kz = data[:, 3]
    rate = data[:, 4]

    unique_E = np.unique(E)
    unique_kx = np.unique(kx)

    kz_scatt_rate_integral = np.zeros((len(unique_E), len(unique_kx)))

    for i, e in enumerate(unique_E):
        for j, k in enumerate(unique_kx):
            mask = (E == e) & (kx == k)
            kz_scatt_rate_integral[i, j] = np.sum(rate[mask])

    scatt_rate_vs_E = np.sum(kz_scatt_rate_integral, axis=1)

    smallest_E_index = np.argmin(unique_E)
    largest_E_index = np.argmax(unique_E)

    plt.figure()
    plt.plot(unique_kx, kz_scatt_rate_integral[smallest_E_index, :], label="Smallest Energy")
    plt.plot(unique_kx, kz_scatt_rate_integral[largest_E_index, :], label="Largest Energy")
    plt.xlabel("kx [1/nm]")
    plt.ylabel("Scattering Rate Integral [1/s]")
    plt.title(flag + " Scattering Rate vs kx for Smallest and Largest Energy")
    plt.legend()
    plt.savefig(flag + "_scatt_rate_vs_kx.png")
    plt.show()

    plt.figure()
    plt.plot(unique_E, scatt_rate_vs_E)
    plt.xlabel("Energy [eV]")
    plt.ylabel("Scattering Rate [1/s]")
    plt.title(flag + " Scattering Rate vs Energy")
    plt.savefig(flag + "_scatt_rate.png")
    plt.show()

def plot(real_folder, complex_folder):
    try:
        scattrate_plotter(real_folder, complex_folder)
    except:
        print(f"Error in plotter.")