Plots to compare between quantile regression and diffusion

src/trainers/diffusion_trainer.py

@@ -96,7 +96,16 @@ class DiffusionTrainer:
         else:
             loader = test_loader
 
-        indices = np.random.randint(0, len(loader.dataset) - 1, size=num_samples)
+        # set seed
+        np.random.seed(42)
+
+        actual_indices = np.random.choice(loader.dataset.full_day_valid_indices, num_samples, replace=False)
+        indices = {}
+        for i in actual_indices:
+            indices[i] = loader.dataset.valid_indices.index(i)
+
+        print(actual_indices)
+        
         return indices
     
     def init_clearml_task(self, task):
@@ -171,7 +180,7 @@ class DiffusionTrainer:
 
         
     def debug_plots(self, task, training: bool, data_loader, sample_indices, epoch):
-        for i, idx in enumerate(sample_indices):
+        for actual_idx, idx in sample_indices.items():
             features, target, _ = data_loader.dataset[idx]
 
             features = features.to(self.device)
@@ -180,6 +189,8 @@ class DiffusionTrainer:
             self.model.eval()
             with torch.no_grad():
                 samples = self.sample(self.model, 100, features).cpu().numpy()
+                samples = self.data_processor.inverse_transform(samples)
+                target = self.data_processor.inverse_transform(target)
                
             ci_99_upper = np.quantile(samples, 0.995, axis=0)
             ci_99_lower = np.quantile(samples, 0.005, axis=0)
@@ -218,7 +229,7 @@ class DiffusionTrainer:
 
             task.get_logger().report_matplotlib_figure(
                 title="Training" if training else "Testing",
-                series=f'Sample {i}',
+                series=f'Sample {actual_idx}',
                 iteration=epoch,
                 figure=fig,
             )