added code

code/utils/cycle_utils.py

@@ -0,0 +1,21 @@
+def get_ovulation_day(cycle: dict) -> int:
+    """
+    Get the ovulation day of a cycle.
+    :param cycle: cycle dict
+    :return: ovulation day
+    """
+    try:
+        is_biphasic = "classification_results" in cycle and cycle["classification_results"][0]["results"][
+            "predicted_class"] == "biphasic"
+        if not is_biphasic:
+            return None
+
+        has_ov_detection = "ov_detection_results" in cycle
+        if not has_ov_detection:
+            return None
+
+        ov_detection_results = cycle["ov_detection_results"][0]["results"]["ovulation_day"]
+        return ov_detection_results
+    except Exception as e:
+        # print(e)
+        return None

code/utils/data_utils.py

@@ -2,11 +2,12 @@ import copy
 import logging
 import os
 import random
+from datetime import datetime
 
 import lmdb
 import numpy as np
 from bson import ObjectId
-from torch.utils.data import IterableDataset
+from torch.utils.data import IterableDataset, get_worker_info
 from tqdm import tqdm
 
 from utils.dataset_creation import get_features, load_scalers, scale_item, combine_features
@@ -52,20 +53,24 @@ def get_number_of_windows(base_length: int, model_configuration: dict) -> int:
     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
+        return (base_length - input_window_length) // max(window_shift, 1) + 1
     else:
-        return (base_length - output_window_length - output_window_offset) // window_shift + 1
+        return (base_length - output_window_length - output_window_offset) // max(window_shift, 1) + 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:
+                               lmdb_env=None,
+                               item_subset_source: str = "train",
+                               for_inference: bool = False) -> 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)
+                                             lmdb_env=lmdb_env,
+                                             item_subset_source=item_subset_source,
+                                             for_inference=for_inference)
     collated = model_configuration["collate_fn"](sample_batch_for_key)
     return collated
 
@@ -74,7 +79,9 @@ def get_batch_for_key(key,
                       model_configuration: dict,
                       start_cutoff: int = None,
                       end_cutoff: int = None,
-                      lmdb_env=None) -> np.ndarray:
+                      lmdb_env=None,
+                      item_subset_source: str = "train",
+                      for_inference: bool = False) -> np.ndarray:
     """
     Get the batch for a given key from the lmdb database or compute it directly.
 
@@ -85,6 +92,8 @@ def get_batch_for_key(key,
         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
+        item_subset_source: where the item "came" from, test, train or val, used to determine the scaler to use.
+        for_inference: whether to compute batched features directly from the database and ignore cycle filters
 
     Returns:
         batch: batch for the given key, as returned by the batch_fn in the model configuration
@@ -92,7 +101,10 @@ def get_batch_for_key(key,
     """
     if lmdb_env is None:
         # compute and scale features, ignored features are handled internally
-        features = get_scaled_feature_for_key(key, model_configuration)
+        features = get_scaled_feature_for_key(key,
+                                              model_configuration,
+                                              item_subset_source=item_subset_source,
+                                              for_inference=for_inference)  # TODO: implement proper cutoff usage, otherwise the whole user history will be loaded every time
     else:
         # load features from lmdb
         features = load_from_lmdb(lmdb_env, str(key))
@@ -141,12 +153,17 @@ def get_batch_for_key(key,
 
 
 def get_scaled_feature_for_key(key: str,
-                               model_configuration: dict) -> tuple:
+                               model_configuration: dict,
+                               item_subset_source: str = "train",
+                               for_inference: bool = False) -> 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))
+    if for_inference:
+        user_cycles = list(
+            get_cycles_collection().find({"user_id": ObjectId(key)}).sort("starts_at", 1))
+    else:
+        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)
@@ -159,7 +176,7 @@ def get_scaled_feature_for_key(key: str,
     scaler_dir = os.path.join(feature_config["dataset_dir"], "scalers")
     scalers = load_scalers(scaler_dir)
     # scale features
-    scaled_features = scale_item(features, scalers)
+    scaled_features = scale_item(features, subset_name=item_subset_source, scalers=scalers)
     return scaled_features
 
 
@@ -476,7 +493,7 @@ class LMDBIterableDataset(IterableDataset):
         Args:
             key_subset: list of keys to use
         """
-        self.key_subset = key_subset
+        self.key_subset = self.get_worker_keyset(key_subset)
         # reset length
         self.len = None
         # reset lmdb env
@@ -484,6 +501,26 @@ class LMDBIterableDataset(IterableDataset):
             self.lmdb_env.close()
             self.lmdb_env = None
 
+    def get_worker_keyset(self, key_set: list[str]):
+        """
+        Adjust the key set to the current worker.
+        Args:
+            key_set: list of keys to use
+        """
+        worker_info = get_worker_info()
+        if worker_info is not None:
+            # in a worker process
+            num_workers = worker_info.num_workers
+            worker_id = worker_info.id
+            # split the key set into chunks for each worker
+            chunk_size = len(key_set) // num_workers
+            start = worker_id * chunk_size
+            end = (worker_id + 1) * chunk_size if worker_id != num_workers - 1 else len(key_set)
+            return key_set[start:end]
+        else:
+            # use the whole key set
+            return key_set
+
     def get_length_of_data_subset(self, key_set: list[str]):
         """
         Get the length of the data subset.
@@ -495,6 +532,7 @@ class LMDBIterableDataset(IterableDataset):
         num_steps = 0
         i = 0
         num_batch_items = 0
+        padding_length = get_padding_length(self.model_configuration, resampled=False)
         while True:
             if num_batch_items >= self.batch_size:
                 num_steps += 1
@@ -512,9 +550,8 @@ class LMDBIterableDataset(IterableDataset):
                     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)
