added code

code/new_realtime/models/cnn.py

@@ -0,0 +1,58 @@
+import math
+
+import torch
+from torch import nn
+
+
+class CNNTransformer(nn.Module):
+    def __init__(self,
+                 input_dim,
+                 output_dim,
+                 seq_len,
+                 cnn_channels,
+                 kernel_size,
+                 embed_dim,
+                 num_enc_layers,
+                 num_heads):
+        super().__init__()
+        self.cnn = nn.Sequential(
+            nn.Conv1d(input_dim, cnn_channels, kernel_size=kernel_size, padding=1),
+            nn.BatchNorm1d(cnn_channels),
+            nn.ReLU(),
+            nn.Conv1d(cnn_channels, embed_dim, kernel_size=kernel_size, padding=1),
+            nn.BatchNorm1d(embed_dim),
+            nn.ReLU()
+        )
+
+        # Compute positional embedding ONCE at init
+        pe = self._get_sinusoidal_embedding(seq_len, embed_dim)  # (seq_len, embed_dim)
+        self.register_buffer('pos_embed', pe.unsqueeze(0))  # (1, seq_len, embed_dim)
+
+        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads)
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(embed_dim, output_dim)
+
+    def _get_sinusoidal_embedding(self, seq_len, embed_dim):
+        position = torch.arange(0, seq_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim))
+        pe = torch.zeros(seq_len, embed_dim)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe  # (seq_len, embed_dim)
+
+    def forward(self, x):
+        # x: (batch, seq_len, input_dim)
+        x = x.permute(0, 2, 1)  # (B, input_dim, seq_len)
+        cnn_out = self.cnn(x)  # (B, embed_dim, seq_len)
+        cnn_out = cnn_out.permute(2, 0, 1)  # (S, B, E) for Transformer
+
+        # Add positional embedding
+        pos_embed = self.pos_embed[:, :cnn_out.size(0), :]  # (1, seq_len, embed_dim)
+        pos_embed = pos_embed.transpose(0, 1)  # → (seq_len, 1, embed_dim)
+        cnn_out = cnn_out + pos_embed  # broadcast over batch
+
+        # Apply Transformer encoder
+        enc = self.encoder(cnn_out)  # (S, B, E)
+        pooled = enc.mean(0)  # (B, E)
+        return self.head(pooled)

code/new_realtime/models/collation.py

