model.py

import torch.nn as nn
import torch
from torch import cuda
from torch.nn.init import xavier_normal_

device = 'cuda:0' if cuda.is_available() else 'cpu'

class BiLSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, num_layers):
        super(BiLSTMModel, self).__init__()
        self.input_dim = input_size
        self.hidden_dim = hidden_size
        self.num_layers = num_layers

        self.bilstm = nn.LSTM(self.input_dim, self.hidden_dim, self.num_layers, batch_first=True, bidirectional=True)
        # self.fc = nn.Linear(self.hidden_dim * 2, output_size)

    def forward(self, x):
        # print("x shape", x.shape) # [batch_size, sequence_length, input_dim]

        batch_size = x.size(0)
        num_directions = 2
        h_0 = torch.zeros(self.num_layers * num_directions, batch_size, self.hidden_dim).to(device)
        c_0 = torch.zeros(self.num_layers * num_directions, batch_size, self.hidden_dim).to(device)

        # Pass through BiLSTM
        lstm_out, (hidden, _) = self.bilstm(x, (h_0, c_0))

        # print(hidden.shape)  [4, 16, 128]

        forward_hidden = hidden[-2]
        backward_hidden = hidden[-1]
        # hidden_cat = torch.cat((forward_hidden, backward_hidden), dim=1)
        hidden_avg = (forward_hidden + backward_hidden) / 2

        # output = self.fc(hidden_cat)

        # print("lstm_out shape", lstm_out.size()) # lstm_out shape: [batch_size, sequence_length, hidden_dim * num_directions]
        # print("hidden_cat shape", hidden_avg.size()) # hidden shape : [batch_size, hidden_dim]
        return hidden_avg  