+                    item_length = len(range(0, raw_length, take_every_nth))
                 except:
                     item = load_from_lmdb(self.lmdb_env, key)
                     item_length = get_prepared_sequence_length(
@@ -529,7 +566,7 @@ class LMDBIterableDataset(IterableDataset):
 
         # if no key subset is set, use all keys
         if self.key_subset is None:
-            self.key_subset = self.lmdb_keys
+            self.set_key_subset(self.lmdb_keys)
 
         self.init_lmdb_env()
 
@@ -548,14 +585,18 @@ class LMDBIterableDataset(IterableDataset):
             else:
                 if counter >= len(self.key_subset):
                     if len(batch) > 0:
-                        yield collate_fn(batch)
+                        if collate_fn is None:
+                            yield batch
+                        else:
+                            yield collate_fn(batch)
                     break
-                key = self.lmdb_keys[counter]
+                key = self.key_subset[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
+                except ValueError as e:
+                    # if batch is empty, try next
+
                     continue
                 batch.extend(current_batch)
 
@@ -570,7 +611,7 @@ class LMDBIterableDataset(IterableDataset):
     def __len__(self):
         # if no key subset is set, use all keys
         if self.key_subset is None:
-            self.key_subset = self.lmdb_keys
+            self.set_key_subset(self.lmdb_keys)
 
         self.init_lmdb_env()
 

code/utils/dataset_creation.py

@@ -4,6 +4,7 @@ import lmdb
 import numpy as np
 import pandas as pd
 import torch
+from sklearn.base import BaseEstimator
 from sklearn.preprocessing import StandardScaler, MinMaxScaler
 import joblib
 
@@ -115,11 +116,11 @@ def get_feature_values(feature_type: str | None,
     return feature_values
 
 
-def train_scalers(feature_type: str,
-                  feature_name: str,
-                  scaler_type,
-                  sample,
-                  env) -> dict:
+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()
@@ -147,7 +148,33 @@ def train_scalers(feature_type: str,
     return scalers
 
 
+def train_scalers(feature_type: str,
+                  feature_name: str,
+                  feature_config: dict,
+                  env) -> BaseEstimator | None:
+    feature_values = get_feature_values(feature_type, feature_name, env)
+    if len(feature_values) == 0:
+        raise ValueError(f"No feature values found for feature {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
+
+    current_feature_config = feature_config[feature_type][feature_name]
+    scaler_type = current_feature_config.get("scaler", None)
+
+    if scaler_type is not None:
+        scaler = scaler_type()
+        scaler.fit(features_reshaped)
+    else:
+        scaler = None
+
+    return scaler
+
+
 def scale_item(data: dict,
+               subset_name: str,
                scalers: dict):
     data_format = dict()
     for feature_set in data:
@@ -161,12 +188,13 @@ def scale_item(data: dict,
     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:
+            subset_feature_name = f"{subset_name}_{feature_name}"
+            if subset_feature_name in scalers and scalers[subset_feature_name] is not None:
                 if isinstance(feature_values, list) or isinstance(feature_values, np.ndarray):
-                    scaled_feature_values = scalers[feature_name].transform(
+                    scaled_feature_values = scalers[subset_feature_name].transform(
                         feature_values.reshape(-1, 1)).flatten()
                 else:
-                    scaled_feature_values = scalers[feature_name].transform(
+                    scaled_feature_values = scalers[subset_feature_name].transform(
                         np.array(feature_values).reshape(-1, 1)).flatten()
             else:
                 scaled_feature_values = feature_values
@@ -177,13 +205,15 @@ def scale_item(data: dict,
 
 def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor,
                           feature_names: np.ndarray | list | str,
-                          scalers: dict) -> np.ndarray:
+                          scalers: dict,
+                          subset_name: str = "train") -> 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
+        subset_name: subset name, either train, val or test for proper inverse scaling
 
     Returns:
         scaled input features as numpy array
@@ -196,7 +226,7 @@ def inverse_scale_feature(input_feature: np.ndarray | int | float | 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]
+        input_scaled = scalers[f"{subset_name}_{feature_names[0]}"].inverse_transform([[input_feature]])[0][0]
         return input_scaled
     elif isinstance(input_feature, list):
         input_feature = np.array(input_feature)
@@ -210,7 +240,8 @@ def inverse_scale_feature(input_feature: np.ndarray | int | float | torch.Tensor
 
     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]]))
+        input_scaled[:, i] = scalers[f"{subset_name}_{feature_names[i]}"].inverse_transform(
+            np.array([input_scaled[:, i]]))
 
     return input_scaled
 

code/utils/dataset_utils.py

@@ -26,6 +26,7 @@ def load_key_stats(lmdb_dir: str) -> dict:
     :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):

code/utils/evaluation.py

@@ -1,3 +1,8 @@
+import copy
+import logging
+import math
+from typing import Callable
+
 import lmdb
 import numpy as np
 import sklearn
@@ -5,9 +10,12 @@ import torch
 from torch import nn
 from tqdm import tqdm
 
+from utils.inference import scale_features, apply_sigmoid_if_necessary
+from utils.utils import get_logger
 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 utils.training_utils import load_model_for_usage
 
 from vsm_datascience_common import constants
 
@@ -15,7 +23,20 @@ from vsm_datascience_common import constants
 def evaluate_model(model_configuration: dict,
                    training_configuration: dict,
                    test_ids: list,
-                   evaluation_functions: list) -> dict:
+                   evaluation_functions: list,
+                   log_dir: str,
+                   aggregate: bool = True,
+                   logger: logging.Logger = None,
+                   print_progress: bool = False) -> dict:
+    if logger is None:
+        logger = get_logger(__name__, f"{log_dir}/{model_configuration['id']}_{training_configuration['id']}.log")
+
+    model, training_configuration = load_model_for_usage(
+        model_configuration,
+        training_configuration,
+        test_ids[0]
+    )
+
     dataset_dir = model_configuration["feature_config"]["dataset_dir"]
     lmdb_env = lmdb.open(dataset_dir, readonly=True)
     # get computation rank
@@ -24,18 +45,18 @@ def evaluate_model(model_configuration: dict,
         torch.cuda.set_device(local_rank)
         device = torch.device(f"cuda:{local_rank}")
     else:
+        local_rank = 0
         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)
+    model.to(device)
+
+    batch_size = training_configuration["batch_size"]
+    logger.info(
+        f"Rank {local_rank}: Estimated batch size: {batch_size} for GPU {torch.cuda.get_device_name(local_rank)}")
+
     # 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"]
 
@@ -45,10 +66,13 @@ def evaluate_model(model_configuration: dict,
     scalers = get_scalers_for_model(model_configuration)
 
     errors = dict()
-    for test_id in tqdm(test_ids):
+
+    for test_id in tqdm(test_ids, disable=not print_progress):
         # fetch stats for key
         current_key_stats = keys_stats["by_key"][test_id]
         cycle_stats = current_key_stats["cycle_stats"]
+        # make sure cycle stats are sorted by start time
+        # cycle_stats = sorted(cycle_stats, key=lambda x: x["starts_at"])
         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]])
@@ -70,62 +94,54 @@ def evaluate_model(model_configuration: dict,
                                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()
+            processed_preds = apply_sigmoid_if_necessary(preds, model_configuration)
+            scaled_preds = scale_features(
+                processed_preds,
+                used_targets,
+                model_configuration,
+                scalers,
+            )
 
-                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)
+            scaled_actuals = scale_features(
+                actuals,
+                used_targets,
+                model_configuration,
+                scalers,
+            )
 
             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()
+                    eval_function = eval_fn["eval_fn"]
+                    eval_fn_indices = [eval_fn["input_index"]] if "input_index" in eval_fn else range(num_outputs)
+                    # skip error fn if actuals are nan, since they are ignored
+                    for output_index in eval_fn_indices:
+                        if any(np.isnan(scaled_actuals[:, output_index])):
+                            continue
+                        error = eval_function(scaled_preds[:, output_index],
+                                              scaled_actuals[:, output_index])
 
-                    if f"after_{i}" not in errors[eval_fn_name]:
-                        errors[eval_fn_name][f"after_{i}"] = list()
+                        if np.isnan(error):
+                            # skip if error is nan
+                            continue
 
-                    errors[eval_fn_name][f"after_{i}"].append(error)
+                        if f"after_{i}" not in errors[eval_fn_name]:
+                            errors[eval_fn_name][f"after_{i}"] = dict()
 
