is there a way to implement the TimeRepresentationAgent into this iteration of my model?
import torch
import torch.nn as nn
import numpy as np

Chomp1d remains unchanged
class Chomp1d(nn.Module):
def init(self, chomp_size):
super(Chomp1d, self).init()
self.chomp_size = chomp_size

def forward(self, x):
if self.chomp_size > 0:
return x[:, :, :-self.chomp_size]
else:
return x
Residual Block remains unchanged
class ResidualBlock(nn.Module):
def init(self, in_channels, out_channels, kernel_size, dilation, dropout):
super(ResidualBlock, self).init()
padding = (kernel_size - 1) * dilation


self.conv1 = nn.utils.weight_norm(
    nn.Conv1d(in_channels, out_channels, kernel_size,
              padding=padding, dilation=dilation)
)
self.chomp1 = Chomp1d(padding)
self.relu1 = nn.ReLU()
self.dropout1 = nn.Dropout(dropout)

self.conv2 = nn.utils.weight_norm(
    nn.Conv1d(out_channels, out_channels, kernel_size,
              padding=padding, dilation=dilation)
)
self.chomp2 = Chomp1d(padding)
self.relu2 = nn.ReLU()
self.dropout2 = nn.Dropout(dropout)

# Adjust with a 1x1 convolution if in_channels != out_channels
self.downsample = nn.Conv1d(in_channels, out_channels, kernel_size=1) \
                  if in_channels != out_channels else None

self.relu = nn.ReLU()
self.init_weights()
def init_weights(self):
nn.init.normal_(self.conv1.weight, std=0.01)
nn.init.normal_(self.conv2.weight, std=0.01)
if self.downsample is not None:
nn.init.normal_(self.downsample.weight, std=0.01)

def forward(self, x):
residual = x # Save the input for the residual connection
out = self.conv1(x)
out = self.chomp1(out) # Remove extra padding
out = self.relu1(out)
out = self.dropout1(out)
out = self.conv2(out)
out = self.chomp2(out) # Remove extra padding
out = self.relu2(out)
out = self.dropout2(out)
if self.downsample is not None:
residual = self.downsample(residual)
return self.relu(out + residual) # Residual connection with addition
class TCNEncoder(nn.Module):
def init(self, input_size, hidden_size, num_layers, kernel_size, dropout):
super(TCNEncoder, self).init()
layers = []
in_channels = input_size


for i in range(num_layers):
    dilation = 2 ** i
    out_channels = hidden_size
    layers.append(ResidualBlock(in_channels, out_channels, kernel_size, dilation, dropout))
    in_channels = out_channels  # For next layer

self.network = nn.Sequential(*layers)
def forward(self, x):
# x shape: (batch_size, seq_length, input_size)
x = x.transpose(1, 2) # Change to (batch_size, input_size, seq_length)
x = self.network(x)
x = x.transpose(1, 2) # Back to (batch_size, seq_length, hidden_size)
return x
class TransformerDecoder(nn.Module):
def init(self, input_size, hidden_size, num_layers, num_heads, dropout):
super(TransformerDecoder, self).init()
self.embedding = nn.Linear(input_size, hidden_size)
self.positional_encoding = PositionalEncoding(hidden_size, dropout)
decoder_layer = nn.TransformerDecoderLayer(d_model=hidden_size, nhead=num_heads, dropout=dropout)
self.transformer_decoder = nn.TransformerDecoder(decoder_layer, num_layers=num_layers)

self.fc_out = nn.Linear(hidden_size, input_size) # Not needed here
def forward(self, tgt, memory, tgt_mask):
# tgt: (batch_size, tgt_seq_length, input_size)
tgt_emb = self.embedding(tgt)
tgt_emb = self.positional_encoding(tgt_emb)
memory = memory.transpose(0, 1) # (src_seq_length, batch_size, hidden_size)
tgt_emb = tgt_emb.transpose(0, 1) # (tgt_seq_length, batch_size, hidden_size)
output = self.transformer_decoder(tgt_emb, memory, tgt_mask=tgt_mask)
output = output.transpose(0, 1) # Back to (batch_size, tgt_seq_length, hidden_size)
return output
#parallel model
class TCNEncoderTransformerDecoderEnsemble(nn.Module):
def init(self, input_size, hidden_size, num_layers, kernel_size,
num_heads, dropout, subset_size):
super(TCNEncoderTransformerDecoderEnsemble, self).init()


self.input_size = input_size
self.subset_size = subset_size
self.hidden_size = hidden_size

# Calculate the number of feature subsets
self.num_subsets = (input_size + subset_size - 1) // subset_size  # Ceiling division

