base.py

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau
from torch.utils.data import DataLoader
import numpy as np

from utils import get_data, TimeSeriesDataset, signal_characteristic
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from models import RNNM, LSTMM, LTCLinear


def get_change_indices(arr):

    fifth_feature = arr[:, -1]

    change_indices = np.where(fifth_feature[:-1] != fifth_feature[1:])[0] + 1

    return change_indices


def split_on_change(arr):
    change_indices = get_change_indices(arr)
    split_arrays = np.split(arr, change_indices)
    split_arrays = np.delete(split_arrays, -1, axis=-1)
    return split_arrays


def smape(y_true, y_pred):
    return (
        1
        / len(y_true)
        * torch.sum(
            2 * torch.abs(y_pred - y_true) / (torch.abs(y_true) + torch.abs(y_pred))
        )
    )


def train_model(
    model: nn.Module,
    train_data: DataLoader,
    val_data,
    val_freq: int,
    val_steps: int,
    n_epochs: int,
    model_name: str,
    lr: float,
    device: torch.device,
    reverse: bool = False,
):
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=lr)
    # scheduler = StepLR(optimizer=optimizer, step_size=5000, gamma=0.1)
    scheduler = ReduceLROnPlateau(optimizer=optimizer, patience=500, factor=0.1)

    losses = []
    val_losses = []

    best_loss = 1e20
    best_val_loss = 1e20
    best_epoch = 0
    epoch_val_loss = 10
    model.to(device)
    for epoch in range(n_epochs):
        model.train()
        epoch_loss = 0.0
        # mean_norm = []
        for i, (inputs, targets) in enumerate(train_data):
            if reverse:
                inputs = inputs.flip((0,))
                targets = targets.flip((0,))

            inputs = inputs.unsqueeze(1).to(device)
            targets = targets.to(device)

            outputs = model(inputs).to(device)

            loss = torch.sqrt(criterion(outputs, targets))
            # loss = smape(targets, outputs)

            epoch_loss += loss.item()
            optimizer.zero_grad()
            loss.backward()

            # total_norm = 0
            # for p in model.parameters():
            #     param_norm = p.grad.data.norm(2)
            #     total_norm += param_norm.item() ** 2
            # total_norm = total_norm ** (1.0 / 2)
            # mean_norm.append(total_norm)

            optimizer.step()
            # scheduler.step(metrics=loss)

        epoch_loss /= len(train_data)

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            torch.save(model.state_dict(), f"models/{model_name}_train.pth")
            best_epoch = epoch
            print(f"New Best Loss: {best_loss}")

        if (epoch + 1) % val_freq == 0:
            model.eval()
            with torch.no_grad():
                val_pred = []
                init_val = (
                    torch.tensor(val_data[0].reshape(1, -1), dtype=torch.float32)
                    .unsqueeze(0)
                    .to(device)
                )
                for i in range(val_steps - 1):
                    output = model(init_val).to(device)
                    val_pred.append(output.squeeze(0))
                    init_val = output.unsqueeze(0).to(device)

            val_pred = torch.stack(val_pred)
            epoch_val_loss = torch.sqrt(
                criterion(val_pred, torch.tensor(val_data[1:], dtype=torch.float32))
            )

            if epoch_val_loss < best_val_loss:
                best_val_loss = epoch_val_loss
                torch.save(model.state_dict(), f"models/{model_name}_val.pth")
                print(f"New Best Val Loss: {best_val_loss}")

        val_losses.append(epoch_val_loss)
        scheduler.step(epoch_loss)
        losses.append(epoch_loss)

        print(
            f"Epoch [{epoch+1}/{n_epochs}], Loss: {epoch_loss:.8f}, Best Loss: {best_loss:.8f}, Best Epoch: {best_epoch}, Val Loss: {epoch_val_loss:.8f}"
        )

    return np.stack((losses, val_losses)), best_loss


