fixes

code/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/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/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/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/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/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/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/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/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/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/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/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/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}")