Rewrote dataset to be able to include new features

src/trainers/autoregressive_trainer.py

@@ -10,19 +10,16 @@ from plotly.subplots import make_subplots
 from trainers.trainer import Trainer
 
 class AutoRegressiveTrainer(Trainer):
-    def debug_plots(self, task, train: bool, samples, epoch):
-        X, y = samples
-        X = X.to(self.device)
-
-        num_samples = len(X)
+    def debug_plots(self, task, train: bool, data_loader, sample_indices, epoch):
+        num_samples = len(sample_indices)
         rows = num_samples  # One row per sample since we only want one column
         cols = 1
         
         fig = make_subplots(rows=rows, cols=cols, subplot_titles=[f'Sample {i+1}' for i in range(num_samples)])
 
-        for i, (current_day, next_day) in enumerate(zip(X, y)):
-            predictions = self.predict_auto_regressive(current_day)
-            sub_fig = self.get_plot(current_day, next_day, predictions, show_legend=(i == 0))
+        for i, idx in enumerate(sample_indices):
+            initial, predictions, target = self.auto_regressive(data_loader, idx)
+            sub_fig = self.get_plot(initial, target, predictions, show_legend=(i == 0))
             
             row = i + 1
             col = 1
@@ -30,7 +27,7 @@ class AutoRegressiveTrainer(Trainer):
             for trace in sub_fig.data:
                 fig.add_trace(trace, row=row, col=col)
 
-            loss = self.criterion(predictions.to(self.device), next_day.to(self.device)).item()
+            loss = self.criterion(predictions.to(self.device), target.to(self.device)).item()
 
             fig['layout']['annotations'][i].update(text=f"{loss.__class__.__name__}: {loss:.6f}")
 
@@ -46,14 +43,38 @@ class AutoRegressiveTrainer(Trainer):
             figure=fig
         )
 
+    def auto_regressive(self, data_loader, idx, sequence_length: int = 96):
+        self.model.eval()
+        target_full = []
+        predictions_full = []
 
-    def predict_auto_regressive(self, initial_sequence: torch.Tensor, sequence_length: int = 96):
-        initial_sequence = initial_sequence.to(self.device)
+        prev_features, target = data_loader.dataset[idx]
+        prev_features = prev_features.to(self.device)
 
-        return predict_auto_regressive(self.model, initial_sequence, sequence_length)
+        initial_sequence = prev_features[:96]
 
-    def random_day_prediction(self):
-        current_day_features, next_day_features = self.data_processor.get_random_test_day()
+        target_full.append(target)
+        with torch.no_grad():
+            prediction = self.model(prev_features.unsqueeze(0))
+        predictions_full.append(prediction.squeeze(-1))
 
-        predictions = self.predict_auto_regressive(current_day_features)
-        return current_day_features, next_day_features, predictions
+        for i in range(sequence_length - 1):
+            new_features = torch.cat((prev_features[1:97].cpu(), prediction.squeeze(-1).cpu()), dim=0)
+
+            # get the other needed features
+            other_features, new_target = data_loader.dataset.random_day_autoregressive(idx + i + 1)
+
+            if other_features is not None:
+                prev_features = torch.cat((new_features, other_features), dim=0)
+            else:
+                prev_features = new_features
+
+            # add target to target_full
+            target_full.append(new_target)
+
+            # predict
+            with torch.no_grad():
+                prediction = self.model(new_features.unsqueeze(0).to(self.device))
+            predictions_full.append(prediction.squeeze(-1))
+
+        return initial_sequence.cpu(), torch.stack(predictions_full).cpu(), torch.stack(target_full).cpu()