def predict(
    model: nn.Module, model_name: str, test_data, pred_steps: int, device: torch.device
):
    model.load_state_dict(torch.load(f"models/{model_name}_train.pth"))
    model.to(device)
    model.eval()

    full_pred = []

    for i in range(test_data.shape[0]):
        init_input = (
            torch.tensor(test_data[i][0].reshape(1, -1), dtype=torch.float32)
            .unsqueeze(0)
            .to(device)
        )
        criterion = nn.MSELoss()
        predictions = []
        with torch.no_grad():
            for j in range(pred_steps - 1):
                pred = model(init_input).to(device)
                predictions.append(pred.squeeze(0))
                init_input = pred.unsqueeze(0).to(device)

        predictions = torch.stack(predictions)
        pred_loss = torch.sqrt(
            criterion(
                predictions,
                torch.tensor(test_data[i][1:pred_steps], dtype=torch.float32),
            )
        )
        print(f"Prediction {i} RMSE: {pred_loss}")
        full_pred.append(predictions)

    return torch.stack(full_pred)


def create_plots(test_data, predictions, plot_range):
    plots = []

    sc_t = []
    sc_p = []
    criterion = nn.MSELoss()

    for i in range(test_data.shape[0]):

        fig, ax = plt.subplots(3, sharex="col", figsize=(8.27, 11.69))

        rmse = torch.sqrt(
            criterion(torch.tensor(test_data[i][1:]), torch.tensor(predictions[i]))
        )
        sc_true = signal_characteristic(test_data[i][:, 0], 200)
        sc_pred = signal_characteristic(predictions[i][:, 0], 199)
        sc_t.append(sc_true)
        sc_p.append(sc_pred)
        ae = criterion(torch.tensor(sc_true[0]), torch.tensor(sc_pred[0]))
        me = criterion(torch.tensor(sc_true[1]), torch.tensor(sc_pred[1]))

        fig.suptitle(
            f"{test_data[i][0, -1]}, RMSE: {rmse:.4f}, RAE: {ae:.4f}, RME: {me:.4f}",
            fontsize=16,
        )

        ax[0].plot(test_data[i][1:plot_range, 0], label="True Values", marker="o")
        ax[0].plot(predictions[i][:plot_range, 0], label="Predictions", marker="x")
        ax[0].set(title=f"f_a: {test_data[i][0, -1]} q1")
        ax[0].set_ylabel("q_1")
        ax[0].grid(True)
        ax[0].legend()

        ax[1].plot(test_data[i][1:plot_range, 1], label="True Values", marker="o")
        ax[1].plot(predictions[i][:plot_range, 1], label="Predictions", marker="x")
        ax[1].set(title=f"f_a: {test_data[i][0, -1]} q2")
        ax[1].set_ylabel("q_2")
        ax[1].grid(True)
        ax[1].legend()

        ax[2].plot(test_data[i][1:plot_range, 2], label="True Values", marker="o")
        ax[2].plot(predictions[i][:plot_range, 2], label="Predictions", marker="x")
        ax[2].set(title=f"f_a: {test_data[i][0, -1]} f_a")
        ax[2].grid(True)
        ax[2].set_ylabel("f_a")
        ax[2].set_xlabel("Timesteps")
        ax[2].legend()

        fig.tight_layout()

        plots.append(fig)

    sc_t = torch.tensor(np.stack(sc_t, dtype=np.float32).squeeze())
    sc_p = torch.tensor(np.stack(sc_p, dtype=np.float32).squeeze())

    rae = criterion(sc_t[:, 0], sc_p[:, 0])
    rme = criterion(sc_t[:, 1], sc_p[:, 1])

    metric = {"RAE": rae, "RME": rme}

    plt.figure()
    plt.axis("off")
    text = "\n".join([f"{key}: {value}" for key, value in metric.items()])
    plt.text(
        0.1,
        0.5,
        text,
        ha="left",
        va="center",
        fontsize=12,
        transform=plt.gca().transAxes,
    )
    metric_dict = plt.gcf()

    plots.append(metric_dict)

    return plots