@@ -0,0 +1,94 @@
+import numpy as np
+import torch
+
+
+def flat_x_y_collate(batch, *args, **kwargs) -> tuple:
+    """
+    Collate function for a classical model with flat x and y feature vectors
+    """
+    # Flatten each item in the batch
+    flattened_x = list()
+    flattened_y = list()
+    for i in range(len(batch)):
+        item = batch[i]
+        flattened_x_item = list()
+        flattened_y_item = list()
+        for j, (category, windows) in enumerate(item.items()):
+            for feature in windows.T:
+                if category == "target":
+                    flattened_y_item.append(feature)
+                else:
+                    flattened_x_item.append(feature)
+        # create numpy arrays with t_1_f1, t_1_f2, t_2_f1, t_2_f2 and so on
+        reordered_x = np.empty((len(flattened_x_item) * len(flattened_x_item[0])))
+        reordered_y = np.empty((len(flattened_y_item) * len(flattened_y_item[0])))
+        for j in range(len(flattened_x_item)):
+            for k in range(len(flattened_x_item[j])):
+                reordered_x[j + k * len(flattened_x_item)] = flattened_x_item[j][k]
+
+        for j in range(len(flattened_y_item)):
+            for k in range(len(flattened_y_item[j])):
+                reordered_y[j + k * len(flattened_y_item)] = flattened_y_item[j][k]
+
+        # append to the list
+        flattened_x.append(reordered_x)
+        flattened_y.append(reordered_y)
+
+    return np.array(flattened_x), np.array(flattened_y)
+
+
+def simple_x_y_collate(batch):
+    """
+    Collate function for LSTM model
+    """
+    collated_x = list()
+    collated_y = list()
+    for item in batch:
+        current_x = list()
+        current_y = list()
+        for category, windows in item.items():
+            if category == "target_features":
+                current_y.append(windows)
+            else:
+                current_x.append(windows)
+
+        collated_x.append(np.concatenate(current_x, axis=1))
+        collated_y.append(np.concatenate(current_y, axis=1))
+
+    return torch.tensor(collated_x, dtype=torch.float32), torch.tensor(collated_y, dtype=torch.float32)
+
+
+def collate_with_padding(batch,
+                         padding_value: float = 0.0, ):
+    """
+    Collate the batch with padding
+    """
+
+    # create simple lists for x and y
+    collated_x = list()
+    collated_y = list()
+    for item in batch:
+        current_x = list()
+        current_y = list()
+        for category, windows in item.items():
+            if category == "target_features":
+                current_y.append(windows)
+            else:
+                current_x.append(windows)
+
+        collated_x.append(np.concatenate(current_x, axis=1))
+        collated_y.append(np.concatenate(current_y, axis=1))
+
+    # get the max length of the x and y
+    max_x_length = max([x.shape[0] for x in collated_x])
+    max_y_length = max([y.shape[0] for y in collated_y])
+    # pad the x and y
+    padded_x = list()
+    padded_y = list()
+    for x, y in zip(collated_x, collated_y):
+        padded_x.append(
+            np.pad(x, ((max_x_length - x.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value))
+        padded_y.append(
+            np.pad(y, ((max_y_length - y.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value))
+
+    return torch.tensor(padded_x, dtype=torch.float32), torch.tensor(padded_y, dtype=torch.float32)

code/new_realtime/models/lstm.py

@@ -0,0 +1,46 @@
+from torch import nn
+
+
+class LSTMModel(nn.Module):
+    def __init__(self,
+                 input_dim: int,
+                 output_dim: int,
+                 cnn_channels: int,
+                 cnn_kernel_size: int,
+                 embed_dim: int,
+                 lstm_hidden_size=64,
+                 num_layers=1):
+        super().__init__()
+        self.cnn = nn.Sequential(
+            nn.Conv1d(input_dim, cnn_channels, kernel_size=cnn_kernel_size, padding=1),
+            nn.BatchNorm1d(cnn_channels),
+            nn.ReLU(),
+            nn.Conv1d(cnn_channels, embed_dim, kernel_size=cnn_kernel_size, padding=1),
+            nn.BatchNorm1d(embed_dim),
+            nn.ReLU()
+        )
+        self.lstm = nn.LSTM(
+            input_size=embed_dim,
+            hidden_size=lstm_hidden_size,
+            num_layers=num_layers,
+            batch_first=True,
+        )
+        self.head = nn.Sequential(
+            nn.Linear(lstm_hidden_size, output_dim)  # Output is a scalar Δt
+        )
+
+    def forward(self, x):
+        # x: (batch_size, seq_len, input_size)
+        x = x.permute(0, 2, 1)
+        # x: (batch_size, input_size, seq_len)
+        x = self.cnn(x)
+        # x: (batch_size, embed_dim, seq_len)
+        x = x.permute(0, 2, 1)
+        # x: (batch_size, seq_len, embed_dim)
+        x, _ = self.lstm(x)
+        # x: (batch_size, seq_len, lstm_hidden_size)
+        x = x[:, -1, :]  # Get the last time step
+        # x: (batch_size, lstm_hidden_size)
+        x = self.head(x)
+        # x: (batch_size, output_size)
+        return x

code/new_realtime/models/tft_utils.py

@@ -0,0 +1,155 @@
+import datetime
+import os
+
+import torch
+from torch import nn
+from torch.optim import AdamW
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.utils.data import IterableDataset, DataLoader
+
+from models.third_party.tft_model import TemporalFusionTransformer
+
+
+def get_tft_model(model_configuration: dict,
+                  sample_item: dict,
+                  device: str) -> nn.Module:
+    config_class = create_config_class(model_configuration, sample_item)
+    model = TemporalFusionTransformer(config_class)
+    model.to(device)
+    return model
+
+
+def create_training_state(model_configuration: dict,
+                          training_configuration: dict,
+                          train_dataset: IterableDataset | DataLoader,
+                          device: str) -> dict:
+    training_state = dict()
+
+    sample_item = next(iter(train_dataset))
+
+    if model_configuration["model_type"] == "TemporalFusionTransformer":
+        training_state["model"] = get_tft_model(model_configuration,
+                                                sample_item,
+                                                device)
+    else:
+        raise NotImplementedError(f"Model type {model_configuration['type']} not implemented")
+
+    training_state["optimizer"] = AdamW(training_state["model"].parameters(),
+                                        lr=training_configuration["learning_rate"])
+    training_state["scheduler"] = OneCycleLR(training_state["optimizer"],
+                                             max_lr=training_configuration["learning_rate"],
+                                             total_steps=len(train_dataset) *
+                                                         training_configuration[
+                                                             "epochs"])
+    training_state["current_epoch"] = 1
+    return training_state
+
+
+def get_checkpoints(model_configuration: dict, training_configuration: dict) -> list:
+    checkpoints_dir = f"{model_configuration['model_dir']}/trainings/{training_configuration['id']}/checkpoints"
+    if os.path.exists(checkpoints_dir):
+        checkpoints = [os.path.join(checkpoints_dir, f) for f in os.listdir(checkpoints_dir) if
+                       f.endswith('.pt') and "checkpoint" in f]
+        checkpoints.sort(key=os.path.getmtime, reverse=False)
+        return checkpoints
+    else:
+        return []
+
+
+def load_checkpoint(checkpoint_path: str,
+                    training_configuration: dict,
+                    model_configuration: dict,
+                    test_data_loader: IterableDataset | DataLoader,
+                    device: str) -> tuple:
+    checkpoint = torch.load(checkpoint_path)
+    training_state = torch.load(checkpoint_path)
+
+    sample_item = next(iter(test_data_loader))
+
+    # load model
+    if model_configuration["model_type"] == "TemporalFusionTransformer":
+        model = get_tft_model(model_configuration,
+                              sample_item,
+                              device)
+        model_configuration["model"] = model
+    else:
+        raise NotImplementedError(f"Model type {model_configuration['model_type']} not implemented")
+
+    # load optimizer
+    optimizer = AdamW(model.parameters(), lr=training_configuration["learning_rate"])
+    optimizer.load_state_dict(checkpoint["optimizer"])
+    training_state["optimizer"] = optimizer
+
+    # load scheduler
+    scheduler = OneCycleLR(optimizer,
+                           max_lr=training_configuration["learning_rate"],
+                           total_steps=len(test_data_loader) * training_configuration["epochs"])
+    scheduler.load_state_dict(checkpoint["scheduler"])
+    training_state["scheduler"] = scheduler
+
+    return training_state
+
+
+def save_checkpoint(training_config: dict,
+                    training_state: dict,
+                    model_config: dict) -> None:
+    checkpoints_dir = f"{model_config['model_dir']}/trainings/{training_config['id']}/checkpoints"
+    checkpoint_id = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    checkpoint_path = os.path.join(checkpoints_dir, f"checkpoint_{checkpoint_id}.pt")
+    if not os.path.exists(checkpoints_dir):
+        os.makedirs(checkpoints_dir)
+
+    config_to_save = training_state.copy()
+
+    # replace training parts with their state dicts
+    config_to_save["model"] = config_to_save["model"].state_dict()
+    config_to_save["optimizer"] = config_to_save["optimizer"].state_dict()
+    config_to_save["scheduler"] = config_to_save["scheduler"].state_dict()
+
+    # save training state
+    torch.save(config_to_save, checkpoint_path)
+
+
+def create_config_class(config: dict, sample_batch: dict) -> object:
+    class ConfigClass:
+        def __init__(self):
+            # Feature sizes
+            self.static_categorical_inp_lens = []
+            self.temporal_known_categorical_inp_lens = []
+            self.temporal_observed_categorical_inp_lens = []
+
+            model_parameters = config["model_parameters"]
+
+            self.example_length = model_parameters["encoder_length"] + model_parameters["decoder_length"]
+            self.encoder_length = model_parameters["encoder_length"]
+
+            self.n_head = model_parameters["attention_heads"]
+            self.hidden_size = model_parameters["state_size"]
+            self.dropout = model_parameters["dropout"]
+            self.attn_dropout = model_parameters["attention_dropout"]
+            self.quantiles = model_parameters["output_quantiles"]
+            self.use_past_targets = model_parameters["use_past_targets"]
+
+            #### Derived variables ####
+            self.temporal_known_continuous_inp_size = sample_batch["k_cont"].shape[2]
+            self.temporal_observed_continuous_inp_size = sample_batch["o_cont"].shape[2]
+            self.temporal_target_size = sample_batch["target"].shape[2]
+            self.static_continuous_inp_size = sample_batch["s_cont"].shape[2]
+
+            self.num_static_vars = self.static_continuous_inp_size + len(self.static_categorical_inp_lens)
+            self.num_future_vars = self.temporal_known_continuous_inp_size + len(
+                self.temporal_known_categorical_inp_lens)
+            if self.use_past_targets:
+                self.num_historic_vars = self.num_future_vars + self.temporal_observed_continuous_inp_size + self.temporal_target_size + len(
+                    self.temporal_observed_categorical_inp_lens)
+            else:
+                self.num_historic_vars = self.num_future_vars + self.temporal_observed_continuous_inp_size + len(
+                    self.temporal_observed_categorical_inp_lens)
+            # self.num_historic_vars = sum([self.num_future_vars,
+            #                               self.temporal_observed_continuous_inp_size,
+            #                               self.temporal_target_size,
+            #                               len(self.temporal_observed_categorical_inp_lens),
+            #                               ])
+            self.target_size = self.temporal_target_size
+
+    return ConfigClass()

code/new_realtime/models/third_party/patch_tst/layers/AutoCorrelation.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+import numpy as np
+import math
+from math import sqrt
+import os
+
+
+class AutoCorrelation(nn.Module):
+    """
+    AutoCorrelation Mechanism with the following two phases:
+    (1) period-based dependencies discovery
+    (2) time delay aggregation
+    This block can replace the self-attention family mechanism seamlessly.
+    """
+    def __init__(self, mask_flag=True, factor=1, scale=None, attention_dropout=0.1, output_attention=False):
+        super(AutoCorrelation, self).__init__()
+        self.factor = factor
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.output_attention = output_attention
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def time_delay_agg_training(self, values, corr):
+        """
+        SpeedUp version of Autocorrelation (a batch-normalization style design)
+        This is for the training phase.
+        """
+        head = values.shape[1]
+        channel = values.shape[2]
+        length = values.shape[3]
+        # find top k
+        top_k = int(self.factor * math.log(length))
+        mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
+        index = torch.topk(torch.mean(mean_value, dim=0), top_k, dim=-1)[1]
+        weights = torch.stack([mean_value[:, index[i]] for i in range(top_k)], dim=-1)
+        # update corr
+        tmp_corr = torch.softmax(weights, dim=-1)
+        # aggregation
+        tmp_values = values
+        delays_agg = torch.zeros_like(values).float()
+        for i in range(top_k):
+            pattern = torch.roll(tmp_values, -int(index[i]), -1)
+            delays_agg = delays_agg + pattern * \
+                         (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
+        return delays_agg
+
+    def time_delay_agg_inference(self, values, corr):
+        """
+        SpeedUp version of Autocorrelation (a batch-normalization style design)
+        This is for the inference phase.
+        """
+        batch = values.shape[0]
+        head = values.shape[1]
+        channel = values.shape[2]
+        length = values.shape[3]
+        # index init
+        init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
+        # find top k
+        top_k = int(self.factor * math.log(length))
+        mean_value = torch.mean(torch.mean(corr, dim=1), dim=1)
+        weights = torch.topk(mean_value, top_k, dim=-1)[0]
+        delay = torch.topk(mean_value, top_k, dim=-1)[1]
+        # update corr
+        tmp_corr = torch.softmax(weights, dim=-1)
+        # aggregation
+        tmp_values = values.repeat(1, 1, 1, 2)
+        delays_agg = torch.zeros_like(values).float()
+        for i in range(top_k):
+            tmp_delay = init_index + delay[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length)
+            pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
+            delays_agg = delays_agg + pattern * \
+                         (tmp_corr[:, i].unsqueeze(1).unsqueeze(1).unsqueeze(1).repeat(1, head, channel, length))
+        return delays_agg
+
+    def time_delay_agg_full(self, values, corr):
+        """
+        Standard version of Autocorrelation
+        """
+        batch = values.shape[0]
+        head = values.shape[1]
+        channel = values.shape[2]
+        length = values.shape[3]
+        # index init
+        init_index = torch.arange(length).unsqueeze(0).unsqueeze(0).unsqueeze(0).repeat(batch, head, channel, 1).cuda()
+        # find top k
+        top_k = int(self.factor * math.log(length))
+        weights = torch.topk(corr, top_k, dim=-1)[0]
+        delay = torch.topk(corr, top_k, dim=-1)[1]
+        # update corr
+        tmp_corr = torch.softmax(weights, dim=-1)
+        # aggregation
+        tmp_values = values.repeat(1, 1, 1, 2)
+        delays_agg = torch.zeros_like(values).float()
+        for i in range(top_k):
+            tmp_delay = init_index + delay[..., i].unsqueeze(-1)
+            pattern = torch.gather(tmp_values, dim=-1, index=tmp_delay)
+            delays_agg = delays_agg + pattern * (tmp_corr[..., i].unsqueeze(-1))
+        return delays_agg
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, E = queries.shape
+        _, S, _, D = values.shape
+        if L > S:
+            zeros = torch.zeros_like(queries[:, :(L - S), :]).float()
+            values = torch.cat([values, zeros], dim=1)
+            keys = torch.cat([keys, zeros], dim=1)
+        else:
+            values = values[:, :L, :, :]
+            keys = keys[:, :L, :, :]
+
+        # period-based dependencies
+        q_fft = torch.fft.rfft(queries.permute(0, 2, 3, 1).contiguous(), dim=-1)
+        k_fft = torch.fft.rfft(keys.permute(0, 2, 3, 1).contiguous(), dim=-1)
+        res = q_fft * torch.conj(k_fft)
+        corr = torch.fft.irfft(res, dim=-1)
+
+        # time delay agg
+        if self.training:
+            V = self.time_delay_agg_training(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+        else:
+            V = self.time_delay_agg_inference(values.permute(0, 2, 3, 1).contiguous(), corr).permute(0, 3, 1, 2)
+
+        if self.output_attention:
+            return (V.contiguous(), corr.permute(0, 3, 1, 2))
+        else:
+            return (V.contiguous(), None)
+
+
+class AutoCorrelationLayer(nn.Module):
+    def __init__(self, correlation, d_model, n_heads, d_keys=None,
+                 d_values=None):
+        super(AutoCorrelationLayer, self).__init__()
+
+        d_keys = d_keys or (d_model // n_heads)
+        d_values = d_values or (d_model // n_heads)
+
+        self.inner_correlation = correlation
+        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.value_projection = nn.Linear(d_model, d_values * n_heads)
+        self.out_projection = nn.Linear(d_values * n_heads, d_model)
+        self.n_heads = n_heads
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        H = self.n_heads
+
+        queries = self.query_projection(queries).view(B, L, H, -1)
+        keys = self.key_projection(keys).view(B, S, H, -1)
+        values = self.value_projection(values).view(B, S, H, -1)
+
+        out, attn = self.inner_correlation(
+            queries,
+            keys,
+            values,
+            attn_mask
+        )
+        out = out.view(B, L, -1)
+
+        return self.out_projection(out), attn

code/new_realtime/models/third_party/patch_tst/layers/Autoformer_EncDec.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class my_Layernorm(nn.Module):
+    """
+    Special designed layernorm for the seasonal part
+    """
+    def __init__(self, channels):
+        super(my_Layernorm, self).__init__()
+        self.layernorm = nn.LayerNorm(channels)
+
+    def forward(self, x):
+        x_hat = self.layernorm(x)
+        bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
+        return x_hat - bias
+
+
+class moving_avg(nn.Module):
+    """
+    Moving average block to highlight the trend of time series
+    """
+    def __init__(self, kernel_size, stride):
+        super(moving_avg, self).__init__()
+        self.kernel_size = kernel_size
+        self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
+
+    def forward(self, x):
+        # padding on the both ends of time series
+        front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+        x = torch.cat([front, x, end], dim=1)
+        x = self.avg(x.permute(0, 2, 1))
+        x = x.permute(0, 2, 1)
+        return x
+
+
+class series_decomp(nn.Module):
+    """
+    Series decomposition block
+    """
+    def __init__(self, kernel_size):
+        super(series_decomp, self).__init__()
+        self.moving_avg = moving_avg(kernel_size, stride=1)
+
+    def forward(self, x):
+        moving_mean = self.moving_avg(x)
+        res = x - moving_mean
+        return res, moving_mean
+
+
+class EncoderLayer(nn.Module):
+    """
+    Autoformer encoder layer with the progressive decomposition architecture
+    """
+    def __init__(self, attention, d_model, d_ff=None, moving_avg=25, dropout=0.1, activation="relu"):
+        super(EncoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.attention = attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
+        self.decomp1 = series_decomp(moving_avg)
+        self.decomp2 = series_decomp(moving_avg)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, attn_mask=None):
+        new_x, attn = self.attention(
+            x, x, x,
+            attn_mask=attn_mask
+        )
+        x = x + self.dropout(new_x)
+        x, _ = self.decomp1(x)
+        y = x
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+        res, _ = self.decomp2(x + y)
+        return res, attn
+
+
+class Encoder(nn.Module):
+    """
+    Autoformer encoder
+    """
+    def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
+        super(Encoder, self).__init__()
+        self.attn_layers = nn.ModuleList(attn_layers)
+        self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+        self.norm = norm_layer
+
+    def forward(self, x, attn_mask=None):
+        attns = []
+        if self.conv_layers is not None:
+            for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
+                x, attn = attn_layer(x, attn_mask=attn_mask)
+                x = conv_layer(x)
+                attns.append(attn)
+            x, attn = self.attn_layers[-1](x)
+            attns.append(attn)
+        else:
+            for attn_layer in self.attn_layers:
+                x, attn = attn_layer(x, attn_mask=attn_mask)
+                attns.append(attn)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x, attns
+
+
+class DecoderLayer(nn.Module):
+    """
+    Autoformer decoder layer with the progressive decomposition architecture
+    """
+    def __init__(self, self_attention, cross_attention, d_model, c_out, d_ff=None,
+                 moving_avg=25, dropout=0.1, activation="relu"):
+        super(DecoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1, bias=False)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1, bias=False)
+        self.decomp1 = series_decomp(moving_avg)
+        self.decomp2 = series_decomp(moving_avg)
+        self.decomp3 = series_decomp(moving_avg)
+        self.dropout = nn.Dropout(dropout)
+        self.projection = nn.Conv1d(in_channels=d_model, out_channels=c_out, kernel_size=3, stride=1, padding=1,
+                                    padding_mode='circular', bias=False)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        x = x + self.dropout(self.self_attention(
+            x, x, x,
+            attn_mask=x_mask
+        )[0])
+        x, trend1 = self.decomp1(x)
+        x = x + self.dropout(self.cross_attention(
+            x, cross, cross,
+            attn_mask=cross_mask
+        )[0])
+        x, trend2 = self.decomp2(x)
+        y = x
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+        x, trend3 = self.decomp3(x + y)
+
+        residual_trend = trend1 + trend2 + trend3
+        residual_trend = self.projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
+        return x, residual_trend
+
+
+class Decoder(nn.Module):
+    """
+    Autoformer encoder
+    """
+    def __init__(self, layers, norm_layer=None, projection=None):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm_layer
+        self.projection = projection
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None, trend=None):
+        for layer in self.layers:
+            x, residual_trend = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+            trend = trend + residual_trend
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        if self.projection is not None:
+            x = self.projection(x)
+        return x, trend

code/new_realtime/models/third_party/patch_tst/layers/Embed.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils import weight_norm
+import math
+
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEmbedding, self).__init__()
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model).float()
+        pe.require_grad = False
+
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return self.pe[:, :x.size(1)]
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(TokenEmbedding, self).__init__()
+        padding = 1 if torch.__version__ >= '1.5.0' else 2
+        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,
+                                   kernel_size=3, padding=padding, padding_mode='circular', bias=False)
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
+
+    def forward(self, x):
+        x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)
+        return x
+
+
+class FixedEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(FixedEmbedding, self).__init__()
+
+        w = torch.zeros(c_in, d_model).float()
+        w.require_grad = False
+
+        position = torch.arange(0, c_in).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
+
+        w[:, 0::2] = torch.sin(position * div_term)
+        w[:, 1::2] = torch.cos(position * div_term)
+
+        self.emb = nn.Embedding(c_in, d_model)
+        self.emb.weight = nn.Parameter(w, requires_grad=False)
+
+    def forward(self, x):
+        return self.emb(x).detach()
+
+
+class TemporalEmbedding(nn.Module):
+    def __init__(self, d_model, embed_type='fixed', freq='h'):
+        super(TemporalEmbedding, self).__init__()
+
+        minute_size = 4
+        hour_size = 24
+        weekday_size = 7
+        day_size = 32
+        month_size = 13
+
+        Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
+        if freq == 't':
+            self.minute_embed = Embed(minute_size, d_model)
+        self.hour_embed = Embed(hour_size, d_model)
+        self.weekday_embed = Embed(weekday_size, d_model)
+        self.day_embed = Embed(day_size, d_model)
+        self.month_embed = Embed(month_size, d_model)
+
+    def forward(self, x):
+        x = x.long()
+
+        minute_x = self.minute_embed(x[:, :, 4]) if hasattr(self, 'minute_embed') else 0.
+        hour_x = self.hour_embed(x[:, :, 3])
+        weekday_x = self.weekday_embed(x[:, :, 2])
+        day_x = self.day_embed(x[:, :, 1])
+        month_x = self.month_embed(x[:, :, 0])
+
+        return hour_x + weekday_x + day_x + month_x + minute_x
+
+
+class TimeFeatureEmbedding(nn.Module):
+    def __init__(self, d_model, embed_type='timeF', freq='h'):
+        super(TimeFeatureEmbedding, self).__init__()
+
+        freq_map = {'h': 4, 't': 5, 's': 6, 'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}
+        d_inp = freq_map[freq]
+        self.embed = nn.Linear(d_inp, d_model, bias=False)
+
+    def forward(self, x):
+        return self.embed(x)
+
+
+class DataEmbedding(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
+        return self.dropout(x)
+
+
+class DataEmbedding_wo_pos(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding_wo_pos, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x) + self.temporal_embedding(x_mark)
+        return self.dropout(x)
+
+class DataEmbedding_wo_pos_temp(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding_wo_pos_temp, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x)
+        return self.dropout(x)
+
+class DataEmbedding_wo_temp(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
+        super(DataEmbedding_wo_temp, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x) + self.position_embedding(x)
+        return self.dropout(x)

code/new_realtime/models/third_party/patch_tst/layers/PatchTST_backbone.py

@@ -0,0 +1,429 @@
+__all__ = ['PatchTST_backbone']
+
+# Cell
+from typing import Callable, Optional
+import torch
+from torch import nn
+from torch import Tensor
+import torch.nn.functional as F
+import numpy as np
+
+# from collections import OrderedDict
+from models.third_party.patch_tst.layers.PatchTST_layers import *
+from models.third_party.patch_tst.layers.RevIN import RevIN
+
+
+class CustomHead(nn.Module):
+    def __init__(self, output_dim, n_vars, target_window, nf, head_dropout=0):
+        super().__init__()
+        self.flatten = nn.Flatten(start_dim=-3)
+        self.linear = nn.Linear(nf * n_vars, output_dim * target_window)
+        self.dropout = nn.Dropout(head_dropout)
+        self.target_window = target_window
+        self.output_dim = output_dim
+
+    def forward(self, x):  # x: [bs x nvars x d_model x patch_num]
+        x = self.flatten(x)  # [bs x (nf * nvars)]
+        x = self.linear(x)  # [bs x (target_window * output_dim)]
+        x = self.dropout(x)
+        x = x.view(x.size(0), self.target_window, self.output_dim)  # [bs x target_window x output_dim]
+        # permute to match intended structure
+        x = x.permute(0, 2, 1)  # [bs x output_dim x target_window]
+        return x
+
+
+# Cell
+class PatchTST_backbone(nn.Module):
+    def __init__(self, c_in: int,
+                 context_window: int, target_window: int, patch_len: int, stride: int,
+                 # extras
+                 dec_out: int = 1,
+                 seq_pred: bool = False,
+                 #
+                 max_seq_len: Optional[int] = 1024,
+                 n_layers: int = 3, d_model=128, n_heads=16, d_k: Optional[int] = None, d_v: Optional[int] = None,
+                 d_ff: int = 256, norm: str = 'BatchNorm', attn_dropout: float = 0., dropout: float = 0.,
+                 act: str = "gelu", key_padding_mask: bool = 'auto',
+                 padding_var: Optional[int] = None, attn_mask: Optional[Tensor] = None, res_attention: bool = True,
+                 pre_norm: bool = False, store_attn: bool = False,
+                 pe: str = 'zeros', learn_pe: bool = True, fc_dropout: float = 0., head_dropout=0, padding_patch=None,
+                 pretrain_head: bool = False, head_type='flatten', individual=False, revin=True, affine=True,
+                 subtract_last=False,
+                 verbose: bool = False, **kwargs):
+
+        super().__init__()
+
+        # RevIn
+        self.revin = revin
+        if self.revin: self.revin_layer = RevIN(c_in, affine=affine, subtract_last=subtract_last)
+
+        # Patching
+        self.patch_len = patch_len
+        self.stride = stride
+        self.padding_patch = padding_patch
+        patch_num = int((context_window - patch_len) / stride + 1)
+        if padding_patch == 'end':  # can be modified to general case
+            self.padding_patch_layer = nn.ReplicationPad1d((0, stride))
+            patch_num += 1
+
+        # Backbone 
+        self.backbone = TSTiEncoder(c_in, patch_num=patch_num, patch_len=patch_len, max_seq_len=max_seq_len,
+                                    n_layers=n_layers, d_model=d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff,
+                                    attn_dropout=attn_dropout, dropout=dropout, act=act,
+                                    key_padding_mask=key_padding_mask, padding_var=padding_var,
+                                    attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
+                                    store_attn=store_attn,
+                                    pe=pe, learn_pe=learn_pe, verbose=verbose, **kwargs)
+
+        # Head
+        self.head_nf = d_model * patch_num
+        self.n_vars = c_in
+        self.pretrain_head = pretrain_head
+        self.head_type = head_type
+        self.individual = individual
+        # extras for non-sequence prediction
+        self.seq_pred = seq_pred
+        self.dec_out = dec_out
+
+        if self.pretrain_head:
+            self.head = self.create_pretrain_head(self.head_nf, c_in,
+                                                  fc_dropout)  # custom head passed as a partial func with all its kwargs
+        elif not self.seq_pred:
+            self.head = CustomHead(output_dim=self.dec_out,
+                                   n_vars=self.n_vars,
+                                   target_window=target_window,
+                                   nf=self.head_nf,
+                                   head_dropout=head_dropout)
+        elif head_type == 'flatten':
+            self.head = Flatten_Head(self.individual, self.n_vars, self.head_nf, target_window,
+                                     head_dropout=head_dropout)
+
+    def forward(self, z):  # z: [bs x nvars x seq_len]
+        # norm
+        if self.revin:
+            z = z.permute(0, 2, 1)
+            z = self.revin_layer(z, 'norm')
+            z = z.permute(0, 2, 1)
+
+        # do patching
+        if self.padding_patch == 'end':
+            z = self.padding_patch_layer(z)
+        z = z.unfold(dimension=-1, size=self.patch_len, step=self.stride)  # z: [bs x nvars x patch_num x patch_len]
+        z = z.permute(0, 1, 3, 2)  # z: [bs x nvars x patch_len x patch_num]
+
+        # model
+        z = self.backbone(z)  # z: [bs x nvars x d_model x patch_num]
+        z = self.head(z)  # z: [bs x nvars x target_window]
+
+        # denorm
+        if self.revin:
+            z = z.permute(0, 2, 1)
+            z = self.revin_layer(z, 'denorm')
+            z = z.permute(0, 2, 1)
+
+        return z
+
+    def create_pretrain_head(self, head_nf, vars, dropout):
+        return nn.Sequential(nn.Dropout(dropout),
+                             nn.Conv1d(head_nf, vars, 1)
+                             )
+
+
+class Flatten_Head(nn.Module):
+    def __init__(self, individual, n_vars, nf, target_window, head_dropout=0):
+        super().__init__()
+
+        self.individual = individual
+        self.n_vars = n_vars
+
+        if self.individual:
+            self.linears = nn.ModuleList()
+            self.dropouts = nn.ModuleList()
+            self.flattens = nn.ModuleList()
+            for i in range(self.n_vars):
+                self.flattens.append(nn.Flatten(start_dim=-2))
+                self.linears.append(nn.Linear(nf, target_window))
+                self.dropouts.append(nn.Dropout(head_dropout))
+        else:
+            self.flatten = nn.Flatten(start_dim=-2)
+            self.linear = nn.Linear(nf, target_window)
+            self.dropout = nn.Dropout(head_dropout)
+
+    def forward(self, x):  # x: [bs x nvars x d_model x patch_num]
+        if self.individual:
+            x_out = []
+            for i in range(self.n_vars):
+                z = self.flattens[i](x[:, i, :, :])  # z: [bs x d_model * patch_num]
+                z = self.linears[i](z)  # z: [bs x target_window]
+                z = self.dropouts[i](z)
+                x_out.append(z)
+            x = torch.stack(x_out, dim=1)  # x: [bs x nvars x target_window]
+        else:
+            x = self.flatten(x)
+            x = self.linear(x)
+            x = self.dropout(x)
+        return x
+
+
+class TSTiEncoder(nn.Module):  # i means channel-independent
+    def __init__(self, c_in, patch_num, patch_len, max_seq_len=1024,
+                 n_layers=3, d_model=128, n_heads=16, d_k=None, d_v=None,
+                 d_ff=256, norm='BatchNorm', attn_dropout=0., dropout=0., act="gelu", store_attn=False,
+                 key_padding_mask='auto', padding_var=None, attn_mask=None, res_attention=True, pre_norm=False,
+                 pe='zeros', learn_pe=True, verbose=False, **kwargs):
+        super().__init__()
+
+        self.patch_num = patch_num
+        self.patch_len = patch_len
+
+        # Input encoding
+        q_len = patch_num
+        self.W_P = nn.Linear(patch_len, d_model)  # Eq 1: projection of feature vectors onto a d-dim vector space
+        self.seq_len = q_len
+
+        # Positional encoding
+        self.W_pos = positional_encoding(pe, learn_pe, q_len, d_model)
+
+        # Residual dropout
+        self.dropout = nn.Dropout(dropout)
+
+        # Encoder
+        self.encoder = TSTEncoder(q_len, d_model, n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
+                                  attn_dropout=attn_dropout, dropout=dropout,
+                                  pre_norm=pre_norm, activation=act, res_attention=res_attention, n_layers=n_layers,
+                                  store_attn=store_attn)
+
+    def forward(self, x) -> Tensor:  # x: [bs x nvars x patch_len x patch_num]
+
+        n_vars = x.shape[1]
+        # Input encoding
+        x = x.permute(0, 1, 3, 2)  # x: [bs x nvars x patch_num x patch_len]
+        x = self.W_P(x)  # x: [bs x nvars x patch_num x d_model]
+
+        u = torch.reshape(x, (x.shape[0] * x.shape[1], x.shape[2], x.shape[3]))  # u: [bs * nvars x patch_num x d_model]
+        u = self.dropout(u + self.W_pos)  # u: [bs * nvars x patch_num x d_model]
+
+        # Encoder
+        z = self.encoder(u)  # z: [bs * nvars x patch_num x d_model]
+        z = torch.reshape(z, (-1, n_vars, z.shape[-2], z.shape[-1]))  # z: [bs x nvars x patch_num x d_model]
+        z = z.permute(0, 1, 3, 2)  # z: [bs x nvars x d_model x patch_num]
+
+        return z
+
+    # Cell
+
+
+class TSTEncoder(nn.Module):
+    def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=None,
+                 norm='BatchNorm', attn_dropout=0., dropout=0., activation='gelu',
+                 res_attention=False, n_layers=1, pre_norm=False, store_attn=False):
+        super().__init__()
+
+        self.layers = nn.ModuleList(
+            [TSTEncoderLayer(q_len, d_model, n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
+                             attn_dropout=attn_dropout, dropout=dropout,
+                             activation=activation, res_attention=res_attention,
+                             pre_norm=pre_norm, store_attn=store_attn) for i in range(n_layers)])
+        self.res_attention = res_attention
+
+    def forward(self, src: Tensor, key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
+        output = src
+        scores = None
+        if self.res_attention:
+            for mod in self.layers: output, scores = mod(output, prev=scores, key_padding_mask=key_padding_mask,
+                                                         attn_mask=attn_mask)
+            return output
+        else:
+            for mod in self.layers: output = mod(output, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+            return output
+
+
+class TSTEncoderLayer(nn.Module):
+    def __init__(self, q_len, d_model, n_heads, d_k=None, d_v=None, d_ff=256, store_attn=False,
+                 norm='BatchNorm', attn_dropout=0, dropout=0., bias=True, activation="gelu", res_attention=False,
+                 pre_norm=False):
+        super().__init__()
+        assert not d_model % n_heads, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
+        d_k = d_model // n_heads if d_k is None else d_k
+        d_v = d_model // n_heads if d_v is None else d_v
+
+        # Multi-Head attention
+        self.res_attention = res_attention
+        self.self_attn = _MultiheadAttention(d_model, n_heads, d_k, d_v, attn_dropout=attn_dropout,
+                                             proj_dropout=dropout, res_attention=res_attention)
+
+        # Add & Norm
+        self.dropout_attn = nn.Dropout(dropout)
+        if "batch" in norm.lower():
+            self.norm_attn = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2))
+        else:
+            self.norm_attn = nn.LayerNorm(d_model)
+
+        # Position-wise Feed-Forward
+        self.ff = nn.Sequential(nn.Linear(d_model, d_ff, bias=bias),
+                                get_activation_fn(activation),
+                                nn.Dropout(dropout),
+                                nn.Linear(d_ff, d_model, bias=bias))
+
+        # Add & Norm
+        self.dropout_ffn = nn.Dropout(dropout)
+        if "batch" in norm.lower():
+            self.norm_ffn = nn.Sequential(Transpose(1, 2), nn.BatchNorm1d(d_model), Transpose(1, 2))
+        else:
+            self.norm_ffn = nn.LayerNorm(d_model)
+
+        self.pre_norm = pre_norm
+        self.store_attn = store_attn
+
+    def forward(self, src: Tensor, prev: Optional[Tensor] = None, key_padding_mask: Optional[Tensor] = None,
+                attn_mask: Optional[Tensor] = None) -> Tensor:
+
+        # Multi-Head attention sublayer
+        if self.pre_norm:
+            src = self.norm_attn(src)
+        ## Multi-Head attention
+        if self.res_attention:
+            src2, attn, scores = self.self_attn(src, src, src, prev, key_padding_mask=key_padding_mask,
+                                                attn_mask=attn_mask)
+        else:
+            src2, attn = self.self_attn(src, src, src, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        if self.store_attn:
+            self.attn = attn
+        ## Add & Norm
+        src = src + self.dropout_attn(src2)  # Add: residual connection with residual dropout
+        if not self.pre_norm:
+            src = self.norm_attn(src)
+
+        # Feed-forward sublayer
+        if self.pre_norm:
+            src = self.norm_ffn(src)
+        ## Position-wise Feed-Forward
+        src2 = self.ff(src)
+        ## Add & Norm
+        src = src + self.dropout_ffn(src2)  # Add: residual connection with residual dropout
+        if not self.pre_norm:
+            src = self.norm_ffn(src)
+
+        if self.res_attention:
+            return src, scores
+        else:
+            return src
+
+
+class _MultiheadAttention(nn.Module):
+    def __init__(self, d_model, n_heads, d_k=None, d_v=None, res_attention=False, attn_dropout=0., proj_dropout=0.,
+                 qkv_bias=True, lsa=False):
+        """Multi Head Attention Layer
+        Input shape:
+            Q:       [batch_size (bs) x max_q_len x d_model]
+            K, V:    [batch_size (bs) x q_len x d_model]
+            mask:    [q_len x q_len]
+        """
+        super().__init__()
+        d_k = d_model // n_heads if d_k is None else d_k
+        d_v = d_model // n_heads if d_v is None else d_v
+
+        self.n_heads, self.d_k, self.d_v = n_heads, d_k, d_v
+
+        self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
+        self.W_K = nn.Linear(d_model, d_k * n_heads, bias=qkv_bias)
+        self.W_V = nn.Linear(d_model, d_v * n_heads, bias=qkv_bias)
+
+        # Scaled Dot-Product Attention (multiple heads)
+        self.res_attention = res_attention
+        self.sdp_attn = _ScaledDotProductAttention(d_model, n_heads, attn_dropout=attn_dropout,
+                                                   res_attention=self.res_attention, lsa=lsa)
+
+        # Poject output
+        self.to_out = nn.Sequential(nn.Linear(n_heads * d_v, d_model), nn.Dropout(proj_dropout))
+
+    def forward(self, Q: Tensor, K: Optional[Tensor] = None, V: Optional[Tensor] = None, prev: Optional[Tensor] = None,
+                key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
+
+        bs = Q.size(0)
+        if K is None: K = Q
+        if V is None: V = Q
+
+        # Linear (+ split in multiple heads)
+        q_s = self.W_Q(Q).view(bs, -1, self.n_heads, self.d_k).transpose(1,
+                                                                         2)  # q_s    : [bs x n_heads x max_q_len x d_k]
+        k_s = self.W_K(K).view(bs, -1, self.n_heads, self.d_k).permute(0, 2, 3,
+                                                                       1)  # k_s    : [bs x n_heads x d_k x q_len] - transpose(1,2) + transpose(2,3)
+        v_s = self.W_V(V).view(bs, -1, self.n_heads, self.d_v).transpose(1, 2)  # v_s    : [bs x n_heads x q_len x d_v]
+
+        # Apply Scaled Dot-Product Attention (multiple heads)
+        if self.res_attention:
+            output, attn_weights, attn_scores = self.sdp_attn(q_s, k_s, v_s, prev=prev,
+                                                              key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        else:
+            output, attn_weights = self.sdp_attn(q_s, k_s, v_s, key_padding_mask=key_padding_mask, attn_mask=attn_mask)
+        # output: [bs x n_heads x q_len x d_v], attn: [bs x n_heads x q_len x q_len], scores: [bs x n_heads x max_q_len x q_len]
+
+        # back to the original inputs dimensions
+        output = output.transpose(1, 2).contiguous().view(bs, -1,
+                                                          self.n_heads * self.d_v)  # output: [bs x q_len x n_heads * d_v]
+        output = self.to_out(output)
+
+        if self.res_attention:
+            return output, attn_weights, attn_scores
+        else:
+            return output, attn_weights
+
+
+class _ScaledDotProductAttention(nn.Module):
+    r"""Scaled Dot-Product Attention module (Attention is all you need by Vaswani et al., 2017) with optional residual attention from previous layer
+    (Realformer: Transformer likes residual attention by He et al, 2020) and locality self sttention (Vision Transformer for Small-Size Datasets
+    by Lee et al, 2021)"""
+
+    def __init__(self, d_model, n_heads, attn_dropout=0., res_attention=False, lsa=False):
+        super().__init__()
+        self.attn_dropout = nn.Dropout(attn_dropout)
+        self.res_attention = res_attention
+        head_dim = d_model // n_heads
+        self.scale = nn.Parameter(torch.tensor(head_dim ** -0.5), requires_grad=lsa)
+        self.lsa = lsa
+
+    def forward(self, q: Tensor, k: Tensor, v: Tensor, prev: Optional[Tensor] = None,
+                key_padding_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None):
+        '''
+        Input shape:
+            q               : [bs x n_heads x max_q_len x d_k]
+            k               : [bs x n_heads x d_k x seq_len]
+            v               : [bs x n_heads x seq_len x d_v]
+            prev            : [bs x n_heads x q_len x seq_len]
+            key_padding_mask: [bs x seq_len]
+            attn_mask       : [1 x seq_len x seq_len]
+        Output shape:
+            output:  [bs x n_heads x q_len x d_v]
+            attn   : [bs x n_heads x q_len x seq_len]
+            scores : [bs x n_heads x q_len x seq_len]
+        '''
+
+        # Scaled MatMul (q, k) - similarity scores for all pairs of positions in an input sequence
+        attn_scores = torch.matmul(q, k) * self.scale  # attn_scores : [bs x n_heads x max_q_len x q_len]
+
+        # Add pre-softmax attention scores from the previous layer (optional)
+        if prev is not None: attn_scores = attn_scores + prev
+
+        # Attention mask (optional)
+        if attn_mask is not None:  # attn_mask with shape [q_len x seq_len] - only used when q_len == seq_len
+            if attn_mask.dtype == torch.bool:
+                attn_scores.masked_fill_(attn_mask, -np.inf)
+            else:
+                attn_scores += attn_mask
+
+        # Key padding mask (optional)
+        if key_padding_mask is not None:  # mask with shape [bs x q_len] (only when max_w_len == q_len)
+            attn_scores.masked_fill_(key_padding_mask.unsqueeze(1).unsqueeze(2), -np.inf)
+
+        # normalize the attention weights
+        attn_weights = F.softmax(attn_scores, dim=-1)  # attn_weights   : [bs x n_heads x max_q_len x q_len]
+        attn_weights = self.attn_dropout(attn_weights)
+
+        # compute the new values given the attention weights
+        output = torch.matmul(attn_weights, v)  # output: [bs x n_heads x max_q_len x d_v]
+
+        if self.res_attention:
+            return output, attn_weights, attn_scores
+        else:
+            return output, attn_weights

code/new_realtime/models/third_party/patch_tst/layers/PatchTST_layers.py

@@ -0,0 +1,121 @@
+__all__ = ['Transpose', 'get_activation_fn', 'moving_avg', 'series_decomp', 'PositionalEncoding', 'SinCosPosEncoding', 'Coord2dPosEncoding', 'Coord1dPosEncoding', 'positional_encoding']           
+
+import torch
+from torch import nn
+import math
+
+class Transpose(nn.Module):
+    def __init__(self, *dims, contiguous=False): 
+        super().__init__()
+        self.dims, self.contiguous = dims, contiguous
+    def forward(self, x):
+        if self.contiguous: return x.transpose(*self.dims).contiguous()
+        else: return x.transpose(*self.dims)
+
+    
+def get_activation_fn(activation):
+    if callable(activation): return activation()
+    elif activation.lower() == "relu": return nn.ReLU()
+    elif activation.lower() == "gelu": return nn.GELU()
+    raise ValueError(f'{activation} is not available. You can use "relu", "gelu", or a callable') 
+    
+    
+# decomposition
+
+class moving_avg(nn.Module):
+    """
+    Moving average block to highlight the trend of time series
+    """
+    def __init__(self, kernel_size, stride):
+        super(moving_avg, self).__init__()
+        self.kernel_size = kernel_size
+        self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
+
+    def forward(self, x):
+        # padding on the both ends of time series
+        front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+        x = torch.cat([front, x, end], dim=1)
+        x = self.avg(x.permute(0, 2, 1))
+        x = x.permute(0, 2, 1)
+        return x
+
+
+class series_decomp(nn.Module):
+    """
+    Series decomposition block
+    """
+    def __init__(self, kernel_size):
+        super(series_decomp, self).__init__()
+        self.moving_avg = moving_avg(kernel_size, stride=1)
+
+    def forward(self, x):
+        moving_mean = self.moving_avg(x)
+        res = x - moving_mean
+        return res, moving_mean
+    
+    
+    
+# pos_encoding
+
+def PositionalEncoding(q_len, d_model, normalize=True):
+    pe = torch.zeros(q_len, d_model)
+    position = torch.arange(0, q_len).unsqueeze(1)
+    div_term = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model))
+    pe[:, 0::2] = torch.sin(position * div_term)
+    pe[:, 1::2] = torch.cos(position * div_term)
+    if normalize:
+        pe = pe - pe.mean()
+        pe = pe / (pe.std() * 10)
+    return pe
+
+SinCosPosEncoding = PositionalEncoding
+
+def Coord2dPosEncoding(q_len, d_model, exponential=False, normalize=True, eps=1e-3, verbose=False):
+    x = .5 if exponential else 1
+    i = 0
+    for i in range(100):
+        cpe = 2 * (torch.linspace(0, 1, q_len).reshape(-1, 1) ** x) * (torch.linspace(0, 1, d_model).reshape(1, -1) ** x) - 1
+        pv(f'{i:4.0f}  {x:5.3f}  {cpe.mean():+6.3f}', verbose)
+        if abs(cpe.mean()) <= eps: break
+        elif cpe.mean() > eps: x += .001
+        else: x -= .001
+        i += 1
+    if normalize:
+        cpe = cpe - cpe.mean()
+        cpe = cpe / (cpe.std() * 10)
+    return cpe
+
+def Coord1dPosEncoding(q_len, exponential=False, normalize=True):
+    cpe = (2 * (torch.linspace(0, 1, q_len).reshape(-1, 1)**(.5 if exponential else 1)) - 1)
+    if normalize:
+        cpe = cpe - cpe.mean()
+        cpe = cpe / (cpe.std() * 10)
+    return cpe
+
+def positional_encoding(pe, learn_pe, q_len, d_model):
+    # Positional encoding
+    if pe == None:
+        W_pos = torch.empty((q_len, d_model)) # pe = None and learn_pe = False can be used to measure impact of pe
+        nn.init.uniform_(W_pos, -0.02, 0.02)
+        learn_pe = False
+    elif pe == 'zero':
+        W_pos = torch.empty((q_len, 1))
+        nn.init.uniform_(W_pos, -0.02, 0.02)
+    elif pe == 'zeros':
+        W_pos = torch.empty((q_len, d_model))
+        nn.init.uniform_(W_pos, -0.02, 0.02)
+    elif pe == 'normal' or pe == 'gauss':
+        W_pos = torch.zeros((q_len, 1))
+        torch.nn.init.normal_(W_pos, mean=0.0, std=0.1)
+    elif pe == 'uniform':
+        W_pos = torch.zeros((q_len, 1))
+        nn.init.uniform_(W_pos, a=0.0, b=0.1)
+    elif pe == 'lin1d': W_pos = Coord1dPosEncoding(q_len, exponential=False, normalize=True)
+    elif pe == 'exp1d': W_pos = Coord1dPosEncoding(q_len, exponential=True, normalize=True)
+    elif pe == 'lin2d': W_pos = Coord2dPosEncoding(q_len, d_model, exponential=False, normalize=True)
+    elif pe == 'exp2d': W_pos = Coord2dPosEncoding(q_len, d_model, exponential=True, normalize=True)
+    elif pe == 'sincos': W_pos = PositionalEncoding(q_len, d_model, normalize=True)
+    else: raise ValueError(f"{pe} is not a valid pe (positional encoder. Available types: 'gauss'=='normal', \
+        'zeros', 'zero', uniform', 'lin1d', 'exp1d', 'lin2d', 'exp2d', 'sincos', None.)")
+    return nn.Parameter(W_pos, requires_grad=learn_pe)

code/new_realtime/models/third_party/patch_tst/layers/RevIN.py

@@ -0,0 +1,63 @@
+# code from https://github.com/ts-kim/RevIN, with minor modifications
+
+import torch
+import torch.nn as nn
+
+class RevIN(nn.Module):
+    def __init__(self, num_features: int, eps=1e-5, affine=True, subtract_last=False):
+        """
+        :param num_features: the number of features or channels
+        :param eps: a value added for numerical stability
+        :param affine: if True, RevIN has learnable affine parameters
+        """
+        super(RevIN, self).__init__()
+        self.num_features = num_features
+        self.eps = eps
+        self.affine = affine
+        self.subtract_last = subtract_last
+        if self.affine:
+            self._init_params()
+
+    def forward(self, x, mode:str):
+        if mode == 'norm':
+            self._get_statistics(x)
+            x = self._normalize(x)
+        elif mode == 'denorm':
+            x = self._denormalize(x)
+        else: raise NotImplementedError
+        return x
+
+    def _init_params(self):
+        # initialize RevIN params: (C,)
+        self.affine_weight = nn.Parameter(torch.ones(self.num_features))
+        self.affine_bias = nn.Parameter(torch.zeros(self.num_features))
+
+    def _get_statistics(self, x):
+        dim2reduce = tuple(range(1, x.ndim-1))
+        if self.subtract_last:
+            self.last = x[:,-1,:].unsqueeze(1)
+        else:
+            self.mean = torch.mean(x, dim=dim2reduce, keepdim=True).detach()
+        self.stdev = torch.sqrt(torch.var(x, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
+
+    def _normalize(self, x):
+        if self.subtract_last:
+            x = x - self.last
+        else:
+            x = x - self.mean
+        x = x / self.stdev
+        if self.affine:
+            x = x * self.affine_weight
+            x = x + self.affine_bias
+        return x
+
+    def _denormalize(self, x):
+        if self.affine:
+            x = x - self.affine_bias
+            x = x / (self.affine_weight + self.eps*self.eps)
+        x = x * self.stdev
+        if self.subtract_last:
+            x = x + self.last
+        else:
+            x = x + self.mean
+        return x

code/new_realtime/models/third_party/patch_tst/layers/SelfAttention_Family.py

@@ -0,0 +1,166 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import matplotlib.pyplot as plt
+
+import numpy as np
+import math
+from math import sqrt
+from utils.masking import TriangularCausalMask, ProbMask
+import os
+
+
+class FullAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
+        super(FullAttention, self).__init__()
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.output_attention = output_attention
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, E = queries.shape
+        _, S, _, D = values.shape
+        scale = self.scale or 1. / sqrt(E)
+
+        scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+
+        if self.mask_flag:
+            if attn_mask is None:
+                attn_mask = TriangularCausalMask(B, L, device=queries.device)
+
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        A = self.dropout(torch.softmax(scale * scores, dim=-1))
+        V = torch.einsum("bhls,bshd->blhd", A, values)
+
+        if self.output_attention:
+            return (V.contiguous(), A)
+        else:
+            return (V.contiguous(), None)
+
+
+class ProbAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0.1, output_attention=False):
+        super(ProbAttention, self).__init__()
+        self.factor = factor
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.output_attention = output_attention
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def _prob_QK(self, Q, K, sample_k, n_top):  # n_top: c*ln(L_q)
+        # Q [B, H, L, D]
+        B, H, L_K, E = K.shape
+        _, _, L_Q, _ = Q.shape
+
+        # calculate the sampled Q_K
+        K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)
+        index_sample = torch.randint(L_K, (L_Q, sample_k))  # real U = U_part(factor*ln(L_k))*L_q
+        K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :]
+        Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
+
+        # find the Top_k query with sparisty measurement
+        M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)
+        M_top = M.topk(n_top, sorted=False)[1]
+
+        # use the reduced Q to calculate Q_K
+        Q_reduce = Q[torch.arange(B)[:, None, None],
+                   torch.arange(H)[None, :, None],
+                   M_top, :]  # factor*ln(L_q)
+        Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1))  # factor*ln(L_q)*L_k
+
+        return Q_K, M_top
+
+    def _get_initial_context(self, V, L_Q):
+        B, H, L_V, D = V.shape
+        if not self.mask_flag:
+            # V_sum = V.sum(dim=-2)
+            V_sum = V.mean(dim=-2)
+            contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
+        else:  # use mask
+            assert (L_Q == L_V)  # requires that L_Q == L_V, i.e. for self-attention only
+            contex = V.cumsum(dim=-2)
+        return contex
+
+    def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
+        B, H, L_V, D = V.shape
+
+        if self.mask_flag:
+            attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        attn = torch.softmax(scores, dim=-1)  # nn.Softmax(dim=-1)(scores)
+
+        context_in[torch.arange(B)[:, None, None],
+        torch.arange(H)[None, :, None],
+        index, :] = torch.matmul(attn, V).type_as(context_in)
+        if self.output_attention:
+            attns = (torch.ones([B, H, L_V, L_V]) / L_V).type_as(attn).to(attn.device)
+            attns[torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], index, :] = attn
+            return (context_in, attns)
+        else:
+            return (context_in, None)
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L_Q, H, D = queries.shape
+        _, L_K, _, _ = keys.shape
+
+        queries = queries.transpose(2, 1)
+        keys = keys.transpose(2, 1)
+        values = values.transpose(2, 1)
+
+        U_part = self.factor * np.ceil(np.log(L_K)).astype('int').item()  # c*ln(L_k)
+        u = self.factor * np.ceil(np.log(L_Q)).astype('int').item()  # c*ln(L_q)
+
+        U_part = U_part if U_part < L_K else L_K
+        u = u if u < L_Q else L_Q
+
+        scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
+
+        # add scale factor
+        scale = self.scale or 1. / sqrt(D)
+        if scale is not None:
+            scores_top = scores_top * scale
+        # get the context
+        context = self._get_initial_context(values, L_Q)
+        # update the context with selected top_k queries
+        context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
+
+        return context.contiguous(), attn
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, attention, d_model, n_heads, d_keys=None,
+                 d_values=None):
+        super(AttentionLayer, self).__init__()
+
+        d_keys = d_keys or (d_model // n_heads)
+        d_values = d_values or (d_model // n_heads)
+
+        self.inner_attention = attention
+        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.value_projection = nn.Linear(d_model, d_values * n_heads)
+        self.out_projection = nn.Linear(d_values * n_heads, d_model)
+        self.n_heads = n_heads
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        H = self.n_heads
+
+        queries = self.query_projection(queries).view(B, L, H, -1)
+        keys = self.key_projection(keys).view(B, S, H, -1)
+        values = self.value_projection(values).view(B, S, H, -1)
+
+        out, attn = self.inner_attention(
+            queries,
+            keys,
+            values,
+            attn_mask
+        )
+        out = out.view(B, L, -1)
+
+        return self.out_projection(out), attn

code/new_realtime/models/third_party/patch_tst/layers/Transformer_EncDec.py

@@ -0,0 +1,131 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, c_in):
+        super(ConvLayer, self).__init__()
+        self.downConv = nn.Conv1d(in_channels=c_in,
+                                  out_channels=c_in,
+                                  kernel_size=3,
+                                  padding=2,
+                                  padding_mode='circular')
+        self.norm = nn.BatchNorm1d(c_in)
+        self.activation = nn.ELU()
+        self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        x = self.downConv(x.permute(0, 2, 1))
+        x = self.norm(x)
+        x = self.activation(x)
+        x = self.maxPool(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, attention, d_model, d_ff=None, dropout=0.1, activation="relu"):
+        super(EncoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.attention = attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, attn_mask=None):
+        new_x, attn = self.attention(
+            x, x, x,
+            attn_mask=attn_mask
+        )
+        x = x + self.dropout(new_x)
+
+        y = x = self.norm1(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+
+        return self.norm2(x + y), attn
+
+
+class Encoder(nn.Module):
+    def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
+        super(Encoder, self).__init__()
+        self.attn_layers = nn.ModuleList(attn_layers)
+        self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+        self.norm = norm_layer
+
+    def forward(self, x, attn_mask=None):
+        # x [B, L, D]
+        attns = []
+        if self.conv_layers is not None:
+            for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
+                x, attn = attn_layer(x, attn_mask=attn_mask)
+                x = conv_layer(x)
+                attns.append(attn)
+            x, attn = self.attn_layers[-1](x)
+            attns.append(attn)
+        else:
+            for attn_layer in self.attn_layers:
+                x, attn = attn_layer(x, attn_mask=attn_mask)
+                attns.append(attn)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x, attns
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
+                 dropout=0.1, activation="relu"):
+        super(DecoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        x = x + self.dropout(self.self_attention(
+            x, x, x,
+            attn_mask=x_mask
+        )[0])
+        x = self.norm1(x)
+
+        x = x + self.dropout(self.cross_attention(
+            x, cross, cross,
+            attn_mask=cross_mask
+        )[0])
+
+        y = x = self.norm2(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+
+        return self.norm3(x + y)
+
+
+class Decoder(nn.Module):
+    def __init__(self, layers, norm_layer=None, projection=None):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm_layer
+        self.projection = projection
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        for layer in self.layers:
+            x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        if self.projection is not None:
+            x = self.projection(x)
+        return x

code/new_realtime/models/third_party/patch_tst/models/Autoformer.py

@@ -0,0 +1,121 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from layers.Embed import DataEmbedding, DataEmbedding_wo_pos,DataEmbedding_wo_pos_temp,DataEmbedding_wo_temp
+from layers.AutoCorrelation import AutoCorrelation, AutoCorrelationLayer
+from layers.Autoformer_EncDec import Encoder, Decoder, EncoderLayer, DecoderLayer, my_Layernorm, series_decomp
+import math
+import numpy as np
+
+
+class Model(nn.Module):
+    """
+    Autoformer is the first method to achieve the series-wise connection,
+    with inherent O(LlogL) complexity
+    """
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.seq_len = configs.seq_len
+        self.label_len = configs.label_len
+        self.pred_len = configs.pred_len
+        self.output_attention = configs.output_attention
+
+        # Decomp
+        kernel_size = configs.moving_avg
+        self.decomp = series_decomp(kernel_size)
+
+        # Embedding
+        # The series-wise connection inherently contains the sequential information.
+        # Thus, we can discard the position embedding of transformers.
+        if configs.embed_type == 0:
+            self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 1:
+            self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 2:
+            self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+
+        elif configs.embed_type == 3:
+            self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 4:
+            self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        
+        # Encoder
+        self.encoder = Encoder(
+            [
+                EncoderLayer(
+                    AutoCorrelationLayer(
+                        AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
+                                        output_attention=configs.output_attention),
+                        configs.d_model, configs.n_heads),
+                    configs.d_model,
+                    configs.d_ff,
+                    moving_avg=configs.moving_avg,
+                    dropout=configs.dropout,
+                    activation=configs.activation
+                ) for l in range(configs.e_layers)
+            ],
+            norm_layer=my_Layernorm(configs.d_model)
+        )
+        # Decoder
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AutoCorrelationLayer(
+                        AutoCorrelation(True, configs.factor, attention_dropout=configs.dropout,
+                                        output_attention=False),
+                        configs.d_model, configs.n_heads),
+                    AutoCorrelationLayer(
+                        AutoCorrelation(False, configs.factor, attention_dropout=configs.dropout,
+                                        output_attention=False),
+                        configs.d_model, configs.n_heads),
+                    configs.d_model,
+                    configs.c_out,
+                    configs.d_ff,
+                    moving_avg=configs.moving_avg,
+                    dropout=configs.dropout,
+                    activation=configs.activation,
+                )
+                for l in range(configs.d_layers)
+            ],
+            norm_layer=my_Layernorm(configs.d_model),
+            projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
+        )
+
+    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
+                enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
+        # decomp init
+        mean = torch.mean(x_enc, dim=1).unsqueeze(1).repeat(1, self.pred_len, 1)
+        zeros = torch.zeros([x_dec.shape[0], self.pred_len, x_dec.shape[2]], device=x_enc.device)
+        seasonal_init, trend_init = self.decomp(x_enc)
+        # decoder input
+        trend_init = torch.cat([trend_init[:, -self.label_len:, :], mean], dim=1)
+        seasonal_init = torch.cat([seasonal_init[:, -self.label_len:, :], zeros], dim=1)
+        # enc
+        enc_out = self.enc_embedding(x_enc, x_mark_enc)
+        enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
+        # dec
+        dec_out = self.dec_embedding(seasonal_init, x_mark_dec)
+        seasonal_part, trend_part = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask,
+                                                 trend=trend_init)
+        # final
+        dec_out = trend_part + seasonal_part
+
+        if self.output_attention:
+            return dec_out[:, -self.pred_len:, :], attns
+        else:
+            return dec_out[:, -self.pred_len:, :]  # [B, L, D]

