base_node.py

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torchdiffeq
from matplotlib.backends.backend_pdf import PdfPages
from scipy.integrate import solve_ivp
from sklearn.preprocessing import StandardScaler
from torch.optim.lr_scheduler import StepLR, ReduceLROnPlateau

from base import split_on_change
from ltc import LTC
from utils import get_data


class NeuralODEFunc(nn.Module):
    def __init__(
        self, input_dim, hidden_dim, output_dim, num_layers, aug_dim, ltc: bool = False
    ):
        super(NeuralODEFunc, self).__init__()

        self.net = None
        self.ltc = ltc
        self.aug_dig = aug_dim

        if self.ltc:
            ltc_cell = LTC(
                input_size=input_dim,
                units=hidden_dim,
                return_sequences=True,
                mixed_memory=False,
            )

            linear = nn.Linear(hidden_dim, input_dim)

            layers = [ltc_cell, linear]

            self.net = layers

        else:
            layers = [nn.Linear(input_dim, hidden_dim), nn.Tanh()]

            for _ in range(num_layers - 1):
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                layers.append(nn.Tanh())

            layers.append(nn.Linear(hidden_dim, input_dim))

            self.net = nn.Sequential(*layers)

    def forward(self, t, y):
        if self.ltc:
            ltc = self.net[0]
            linear = self.net[1]

            out, _ = ltc(y)

            return linear(out)

        else:
            return self.net(y)


class ODEBlock(nn.Module):
    def __init__(self, odefunc):
        super(ODEBlock, self).__init__()
        self.odefunc = odefunc

    def forward(self, y0, t):
        # sol = torchdiffeq.odeint_adjoint(self.odefunc, y0, t)
        if self.odefunc.aug_dig > 0:
            # Add augmentation
            aug = torch.zeros(y0.shape[0], self.odefunc.aug_dig)
            # Shape (batch_size, data_dim + augment_dim)
            x_aug = torch.cat([y0, aug], 1)
        else:
            x_aug = y0

        sol = torchdiffeq.odeint(self.odefunc, x_aug, t)
        return sol


class NeuralODE(nn.Module):
    def __init__(self, odefunc, input_dim, output_dim, aug_dim):
        super(NeuralODE, self).__init__()
        self.odeblock = ODEBlock(odefunc)
        self.fc = nn.Linear(input_dim + aug_dim, output_dim)

    def forward(self, y0, t):
        out = self.odeblock(y0, t)
        out = self.fc(out)
        return out