+                        if output_index not in errors[eval_fn_name][f"after_{i}"]:
+                            errors[eval_fn_name][f"after_{i}"][output_index] = list()
+
+                        errors[eval_fn_name][f"after_{i}"][output_index].append(error)
+
+    if aggregate:
+        errors = aggregate_errors(errors, evaluation_functions)
+    return errors
+
+
+def aggregate_errors(errors: dict, evaluation_functions: list) -> dict:
+    errors = copy.deepcopy(errors)
     for eval_fn in evaluation_functions:
         if eval_fn is not None:
             eval_fn_name = eval_fn["name"]
@@ -133,13 +149,32 @@ def evaluate_model(model_configuration: dict,
             if eval_fn_name not in errors:
                 continue
             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])
+                for output_index in errors[eval_fn_name][key]:
+                    if len(errors[eval_fn_name][key][output_index]) == 0:
+                        errors[eval_fn_name][key][output_index] = np.nan
+                    else:
+                        # only use non nan values
+                        without_nans = [val for val in errors[eval_fn_name][key][output_index] if
+                                        val is not None and val is not np.nan]
+                        errors[eval_fn_name][key][output_index] = accumulation_fn(without_nans)
+
     return errors
 
 
+def get_index_offset(offset_in_hours: int,
+                     model_configuration: dict, ) -> float:
+    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
+
+    if offset_in_hours == 0:
+        return 0
+
+    index_offset_in_hours = math.ceil(offset_in_hours / shift_hour_factor)
+    return index_offset_in_hours
+
+
 def pre_ov_error(preds, actuals, *args, **kwargs):
     if any(np.isnan(actuals)):
         return np.nan
@@ -207,8 +242,8 @@ def ov_error(preds, actuals, model_configuration, *args, **kwargs):
     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
+    shift_hour_factor = 24 / (measurements_per_day // downsampling_factor) * step_size
+    error_in_days = error * shift_hour_factor / 24
 
     return error_in_days
 
@@ -246,12 +281,7 @@ def day_relative_to_ov_error(preds: np.ndarray | torch.Tensor,
         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))
+    index_offset = get_index_offset(day_relative_to_ov * 24, model_configuration)
 
     if ov_day_index + index_offset >= preds.shape[0]:
         relative_day_pred = preds[-1] if len(preds.shape) > 1 else preds[-1]
