Updated some stuff

src/trainers/quantile_trainer.py

@@ -558,8 +558,7 @@ class NonAutoRegressiveQuantileRegression(Trainer):
 
                 outputs = self.model(inputs)
                 outputted_samples = [
-                    sample_from_dist(self.quantiles, output.cpu().numpy())
-                    for output in outputs
+                    sample_from_dist(self.quantiles, output.cpu()) for output in outputs
                 ]
 
                 outputted_samples = torch.tensor(outputted_samples)
@@ -618,20 +617,24 @@ class NonAutoRegressiveQuantileRegression(Trainer):
 
     def debug_plots(self, task, train: bool, data_loader, sample_indices, epoch):
         for actual_idx, idx in sample_indices.items():
-            initial, target, _ = data_loader.dataset[idx]
+            features, target, _ = data_loader.dataset[idx]
 
-            # get predictions
-            initial = initial.to(self.device)
+            features = features.to(self.device)
+            target = target.to(self.device)
 
-            predicted_quantiles = self.model(initial)
-            predictions = predicted_quantiles.reshape(-1, len(self.quantiles))
+            self.model.eval()
+            with torch.no_grad():
+                predicted_quantiles = self.model(features)
+                predictions = predicted_quantiles.reshape(-1, len(self.quantiles))
 
             samples = [
                 sample_from_dist(self.quantiles, predictions) for _ in range(100)
             ]
             samples = torch.tensor(samples)
 
-            fig = self.get_plot(initial, target, samples, show_legend=(0 == 0))
+            fig, fig2 = self.get_plot(
+                features[:96], target, samples, show_legend=(0 == 0)
+            )
 
             task.get_logger().report_matplotlib_figure(
                 title="Training" if train else "Testing",
@@ -640,17 +643,12 @@ class NonAutoRegressiveQuantileRegression(Trainer):
                 figure=fig,
             )
 
-            fig, ax = plt.subplots(figsize=(20, 10))
-            for i in range(10):
-                ax.plot(samples[i], label=f"Sample {i}")
-
-            ax.plot(target, label="Real NRV", linewidth=3)
-            ax.legend()
             task.get_logger().report_matplotlib_figure(
-                title="Training" if train else "Testing",
-                series=f"Sample {actual_idx} Samples",
+                title="Training Samples" if train else "Testing Samples",
+                series=f"Sample {actual_idx} samples",
                 iteration=epoch,
-                figure=fig,
+                figure=fig2,
+                report_interactive=False,
             )
 
             plt.close()
@@ -750,6 +748,8 @@ class NonAutoRegressiveQuantileRegression(Trainer):
             ]
         )
 
+        ax.set_ylim(-1500, 1500)
+
         fig2, ax2 = plt.subplots(figsize=(20, 10))
         for i in range(10):
             ax2.plot(predictions_np[i], label=f"Sample {i}")
@@ -757,6 +757,8 @@ class NonAutoRegressiveQuantileRegression(Trainer):
         ax2.plot(next_day_np, label="Real NRV", linewidth=3)
         ax2.legend()
 
+        ax2.set_ylim(-1500, 1500)
+
         return fig, fig2
 
     def calculate_crps_from_samples(self, task, dataloader, epoch: int):
@@ -812,26 +814,36 @@ class NonAutoRegressiveQuantileRegression(Trainer):
 
         #  using the policy evaluator, evaluate the policy with the generated samples
         if self.policy_evaluator is not None:
-            _, test_loader = self.data_processor.get_dataloaders(
-                predict_sequence_length=self.model.output_size, full_day_skip=True
+            optimal_penalty, profit, charge_cycles = (
+                self.policy_evaluator.optimize_penalty_for_target_charge_cycles(
+                    idx_samples=generated_samples,
+                    test_loader=dataloader,
+                    initial_penalty=500,
+                    target_charge_cycles=283,
+                    learning_rate=2,
+                    max_iterations=100,
+                    tolerance=1,
+                )
             )
-            self.policy_evaluator.evaluate_test_set(generated_samples, test_loader)
-            df = self.policy_evaluator.get_profits_as_scalars()
 
-            # for each row, report the profits
-            for idx, row in df.iterrows():
-                task.get_logger().report_scalar(
-                    title="Profit",
-                    series=f"penalty_{row['Penalty']}",
-                    value=row["Total Profit"],
-                    iteration=epoch,
-                )
+            print(
+                f"Optimal Penalty: {optimal_penalty}, Profit: {profit}, Charge Cycles: {charge_cycles}"
+            )
 
-            df = self.policy_evaluator.get_profits_till_400()
-            for idx, row in df.iterrows():
-                task.get_logger().report_scalar(
-                    title="Profit_till_400",
-                    series=f"penalty_{row['Penalty']}",
-                    value=row["Profit_till_400"],
-                    iteration=epoch,
-                )
+            task.get_logger().report_scalar(
+                title="Optimal Penalty",
+                series="test",
+                value=optimal_penalty,
+                iteration=epoch,
+            )
+
+            task.get_logger().report_scalar(
+                title="Optimal Profit", series="test", value=profit, iteration=epoch
+            )
+
+            task.get_logger().report_scalar(
+                title="Optimal Charge Cycles",
+                series="test",
+                value=charge_cycles,
+                iteration=epoch,
+            )

src/trainers/trainer.py

@@ -7,6 +7,7 @@ import numpy as np
 from plotly.subplots import make_subplots
 from clearml.config import running_remotely
 from torchinfo import summary
+import matplotlib.pyplot as plt
 
 
 class Trainer:
@@ -329,18 +330,7 @@ class Trainer:
         return fig
 
     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, idx in enumerate(sample_indices):
+        for actual_idx, idx in sample_indices.items():
             features, target, _ = data_loader.dataset[idx]
 
             features = features.to(self.device)
@@ -350,30 +340,26 @@ class Trainer:
             with torch.no_grad():
                 predictions = self.model(features).cpu()
 
-            sub_fig = self.get_plot(
-                features[:96], target, predictions, show_legend=(i == 0)
+            fig, fig2 = self.get_plot(
+                features[:96], target, predictions, show_legend=(0 == 0)
             )
 
-            row = i + 1
-            col = 1
+            task.get_logger().report_matplotlib_figure(
+                title="Training" if train else "Testing",
+                series=f"Sample {actual_idx}",
+                iteration=epoch,
+                figure=fig,
+            )
 
-            for trace in sub_fig.data:
-                fig.add_trace(trace, row=row, col=col)
+            task.get_logger().report_matplotlib_figure(
+                title="Training Samples" if train else "Testing Samples",
+                series=f"Sample {actual_idx} samples",
+                iteration=epoch,
+                figure=fig2,
+                report_interactive=False,
+            )
 
-            # loss = self.criterion(predictions.to(self.device), target.squeeze(-1).to(self.device)).item()
-
-            # fig['layout']['annotations'][i].update(text=f"{loss.__class__.__name__}: {loss:.6f}")
-
-        # 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,
-        )
+            plt.close()
 
     def debug_scatter_plot(self, task, train: bool, samples, epoch):
         X, y = samples