seq_node.py

import torch

from ltc_node.node_node import *
from ltc_node.ltc_utils import *
from ltc_node.ltc_models import *

import time


class LTCODESEQ(nn.Module):
    def __init__(
        self,
        units,
        input_size,
        output_size,
        wiring1,
        wiring2,
        ode_unfolds=6,
        epsilon=1e-8,
    ):
        super(LTCODESEQ, self).__init__()
        self.units = units
        self.input_dim = input_size
        self.output_dim = output_size
        self.ode_unfolds = ode_unfolds
        self.epsilon = epsilon
        self.wiring1 = wiring1
        self.wiring2 = wiring2

        self.rnn_cell1 = LTCCell(
            wiring=self.wiring1,
            in_features=self.input_dim,
            input_mapping="affine",
            output_mapping="affine",
            ode_unfolds=self.ode_unfolds,
            epsilon=self.epsilon,
            implicit_param_constraints=True,
        )

        self.rnn_cell2 = LTCCell(
            wiring=self.wiring2,
            in_features=1,
            input_mapping="affine",
            output_mapping="affine",
            ode_unfolds=self.ode_unfolds,
            epsilon=self.epsilon,
            implicit_param_constraints=True,
        )
        self.linear1 = nn.Linear(16, 64)
        self.linear2 = nn.Linear(64, 16)
        self.linear3 = nn.Linear(16, 1)
        self.act = nn.ReLU()

    def forward(self, inputs, timespan, h_state=None, ts=0.1):
        if not h_state:
            h_state = torch.zeros((inputs.shape[0], self.wiring1.units))

        h_state2 = torch.zeros((inputs.shape[0], self.wiring2.units))

        outputs1 = torch.empty((timespan, inputs.shape[0], 1))
        outputs2 = torch.empty((timespan, inputs.shape[0], 1))
        outputs1[0] = inputs[:, 1].view(-1, 1)
        for t in range(1, timespan):
            h_out, h_state = self.rnn_cell1.forward(inputs, h_state, ts)
            inputs = h_out
            outputs1[t, :, :] = h_out

            h_out2 = self.act(self.linear1(h_state))
            h_out2 = self.act(self.linear2(h_out2))
            h_out2 = self.linear3(h_out2)

            h_out2, h_state2 = self.rnn_cell2.forward(h_out2, h_state2, ts)
            outputs2[t, :, :] = h_out2

        # h_out2 = self.act(self.linear1(h_state))
        # h_out2 = self.linear2(h_out2)
        #
        # for t in range(1, timespan):
        #
        #     h_out2, h_state2 = self.rnn_cell2.forward(h_out2, h_state2, ts)
        #     outputs2[t, :, :] = h_out2

        outputs = torch.cat((outputs2, outputs1), dim=2)

        return outputs


class DataLoaderSeq:
    def __init__(self, x, num_times_per_obs, mini_batch_size):
        self.x = x
        self.num_timesteps = x.shape[0]
        self.num_times_per_obs = num_times_per_obs
        self.mini_batch_size = mini_batch_size
        self.indices = np.arange(self.num_timesteps - num_times_per_obs + 1)
        self.current_index = 0
        np.random.shuffle(self.indices)

    def __iter__(self):
        return self

    def __next__(self):
        if self.current_index >= len(self.indices):
            self.current_index = 0
            np.random.shuffle(self.indices)
            raise StopIteration

        s = self.indices[self.current_index : self.current_index + self.mini_batch_size]
        self.current_index += self.mini_batch_size

        x0_train = self.x[s, :]

        targets = np.zeros((self.num_times_per_obs, len(s), self.x.shape[1] - 1))

        for i, start_index in enumerate(s):
            targets[:, i, :] = self.x[
                start_index : start_index + self.num_times_per_obs, :2
            ]

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

    def __len__(self):
        return int(np.ceil(len(self.indices) / self.mini_batch_size))