@@ -266,3 +296,237 @@ def day_relative_to_ov_error(preds: np.ndarray | torch.Tensor,
     # calculate error
     error = np.abs(relative_day_pred - relative_day_actual)
     return error
+
+
+def get_ov_over_ov_index(input_sequence: np.ndarray) -> int:
+    """
+    Returns the index of the ovulation based on the ov_over feature curve.
+
+    Ov index is first non-zero index in actual curve
+    Args:
+        input_sequence: input sequence
+    Returns:
+        index of the ovulation based on the ov_over feature curve
+    """
+
+    try:
+        ov_index = int(np.where(input_sequence > 0)[0][0])
+    except IndexError:
+        return None
+
+    return ov_index
+
+
+def get_ov_over_pre_ov_error(preds: np.ndarray | torch.Tensor,
+                             actual: np.ndarray | torch.Tensor,
+                             error_fn: Callable) -> float:
+    ov_index = get_ov_over_ov_index(actual)
+
+    if ov_index is None:
+        return np.nan
+
+    if any(np.isnan(actual)):
+        return np.nan
+
+    pre_ov_preds = preds[:ov_index]
+    pre_ov_actuals = actual[:ov_index]
+
+    if len(pre_ov_preds) == 0 or len(pre_ov_actuals) == 0:
+        return np.nan
+
+    error = error_fn(pre_ov_preds, pre_ov_actuals)
+    return error
+
+
+def get_ov_over_post_ov_error(preds: np.ndarray | torch.Tensor,
+                              actual: np.ndarray | torch.Tensor,
+                              error_fn: Callable) -> float:
+    ov_index = get_ov_over_ov_index(actual)
+    if any(np.isnan(actual)):
+        return np.nan
+    post_ov_preds = preds[ov_index:]
+    post_ov_actuals = actual[ov_index:]
+    if len(post_ov_preds) == 0 or len(post_ov_actuals) == 0:
+        return np.nan
+    error = error_fn(post_ov_preds, post_ov_actuals)
+    return error
+
+
+def get_pre_fertility_error(preds: np.ndarray | torch.Tensor,
+                            actual: np.ndarray | torch.Tensor,
+                            error_fn: Callable) -> float:
+    try:
+        fertility_start_index = np.where(actual > 0)[0][0]
+    except IndexError:
+        return np.nan
+
+    pre_fertility_preds = preds[:fertility_start_index]
+    pre_fertility_actuals = actual[:fertility_start_index]
+
+    if len(pre_fertility_preds) == 0 or len(pre_fertility_actuals) == 0:
+        return np.nan
+
+    error = error_fn(pre_fertility_preds, pre_fertility_actuals)
+    return error
+
+
+def get_post_fertility_error(preds: np.ndarray | torch.Tensor,
+                             actual: np.ndarray | torch.Tensor,
+                             error_fn: Callable) -> float:
+    try:
+        fertility_end_index = np.where(actual > 0)[-1][-1]
+    except IndexError:
+        return np.nan
+
+    pre_fertility_preds = preds[fertility_end_index:]
+    pre_fertility_actuals = actual[fertility_end_index:]
+
+    if len(pre_fertility_preds) == 0 or len(pre_fertility_actuals) == 0:
+        return np.nan
+
+    error = error_fn(pre_fertility_preds, pre_fertility_actuals)
+    return error
+
+
+def get_during_fertility_error(preds: np.ndarray | torch.Tensor,
+                               actual: np.ndarray | torch.Tensor,
+                               error_fn: Callable) -> float:
+    try:
+        fertility_indices = np.where(actual > 0)[0]
+    except IndexError:
+        return np.nan
+
+    during_fertility_preds = preds[fertility_indices]
+    during_fertility_actuals = actual[fertility_indices]
+
+    if len(during_fertility_preds) == 0 or len(during_fertility_actuals) == 0:
+        return np.nan
+
+    error = error_fn(during_fertility_preds, during_fertility_actuals)
+    return error
+
+
+def get_non_fertility_error(preds: np.ndarray | torch.Tensor,
+                            actual: np.ndarray | torch.Tensor,
+                            error_fn: Callable) -> float:
+    """
+    Calculate the error of the model predictions for non-fertility periods.
+    Args:
+        preds: predictions
+        actual: actual values
+    Returns:
+        error: error of the model predictions for non-fertility periods
+    """
+
+    if any(np.isnan(actual)):
+        return np.nan
+
+    try:
+        non_fertility_indices = np.where(actual == 0)[0]
+    except IndexError:
+        return np.nan
+
+    if len(non_fertility_indices) == 0:
+        return np.nan
+
+    # get non-fertility predictions
+    non_fertility_preds = preds[non_fertility_indices]
+    non_fertility_actuals = actual[non_fertility_indices]
+    if len(non_fertility_preds) == 0 or len(non_fertility_actuals) == 0:
+        return np.nan
+
+    # calculate error
+    error = error_fn(non_fertility_preds, non_fertility_actuals)
+    return error
+
+
+def get_fertility_over_index(input_sequence: np.ndarray) -> int:
+    """
+    Returns the index of the fertility based on the fertility feature curve.
+
+    Fertility index is first non-zero index in actual curve
+    Args:
+        input_sequence: input sequence
+    Returns:
+        index of the fertility based on the fertility feature curve
+    """
+
+    try:
+        fertility_index = np.where(input_sequence > 0)[-1][-1]
+    except IndexError:
+        return None
+
+    return fertility_index
+
+
+def get_error_relative_to_x(preds: np.ndarray | torch.Tensor,
+                            actual: np.ndarray | torch.Tensor,
+                            x_fn: Callable,
+                            offset_in_hours: int,
+                            model_configuration: dict) -> float:
+    """
+    Calculate the error at the index relative
+    Args:
+        preds: predictions
+        actual: actual values
+        x_fn: function to get the index of the x value
+        offset_in_hours: offset in hours to apply to the x index
+        model_configuration: model configuration
+
+    Returns:
+        error: error of the model predictions relative to the x value
+    """
+
+    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 x index
+    try:
+        x_index = x_fn(actual)
+    except Exception:
+        return np.nan
+
+    if x_index is None:
+        return np.nan
+
+    # get index offset factor -> how much time between each step
+    index_offset = get_index_offset(offset_in_hours, model_configuration)
+
+    if x_index + index_offset >= preds.shape[0]:
+        offset_pred = preds[-1]
+        offset_actual = actual[-1]
+    elif x_index + index_offset < 0:
+        offset_pred = preds[0]
+        offset_actual = actual[0]
+    else:
+        offset_pred = preds[x_index + index_offset]
+        offset_actual = actual[x_index + index_offset]
+
+    # calculate error
+    error = np.abs(offset_pred - offset_actual)
+    return error
+
+
+def r2_wrapper(preds: np.ndarray | torch.Tensor,
+               actual: np.ndarray | torch.Tensor,
+               var_tol: float = 1e-8,
+               max_allowed_var: float = 1.0):
+    if len(preds) <= 1 or len(actual) <= 1:
+        return np.nan
+
+    # make sure to capture the case where actuals have no variance
+    if np.var(actual) < var_tol:
+        # return variance of predictions here, as it should be at best zero
+        # cap at 1.0 to make it comparable
+        error = min(np.var(preds), max_allowed_var)
+        return error
+
+    # return 1 - r2 to get an error based metric, where lower is better
+    error = 1 - sklearn.metrics.r2_score(actual, preds)
+    return error

code/utils/feature_functions.py

@@ -3,7 +3,7 @@ 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.cycle_database_connection.db_utils import get_cycles_collection, get_collection
 from vsm_datascience_common.cycles.sequences import get_timestamps, get_values
 
 from utils.smoothing import get_curve_composition
@@ -243,6 +243,67 @@ def get_window_fn(cycle: dict,
     return fn(windows)
 
 
+def get_user_age(cycle: dict) -> int:
+    """
+    Returns the age of the user in years the cycle belongs to as array. Returns average age, if user has no age specified
+    Args:
+        cycle: cycle belonging to the user
+
+    Returns:
+        age_array: array with the age of th user, same length as cycle
+
+    """
+
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    pii = get_collection("personal_informations", constants.METADATA_DB_NAME).find_one({"user_id": user_id})
+    if pii is None:
+        return np.full((cycle_length,), 37)  # 37 is mean of all users
+    else:
+        return np.full((cycle_length,), pii["age"])
+
+
+def get_user_weight(cycle: dict) -> np.array:
+    """
+    Returns the weight of the user in KG the cycle belongs to as array. Returns average weight of user if user has no weight specified
+    Args:
+        cycle: cycle belonging to the user
+
+    Returns:
+        weight_array: weight of the user for each datapoint in the cycle
+
+    """
+
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    pii = get_collection("personal_informations", constants.METADATA_DB_NAME).find_one({"user_id": user_id})
+    if pii is None:
+        return np.full((cycle_length,), 70)  # 70 is mean for all users
+    else:
+        return np.full((cycle_length,), pii["weight"])
+
+
+def get_user_height(cycle: dict) -> np.array:
+    """
+    Returns the height of the user in cm the cycle belongs to. Returns average height of user if user has no height specified
+    Args:
+        cycle: cycle belonging to the user
+
+    Returns:
+        height_array: height of the user for each datapoint in the cycle
+
+    """
+
+    user_id = cycle["user_id"]
+    cycle_length = len(get_values(cycle))
+    pii = get_collection("personal_informations", constants.METADATA_DB_NAME).find_one({"user_id": user_id})
+
+    if pii is None:
+        return np.full((cycle_length,), 167)  # 167 is mean for all users
+    else:
+        return np.full((cycle_length,), pii["height"])
+
+
 def get_cycle_length_stats(cycle: dict) -> dict:
     user_id = cycle["user_id"]
     cycle_length = len(get_values(cycle))

code/utils/inference.py

@@ -0,0 +1,213 @@
+import os
+
+import lmdb
+import numpy as np
+import torch
+from torch import nn
+
+from utils.dataset_creation import get_scalers_for_model, inverse_scale_feature
+from utils.data_utils import get_collated_batch_for_key
+from utils.dataset_utils import get_dataset_path
+from utils.model_utils import get_model_config_from_file
+from utils.training_utils import get_training_config_from_file, get_data_ids, find_max_batch_size
+
+
+def load_model(model_dir: str,
+               results_base_dir: str,
+               lmdb_base_dir: str,
+               device: str = "cuda") -> tuple:
+    """
+    Load the model from the model directory.
+    :param model_dir: model directory
+    :param results_base_dir: results base directory
+    :param lmdb_base_dir: LMDB base directory
+    :param device: device to use
+    :return: model configuration, training configuration, dataset directory
+    """
+
+    training_id = None
+    model_configuration = get_model_config_from_file(model_dir, results_base_dir, lmdb_base_dir)
+    feature_config = model_configuration["feature_config"]
+    dataset_dir = get_dataset_path(lmdb_base_dir, feature_config["feature_set_name"])
+
+    # get list of trainings
+    training_base_dir = os.path.join(model_dir, "trainings")
+    # if no training id is given, use the latest training
+    if training_id is None:
+        training_dirs = os.listdir(training_base_dir)
+        training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
+                         os.path.isdir(os.path.join(training_base_dir, dir))]
+        training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
+        training_dir = training_dirs[0]
+    else:
+        training_dir = os.path.join(training_base_dir, training_id)
+
+    training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
+
+    limit = None
+    train_ids, val_ids, test_ids = get_data_ids(model_configuration, training_configuration, dataset_dir, limit)
+
+    dataset_dir = f"{lmdb_base_dir}/{feature_config['feature_set_name']}"
+    env = lmdb.open(dataset_dir, readonly=True)
+    load_fn = model_configuration["model_load_fn"]
+    model = load_fn(model_configuration, training_configuration, test_ids[0], device, lmdb_env=env)
+
+    # estimate batch size
+    sample_batch = get_collated_batch_for_key(train_ids[0], model_configuration, lmdb_env=env)
+    sample_input = sample_batch[0][0]  # get the first item in the batch
+    input_shape = sample_input.shape
+    # send model to proper device
+    model.to(device)
+    batch_size = find_max_batch_size(model,
+                                     input_shape,
+                                     device=device)
+    training_configuration["batch_size"] = batch_size
+    return model, model_configuration, training_configuration, train_ids, val_ids, test_ids
+
+
+def get_training(model_configuration: dict,
+                 results_base_dir: str,
+                 training_id: str = None) -> dict:
+    """
+    Get the training configuration from the model directory.
+    :param model_configuration: model configuration
+    :param training_id: training name
+    :return: training configuration
+    """
+    model_dir = model_configuration["model_dir"]
+    # get list of trainings
+    training_base_dir = os.path.join(model_dir, "trainings")
+    # if no training id is given, use the latest training
+    if training_id is None:
+        training_dirs = os.listdir(training_base_dir)
+        training_dirs = [os.path.join(training_base_dir, dir) for dir in training_dirs if
+                         os.path.isdir(os.path.join(training_base_dir, dir))]
+        training_dirs = sorted(training_dirs, key=os.path.getmtime, reverse=True)
+        training_dir = training_dirs[0]
+    else:
+        training_dir = os.path.join(training_base_dir, training_id)
+
+    training_configuration = get_training_config_from_file(training_dir, results_base_dir, model_configuration)
+    return training_configuration
+
+
+def find_models(base_dir: str, current_dir: str = None) -> dict:
+    """
+    Find all models in the base directory.
+    :param base_dir: base directory
+    :return: list of model directories
+    """
+    model_dirs = dict()
+    is_root = base_dir == current_dir
+    if current_dir is None:
+        current_dir = base_dir
+    for sub_dir_name in os.listdir(current_dir):
+        sub_dir = os.path.join(current_dir, sub_dir_name)
+        if not os.path.isdir(sub_dir):
+            continue
+        model_config_path = os.path.join(sub_dir, "model_configuration.pickle")
+
+        if current_dir not in model_dirs:
+            model_dirs[current_dir] = list()
+        if os.path.exists(model_config_path) and is_root:
+            model_dirs[current_dir].append(sub_dir)
+        elif os.path.exists(model_config_path) and not is_root:
+            model_dirs[current_dir].append(sub_dir)
+        else:
+            model_dirs = model_dirs | find_models(base_dir, sub_dir)
+
+    return model_dirs
+
+
+def apply_sigmoid_if_necessary(model_outputs: np.ndarray | torch.Tensor,
+                               model_configuration: dict) -> np.ndarray:
+    """
+    Applies a sigmoid to the outputs of a model, if they use a BCE loss.
+    Args:
+        model_outputs: models outputs
+        model_configuration: configuration of the model that produced the outputs
+
+    Returns:
+        processed_model_outputs: processed model outputs with sigmoid applied where necessary
+    """
+
+    target_features = model_configuration["feature_config"]["target_features"]
+    ignored_features = model_configuration["feature_config"]["ignored_features"]
+    used_targets = [feature for feature in target_features if feature["name"] not in ignored_features]
+
+    is_3d = len(model_outputs.shape) == 3
+
+    processed_outputs = np.empty_like(model_outputs)
+    for i, used_target in enumerate(used_targets):
+        if used_target["loss_fn"] == nn.BCEWithLogitsLoss:
+            if is_3d:
+                target_values = model_outputs[:, 0, i]
+                if isinstance(target_values, np.ndarray):
+                    target_values = torch.from_numpy(target_values)
+                processed_outputs[:, 0, i] = torch.sigmoid(target_values).numpy()
+            else:
+                target_values = model_outputs[:, i]
+                if isinstance(target_values, np.ndarray):
+                    target_values = torch.from_numpy(target_values)
+                processed_outputs[:, i] = torch.sigmoid(target_values).numpy()
+        else:
+            if is_3d:
+                processed_outputs[:, 0, i] = model_outputs[:, 0, i].numpy() if isinstance(model_outputs,
+                                                                                          torch.Tensor) else model_outputs[
+                                                                                                             :, 0, i]
+            else:
+                processed_outputs[:, i] = model_outputs[:, i].numpy() if isinstance(model_outputs,
+                                                                                    torch.Tensor) else model_outputs[:,
+                                                                                                       i]
+
+    return processed_outputs
+
+
+def scale_features(features: np.ndarray | torch.Tensor,
+                   feature_names: list | None,
+                   model_configuration: dict,
+                   scalers: dict = None,
+                   subset_name: str = "train") -> np.ndarray:
+    """
+    Scales the given features based on the scalers associated with the given model configuration
+    Args:
+        features: feature arrays, may be up to 3D (but may only have 2 usable dimensions, (B, 1, F) is allowed)
+        feature_names: names of the features in the same order as in the input array, may be None, then no scaling is performed
+        model_configuration: model configuration of the model the features belong to
+        scalers: dict of scales associated by name, will be fetched from model configuration, if None
+        subset_name: subset name, either "train", "test" or "val", used for determining the correct scalers
+
+    Returns:
+        features: scaled features
+    """
+
+    if features.shape[-1] != len(feature_names):
+        raise ValueError("Number of features does not match number of given feature names")
+
+    # remove singleton dimension, if 3D
+    if len(features.shape) == 3:
+        if features.shape[1] != 1:
+            raise ValueError(f"Input has 3 dimensions, expected (B, 1, F), but got {features.shape}")
+        features = features[:, 0, :]
+
+    if scalers is None:
+        scalers = get_scalers_for_model(model_configuration)
+
+    if isinstance(features, torch.Tensor):
+        features = features.detach().cpu().numpy()
+
+    scaled_features = np.empty_like(features)
+    for i in range(len(feature_names)):
+        # skip scaling, if feature name is None
+        if feature_names[i] is not None:
+            scaled_output = inverse_scale_feature(features[:, i],
+                                                  feature_names[i],
+                                                  scalers,
+                                                  subset_name=subset_name)
+        else:
+            scaled_output = features[:, i]
+
+        # make sure to make into 1D array when assigning
+        scaled_features[:, i] = scaled_output.ravel()
+
+    return scaled_features

code/utils/infrastructure_utils.py

@@ -0,0 +1,20 @@
+import time
+
+import torch
+
+
+def custom_barrier_with_timeout(timeout_sec=60 * 60 * 2, check_interval=60 * 5):
+    """
+    Repeatedly attempts to synchronize processes using torch.distributed.barrier().
+    Retries until timeout_sec is exceeded.
+    """
+    start_time = time.time()
+    while True:
+        try:
+            torch.distributed.barrier()
+            break  # Success
+        except Exception as e:
+            elapsed = time.time() - start_time
+            if elapsed > timeout_sec:
+                raise TimeoutError(f"Barrier timed out after {timeout_sec} seconds") from e
+            time.sleep(check_interval)  # Wait before retrying

code/utils/model_utils.py

@@ -9,9 +9,9 @@ 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
+    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

code/utils/training.py

@@ -1,11 +1,13 @@
 import math
+import inspect
 
 import torch
-from torch import nn
+from torch import nn, Gradient
 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 torch.amp import autocast, GradScaler
 from tqdm import tqdm
 
 from utils.data_utils import LMDBIterableDataset
@@ -193,11 +195,13 @@ def train_model(model: nn.Module,
             train_dataset,
             batch_size=None,
             num_workers=num_dataloader_workers,
+            # multiprocessing_context="forkserver",
         )
         val_dataloader = DataLoader(
             val_dataset,
             batch_size=None,
             num_workers=num_dataloader_workers,
+            # multiprocessing_context="forkserver",
         )
 
     logger.info(f"Rank {local_rank}: Training {training_id} with {num_epochs} epochs")
@@ -219,10 +223,18 @@ def train_model(model: nn.Module,
                            max_lr=learning_parameters["learning_rate"],
                            # make sure to use length of full dataset here
                            total_steps=total_train_steps)
+    # mixed precision grad scaler to prevent underflow
+    scaler = GradScaler()
     current_epoch = 1
 
-    loss_functions = training_configuration["loss_functions"]
-    # loss_fn = nn.MSELoss()
+    # get loss functions from feature config
+    used_targets = [feat for feat in model_configuration["feature_config"]["target_features"] if
+                    feat not in model_configuration["feature_config"]["ignored_features"]]
+    loss_functions = [target_feat["loss_fn"] for target_feat in used_targets]
+    # instantiate class of loss functions if they are not already
+    for i, loss_fn in enumerate(loss_functions):
+        if inspect.isclass(loss_fn):
+            loss_functions[i] = loss_fn()
     writer = SummaryWriter(log_dir=f'{log_dir}/{model_configuration["id"]}_{training_id}', )
 
     best_val_loss = math.inf
@@ -239,9 +251,9 @@ def train_model(model: nn.Module,
         # 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])
-
+            logger.info(f"Rank {local_rank}: Train subset length: {len(train_subsets[epoch - 1])}")
+            train_subset = train_subsets[epoch - 1]
+            train_dataset.set_key_subset(train_subset)
             # recreate data loaders
             train_dataloader = DataLoader(
                 train_dataset,
@@ -249,104 +261,160 @@ def train_model(model: nn.Module,
                 num_workers=num_dataloader_workers,
                 # set multiprocessing start method to spawn
                 # multiprocessing_context="forkserver",
-                multiprocessing_context="spawn",
+                # multiprocessing_context="spawn",
             )
             train_length = train_epoch_lengths[epoch - 1]
-            val_dataloader = DataLoader(
-                val_dataset,
-                batch_size=None,
-                num_workers=num_dataloader_workers,
-                # set multiprocessing start method to spawn
-                # multiprocessing_context="forkserver",
-                multiprocessing_context="spawn",
-            )
-            val_length = val_epoch_lengths[epoch - 1]
+            logger.info(f"Rank {local_rank}: Current epoch train length: {train_length}")
+            logger.info(f"Rank {local_rank}: Val subset length: {len(val_subsets[epoch - 1])}")
+        else:
+            train_subset = train_dataset.lmdb_keys
+            train_length = len(train_dataset)
 
-        iterator = iter(train_dataloader)
+        train_iter = iter(train_dataloader)
         for step in tqdm(range(train_length)):
+            optimizer.zero_grad()
             try:
-                loss = batch_loss_fn(model,
-                                     iterator,
-                                     loss_functions,
-                                     device,
-                                     model_configuration)
+                with autocast("cuda"):
+                    loss = batch_loss_fn(model,
+                                         train_iter,
+                                         loss_functions,
+                                         device,
+                                         model_configuration)
             except StopIteration:
                 # if the iterator is exhausted, reset it
-                logger.info(f"Rank {local_rank}: Iterator exhausted, continuing to next epoch.")
-                break
+                # recreate data loaders
+                train_dataset.set_key_subset(train_subset)
+                train_dataloader = DataLoader(
+                    train_dataset,
+                    batch_size=None,
+                    num_workers=num_dataloader_workers,
+                    # set multiprocessing start method to spawn
+                    # multiprocessing_context="forkserver",
+                    # multiprocessing_context="spawn",
+                )
+                train_iter = iter(train_dataloader)
+                with autocast("cuda"):
+                    loss = batch_loss_fn(model,
+                                         train_iter,
+                                         loss_functions,
+                                         device,
+                                         model_configuration)
 
-            optimizer.zero_grad()
-            loss.backward()
-            optimizer.step()
+            # scale the loss and backpropagate
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+
+            # loss.backward()
+            # optimizer.step()
             scheduler.step()
+            scaler.update()
 
             total_train_loss += loss.item()
             if local_rank == 0:
                 if step % log_every_n_steps == 0:
+                    normalized_step = ((epoch - 1) * train_length + step) * world_size * \
+                                      training_configuration["batch_size"]
                     writer.add_scalar("Loss/Train_Step", loss.item(),
-                                      ((epoch - 1) * len(train_dataloader) + step) * training_configuration[
-                                          "batch_size"])
+                                      normalized_step)
                     writer.add_scalar("LR", scheduler.get_last_lr()[0],
-                                      ((epoch - 1) * len(train_dataloader) + step) * training_configuration[
-                                          "batch_size"])
+                                      normalized_step)
                     writer.flush()
 
         avg_train_loss = total_train_loss / len(train_dataloader)
         logger.info(f"Rank {local_rank}: Epoch {epoch}/{num_epochs} done. Local Train loss: {avg_train_loss:.4f}")
 
-        # sync before validation
-        if torch.distributed.is_initialized():
-            torch.distributed.barrier()
+        # clear up memory and dataloaders
+        del train_iter
+        del train_dataloader
 
         # Validation
-        # if local_rank == 0 or not torch.distributed.is_initialized():
-        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(val_length)):
-                try:
-                    loss = batch_loss_fn(model,
-                                         val_iter,
-                                         loss_functions,
-                                         device,
-                                         model_configuration)
-                except StopIteration:
-                    # if the iterator is exhausted, reset it
-                    logger.info(f"Rank {local_rank}: Iterator exhausted, continuing to next epoch.")
-                    break
-
-                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)
-        # else:
-        #     avg_val_loss = None
-        #     logger.info(f"Rank {local_rank}: Validation skipped, using 0 as validation loss.")
-
-        # sync before logging
-        if torch.distributed.is_initialized():
-            torch.distributed.barrier()
-
-        # 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)
+        try:
+            # create validation dataloader with proper subset
+            if torch.distributed.is_initialized():
+                val_subset = val_subsets[epoch - 1]
+                val_dataset.set_key_subset(val_subset)
+                val_dataloader = DataLoader(
+                    val_dataset,
+                    batch_size=None,
+                    num_workers=num_dataloader_workers,
+                    # set multiprocessing start method to spawn
+                    # multiprocessing_context="forkserver",
+                    # multiprocessing_context="spawn",
+                )
+                val_length = val_epoch_lengths[epoch - 1]
+                logger.info(f"Rank {local_rank}: Current epoch val length: {val_length}")
             else:
