added code

code/new_realtime/.Rhistory

@@ -0,0 +1,11 @@
+install.package("cpm")
+install.packages("cpm")
+install.packages("cpm")
+install.packages("signal", "zoo", "pracma")
+install.packages("signal", "zoo", "pracma")
+install.packages("signal", "pracma")
+install.packages("signal")
+install.packages("signal")
+install.packages("signal", "zoo", "pracma")
+install.packages("zoo", "pracma")
+install.packages("pracma")

code/new_realtime/.env

@@ -0,0 +1,3 @@
+LOG_DIR=/home/alex/projects/datascience-analysis/notebooks/logs
+LMDB_ROOT_DIR=/home/alex/projects/datascience-analysis/notebooks/lmdb_datasets
+RESULTS_ROOT_DIR=/home/alex/projects/datascience-analysis/notebooks/results

code/new_realtime/configs/cnn_run_config.py

@@ -0,0 +1,127 @@
+from functools import partial
+
+from configs.feature_config import feature_config
+from models.cnn import CNNTransformer
+
+from utils.data_utils import *
+from utils.training import *
+
+from models.lstm import *
+from models.utils import *
+from models.collation import *
+
+take_every_nth = int(288 / 12)
+shift_in_hours = 12
+input_window_length = (288 // take_every_nth) * 80
+# output_window_length = (288 // take_every_nth) * 1
+output_window_length = 1
+output_window_offset = input_window_length + (288 // take_every_nth) * 0
+window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
+min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
+
+max_lr = 1e-5
+batch_size = 256
+
+
+def get_cnn_run_config(
+        run_name: str,
+        run_description: str,
+        input_window_length: int,
+        output_window_length: int,
+        take_every_nth: int,
+        shift_in_hours: int,
+        output_window_offset: int,
+        batch_size: int,
+        model_parameters: dict,
+        max_lr: float,
+        num_epochs: int,
+        patience: int,
+        feature_config: dict):
+    """
+    Get the configuration for the CNN model
+    Returns:
+        run_configuration: configuration for the CNN model
+    """
+    base_model_config = {
+        "model_name": "cnn_regressor",
+        "version": "1.0.0",
+        "model_class": CNNTransformer,
+        "feature_config": feature_config | {"ignored_features":
+            [
+                "fertility_probability",
+                "ov_over_probability",
+                # "days_relative_to_ov",
+                # "is_biphasic",
+                "ov_day",
+            ]},
+        "preprocessing": {
+            "window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
+            "take_every_nth": take_every_nth,
+            "min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
+        },
+        "batch_fn": produce_window_batches,
+        "collate_fn": simple_x_y_collate,
+        "model_creation_fn": simple_model_creation,
+        "model_save_fn": simple_model_save,
+        "model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
+        "batch_loss_fn": get_model_loss,
+        "actual_fn": simple_get_y,
+        "predict_fn": simple_x_y_predict,
+        "model_parameters": {
+            **model_parameters | {"seq_len": input_window_length},
+        },
+        "input_window_length": input_window_length,
+        "output_window_length": output_window_length,
+        "output_window_offset": output_window_offset,
+    }
+
+    base_training_config = {
+        "batch_size": batch_size,
+        "model_class": CNNTransformer,
+        "learning_parameters": {
+            "learning_rate": max_lr * (batch_size / 4),
+            "epochs": num_epochs,
+            "patience": patience
+        },
+        "loss_functions": [
+            nn.MSELoss(),
+            nn.BCEWithLogitsLoss(),
+        ],
+        "max_grad_norm": 1.0,
+        "train_size": 0.7,
+        "val_size": 0.15,
+        "test_size": 0.15,
+    }
+    return {
+        "name": run_name,
+        "description": run_description,
+        "model_configuration": base_model_config,
+        "training_configuration": base_training_config,
+    }
+
+
+run_configuration = {
+    "runs": [
+        get_cnn_run_config(
+            run_name="cnn_ovulation_regression",
+            run_description="CNN model for ovulation regression",
+            input_window_length=input_window_length,
+            output_window_length=output_window_length,
+            take_every_nth=take_every_nth,
+            shift_in_hours=shift_in_hours,
+            output_window_offset=output_window_offset,
+            batch_size=batch_size,
+            model_parameters={
+                "cnn_channels": 64,
+                "kernel_size": 3,
+                "embed_dim": 64,
+                "num_enc_layers": 2,
+                "num_heads": 2,
+            },
+            max_lr=max_lr,
+            num_epochs=10,
+            patience=3,
+            feature_config=feature_config,
+        )
+    ]
+}

code/new_realtime/configs/combined.py

@@ -0,0 +1,132 @@
+from typing import Callable
+
+import configs
+from configs.feature_config import feature_config
+from configs.cnn_run_config import get_cnn_run_config
+from configs.transformer_run_config import get_transformer_run_config
+
+take_every_nth = int(288 / 12)
+shift_in_hours = 12
+input_window_length = (288 // take_every_nth) * 80
+# output_window_length = (288 // take_every_nth) * 1
+output_window_length = 1
+output_window_offset = input_window_length + (288 // take_every_nth) * 0
+window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
+min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
+
+
+def config_generator(config_gen_fn: Callable,
+                     fixed_params: dict,
+                     variable_param_configs: list,
+                     ) -> list:
+    configs = list()
+    for variable_config in variable_param_configs:
+        config = fixed_params.copy()
+        for param_name, param_value in variable_config.items():
+            # update the config with the variable parameter, values can be None, if default should be used
+            if param_value is not None:
+                config[param_name] = param_value
+        configs.append(config_gen_fn(**config))
+
+    return configs
+
+
+run_configuration = {
+    "name": "ovulation_regression",
+    "runs": config_generator(
+        get_cnn_run_config,
+        fixed_params={
+            "run_name": "cnn_ovulation_regression",
+            "run_description": "CNN model for ovulation regression",
+            "input_window_length": input_window_length,
+            "output_window_length": output_window_length,
+            "take_every_nth": take_every_nth,
+            "shift_in_hours": shift_in_hours,
+            "output_window_offset": output_window_offset,
+            "batch_size": 256,
+            "max_lr": 1e-5,
+            "num_epochs": 10,
+            "patience": 3,
+            "feature_config": feature_config,
+        },
+        variable_param_configs=[
+            {
+                "model_parameters": {
+                    "cnn_channels": 32,
+                    "kernel_size": 3,
+                    "embed_dim": 32,
+                    "num_enc_layers": 2,
+                    "num_heads": 2,
+                },
+            },
+            {
+                "model_parameters": {
+                    "cnn_channels": 64,
+                    "kernel_size": 3,
+                    "embed_dim": 64,
+                    "num_enc_layers": 2,
+                    "num_heads": 2,
+                },
+            },
+            {
+                "batch_size": 128,
+                "model_parameters": {
+                    "cnn_channels": 128,
+                    "kernel_size": 5,
+                    "embed_dim": 128,
+                    "num_enc_layers": 4,
+                    "num_heads": 4,
+                },
+            },
+        ]
+    ) + config_generator(
+        get_transformer_run_config,
+        fixed_params={
+            "run_name": "transformer_ovulation_regression",
+            "run_description": "Transformer model for ovulation regression",
+            "input_window_length": input_window_length,
+            "output_window_length": output_window_length,
+            "take_every_nth": take_every_nth,
+            "shift_in_hours": shift_in_hours,
+            "output_window_offset": output_window_offset,
+            "batch_size": 128,
+            "max_lr": 1e-5,
+            "num_epochs": 10,
+            "patience": 3,
+            "feature_config": feature_config,
+        },
+        variable_param_configs=[
+            {
+                "model_parameters": {
+                    "embed_dim": 64,
+                    "num_enc_layers": 2,
+                    "num_heads": 2,
+                },
+            },
+            {
+                "batch_size": 64,
+                "model_parameters": {
+                    "embed_dim": 128,
+                    "num_enc_layers": 4,
+                    "num_heads": 4,
+                },
+            },
+            {
+                "batch_size": 64,
+                "model_parameters": {
+                    "embed_dim": 256,
+                    "num_enc_layers": 4,
+                    "num_heads": 4,
+                },
+            },
+            {
+                "batch_size": 32,
+                "model_parameters": {
+                    "embed_dim": 512,
+                    "num_enc_layers": 4,
+                    "num_heads": 4,
+                },
+            },
+        ]
+    )
+}

code/new_realtime/configs/feature_config.py

@@ -0,0 +1,167 @@
+import os
+from functools import partial
+import math
+
+import torch
+from torch import nn
+import numpy as np
+
+from sklearn.preprocessing import StandardScaler, MinMaxScaler, RobustScaler
+from torch.optim.lr_scheduler import OneCycleLR
+from torch.optim import AdamW
+
+from data_analysis.models.ov_detection.config import model_config
+from utils.feature_functions import *
+from utils.loss_functions import weighted_bce_loss_fn
+from utils.training_utils import collate
+
+from vsm_datascience_common import constants
+
+feature_config = {
+    "feature_set_name": "full_feature_set",
+    "filter_criteria": {
+        "measurements.length": {
+            "$gt": constants.MEASUREMENTS_PER_DAY * 10,
+            "$lt": constants.MEASUREMENTS_PER_DAY * 150
+        },
+        "ends_at": {
+            "$exists": True,
+            "$ne": None
+        },
+        "$and": [
+            {"measurements.values": {"$not": {"$elemMatch": {"$lt": 35}}}},
+            {"measurements.values": {"$not": {"$elemMatch": {"$gt": 43}}}}
+        ]
+    },
+    "augmentation": {
+        "use_augmentation": True,
+        "max_lookback": 4,
+    },
+    "feature_sets": [
+        "static_categorical_features",
+        "static_continuous_features",
+        "known_categorical_features",
+        "known_continuous_features",
+        "observed_categorical_features",
+        "observed_continuous_features",
+        "target_features"
+    ],
+    "static_categorical_features": [
+    ],
+    "static_continuous_features": [
+        {
+            "name": "average_cycle_length",
+            "fn": get_cycle_length_stats,
+            "scaler": RobustScaler
+        },
+        {
+            "name": "average_ovulation_day",
+            "fn": get_average_ovulation_day,
+            "scaler": RobustScaler
+        },
+        {
+            "name": "ovulation_std",
+            "fn": get_ovulation_std,
+            "scaler": StandardScaler
+        },
+        {
+            "name": "biphasic_fraction",
+            "fn": get_biphasic_fraction,
+            "scaler": MinMaxScaler
+        },
+        {
+            "name": "num_cycles",
+            "fn": get_num_cycles,
+            "scaler": RobustScaler
+        },
+        {
+            "name": "temperature_averages",
+            "fn": get_average_temperatures,
+            "scaler": StandardScaler
+        }
+    ],
+    "known_categorical_features": [
+    ],
+    "known_continuous_features": [
+        {
+            "name": "hours_from_start",
+            "fn": partial(get_hours_from_start, shift=0),
+            "scaler": MinMaxScaler,
+            "accumulation_fn": np.max
+        },
+        {
+            "name": "hour_of_day",
+            "fn": partial(get_hour_of_day_encoded, shift=0),
+            "scaler": None,
+            "accumulation_fn": np.max
+        },
+        {
+            "name": "day_of_week",
+            "fn": partial(get_day_of_week_encoded, shift=0),
+            "scaler": None,
+            "accumulation_fn": np.max
+        },
+        {
+            "name": "month_of_year",
+            "fn": partial(get_month_of_year_encoded, shift=0),
+            "scaler": None,
+            "accumulation_fn": np.max
+        }
+    ],
+    "observed_categorical_features": [
+    ],
+    "observed_continuous_features": [
+        {
+            "name": "temperature",
+            "fn": partial(get_temperature, shift=0),
+            "scaler": StandardScaler
+        },
+        {
+            "name": "rolling_average_temperature",
+            "fn": partial(get_rolling_average_with_padding, shift=0),
+            "scaler": StandardScaler
+        },
+        {
+            "name": "rolling_window_temperature_min",
+            "fn": partial(get_window_fn, window_size=constants.MEASUREMENTS_PER_DAY, fn=partial(np.min, axis=1)),
+            "scaler": StandardScaler
+        },
+        {
+            "name": "rolling_window_temperature_max",
+            "fn": partial(get_window_fn, window_size=constants.MEASUREMENTS_PER_DAY, fn=partial(np.max, axis=1)),
+            "scaler": StandardScaler
+        },
+    ],
+    "target_features": [
+        # {
+        #     "name": "fertility_probability",
+        #     "fn": partial(get_fertility_probability, shift=0),
+        #     "scaler": MinMaxScaler,
+        #     "accumulation_fn": np.max
+        # },
+        # {
+        #     "name": "ov_over_probability",
+        #     "fn": partial(get_ov_over_probability, shift=0),
+        #     "scaler": MinMaxScaler,
+        #     "accumulation_fn": np.max
+        # },
+        {
+            "name": "days_relative_to_ov",
+            "fn": get_days_relative_to_ov,
+            "scaler": RobustScaler,
+            "accumulation_fn": np.max
+        },
+        {
+            "name": "ov_day",
+            "fn": get_ov_day,
+            "scaler": RobustScaler,
+            "accumulation_fn": np.max
+        },
+        {
+            "name": "is_biphasic",
+            "fn": get_is_biphasic,
+            "scaler": MinMaxScaler,
+            "accumulation_fn": np.max
+        }
+    ]
+}

code/new_realtime/configs/lstm_run_config.py

@@ -0,0 +1,127 @@
+from functools import partial
+
+from configs.feature_config import feature_config
+
+from utils.data_utils import *
+from utils.training import *
+
+from models.lstm import *
+from models.utils import *
+from models.collation import *
+
+take_every_nth = int(288 / 12)
+shift_in_hours = 12
+input_window_length = (288 // take_every_nth) * 80
+# output_window_length = (288 // take_every_nth) * 1
+output_window_length = 1
+output_window_offset = input_window_length + (288 // take_every_nth) * 0
+window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
+min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
+
+max_lr = 1e-5
+batch_size = 128
+
+
+def get_lstm_run_config(
+        run_name: str,
+        run_description: str,
+        input_window_length: int,
+        output_window_length: int,
+        take_every_nth: int,
+        shift_in_hours: int,
+        output_window_offset: int,
+        batch_size: int,
+        model_parameters: dict,
+        max_lr: float,
+        num_epochs: int,
+        patience: int,
+        feature_config: dict):
+    """
+    Get the configuration for the LSTM model
+    Returns:
+        run_configuration: configuration for the LSTM model
+    """
+    base_model_config = {
+        "model_name": "lstm_regressor",
+        "version": "1.0.0",
+        "model_class": LSTMModel,
+        "feature_config": feature_config | {"ignored_features":
+            [
+                "fertility_probability",
+                "ov_over_probability",
+                # "days_relative_to_ov",
+                # "is_biphasic",
+                "ov_day",
+            ]},
+        "preprocessing": {
+            "window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
+            "take_every_nth": take_every_nth,
+            "min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
+        },
+        "batch_fn": produce_window_batches,
+        "collate_fn": simple_x_y_collate,
+        "model_creation_fn": simple_model_creation,
+        "model_save_fn": simple_model_save,
+        "model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
+        "batch_loss_fn": get_model_loss,
+        "actual_fn": simple_get_y,
+        "predict_fn": simple_x_y_predict,
+        "model_parameters": {
+            **model_parameters,
+        },
+        "input_window_length": input_window_length,
+        "output_window_length": output_window_length,
+        "output_window_offset": output_window_offset,
+    }
+    base_training_config = {
+        "batch_size": batch_size,
+        "model_class": LSTMModel,
+        "learning_parameters": {
+            "learning_rate": max_lr * (batch_size / 4),
+            # "learning_rate": base_lr * (batch_size / 4),
+            "epochs": num_epochs,
+            "patience": patience,
+        },
+        "loss_functions": [
+            nn.MSELoss(),
+            nn.BCEWithLogitsLoss(),
+        ],
+        "max_grad_norm": 1.0,
+        "train_size": 0.7,
+        "val_size": 0.15,
+        "test_size": 0.15,
+    }
+    return {
+        "name": run_name,
+        "description": run_description,
+        "model_configuration": base_model_config,
+        "training_configuration": base_training_config,
+    }
+
+
+run_configuration = {
+    "item_limit": 100,
+    "runs": [
+        get_lstm_run_config(
+            run_name="lstm_regressor",
+            run_description="LSTM regressor for fertility prediction",
+            input_window_length=input_window_length,
+            output_window_length=output_window_length,
+            take_every_nth=take_every_nth,
+            shift_in_hours=shift_in_hours,
+            output_window_offset=output_window_offset,
+            batch_size=batch_size,
+            model_parameters={
+                "cnn_channels": 64,
+                "cnn_kernel_size": 3,
+                "embed_dim": 128,
+                "lstm_hidden_size": 128,
+                "num_layers": 4,
+            },
+            max_lr=max_lr,
+            num_epochs=1000,
+            patience=50,
+            feature_config=feature_config
+        )
+    ]
+}

code/new_realtime/configs/patch_run_config.py

@@ -0,0 +1,109 @@
+from functools import partial
+
+from configs.feature_config import feature_config
+
+from utils.data_utils import *
+from utils.training import *
+
+from models.third_party.patch_tst.models.PatchTST import Model as PatchTST
+from models.utils import *
+from models.collation import *
+
+take_every_nth = int(288 / 12)
+shift_in_hours = 12
+input_window_length = (288 // take_every_nth) * 80
+# output_window_length = (288 // take_every_nth) * 1
+output_window_length = 1
+output_window_offset = input_window_length + (288 // take_every_nth) * 0
+window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
+min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
+
+max_lr = 1e-5
+batch_size = 64
+
+run_configuration = {
+    "item_limit": 100,
+    "runs": [
+        {
+            "name": "lstm_ovulation_regression",
+            "description": "LSTM model for ovulation regression",
+            "model_configuration": {
+                "model_name": "patch_tst_regressor",
+                "version": "1.0.0",
+                "model_class": PatchTST,
+                "feature_config": feature_config | {"ignored_features":
+                    [
+                        "fertility_probability",
+                        "ov_over_probability",
+                        # "days_relative_to_ov",
+                        # "is_biphasic",
+                        "ov_day",
+                    ]},
+                "preprocessing": {
+                    "window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
+                    "take_every_nth": take_every_nth,
+                    "min_input_length_fraction_for_padding": ((288 // take_every_nth) * 4) / input_window_length,
+                },
+                "batch_fn": produce_window_batches,
+                "collate_fn": simple_x_y_collate,
+                "model_creation_fn": simple_model_creation,
+                "model_save_fn": simple_model_save,
+                "model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
+                "batch_loss_fn": get_model_loss,
+                "actual_fn": simple_get_y,
+                "predict_fn": simple_x_y_predict,
+                "model_parameters": {
+                    "configs": {
+                        # core
+                        "seq_len": input_window_length,
+                        "pred_len": output_window_length,
+                        "seq_pred": False,
+                        # model
+                        "e_layers": 4,
+                        "n_heads": 4,
+                        "d_model": 128,
+                        "d_ff": 128,
+                        "dropout": 0.2,
+                        "fc_dropout": 0.2,
+                        "head_dropout": 0.0,
+                        "individual": True,
+                        # patch
+                        # "patch_len": input_window_length,
+                        "patch_len": int(288 / take_every_nth),
+                        "stride": int(288 / take_every_nth / 2),
+                        "padding_patch": 0,
+                        # preprocessing
+                        "revin": False,
+                        "affine": False,
+                        "subtract_last": False,
+                        # decomp
+                        "decomposition": True,
+                        "kernel_size": 3,
+                    }
+                },
+                "input_window_length": input_window_length,
+                "output_window_length": output_window_length,
+                "output_window_offset": output_window_offset,
+            },
+            "training_configuration": {
+                "batch_size": batch_size,
+                "model_class": PatchTST,
+                "learning_parameters": {
+                    # "learning_rate": base_lr,
+                    "learning_rate": max_lr * (batch_size / 4),
+                    "epochs": 10,
+                    "patience": 3,
+                },
+                "loss_functions": [
+                    nn.MSELoss(),
+                    nn.BCEWithLogitsLoss(),
+                ],
+                "max_grad_norm": 1.0,
+                "train_size": 0.7,
+                "val_size": 0.15,
+                "test_size": 0.15,
+            }
+
+        }
+    ]
+}

code/new_realtime/configs/transformer_run_config.py

@@ -0,0 +1,123 @@
+from functools import partial
+
+from configs.feature_config import feature_config
+
+from utils.data_utils import *
+from utils.training import *
+
+from models.transformer import *
+from models.utils import *
+from models.collation import *
+
+take_every_nth = int(288 / 12)
+shift_in_hours = 12
+input_window_length = (288 // take_every_nth) * 80
+# output_window_length = (288 // take_every_nth) * 1
+output_window_length = 1
+output_window_offset = input_window_length + (288 // take_every_nth) * 0
+window_shift = int((288 // take_every_nth) / 24 * shift_in_hours)
+min_input_length_fraction_for_padding = ((288 // take_every_nth) * 4) / input_window_length
+
+transformer_batch_size = 128
+
+
+def get_transformer_run_config(
+        run_name: str,
+        run_description: str,
+        input_window_length: int,
+        output_window_length: int,
+        take_every_nth: int,
+        shift_in_hours: int,
+        output_window_offset: int,
+        batch_size: int,
+        model_parameters: dict,
+        max_lr: float,
+        num_epochs: int,
+        patience: int,
+        feature_config: dict):
+    """
+    Get the configuration for the Transformer model
+    Returns:
+        run_configuration: configuration for the Transformer model
+    """
+    base_model_config = {
+        "model_name": "transformer_regressor",
+        "version": "1.0.0",
+        "model_class": TransformerModel,
+        "feature_config": feature_config | {"ignored_features":
+            [
+                "fertility_probability",
+                "ov_over_probability",
+                # "days_relative_to_ov",
+                # "is_biphasic",
+                "ov_day",
+            ]},
+        "preprocessing": {
+            "window_shift": int((288 // take_every_nth) / 24 * shift_in_hours),
+            "take_every_nth": take_every_nth,
+            "min_input_length_fraction_for_padding": min_input_length_fraction_for_padding,
+        },
+        "batch_fn": produce_window_batches,
+        "model_creation_fn": simple_model_creation,
+        "collate_fn": simple_x_y_collate,
+        "model_save_fn": simple_model_save,
+        "model_load_fn": partial(simple_model_load, model_creation_fn=simple_model_creation),
+        "batch_loss_fn": get_model_loss,
+        "actual_fn": simple_get_y,
+        "predict_fn": simple_x_y_predict,
+        "model_parameters": {
+            **model_parameters | {"seq_len": input_window_length},
+        },
+        "input_window_length": input_window_length,
+        "output_window_length": output_window_length,
+        "output_window_offset": output_window_offset,
+    }
+    base_training_config = {
+        "batch_size": batch_size,
+        "model_class": TransformerModel,
+        "learning_parameters": {
+            "learning_rate": max_lr * (batch_size / 4),
+            # "learning_rate": base_lr * (batch_size / 4),
+            "epochs": num_epochs,
+            "patience": patience,
+        },
+        "loss_functions": [
+            nn.MSELoss(),
+            nn.BCEWithLogitsLoss(),
+        ],
+        "max_grad_norm": 1.0,
+        "train_size": 0.7,
+        "val_size": 0.15,
+        "test_size": 0.15,
+    }
+    return {
+        "name": run_name,
+        "description": run_description,
+        "model_configuration": base_model_config,
+        "training_configuration": base_training_config
+    }
+
+
+run_configuration = {
+    "runs": [
+        get_transformer_run_config(
+            run_name="transformer_ovulation_regression",
+            run_description="Transformer regressor for ovulation prediction",
+            input_window_length=input_window_length,
+            output_window_length=output_window_length,
+            take_every_nth=take_every_nth,
+            shift_in_hours=shift_in_hours,
+            output_window_offset=output_window_offset,
+            batch_size=transformer_batch_size,
+            model_parameters={
+                "embed_dim": 128,
+                "num_heads": 4,
+                "num_enc_layers": 4,
+            },
+            max_lr=1e-5,
+            num_epochs=20,
+            patience=3,
+            feature_config=feature_config,
+        )
+    ]
+}

code/new_realtime/dataset_wrapper.py

@@ -0,0 +1,156 @@
+import sys
+import argparse
+from concurrent.futures import ProcessPoolExecutor
+from functools import partial
+import logging
+
+import lmdb
+import dotenv
+from tqdm import tqdm
+
+dotenv.load_dotenv()
+
+from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id
+from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
+
+from utils.dataset_creation import get_features, train_scalers, save_scalers, scale_item
+from utils.lmdb_utils import save_to_lmdb, load_from_lmdb
+from utils.utils import get_variable_from_module
+
+MAX_LMDB_SIZE_IN_MB = 200_000
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+logger.addHandler(logging.StreamHandler(sys.stdout))
+
+
+def process_id_wrapper(cycle_id: str, feature_config: dict, env: lmdb.Environment):
+    try:
+        cycle = get_cycle_by_id(cycle_id)
+        features = get_features(cycle, feature_config)
+        save_to_lmdb(env, key=str(cycle_id), dataset=features)
+    except:
+        pass
+
+
+def scaling_wrapper(key_batch: str,
+                    scalers: dict,
+                    max_lmdb_size_in_mb: int,
+                    lmdb_env_dir: str):
+    env = lmdb.open(lmdb_env_dir, readonly=True, lock=False)
+    scaled_items = list()
+    for key in key_batch:
+        item = load_from_lmdb(env, key)
+        scaled_item = scale_item(item, scalers)
+        scaled_items.append(scaled_item)
+    env.close()
+    env = lmdb.open(lmdb_env_dir, map_size=max_lmdb_size_in_mb * 1024 * 1024)
+    for key, scaled_item in zip(key_batch, scaled_items):
+        save_to_lmdb(env, key=key, dataset=scaled_item)
+    env.close()
+
+
+def create_dataset(model_configuration: dict,
+                   lmdb_root_dir: str,
+                   max_lmdb_size_in_mb: int,
+                   max_workers: int) -> None:
+    feature_config = model_configuration["feature_config"]
+
+    logger.info(f"Creating dataset for feature set {feature_config['feature_set_name']}")
+
+    # fetch valid cycle ids from database
+    valid_cycle_ids = [x["_id"] for x in get_cycles_collection().aggregate(
+        feature_config["filter_criteria_pipeline"] + [
+            {
+                "$project": {
+                    "_id": 1
+                }
+            }
+        ]
+    )]
+
+    logger.info(f"Fetched {len(valid_cycle_ids)} valid cycle ids from database")
+
+    env = lmdb.open(f"{lmdb_root_dir}/{feature_config['feature_set_name']}", map_size=max_lmdb_size_in_mb * 1024 * 1024)
+
+    # create features for items
+    logger.info(f"Creating features for {len(valid_cycle_ids)} cycles")
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        list(tqdm(executor.map(partial(process_id_wrapper, feature_config=feature_config, env=env), valid_cycle_ids),
+                  total=len(valid_cycle_ids)))
+
+    # convert ids (here bson objectids) to keys for use in lmdb
+    keys = [str(cycle_id) for cycle_id in valid_cycle_ids]
+
+    # train scalers for featues
+    logger.info("Training scalers for features")
+    all_scalers = dict()
+    sample = load_from_lmdb(env, keys[0])
+    for feature_set in tqdm(feature_config["feature_sets"]):
+        for feature in feature_config[feature_set]:
+            print(f"Training scalers for feature {feature['name']} in feature set {feature_set}")
+            scaler = train_scalers(feature_set, feature["name"], feature["scaler"], sample, env)
+            if feature_set not in all_scalers:
+                all_scalers[feature_set] = dict()
+            all_scalers[feature_set] = all_scalers[feature_set] | scaler
+
+    # save scalers
+    logger.info("Saving scalers to disk")
+    scaler_dir = f"{lmdb_root_dir}/{feature_config['feature_set_name']}/scalers"
+    save_scalers(all_scalers, scaler_dir)
+
+    # creat scaling batches for less burden on lmdb
+    batch_size = 1_000
+    key_batches = [keys[i:i + batch_size] for i in range(0, len(keys), batch_size)]
+
+    # scale items
+    logger.info("Scaling items")
+    with ProcessPoolExecutor(max_workers=max_workers) as executor:
+        list(tqdm(executor.map(partial(scaling_wrapper, scalers=all_scalers,
+                                       max_lmdb_size_in_mb=max_lmdb_size_in_mb,
+                                       lmdb_env_dir=f"{lmdb_root_dir}/{feature_config['feature_set_name']}"),
+                               key_batches),
+                  total=len(key_batches)))
+
+    logger.info(f"Dataset creation finished. LMDB saved in {lmdb_root_dir}/{feature_config['feature_set_name']}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Dataset Creation Wrapper")
+    parser.add_argument("--model_config_module",
+                        type=str,
+                        required=True,
+                        help="Path to the model config module")
+    parser.add_argument("--model_config_variable",
+                        type=str,
+                        required=False,
+                        default="model_configuration",
+                        help="Name of the model config variable, default is 'model_configuration'")
+    parser.add_argument("--lmdb_dir",
+                        type=str,
+                        required=False,
+                        default="./lmdb_datasets",
+                        help="Path to the lmdb directory")
+    parser.add_argument("--lmdb_size",
+                        type=int,
+                        required=False,
+                        default=MAX_LMDB_SIZE_IN_MB,
+                        help="Size of the lmdb in MB, default is 200_000")
+    parser.add_argument("--max_workers",
+                        type=int,
+                        required=False,
+                        default=None,
+                        help="Number of workers for multiprocessing, default is None (use all available cores)")
+
+    args = parser.parse_args()
+
+    # import and load the model config
+    model_configuration = get_variable_from_module(
+        module_path=args.model_config_module,
+        variable_name=args.model_config_variable
+    )
+
+    create_dataset(model_configuration,
+                   lmdb_root_dir=args.lmdb_dir,
+                   max_lmdb_size_in_mb=args.lmdb_size,
+                   max_workers=args.max_workers)

code/new_realtime/experiment_setup.py

@@ -0,0 +1,48 @@
+from functools import partial
+
+import numpy as np
+import sklearn
+
+from utils.evaluation import *
+
+
+def get_eval_functions(model_configuration: dict):
+    """
+    Returns the evaluation functions for the model
+    :return: list of evaluation functions
+    """
+    eval_functions = [
+        {
+            "name": "mean_absolute_error_overall",
+            "input_index": 0,
+            "eval_fn": sklearn.metrics.mean_absolute_error,
+            "accumulation_fn": np.mean,
+        },
+        {
+            "name": "mean_absolute_error_pre_ov",
+            "input_index": 0,
+            "eval_fn": pre_ov_error,
+            "accumulation_fn": np.mean,
+        },
+        {
+            "name": "mean_absolute_error_after_ov",
+            "input_index": 0,
+            "eval_fn": after_ov_error,
+            "accumulation_fn": np.mean,
+        },
+        {
+            "name": "mean_absolute_error_ov_in_days",
+            "input_index": 0,
+            "eval_fn": partial(ov_error, model_configuration=model_configuration),
+            "accumulation_fn": np.mean,
+        },
+        {
+            "name": "mean_absolute_error_five_days_before_ov",
+            "input_index": 0,
+            "eval_fn": partial(day_relative_to_ov_error,
+                               day_relative_to_ov=-5,
+                               model_configuration=model_configuration),
+            "accumulation_fn": np.mean,
+        }
+    ]
+    return eval_functions

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]

code/new_realtime/training_wrapper.py

@@ -0,0 +1,321 @@
+import json
+import sys
+import os
+import argparse
+import logging
+from datetime import datetime
+
+import lmdb
+import dotenv
+import torch
+from torch.utils.data import DataLoader
+import torch.distributed as dist
+
+from experiment_setup import get_eval_functions
+from utils.evaluation import evaluate_model
+from utils.model_utils import get_model_config
+from utils.training_utils import get_training_config, get_data_ids
+from utils.data_utils import LMDBIterableDataset
+from utils.utils import get_variable_from_module, get_logger, convert_for_json
+from utils.training import train_model
+
+dotenv.load_dotenv()
+
+logger = None
+
+# set up distributed training
+dist.init_process_group(backend="nccl", init_method="env://")
+local_rank = torch.distributed.get_rank()
+torch.cuda.set_device(local_rank)
+
+
+def prepare_run(results_dir: str,
+                lmdb_root_dir: str,
+                base_model_configuration: dict,
+                base_training_configuration: dict):
+    # create model configuration from base
+    logger.info("Creating model configuration")
+    model_configuration = get_model_config(base_model_configuration, results_dir, lmdb_root_dir)
+    feature_config = model_configuration["feature_config"]
+    dataset_dir = f"{lmdb_root_dir}/{feature_config['feature_set_name']}"
+    logger.info(f"Model configuration name: {model_configuration['id']}")
+
+    # create training configuration
+    logger.info("Creating training configuration")
+    training_configuration = get_training_config(base_training_configuration, model_configuration)
+    logger.info(f"Training configuration: {training_configuration['id']}")
+
+    logger.info(f"Fetching data ids, limit: {item_limit}")
+    train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, item_limit)
+
+    logger.info(f"Train ids: {len(train_ids)}, Val ids: {len(val_ids)}, Test ids: {len(test_ids)}")
+
+    return model_configuration, training_configuration, train_ids, val_ids, test_ids, dataset_dir
+
+
+def train(
+        model_configuration: dict,
+        training_configuration: dict,
+        train_ids: list,
+        val_ids: list,
+        dataset_dir: str,
+        log_dir: str) -> None:
+    # create training and validation loaders
+    logger.info("Creating training loaders")
+    train_loader = LMDBIterableDataset(dataset_dir,
+                                       train_ids,
+                                       model_configuration=model_configuration,
+                                       batch_size=training_configuration["batch_size"])
+    logger.info("Creating validation loaders")
+    val_loader = LMDBIterableDataset(dataset_dir,
+                                     val_ids,
+                                     model_configuration=model_configuration,
+                                     batch_size=training_configuration["batch_size"])
+
+    # instantiate model
+    logger.info("Creating model")
+    lmdb_env = lmdb.open(dataset_dir, readonly=True)
+    model_creation_fn = model_configuration["model_creation_fn"]
+    model = model_creation_fn(model_configuration,
+                              train_ids[0],
+                              lmdb_env=lmdb_env)
+
+    logger.info(f"Started training for model {model_configuration['id']}")
+    train_model(
+        model=model,
+        model_configuration=model_configuration,
+        training_configuration=training_configuration,
+        train_dataset=train_loader,
+        val_dataset=val_loader,
+        log_dir=log_dir,
+        logger=logger
+    )
+
+
+def evaluate(model_configuration: dict,
+             training_configuration: dict,
+             test_ids: list,
+             device) -> dict:
+    logger.info(f"Evaluating model {model_configuration['id']}")
+    eval_functions = get_eval_functions(model_configuration)
+    results = evaluate_model(model_configuration,
+                             training_configuration,
+                             test_ids,
+                             eval_functions)
+
+    # save results to file
+    results_dir = training_configuration["training_dir"]
+    if not os.path.exists(results_dir):
+        os.makedirs(results_dir)
+    results_file = os.path.join(results_dir, "evaluation_results.json")
+    with open(results_file, "w") as f:
+        json.dump(convert_for_json(results), f, indent=4)
+    return results
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Training wrapper for model training")
+    parser.add_argument("run_configuration_module",
+                        type=str,
+                        help="Path to the run configuration module")
+    parser.add_argument("--run_configuration_variable",
+                        type=str,
+                        required=False,
+                        default="run_configuration",
+                        help="Name of the run configuration variable in the module")
+    parser.add_argument("--results_dir",
+                        type=str,
+                        required=False,
+                        default=None,
+                        help="Directory to save the results")
+    parser.add_argument("--lmdb_root_dir",
+                        type=str,
+                        required=False,
+                        default=None,
+                        help="Path to the lmdb directory")
+    parser.add_argument("--log_dir",
+                        type=str,
+                        required=False,
+                        default=None,
+                        help="Directory to save the logs")
+
+    parser.add_argument("--item_limit",
+                        type=int,
+                        required=False,
+                        default=None,
+                        help="Limit the number of items to process, default is None (no limit)")
+    parser.add_argument("--device",
+                        type=str,
+                        required=False,
+                        default="cuda",
+                        help="Device to use for training, default is cuda")
+
+    args = parser.parse_args()
+    # load run configuration from module
+    run_configuration = get_variable_from_module(
+        # make sure to replace / with . and remove .py to get proper module tree
+        module_path=args.run_configuration_module.replace(".py", "").replace("/", "."),
+        variable_name=args.run_configuration_variable)
+
+    # check for variables in run_configuration
+    if not args.results_dir:
+        if "base_results_dir" not in run_configuration:
+            results_dir = os.getenv("RESULTS_ROOT_DIR")
+        else:
+            results_dir = run_configuration["base_results_dir"]
+    else:
+        results_dir = args.results_dir
+
+    if results_dir is None:
+        raise ValueError(
+            "No results directory specified. Please set the RESULTS_ROOT_DIR environment variable or provide a results_dir argument.")
+
+    if not os.path.exists(results_dir):
+        os.makedirs(results_dir)
+
+    if not args.lmdb_root_dir:
+        if "base_lmdb_root_dir" not in run_configuration:
+            lmdb_root_dir = os.getenv("LMDB_ROOT_DIR")
+        else:
+            lmdb_root_dir = run_configuration["base_lmdb_root_dir"]
+    else:
+        lmdb_root_dir = args.lmdb_root_dir
+
+    if lmdb_root_dir is None:
+        raise ValueError(
+            "No LMDB root directory specified. Please set the LMDB_ROOT_DIR environment variable or provide a lmdb_root_dir argument.")
+
+    if not os.path.exists(lmdb_root_dir):
+        os.makedirs(lmdb_root_dir)
+
+    if not args.log_dir:
+        if "base_log_dir" not in run_configuration:
+            log_dir = os.getenv("LOG_DIR")
+        else:
+            log_dir = run_configuration["log_dir"]
+    else:
+        log_dir = args.log_dir
+    if log_dir is None:
+        raise ValueError(
+            "No log directory specified. Please set the LOG_DIR environment variable or provide a log_dir argument.")
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+
+    item_limit = args.item_limit if args.item_limit else run_configuration["item_limit"]
+    if item_limit == -1:
+        item_limit = None
+
+    run_name = run_configuration["name"]
+    runs = run_configuration["runs"]
+
+    # set up run_id on rank 0
+    if local_rank == 0:
+        run_id = f"{run_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    else:
+        run_id = None
+
+    # broadcast run_id to all processes
+    if dist.is_initialized():
+        run_id_list = [run_id]
+        torch.distributed.broadcast_object_list(run_id_list, src=0)
+        run_id = run_id_list[0]
+
+    # make sure run_id is a string
+    run_id = str(run_id)
+
+
+    # append run_id to results_dir and log_dir
+    results_dir = os.path.join(results_dir, run_id)
+    log_dir = os.path.join(log_dir, run_id)
+
+    # create directories if they do not exist, only on the main process
+    if torch.distributed.get_rank() == 0 or not dist.is_initialized():
+        if not os.path.exists(results_dir):
+            os.makedirs(results_dir)
+        if not os.path.exists(log_dir):
+            os.makedirs(log_dir)
+
+    # set up logger
+    logger = get_logger(module_name=run_name, filename=os.path.join(log_dir, "main.log"))
+
+    logger.info(f"Rank {local_rank}: Starting run {run_name}")
+    logger.info(f"Rank {local_rank}: Run ID: {run_id}")
+
+    # print parameter values
+    logger.info(f"Rank {local_rank}: Results directory: {results_dir}")
+    logger.info(f"Rank {local_rank}: LMDB root directory: {lmdb_root_dir}")
+    logger.info(f"Rank {local_rank}: Log directory: {log_dir}")
+    logger.info(f"Rank {local_rank}: Item limit: {item_limit}")
+
+    if not runs or len(runs) == 0:
+        raise ValueError("No runs specified in the run configuration. Please provide a list of runs to train.")
+
+    try:
+        for run in runs:
+            run_step_name = run["name"]
+            run_description = run["description"]
+            run_model_configuration = run["model_configuration"]
+            run_training_configuration = run["training_configuration"]
+
+            logger.info(f"Rank {local_rank}: Running: {run_step_name}")
+            logger.info(f"Rank {local_rank}: Description: {run_description}")
+
+            # prepare run, make sure rank 0 is the first to avoid race conditions
+            if local_rank == 0:
+                logger.info(f"Rank {local_rank}: Preparing run {run_step_name}")
+                run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run(
+                    results_dir=results_dir,
+                    lmdb_root_dir=lmdb_root_dir,
+                    base_model_configuration=run_model_configuration,
+                    base_training_configuration=run_training_configuration
+                )
+                logger.info(f"Rank {local_rank}: Finished preparing run {run_step_name}")
+                # sync after preparing run
+                if dist.is_initialized():
+                    dist.barrier()
+            else:
+                # wait for rank 0 to finish preparing run
+                if dist.is_initialized():
+                    dist.barrier()
+                    logger.info(f"Rank {local_rank}: Waiting for rank 0 to finish preparing run {run_step_name}")
+                run_model_configuration, run_training_configuration, train_ids, val_ids, test_ids, dataset_dir = prepare_run(
+                    results_dir=results_dir,
+                    lmdb_root_dir=lmdb_root_dir,
+                    base_model_configuration=run_model_configuration,
+                    base_training_configuration=run_training_configuration
+                )
+
+            train(
+                model_configuration=run_model_configuration,
+                training_configuration=run_training_configuration,
+                train_ids=train_ids,
+                val_ids=val_ids,
+                dataset_dir=dataset_dir,
+                log_dir=log_dir,
+            )
+            # sync after training
+            if dist.is_initialized():
+                dist.barrier()
+
+            logger.info(f"Rank {torch.distributed.get_rank()} finished training {run_step_name}")
+
+            # run evaluation, only on rank 0
+            local_rank = torch.distributed.get_rank()
+            if local_rank == 0:
+                logger.info(f"Rank {local_rank} starting evaluation for {run_step_name}")
+                results = evaluate(
+                    model_configuration=run_model_configuration,
+                    training_configuration=run_training_configuration,
+                    test_ids=test_ids,
+                )
+                logger.info(
+                    f"Rank {local_rank} finished evaluation for {run_step_name} with model {run_model_configuration['id']}")
+                logger.info(f"Results: {results}")
+
+            # sync after evaluation
+            if dist.is_initialized():
+                dist.barrier()
+    finally:
+        # clean up
+        if dist.is_initialized():
+            dist.destroy_process_group()

code/new_realtime/utils/data_utils.py

@@ -0,0 +1,576 @@
+import copy
+import logging
+import os
+import random
+
+import lmdb
+import numpy as np
+from bson import ObjectId
+from torch.utils.data import IterableDataset
+from tqdm import tqdm
+
+from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features
+from utils.dataset_utils import load_key_stats
+from utils.lmdb_utils import load_from_lmdb
+from vsm_datascience_common.cycle_database_connection.cycle_data import get_cycle_by_id
+from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
+
+logger = logging.getLogger(__name__)
+
+
+def get_prepared_sequence_length(input_sequence: np.ndarray, model_configuration: dict) -> int:
+    """
+    Returns the length of the input sequence after padding and resampling
+    """
+
+    initial_length = len(input_sequence)
+    take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
+    padding_length = get_padding_length(model_configuration, resampled=False)
+    raw_length = initial_length + padding_length
+    adjusted = int(raw_length // take_every_nth)
+    return adjusted
+
+
+def get_padding_length(model_configuration: dict, resampled: bool = True) -> int:
+    """
+    Returns the length of the input sequence after padding and resampling
+    """
+    take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
+    padding_length = int(model_configuration["input_window_length"] * (
+            1 - model_configuration["preprocessing"]["min_input_length_fraction_for_padding"])) * (
+                         take_every_nth if not resampled else 1)
+    return padding_length
+
+
+def get_number_of_windows(base_length: int, model_configuration: dict) -> int:
+    """
+    Returns the number of windows for a given sequence length
+    """
+    window_shift = model_configuration["preprocessing"]["window_shift"]
+    input_window_length = model_configuration["input_window_length"]
+    output_window_length = model_configuration["output_window_length"]
+    output_window_offset = model_configuration["output_window_offset"]
+
+    if input_window_length > output_window_length + output_window_offset:
+        return (base_length - input_window_length) // window_shift + 1
+    else:
+        return (base_length - output_window_length - output_window_offset) // window_shift + 1
+
+
+def get_collated_batch_for_key(sample_key: ObjectId | str,
+                               model_configuration: dict,
+                               start_cutoff: int = None,
+                               end_cutoff: int = None,
+                               lmdb_env=None) -> dict:
+    sample_batch_for_key = get_batch_for_key(sample_key, model_configuration,
+                                             start_cutoff=start_cutoff,
+                                             end_cutoff=end_cutoff,
+                                             lmdb_env=lmdb_env)
+    collated = model_configuration["collate_fn"](sample_batch_for_key)
+    return collated
+
+
+def get_batch_for_key(key,
+                      model_configuration: dict,
+                      start_cutoff: int = None,
+                      end_cutoff: int = None,
+                      lmdb_env=None) -> np.ndarray:
+    """
+    Get the batch for a given key from the lmdb database or compute it directly.
+
+    The return has the same format as the batch produced by the batch_fn in the model configuration.
+    Args:
+        key: key to get the batch for, can also be object id from database
+        model_configuration: model configuration to use
+        start_cutoff: start cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points
+        end_cutoff: end cutoff for the batch, if None, the whole batch is used, CAUTION: cutoff should not be normalized -> in raw data points
+        lmdb_env: the lmdb environment to use, if None, the features are computed directly from the database
+
+    Returns:
+        batch: batch for the given key, as returned by the batch_fn in the model configuration
+
+    """
+    if lmdb_env is None:
+        # compute and scale features, ignored features are handled internally
+        features = get_scaled_feature_for_key(key, model_configuration)
+    else:
+        # load features from lmdb
+        features = load_from_lmdb(lmdb_env, str(key))
+        # ignored features are handled internally, thus we need to remove them here
+        ignored_features = model_configuration["feature_config"]["ignored_features"] if "ignored_features" in \
+                                                                                        model_configuration[
+                                                                                            "feature_config"] else []
+        # add scaled versions of features
+        ignored_features = ignored_features + [feat + "_scaled" for feat in ignored_features]
+        for feature_set in features:
+            for feature_name in list(features[feature_set].keys()):
+                if feature_name in ignored_features:
+                    del features[feature_set][feature_name]
+
+    processed_chunk = process_chunk([key], [features], model_configuration)
+
+    # use cutoff, if provided
+    if start_cutoff is not None or end_cutoff is not None:
+
+        item_length = len(processed_chunk[0]["target_features"])
+
+        if start_cutoff is None:
+            start_cutoff = 0
+        if end_cutoff is None:
+            end_cutoff = item_length
+
+        # "normalize" cutoff to adjust for added padding and resampling
+        padding_length = get_padding_length(model_configuration, resampled=False)
+        take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
+        # make sure that there is no padding added to the start cutoff, so that it includes the padding added to the sequence
+        start_cutoff_normalized = max(int(start_cutoff // take_every_nth), 0)
+        # end cutoff must be adjusted to include the padding added to the sequence
+        end_cutoff_normalized = min(int((end_cutoff + padding_length) // take_every_nth), item_length)
+        for feature_set in processed_chunk[0]:
+            if feature_set in model_configuration["feature_config"]["feature_sets"]:
+                processed_chunk[0][feature_set] = processed_chunk[0][feature_set][
+                                                  start_cutoff_normalized:end_cutoff_normalized + 1]
+
+    if "batch_fn" in model_configuration and model_configuration["batch_fn"] is not None:
+        batch = model_configuration["batch_fn"](processed_chunk, model_configuration)
+    else:
+        batch = processed_chunk
+    if batch is None or len(batch) == 0:
+        raise ValueError(f"Batch is empty for key {key}")
+    return batch
+
+
+def get_scaled_feature_for_key(key: str,
+                               model_configuration: dict) -> tuple:
+    feature_config = model_configuration["feature_config"]
+
+    user_cycles = list(
+        get_cycles_collection().find({"user_id": ObjectId(key)} | feature_config["filter_criteria"]).sort("starts_at",
+                                                                                                          1))
+    cycle_features = list()
+    for cycle in user_cycles:
+        features = get_features(cycle, feature_config)
+        cycle_features.append(features)
+
+    # combine all cycles for the user
+    if len(cycle_features) == 0:
+        raise ValueError(f"No features found for key {key}")
+    features = combine_features(cycle_features, feature_config)
+    scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
+    scalers = load_scalers(scaler_dir)
+    # scale features
+    scaled_features = scale_item(features, scalers)
+    return scaled_features
+
+
+def augment_items(identifiers: list,
+                  items: list[dict],
+                  feature_config: dict,
+                  lmdb_env: lmdb.Environment = None) -> list:
+    """
+    Augments the given sequences by attaching previous sequences
+    Args:
+        identifiers: identifiers of items for finding previous sequences
+        items: actual items
+        feature_config: feature config of the dataset
+        lmdb_env: lmdb environment for loading previous sequences, can be None, in this case the items are computed directly
+
+    Returns:
+        list of augmented items
+    """
+
+    if "augmentation" not in feature_config or \
+            feature_config["augmentation"]["use_augmentation"] is False:
+        return items
+
+    augmented_items = list()
+    for i, item in enumerate(items):
+        item_id = identifiers[i]
+        try:
+            item_data = get_cycles_collection().find_one({"_id": ObjectId(item_id)}, {"user_id": 1, "starts_at": 1})
+            item_user_id = item_data["user_id"]
+            item_starts_at = item_data["starts_at"]
+        except KeyError:
+            print(f"User ID not found for item {item_id}")
+            continue
+        previous_item_identifiers = list(get_cycles_collection().aggregate(
+            # make sure to filter by user first to significantly reduce the number of items
+            [
+                {
+                    "$match": {
+                        "user_id": item_user_id,
+                        "starts_at": {"$lt": item_starts_at},
+                    }
+                }
+            ] + feature_config["filter_criteria_pipeline"] + [
+                {
+                    "$sort": {
+                        "starts_at": -1
+                    }
+                },
+                {
+                    "$project": {
+                        "_id": 1,
+                        "starts_at": 1
+                    }
+                }
+            ]
+        ))
+        max_lookback = feature_config["augmentation"]["max_lookback"]
+        previous_items = list()
+        for previous_item in previous_item_identifiers:
+            if len(previous_items) >= max_lookback:
+                break
+            previous_item_id = previous_item["_id"]
+            if lmdb_env is None:
+                # compute item
+                cycle = get_cycle_by_id(previous_item_id)
+                features = get_features(cycle, feature_config)
+                scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
+                scalers = load_scalers(scaler_dir)
+                previous_item = scale_item(features, scalers)
+                previous_items.append(previous_item)
+            else:
+                # take item from lmdb
+                try:
+                    # make sure to parse object id
+                    previous_item = load_from_lmdb(lmdb_env, str(previous_item_id))
+                except KeyError:
+                    continue
+                previous_items.append(previous_item)
+
+        # merge items into one
+        feature_sets = feature_config["feature_sets"]
+        augmented_item = copy.deepcopy(item)
+        for feature_set in feature_sets:
+            if feature_set not in augmented_item:
+                continue
+            if "static" in feature_set:
+                continue
+            for feature_name in augmented_item[feature_set]:
+                if feature_name not in augmented_item[feature_set]:
+                    continue
+                augmented_item[feature_set][feature_name] = np.concatenate(
+                    [previous_item[feature_set][feature_name] for previous_item in previous_items]
+                    + [augmented_item[feature_set][feature_name]])
+
+        augmented_items.append(augmented_item)
+
+    return augmented_items
+
+
+def process_chunk(ids: list,
+                  chunk: list,
+                  model_configuration: dict,
+                  ignored_features: list = None,
+                  pad_sequences: bool = True,
+                  statics_as_list: bool = True):
+    """
+    Creates a chunk of data as dataframe for training and inference of a tft model.
+    :param ids: list of ids of the data points in the chunk
+    :param chunk: list of data points
+
+    :return: dataframe with the data points in the chunk
+    """
+
+    if ignored_features is None:
+        ignored_features = list()
+
+    input_features = list(chunk[0]["target_features"].keys())
+    feature_config = model_configuration["feature_config"]
+
+    # variables for padding
+    padding_value = 0
+
+    # create dataframe from items in chunk
+    data = []
+    for i, item in enumerate(chunk):
+        item_length = item["target_features"][input_features[0]].shape[0]
+        take_every_nth = model_configuration["preprocessing"]["take_every_nth"]
+
+        data_item = dict()
+        for feature_set in feature_config["feature_sets"]:
+            if feature_set in item and len(item[feature_set]) > 0:
+                # fill in the data
+                source = item[feature_set]
+                features_to_use = [feature for feature in source.keys() if "_scaled" in feature]
+                feature_set_data = [source[feature] for feature in features_to_use]
+
+                padding_length = get_padding_length(model_configuration, resampled=False)
+
+                if len(feature_set_data) == 0:
+                    raise ValueError(f"No features found for key {ids[i]} in feature set {feature_set}")
+
+                if len(feature_set_data[0]) == 1:
+                    # if the feature is constant, we need to repeat it for all time points
+                    if statics_as_list:
+                        if pad_sequences:
+                            desired_length = item_length + padding_length
+                            feature_set_data = np.array(
+                                [np.full((desired_length,), val) for val in feature_set_data]).T[::take_every_nth]
+                        else:
+                            feature_set_data = np.array(
+                                [np.full((item_length,), val) for val in feature_set_data]).T[::take_every_nth]
+                    else:
+                        feature_set_data = np.array(feature_set_data).flatten()
+                else:
+                    if pad_sequences:
+                        padding_values = np.full((padding_length, len(feature_set_data)),
+                                                 [padding_value for vals in feature_set_data]).T
+                        feature_set_data = np.concatenate([padding_values, np.array(feature_set_data)], axis=1)
+                        # create bins of n values and apply mean
+
+                    # check, if special accumulation function has been specified
+                    individual_feature_data = list()
+                    for i, feature_to_use in enumerate(features_to_use):
+                        try:
+                            # get config for feature, make sure to replace scaled to get actual config
+                            current_feature_config = \
+                                [x for x in feature_config[feature_set] if
+                                 x["name"] == feature_to_use.replace("_scaled", "")][0]
+                            current_accumulation_fn = current_feature_config["accumulation_fn"]
+                        except (KeyError, IndexError):
+                            # default to mean
+                            current_accumulation_fn = np.mean
+
+                        current_feature_data = apply_fn_to_bins(feature_set_data[i], take_every_nth,
+                                                                current_accumulation_fn)
+                        individual_feature_data.append(current_feature_data)
+
+                    # stack individual features
+                    feature_set_data = np.stack(individual_feature_data).T
+
+                if feature_set not in data_item:
+                    data_item[feature_set] = feature_set_data
+                else:
+                    data_item[feature_set] = np.concatenate([data_item[feature_set], feature_set_data], axis=1)
+
+        data.append(data_item)
+
+    return data
+
+
+def apply_fn_to_bins(input_sequence: np.ndarray,
+                     bin_size: int,
+                     fn: callable) -> np.ndarray:
+    """
+    Applies a function to the bins of the input sequence along the last dimension and returns the results.
+    Args:
+        input_sequence: sequence to apply the function to (1D or 2D)
+        bin_size: size of the bins
+        fn: function to apply to the bins
+
+    Returns:
+        np.ndarray: result of the function applied to the bins
+
+    """
+    if input_sequence.ndim == 1:
+        input_sequence = input_sequence.reshape(1, -1)
+        dims = 1
+    else:
+        dims = input_sequence.shape[-1]
+
+    result = []
+    for row in input_sequence:
+        row_result = []
+        for start in range(0, len(row), bin_size):
+            end = min(start + bin_size, len(row))
+            bin_slice = row[start:end]
+            row_result.append(fn(bin_slice))
+        result.append(row_result)
+
+    if dims > 1:
+        return np.array(result)
+    else:
+        return np.array(result[0])
+
+
+def produce_window_batches(data_chunk: list,
+                           model_configuration: dict,
+                           offsets: list | np.ndarray = None) -> list:
+    """
+    Produces batches of windows from the data chunk.
+    Args:
+        data_chunk: data chunk
+        model_configuration: configuration to use
+        offsets: offsets used to determine the length of data to use, uses item[-offset:] of data
+
+    Returns:
+        data: list of windows
+    """
+    input_window_length = model_configuration["input_window_length"]
+    output_window_length = model_configuration["output_window_length"]
+    output_window_offset = model_configuration["output_window_offset"]
+    window_shift = model_configuration["preprocessing"]["window_shift"]
+
+    if offsets is None:
+        offsets = np.zeros(len(data_chunk), dtype=int)
+
+    data = list()
+    for i, item in enumerate(data_chunk):
+        item_length = item["target_features"][-offsets[i]:].shape[0]
+        total_window_length = input_window_length if input_window_length > output_window_length + output_window_offset \
+            else output_window_length + output_window_offset
+
+        if item_length < total_window_length:
+            continue
+
+        # number of windows is defined by the input window length and output window length with offset
+        num_windows = get_number_of_windows(item_length, model_configuration)
+        for j in range(num_windows):
+            window = dict()
+            for key, value in item.items():
+                if value is None:
+                    window[key] = None
+                else:
+                    if key == "target_features":
+                        window[key] = value[-offsets[i]:][j * window_shift + output_window_offset:
+                                                          j * window_shift + output_window_length + output_window_offset]
+                    else:
+                        window[key] = value[-offsets[i]:][j * window_shift:j * window_shift + input_window_length]
+
+            data.append(window)
+
+    return data
+
+
+def produce_simple_batches(data_chunk: list,
+                           model_configuration: dict,
+                           offsets: list | np.ndarray = None) -> list:
+    """
+    Produces batches without windowing.
+    Args:
+         data_chunk: data chunk
+         model_configuration: configuration to use
+         offsets: offsets used to determine the length of data to use, uses item[-offset:] of data
+
+    Returns:
+           data: list of items
+    """
+    pass
+
+
+class LMDBIterableDataset(IterableDataset):
+    def __init__(self,
+                 lmdb_env_path: str,
+                 lmdb_keys: list[str],
+                 model_configuration: dict,
+                 batch_size: int = 32):
+        self.lmdb_path = lmdb_env_path
+        self.lmdb_env = None
+        self.lmdb_keys = lmdb_keys
+        self.key_subset = None
+        self.model_configuration = model_configuration
+        self.batch_size = batch_size
+        self.keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"])
+        self.len = None
+
+        try:
+            self.model_configuration["batch_fn"]
+        except KeyError:
+            raise KeyError("Batch function not found in model configuration")
+
+    def set_key_subset(self, key_subset: list[str]):
+        """
+        Set the key subset to use for the dataset.
+        Args:
+            key_subset: list of keys to use
+        """
+        self.key_subset = key_subset
+        # reset length
+        self.len = None
+
+    def get_length_of_data_subset(self, key_set: list[str]):
+        """
+        Get the length of the data subset.
+        Args:
+            key_set: list of keys to use
+        """
+        # run simplified version of __iter__ to get the length
+
+        num_steps = 0
+        i = 0
+        num_batch_items = 0
+        while True:
+            if num_batch_items >= self.batch_size:
+                num_steps += 1
+                num_batch_items -= self.batch_size
+            else:
+                if i >= len(key_set):
+                    if num_batch_items > 0:
+                        num_steps += 1
+                    break
+                key = key_set[i]
+                i += 1
+
+                # try to fetch stats from key_stats
+                try:
+                    current_key_stats = self.keys_stats["by_key"][key]
+                    base_length = current_key_stats["item_length"]
+                    take_every_nth = self.model_configuration["preprocessing"]["take_every_nth"]
+                    padding_length = get_padding_length(self.model_configuration, resampled=False)
+                    raw_length = base_length + padding_length
+                    item_length = int(raw_length // take_every_nth)
+                except:
+                    item = load_from_lmdb(self.lmdb_env, key)
+                    item_length = get_prepared_sequence_length(
+                        item["target_features"][list(item["target_features"])[0]],
+                        self.model_configuration)
+                num_windows = get_number_of_windows(item_length, self.model_configuration)
+                num_batch_items += num_windows
+
+        return num_steps
+
+    def __iter__(self):
+        self.init_lmdb_env()
+
+        # if no key subset is set, use all keys
+        if self.key_subset is None:
+            self.key_subset = self.lmdb_keys
+
+        random.shuffle(self.key_subset)
+
+        batch = list()
+        counter = 0
+        collate_fn = self.model_configuration["collate_fn"]
+        while True:
+            if len(batch) >= self.batch_size:
+                if collate_fn is None:
+                    yield batch[:self.batch_size]
+                else:
+                    yield collate_fn(batch[:self.batch_size])
+                batch = batch[self.batch_size:]
+            else:
+                if counter >= len(self.key_subset):
+                    if len(batch) > 0:
+                        yield collate_fn(batch)
+                    break
+                key = self.lmdb_keys[counter]
+                counter += 1
+                try:
+                    current_batch = get_batch_for_key(key, self.model_configuration, lmdb_env=self.lmdb_env)
+                except ValueError:
+                    # if betch is empty, try next
+                    continue
+                batch.extend(current_batch)
+
+    def init_lmdb_env(self):
+        if self.lmdb_env is None:
+            self.lmdb_env = lmdb.open(self.lmdb_path,
+                                      readonly=True,
+                                      lock=False,
+                                      readahead=False,
+                                      meminit=False)
+
+    def __len__(self):
+        self.init_lmdb_env()
+
+        # if no key subset is set, use all keys
+        if self.key_subset is None:
+            self.key_subset = self.lmdb_keys
+
+        if self.len is None:
+            num_steps = self.get_length_of_data_subset(self.key_subset)
+            self.len = num_steps
+
+        return self.len

code/new_realtime/utils/dataset_creation.py

@@ -0,0 +1,235 @@
+import os
+
+import lmdb
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.preprocessing import StandardScaler, MinMaxScaler
+import joblib
+
+from utils.lmdb_utils import load_from_lmdb, get_lmdb_keys
+
+
+def get_features(cycle: dict,
+                 feature_config: dict) -> dict:
+    """
+    Computes the feature for a given cycle defined by the feature_config
+    :param cycle: cycle data as dictionary
+    :param feature_config: feature config as dictionary
+    :return: dict with feature according feature config
+    """
+    features = {}
+    feature_sets = feature_config["feature_sets"]
+    for feature_set in feature_sets:
+        for feature_def in feature_config[feature_set]:
+            if feature_set not in features:
+                features[feature_set] = {}
+
+            # skip features that are marked as ignored
+            if "ignored_features" in feature_config and feature_def["name"] in feature_config["ignored_features"]:
+                continue
+
+            feature_return = feature_def["fn"](cycle=cycle)
+            if isinstance(feature_return, dict):
+                if len(feature_return) > 1:
+                    for feature_name, feature in feature_return.items():
+                        features[feature_set][f"{feature_def['name']}_{feature_name}"] = feature
+                else:
+                    features[feature_set][feature_def["name"]] = list(feature_return.values())[0]
+            else:
+                features[feature_set][feature_def["name"]] = feature_return
+
+        # check, if all features have same length
+        if feature_set in features and len(features[feature_set]) > 0 and isinstance(
+                list(features[feature_set].values())[0], np.ndarray):
+            feature_lengths = [len(x) for x in features[feature_set].values()]
+            if len(set(feature_lengths)) > 1:
+                raise ValueError(f"Feature set {feature_set} has features of different lengths: {feature_lengths}")
+
+    return features
+
+
+def save_scalers(scalers: dict,
+                 scaler_dir: str):
+    if not os.path.exists(scaler_dir):
+        os.makedirs(scaler_dir)
+    for feature_type in scalers:
+        for feature_name in scalers[feature_type]:
+            joblib.dump(scalers[feature_type][feature_name],
+                        f"{scaler_dir}/{feature_name}.pkl")
+
+
+def load_scalers(scaler_dir: str) -> dict:
+    scalers = dict()
+    for scaler_file in os.listdir(scaler_dir):
+        if scaler_file.endswith(".pkl"):
+            feature_name = scaler_file.replace(".pkl", "")
+            scalers[feature_name] = joblib.load(f"{scaler_dir}/{scaler_file}")
+    return scalers
+
+
+def get_scalers_for_model(model_configuration: dict):
+    """
+    Get the scalers for the features of a model configuration
+    Args:
+        model_configuration: configuration of the model
+
+    Returns:
+        scalers: scalers for the model
+
+    """
+
+    scaler_dir = os.path.join(model_configuration["feature_config"]["dataset_dir"], "scalers")
+    scalers = load_scalers(scaler_dir)
+    return scalers
+
+
+def get_feature_values(feature_type: str | None,
+                       feature_name: str,
+                       env: lmdb.Environment,
+                       keys: list = None) -> list:
+    """
+    Get the values of a feature from the lmdb dataset for all keys
+    :param feature_type: type of feature, can be None, then the first feature with the given name will be used
+    :param feature_name: feature name to extract
+    :param env: lmdb environment
+    :param keys: keys to extract the feature from, if None, all keys will be used / fetched from the database
+    :return: list with feature values
+    """
+    if keys is None:
+        keys = get_lmdb_keys(env)
+    feature_values = []
+    for key in keys:
+        data = load_from_lmdb(env, key)
+        if feature_type is None:
+            for feature_type, feature_data in data.items():
+                if feature_name in feature_data:
+                    feature_values.append(feature_data[feature_name])
+                    break
+        else:
+            for current_feature_name, feature_data in data[feature_type].items():
+                # catch sub features that have been prefixed with the feature name
+                if current_feature_name.startswith(feature_name):
+                    feature_values.append(feature_data)
+                    break
+    return feature_values
+
+
+def train_scalers(feature_type: str,
+                  feature_name: str,
+                  scaler_type,
+                  sample,
+                  env) -> dict:
+    scalers = dict()
+    all_features = sample[feature_type]
+    individual_feature_names = list()
+    for individual_feature_name in all_features:
+        if individual_feature_name.startswith(feature_name) and not individual_feature_name.endswith("_scaled"):
+            individual_feature_names.append(individual_feature_name)
+
+    for individual_feature_name in individual_feature_names:
+        feature_values = get_feature_values(feature_type, individual_feature_name, env)
+        if len(feature_values) == 0:
+            raise ValueError(f"No feature values found for feature {individual_feature_name}")
+        if isinstance(feature_values[0], list) or isinstance(feature_values[0], np.ndarray):
+            features_reshaped = np.concatenate(feature_values).reshape(-1, 1)
+        else:
+            features_reshaped = np.array(feature_values).reshape(-1, 1)
+        del feature_values
+
+        if scaler_type is not None:
+            scaler = scaler_type()
+            scaler.fit(features_reshaped)
+        else:
+            scaler = None
+
+        scalers[individual_feature_name] = scaler
+    return scalers
+
+
+def scale_item(data: dict,
+               scalers: dict):
+    data_format = dict()
+    for feature_set in data:
+        if feature_set not in data_format:
+            data_format[feature_set] = dict()
+        for feature_name in data[feature_set]:
+            if feature_name.endswith("_scaled"):
+                continue
+            data_format[feature_set][feature_name] = data[feature_set][feature_name]
+
+    for feature_set in data_format:
+        for feature_name in data_format[feature_set]:
+            feature_values = data[feature_set][feature_name]
+            if scalers[feature_name] is not None:
+                if isinstance(feature_values, list) or isinstance(feature_values, np.ndarray):
+                    scaled_feature_values = scalers[feature_name].transform(
+                        feature_values.reshape(-1, 1)).flatten()
+                else:
+                    scaled_feature_values = scalers[feature_name].transform(
+                        np.array(feature_values).reshape(-1, 1)).flatten()
+            else:
+                scaled_feature_values = feature_values
+            data[feature_set][f"{feature_name}_scaled"] = scaled_feature_values
+
+    return data
+
+
+def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor,
+                          feature_names: np.ndarray | list | str,
+                          scalers: dict) -> np.ndarray:
+    """
+    Inverse scales the input features, can handle single and multi feature input
+    Args:
+        input_feature: input feature as a numpy array
+        feature_names: names of input features as reference for scalers
+        scalers: dict of feature scalers
+
+    Returns:
+        scaled input features as numpy array
+    """
+
+    if isinstance(feature_names, str):
+        feature_names = [feature_names]
+
+    if isinstance(input_feature, np.ndarray) or isinstance(input_feature, torch.Tensor):
+        input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1
+    elif isinstance(input_feature, float) or isinstance(input_feature, int):
+        # handle case, where input is scalar
+        input_scaled = scalers[feature_names[0]].inverse_transform([[input_feature]])[0][0]
+        return input_scaled
+    elif isinstance(input_feature, list):
+        input_feature = np.array(input_feature)
+        input_dim = input_feature.shape[1] if len(input_feature.shape) > 1 else 1
+    else:
+        raise ValueError(f"Unsupported input type: {type(input_feature)}")
+
+    # reshape, if input is one dimensional
+    if input_dim == 1:
+        input_feature = input_feature.reshape(-1, 1)
+
+    input_scaled = input_feature.copy()
+    for i in range(input_dim):
+        input_scaled[:, i] = scalers[feature_names[i]].inverse_transform(np.array([input_scaled[:, i]]))
+
+    return input_scaled
+
+
+def combine_features(cycles: list[dict], feature_config: dict) -> dict:
+    combined_features = dict()
+    for feature_set in feature_config["feature_sets"]:
+        combined_features[feature_set] = dict()
+        if feature_set not in cycles[0]:
+            continue
+        features_in_set = cycles[0][feature_set].keys()
+        for feature in features_in_set:
+            feature_values = [cycles[i][feature_set][feature] for i in range(len(cycles))]
+            if isinstance(feature_values[0], dict):
+                for key in feature_values[0].keys():
+                    feature_array = np.concatenate([feature_values[i][key] for i in range(len(feature_values))])
+                    combined_features[feature_set][key] = feature_array
+            else:
+                feature_array = np.concatenate(feature_values)
+                combined_features[feature_set][feature] = feature_array
+
+    return combined_features

code/new_realtime/utils/dataset_utils.py

@@ -0,0 +1,36 @@
+import os
+import pickle
+
+
+def get_dataset_path(lmdb_base_dir: str,
+                     dataset_name: str) -> str:
+    """
+    Get the path to the dataset in the lmdb directory
+    Args:
+        lmdb_base_dir: base directory of the lmdb dataset
+        dataset_name: name of the dataset
+
+    Returns:
+        path to the dataset
+    """
+
+    dataset_path = os.path.join(lmdb_base_dir, dataset_name)
+    if not os.path.exists(dataset_path):
+        raise FileNotFoundError(f"Dataset {dataset_name} not found in {lmdb_base_dir}")
+    return dataset_path
+
+
+def load_key_stats(lmdb_dir: str) -> dict:
+    """
+    Load key statistics from the LMDB database.
+    :param lmdb_dir: Directory of the LMDB database.
+    :return: Dictionary with statistics.
+    """
+    key_stats_path = f"{lmdb_dir}/key_stats.pickle"
+
+    if not os.path.exists(key_stats_path):
+        return dict()
+
+    with open(key_stats_path, "rb") as f:
+        key_stats = pickle.load(f)
+    return key_stats

code/new_realtime/utils/evaluation.py

@@ -0,0 +1,266 @@
+import lmdb
+import numpy as np
+import sklearn
+import torch
+from torch import nn
+from tqdm import tqdm
+
+from utils.data_utils import get_collated_batch_for_key, get_padding_length
+from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature
+from utils.dataset_utils import load_key_stats
+
+from vsm_datascience_common import constants
+
+
+def evaluate_model(model_configuration: dict,
+                   training_configuration: dict,
+                   test_ids: list,
+                   evaluation_functions: list) -> dict:
+    dataset_dir = model_configuration["feature_config"]["dataset_dir"]
+    lmdb_env = lmdb.open(dataset_dir, readonly=True)
+    # get computation rank
+    if torch.distributed.is_initialized():
+        local_rank = torch.distributed.get_rank()
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    load_fn = model_configuration["model_load_fn"]
+    model = load_fn(model_configuration=model_configuration,
+                    training_configuration=training_configuration,
+                    sample_key=test_ids[0],
+                    device=device,
+                    lmdb_env=lmdb_env)
+    # load key stats to retrieve individual cycles
+    keys_stats = load_key_stats(model_configuration["feature_config"]["dataset_dir"])
+
+    batch_size = training_configuration["batch_size"]
+    predict_fn = model_configuration["predict_fn"]
+    actual_fn = model_configuration["actual_fn"]
+
+    target_features = [x["name"] for x in model_configuration["feature_config"]["target_features"]]
+    ignored_features = model_configuration["feature_config"]["ignored_features"]
+    used_targets = [x for x in target_features if x not in ignored_features]
+    scalers = get_scalers_for_model(model_configuration)
+
+    errors = dict()
+    for test_id in tqdm(test_ids):
+        # fetch stats for key
+        current_key_stats = keys_stats["by_key"][test_id]
+        cycle_stats = current_key_stats["cycle_stats"]
+        for i in range(len(cycle_stats)):
+            # compute cutoffs to isolate current cycle
+            current_cycle_start_cutoff = sum([x["cycle_length"] for x in cycle_stats[:i]])
+            current_cycle_end_cutoff = sum([x["cycle_length"] for x in cycle_stats[:i + 1]])
+            try:
+                batch = get_collated_batch_for_key(test_id, model_configuration,
+                                                   start_cutoff=current_cycle_start_cutoff,
+                                                   end_cutoff=current_cycle_end_cutoff,
+                                                   lmdb_env=lmdb_env)
+            except ValueError:
+                print(f"Skipping {test_id}")
+                continue
+            if batch is None:
+                print(f"Empty batch, skipping {test_id}")
+                continue
+
+            preds = predict_fn(model, batch,
+                               batch_size=batch_size,
+                               device=device)
+            actuals = actual_fn(batch)
+            num_outputs = preds.shape[-1] if len(preds.shape) > 1 else 1
+            scaled_preds = list()
+            for j in range(num_outputs):
+                if isinstance(preds, torch.Tensor):
+                    preds = preds.cpu().numpy()
+
+                if len(preds.shape) == 3:
+                    output = preds[:, 0, j].squeeze()
+                else:
+                    output = preds[:, j] if num_outputs > 1 else preds
+
+                if "loss_functions" in training_configuration and len(training_configuration["loss_functions"]) > j:
+                    if isinstance(training_configuration["loss_functions"][j], nn.BCEWithLogitsLoss):
+                        output = torch.sigmoid(torch.tensor(output)).numpy()
+
+                scaled_output = inverse_scale_feature(output,
+                                                      used_targets[j],
+                                                      scalers)
+
+                scaled_preds.append(scaled_output)
+
+            scaled_actuals = list()
+            for j in range(num_outputs):
+                if len(actuals.shape) == 3:
+                    output = actuals[:, 0, j].squeeze()
+                else:
+                    output = actuals[:, j] if num_outputs > 1 else actuals
+                if isinstance(output, torch.Tensor):
+                    output = output.cpu().numpy()
+                scaled_output = inverse_scale_feature(output,
+                                                      used_targets[j],
+                                                      scalers).ravel()
+
+                scaled_actuals.append(scaled_output)
+
+            for eval_fn in evaluation_functions:
+                if eval_fn is not None:
+                    eval_fn_name = eval_fn["name"]
+                    eval_function = eval_fn["eval_fn"]
+                    eval_fn_index = eval_fn["input_index"]
+                    # skip error fn if actuals are nan, since they are ignored
+                    if any(np.isnan(scaled_actuals[eval_fn_index])):
+                        continue
+                    error = eval_function(scaled_preds[eval_fn_index], scaled_actuals[eval_fn_index])
+
+                    if np.isnan(error):
+                        # skip if error is nan
+                        continue
+
+                    if eval_fn_name not in errors:
+                        errors[eval_fn_name] = dict()
+
+                    if f"after_{i}" not in errors[eval_fn_name]:
+                        errors[eval_fn_name][f"after_{i}"] = list()
+
+                    errors[eval_fn_name][f"after_{i}"].append(error)
+
+    for eval_fn in evaluation_functions:
+        if eval_fn is not None:
+            eval_fn_name = eval_fn["name"]
+            accumulation_fn = eval_fn["accumulation_fn"]
+            for key in errors[eval_fn_name]:
+                if len(errors[eval_fn_name][key]) == 0:
+                    errors[eval_fn_name][key] = np.nan
+                else:
+                    errors[eval_fn_name][key] = accumulation_fn(errors[eval_fn_name][key])
+    return errors
+
+
+def pre_ov_error(preds, actuals, *args, **kwargs):
+    if any(np.isnan(actuals)):
+        return np.nan
+    # get ov day index
+    try:
+        ov_day_index = np.where(actuals == 0)[0][0]
+    except IndexError:
+        # no ov day in actuals
+        return np.nan
+
+    # get pre ov predictions
+    pre_ov_preds = preds[:ov_day_index]
+    pre_ov_actuals = actuals[:ov_day_index]
+
+    if len(pre_ov_preds) == 0 or len(pre_ov_actuals) == 0:
+        return np.nan
+
+    # calculate error
+    error = sklearn.metrics.mean_absolute_error(pre_ov_preds, pre_ov_actuals)
+    return error
+
+
+def after_ov_error(preds, actuals, *args, **kwargs):
+    if any(np.isnan(actuals)):
+        return np.nan
+    # get ov day index
+    try:
+        ov_day_index = np.where(actuals == 0)[0][0]
+    except IndexError:
+        # no ov day in actuals
+        return np.nan
+    # get pre ov predictions
+    after_ov_preds = preds[ov_day_index:]
+    after_ov_actuals = actuals[ov_day_index:]
+
+    if len(after_ov_preds) == 0 or len(after_ov_actuals) == 0:
+        return np.nan
+
+    # calculate error
+    error = sklearn.metrics.mean_absolute_error(after_ov_preds, after_ov_actuals)
+    return error
+
+
+def ov_error(preds, actuals, model_configuration, *args, **kwargs):
+    if any(np.isnan(actuals)):
+        return np.nan
+    # get ov day index
+    try:
+        ov_day_index = np.where(actuals == 0)[0][0]
+    except IndexError:
+        # no ov day in actuals
+        return np.nan
+
+    # get predicted ov index
+    try:
+        pred_ov_index = np.where(preds >= 0)[0][0]
+    except IndexError:
+        # no ov day in actuals
+        return np.nan
+
+    # calculate error
+    error = abs(pred_ov_index - ov_day_index)
+
+    # scale error to account for step size
+    step_size = model_configuration["preprocessing"]["window_shift"]
+    measurements_per_day = constants.MEASUREMENTS_PER_DAY
+    downsampling_factor = model_configuration["preprocessing"]["take_every_nth"]
+    shift_hour_factor = 24 // (measurements_per_day // downsampling_factor) * step_size
+    error_in_days = error * shift_hour_factor // 24
+
+    return error_in_days
+
+
+def day_relative_to_ov_error(preds: np.ndarray | torch.Tensor,
+                             actual: np.ndarray | torch.Tensor,
+                             day_relative_to_ov: int,
+                             model_configuration: dict) -> float:
+    """
+    Calculate the error of the model predictions relative to the ov day
+    Args:
+        preds: predictions
+        actual: actual values
+        day_relative_to_ov: day relative to ov day
+        model_configuration: model configuration
+
+    Returns:
+        error: error of the model predictions relative to the ov day
+    """
+
+    if isinstance(preds, torch.Tensor):
+        preds = preds.cpu().numpy()
+
+    if isinstance(actual, torch.Tensor):
+        actual = actual.cpu().numpy()
+
+    if any(np.isnan(actual)):
+        return np.nan
+
+    # get ov day index
+    try:
+        ov_day_index = np.where(actual == 0)[0][0]
+    except IndexError:
+        # no ov day in actuals
+        return np.nan
+
+    # get index offset factor -> how much time between each step
+    step_size = model_configuration["preprocessing"]["window_shift"]
+    measurements_per_day = constants.MEASUREMENTS_PER_DAY
+    downsampling_factor = model_configuration["preprocessing"]["take_every_nth"]
+    shift_hour_factor = 24 // (measurements_per_day // downsampling_factor) * step_size
+
+    index_offset = int(day_relative_to_ov * (shift_hour_factor // 24))
+
+    if ov_day_index + index_offset >= preds.shape[0]:
+        relative_day_pred = preds[-1] if len(preds.shape) > 1 else preds[-1]
+    elif ov_day_index + index_offset < 0:
+        relative_day_pred = preds[0][0] if len(preds.shape) > 1 else preds[0]
+    else:
+        relative_day_pred = preds[ov_day_index + index_offset][0] if len(preds.shape) > 1 else \
+            preds[ov_day_index + index_offset][0]
+
+    relative_day_actual = day_relative_to_ov
+
+    # calculate error
+    error = np.abs(relative_day_pred - relative_day_actual)
+    return error

code/new_realtime/utils/feature_functions.py

@@ -0,0 +1,372 @@
+from datetime import timedelta
+from typing import Callable
+
+import numpy as np
+from vsm_datascience_common import constants
+from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
+from vsm_datascience_common.cycles.sequences import get_timestamps, get_values
+
+from utils.smoothing import get_curve_composition
+
+
+def get_day_of_week_encoded(cycle: dict, shift: int = 0) -> dict:
+    timestamps = get_timestamps(cycle)
+    weekdays = np.array([x.weekday() for x in timestamps])
+    sine_encoded = np.sin(weekdays * (2 * np.pi / 7))
+    cosine_encoded = np.cos(weekdays * (2 * np.pi / 7))
+    return {
+        "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None],
+        "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None]
+    }
+
+
+def get_hour_of_day_encoded(cycle: dict, shift: int = 0) -> dict:
+    timestamps = get_timestamps(cycle)
+    hours = np.array([x.hour for x in timestamps])
+    sine_encoded = np.sin(hours * (2 * np.pi / 24))
+    cosine_encoded = np.cos(hours * (2 * np.pi / 24))
+    return {
+        "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None],
+        "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None]
+    }
+
+
+def get_month_of_year_encoded(cycle: dict, shift: int = 0) -> dict:
+    timestamps = get_timestamps(cycle)
+    months = np.array([x.month for x in timestamps])
+    sine_encoded = np.sin(months * (2 * np.pi / 12))
+    cosine_encoded = np.cos(months * (2 * np.pi / 12))
+    return {
+        "sine": sine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None],
+        "cosine": cosine_encoded[shift if shift > 0 else 0:shift if shift < 0 else None]
+    }
+
+
+def get_hours_from_start(cycle: dict, shift: int = 0) -> np.ndarray:
+    timestamps = get_timestamps(cycle)
+    return np.arange(len(timestamps))[shift if shift > 0 else 0:shift if shift < 0 else None]
+
+
+def get_temperature(cycle: dict, shift: int = 0) -> np.ndarray:
+    return get_values(cycle)[shift if shift > 0 else 0:shift if shift < 0 else None]
+
+
+def get_fertility_curve(cycle_length: int,
+                        ov_index: int,
+                        fertility_base_curve: np.ndarray | list,
+                        base_curve_offset: int) -> np.ndarray:
+    fert_prob = np.zeros((cycle_length,))
+    if ov_index + base_curve_offset < 0:
+        length_in_cycle = len(fertility_base_curve) + ov_index + base_curve_offset
+        fert_prob[:length_in_cycle] = fertility_base_curve[-length_in_cycle:]
+    elif ov_index + base_curve_offset + len(fertility_base_curve) > cycle_length:
+        length_in_cycle = cycle_length - ov_index - base_curve_offset
+        fert_prob[ov_index + base_curve_offset:] = fertility_base_curve[:length_in_cycle]
+    else:
+        fert_prob[ov_index + base_curve_offset:
+                  ov_index + base_curve_offset + len(fertility_base_curve)] = fertility_base_curve
+    return fert_prob
+
+
+def get_fertility_probability_base_curve() -> (np.ndarray, int):
+    """
+    This function returns the base curve for the fertility probability as well as the offset to the ovulation day in measurements
+    :return: numpy array with the base curve resampled to the number of measurements per day
+    """
+    fertilization_chance_curve = [0, 0.07, 0.12, 0.25, 0.3, 0.18, 0]
+    offset_to_ov = -6
+
+    new_indices = np.arange(len(fertilization_chance_curve) * constants.MEASUREMENTS_PER_DAY)
+    fertility_chance_curve_resampled = np.interp(new_indices,
+                                                 np.linspace(0, len(new_indices), num=len(fertilization_chance_curve)),
+                                                 fertilization_chance_curve)
+    ov_offset_resampled = offset_to_ov * constants.MEASUREMENTS_PER_DAY
+    return fertility_chance_curve_resampled, ov_offset_resampled
+
+
+def get_fertility_probability(cycle: dict, shift: int = 0) -> np.ndarray:
+    # check, if biphasic
+    if "classification_results" in cycle and "results" in cycle["classification_results"][0]:
+        if cycle["classification_results"][0]["results"]["predicted_class"] != "biphasic":
+            return np.zeros((len(get_values(cycle)) - abs(shift, )))
+    else:
+        return np.zeros((len(get_values(cycle)) - abs(shift, )))
+
+    if "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]:
+        try:
+            ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"]
+            timestamps = get_timestamps(cycle)
+            ov_index = np.where(timestamps >= ov_timestamp)[0][0]
+            fertility_chance_curve_resampled, ov_offset_resampled = get_fertility_probability_base_curve()
+            return get_fertility_curve(len(get_values(cycle)), ov_index, fertility_chance_curve_resampled,
+                                       ov_offset_resampled)[shift if shift > 0 else 0:shift if shift < 0 else None]
+        except:
+            return np.zeros((len(get_values(cycle)) - abs(shift, )))
+    else:
+        return np.zeros((len(get_values(cycle)) - abs(shift, )))
+
+
+def get_ov_over_probability(cycle: dict, shift: int = 0) -> np.ndarray:
+    """
+    Get the ovulation over probability for a cycle
+    :param cycle: cycle data
+    :param shift: shift to apply
+    :return: ovulation over probability
+    """
+    if "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]:
+        try:
+            ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"]
+            timestamps = get_timestamps(cycle)
+            ov_index = np.where(timestamps > ov_timestamp)[0][0]
+            ov_over_probability = np.zeros(len(get_values(cycle)))
+            ov_over_probability[ov_index:] = 1
+            return ov_over_probability[shift if shift > 0 else 0:shift if shift < 0 else None]
+        except:
+            return np.zeros(len(get_values(cycle)))[shift if shift > 0 else 0:shift if shift < 0 else None]
+    else:
+        return np.zeros(len(get_values(cycle)))[shift if shift > 0 else 0:shift if shift < 0 else None]
+
+
+def get_days_relative_to_ov(cycle: dict) -> np.ndarray:
+    """
+    Get the time in days until the next OV event.
+    """
+    timestamps = get_timestamps(cycle)
+    if ("ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]) and (
+            "classification_results" in cycle and "results" in cycle["classification_results"][0] and
+            cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"):
+        try:
+            ov_timestamp = cycle["ov_detection_results"][0]["results"]["ovulation_timestamp"]
+            ov_index = np.where(timestamps > ov_timestamp)[0][0]
+            ov_timestamp = timestamps[ov_index]
+            days_until_ov = np.empty(len(timestamps))
+            for i in range(len(timestamps)):
+                current_day = timestamps[i]
+                day_diff = ov_timestamp - current_day
+                if day_diff.days < -1:
+                    pass
+                # days_until_ov[i] = max(day_diff.days -1)
+                # use - days, as days relative to event
+                days_until_ov[i] = - day_diff.days
+
+            return np.array(days_until_ov)
+        except:
+            pass
+
+    return np.full((len(timestamps),), np.nan)
+
+
+def get_ov_day(cycle: dict) -> np.ndarray:
+    """
+    Get the ovulation day for a cycle
+    :param cycle: cycle data
+    :return: ovulation day
+    """
+    values = get_values(cycle)
+    seq_len = len(values)
+    is_biphasic = "classification_results" in cycle and "results" in cycle["classification_results"][0] and \
+                  cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"
+    has_ov_results = "ov_detection_results" in cycle and "results" in cycle["ov_detection_results"][0]
+    if is_biphasic and has_ov_results:
+        try:
+            ov_day = cycle["ov_detection_results"][0]["results"]["ovulation_day"]
+            return np.full((seq_len,), ov_day)
+        except:
+            pass
+
+    return np.full((seq_len,), np.nan)
+
+
+def get_is_biphasic(cycle: dict) -> np.ndarray:
+    """
+    Get the is_biphasic flag for a cycle
+    :param cycle: cycle data
+    :return: is_biphasic flag
+    """
+    values = get_values(cycle)
+    seq_len = len(values)
+    is_biphasic = "classification_results" in cycle and "results" in cycle["classification_results"][0] and \
+                  cycle["classification_results"][0]["results"]["predicted_class"] == "biphasic"
+    if is_biphasic:
+        return np.full((seq_len,), 1)
+
+    return np.full((seq_len,), 0)
+
+
+def get_curve_composition_as_features(cycle: dict, shift: int = 0) -> dict:
+    composition = get_curve_composition(get_values(cycle))
+    return {
+        "trend": composition[0][shift if shift > 0 else 0:shift if shift < 0 else None],
+        "seasonal": composition[1][shift if shift > 0 else 0:shift if shift < 0 else None],
+        "residual": composition[2][shift if shift > 0 else 0:shift if shift < 0 else None],
+        "smoothed": composition[3][shift if shift > 0 else 0:shift if shift < 0 else None]
+    }
+
+
+def get_rolling_average_with_padding(cycle: dict,
+                                     shift: int = 0,
+                                     rolling_average_window_length: int = constants.MEASUREMENTS_PER_DAY) -> np.ndarray:
+    """
+    Get the rolling average of the temperature values with padding, so that the rolling average has the same length as the
+    original temperature values
+    Args:
+        cycle (dict): cycle data
+        shift (int): shift to apply
+        rolling_average_window_length (int): window length for the rolling average, default is one day (288 measurements)
+
+    Returns:
+        np.ndarray: rolling average of the temperature values with padding
+
+    """
+
+    values = get_values(cycle)
+    values_padded = np.concatenate(
+        [np.full((constants.MEASUREMENTS_PER_DAY,), np.mean(values[:constants.MEASUREMENTS_PER_DAY])), values])
+    rolling_average = np.convolve(values_padded, np.ones(rolling_average_window_length) / rolling_average_window_length,
+                                  mode="valid")[-len(values):]
+    return rolling_average[shift if shift > 0 else 0:shift if shift < 0 else None]
+
+
+def get_window_fn(cycle: dict,
+                  window_size: int,
+                  fn: Callable) -> np.ndarray:
+    """
+    Apply a window function to a time series.
+    :param cycle: cycle data
+    :param window_size: window size
+    :return: array with minimum value in window for each value in input array
+    """
+    values = get_values(cycle)
+    values_padded = np.concatenate([np.full(window_size - 1, values[0]), values])
+
+    windows = np.lib.stride_tricks.sliding_window_view(values_padded, window_shape=window_size)
+    return fn(windows)
+
+
+def get_cycle_length_stats(cycle: dict) -> dict:
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    cycles_dates = list(get_cycles_collection().find({
+        "user_id": user_id, "ends_at": {"$exists": True},
+        "starts_at": {"$lt": cycle["starts_at"]},
+    },
+        {"starts_at": 1, "ends_at": 1}))
+    if len(cycles_dates) == 0:
+        # use the average cycle length as fallback
+        average_cycle_length = 37.13
+        cycle_length_std = 12.42
+    else:
+        cycle_lengths = [x["ends_at"] - x["starts_at"] for x in cycles_dates]
+        average_cycle_length = np.mean(cycle_lengths).total_seconds() / (60 * 60 * 24)
+        cycle_length_std = np.std([x.total_seconds() for x in cycle_lengths]) / (60 * 60 * 24)
+
+    return {
+        "average_cycle_length": np.full((cycle_length,), average_cycle_length),
+        "cycle_length_std": np.full((cycle_length,), cycle_length_std),
+    }
+
+
+def get_average_ovulation_day(cycle: dict) -> float:
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    cycles_dates = list(
+        get_cycles_collection().find({
+            "user_id": user_id,
+            "ov_detection_results.0.results": {"$exists": True},
+            "starts_at": {"$lt": cycle["starts_at"]},
+        },
+            {"ov_detection_results": 1}))
+    if len(cycles_dates) == 0:
+        # use the average ovulation day as fallback
+        average_ovulation_day = 18.9
+    else:
+        ovulation_days = [x["ov_detection_results"][0]["results"]["ovulation_day"] for x in cycles_dates]
+        average_ovulation_day = np.mean(ovulation_days)
+    return np.full((cycle_length,), average_ovulation_day)
+
+
+def get_ovulation_std(cycle: dict) -> float:
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    cycles_dates = list(
+        get_cycles_collection().find({
+            "user_id": user_id,
+            "ov_detection_results.0.results": {"$exists": True},
+            "starts_at": {"$lt": cycle["starts_at"]},
+        },
+            {"ov_detection_results": 1}))
+    if len(cycles_dates) == 0:
+        # use the average ovulation std as fallback
+        ovulation_std = 4.07
+    else:
+        ovulation_days = [x["ov_detection_results"][0]["results"]["ovulation_day"] for x in cycles_dates]
+        ovulation_std = np.std(ovulation_days)
+    return np.full((cycle_length,), ovulation_std)
+
+
+def get_biphasic_fraction(cycle: dict) -> float:
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    cycle_dates = list(
+        get_cycles_collection().find({
+            "user_id": user_id,
+            "classification_results.0.results": {"$exists": True},
+            "starts_at": {"$lt": cycle["starts_at"]},
+        },
+            {"classification_results": 1}))
+    if len(cycle_dates) == 0:
+        # use the average cycle length as fallback
+        biphasic_fraction = 0.9449
+    else:
+        biphasic = [x for x in cycle_dates if
+                    x["classification_results"][0]["results"]["predicted_class"] == "biphasic"]
+        biphasic_fraction = len(biphasic) / len(cycle_dates)
+    return np.full((cycle_length,), biphasic_fraction)
+
+
+def get_num_cycles(cycle: dict) -> int:
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    cycles_dates = list(get_cycles_collection().find({"user_id": user_id,
+                                                      "ends_at": {"$exists": True},
+                                                      "starts_at": {"$lt": cycle["starts_at"]}
+                                                      },
+                                                     {"starts_at": 1, "ends_at": 1}))
+    num_cycles = len(cycles_dates)
+    return np.full((cycle_length,), num_cycles)
+
+
+def get_average_temperatures(cycle: dict) -> dict:
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    cycles_dates = list(
+        get_cycles_collection().find({"user_id": user_id, "ov_detection_results.0.results": {"$exists": True}},
+                                     {"starts_at": 1, "ov_detection_results": 1, "measurements": 1}))
+
+    if len(cycles_dates) == 0:
+        return {
+            "pre_ov_temperatures": np.full((cycle_length,), 37.07),
+            "post_ov_temperatures": np.full((cycle_length,), 37.37),
+        }
+
+    pre_ov_temperatures = []
+    post_ov_temperatures = []
+    for x in cycles_dates:
+        ov_detection = x["ov_detection_results"][0]["results"]
+        values = get_values(x)
+        timestamps = get_timestamps(x)
+        cutoff_day = x["starts_at"] + timedelta(days=ov_detection["ovulation_day"])
+        cutoff_index = np.where(np.array(timestamps) >= cutoff_day)[0][0]
+        pre_ov_temperatures.append(values[:cutoff_index])
+        post_ov_temperatures.append(values[cutoff_index:])
+
+    pre_ov_temperatures = np.concatenate(pre_ov_temperatures)
+    post_ov_temperatures = np.concatenate(post_ov_temperatures)
+
+    pre_ov_temperatures = np.mean(pre_ov_temperatures)
+    post_ov_temperatures = np.mean(post_ov_temperatures)
+
+    return {
+        "pre_ov_temperatures": np.full((cycle_length,), pre_ov_temperatures),
+        "post_ov_temperatures": np.full((cycle_length,), post_ov_temperatures),
+    }

code/new_realtime/utils/lmdb_utils.py

@@ -0,0 +1,119 @@
+import pickle
+
+import pandas as pd
+
+
+def save_to_lmdb(env, key, dataset):
+    """
+    Saves the given dataset to an LMDB environment with the given key.
+    :param env: LMDB environment
+    :param key: key to save the dataset to
+    :param dataset: tuple of pandas dataframes
+    """
+
+    with env.begin(write=True) as txn:
+        txn.put(key.encode('ascii'), pickle.dumps(dataset))
+
+
+def load_from_lmdb(env, key):
+    """
+    Loads a dataset from an LMDB environment with the given key.
+    :param env: LMDB environment
+    :param key: key of the dataset to load
+    :return: key and tuple of pandas dataframes (input, context, output)
+    """
+    with env.begin(write=False) as txn:
+        try:
+            data = pickle.loads(txn.get(key.encode('ascii')))
+            return data
+        except TypeError:
+            raise KeyError(key)
+
+
+def delete_from_lmdb(env, key):
+    """
+    Deletes a dataset from an LMDB environment with the given key.
+    :param env: LMDB environment
+    :param key: key of the dataset to delete
+    """
+    with env.begin(write=True) as txn:
+        txn.delete(key.encode('ascii'))
+
+
+def clear_lmdb(env):
+    """
+    Clears all datasets from an LMDB environment.
+    :param env: LMDB environment
+    """
+    with env.begin(write=True) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            txn.delete(key)
+
+
+def lmdb_dataset_generator(env):
+    """
+    Generator function to yield datasets from an LMDB environment.
+    :param env: LMDB environment
+    :return: generator
+    """
+    with env.begin(write=False) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            data = pickle.loads(value)
+            yield data
+
+
+def lmdb_contains(env, substring) -> bool:
+    """
+    Checks if the given substring is contained in any of the keys of the LMDB environment.
+    :param env: LMDB environment
+    :param substring: substring to search for
+    :return: boolean
+    """
+    with env.begin(write=False) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            if substring in key.decode('ascii'):
+                return True
+    return False
+
+
+def lmdb_substring_key_search(env, substring):
+    """
+    Searches for keys in the LMDB environment that contain the given substring.
+    :param env: LMDB environment
+    :param substring: substring to search for
+    :return: list of keys
+    """
+    keys = []
+    with env.begin(write=False) as txn:
+        cursor = txn.cursor()
+        for key, value in cursor:
+            if substring in key.decode('ascii'):
+                keys.append(key)
+    return keys
+
+
+def get_lmdb_keys(env, limit: int = None):
+    """
+    Get all keys in the LMDB environment.
+    :param env: LMDB environment
+    :param limit: maximum number of keys to return
+    :return: list of keys
+    """
+    with env.begin(write=False) as txn:
+        with txn.cursor() as cursor:
+            keys = [key.decode("ascii") for key in cursor.iternext(keys=True, values=False)]
+
+    return keys
+
+
+def get_lmdb_keyspace_size(env):
+    """
+    Get the number of keys in the LMDB environment.
+    :param env: LMDB environment
+    :return: number of keys
+    """
+    with env.begin(write=False) as txn:
+        return txn.stat()['entries']

code/new_realtime/utils/loss_functions.py

@@ -0,0 +1,118 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class QuantileLoss(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.register_buffer('q', torch.tensor(config.quantiles))
+
+    def forward(self, predictions, targets):
+        diff = predictions - targets.unsqueeze(-1).expand(-1, -1, -1, predictions.shape[-1])
+        ql = (1 - self.q) * F.relu(diff) + self.q * F.relu(-diff)
+        losses = ql.view(-1, ql.shape[-1]).mean(0)
+        return losses
+
+
+def qrisk(pred, tgt, quantiles):
+    diff = pred - tgt
+    ql = (1 - quantiles) * np.clip(diff, 0, float('inf')) + quantiles * np.clip(-diff, 0, float('inf'))
+    losses = ql.reshape(-1, ql.shape[-1])
+    normalizer = np.abs(tgt).mean()
+    risk = 2 * losses / normalizer
+    return risk.mean(0)
+
+
+def weighted_bce_loss_fn():
+    bce = nn.BCEWithLogitsLoss(reduction="none")  # Use reduction="none" to get per-element loss
+    use_class_weights = False
+    time_weight_factor = 1
+
+    def _loss_fn(y_pred, y_true):
+        loss_list = []  # To accumulate per-sample loss values
+        for i in range(len(y_pred)):
+            current_pred = y_pred[i]
+            current_true = y_true[i]
+            loss_raw = bce(current_pred, current_true)
+
+            if use_class_weights:
+                # Calculate class weights using torch operations
+                try:
+                    factor = torch.max(current_true) / torch.mean(current_true) / 2
+                    non_zero_indices = current_true != 0
+                    class_weights = torch.ones_like(current_pred)
+                    class_weights[non_zero_indices] = factor
+                except Exception:
+                    class_weights = torch.ones_like(current_pred)
+            else:
+                class_weights = torch.ones_like(current_pred)
+
+            # Calculate time weights using torch operations
+            try:
+                # Find indices where the event occurs
+                event_indices = torch.where(current_true == 1)[0]
+                if len(event_indices) > 0:
+                    first_event_index = event_indices[0].item()
+                    time_weights = torch.ones_like(current_pred)
+                    # Weight the loss before the first event higher if needed
+                    time_weights[:first_event_index] = time_weight_factor
+                else:
+                    time_weights = torch.ones_like(current_pred)
+            except Exception:
+                time_weights = torch.ones_like(current_pred)
+
+            # Apply the weights to the raw loss
+            loss = loss_raw * class_weights * time_weights
+            # Optionally, take the mean over the time dimension
+            loss_list.append(loss.mean())
+
+        # Aggregate the loss for the entire batch
+        return torch.stack(loss_list).mean()
+
+    return _loss_fn
+
+
+def weighted_mse_prob_loss_fn():
+    mse = nn.MSELoss(reduce=False)
+    after_event_factor = 1
+    use_weighted = False
+
+    def _loss_fn(y_pred, y_true):
+        nonlocal mse, after_event_factor
+        loss_list = list()
+        for i in range(len(y_pred)):
+            current_pred = y_pred[i]
+            current_true = y_true[i]
+            loss_raw = mse(current_pred, current_true)
+            # calculate factor to use for weighting, only works with probability values
+            if use_weighted:
+                try:
+                    factor = torch.max(current_true) / torch.mean(current_true) / 2
+                    non_zero_indices = current_true != 0
+                    class_weights = np.ones(len(current_true))
+                    class_weights[non_zero_indices] = factor
+                except:
+                    # if there are no non-zero indices, set all weights to 1
+                    class_weights = np.ones(len(current_true))
+            else:
+                # if we don't want to weight the loss, set all weights to 1
+                class_weights = np.ones(len(current_true))
+
+            # calculate time weight, here, after the event, the loss should be higher
+            try:
+                last_event_index = np.where(current_true != 0)[-1]
+                time_weights = np.zeros(len(current_true))
+                # here, we want to weight the loss after the 'fertility curve' higher, so that we don't get the signal too late
+                time_weights[last_event_index:] = after_event_factor
+            except:
+                # if there are no non-zero indices, set all weights to 1
+                time_weights = np.ones(len(current_true))
+
+            loss = loss_raw * torch.tensor(class_weights) * torch.tensor(time_weights)
+            loss_list.append(loss.mean())
+
+        return torch.stack(loss_list).mean()
+
+    return _loss_fn

code/new_realtime/utils/model_utils.py

@@ -0,0 +1,62 @@
+import os
+import pickle
+from datetime import datetime
+
+from utils.utils import get_config_id
+
+
+def get_model_config(base_config: dict,
+                     base_result_dir: str,
+                     dataset_base_dir: str):
+    # get config identifier
+    # name_for_current_config = base_config["model_name"] + "_" + get_config_id(base_config)
+    date_part = datetime.now().strftime("%Y_%m_%d_%H_%M")
+    name_for_current_config = base_config["model_name"] + "_" + date_part
+    model_dir = os.path.abspath(f"{base_result_dir}/{name_for_current_config}")
+    if not os.path.exists(model_dir):
+        # create config
+        os.makedirs(model_dir)
+
+        # append identifier to config
+        config = base_config.copy()
+        config["id"] = name_for_current_config
+
+        # append paths
+        config["model_dir"] = model_dir
+        config["dataset_dir"] = os.path.join(dataset_base_dir, config["feature_config"]["feature_set_name"])
+        config["feature_config"]["dataset_dir"] = config["dataset_dir"]
+
+        # save model configuration
+        with open(f"{model_dir}/model_configuration.pickle", "wb") as f:
+            pickle.dump(config, f)
+    else:
+        # fetch config
+        with open(f"{model_dir}/model_configuration.pickle", "rb") as f:
+            config = pickle.load(f)
+
+        # append paths
+        config["model_dir"] = model_dir
+        config["dataset_dir"] = os.path.join(dataset_base_dir, config["feature_config"]["feature_set_name"])
+        config["feature_config"]["dataset_dir"] = config["dataset_dir"]
+    return config
+
+
+def get_model_config_from_file(model_dir: str,
+                               model_base_dir: str,
+                               lmdb_base_dir: str):
+    # fetch config
+    with open(f"{model_dir}/model_configuration.pickle", "rb") as f:
+        config = pickle.load(f)
+
+    # update paths based on base directories
+    config["model_dir"] = os.path.abspath(f"{model_base_dir}/{config['id']}")
+    config["feature_config"]["dataset_dir"] = os.path.abspath(
+        f"{lmdb_base_dir}/{config['feature_config']['feature_set_name']}")
+
+    return config
+
+
+def save_model_config(model_dir: str, model_config: dict):
+    # save model configuration
+    with open(f"{model_dir}/model_configuration.pickle", "wb") as f:
+        pickle.dump(model_config, f)

code/new_realtime/utils/smoothing.py

@@ -0,0 +1,50 @@
+import numpy as np
+from scipy.signal import butter, filtfilt
+from statsmodels.tsa.stl._stl import STL
+
+
+def highpass_filter(data, cutoff_freq, fs=288):
+    nyquist = 0.5 * fs
+    normal_cutoff = cutoff_freq / nyquist
+    b, a = butter(N=3, Wn=normal_cutoff, btype="high", analog=False)
+    return filtfilt(b, a, data)
+
+
+def mirror_extend(series, extend_len):
+    """Mirrors the beginning and end of the time series to stabilize smoothing."""
+    # Mirror extension
+    start_extension = series[:extend_len][::-1]  # Reverse first part
+    end_extension = series[-extend_len:][::-1]  # Reverse last part
+
+    extended_series = np.concatenate([start_extension, series, end_extension])
+    return extended_series
+
+
+def get_trend(input_curve: np.ndarray | list, measurements_per_day: int = 288) -> np.ndarray:
+    extension_len = 3
+    extended_input_curve = mirror_extend(input_curve, extension_len * measurements_per_day)
+    stl = STL(extended_input_curve, period=measurements_per_day, robust=False, trend=measurements_per_day * 14 + 1)
+    trend = stl.fit().trend
+    return trend[extension_len * measurements_per_day:-extension_len * measurements_per_day]
+
+
+def get_curve_composition(input_curve: np.ndarray | list, measurements_per_day: int = 288) -> tuple:
+    """
+    Decomposes the input curve into trend, seasonal, residual and smoothed components.
+    :param input_curve: raw input curve
+    :param measurements_per_day: seasonal period, here: measurements per day -> 288
+    :return: composition of curve as tuple (trend, seasonal, residual, smoothed)
+    """
+    extension_len = 3
+    extended_input_curve = mirror_extend(input_curve, extension_len * measurements_per_day)
+    stl_results = STL(extended_input_curve, period=measurements_per_day, robust=False).fit()
+    long_term = (extended_input_curve - stl_results.seasonal)
+    wiggles = highpass_filter(long_term, 0.1)
+    smoothed = long_term - wiggles
+
+    return (
+        stl_results.trend[extension_len * measurements_per_day:-extension_len * measurements_per_day],
+        stl_results.seasonal[extension_len * measurements_per_day:-extension_len * measurements_per_day],
+        stl_results.resid[extension_len * measurements_per_day:-extension_len * measurements_per_day],
+        smoothed[extension_len * measurements_per_day:-extension_len * measurements_per_day],
+    )

code/new_realtime/utils/training.py

@@ -0,0 +1,330 @@
+import math
+
+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 torch.utils.tensorboard import SummaryWriter
+from tqdm import tqdm
+
+from utils.data_utils import LMDBIterableDataset
+from utils.utils import get_logger
+
+
+def process_tft_batch(model: nn.Module,
+                      data_iterator: IterableDataset,
+                      loss_functions: list,
+                      device: torch.device,
+                      model_configuration: dict) -> torch.Tensor:
+    batch = next(data_iterator)
+    batch = {k: v.to(device) for k, v in batch.items() if v is not None}
+    input_window_length = model_configuration["model_parameters"]["encoder_length"]
+
+    preds = model(batch).cpu()  # [B, decoder_len, Q]
+    target = batch["target"][:, input_window_length:, :].cpu()  # match decoder segment
+    loss = get_x_y_loss(preds, target, loss_functions)
+    return loss
+
+
+def get_x_y_loss(pred: torch.Tensor,
+                 target: torch.Tensor,
+                 loss_functions: list,
+                 *args, **kwargs) -> torch.Tensor:
+    if len(loss_functions) > 1:
+        losses = list()
+        for dim in range(target.shape[-1]):
+            if len(pred.shape) > 2:
+                current_preds = pred[:, :, dim].ravel()
+            else:
+                current_preds = pred[:, dim]
+            if len(target.shape) > 2:
+                current_target = target[:, :, dim].ravel()
+            else:
+                current_target = target[:, dim].ravel()
+            # skip dimension, if it contains only NaN values, as loss cens
+            nan_indices = torch.isnan(current_target)
+            if torch.all(nan_indices):
+                continue
+            current_target = current_target[~nan_indices]
+            current_preds = current_preds[~nan_indices]
+            if len(current_target) == 0:
+                continue
+
+            loss = loss_functions[dim](current_preds, current_target)
+            losses.append(loss)
+        loss = torch.stack(losses).mean()
+    else:
+        nan_indices = torch.isnan(target)
+        if torch.all(nan_indices):
+            return torch.tensor(0.0)
+        current_target = target[~nan_indices]
+        current_preds = pred[~nan_indices]
+        if len(current_target) == 0:
+            return torch.tensor(0.0)
+        loss = loss_functions[0](current_preds, current_target)
+
+    return loss
+
+
+def get_model_loss(model: nn.Module,
+                   data_iterator: IterableDataset,
+                   loss_functions: list,
+                   device: str,
+                   *args, **kwargs) -> torch.Tensor:
+    batch_x, batch_y = next(data_iterator)
+    batch_x = batch_x.to(device).float()
+    target = batch_y.to(device).float()
+    pred = model(batch_x)
+    loss = get_x_y_loss(pred, target, loss_functions)
+    return loss
+
+
+def get_ranked_ids(all_ids, epoch, rank, world_size, base_seed=42):
+    """
+    Get ranked ids for distributed training.
+    Args:
+        all_ids: list of all available ids
+        epoch: current epoch
+        rank: rank of the current process
+        world_size: number of processes
+        base_seed: base seed for random number generator
+    Returns:
+        list of ids for the current process
+    """
+    g = torch.Generator()
+    g.manual_seed(base_seed + epoch)
+    permuted = torch.randperm(len(all_ids), generator=g).tolist()
+    return [all_ids[i] for i in permuted[rank::world_size]]
+
+
+def train_model(model: nn.Module,
+                model_configuration: dict,
+                training_configuration: dict,
+                train_dataset: LMDBIterableDataset,
+                val_dataset: LMDBIterableDataset,
+                log_dir: str = "./logs",
+                logger=None) -> torch.nn.Module:
+    if logger is None:
+        logger = get_logger(__name__, f"{log_dir}/{model_configuration['id']}_{training_configuration['id']}.log")
+
+    learning_parameters = training_configuration["learning_parameters"]
+    num_epochs = learning_parameters["epochs"]
+    patience = learning_parameters["patience"]
+
+    training_id = training_configuration["id"]
+
+    # get computation rank
+    if torch.distributed.is_initialized():
+        local_rank = torch.distributed.get_rank()
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+        world_size = torch.distributed.get_world_size()
+    else:
+        local_rank = 0
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        world_size = 1
+
+    logger.info(f"Rank {local_rank}: Using device: {device}, world size: {world_size}")
+
+    # get subsets for distributed training
+    if torch.distributed.is_initialized():
+        all_train_ids = train_dataset.lmdb_keys
+        train_subsets = [get_ranked_ids(all_train_ids, i, local_rank, world_size) for i in range(num_epochs)]
+        # calc total number of steps for gpu, as it is dependent on subsets
+        total_train_steps = sum([train_dataset.get_length_of_data_subset(subset) for subset in train_subsets])
+
+        all_val_ids = val_dataset.lmdb_keys
+        val_subsets = [get_ranked_ids(all_val_ids, i, local_rank, world_size) for i in range(num_epochs)]
+    else:
+        train_subsets = [train_dataset.lmdb_keys] * num_epochs
+        val_subsets = [val_dataset.lmdb_keys] * num_epochs
+        total_train_steps = len(train_dataset)
+
+    # log the number of training steps for each epoch
+    train_subset_lengths = {f"epoch_{i}": len(subset) for i, subset in enumerate(train_subsets)}
+    logger.info(f"Train_subsets:  {train_subset_lengths}")
+    logger.info(f"Rank {local_rank}: Total training steps: {total_train_steps}")
+
+    # initialize loaders for non distributed
+    if not torch.distributed.is_initialized():
+        # set the subsets for the datasets
+        train_dataset.set_key_subset(train_subsets[0])
+        val_dataset.set_key_subset(val_subsets[0])
+
+        # create data loaders
+        train_dataloader = DataLoader(
+            train_dataset,
+            batch_size=None,
+            num_workers=4,
+        )
+        val_dataloader = DataLoader(
+            val_dataset,
+            batch_size=None,
+            num_workers=4,
+        )
+
+    logger.info(f"Rank {local_rank}: Training {training_id} with {num_epochs} epochs")
+    logger.info(f"Rank {local_rank}: Training on {torch.cuda.device_count()} GPUs")
+    logger.info(
+        f"Rank {local_rank}: Current device: {torch.cuda.get_device_name(local_rank)} on local rank {local_rank}")
+
+    model.to(device)
+
+    # load training state from training configuration, if available
+    optimizer = AdamW(model.parameters(), lr=learning_parameters["learning_rate"])
+    scheduler = OneCycleLR(optimizer,
+                           max_lr=learning_parameters["learning_rate"],
+                           # make sure to use length of full dataset here
+                           total_steps=total_train_steps)
+    current_epoch = 1
+
+    loss_functions = training_configuration["loss_functions"]
+    # loss_fn = nn.MSELoss()
+    writer = SummaryWriter(log_dir=f'{log_dir}/{model_configuration["id"]}_{training_id}', )
+
+    best_val_loss = math.inf
+    epochs_no_improve = 0
+    log_every_n_steps = max(len(train_dataset) // 500, 1)
+
+    batch_loss_fn = model_configuration["batch_loss_fn"]
+
+    for epoch in range(current_epoch, num_epochs + 1):
+        logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs}")
+        model.train()
+        total_train_loss = 0
+
+        # on distributed training, reshuffle the data
+        if torch.distributed.is_initialized():
+            # update the datasets with the new ids
+            train_dataset.set_key_subset(train_subsets[epoch - 1])
+            val_dataset.set_key_subset(val_subsets[epoch - 1])
+
+            # recreate data loaders
+            train_dataloader = DataLoader(
+                train_dataset,
+                batch_size=None,
+                num_workers=4,
+            )
+            val_dataloader = DataLoader(
+                val_dataset,
+                batch_size=None,
+                num_workers=4,
+            )
+
+        iterator = iter(train_dataloader)
+        for step in tqdm(range(len(train_dataloader)), total=len(train_dataloader)):
+            loss = batch_loss_fn(model,
+                                 iterator,
+                                 loss_functions,
+                                 device,
+                                 model_configuration)
+
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+            scheduler.step()
+
+            total_train_loss += loss.item()
+            if local_rank == 0:
+                if step % log_every_n_steps == 0:
+                    writer.add_scalar("Loss/Train_Step", loss.item(),
+                                      ((epoch - 1) * len(train_dataloader) + step) * training_configuration[
+                                          "batch_size"])
+                    writer.add_scalar("LR", scheduler.get_last_lr()[0],
+                                      ((epoch - 1) * len(train_dataloader) + step) * training_configuration[
+                                          "batch_size"])
+                    writer.flush()
+
+        avg_train_loss = total_train_loss / len(train_dataloader)
+        logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs} done. Train loss: {avg_train_loss:.4f}")
+
+        # Validation
+        model.eval()
+        total_val_loss = 0
+        logger.info(f"Rank {local_rank}: Validation")
+        with torch.no_grad():
+            val_iter = iter(val_dataloader)
+            for step in tqdm(range(len(val_dataloader))):
+                loss = batch_loss_fn(model,
+                                     val_iter,
+                                     loss_functions,
+                                     device,
+                                     model_configuration)
+
+                total_val_loss += loss.item()
+
+        if len(val_dataloader) == 0:
+            logger.info(f"Rank {local_rank}: Validation set is empty, using 0 as validation loss.")
+            avg_val_loss = None
+        else:
+            avg_val_loss = total_val_loss / len(val_dataloader)
+
+        # publish validation loss and wait for other gpus
+        if torch.distributed.is_initialized():
+            if avg_val_loss is not None:
+                avg_val_loss_global = torch.tensor(avg_val_loss, device=device, dtype=torch.float32)
+            torch.distributed.all_reduce(avg_val_loss_global)
+            avg_val_loss_global /= torch.distributed.get_world_size()
+
+            avg_train_loss_global = torch.tensor(avg_train_loss).to(device)
+            torch.distributed.all_reduce(avg_train_loss_global)
+            avg_train_loss_global /= torch.distributed.get_world_size()
+        else:
+            avg_val_loss_global = torch.tensor(avg_val_loss)
+            avg_train_loss_global = torch.tensor(avg_train_loss)
+
+        # only rank 0 checks for early stopping
+        if local_rank == 0:
+            logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs} done. Val loss: {avg_val_loss:.4f}")
+            writer.add_scalar("Loss/Train_Epoch", avg_train_loss_global, epoch)
+            writer.add_scalar("Loss/Val_Epoch", avg_val_loss_global, epoch)
+            writer.flush()
+
+            # Early stopping
+            should_stop = False
+            if avg_val_loss_global < best_val_loss:
+                logger.info(
+                    f"Rank {local_rank}: Validation loss improved from {best_val_loss:.4f} to {avg_val_loss_global:.4f}.")
+                best_val_loss = avg_val_loss_global
+                epochs_no_improve = 0
+                # torch.save(model.state_dict(), os.path.join(model_configuration["id"], "model.pt"))
+                save_fn = model_configuration["model_save_fn"]
+                save_fn(model, training_configuration)
+            else:
+                epochs_no_improve += 1
+                logger.info(
+                    f"Rank {local_rank}: No improvement in validation loss, no-improve count: {epochs_no_improve}")
+                if epochs_no_improve >= patience:
+                    logger.info("Early stopping triggered.")
+                    # broadcast stop signal to all gpus
+                    should_stop = True
+        else:
+            should_stop = None
+
+        if torch.distributed.is_initialized():
+            if local_rank == 0:
+                should_stop_tensor = torch.tensor([int(should_stop)], device=device)
+            else:
+                should_stop_tensor = torch.zeros(1, dtype=torch.uint8, device=device)  # safe default
+            torch.distributed.broadcast(should_stop_tensor, src=0)
+            should_stop = bool(should_stop_tensor.item())
+
+        if should_stop:
+            logger.info(f"Rank {local_rank}: Stopping training.")
+            break
+
+        # ensure sync between epochs
+        if torch.distributed.is_initialized():
+            torch.distributed.barrier()
+
+    # clean up
+    del loss
+    del train_dataloader
+    del val_dataloader
+    del model
+    del optimizer
+    del scheduler
+
+    # free up memory
+    torch.cuda.synchronize()

code/new_realtime/utils/training_utils.py

@@ -0,0 +1,261 @@
+import os
+import pickle
+import random
+from datetime import datetime
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import init
+from sklearn.model_selection import train_test_split
+import lmdb
+from bson import ObjectId
+from tqdm import tqdm
+
+from utils.lmdb_utils import get_lmdb_keys
+from utils.utils import get_config_id
+from utils.data_utils import produce_window_batches
+
+from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
+
+
+def get_training_config(base_config: dict, model_config: dict):
+    try:
+        if base_config["model_class"] != model_config["model_class"]:
+            raise Exception("Model type differ in model config and training config")
+    except KeyError:
+        raise Exception("Model class missing in training or model config")
+
+    # training_config_id = get_config_id(base_config)
+    # hash = training_config_id[-5:]
+
+    timestamp = datetime.now().strftime("%Y%m%d-%H%M")
+    training_config_id = f"{timestamp}"
+
+    training_dir = os.path.abspath(f"{model_config['model_dir']}/trainings/{training_config_id}")
+
+    # append identifier to config
+    training_config = base_config.copy()
+    training_config["id"] = training_config_id
+    training_config["training_dir"] = training_dir
+
+    if not os.path.isdir(training_dir):
+        os.makedirs(training_dir)
+        # save training configuration
+        with open(f"{training_dir}/training_configuration.pickle", "wb") as f:
+            pickle.dump(training_config, f)
+    else:
+        # fetch config
+        with open(f"{training_dir}/training_configuration.pickle", "rb") as f:
+            training_config = pickle.load(f)
+
+    return training_config
+
+
+def get_training_config_from_file(training_dir: str,
+                                  base_model_dir: str,
+                                  model_configuration: dict) -> dict:
+    # fetch config
+    with open(f"{training_dir}/training_configuration.pickle", "rb") as f:
+        training_config = pickle.load(f)
+
+    # update paths based on base directories
+    training_config["training_dir"] = os.path.join(os.path.abspath(base_model_dir),
+                                                   model_configuration["id"],
+                                                   "trainings",
+                                                   training_config["id"])
+
+    return training_config
+
+
+def weight_init(m):
+    """
+    Usage:
+        model = Model()
+        model.apply(weight_init)
+    """
+    if isinstance(m, nn.Conv1d):
+        init.normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.Conv2d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.Conv3d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose1d):
+        init.normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose2d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.ConvTranspose3d):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.BatchNorm1d):
+        init.normal_(m.weight.data, mean=1, std=0.02)
+        init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.BatchNorm2d):
+        init.normal_(m.weight.data, mean=1, std=0.02)
+        init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.BatchNorm3d):
+        init.normal_(m.weight.data, mean=1, std=0.02)
+        init.constant_(m.bias.data, 0)
+    elif isinstance(m, nn.Linear):
+        init.xavier_normal_(m.weight.data)
+        if m.bias is not None:
+            init.normal_(m.bias.data)
+    elif isinstance(m, nn.LSTM):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+    elif isinstance(m, nn.LSTMCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+    elif isinstance(m, nn.GRU):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+        for names in m._all_weights:
+            for name in filter(lambda n: "bias" in n, names):
+                bias = getattr(m, name)
+                n = bias.size(0)
+                bias.data[:n // 3].fill_(-1.)
+    elif isinstance(m, nn.GRUCell):
+        for param in m.parameters():
+            if len(param.shape) >= 2:
+                init.orthogonal_(param.data)
+            else:
+                init.normal_(param.data)
+
+
+def collate(batch_items: list) -> dict:
+    batch = dict()
+    for key in batch_items[0].keys():
+        if key in ["combination_id", "time_index"]:
+            continue
+        else:
+            if batch_items[0][key] is None:
+                batch[key] = None
+            else:
+                batch[key] = np.stack([item[key] for item in batch_items])
+
+    for key in batch.keys():
+        if batch[key] is not None:
+            batch[key] = torch.tensor(batch[key], dtype=torch.float32)
+
+    return batch
+
+
+def get_splits_by_user(input_keys: list, train_size: float, val_size: float, test_size: float):
+    if train_size + val_size + test_size != 1:
+        raise ValueError("Train, val and test sizes must sum to 1")
+
+    if len(input_keys) == 0:
+        raise ValueError("Input keys list is empty")
+
+    items_by_use = dict()
+    for input_key in tqdm(input_keys):
+        try:
+            user_id = get_cycles_collection().find_one({"_id": ObjectId(input_key)})["user_id"]
+            if user_id not in items_by_use:
+                items_by_use[user_id] = []
+            items_by_use[user_id].append(input_key)
+        except Exception:
+            print(f"Error getting user id for key {input_key}")
+            continue
+
+    user_ids = list(items_by_use.keys())
+    random.shuffle(user_ids)
+    train_users, temp_users = train_test_split(user_ids, train_size=train_size, test_size=test_size + val_size)
+    # compute relative test size, as it must be relative to the remaining users
+    relative_test_size = test_size / (1 - train_size)
+    val_users, test_users = train_test_split(temp_users, test_size=relative_test_size)
+
+    train_keys = []
+    val_keys = []
+    test_keys = []
+    for user_id in train_users:
+        train_keys.extend(items_by_use[user_id])
+    for user_id in val_users:
+        val_keys.extend(items_by_use[user_id])
+    for user_id in test_users:
+        test_keys.extend(items_by_use[user_id])
+
+    return train_keys, val_keys, test_keys
+
+
+def save_splits(train_keys: list, val_keys: list, test_keys: list, base_dir: str):
+    if not os.path.exists(base_dir):
+        os.makedirs(base_dir)
+
+    with open(f"{base_dir}/train_keys.pickle", "wb") as f:
+        pickle.dump(train_keys, f)
+    with open(f"{base_dir}/val_keys.pickle", "wb") as f:
+        pickle.dump(val_keys, f)
+    with open(f"{base_dir}/test_keys.pickle", "wb") as f:
+        pickle.dump(test_keys, f)
+
+
+def load_splits(base_dir: str):
+    with open(f"{base_dir}/train_keys.pickle", "rb") as f:
+        train_keys = pickle.load(f)
+    with open(f"{base_dir}/val_keys.pickle", "rb") as f:
+        val_keys = pickle.load(f)
+    with open(f"{base_dir}/test_keys.pickle", "rb") as f:
+        test_keys = pickle.load(f)
+
+    return train_keys, val_keys, test_keys
+
+
+def get_data_ids(model_configuration: dict,
+                 training_configuration: dict,
+                 env_path: str,
+                 limit: int = None) -> tuple:
+    env = lmdb.open(f"{env_path}", readonly=True)
+    if os.path.exists(f"{model_configuration['feature_config']['dataset_dir']}/train_keys.pickle"):
+        train_ids, val_ids, test_ids = load_splits(model_configuration["feature_config"]["dataset_dir"])
+    else:
+        lmdb_keys = get_lmdb_keys(env, limit)
+        # train_ids, val_ids, test_ids = get_splits_by_user(lmdb_keys,
+        #                                                   training_configuration["train_size"],
+        #                                                   training_configuration["val_size"],
+        #                                                   training_configuration["test_size"])
+        train_ids, temp_ids = train_test_split(lmdb_keys,
+                                               train_size=training_configuration["train_size"],
+                                               test_size=training_configuration["val_size"] + training_configuration[
+                                                   "test_size"])
+        # compute relative test size, as it must be relative to the remaining users
+        relative_test_size = training_configuration["test_size"] / (1 - training_configuration["train_size"])
+        val_ids, test_ids = train_test_split(temp_ids,
+                                             test_size=relative_test_size)
+        # save splits to file
+        save_splits(train_ids, val_ids, test_ids, model_configuration["feature_config"]["dataset_dir"])
+
+    if limit is not None:
+        train_size = training_configuration["train_size"]
+        val_size = training_configuration["val_size"]
+        test_size = training_configuration["test_size"]
+
+        train_lim = int(limit * train_size)
+        val_lim = int(limit * val_size)
+        test_lim = int(limit * test_size)
+
+        train_ids = train_ids[:train_lim]
+        val_ids = val_ids[:val_lim]
+        test_ids = test_ids[:test_lim]
+
+    return train_ids, val_ids, test_ids

code/new_realtime/utils/utils.py

@@ -0,0 +1,152 @@
+import sys
+import inspect
+import hashlib
+import logging
+from functools import partial
+import random
+from typing import Callable
+import importlib
+
+import numpy as np
+
+
+def get_logger(module_name: str, filename: str = "main.log") -> logging.Logger:
+    """
+    Returns a logger for the given module name and filename.
+    :param module_name: name of the module, as string
+    :param filename: name of the logging file, as string
+    :return: the logger, as logging.Logger object
+    """
+    logger = logging.getLogger(module_name)
+    logger.setLevel(logging.DEBUG)
+    logger.propagate = False
+    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+
+    file_handler = logging.FileHandler(filename)
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(formatter)
+
+    logger.addHandler(file_handler)
+
+    # add system out handler
+    stream_handler = logging.StreamHandler(sys.stdout)
+    stream_handler.setLevel(logging.INFO)
+    stream_handler.setFormatter(formatter)
+
+    logger.addHandler(stream_handler)
+
+    return logger
+
+
+def get_variable_from_module(module_path: str, variable_name: str):
+    """
+    Get a variable from a module by its name during runtime
+
+    Args:
+        module_path: module to fetch variable from
+        variable_name: variable to fetch from module
+
+    Returns:
+        variable from module
+
+    """
+    module = importlib.import_module(module_path)
+    variable = getattr(module, variable_name)
+    if variable is None:
+        raise ValueError(f"Variable {variable_name} not found in module {module_path}")
+    return variable
+
+
+def get_callable_name(callable_obj):
+    """
+    Get the name of the callable, handling `functools.partial`.
+
+    :param callable_obj: callable object
+    :return: name of the callable
+    """
+    if isinstance(callable_obj, partial):
+        func_name = callable_obj.func.__name__
+        args = ", ".join(object_to_string(arg) for arg in callable_obj.args)
+        kwargs = ", ".join(f"{object_to_string(k)}={object_to_string(v)!r}" for k, v in callable_obj.keywords.items())
+        return f"partial({func_name}, {args}, {kwargs})"
+    else:
+        if hasattr(callable_obj, '__name__'):
+            return callable_obj.__name__
+        elif hasattr(callable_obj, '__class__'):
+            return callable_obj.__class__.__name__
+        elif hasattr(callable_obj, '__hash__'):
+            return callable_obj.__hash__
+        else:
+            raise ValueError(f"Could not determine name of callable object {callable_obj}")
+
+
+def object_to_string(value, skip_types=None):
+    """
+    Convert any object to a string representation that avoids memory addresses.
+    Handles complex data types recursively.
+
+    :param value: object to convert
+    :param skip_types: types to skip during conversion
+    :return: string representation of the object
+    """
+    if skip_types is None:
+        skip_types = []
+
+    if any(isinstance(value, t) for t in skip_types):
+        return 'skipped_type'
+    elif isinstance(value, (str, int, float, bool)):  # Handle primitive data types directly
+        return repr(value)
+    elif isinstance(value, dict):
+        return '{' + ', '.join(f"{k}: {object_to_string(v, skip_types)}" for k, v in value.items()) + '}'
+    elif isinstance(value, (list, tuple)):
+        return '[' + ', '.join(object_to_string(item, skip_types) for item in value) + ']'
+    elif isinstance(value, partial):
+        return get_callable_name(value)
+    elif inspect.isclass(value):
+        return f"<class '{value.__name__}'>"
+    elif hasattr(value,
+                 '__class__') and not value.__class__ != "function":  # Correct handling for instances of classes, but not functions
+        return f"<instance of class '{value.__class__.__name__}'>"
+    elif isinstance(value, Callable):
+        return f"<callable '{get_callable_name(value)}'>"
+    else:
+        return repr(value)
+
+
+def get_config_id(configuration: dict) -> str:
+    """
+    Generate a somewhat unique human-readable model name from the model and training parameters.
+    :param configuration: dictionary containing model and training parameters
+    :return: human-readable model name
+    """
+    adjectives = ["autumn", "hidden", "bitter", "misty", "silent", "empty", "dry", "dark", "summer", "icy", "delicate",
+                  "quiet", "white", "black", "blue", "green", "red", "yellow",
+                  "purple", "orange", "pink", "golden", "silver", "crimson", "violet", "azure", "amber", "sapphire",
+                  "emerald", "ruby", "pearl", "topaz", "onyx", "turquoise", "citrine", ]
+    nouns = ["waterfall", "river", "breeze", "moon", "rain", "wind", "sea", "morning", "snow", "lake", "sunset", "pine",
+             "shadow", "leaf", "dawn", "glitter", "forest", "cloud", "sky", "sun", "butterfly",
+             "flower", "bird", "mountain", "valley", "ocean", "star", "night", "dream", "whisper", "echo", "horizon",
+             "wave", "petal", "dew", "mist"]
+
+    # also add dataset config, but skip functions
+    base_name = object_to_string(configuration)
+
+    # hash long name
+    basename_hash = hashlib.md5(base_name.encode()).hexdigest()
+
+    # select adjective and noun based on hash
+    random.seed(int(basename_hash, 16))
+    model_name = f"{random.choice(adjectives)}_{random.choice(nouns)}_{basename_hash[:5]}"
+
+    return model_name
+
+
+def convert_for_json(obj):
+    if isinstance(obj, dict):
+        return {k: convert_for_json(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_for_json(v) for v in obj]
+    elif isinstance(obj, np.generic):
+        return obj.item()
+    else:
+        return obj

code/new_realtime/utils/visualization.py

@@ -0,0 +1,61 @@
+import numpy as np
+import torch
+from torch import nn
+
+from utils.dataset_creation import inverse_scale_feature
+
+
+def plot_prediction_windows(index: int,
+                            training_configuration: dict,
+                            item_x: np.ndarray,
+                            item_y: np.ndarray,
+                            preds: np.ndarray,
+                            scalers: dict,
+                            features_to_plot: list,
+                            output_feature_names: list,
+                            fig_widget):
+    window_features = item_x[index]
+    indices = np.arange(window_features.shape[0])
+    actual = item_y[index].ravel()
+    predicted = preds[index]
+
+    # clear previous traces
+    fig_widget.data = []
+
+    for feature in features_to_plot:
+        feature_index = feature["index"]
+        feature_name = feature["name"]
+        fig_widget.add_scatter(
+            x=indices,
+            y=window_features[:, feature_index],
+            mode="lines",
+            name=feature_name,
+        )
+
+    num_outputs = predicted.shape[-1]
+    scaled_preds = list()
+    for i in range(num_outputs):
+        if isinstance(training_configuration["loss_functions"][i], nn.BCEWithLogitsLoss):
+            output = torch.sigmoid(torch.tensor(predicted[i]))
+        else:
+            output = predicted[i]
+
+        output = float(output)
+
+        scaled_output = inverse_scale_feature(output,
+                                              output_feature_names[i],
+                                              scalers)
+
+        scaled_preds.append(scaled_output)
+
+    scaled_actuals = list()
+    for i in range(num_outputs):
+        output = float(actual[i].numpy())
+        scaled_output = inverse_scale_feature(output,
+                                              output_feature_names[i],
+                                              scalers)
+
+        scaled_actuals.append(scaled_output)
+
+    # add actual and predicted values
+    print(f"actual: {scaled_actuals}, predicted: {scaled_preds}")