Plots to compare between quantile regression and diffusion

src/trainers/autoregressive_trainer.py

@@ -33,67 +33,29 @@ class AutoRegressiveTrainer(Trainer):
         self.model.output_size = 1
 
     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
-
-        #  check if self has get_plot_error
-        if hasattr(self, "get_plot_error"):
-            cols = 2
-            print("Using get_plot_error")
-        else:
-            cols = 1
-            print("Using get_plot")
-
-        fig = make_subplots(
-            rows=rows,
-            cols=cols,
-            subplot_titles=[f"Sample {i+1}" for i in range(num_samples)],
-        )
-
-        for i, idx in enumerate(sample_indices):
-            auto_regressive_output = self.auto_regressive(data_loader.dataset, [idx])
+        for actual_idx, idx in sample_indices.items():
+            auto_regressive_output = self.auto_regressive(data_loader.dataset, [idx]*1000)
             if len(auto_regressive_output) == 3:
                 initial, predictions, target = auto_regressive_output
             else:
-                initial, predictions, _, target = auto_regressive_output
+                initial, _, predictions, target = auto_regressive_output
 
-            initial = initial.squeeze(0)
-            predictions = predictions.squeeze(0)
-            target = target.squeeze(0)
+            
+            # keep one initial
+            initial = initial[0]
+            target = target[0]
 
-            sub_fig = self.get_plot(initial, target, predictions, show_legend=(i == 0))
+            predictions = predictions
 
-            row = i + 1
-            col = 1
+            fig = self.get_plot(initial, target, predictions, show_legend=(0 == 0))
 
-            for trace in sub_fig.data:
-                fig.add_trace(trace, row=row, col=col)
-
-            if cols == 2:
-                error_sub_fig = self.get_plot_error(
-                    target, predictions
-                )
-                for trace in error_sub_fig.data:
-                    fig.add_trace(trace, row=row, col=col + 1)
-
-            loss = self.criterion(
-                predictions.to(self.device), target.to(self.device)
-            ).item()
-
-            fig["layout"]["annotations"][i].update(
-                text=f"{self.criterion.__class__.__name__}: {loss:.6f}"
+            task.get_logger().report_matplotlib_figure(
+                title="Training" if train else "Testing",
+                series=f'Sample {actual_idx}',
+                iteration=epoch,
+                figure=fig,
             )
 
-        # y axis same for all plots
-        # fig.update_yaxes(range=[-1, 1], col=1)
-
-        fig.update_layout(height=1000 * rows)
-        task.get_logger().report_plotly(
-            title=f"{'Training' if train else 'Test'} Samples",
-            series="full_day",
-            iteration=epoch,
-            figure=fig,
-        )
 
     def auto_regressive(self, data_loader, idx, sequence_length: int = 96):
         self.model.eval()