def make_report(
    model_name: str,
    param_dict: dict,
    losses: np.ndarray,
    predictions: np.ndarray,
    test_data: np.ndarray,
    plot_range: int,
):

    plt.figure()
    plt.axis("off")
    text = "\n".join([f"{key}: {value}" for key, value in param_dict.items()])
    plt.text(
        0.1,
        0.5,
        text,
        ha="left",
        va="center",
        fontsize=12,
        transform=plt.gca().transAxes,
    )
    dict_page = plt.gcf()

    plots = create_plots(test_data, predictions, plot_range)

    plt.figure()
    plt.plot(range(1, param_dict["Epochs"] + 1), losses[0], label="Training Loss")
    # plt.plot(range(1, param_dict["Epochs"] + 1), losses[1], label="Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.yscale("log")
    plt.legend()
    plt.grid(True)
    curve_page = plt.gcf()

    with PdfPages(f"evaluation/{model_name}.pdf") as pdf:
        pdf.savefig(dict_page)
        pdf.savefig(curve_page)

        for i in range(len(plots)):
            pdf.savefig(plots[i])


def main():

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model_name = "lstm_h20_l2_e200k_b500"
    architecture = "LSTM"
    input_dim = 3
    output_dim = 3

    batch_size = 500
    hidden_dim = 20
    num_layers = 2
    lr = 0.001
    num_epochs = 200000

    pred_steps = 2500
    val_freq = 500
    plot_range = 350

    # model = LSTMM(
    #     input_dim=input_dim,
    #     hidden_dim=hidden_dim,
    #     output_dim=output_dim,
    #     num_layers=num_layers,
    #     dropout=0.2,
    # )
    model = LSTMM(
        input_dim=input_dim,
        hidden_dim=hidden_dim,
        output_dim=output_dim,
        num_layers=num_layers,
    )
    # model.load_state_dict(torch.load(f"models/ltc_h20_l2_e200k_train.pth"))

    param_dict = {
        "Model Name": model_name,
        "Architecture": architecture,
        "Input Dim": input_dim,
        "Output Dim": output_dim,
        "Hidden Dim": hidden_dim,
        "Num Layers": num_layers,
        "Batch Size": batch_size,
        "Epochs": num_epochs,
        "Learning Rate": lr,
        "Loss": "RMSE",
        "Shuffle Train": "True",
        "Scaler": "StandardScaler",
    }

    torch.manual_seed(42)

    train_data = get_data("DORA_Train.csv")[:, 1:4]
    test_data = get_data("DORA_Test.csv")[:, 1:5]

    splitted_test = split_on_change(test_data)

    scaler = StandardScaler()

    scaler.fit(train_data)
    train_scaled = scaler.transform(train_data)

    for i in range(len(splitted_test)):
        splitted_test[i] = scaler.transform(splitted_test[i])

    train_dataset = TimeSeriesDataset(train_scaled)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    val_data = splitted_test[2]
    val_dataset = TimeSeriesDataset(val_data)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
    test_data = splitted_test

    losses, best_loss = train_model(
        model=model,
        train_data=train_loader,
        val_data=val_data,
        val_freq=val_freq,
        val_steps=pred_steps,
        n_epochs=num_epochs,
        model_name=model_name,
        lr=lr,
        device=device,
        reverse=True,
    )

    param_dict["best_loss"] = best_loss

    predictions = predict(
        model=model,
        model_name=model_name,
        test_data=test_data,
        pred_steps=pred_steps,
        device=device,
    )

    predictions = predictions.numpy()
    for i in range(len(predictions)):
        predictions[i] = scaler.inverse_transform(predictions[i])

    for i in range(len(test_data)):
        test_data[i] = scaler.inverse_transform(test_data[i])

    make_report(
        model_name=model_name,
        param_dict=param_dict,
        losses=losses,
        test_data=test_data,
        predictions=predictions,
        plot_range=plot_range,
    )


if __name__ == "__main__":
    main()