# Create sub-models for each subset
self.encoders = nn.ModuleList()
self.decoders = nn.ModuleList()
self.attention_layers = nn.ModuleList()
self.final_fc_layers = nn.ModuleList()

# Shared positional encoding
self.positional_encoding = PositionalEncoding(hidden_size, dropout)
self.dropout = nn.Dropout(dropout)

# Separate alpha parameters for each subset (optional)
self.alphas = nn.ParameterList([nn.Parameter(torch.full((hidden_size,), 0.5)) for _ in range(self.num_subsets)])

# Create mini-models for each feature subset
for i in range(self.num_subsets):
    # Calculate the size of the current subset
    current_input_size = min(subset_size, input_size - i * subset_size)

    # Initialize sub-models
    encoder = TCNEncoder(current_input_size, hidden_size, num_layers, kernel_size, dropout)
    decoder = TransformerDecoder(current_input_size, hidden_size, num_layers, num_heads, dropout)
    attention_layer = nn.MultiheadAttention(embed_dim=hidden_size, num_heads=num_heads, dropout=dropout)
    final_fc_layer = nn.Linear(hidden_size, current_input_size)

    self.encoders.append(encoder)
    self.decoders.append(decoder)
    self.attention_layers.append(attention_layer)
    self.final_fc_layers.append(final_fc_layer)
def forward(self, src, tgt):
# src: (batch_size, src_seq_length, input_size)
# tgt: (batch_size, tgt_seq_length, input_size)


batch_size = src.size(0)
src_seq_len = src.size(1)
tgt_seq_len = tgt.size(1)

# Split the source and target into subsets
src_splits = torch.chunk(src, self.num_subsets, dim=2)
tgt_splits = torch.chunk(tgt, self.num_subsets, dim=2)

# Initialize lists to store outputs
encoder_outputs = []
decoder_outputs = []
attn_outputs = []
final_outputs = []

# Process all subsets in parallel using list comprehensions
for i in range(self.num_subsets):
    # Extract feature subsets
    src_sub = src_splits[i]
    tgt_sub = tgt_splits[i]

    # Get the corresponding sub-models
    encoder = self.encoders[i]
    decoder = self.decoders[i]
    attention_layer = self.attention_layers[i]
    final_fc_layer = self.final_fc_layers[i]
    alpha = torch.sigmoid(self.alphas[i]).to(src.device).view(1, 1, -1)

    # Encoder
    memory = encoder(src_sub)  # (batch_size, src_seq_len, hidden_size)
    memory = self.positional_encoding(memory)  # Add positional encoding

    # Decoder
    tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt_sub.size(1)).to(tgt_sub.device)
    decoder_output = decoder(tgt_sub, memory, tgt_mask)  # (batch_size, tgt_seq_len, hidden_size)

    # Reshape for MultiheadAttention
    decoder_output_t = decoder_output.transpose(0, 1)  # (tgt_seq_len, batch_size, hidden_size)
    memory_t = memory.transpose(0, 1)  # (src_seq_len, batch_size, hidden_size)

    # Additional Attention Layer
    attn_output, _ = attention_layer(decoder_output_t, memory_t, memory_t)
    attn_output = attn_output.transpose(0, 1)  # (batch_size, tgt_seq_len, hidden_size)

    # Fusion Strategy 2: Weighted Sum
    combined = alpha * decoder_output + (1 - alpha) * attn_output
    combined = self.dropout(combined)

    # Final Output Layer
    final_output_sub = final_fc_layer(combined)  # (batch_size, tgt_seq_len, subset_size)
    final_outputs.append(final_output_sub)

# Concatenate outputs along the feature dimension
final_output = torch.cat(final_outputs, dim=2)  # (batch_size, tgt_seq_len, input_size)

return final_output
Function to train the model with teacher forcing
def train_model(model, train_loader, val_loader, num_epochs, learning_rate, gradient_clip_val, train_only = False):
criterion = nn.MSELoss()
optimizer = torch.optim.AdamW([
{'params': model.parameters(), 'weight_decay' : 0.01},
#{'params': criterion.parameters(), 'weight_decay' : 0.0}
], lr=learning_rate)
processor_type= 'cuda' if torch.cuda.is_available() else 'cpu'
scaler = GradScaler()
device = torch.device(processor_type)
model.to(device)

#scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=learning_rate*2, steps_per_epoch=len(train_loader), epochs=num_epochs)