code/new_realtime/models/third_party/patch_tst/models/DLinear.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class moving_avg(nn.Module):
+    """
+    Moving average block to highlight the trend of time series
+    """
+    def __init__(self, kernel_size, stride):
+        super(moving_avg, self).__init__()
+        self.kernel_size = kernel_size
+        self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)
+
+    def forward(self, x):
+        # padding on the both ends of time series
+        front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
+        x = torch.cat([front, x, end], dim=1)
+        x = self.avg(x.permute(0, 2, 1))
+        x = x.permute(0, 2, 1)
+        return x
+
+
+class series_decomp(nn.Module):
+    """
+    Series decomposition block
+    """
+    def __init__(self, kernel_size):
+        super(series_decomp, self).__init__()
+        self.moving_avg = moving_avg(kernel_size, stride=1)
+
+    def forward(self, x):
+        moving_mean = self.moving_avg(x)
+        res = x - moving_mean
+        return res, moving_mean
+
+class Model(nn.Module):
+    """
+    Decomposition-Linear
+    """
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.seq_len = configs.seq_len
+        self.pred_len = configs.pred_len
+
+        # Decompsition Kernel Size
+        kernel_size = 25
+        self.decompsition = series_decomp(kernel_size)
+        self.individual = configs.individual
+        self.channels = configs.enc_in
+
+        if self.individual:
+            self.Linear_Seasonal = nn.ModuleList()
+            self.Linear_Trend = nn.ModuleList()
+            
+            for i in range(self.channels):
+                self.Linear_Seasonal.append(nn.Linear(self.seq_len,self.pred_len))
+                self.Linear_Trend.append(nn.Linear(self.seq_len,self.pred_len))
+
+                # Use this two lines if you want to visualize the weights
+                # self.Linear_Seasonal[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
+                # self.Linear_Trend[i].weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
+        else:
+            self.Linear_Seasonal = nn.Linear(self.seq_len,self.pred_len)
+            self.Linear_Trend = nn.Linear(self.seq_len,self.pred_len)
+            
+            # Use this two lines if you want to visualize the weights
+            # self.Linear_Seasonal.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
+            # self.Linear_Trend.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
+
+    def forward(self, x):
+        # x: [Batch, Input length, Channel]
+        seasonal_init, trend_init = self.decompsition(x)
+        seasonal_init, trend_init = seasonal_init.permute(0,2,1), trend_init.permute(0,2,1)
+        if self.individual:
+            seasonal_output = torch.zeros([seasonal_init.size(0),seasonal_init.size(1),self.pred_len],dtype=seasonal_init.dtype).to(seasonal_init.device)
+            trend_output = torch.zeros([trend_init.size(0),trend_init.size(1),self.pred_len],dtype=trend_init.dtype).to(trend_init.device)
+            for i in range(self.channels):
+                seasonal_output[:,i,:] = self.Linear_Seasonal[i](seasonal_init[:,i,:])
+                trend_output[:,i,:] = self.Linear_Trend[i](trend_init[:,i,:])
+        else:
+            seasonal_output = self.Linear_Seasonal(seasonal_init)
+            trend_output = self.Linear_Trend(trend_init)
+
+        x = seasonal_output + trend_output
+        return x.permute(0,2,1) # to [Batch, Output length, Channel]

