psd.py

import pathlib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


def read_data(path_to_csv):
    df = pd.read_csv(path_to_csv, skiprows=3)
    unit = df.iloc[:, 1].name.split("(")[-1][:-1]

    return df.iloc[:, 0], df.iloc[:, 1], unit


def calculate_psd(path_to_folder: str or pathlib.Path):
    folder_path = pathlib.Path(path_to_folder)
    assert folder_path.exists(), "Path to data folder doesn't exist"

    csv_files = folder_path.glob("*.csv")

    results = {}

    for pth in csv_files:
        t, x, unit = read_data(pth)
        fs = 1 / (t.diff().mean())

        # remove zero-shift
        x -= x.mean()

        # apply windowing function
        w = np.hamming(len(x))
        xw = x * w

        # calculate fft
        nfft = len(x)
        freq = np.arange(0, fs / 2, fs / nfft)[:-1]
        X = np.fft.fft(xw, nfft)

        # convert to PSD
        psdX = np.abs(X) ** 2
        psdX = psdX[:nfft // 2]

        # scale the PSD
        # scaling factor (from window function)
        S = sum(w ** 2)
        psdX = 2 * psdX / (fs * S)

        # from scipy.signal import periodogram
        # Calculate PSD using periodogram function (gives same results)
        # freq2, psd2 = periodogram(x, fs, "hamming", nfft)[:nfft//2]

        results[pth.stem] = {"freq": freq, "PSD": psdX, "unit": unit}

    return results


def plot_results(results: dict[str: dict], save_folder: str or pathlib.Path, f_start: float, f_finish: float,
                 title: str, xlabel: str, ylabel: str):
    save_folder = pathlib.Path(save_folder)
    plt.figure()
    global_unit = None
    for label, data in results.items():
        freq = data["freq"]
        psdX = data["PSD"]
        if global_unit is None:
            global_unit = data["unit"]
        else:
            assert global_unit == data["unit"], "All data must have same unit!"
        plt.plot(freq, psdX, label=label, linestyle="solid")

    plt.title(title)
    plt.xlabel(xlabel)
    lab = ylabel.format(global_unit)
    plt.ylabel(lab)

    save_folder.mkdir(parents=True, exist_ok=True)

    plt.xlim(f_start, f_finish)
    plt.legend()
    plt.savefig(save_folder.joinpath("psd_comparison.png"), dpi=200)
    plt.savefig(save_folder.joinpath("psd_comparison.pdf"), )


if __name__ == '__main__':
    # --- set these and then run --------------
    path_to_data_folder = "./data/test"
    f_start = 0.
    f_finish = 0.5
    # --- optional:
    result_folder = path_to_data_folder  # by default, results are saved to same folder as data
    title_of_plot = "Power Spectral Density"
    xlabel_of_plot = 'Frequency (Hz)'
    ylabel_of_plot = "Power Spectral Density ({}$^2$/Hz)"  # {} will be replaced by unit of the time series
    # -----------------------------------------

    results = calculate_psd(path_to_data_folder)
    plot_results(results, result_folder, f_start, f_finish, title_of_plot, xlabel_of_plot, ylabel_of_plot)
    plt.show()