added code

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