code/new_realtime/models/third_party/patch_tst/models/Informer.py

@@ -0,0 +1,101 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from utils.masking import TriangularCausalMask, ProbMask
+from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
+from layers.SelfAttention_Family import FullAttention, ProbAttention, AttentionLayer
+from layers.Embed import DataEmbedding,DataEmbedding_wo_pos,DataEmbedding_wo_temp,DataEmbedding_wo_pos_temp
+import numpy as np
+
+
+class Model(nn.Module):
+    """
+    Informer with Propspare attention in O(LlogL) complexity
+    """
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.pred_len = configs.pred_len
+        self.output_attention = configs.output_attention
+
+        # Embedding
+        if configs.embed_type == 0:
+            self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                            configs.dropout)
+            self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                           configs.dropout)
+        elif configs.embed_type == 1:
+            self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 2:
+            self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+
+        elif configs.embed_type == 3:
+            self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 4:
+            self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        # Encoder
+        self.encoder = Encoder(
+            [
+                EncoderLayer(
+                    AttentionLayer(
+                        ProbAttention(False, configs.factor, attention_dropout=configs.dropout,
+                                      output_attention=configs.output_attention),
+                        configs.d_model, configs.n_heads),
+                    configs.d_model,
+                    configs.d_ff,
+                    dropout=configs.dropout,
+                    activation=configs.activation
+                ) for l in range(configs.e_layers)
+            ],
+            [
+                ConvLayer(
+                    configs.d_model
+                ) for l in range(configs.e_layers - 1)
+            ] if configs.distil else None,
+            norm_layer=torch.nn.LayerNorm(configs.d_model)
+        )
+        # Decoder
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AttentionLayer(
+                        ProbAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
+                        configs.d_model, configs.n_heads),
+                    AttentionLayer(
+                        ProbAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
+                        configs.d_model, configs.n_heads),
+                    configs.d_model,
+                    configs.d_ff,
+                    dropout=configs.dropout,
+                    activation=configs.activation,
+                )
+                for l in range(configs.d_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(configs.d_model),
+            projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
+        )
+
+    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
+                enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
+
+        enc_out = self.enc_embedding(x_enc, x_mark_enc)
+        enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
+
+        dec_out = self.dec_embedding(x_dec, x_mark_dec)
+        dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
+
+        if self.output_attention:
+            return dec_out[:, -self.pred_len:, :], attns
+        else:
+            return dec_out[:, -self.pred_len:, :]  # [B, L, D]

