fixes

code/utils/loss_functions.py

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