def predict(model, x_test, t):
    model.eval()
    with torch.no_grad():
        pred = model(x_test, t)

    return pred.view(7, 2500, 3).numpy()


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

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

        fig, ax = plt.subplots(3, sharex="col", figsize=(8.27, 11.69))
        criterion = nn.MSELoss()
        rmse = torch.sqrt(
            criterion(torch.tensor(test_data[i]), torch.tensor(predictions[i]))
        )

        fig.suptitle(
            f"{test_data[i][0, -1]}, RMSE: {rmse:.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)

    return plots


def make_report(
    model_name: str,
    param_dict: dict,
    prediction: 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_full_plots(test_data, prediction, plot_range)

    plt.figure()
    plt.plot(
        test_data[3][:, 0],
        test_data[3][:, 1],
        label="Truth",
    )
    plt.plot(prediction[3][:, 0], prediction[3][:, 1], label="Prediction")
    plt.xlabel("x_1")
    plt.ylabel("x_2")
    plt.legend()
    plt.grid(True)
    q1_q2 = plt.gcf()

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

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


def assert_test_data(train_data, splitted_test_data):
    train_data = torch.tensor(train_data, dtype=torch.float32)
    splitted_test_data = torch.tensor(splitted_test_data, dtype=torch.float32)

    x_test = torch.zeros(7, 3)
    x_test[0] = splitted_test_data[0][0, :]
    x_test[1] = splitted_test_data[1][0, :]
    x_test[2] = train_data[0, :]
    x_test[3] = splitted_test_data[2][0, :]
    x_test[4] = train_data[2500, :]
    x_test[5] = splitted_test_data[3][0, :]
    x_test[6] = splitted_test_data[4][0, :]

    y_test = torch.zeros(7, 2500, 3)
    y_test[0] = splitted_test_data[0]
    y_test[1] = splitted_test_data[1]
    y_test[2] = train_data[0:2500, :]
    y_test[3] = splitted_test_data[2]
    y_test[4] = train_data[2500:, :]
    y_test[5] = splitted_test_data[3]
    y_test[6] = splitted_test_data[4]

    return x_test, y_test


def create_batch(num_timesteps, num_times_per_obs, mini_batch_size, x, out_dim):

    s = np.random.permutation(num_timesteps - num_times_per_obs + 1)[:mini_batch_size]

    x0_train = x[s, :]

    targets = np.zeros((num_times_per_obs, mini_batch_size, out_dim))

    for i, start_index in enumerate(s):
        if out_dim == 2:
            targets[:, i, (0, 1)] = x[
                start_index : start_index + num_times_per_obs, (0, 1)
            ]
        else:
            targets[:, i, :] = x[start_index : start_index + num_times_per_obs, :]

    return torch.tensor(x0_train, dtype=torch.float32), torch.tensor(
        targets, dtype=torch.float32
    )


def create_full_batch(num_times_per_obs, mini_batch_size, x):
    if x.shape[0] > 2500:

        i1 = np.arange(0, 2500 - num_times_per_obs)
        i2 = np.arange(2500, 5000 - num_times_per_obs)
        indices = np.hstack((i1, i2))
        perm_indices = np.random.permutation(indices)[:mini_batch_size]

        x_full_train = x[perm_indices, :]

        targets = np.zeros((num_times_per_obs, mini_batch_size, x.shape[1]))
        for i, start_index in enumerate(perm_indices):
            targets[:, i, :] = x[start_index : start_index + num_times_per_obs, :]

        return torch.tensor(x_full_train, dtype=torch.float32), torch.tensor(
            targets, dtype=torch.float32
        )

    else:
        print(ValueError("Please use create batch if only one time series is used."))


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

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

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

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

    ax[2].plot(test_data[1:plot_range, 2], label="True Values", marker="o")
    ax[2].plot(predictions[:plot_range, 2], label="Predictions", marker="x")
    ax[2].set(title=f"f_a: {test_data[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)

    return plots


def make_full_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():

    train_data = get_data(r"C:\Users\micha\Desktop\Programming\LTC\DORA_Train.csv")
    test_data = get_data(r"C:\Users\micha\Desktop\Programming\LTC\DORA_Test.csv")
    t = train_data[:2500, 0]
    dt = t[1]
    t = torch.tensor(t, dtype=torch.float32)

    train_data = train_data[:, 1:4]

    splitted_test = split_on_change(test_data[:, 1:5])
    x_test, y_test = assert_test_data(train_data, splitted_test)
    train_data = train_data[:2500]
    val_data = splitted_test[2]
    y0_val = val_data[0]

    neural_ode_timesteps = 5
    input_dim = 3
    hidden_dim = 20
    output_dim = 2
    num_layers = 2
    model_name = "try_node_h20_lay2_ts5_rmse_b500_e300k"
    num_epochs = 10
    batch_size = 500
    lr = 0.001
    timesteps = np.arange(0, neural_ode_timesteps) * dt

    param_dict = {
        "Model Name": model_name,
        "Architecture": "NODE",
        "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",
        "Scaler": "None",
    }

    odefunc = NeuralODEFunc(
        input_dim, hidden_dim, output_dim, num_layers, aug_dim=0, ltc=False
    )
    neural_ode = NeuralODE(odefunc, input_dim, output_dim, aug_dim=0)

    # neural_ode.load_state_dict(
    #     torch.load("models/node_h2_lay2_ts5_rmse_b500_e500k.pth")
    # )

    criterion = nn.MSELoss()
    optimizer = optim.Adam(neural_ode.parameters(), lr=lr)
    scheduler = ReduceLROnPlateau(patience=50000, factor=0.1, optimizer=optimizer)

    losses = []

    best_loss = 1e20
    best_epoch = 0
    for epoch in range(num_epochs):
        neural_ode.train()
        train_x, train_y = create_batch(
            num_timesteps=train_data.shape[0],
            num_times_per_obs=neural_ode_timesteps,
            mini_batch_size=batch_size,
            x=train_data,
            out_dim=2,
        )
        # train_x, train_y = create_full_batch(
        #     num_times_per_obs=neural_ode_timesteps,
        #     mini_batch_size=batch_size,
        #     x=train_data,
        # )

        optimizer.zero_grad()
        output = neural_ode(train_x, t[:neural_ode_timesteps])
        loss = torch.sqrt(criterion(output, train_y))
        loss.backward()
        optimizer.step()
        # scheduler.step(loss)

        losses.append(loss.item())

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

        print(
            f"Epoch {epoch}, Loss: {loss.item():.8f}, Best Loss: {best_loss:.8f}, Best Epoch: {best_epoch}"
        )

    plt.figure(figsize=(10, 5))
    plt.plot(range(1, num_epochs + 1), losses, label="Training Loss")
    plt.title("Training Loss Curve")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.yscale("log")
    plt.legend()
    plt.grid(True)
    plt.show()

    model = NeuralODE(odefunc, input_dim, output_dim, aug_dim=0)
    model.load_state_dict(torch.load(f"models/{model_name}.pth"))

    print(f"Best Loss: {best_loss}")

    param_dict["Best Loss"] = best_loss
    param_dict["Best Epoch"] = best_epoch

    predictions = predict(model=model, x_test=x_test, t=t)

    make_report(
        model_name=model_name,
        test_data=y_test.numpy(),
        prediction=predictions,
        param_dict=param_dict,
        plot_range=300,
    )


if __name__ == "__main__":
    main()