code/new_realtime/models/third_party/patch_tst/models/Linear.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class Model(nn.Module):
+    """
+    Just one Linear layer
+    """
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.seq_len = configs.seq_len
+        self.pred_len = configs.pred_len
+        self.Linear = nn.Linear(self.seq_len, self.pred_len)
+        # Use this line if you want to visualize the weights
+        # self.Linear.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
+
+    def forward(self, x):
+        # x: [Batch, Input length, Channel]
+        x = self.Linear(x.permute(0,2,1)).permute(0,2,1)
+        return x # [Batch, Output length, Channel]

code/new_realtime/models/third_party/patch_tst/models/NLinear.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+class Model(nn.Module):
+    """
+    Normalization-Linear
+    """
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.seq_len = configs.seq_len
+        self.pred_len = configs.pred_len
+        self.Linear = nn.Linear(self.seq_len, self.pred_len)
+        # Use this line if you want to visualize the weights
+        # self.Linear.weight = nn.Parameter((1/self.seq_len)*torch.ones([self.pred_len,self.seq_len]))
+
+    def forward(self, x):
+        # x: [Batch, Input length, Channel]
+        seq_last = x[:,-1:,:].detach()
+        x = x - seq_last
+        x = self.Linear(x.permute(0,2,1)).permute(0,2,1)
+        x = x + seq_last
+        return x # [Batch, Output length, Channel]

