Added training script to remotely train diffusion model

src/losses/crps_metric.py

@@ -43,4 +43,33 @@ class CRPSLoss(nn.Module):
         crps_per_sample = np.trapz(tiled_probs ** 2, imbalances, axis=-1)
         crps = np.mean(crps_per_sample)
 
-        return crps
+        return crps
+
+def crps_from_samples(samples, targets):
+    """
+    Compute the Continuous Ranked Probability Score (CRPS) from multi-day samples and targets
+    using a vectorized approach with PyTorch tensors.
+    
+    :param samples: (day, n_samples, n_timesteps) tensor of forecasted samples
+    :param targets: (day, n_timesteps) tensor of observed values
+    :return: (day, n_timesteps) tensor of CRPS for each timestep for each day
+    """
+    days, n_samples, n_timesteps = samples.shape
+    
+    # Reshape targets to broadcast along the samples dimension (n_samples)
+    targets_reshaped = targets.unsqueeze(1)
+
+    # Compute the absolute differences of forecasts and observations
+    abs_diff = torch.abs(samples - targets_reshaped)
+    # Compute the average of the absolute differences along the samples dimension
+    term1 = torch.mean(abs_diff, dim=1)
+    
+    # Compute the pairwise absolute differences between all samples for each day
+    pairwise_abs_diff = torch.abs(samples.unsqueeze(2) - samples.unsqueeze(1))
+    # Compute the average of these differences along the sample dimensions
+    term2 = torch.mean(pairwise_abs_diff, dim=(1, 2)) / 2
+    
+    # CRPS for each timestep for each day
+    crps = term1 - term2
+    
+    return crps