-                avg_val_loss_global = torch.tensor(0.0, device=device, dtype=torch.float32)
-            torch.distributed.all_reduce(avg_val_loss_global)
-            avg_val_loss_global /= torch.distributed.get_world_size()
+                val_subset = val_dataset.lmdb_keys
+                val_length = len(val_dataset)
 
-            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)
+                # sync before validation
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+            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(val_length)):
+                    try:
+                        loss = batch_loss_fn(model,
+                                             val_iter,
+                                             loss_functions,
+                                             device,
+                                             model_configuration)
+                    except StopIteration:
+                        # if the iterator is exhausted, reset it
+                        val_dataset.set_key_subset(val_subset)
+                        val_dataloader = DataLoader(
+                            val_dataset,
+                            batch_size=None,
+                            num_workers=num_dataloader_workers,
+                            # set multiprocessing start method to spawn
+                            # multiprocessing_context="forkserver",
+                            # multiprocessing_context="spawn",
+                        )
+                        val_iter = iter(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)
+
+            # clean up
+            del val_iter
+            del val_dataloader
+
+            # sync before logging
+            if torch.distributed.is_initialized():
+                torch.distributed.barrier()
+
+            # 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)
+                else:
+                    avg_val_loss_global = torch.tensor(0.0, 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)
+        except Exception as e:
+            logger.error(f"Rank {local_rank}: Validation failed: {e}")
+            raise e
 
         # only rank 0 checks for early stopping
         if local_rank == 0:
@@ -364,6 +432,7 @@ def train_model(model: nn.Module,
                 epochs_no_improve = 0
                 # torch.save(model.state_dict(), os.path.join(model_configuration["id"], "model.pt"))
                 save_fn = model_configuration["model_save_fn"]
+                logger.info(f"Rank {local_rank}: Saving model to {training_configuration['training_dir']}")
                 if torch.distributed.is_initialized():
                     save_fn(model.module, training_configuration)
                 else:
@@ -371,7 +440,7 @@ def train_model(model: nn.Module,
             else:
                 epochs_no_improve += 1
                 logger.info(
-                    f"Rank {local_rank}: No improvement in validation loss, no-improve count: {epochs_no_improve}")
+                    f"Rank {local_rank}: No improvement in validation loss, no-improve count: {epochs_no_improve} of max {patience}")
                 if epochs_no_improve >= patience:
                     logger.info("Early stopping triggered.")
                     # broadcast stop signal to all gpus
@@ -395,13 +464,5 @@ def train_model(model: nn.Module,
         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

@@ -1,3 +1,4 @@
+import math
 import os
 import pickle
 import random
@@ -14,7 +15,7 @@ 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 utils.data_utils import produce_window_batches, get_collated_batch_for_key
 
 from vsm_datascience_common.cycle_database_connection.db_utils import get_cycles_collection
 
@@ -26,11 +27,11 @@ def get_training_config(base_config: dict, model_config: dict):
     except KeyError:
         raise Exception("Model class missing in training or model config")
 
-    # training_config_id = get_config_id(base_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}"
+    # 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}")
 
@@ -49,6 +50,10 @@ def get_training_config(base_config: dict, model_config: dict):
         with open(f"{training_dir}/training_configuration.pickle", "rb") as f:
             training_config = pickle.load(f)
 
+        # update paths
+        training_config["training_dir"] = training_dir
+        training_config["model_dir"] = model_config["model_dir"]
+
     return training_config
 
 
@@ -68,6 +73,78 @@ def get_training_config_from_file(training_dir: str,
     return training_config
 
 
+def load_model_for_usage(model_configuration: dict,
+                         training_configuration: dict,
+                         sample_id: str) -> tuple[nn.Module, dict]:
+    """
+    Load the model associated with the given model configuration for the usage on the current node, given its training configuration
+
+
+    If the model does not exist, create it
+    Args:
+        model_configuration: configuration of the model
+        training_configuration: training configuration
+        sample_id: id of a sample data item
+    Returns:
+        model: model
+        training_config: training configuration with updated variables such as batch size
+
+    """
+
+    # 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:
+        local_rank = 0
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    dataset_dir = model_configuration["feature_config"]["dataset_dir"]
+    lmdb_env = lmdb.open(dataset_dir, readonly=True)
+
+    try:
+        load_fn = model_configuration["model_load_fn"]
+        model = load_fn(model_configuration=model_configuration,
+                        training_configuration=training_configuration,
+                        sample_key=sample_id,
+                        device=device,
+                        lmdb_env=lmdb_env)
+    except FileNotFoundError:
+        model_creation_fn = model_configuration["model_creation_fn"]
+        model = model_creation_fn(
+            model_configuration=model_configuration,
+            sample_key=sample_id,
+            lmdb_env=lmdb_env
+        )
+
+    # estimate batch size
+    sample_batch = get_collated_batch_for_key(sample_id, model_configuration, lmdb_env=lmdb_env)
+    if sample_batch is None or len(sample_batch) == 0:
+        raise ValueError(
+            f"Rank {local_rank}: Could not get sample input for key {sample_id} in dataset {dataset_dir}. "
+            f"Please check if the key exists in the LMDB dataset.")
+    sample_input = sample_batch[0][0]  # get the first item in the batch
+    input_shape = sample_input.shape
+    # send model to proper device
+    model.to(device)
+    batch_size = find_max_batch_size(model,
+                                     input_shape,
+                                     device=device)
+    # update max lr in training configuration, use square root scaling law based on initial batch size
+    training_configuration["learning_parameters"]["learning_rate"] = training_configuration["learning_parameters"][
+                                                                         "max_lr"] * math.sqrt(
+        batch_size / training_configuration["batch_size"])
+
+    # update batch size in training configuration
+    training_configuration["batch_size"] = batch_size
+
+    # close lmdb env
+    lmdb_env.close()
+
+    return model, training_configuration
+
+
 def weight_init(m):
     """
     Usage:
@@ -259,3 +336,44 @@ def get_data_ids(model_configuration: dict,
         test_ids = test_ids[:test_lim]
 
     return train_ids, val_ids, test_ids
+
+
+def find_max_batch_size(model,
+                        input_shape,
+                        device,
+                        dtype=torch.float32,
+                        max_batch=2048):
+    """
+    Find the maximum batch size that can be processed by the model without running out of memory.
+
+    Model must already be moved to the appropriate device (e.g., GPU).
+    Args:
+        model: model to test
+        input_shape: shape of the input tensor (excluding batch size)
+        device: device to work on
+        dtype: data type of the input tensor (default: torch.float32)
+        max_batch: maximum batch size to test (default: 2048)
+
+    Returns:
+        last_successful: the largest batch size that did not cause an out of memory error
+
+    """
+    batch_size = 1
+    last_successful = 1
+
+    while last_successful < max_batch:
+        try:
+            inputs = torch.randn((batch_size, *input_shape), device=device, dtype=dtype)
+            outputs = model(inputs)
+            loss = outputs.sum()
+            loss.backward()  # simulate training
+            torch.cuda.empty_cache()
+            last_successful = batch_size
+            batch_size *= 2
+        except RuntimeError as e:
+            if 'out of memory' in str(e):
+                torch.cuda.empty_cache()
+                break
+            else:
+                raise e
+    return last_successful

code/utils/utils.py

@@ -150,3 +150,22 @@ def convert_for_json(obj):
         return obj.item()
     else:
         return obj
+
+
+def recursive_dict_update(current_dict: dict, update_dict: dict) -> dict:
+    for key, value in update_dict.items():
+        if key not in current_dict:
+            current_dict[key] = value
+        else:
+            if isinstance(value, list):
+                if not isinstance(current_dict[key], list):
+                    raise ValueError(f"Type mismatch, {type(current_dict[key])} is not a list")
+                current_dict[key] = current_dict[key] + value
+            elif isinstance(value, dict):
+                if not isinstance(current_dict[key], dict):
+                    raise ValueError(f"Type mismatch, {type(current_dict[key])} is not a dict")
+                current_dict[key] = recursive_dict_update(current_dict[key], value)
+            else:
+                current_dict[key] = value
+
+    return current_dict

code/utils/visualization.py

@@ -1,61 +1,133 @@
 import numpy as np
 import torch
 from torch import nn
+import plotly
+import plotly.graph_objs as go
 
+from utils.inference import apply_sigmoid_if_necessary, scale_features
+from utils.model_utils import *
 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]
+def get_result_plotting_function(
+        inputs: np.ndarray | torch.Tensor,
+        predictions: np.ndarray | torch.Tensor,
+        actuals: np.ndarray | torch.Tensor,
+        output_feature_names: list,
+        output_colors: list,
+        inputs_to_plot: list,
+        model_configuration: dict,
+        skip_scaling: list = None) -> tuple:
+    fig_widget = go.FigureWidget()
 
-    # clear previous traces
-    fig_widget.data = []
+    processed_predictions = apply_sigmoid_if_necessary(predictions, model_configuration)
 
-    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,
-        )
+    if skip_scaling is not None:
+        features_to_scale = [feature if feature not in skip_scaling else None for feature in output_feature_names]
+    else:
+        features_to_scale = output_feature_names
+    scaled_predictions = scale_features(processed_predictions, features_to_scale, model_configuration,
+                                        subset_name="test")
+    scaled_actuals = scale_features(actuals, features_to_scale, model_configuration, subset_name="test")
 
-    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]
+    target_features = model_configuration["feature_config"]["target_features"]
+    ignored_features = model_configuration["feature_config"]["ignored_features"]
+    used_targets = [feature for feature in target_features if feature['name'] not in ignored_features]
+    output_feature_indices = list()
+    for i, feature in enumerate(used_targets):
+        if feature["name"] in output_feature_names:
+            output_feature_indices.append(i)
 
-        output = float(output)
+    def plot_prediction_windows(index: int):
+        window_features = inputs[index]
+        indices = np.arange(window_features.shape[0])
+        actual = scaled_actuals[index]
+        predicted = scaled_predictions[index]
 
-        scaled_output = inverse_scale_feature(output,
-                                              output_feature_names[i],
-                                              scalers)
+        # clear previous traces
+        fig_widget.data = []
+        fig_widget.layout.shapes = []
 
-        scaled_preds.append(scaled_output)
+        downsampling_rate = 1
+        step_size = model_configuration["preprocessing"]["window_shift"] * downsampling_rate
+        actuals_raw_until_now = scaled_actuals[:index + 1, output_feature_indices]
+        predicted_raw_until_now = scaled_predictions[:index + 1, output_feature_indices]
+        out_indices = np.arange(0, len(actuals_raw_until_now) * step_size, step_size)
 
-    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)
+        offset = model_configuration["input_window_length"] - (index * step_size)
 
-        scaled_actuals.append(scaled_output)
+        # if find_ovs:
+        #     # calculate the number of cycles until now, only use first occurence of 0, not consecutive zeros
+        #     actual_ov_indices_raw = [x.item() for x in torch.where(actuals_raw_until_now == 0)[0].numpy()]
+        #     actual_ov_indices = [actual_ov_indices_raw[i] for i in range(len(actual_ov_indices_raw)) if
+        #                          i == 0 or actual_ov_indices_raw[i] - 1 not in actual_ov_indices_raw]
+        #     num_cycles_until_now = len(actual_ov_indices)
+        #
+        #     # find cycle starts, cycle starts are where the actuals jump from positive to negative
+        #     start_offset = 5
+        #     cycle_start_indices = np.where(np.diff(actuals_raw_until_now) < 0)[0] + start_offset
+        #
+        #     current_index = 0
+        #     predicted_ov_indices = list()
+        #     for i in range(num_cycles_until_now):
+        #         first_post_0_predicted = np.where(predicted_raw_until_now[current_index:] >= 0)[0]
+        #         first_post_0_predicted = first_post_0_predicted[0].item() if len(first_post_0_predicted) > 0 else None
+        #         if first_post_0_predicted is not None:
+        #             predicted_ov_indices.append(first_post_0_predicted + current_index)
+        #         if i < len(cycle_start_indices):
+        #             current_index = cycle_start_indices[i]
+        #
+        #     # add step sizes
+        #     actual_ov_indices = [x * step_size for x in actual_ov_indices]
+        #     predicted_ov_indices = [x * step_size for x in predicted_ov_indices]
+        #
+        #     # plot actual and predicted ovs
+        #     for ov_index in actual_ov_indices:
+        #         fig_widget.add_vline(
+        #             x=ov_index + offset,
+        #             line=dict(color='blue', width=2, dash='dot'),
+        #             name="Actual Ovulation",
+        #         )
+        #
+        #     print(len(predicted_ov_indices))
+        #     for ov_index in predicted_ov_indices:
+        #         fig_widget.add_vline(
+        #             x=ov_index + offset,
+        #             line=dict(color='red', width=2, dash='dot'),
+        #             name="Predicted Ovulation",
+        #         )
 
-    # add actual and predicted values
-    print(f"actual: {scaled_actuals}, predicted: {scaled_preds}")
+        for feature in inputs_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,
+            )
+
+        length_limiter = model_configuration["input_window_length"] // step_size
+        print(length_limiter)
+
+        for output_feature_index in output_feature_indices:
+            color = output_colors[output_feature_index]
+            fig_widget.add_scatter(
+                x=(out_indices + offset)[-length_limiter:],
+                y=actuals_raw_until_now[:, output_feature_index][-length_limiter:],
+                mode="lines",
+                name="Actuals Raw",
+                line=dict(color=color, width=2, dash='dot'),
+            )
+            fig_widget.add_scatter(
+                x=(out_indices + offset)[-length_limiter:],
+                y=predicted_raw_until_now[:, output_feature_index][-length_limiter:],
+                mode="lines",
+                name="Predicted Raw",
+                line=dict(color=color, width=2),
+            )
+
+        # add actual and predicted values
+        print(f"actual: {actual}, predicted: {predicted}")
+
+    return fig_widget, plot_prediction_windows