code/new_realtime/models/third_party/patch_tst/models/PatchTST.py

@@ -0,0 +1,127 @@
+__all__ = ['PatchTST']
+
+# Cell
+from typing import Callable, Optional
+import torch
+from torch import nn
+from torch import Tensor
+import torch.nn.functional as F
+import numpy as np
+
+from models.third_party.patch_tst.layers.PatchTST_backbone import PatchTST_backbone
+from models.third_party.patch_tst.layers.PatchTST_layers import series_decomp
+
+
+class Model(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int, configs, max_seq_len: Optional[int] = 1024,
+                 d_k: Optional[int] = None, d_v: Optional[int] = None,
+                 norm: str = 'BatchNorm', attn_dropout: float = 0.,
+                 act: str = "gelu", key_padding_mask: bool = 'auto', padding_var: Optional[int] = None,
+                 attn_mask: Optional[Tensor] = None, res_attention: bool = True,
+                 pre_norm: bool = False, store_attn: bool = False, pe: str = 'zeros', learn_pe: bool = True,
+                 pretrain_head: bool = False, head_type='flatten', verbose: bool = False, **kwargs):
+
+        super().__init__()
+
+        # load parameters
+        c_in = input_dim
+        context_window = configs['seq_len']
+        target_window = configs['pred_len']
+        dec_out = output_dim
+        seq_pred = configs["seq_pred"]
+
+        n_layers = configs['e_layers']
+        n_heads = configs['n_heads']
+        d_model = configs['d_model']
+        d_ff = configs['d_ff']
+        dropout = configs['dropout']
+        fc_dropout = configs['fc_dropout']
+        head_dropout = configs['head_dropout']
+
+        individual = configs['individual']
+
+        patch_len = configs['patch_len']
+        stride = configs['stride']
+        padding_patch = configs['padding_patch']
+
+        revin = configs['revin']
+        affine = configs['affine']
+        subtract_last = configs['subtract_last']
+
+        decomposition = configs['decomposition']
+        kernel_size = configs['kernel_size']
+
+        # model
+        self.decomposition = decomposition
+        if self.decomposition:
+            self.decomp_module = series_decomp(kernel_size)
+            self.model_trend = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
+                                                 # extras
+                                                 dec_out=dec_out,
+                                                 seq_pred=seq_pred,
+                                                 #
+                                                 patch_len=patch_len, stride=stride,
+                                                 max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
+                                                 n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
+                                                 attn_dropout=attn_dropout,
+                                                 dropout=dropout, act=act, key_padding_mask=key_padding_mask,
+                                                 padding_var=padding_var,
+                                                 attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
+                                                 store_attn=store_attn,
+                                                 pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout,
+                                                 head_dropout=head_dropout, padding_patch=padding_patch,
+                                                 pretrain_head=pretrain_head, head_type=head_type,
+                                                 individual=individual, revin=revin, affine=affine,
+                                                 subtract_last=subtract_last, verbose=verbose, **kwargs)
+            self.model_res = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
+                                               # extras
+                                               dec_out=dec_out,
+                                               seq_pred=seq_pred,
+                                               #
+                                               patch_len=patch_len, stride=stride,
+                                               max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
+                                               n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
+                                               attn_dropout=attn_dropout,
+                                               dropout=dropout, act=act, key_padding_mask=key_padding_mask,
+                                               padding_var=padding_var,
+                                               attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
+                                               store_attn=store_attn,
+                                               pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout,
+                                               head_dropout=head_dropout, padding_patch=padding_patch,
+                                               pretrain_head=pretrain_head, head_type=head_type, individual=individual,
+                                               revin=revin, affine=affine,
+                                               subtract_last=subtract_last, verbose=verbose, **kwargs)
+        else:
+            self.model = PatchTST_backbone(c_in=c_in, context_window=context_window, target_window=target_window,
+                                           # extras
+                                           dec_out=dec_out,
+                                           seq_pred=seq_pred,
+                                           #
+                                           patch_len=patch_len, stride=stride,
+                                           max_seq_len=max_seq_len, n_layers=n_layers, d_model=d_model,
+                                           n_heads=n_heads, d_k=d_k, d_v=d_v, d_ff=d_ff, norm=norm,
+                                           attn_dropout=attn_dropout,
+                                           dropout=dropout, act=act, key_padding_mask=key_padding_mask,
+                                           padding_var=padding_var,
+                                           attn_mask=attn_mask, res_attention=res_attention, pre_norm=pre_norm,
+                                           store_attn=store_attn,
+                                           pe=pe, learn_pe=learn_pe, fc_dropout=fc_dropout, head_dropout=head_dropout,
+                                           padding_patch=padding_patch,
+                                           pretrain_head=pretrain_head, head_type=head_type, individual=individual,
+                                           revin=revin, affine=affine,
+                                           subtract_last=subtract_last, verbose=verbose, **kwargs)
+
+    def forward(self, x):  # x: [Batch, Input length, Channel]
+        if self.decomposition:
+            res_init, trend_init = self.decomp_module(x)
+            res_init, trend_init = res_init.permute(0, 2, 1), trend_init.permute(0, 2,
+                                                                                 1)  # x: [Batch, Channel, Input length]
+            res = self.model_res(res_init)
+            trend = self.model_trend(trend_init)
+            x = res + trend
+            x = x.permute(0, 2, 1)  # x: [Batch, Input length, Channel]
+        else:
+            x = x.permute(0, 2, 1)  # x: [Batch, Channel, Input length]
+            x = self.model(x)
+            x = x.permute(0, 2, 1)  # x: [Batch, Input length, Channel]
+        return x

code/new_realtime/models/third_party/patch_tst/models/Stat_models.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from tqdm import tqdm
+import pmdarima as pm
+import threading
+from sklearn.ensemble import GradientBoostingRegressor
+
+class Naive_repeat(nn.Module):
+    def __init__(self, configs):
+        super(Naive_repeat, self).__init__()
+        self.pred_len = configs.pred_len
+        
+    def forward(self, x):
+        B,L,D = x.shape
+        x = x[:,-1,:].reshape(B,1,D).repeat(self.pred_len,axis=1)
+        return x # [B, L, D]
+
+class Naive_thread(threading.Thread):
+    def __init__(self,func,args=()):
+        super(Naive_thread,self).__init__()
+        self.func = func
+        self.args = args
+
+    def run(self):
+        self.results = self.func(*self.args)
+    
+    def return_result(self):
+        threading.Thread.join(self)
+        return self.results
+
+def _arima(seq,pred_len,bt,i):
+    model = pm.auto_arima(seq)
+    forecasts = model.predict(pred_len) 
+    return forecasts,bt,i
+
+class Arima(nn.Module):
+    """
+    Extremely slow, please sample < 0.1
+    """
+    def __init__(self, configs):
+        super(Arima, self).__init__()
+        self.pred_len = configs.pred_len
+        
+    def forward(self, x):
+        result = np.zeros([x.shape[0],self.pred_len,x.shape[2]])
+        threads = []
+        for bt,seqs in tqdm(enumerate(x)):
+            for i in range(seqs.shape[-1]):
+                seq = seqs[:,i]
+                one_seq = Naive_thread(func=_arima,args=(seq,self.pred_len,bt,i))
+                threads.append(one_seq)
+                threads[-1].start()
+        for every_thread in tqdm(threads):
+            forcast,bt,i = every_thread.return_result()
+            result[bt,:,i] = forcast
+
+        return result # [B, L, D]
+
+def _sarima(season,seq,pred_len,bt,i):
+    model = pm.auto_arima(seq, seasonal=True, m=season)
+    forecasts = model.predict(pred_len) 
+    return forecasts,bt,i
+
+class SArima(nn.Module):
+    """
+    Extremely extremely slow, please sample < 0.01
+    """
+    def __init__(self, configs):
+        super(SArima, self).__init__()
+        self.pred_len = configs.pred_len
+        self.seq_len = configs.seq_len
+        self.season = 24
+        if 'Ettm' in configs.data_path:
+            self.season = 12
+        elif 'ILI' in configs.data_path:
+            self.season = 1
+        if self.season >= self.seq_len:
+            self.season = 1
+
+    def forward(self, x):
+        result = np.zeros([x.shape[0],self.pred_len,x.shape[2]])
+        threads = []
+        for bt,seqs in tqdm(enumerate(x)):
+            for i in range(seqs.shape[-1]):
+                seq = seqs[:,i]
+                one_seq = Naive_thread(func=_sarima,args=(self.season,seq,self.pred_len,bt,i))
+                threads.append(one_seq)
+                threads[-1].start()
+        for every_thread in tqdm(threads):
+            forcast,bt,i = every_thread.return_result()
+            result[bt,:,i] = forcast
+        return result # [B, L, D]
+
+def _gbrt(seq,seq_len,pred_len,bt,i):
+    model = GradientBoostingRegressor()
+    model.fit(np.arange(seq_len).reshape(-1,1),seq.reshape(-1,1))
+    forecasts = model.predict(np.arange(seq_len,seq_len+pred_len).reshape(-1,1))  
+    return forecasts,bt,i
+
+class GBRT(nn.Module):
+    def __init__(self, configs):
+        super(GBRT, self).__init__()
+        self.seq_len = configs.seq_len
+        self.pred_len = configs.pred_len
+    
+    def forward(self, x):
+        result = np.zeros([x.shape[0],self.pred_len,x.shape[2]])
+        threads = []
+        for bt,seqs in tqdm(enumerate(x)):
+            for i in range(seqs.shape[-1]):
+                seq = seqs[:,i]
+                one_seq = Naive_thread(func=_gbrt,args=(seq,self.seq_len,self.pred_len,bt,i))
+                threads.append(one_seq)
+                threads[-1].start()
+        for every_thread in tqdm(threads):
+            forcast,bt,i = every_thread.return_result()
+            result[bt,:,i] = forcast
+        return result # [B, L, D]

code/new_realtime/models/third_party/patch_tst/models/Transformer.py

