salary02/data_analysis.py

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

num_time_step = 201
num_iter = 10

env_data = pd.DataFrame(pd.read_excel('env_data.xlsx', engine='openpyxl', sheet_name=0))

assert env_data.shape[0] == num_iter * (num_time_step + 1), f"{env_data.shape[0]}, {num_iter * (num_time_step + 1)}"

lst_df = []
for i in range(num_iter):
    df_tmp = env_data.iloc[i * (num_time_step + 1): (i + 1) * (num_time_step + 1), 1:]
    lst_df.append(df_tmp)

lst_column = lst_df[0].columns

for str_col in lst_column:
    x = np.arange(num_time_step+1)
    for df in lst_df:
        y = np.array(df[str_col]).flatten()
        plt.xlabel('t')
        plt.ylabel(str_col)
        plt.plot(x, y)
    # plt.show()
    # plt.close()
    plt.savefig(f'{str_col}.pdf', bbox_inches="tight")
    plt.close()