class FullDataLoaderSeq:
    def __init__(self, x, num_times_per_obs, mini_batch_size):
        self.x = x
        self.num_timesteps = x.shape[0]
        self.num_times_per_obs = num_times_per_obs
        self.mini_batch_size = mini_batch_size
        self.half_timesteps = self.num_timesteps // 2

        # Create separate indices for each datapoint
        self.indices1 = np.arange(self.half_timesteps - num_times_per_obs + 1)
        self.indices2 = np.arange(
            self.half_timesteps, self.num_timesteps - num_times_per_obs + 1
        )

        self.current_index1 = 0
        self.current_index2 = 0

        np.random.shuffle(self.indices1)
        np.random.shuffle(self.indices2)

    def __iter__(self):
        return self

    def __next__(self):
        if self.current_index1 >= len(self.indices1) and self.current_index2 >= len(
            self.indices2
        ):
            self.current_index1 = 0
            self.current_index2 = 0
            np.random.shuffle(self.indices1)
            np.random.shuffle(self.indices2)
            raise StopIteration

        half_batch_size = self.mini_batch_size // 2

        s1 = self.indices1[self.current_index1 : self.current_index1 + half_batch_size]
        s2 = self.indices2[self.current_index2 : self.current_index2 + half_batch_size]

        self.current_index1 += half_batch_size
        self.current_index2 += half_batch_size

        s = np.concatenate((s1, s2))
        np.random.shuffle(s)  # Shuffle to mix datapoints from both halves

        x0_train = self.x[s, :]

        targets = np.zeros((self.num_times_per_obs, len(s), 2))

        for i, start_index in enumerate(s):
            targets[:, i, :] = self.x[
                start_index : start_index + self.num_times_per_obs, :2
            ]

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

    def __len__(self):
        return int(np.ceil(len(self.indices1) / (self.mini_batch_size // 2)))


def predict(model, x_test, t):
    model.eval()
    predictions = torch.empty((x_test.shape[0], t, 2))
    with torch.no_grad():
        for i in range(x_test.shape[0]):
            pred = model(x_test[i].view(1, 3), t)
            predictions[i] = pred.squeeze(1)

    return predictions


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

    for i in range(test_data.shape[0]):
        fig, ax = plt.subplots(2, sharex="col", figsize=(8.27, 11.69))
        criterion = nn.MSELoss()
        rmse = torch.sqrt(
            criterion(
                torch.tensor(test_data[i, :, :2], dtype=torch.float32), predictions[i]
            )
        )

        fig.suptitle(
            f"{test_data[i][0, -1]}, RMSE: {rmse:.5f}",
            fontsize=16,
        )

        ax[0].plot(test_data[i, :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, :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()

        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,
    losses=None,
):
    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()

    plt.figure()
    plt.plot(range(1, len(losses) + 1), losses, label="Training Loss")
    plt.title("Training Loss Curve")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.yscale("log")
    plt.legend()
    plt.grid(True)
    loss_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/{model_name}_eval.pdf") as pdf:
        pdf.savefig(dict_page)
        pdf.savefig(loss_page)
        pdf.savefig(q1_q2)

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


def main():
    units = 16
    input_size = 3
    output_size = 3
    ode_unfolds = 6
    timespan = 5
    epsilon = 1e-8
    epochs = 10000
    val_range = 500

    batch_size = 1000
    lr = 0.005

    model_name = f"comp_seq_ltc_node_u{units}_ode{ode_unfolds}_ts{timespan}_e{epochs}_bs{batch_size}_lr{lr}"

    param_dict = {
        "Model Name": model_name,
        "Architecture": "NODE",
        "Input Dim": input_size,
        "Output Dim": output_size,
        "Units": units,
        "ODE Unfolds": ode_unfolds,
        "Timespan": timespan,
        "Batch Size": batch_size,
        "Val Range": val_range,
        "Epochs": epochs,
        "Learning Rate": lr,
        "Loss": "RMSE",
        "Scaler": "None",
    }

    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_full = torch.tensor(
        np.linspace(0, train_data[2499, 0], 2500), dtype=torch.float32
    )
    t = torch.tensor(train_data[:timespan, 0], dtype=torch.float32)
    train_data = train_data[:, 1:4]
    val_x0 = torch.tensor(train_data[0], dtype=torch.float32).view(1, 3)
    val_x1 = torch.tensor(train_data[2500], dtype=torch.float32).view(1, 3)

    splitted_test = split_on_change(test_data[:, 1:5])
    x_test, y_test = assert_test_data(train_data, splitted_test)

    data_loader = FullDataLoaderSeq(train_data, timespan, batch_size)
    # data = train_data
    # train_data = train_data[:, :]
    # data_loader = DataLoaderSeq(train_data, timespan, batch_size)

    wiring1 = FullyConnected(units, 1)
    wiring2 = FullyConnected(units, 1)

    ltc_model = LTCODESEQ(
        input_size=input_size,
        units=units,
        output_size=output_size,
        ode_unfolds=ode_unfolds,
        epsilon=epsilon,
        wiring1=wiring1,
        wiring2=wiring2,
    )

    criterion = nn.MSELoss()
    optimizer = optim.Adam(ltc_model.parameters(), lr=lr)

    losses = []
    val_losses = []
    times = []
    tpe = []
    best_loss = 1e20
    best_epoch = 0
    best_val_loss = 1e20
    best_val_epoch = 0
    best_val_time = 0

    start_time = time.time()

    for epoch in range(epochs):
        ltc_model.train()
        epoch_loss = 0
        epoch_time = time.time()
        for batch in data_loader:
            x0_train, targets = batch

            optimizer.zero_grad()
            output = ltc_model(x0_train, timespan)
            loss = torch.sqrt(criterion(output, targets))
            epoch_loss += loss.item()
            loss.backward()
            optimizer.step()

        losses.append(epoch_loss / len(data_loader))
        epoch_loss = epoch_loss / len(data_loader)

        tpe.append(time.time() - epoch_time)

        # ltc_model.eval()
        # with torch.no_grad():
        #     val1 = ltc_model(val_x0, val_range).squeeze(1)
        #     val2 = ltc_model(val_x1, val_range).squeeze(1)
        #
        # val_loss1 = torch.sqrt(
        #     criterion(
        #         val1, torch.tensor(train_data[:val_range, :2], dtype=torch.float32)
        #     )
        # )
        # val_loss2 = torch.sqrt(
        #     criterion(
        #         val2,
        #         torch.tensor(
        #             train_data[2500 : 2500 + val_range, :2], dtype=torch.float32
        #         ),
        #     )
        # )
        # val_loss = (val_loss1.item() + val_loss2.item()) / 2
        # val_losses.append(val_loss)
        #
        # if val_loss < best_val_loss:
        #     best_val_loss = val_loss
        #     best_val_epoch = epoch
        #     torch.save(ltc_model.state_dict(), f"models/val_{model_name}.pth")
        #     print(f"New Best Val Loss: {best_val_loss}")
        #     best_val_time = time.time()
        #     best_time = best_val_time - start_time

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

        e_time = time.time() - epoch_time
        times.append(e_time)

        # print(
        #     f"Epoch {epoch}, Loss: {epoch_loss:.8f}, Val Loss:{val_loss:.8f}, Best Loss: {best_loss:.8f}, Best Epoch: {best_epoch}, Best Val Loss: {best_val_loss:.8f}, Best Val Epoch: {best_val_epoch}, Time: {e_time}"
        # )

        print(
            f"Epoch {epoch}, Loss: {epoch_loss:.8f}, Best Loss: {best_loss:.8f}, Best Epoch: {best_epoch}, Time: {e_time}"
        )

    comp_time = time.time() - start_time

    plt.figure(figsize=(10, 5))
    plt.plot(range(1, epochs + 1), losses, label="Training Loss")
    # plt.plot(range(1, epochs + 1), val_losses, label="Validation Loss")
    plt.title("Training Loss Curve")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.yscale("log")
    plt.legend()
    plt.grid(True)
    plt.show()
    # np.savetxt(
    #     f"losses/loss_{model_name}.txt",
    #     np.stack((np.array(losses), np.array(val_losses)), axis=1),
    # )
    np.savetxt(f"losses/val_loss_{model_name}.txt", np.array(val_losses))

    model = LTCODESEQ(
        input_size=input_size,
        units=units,
        output_size=output_size,
        ode_unfolds=ode_unfolds,
        epsilon=epsilon,
        wiring1=wiring1,
        wiring2=wiring2,
    )
    model.load_state_dict(torch.load(f"models/{model_name}.pth"))
    # val_model = LTCODESEQ(
    #     input_size=input_size,
    #     units=units,
    #     output_size=output_size,
    #     ode_unfolds=ode_unfolds,
    #     epsilon=epsilon,
    #     wiring1=wiring1,
    #     wiring2=wiring2,
    # )
    # val_model.load_state_dict(torch.load(f"models/val_{model_name}.pth"))

    param_dict["Best Loss"] = best_loss
    # param_dict["Best Val Loss"] = best_val_loss
    param_dict["Best Epoch"] = best_epoch
    # param_dict["Best Val Epoch"] = best_val_epoch

    if comp_time:
        param_dict["Duration"] = comp_time
        # param_dict["Best Time"] = best_time
        param_dict["Average Epoch Time"] = np.sum(times) / len(times)
        param_dict["TpE"] = np.sum(tpe) / len(tpe)

    train_data = get_data(r"C:\Users\micha\Desktop\Programming\LTC\DORA_Train.csv")[
        :, 1:4
    ]
    test_data = get_data(r"C:\Users\micha\Desktop\Programming\LTC\DORA_Test.csv")[
        :, 1:5
    ]
    splitted_test = split_on_change(test_data)
    x_test, y_test = assert_test_data(train_data, splitted_test)

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

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


if __name__ == "__main__":
    main()