@@ -0,0 +1,94 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from layers.Transformer_EncDec import Decoder, DecoderLayer, Encoder, EncoderLayer, ConvLayer
+from layers.SelfAttention_Family import FullAttention, AttentionLayer
+from layers.Embed import DataEmbedding,DataEmbedding_wo_pos,DataEmbedding_wo_temp,DataEmbedding_wo_pos_temp
+import numpy as np
+
+
+class Model(nn.Module):
+    """
+    Vanilla Transformer with O(L^2) complexity
+    """
+    def __init__(self, configs):
+        super(Model, self).__init__()
+        self.pred_len = configs.pred_len
+        self.output_attention = configs.output_attention
+
+        # Embedding
+        if configs.embed_type == 0:
+            self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                            configs.dropout)
+            self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                           configs.dropout)
+        elif configs.embed_type == 1:
+            self.enc_embedding = DataEmbedding(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 2:
+            self.enc_embedding = DataEmbedding_wo_pos(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+
+        elif configs.embed_type == 3:
+            self.enc_embedding = DataEmbedding_wo_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        elif configs.embed_type == 4:
+            self.enc_embedding = DataEmbedding_wo_pos_temp(configs.enc_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+            self.dec_embedding = DataEmbedding_wo_pos_temp(configs.dec_in, configs.d_model, configs.embed, configs.freq,
+                                                    configs.dropout)
+        # Encoder
+        self.encoder = Encoder(
+            [
+                EncoderLayer(
+                    AttentionLayer(
+                        FullAttention(False, configs.factor, attention_dropout=configs.dropout,
+                                      output_attention=configs.output_attention), configs.d_model, configs.n_heads),
+                    configs.d_model,
+                    configs.d_ff,
+                    dropout=configs.dropout,
+                    activation=configs.activation
+                ) for l in range(configs.e_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(configs.d_model)
+        )
+        # Decoder
+        self.decoder = Decoder(
+            [
+                DecoderLayer(
+                    AttentionLayer(
+                        FullAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
+                        configs.d_model, configs.n_heads),
+                    AttentionLayer(
+                        FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
+                        configs.d_model, configs.n_heads),
+                    configs.d_model,
+                    configs.d_ff,
+                    dropout=configs.dropout,
+                    activation=configs.activation,
+                )
+                for l in range(configs.d_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(configs.d_model),
+            projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
+        )
+
+    def forward(self, x_enc, x_mark_enc, x_dec, x_mark_dec,
+                enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
+
+        enc_out = self.enc_embedding(x_enc, x_mark_enc)
+        enc_out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
+
+        dec_out = self.dec_embedding(x_dec, x_mark_dec)
+        dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
+
+        if self.output_attention:
+            return dec_out[:, -self.pred_len:, :], attns
+        else:
+            return dec_out[:, -self.pred_len:, :]  # [B, L, D]

code/new_realtime/models/third_party/tft/LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2021-2022 NVIDIA Corporation
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

code/new_realtime/models/third_party/tft/NOTICE

@@ -0,0 +1,5 @@
+TFT for PyTorch
+
+This repository includes software from https://github.com/google-research/google-research/tree/master/tft licensed under the Apache 2.0 License.
+
+This repository contains code from https://github.com/rwightman/pytorch-image-models/blob/master/timm/utils/model_ema.py under the Apache 2.0 License.

code/new_realtime/models/third_party/tft/README.md

@@ -0,0 +1,3 @@
+This folder contains code copied from NVIDIA's Temporal Fusion Transformer implementation, licensed under Apache 2.0.
+All rights belong to NVIDIA Corporation.
+Modifications are noted in the file headers.

code/new_realtime/models/third_party/tft/tft_model.py

@@ -0,0 +1,525 @@
+# Copyright (c) 2021-2022, NVIDIA CORPORATION. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#           http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Modified by Alexander Blank, 2025.
+# Modifications:
+# - added support for multiple outputs
+# - added support for mode configurable targets
+# - added support for single dimension, non-quantile outputs
+# - added support for target agnostic predictions, for cases, where the target does not become known after prediction
+
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch import Tensor
+from torch.nn.parameter import UninitializedParameter
+from typing import Dict, Tuple, Optional, List
+
+
+MAKE_CONVERT_COMPATIBLE = os.environ.get("TFT_SCRIPTING", None) is not None
+from torch.nn import LayerNorm
+
+
+class MaybeLayerNorm(nn.Module):
+    def __init__(self, output_size, hidden_size, eps):
+        super().__init__()
+        if output_size and output_size == 1:
+            self.ln = nn.Identity()
+        else:
+            self.ln = LayerNorm(output_size if output_size else hidden_size, eps=eps)
+
+    def forward(self, x):
+        return self.ln(x)
+
+
+class GLU(nn.Module):
+    def __init__(self, hidden_size, output_size):
+        super().__init__()
+        self.lin = nn.Linear(hidden_size, output_size * 2)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.lin(x)
+        x = F.glu(x)
+        return x
+
+
+class GRN(nn.Module):
+    def __init__(self,
+                 input_size,
+                 hidden_size,
+                 output_size=None,
+                 context_hidden_size=None,
+                 dropout=0.0, ):
+        super().__init__()
+        self.layer_norm = MaybeLayerNorm(output_size, hidden_size, eps=1e-3)
+        self.lin_a = nn.Linear(input_size, hidden_size)
+        if context_hidden_size is not None:
+            self.lin_c = nn.Linear(context_hidden_size, hidden_size, bias=False)
+        else:
+            self.lin_c = nn.Identity()
+        self.lin_i = nn.Linear(hidden_size, hidden_size)
+        self.glu = GLU(hidden_size, output_size if output_size else hidden_size)
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(input_size, output_size) if output_size else None
+
+    def forward(self, a: Tensor, c: Optional[Tensor] = None):
+        x = self.lin_a(a)
+        if c is not None:
+            x = x + self.lin_c(c).unsqueeze(1)
+        x = F.elu(x)
+        x = self.lin_i(x)
+        x = self.dropout(x)
+        x = self.glu(x)
+        y = a if self.out_proj is None else self.out_proj(a)
+        x = x + y
+        return self.layer_norm(x)
+
+    # @torch.jit.script #Currently broken with autocast
+
+
+def fused_pointwise_linear_v1(x, a, b):
+    out = torch.mul(x.unsqueeze(-1), a)
+    out = out + b
+    return out
+
+
+@torch.jit.script
+def fused_pointwise_linear_v2(x, a, b):
+    out = x.unsqueeze(3) * a
+    out = out + b
+    return out
+
+
+class TFTEmbedding(nn.Module):
+    def __init__(self, config, initialize_cont_params=True):
+        # initialize_cont_params=False prevents form initializing parameters inside this class
+        # so they can be lazily initialized in LazyEmbedding module
+        super().__init__()
+        self.s_cat_inp_lens = config.static_categorical_inp_lens
+        self.t_cat_k_inp_lens = config.temporal_known_categorical_inp_lens
+        self.t_cat_o_inp_lens = config.temporal_observed_categorical_inp_lens
+        self.s_cont_inp_size = config.static_continuous_inp_size
+        self.t_cont_k_inp_size = config.temporal_known_continuous_inp_size
+        self.t_cont_o_inp_size = config.temporal_observed_continuous_inp_size
+        self.t_tgt_size = config.temporal_target_size
+
+        self.hidden_size = config.hidden_size
+
+        # There are 7 types of input:
+        # 1. Static categorical
+        # 2. Static continuous
+        # 3. Temporal known a priori categorical
+        # 4. Temporal known a priori continuous
+        # 5. Temporal observed categorical
+        # 6. Temporal observed continuous
+        # 7. Temporal observed targets (time series obseved so far)
+
+        self.s_cat_embed = nn.ModuleList([
+            nn.Embedding(n, self.hidden_size) for n in self.s_cat_inp_lens]) if self.s_cat_inp_lens else None
+        self.t_cat_k_embed = nn.ModuleList([
+            nn.Embedding(n, self.hidden_size) for n in self.t_cat_k_inp_lens]) if self.t_cat_k_inp_lens else None
+        self.t_cat_o_embed = nn.ModuleList([
+            nn.Embedding(n, self.hidden_size) for n in self.t_cat_o_inp_lens]) if self.t_cat_o_inp_lens else None
+
+        if initialize_cont_params:
+            self.s_cont_embedding_vectors = nn.Parameter(
+                torch.Tensor(self.s_cont_inp_size, self.hidden_size)) if self.s_cont_inp_size else None
+            self.t_cont_k_embedding_vectors = nn.Parameter(
+                torch.Tensor(self.t_cont_k_inp_size, self.hidden_size)) if self.t_cont_k_inp_size else None
+            self.t_cont_o_embedding_vectors = nn.Parameter(
+                torch.Tensor(self.t_cont_o_inp_size, self.hidden_size)) if self.t_cont_o_inp_size else None
+            self.t_tgt_embedding_vectors = nn.Parameter(torch.Tensor(self.t_tgt_size, self.hidden_size))
+
+            self.s_cont_embedding_bias = nn.Parameter(
+                torch.zeros(self.s_cont_inp_size, self.hidden_size)) if self.s_cont_inp_size else None
+            self.t_cont_k_embedding_bias = nn.Parameter(
+                torch.zeros(self.t_cont_k_inp_size, self.hidden_size)) if self.t_cont_k_inp_size else None
+            self.t_cont_o_embedding_bias = nn.Parameter(
+                torch.zeros(self.t_cont_o_inp_size, self.hidden_size)) if self.t_cont_o_inp_size else None
+            self.t_tgt_embedding_bias = nn.Parameter(torch.zeros(self.t_tgt_size, self.hidden_size))
+
+            self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.s_cont_embedding_vectors is not None:
+            torch.nn.init.xavier_normal_(self.s_cont_embedding_vectors)
+            torch.nn.init.zeros_(self.s_cont_embedding_bias)
+        if self.t_cont_k_embedding_vectors is not None:
+            torch.nn.init.xavier_normal_(self.t_cont_k_embedding_vectors)
+            torch.nn.init.zeros_(self.t_cont_k_embedding_bias)
+        if self.t_cont_o_embedding_vectors is not None:
+            torch.nn.init.xavier_normal_(self.t_cont_o_embedding_vectors)
+            torch.nn.init.zeros_(self.t_cont_o_embedding_bias)
+        if self.t_tgt_embedding_vectors is not None:
+            torch.nn.init.xavier_normal_(self.t_tgt_embedding_vectors)
+            torch.nn.init.zeros_(self.t_tgt_embedding_bias)
+        if self.s_cat_embed is not None:
+            for module in self.s_cat_embed:
+                module.reset_parameters()
+        if self.t_cat_k_embed is not None:
+            for module in self.t_cat_k_embed:
+                module.reset_parameters()
+        if self.t_cat_o_embed is not None:
+            for module in self.t_cat_o_embed:
+                module.reset_parameters()
+
+    def _apply_embedding(self,
+                         cat: Optional[Tensor],
+                         cont: Optional[Tensor],
+                         cat_emb: Optional[nn.ModuleList],
+                         cont_emb: Tensor,
+                         cont_bias: Tensor,
+                         ) -> Tuple[Optional[Tensor], Optional[Tensor]]:
+        e_cat = torch.stack([embed(cat[..., i]) for i, embed in enumerate(cat_emb)],
+                            dim=-2) if cat is not None else None
+        if cont is not None:
+            # the line below is equivalent to following einsums
+            # e_cont = torch.einsum('btf,fh->bthf', cont, cont_emb)
+            # e_cont = torch.einsum('bf,fh->bhf', cont, cont_emb)
+            if MAKE_CONVERT_COMPATIBLE:
+                e_cont = torch.mul(cont.unsqueeze(-1), cont_emb)
+                e_cont = e_cont + cont_bias
+            else:
+                e_cont = fused_pointwise_linear_v1(cont, cont_emb, cont_bias)
+        else:
+            e_cont = None
+
+        if e_cat is not None and e_cont is not None:
+            return torch.cat([e_cat, e_cont], dim=-2)
+        elif e_cat is not None:
+            return e_cat
+        elif e_cont is not None:
+            return e_cont
+        else:
+            return None
+
+    def forward(self, x: Dict[str, Tensor], use_target: bool = False):
+        # Extract inputs
+        s_cat_inp = x.get('s_cat', None)
+        s_cont_inp = x.get('s_cont', None)
+        t_cat_k_inp = x.get('k_cat', None)
+        t_cont_k_inp = x.get('k_cont', None)
+        t_cat_o_inp = x.get('o_cat', None)
+        t_cont_o_inp = x.get('o_cont', None)
+
+        # Only use target if teacher forcing is enabled.
+        # When disabled, we ignore target values.
+        if use_target:
+            t_tgt_obs = x['target']  # Must be present when using teacher forcing
+        else:
+            t_tgt_obs = None
+
+        # For static inputs, take the first timestep
+        s_cat_inp = s_cat_inp[:, 0, :] if s_cat_inp is not None else None
+        s_cont_inp = s_cont_inp[:, 0, :] if s_cont_inp is not None else None
+
+        # Apply embeddings for static and known/observed temporal features
+        s_inp = self._apply_embedding(s_cat_inp,
+                                      s_cont_inp,
+                                      self.s_cat_embed,
+                                      self.s_cont_embedding_vectors,
+                                      self.s_cont_embedding_bias)
+        t_known_inp = self._apply_embedding(t_cat_k_inp,
+                                            t_cont_k_inp,
+                                            self.t_cat_k_embed,
+                                            self.t_cont_k_embedding_vectors,
+                                            self.t_cont_k_embedding_bias)
+        t_observed_inp = self._apply_embedding(t_cat_o_inp,
+                                               t_cont_o_inp,
+                                               self.t_cat_o_embed,
+                                               self.t_cont_o_embedding_vectors,
+                                               self.t_cont_o_embedding_bias)
+        # Compute the target embedding only if teacher forcing is enabled.
+        if use_target and t_tgt_obs is not None:
+            if MAKE_CONVERT_COMPATIBLE:
+                t_observed_tgt = torch.matmul(t_tgt_obs.unsqueeze(3).unsqueeze(4),
+                                              self.t_tgt_embedding_vectors.unsqueeze(1)).squeeze(3)
+                t_observed_tgt = t_observed_tgt + self.t_tgt_embedding_bias
+            else:
+                t_observed_tgt = fused_pointwise_linear_v2(t_tgt_obs,
+                                                           self.t_tgt_embedding_vectors,
+                                                           self.t_tgt_embedding_bias)
+        else:
+            t_observed_tgt = None
+
+        return s_inp, t_known_inp, t_observed_inp, t_observed_tgt
+
+
+class LazyEmbedding(nn.modules.lazy.LazyModuleMixin, TFTEmbedding):
+    cls_to_become = TFTEmbedding
+
+    def __init__(self, config):
+        super().__init__(config, initialize_cont_params=False)
+
+        if config.static_continuous_inp_size:
+            self.s_cont_embedding_vectors = UninitializedParameter()
+            self.s_cont_embedding_bias = UninitializedParameter()
+        else:
+            self.s_cont_embedding_vectors = None
+            self.s_cont_embedding_bias = None
+
+        if config.temporal_known_continuous_inp_size:
+            self.t_cont_k_embedding_vectors = UninitializedParameter()
+            self.t_cont_k_embedding_bias = UninitializedParameter()
+        else:
+            self.t_cont_k_embedding_vectors = None
+            self.t_cont_k_embedding_bias = None
+
+        if config.temporal_observed_continuous_inp_size:
+            self.t_cont_o_embedding_vectors = UninitializedParameter()
+            self.t_cont_o_embedding_bias = UninitializedParameter()
+        else:
+            self.t_cont_o_embedding_vectors = None
+            self.t_cont_o_embedding_bias = None
+
+        self.t_tgt_embedding_vectors = UninitializedParameter()
+        self.t_tgt_embedding_bias = UninitializedParameter()
+
+    def initialize_parameters(self, x):
+        if self.has_uninitialized_params():
+            s_cont_inp = x.get('s_cont', None)
+            t_cont_k_inp = x.get('k_cont', None)
+            t_cont_o_inp = x.get('o_cont', None)
+            t_tgt_obs = x['target']  # Has to be present
+
+            if s_cont_inp is not None:
+                self.s_cont_embedding_vectors.materialize((s_cont_inp.shape[-1], self.hidden_size))
+                self.s_cont_embedding_bias.materialize((s_cont_inp.shape[-1], self.hidden_size))
+
+            if t_cont_k_inp is not None:
+                self.t_cont_k_embedding_vectors.materialize((t_cont_k_inp.shape[-1], self.hidden_size))
+                self.t_cont_k_embedding_bias.materialize((t_cont_k_inp.shape[-1], self.hidden_size))
+
+            if t_cont_o_inp is not None:
+                self.t_cont_o_embedding_vectors.materialize((t_cont_o_inp.shape[-1], self.hidden_size))
+                self.t_cont_o_embedding_bias.materialize((t_cont_o_inp.shape[-1], self.hidden_size))
+
+            self.t_tgt_embedding_vectors.materialize((t_tgt_obs.shape[-1], self.hidden_size))
+            self.t_tgt_embedding_bias.materialize((t_tgt_obs.shape[-1], self.hidden_size))
+
+            self.reset_parameters()
+
+    # def forward(self, x: Dict[str, Tensor], use_target: bool = True):
+    #     return super().forward(x, use_target=use_target)
+
+
+class VariableSelectionNetwork(nn.Module):
+    def __init__(self, config, num_inputs):
+        super().__init__()
+        self.joint_grn = GRN(config.hidden_size * num_inputs, config.hidden_size, output_size=num_inputs,
+                             context_hidden_size=config.hidden_size)
+        self.var_grns = nn.ModuleList(
+            [GRN(config.hidden_size, config.hidden_size, dropout=config.dropout) for _ in range(num_inputs)])
+
+    def forward(self, x: Tensor, context: Optional[Tensor] = None):
+        Xi = torch.flatten(x, start_dim=-2)
+        grn_outputs = self.joint_grn(Xi, c=context)
+        sparse_weights = F.softmax(grn_outputs, dim=-1)
+        transformed_embed_list = [m(x[..., i, :]) for i, m in enumerate(self.var_grns)]
+        transformed_embed = torch.stack(transformed_embed_list, dim=-1)
+        # the line below performs batched matrix vector multiplication
+        # for temporal features it's bthf,btf->bth
+        # for static features it's bhf,bf->bh
+        variable_ctx = torch.matmul(transformed_embed, sparse_weights.unsqueeze(-1)).squeeze(-1)
+
+        return variable_ctx, sparse_weights
+
+
+class StaticCovariateEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.vsn = VariableSelectionNetwork(config, config.num_static_vars)
+        self.context_grns = nn.ModuleList(
+            [GRN(config.hidden_size, config.hidden_size, dropout=config.dropout) for _ in range(4)])
+
+    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+        variable_ctx, sparse_weights = self.vsn(x)
+
+        # Context vectors:
+        # variable selection context
+        # enrichment context
+        # state_c context
+        # state_h context
+        cs, ce, ch, cc = [m(variable_ctx) for m in self.context_grns]
+
+        return cs, ce, ch, cc
+
+
+class InterpretableMultiHeadAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.n_head = config.n_head
+        assert config.hidden_size % config.n_head == 0
+        self.d_head = config.hidden_size // config.n_head
+        self.qkv_linears = nn.Linear(config.hidden_size, (2 * self.n_head + 1) * self.d_head, bias=False)
+        self.out_proj = nn.Linear(self.d_head, config.hidden_size, bias=False)
+        self.attn_dropout = nn.Dropout(config.attn_dropout)
+        self.out_dropout = nn.Dropout(config.dropout)
+        self.scale = self.d_head ** -0.5
+        self.register_buffer("_mask",
+                             torch.triu(torch.full((config.example_length, config.example_length), float('-inf')),
+                                        1).unsqueeze(0))
+
+    def forward(self, x: Tensor) -> Tuple[Tensor, Tensor]:
+        bs, t, h_size = x.shape
+        qkv = self.qkv_linears(x)
+        q, k, v = qkv.split((self.n_head * self.d_head, self.n_head * self.d_head, self.d_head), dim=-1)
+        q = q.view(bs, t, self.n_head, self.d_head)
+        k = k.view(bs, t, self.n_head, self.d_head)
+        v = v.view(bs, t, self.d_head)
+
+        # attn_score = torch.einsum('bind,bjnd->bnij', q, k)
+        attn_score = torch.matmul(q.permute((0, 2, 1, 3)), k.permute((0, 2, 3, 1)))
+        attn_score.mul_(self.scale)
+
+        attn_score = attn_score + self._mask
+
+        attn_prob = F.softmax(attn_score, dim=3)
+        attn_prob = self.attn_dropout(attn_prob)
+
+        # attn_vec = torch.einsum('bnij,bjd->bnid', attn_prob, v)
+        attn_vec = torch.matmul(attn_prob, v.unsqueeze(1))
+        m_attn_vec = torch.mean(attn_vec, dim=1)
+        out = self.out_proj(m_attn_vec)
+        out = self.out_dropout(out)
+
+        return out, attn_prob
+
+
+class TFTBack(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.encoder_length = config.encoder_length
+        self.history_vsn = VariableSelectionNetwork(config, config.num_historic_vars)
+        self.history_encoder = nn.LSTM(config.hidden_size, config.hidden_size, batch_first=True)
+        self.future_vsn = VariableSelectionNetwork(config, config.num_future_vars)
+        self.future_encoder = nn.LSTM(config.hidden_size, config.hidden_size, batch_first=True)
+
+        self.input_gate = GLU(config.hidden_size, config.hidden_size)
+        self.input_gate_ln = LayerNorm(config.hidden_size, eps=1e-3)
+
+        self.enrichment_grn = GRN(config.hidden_size,
+                                  config.hidden_size,
+                                  context_hidden_size=config.hidden_size,
+                                  dropout=config.dropout)
+        self.attention = InterpretableMultiHeadAttention(config)
+        self.attention_gate = GLU(config.hidden_size, config.hidden_size)
+        self.attention_ln = LayerNorm(config.hidden_size, eps=1e-3)
+
+        self.positionwise_grn = GRN(config.hidden_size,
+                                    config.hidden_size,
+                                    dropout=config.dropout)
+
+        self.decoder_gate = GLU(config.hidden_size, config.hidden_size)
+        self.decoder_ln = LayerNorm(config.hidden_size, eps=1e-3)
+
+        self.quantiles = config.quantiles
+        self.target_size = config.target_size
+        if self.quantiles is not None:
+            self.output = nn.Linear(config.hidden_size, len(config.quantiles) * config.target_size)
+        else:
+            self.output = nn.Linear(config.hidden_size, config.target_size)
+
+    def forward(self, historical_inputs, cs, ch, cc, ce, future_inputs):
+        historical_features, _ = self.history_vsn(historical_inputs, cs)
+        history, state = self.history_encoder(historical_features, (ch, cc))
+        future_features, _ = self.future_vsn(future_inputs, cs)
+        future, _ = self.future_encoder(future_features, state)
+        torch.cuda.synchronize()
+
+        # skip connection
+        input_embedding = torch.cat([historical_features, future_features], dim=1)
+        temporal_features = torch.cat([history, future], dim=1)
+        temporal_features = self.input_gate(temporal_features)
+        temporal_features = temporal_features + input_embedding
+        temporal_features = self.input_gate_ln(temporal_features)
+
+        # Static enrichment
+        enriched = self.enrichment_grn(temporal_features, c=ce)
+
+        # Temporal self attention
+        x, _ = self.attention(enriched)
+
+        # Don't compute hictorical quantiles
+        x = x[:, self.encoder_length:, :]
+        temporal_features = temporal_features[:, self.encoder_length:, :]
+        enriched = enriched[:, self.encoder_length:, :]
+
+        x = self.attention_gate(x)
+        x = x + enriched
+        x = self.attention_ln(x)
+
+        # Position-wise feed-forward
+        x = self.positionwise_grn(x)
+
+        # Final skip connection
+        x = self.decoder_gate(x)
+        x = x + temporal_features
+        x = self.decoder_ln(x)
+
+        out = self.output(x)
+        if self.quantiles is not None:
+            # Reshape to [batch, time, target_size, n_quantiles]
+            out = out.view(out.size(0), out.size(1), self.target_size, len(self.quantiles))
+        else:
+            # Reshape to [batch, time, target_size]
+            out = out.view(out.size(0), out.size(1), self.target_size)
+
+        return out
+
+
+class TemporalFusionTransformer(nn.Module):
+    """
+    Implementation of https://arxiv.org/abs/1912.09363
+    """
+
+    def __init__(self, config):
+        super().__init__()
+
+        if hasattr(config, 'model'):
+            config = config.model
+
+        self.encoder_length = config.encoder_length  # this determines from how distant past we want to use data from
+
+        # self.embedding = LazyEmbedding(config)
+        self.embedding = TFTEmbedding(config)
+        self.static_encoder = StaticCovariateEncoder(config)
+        # if MAKE_CONVERT_COMPATIBLE:
+        self.TFTpart2 = TFTBack(config)
+        # else:
+        #     self.TFTpart2 = torch.jit.script(TFTBack(config))
+
+    def forward(self, x: Dict[str, Tensor]) -> Tensor:
+        # Call embedding with use_target=False to skip target features entirely.
+        s_inp, t_known_inp, t_observed_inp, t_observed_tgt = self.embedding(x, use_target=False)
+
+        # Compute static context
+        cs, ce, ch, cc = self.static_encoder(s_inp)
+        ch, cc = ch.unsqueeze(0), cc.unsqueeze(0)  # Initialize LSTM states
+
+        # Build historical inputs without teacher-forced targets.
+        # Include observed features if available, and the known inputs.
+        historical_inputs = []
+        if t_observed_inp is not None:
+            historical_inputs.append(t_observed_inp[:, :self.encoder_length, :])
+        historical_inputs.append(t_known_inp[:, :self.encoder_length, :])
+        historical_inputs = torch.cat(historical_inputs, dim=-2)
+
+        # Future inputs remain the same
+        future_inputs = t_known_inp[:, self.encoder_length:]
+        return self.TFTpart2(historical_inputs, cs, ch, cc, ce, future_inputs)

code/new_realtime/models/transformer.py

@@ -0,0 +1,40 @@
+import math
+
+import torch
+from torch import nn
+
+
+class TransformerModel(nn.Module):
+    def __init__(self, input_dim: int,
+                 output_dim: int,
+                 seq_len: int,
+                 embed_dim: int,
+                 num_heads: int,
+                 num_enc_layers: int):
+        super().__init__()
+        self.input_proj = nn.Linear(input_dim, embed_dim)
+
+        # Compute positional embedding ONCE at init
+        pe = self._get_sinusoidal_embedding(seq_len, embed_dim)  # (seq_len, embed_dim)
+        self.register_buffer('pos_embed', pe.unsqueeze(0))  # (1, seq_len, embed_dim)
+
+        encoder_layer = nn.TransformerEncoderLayer(embed_dim, num_heads)
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_enc_layers)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Linear(embed_dim, output_dim)
+
+    def _get_sinusoidal_embedding(self, seq_len, embed_dim):
+        position = torch.arange(0, seq_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim))
+        pe = torch.zeros(seq_len, embed_dim)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe  # (seq_len, embed_dim)
+
+    def forward(self, x):
+        # x: (B, seq_len, input_dim)
+        x = self.input_proj(x) + self.pos_embed[:, :x.size(1), :]  # broadcasting
+        x = x.permute(1, 0, 2)  # (S, B, E)
+        enc = self.encoder(x)  # (S, B, E)
+        pooled = enc.mean(0)  # (B, E)
+        return self.head(pooled)

code/new_realtime/models/utils.py

@@ -0,0 +1,137 @@
+import os
+from typing import Callable
+import math
+
+import numpy as np
+from bson import ObjectId
+import torch
+from torch import nn
+
+from utils.data_utils import get_collated_batch_for_key
+
+
+def simple_get_y(collated_batch):
+    """
+    Get y from collated batch
+    Args:
+        collated_batch: collated batch to get y from
+
+    Returns:
+        np.ndarray: y
+    """
+    return collated_batch[1]
+
+
+def simple_model_save(model: nn.Module,
+                      training_configuration: dict) -> None:
+    """
+    Save model to disk
+    Args:
+        model: model to save
+        training_configuration: training configuration
+    Returns:
+        None
+    """
+    model_path = os.path.join(training_configuration["training_dir"], "model.pt")
+    torch.save(model.state_dict(), model_path)
+
+
+def simple_model_load(model_configuration: dict,
+                      training_configuration: dict,
+                      sample_key: str | ObjectId,
+                      device: str,
+                      model_creation_fn: Callable,
+                      *args, **kwargs) -> nn.Module:
+    """
+    Load model from disk
+    Args:
+        model_configuration: model configuration
+        training_configuration: training configuration
+        sample_key: sample key to get sample data batch with
+        device: device to load model on
+        model_creation_fn: function to create model
+    Returns:
+        model: loaded model
+    """
+    training_dir = training_configuration["training_dir"]
+    model_state_path = os.path.join(training_dir, "model.pt")
+    model = model_creation_fn(model_configuration,
+                              sample_key, *args, **kwargs)
+    model.to(device)
+    model.load_state_dict(torch.load(model_state_path))
+    return model
+
+
+def simple_model_creation(model_configuration: dict,
+                          sample_key: str | ObjectId,
+                          lmdb_env=None) -> nn.Module:
+    """
+    Create model from configuration
+    Args:
+        model_configuration: model configuration
+        sample_key: sample key to get sample data batch with
+        lmdb_env: LMDB environment to use for getting sample data batch
+
+    Returns:
+        model: created model
+    """
+    sample_item = get_collated_batch_for_key(sample_key, model_configuration, lmdb_env=lmdb_env)
+    input_size = sample_item[0].shape[2]
+    output_size = sample_item[1].shape[2]
+    model_class = model_configuration["model_class"]
+    cnn_model = model_class(input_dim=input_size,
+                            output_dim=output_size,
+                            **model_configuration["model_parameters"], )
+    return cnn_model
+
+
+def simple_x_y_predict(model: nn.Module,
+                       collated_batch: tuple[torch.Tensor, torch.Tensor],
+                       device: str,
+                       batch_size: int,
+                       *args, **kwargs) -> torch.Tensor:
+    """
+    Predict y from x
+    Args:
+        model: model to use for prediction, must have 'predict' method
+        collated_batch: collated data batch to predict with,
+            must be a tuple of (x, y) where x is the input data and y is the target data
+        device: device to use for prediction
+        *args: additional arguments to pass to the model's predict method
+        **kwargs: additional keyword arguments to pass to the model's predict method
+
+    Returns:
+        torch.Tensor: predicted y as numpy array on CPU
+    """
+
+    model.eval()
+    x = collated_batch[0].to(device).float()
+    batches = list()
+    num_batches = math.ceil(len(x) / batch_size)
+    for i in range(num_batches):
+        start = i * batch_size
+        end = (i + 1) * batch_size
+        if end > len(x):
+            end = len(x)
+        batch_slice = x[start:end]
+        batches.append(batch_slice)
+
+    preds = list()
+    with torch.no_grad():
+        for batch in batches:
+            batch = batch.to(device).float()
+            pred = model(batch, *args, **kwargs)
+            preds.append(pred.cpu().numpy())
+    return np.concatenate(preds)
+
+
+def simple_get_y(collated_batch: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor:
+    """
+    Get y from collated batch
+    Args:
+        collated_batch: collated batch to get y from
+
+    Returns:
+        torch.Tensor: y
+    """
+    return collated_batch[1]