Track best validation loss for early stopping
best_val_loss = float('inf')
patience = 10
patience_counter = 0
pbar = tqdm(range(num_epochs), desc=f"Epochs", unit="epoch")
for epoch in pbar:
model.train()
train_loss = 0.0
train_sMAPE = 0.0
train_MASE = 0.0
for batch_x, batch_y in train_loader:


    batch_x, batch_y = _prepare_batch(batch_x, batch_y, device)
    # Prepare decoder input by shifting targets (teacher forcing)
    decoder_input = _prepare_decoder_input(batch_x, batch_y, True)

    optimizer.zero_grad(set_to_none=True)
    
    with autocast(device_type= processor_type):
        output = model(batch_x, decoder_input)
        loss = criterion(output, batch_y)
        batch_sMAPE = smape(batch_y, output)
        batch_MASE = 0#mase(batch_y, output, torch.tensor(train_data_Weekly.values, dtype=torch.float32).to(device))
    # Scale the loss and call backward
    
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), gradient_clip_val)
    scaler.step(optimizer)
    scaler.update()
    train_loss += loss.item()
    train_sMAPE += batch_sMAPE
    train_MASE += batch_MASE
    #scheduler.step()
    
train_loss /= len(train_loader)
train_sMAPE /= len(train_loader)
train_MASE /= len(train_loader)

if not(train_only):
    Flag, best_val_loss, losses = validate_model(model, val_loader, criterion, best_val_loss, processor_type, patience_counter, patience)
    val_loss, val_sMAPE, val_MASE = losses

    #break if patience runs out
    if not(Flag):
        break
    
# Log to wandb
wandb.log({"Epoch": epoch + 1, "Train Loss": train_loss, "Train sMAPE": train_sMAPE, "Train MASE": train_MASE, 
           "Validation Loss": val_loss, "Validation sMAPE": val_sMAPE, "Validation MASE": val_MASE,
           #"Learning Rate": scheduler.get_last_lr()[0]
           #"Loss alpha": criterion.logits[0],"Loss beta": criterion.logits[1],"Loss gamma": criterion.logits[2]
           #,"Loss dealta": criterion.logits[3],"Loss lambda": criterion.logits[4]
          })
# Update progress bar
pbar.set_description(f"Epoch {epoch + 1}/{num_epochs}" 
+f"- Average Loss: {val_loss:.4f}" 
#+f"- LR: {scheduler.get_last_lr()[0]:.6f}"
)
Finalize wandb logging
#wandb.finish()
return model
def validate_model(model, val_loader, criterion, best_val_loss, processor_type, patience_counter, patience):
model.eval()
val_loss = 0.0
val_sMAPE = 0.0
val_MASE = 0.0
with torch.no_grad():
for batch_x, batch_y in val_loader:
if batch_y.dim() == 1:
batch_x = batch_x.to(device).unsqueeze(-1)
batch_y = batch_y.to(device).unsqueeze(-1)
elif batch_y.dim() > 1:
batch_x = batch_x.to(device)
batch_y = batch_y.to(device)
else:
print("The input doesn't fit the expected shape of a uni or multi variate Timeseries")
decoder_input = torch.zeros_like(batch_y)
decoder_input[:, 1:, :] = batch_y[:, :-1, :]
decoder_input[:, 0, :] = batch_x[:, -1, :]
with autocast(device_type= processor_type):
output = model(batch_x, decoder_input)
loss = criterion(output, batch_y)
batch_sMAPE = smape(batch_y, output)
val_loss += loss.item()
val_sMAPE += batch_sMAPE
#val_MASE += batch_MASE


val_loss /= len(val_loader)
val_sMAPE /= len(val_loader)
val_MASE /= len(val_loader)
    # Early stopping check
if val_loss < best_val_loss:
best_val_loss = val_loss
patience_counter = 0
elif val_loss > (best_val_loss*1.025):
patience_counter += 1
if patience_counter >= patience:
print(f"\nEarly stopping triggered after {epoch + 1} epochs")
return False
return True, best_val_loss, (val_loss, val_sMAPE, val_MASE)

# TimeRepresentationAgent (NEW)
class TimeRepresentationAgent(nn.Module):
    def __init__(self, hidden_size, num_time_slots=24):
        super(TimeRepresentationAgent, self).__init__()
        self.hidden_size = hidden_size
        self.num_time_slots = num_time_slots  # e.g., 24 hours
        self.policy_net = nn.Sequential(
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, num_time_slots)
        )
        self.time_embedding = nn.Embedding(num_time_slots, hidden_size)

    def forward(self, state):
        # State: (batch_size, hidden_size)
        logits = self.policy_net(state)  # (batch_size, num_time_slots)
        m = torch.distributions.Categorical(logits=logits)
        time_indices = m.sample()  # (batch_size,)
        time_embed = self.time_embedding(time_indices)  # (batch_size, hidden_size)
        return time_embed, m.log_prob(time_indices)  # Return embedding and log prob