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,
+            )