seq_lstm.py

import typing

import numpy as np
import torch

from utils import get_data, plot_data, create_sequences
import torch.nn as nn
import torch.optim as optim
from models import LSTMModel, train_model, predict
from torch.utils.data import DataLoader, Dataset
from sklearn.preprocessing import StandardScaler, MinMaxScaler
import matplotlib.pyplot as plt


class TimeSeriesDataset(Dataset):
    def __init__(self, data):
        self.data = torch.tensor(data, dtype=torch.float32)

    def __len__(self):
        return len(self.data) - 1

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.data[idx + 1]
        return x, y


class TimeSeriesSequenceDataset(Dataset):
    def __init__(self, sequences, targets):
        self.sequences = sequences
        self.targets = targets

    def __len__(self):
        return len(self.sequences)

    def __getitem__(self, idx):
        return self.sequences[idx], self.targets[idx]


class LSTMM(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers, dropout=0.0):
        super(LSTMM, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.dropout = dropout

        self.lstm = nn.LSTM(
            input_dim, hidden_dim, num_layers, batch_first=True, dropout=self.dropout
        )

        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):

        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)

        out, _ = self.lstm(x, (h0, c0))

        out = self.fc(out[:, -1, :])
        return out


def main():
    batch_size = 128
    num_epochs = 2
    input_dim = 3
    hidden_dim = 64
    num_layers = 3
    output_dim = 3
    seq_length = 7
    torch.manual_seed(42)

    train_data = get_data(
        "DORA_Train.csv"
    )
    test_data = get_data(
        "DORA_Test.csv"
    )
    train_scaler = StandardScaler()
    train_scaled = train_scaler.fit_transform(train_data[:, 1:4])

    sequences, targets = create_sequences(train_scaled, seq_length)
    seq_dataset = TimeSeriesSequenceDataset(sequences, targets)
    seq_loader = DataLoader(seq_dataset, batch_size=batch_size, shuffle=True)

    train_dataset = TimeSeriesDataset(train_scaled[:2500])
    test_dataset = TimeSeriesDataset(test_data)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    model = LSTMM(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers, dropout=0.0)

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

    losses = []

    best_loss = 1e20
    for epoch in range(num_epochs):
        epoch_loss = 0.0
        mean_norm = []
        for i, (inputs, targets) in enumerate(seq_loader):

            inputs = inputs
            targets = targets

            outputs = model(inputs)

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

            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)
            # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

            optimizer.step()

        epoch_loss /= len(train_loader)

        if epoch_loss < best_loss:
            best_loss = epoch_loss
            torch.save(model.state_dict(), "best_seq_lstm.pth")
            print(f"New Best Loss: {best_loss}")

        losses.append(epoch_loss)
        print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss}, Gradient: {np.array(mean_norm).mean():.4f}")

    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()

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

    predictions = []

    best_model = LSTMM(input_dim=input_dim, hidden_dim=hidden_dim, output_dim=output_dim, num_layers=num_layers, dropout=0.0)
    best_model.load_state_dict(torch.load("best_seq_lstm.pth"))
    best_model.eval()

    initial_input = torch.tensor(train_scaler.transform(train_scaled[:seq_length]), dtype=torch.float32).unsqueeze(0)

    num_predictions = 250

    with torch.no_grad():
        for _ in range(num_predictions):

            output = best_model(initial_input)

            predictions.append(output.squeeze(0))

            initial_input = torch.cat((initial_input[:, 1:, :], output.unsqueeze(0)), dim=1)


    predictions = torch.stack(predictions)
    predictions = train_scaler.inverse_transform(predictions)

    # Plotting
    plt.figure(figsize=(10, 6))
    plt.plot(train_data[1:num_predictions, 1], label="True Values", marker="o")
    plt.plot(predictions[:, 0], label="Predictions", marker="x")
    plt.xlabel("Timestep")
    plt.ylabel("Value")
    plt.title("q1")
    plt.legend()
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(10, 6))
    plt.plot(train_data[1:num_predictions, 2], label="True Values", marker="o")
    plt.plot(predictions[:, 1], label="Predictions", marker="x")
    plt.xlabel("Timestep")
    plt.ylabel("Value")
    plt.title("q2")
    plt.legend()
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(10, 6))
    plt.plot(train_data[1:num_predictions, 3], label="True Values", marker="o")
    plt.plot(predictions[:, 2], label="Predictions", marker="x")
    plt.xlabel("Timestep")
    plt.ylabel("Value")
    plt.title("fa_t")
    plt.legend()
    plt.grid(True)
    plt.show()


if __name__ == "__main__":
    main()