class LSTM(nn.Module):

    def __init__(self, input_dim, hidden_dim, output_size, num_layers, batch_size):
        super(LSTM, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.mlp_dim = 4 * hidden_dim
        self.relu = nn.ReLU()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.linear1 = nn.Linear(self.hidden_dim, self.mlp_dim)
        self.linear2 = nn.Linear(self.mlp_dim, self.mlp_dim)
        # self.linear3 = nn.Linear(self.mlp_dim, self.mlp_dim)
        self.hidden2out = nn.Linear(self.mlp_dim, self.hidden_dim)
        self.output_layer_epi = nn.Linear(self.hidden_dim, output_size)
        self.output_layer_hyp = nn.Linear(self.hidden_dim, output_size)
    
    def init_hidden(self, batch_size=0):
        # initialize both hidden layers
        if batch_size == 0:
            batch_size = self.batch_size
            
        # Create tensors with Xavier initialization
        h0 = xavier_normal_(torch.empty(self.num_layers, batch_size, self.hidden_dim))
        c0 = xavier_normal_(torch.empty(self.num_layers, batch_size, self.hidden_dim))

        # Move the tensors to the specified device (self.device)
        h0 = h0.to(device, non_blocking=True)
        c0 = c0.to(device, non_blocking=True)
        
        return h0, c0
    
    def forward(self, x):
        # print("x shape: ", x.shape)
        # print("embedding shape", embedding.shape)
        # batch = batch.squeeze(0)
        batch_size = x.size(0)

        # h0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(device)
        # c0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(device)

        self.h0, self.c0 = self.init_hidden(batch_size)

        output, (self.h0, self.c0) = self.lstm(x, (self.h0, self.c0))
        # output, (hn, cn) = self.lstm(x, (h0, c0))
        # print("foward:: output: {}".format(output.shape))

        output = self.linear1(output)
        output = self.relu(output)
        output = self.linear2(output)
        output = self.relu(output)
        # output = self.linear3(output)
        # output = self.relu(output)
        output = self.hidden2out(output)
        outputs_epi = self.output_layer_epi(output)
        outputs_hyp = self.output_layer_hyp(output)
        # print("final output: {}".format(output))

        outputs_epi = outputs_epi.squeeze(-1)
        outputs_hyp = outputs_hyp.squeeze(-1)
        # print("outputs epi shape: {}".format(outputs_epi.shape))
        # print("outputs hyp shape: {}".format(outputs_hyp.shape))
        return outputs_epi, outputs_hyp



class LSTM_month(nn.Module):

    def __init__(self, input_dim, hidden_dim, output_size, num_layers, batch_size):
        super(LSTM_month, self).__init__()
        self.input_dim = input_dim
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        self.mlp_dim = 4 * hidden_dim
        self.relu = nn.ReLU()
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True)
        self.linear1 = nn.Linear(self.hidden_dim, self.mlp_dim)
        # self.linear2 = nn.Linear(self.mlp_dim, self.mlp_dim)
        # self.linear3 = nn.Linear(self.mlp_dim, self.mlp_dim)
        self.hidden2out = nn.Linear(self.mlp_dim, self.hidden_dim)
        self.output_layer_epi = nn.Linear(self.hidden_dim, output_size)
        self.output_layer_hyp = nn.Linear(self.hidden_dim, output_size)

        self.layer_norm = nn.LayerNorm(hidden_dim)

    def init_hidden(self, batch_size=0):
        # initialize both hidden layers
        if batch_size == 0:
            batch_size = self.batch_size
            
        # Create tensors with Xavier initialization
        h0 = xavier_normal_(torch.empty(self.num_layers, batch_size, self.hidden_dim))
        c0 = xavier_normal_(torch.empty(self.num_layers, batch_size, self.hidden_dim))

        # Move the tensors to the specified device (self.device)
        h0 = h0.to(device, non_blocking=True)
        c0 = c0.to(device, non_blocking=True)
        
        return h0, c0
    
    def forward(self, x):
        # print("x shape: ", x.shape)
        # print("embedding shape", embedding.shape)
        # batch = batch.squeeze(0)
        batch_size = x.size(0)

        # h0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(device)
        # c0 = torch.zeros(self.num_layers, batch_size, self.hidden_dim).to(device)
        self.h0, self.c0 = self.init_hidden(batch_size)

        output, (self.h0, self.c0) = self.lstm(x, (self.h0, self.c0))
        # output, (hn, cn) = self.lstm(x, (h0, c0))
        # print("foward:: output: {}".format(output.shape))]]

        output, _ = self.lstm(x)
        output = self.layer_norm(output)

        l1_norm = torch.mean(torch.abs(output[:, 1:] - output[:, :-1]))  # L1 norm across timesteps
        l2_norm = torch.mean(torch.norm(output[:, 1:] - output[:, :-1], dim=-1))  # L2 norm across timesteps

        l1_norm_scale = torch.mean(torch.abs(output[:, 1:]))  # L1 norm across timesteps
        l2_norm_scale = torch.mean(torch.norm(output[:, 1:], dim=-1))  # L2 norm across timesteps

        # print(f"L1 norm of LSTM output change: {l1_norm.item():.6f}")
        # print(f"L2 norm of LSTM output change: {l2_norm.item():.6f}")

        # print(f"L1 norm of LSTM output scale: {l1_norm_scale.item():.6f}")
        # print(f"L2 norm of LSTM output scale: {l2_norm_scale.item():.6f}")

        output = self.linear1(output)
        output = self.relu(output)
        # output = self.linear2(output)
        # output = self.relu(output)
        # output = self.linear3(output)
        # output = self.relu(output)
        output = self.hidden2out(output)
        output = self.output_layer_epi(output)
        # print("final output: {}".format(output))

        output = output.squeeze(-1)
        # print("outputs shape: {}".format(outputs_hyp.shape))
        return output


# discriminator model for single month
class Discriminator(nn.Module):
    def __init__(self, input_dim, days_per_month=30):
        """
        Discriminator model to classify whether a given month's data favors a specific model.

        Args:
            input_dim (int): Number of input features per day.
            hidden_dim (int): Number of hidden units in the fully connected layers.
            days_per_month (int): Number of days per month (default: 30).
        """
        super(Discriminator, self).__init__()
        self.days_per_month = days_per_month

        # Fully connected layers for a single month's sequence
        self.fc = nn.Sequential(
            nn.Linear(input_dim * days_per_month, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 64),
            nn.ReLU(),
            nn.Linear(64, 1),  # Binary classification
            nn.Sigmoid()  # Output probability
        )

    def forward(self, month_x):
        """
        Forward pass for the Discriminator.

        Args:
            month_x (torch.Tensor): shape (batch_size, days_per_month, input_dim).
        Returns:
            torch.Tensor: shape (batch_size, 1), probability of the model choice.
        """
        batch_size, days_per_month, input_dim = month_x.size()
        assert days_per_month == self.days_per_month, "Input must have the correct number of days per month."

        # Flatten the data for the fully connected layer
        month_x = month_x.reshape(batch_size, -1)

        # Pass through the fully connected layers
        output = self.fc(month_x)  # Shape: (batch_size, 1)
        # print("output shape for discriminator: ", output.shape)
        # output = output.squeeze(0)

        return output