Implemented Non Autorgressive Quantile Regression

src/data/dataset.py

@@ -12,6 +12,8 @@ class NrvDataset(Dataset):
         self.nrv = torch.tensor(dataframe['nrv'].to_numpy(), dtype=torch.float32)
         self.load_forecast = torch.tensor(dataframe['load_forecast'].to_numpy(), dtype=torch.float32)
         self.total_load = torch.tensor(dataframe['total_load'].to_numpy(), dtype=torch.float32)
+        self.pv_gen_forecast = torch.tensor(dataframe['pv_forecast'].to_numpy(), dtype=torch.float32)
+        self.wind_gen_forecast = torch.tensor(dataframe['wind_forecast'].to_numpy(), dtype=torch.float32)
 
         self.sequence_length = sequence_length
         self.predict_sequence_length = predict_sequence_length
@@ -45,10 +47,26 @@ class NrvDataset(Dataset):
             load_history = self.total_load[actual_idx:actual_idx+self.sequence_length]
             features.append(load_history.view(-1))
 
+        if self.data_config.PV_HISTORY:
+            pv_history = self.pv_gen_forecast[actual_idx:actual_idx+self.sequence_length]
+            features.append(pv_history.view(-1))
+
+        if self.data_config.WIND_HISTORY:
+            wind_history = self.wind_gen_forecast[actual_idx:actual_idx+self.sequence_length]
+            features.append(wind_history.view(-1))
+
         if self.data_config.LOAD_FORECAST:
             load_forecast = self.load_forecast[actual_idx+self.sequence_length:actual_idx+self.sequence_length+self.predict_sequence_length]
             features.append(load_forecast.view(-1))
 
+        if self.data_config.PV_FORECAST:
+            pv_forecast = self.pv_gen_forecast[actual_idx+self.sequence_length:actual_idx+self.sequence_length+self.predict_sequence_length]
+            features.append(pv_forecast.view(-1))
+
+        if self.data_config.WIND_FORECAST:
+            wind_forecast = self.wind_gen_forecast[actual_idx+self.sequence_length:actual_idx+self.sequence_length+self.predict_sequence_length]
+            features.append(wind_forecast.view(-1))
+
         if not features:
             raise ValueError("No features are configured to be included in the dataset.")
 
@@ -75,10 +93,28 @@ class NrvDataset(Dataset):
             load_history = self.total_load[idx:idx+self.sequence_length]
             features.append(load_history.view(-1))
 
+        if self.data_config.PV_HISTORY:
+            pv_history = self.pv_gen_forecast[idx:idx+self.sequence_length]
+            features.append(pv_history.view(-1))
+
+        if self.data_config.WIND_HISTORY:
+            wind_history = self.wind_gen_forecast[idx:idx+self.sequence_length]
+            features.append(wind_history.view(-1))
+
         if self.data_config.LOAD_FORECAST:
             load_forecast = self.load_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
             features.append(load_forecast.view(-1))
 
+        if self.data_config.PV_FORECAST:
+            pv_forecast = self.pv_gen_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
+            features.append(pv_forecast.view(-1))
+
+        if self.data_config.WIND_FORECAST:
+            wind_forecast = self.wind_gen_forecast[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
+            features.append(wind_forecast.view(-1))
+
+
+
         target = self.nrv[idx+self.sequence_length:idx+self.sequence_length+self.predict_sequence_length]
 
         if len(features) == 0:

src/data/preprocessing.py

@@ -9,11 +9,24 @@ import pytz
 
 history_data_path = "../../data/history-quarter-hour-data.csv"
 forecast_data_path = "../../data/load_forecast.csv"
+pv_forecast_data_path = "../../data/pv_gen_forecast.csv"
+wind_forecast_data_path = "../../data/wind_gen_forecast.csv"
 
 class DataConfig:
-    NRV_HISTORY: bool = True
-    LOAD_FORECAST: bool = True
-    LOAD_HISTORY: bool = False
+    def __init__(self):
+        self.NRV_HISTORY: bool = True
+
+        ### LOAD ###
+        self.LOAD_FORECAST: bool = True
+        self.LOAD_HISTORY: bool = False
+
+        ### PV ###
+        self.PV_FORECAST: bool = False
+        self.PV_HISTORY: bool = False
+
+        ### WIND ###
+        self.WIND_FORECAST: bool = False
+        self.WIND_HISTORY: bool = False
 
 class DataProcessor:
     def __init__(self, data_config: DataConfig):
@@ -26,8 +39,12 @@ class DataProcessor:
 
         self.history_features = self.get_nrv_history()
         self.future_features = self.get_load_forecast()
+        self.pv_forecast = self.get_pv_forecast()
+        self.wind_forecast = self.get_wind_forecast()
 
         self.all_features = self.history_features.merge(self.future_features, on='datetime', how='left')
+        self.all_features = self.all_features.merge(self.pv_forecast, on='datetime', how='left')
+        self.all_features = self.all_features.merge(self.wind_forecast, on='datetime', how='left')
 
         self.data_config = data_config
 
@@ -69,6 +86,31 @@ class DataProcessor:
         df.sort_values(by="datetime", inplace=True)
         return df
     
+    def get_pv_forecast(self):
+        df = pd.read_csv(pv_forecast_data_path, delimiter=';')
+
+        df = df.rename(columns={'dayahead11hforecast': 'pv_forecast', 'Datetime': 'datetime'})
+        df = df[['datetime', 'pv_forecast']]
+
+        df = df.groupby('datetime').mean().reset_index()
+        df['datetime'] = pd.to_datetime(df['datetime'], utc=True)
+        df.sort_values(by="datetime", inplace=True)
+        return df
+
+    def get_wind_forecast(self):
+        df = pd.read_csv(wind_forecast_data_path, delimiter=';')
+
+        df = df.rename(columns={'dayaheadforecast': 'wind_forecast', 'datetime': 'datetime'})
+        df = df[['datetime', 'wind_forecast']]
+
+        # remove nan rows
+        df = df[~df['wind_forecast'].isnull()]
+
+        df = df.groupby('datetime').mean().reset_index()
+        df['datetime'] = pd.to_datetime(df['datetime'], utc=True)
+        df.sort_values(by="datetime", inplace=True)
+        return df
+    
     def set_batch_size(self, batch_size: int):
         self.batch_size = batch_size
 

src/losses/__init__.py

@@ -1 +1 @@
-from .pinball_loss import PinballLoss
+from .pinball_loss import PinballLoss, NonAutoRegressivePinballLoss

src/losses/pinball_loss.py

@@ -2,32 +2,30 @@ import torch
 from torch import nn
 
 class PinballLoss(nn.Module):
-    """
-    Calculates the quantile loss function.
-
-    Attributes
-    ----------
-    self.pred : torch.tensor
-        Predictions.
-    self.target : torch.tensor
-        Target to predict.
-    self.quantiles : torch.tensor
-    """
     def __init__(self, quantiles):
         super(PinballLoss, self).__init__()
-        self.quantiles_tensor = quantiles
-        self.quantiles = quantiles.tolist()
+        self.quantiles_tensor = torch.tensor(quantiles, dtype=torch.float32)
         
     def forward(self, pred, target):
-        """
-        Computes the loss for the given prediction.
-        """
-
         error = target - pred
-        upper =  self.quantiles_tensor * error
+        upper = self.quantiles_tensor * error
         lower = (self.quantiles_tensor - 1) * error 
-
         losses = torch.max(lower, upper)
-        loss = torch.mean(torch.sum(losses, dim=1))
+        loss = torch.mean(torch.mean(losses, dim=0))
+        return loss
+    
+
+class NonAutoRegressivePinballLoss(nn.Module):
+    def __init__(self, quantiles):
+        super(NonAutoRegressivePinballLoss, self).__init__()
+        self.quantiles_tensor = torch.tensor(quantiles, dtype=torch.float32)
+
+    def forward(self, pred, target):
+        pred = pred.reshape(-1, 96, len(self.quantiles_tensor))
+        target_expanded = target.unsqueeze(2)
+        error = target_expanded - pred
+        upper = self.quantiles_tensor * error
+        lower = (self.quantiles_tensor - 1) * error
+        losses = torch.max(lower, upper)
+        loss = torch.mean(losses)
         return loss
-    

src/trainers/autoregressive_trainer.py

@@ -19,7 +19,12 @@ class AutoRegressiveTrainer(Trainer):
         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):
-            initial, predictions, target = self.auto_regressive(data_loader, idx)
+            auto_regressive_output = self.auto_regressive(data_loader, idx)
+            if len(auto_regressive_output) == 3:
+                initial, predictions, target = auto_regressive_output
+            else:
+                initial, predictions, _, target = auto_regressive_output
+
             sub_fig = self.get_plot(initial, target, predictions, show_legend=(i == 0))
             
             row = i + 1

src/trainers/quantile_trainer.py

@@ -5,13 +5,15 @@ from trainers.trainer import Trainer
 from trainers.autoregressive_trainer import AutoRegressiveTrainer
 from data.preprocessing import DataProcessor
 from utils.clearml import ClearMLHelper
-from losses import PinballLoss
+from losses import PinballLoss, NonAutoRegressivePinballLoss
 from plotly.subplots import make_subplots
 import plotly.graph_objects as go
 import numpy as np
+from tqdm import tqdm
+import matplotlib.pyplot as plt
 
 
-class QuantileTrainer(AutoRegressiveTrainer):
+class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
     def __init__(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer, data_processor: DataProcessor, quantiles: list, device: torch.device, clearml_helper: ClearMLHelper = None, debug: bool = True):
         quantiles_tensor = torch.tensor(quantiles)
         quantiles_tensor = quantiles_tensor.to(device)
@@ -26,29 +28,29 @@ class QuantileTrainer(AutoRegressiveTrainer):
         transformed_metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }
 
         with torch.no_grad():
-            for inputs, targets in dataloader:
-                inputs, targets = inputs.to(self.device), targets
-                outputs = self.model(inputs)
+            total_amount_samples = len(dataloader.dataset) - 95
+
+            for idx in tqdm(range(total_amount_samples)):
+                _, outputs, samples, targets = self.auto_regressive(dataloader, idx)
 
-                samples = []
-                for output in outputs:
-                    samples.append(self.sample_from_dist(self.quantiles.cpu().numpy(), output.cpu().numpy()))
-                
-                samples = torch.tensor(samples).to(self.device).reshape(-1, 1)
 
                 inversed_samples = torch.tensor(self.data_processor.inverse_transform(samples))
-                inversed_targets = torch.tensor(self.data_processor.inverse_transform(targets.reshape(-1, 1)))
+                inversed_targets = torch.tensor(self.data_processor.inverse_transform(targets))
+
+                outputs = outputs.to(self.device)
+                targets = targets.to(self.device)
+                samples = samples.to(self.device)
 
                 for metric in self.metrics_to_track:
                     if metric.__class__ != PinballLoss:
-                        transformed_metrics[metric.__class__.__name__] += metric(samples, targets.view(-1, 1).to(self.device))
+                        transformed_metrics[metric.__class__.__name__] += metric(samples, targets)
                         metrics[metric.__class__.__name__] += metric(inversed_samples, inversed_targets)
                     else:
-                        transformed_metrics[metric.__class__.__name__] += metric(outputs, targets.view(-1, 1).to(self.device))
+                        transformed_metrics[metric.__class__.__name__] += metric(outputs, targets)
 
             for metric in self.metrics_to_track:
-                metrics[metric.__class__.__name__] /= len(dataloader)
-                transformed_metrics[metric.__class__.__name__] /= len(dataloader)
+                metrics[metric.__class__.__name__] /= total_amount_samples
+                transformed_metrics[metric.__class__.__name__] /= total_amount_samples
 
             for metric_name, metric_value in metrics.items():
                 if PinballLoss.__name__ in metric_name:
@@ -125,12 +127,12 @@ class QuantileTrainer(AutoRegressiveTrainer):
             sample = self.sample_from_dist(self.quantiles.cpu(), prediction.squeeze(-1).cpu().numpy())
             predictions_sampled.append(sample)
 
-        return initial_sequence.cpu(), torch.stack(predictions_full).cpu(), torch.stack(target_full).cpu()
+        return initial_sequence.cpu(), torch.stack(predictions_full).cpu(), torch.tensor(predictions_sampled).reshape(-1, 1), torch.stack(target_full).cpu()
     
     @staticmethod
     def sample_from_dist(quantiles, output_values):
         # Interpolate the inverse CDF
-        inverse_cdf = interp1d(quantiles, output_values, kind='quadratic', bounds_error=False, fill_value="extrapolate")
+        inverse_cdf = interp1d(quantiles, output_values, kind='linear', bounds_error=False, fill_value="extrapolate")
 
         # generate one random uniform number
         uniform_random_numbers = np.random.uniform(0, 1, 1000)
@@ -141,4 +143,141 @@ class QuantileTrainer(AutoRegressiveTrainer):
         # Return the mean of the samples
         return np.mean(samples)
 
-        
+    def plot_quantile_percentages(self, task, data_loader, train: bool = True, iteration: int = None):
+
+        total = 0
+        quantile_counter = {q: 0 for q in self.quantiles.cpu().numpy()}
+
+        with torch.no_grad():
+            for inputs, targets in data_loader:
+                inputs = inputs.to("cuda")
+                output = self.model(inputs)
+
+                # output shape: (batch_size, num_quantiles)
+                # target shape: (batch_size, 1)
+                for i, q in enumerate(self.quantiles.cpu().numpy()):
+                    quantile_counter[q] += np.sum(targets.squeeze(-1).cpu().numpy() < output[:, i].cpu().numpy())
+
+                total += len(targets)
+
+        # to numpy array of length len(quantiles)
+        percentages = np.array([quantile_counter[q] / total for q in self.quantiles.cpu().numpy()])
+
+        bar_width = 0.35
+        index = np.arange(len(self.quantiles.cpu().numpy()))
+
+        # Plotting the bars
+        fig, ax = plt.subplots(figsize=(15, 10))
+
+        bar1 = ax.bar(index, self.quantiles.cpu().numpy(), bar_width, label='Ideal', color='brown')
+        bar2 = ax.bar(index + bar_width, percentages, bar_width, label='NN model', color='blue')
+
+        # Adding the percentage values above the bars for bar2
+        for rect in bar2:
+            height = rect.get_height()
+            ax.text(rect.get_x() + rect.get_width() / 2., 1.005 * height,
+                    f'{height:.2}', ha='center', va='bottom')  # Format the number as a percentage
+
+        series_name = "Training Set" if train else "Test Set"
+
+        # Adding labels and title
+        ax.set_xlabel('Quantile')
+        ax.set_ylabel('Fraction of data under quantile forecast')
+        ax.set_title(f'Quantile Performance Comparison ({series_name})')
+        ax.set_xticks(index + bar_width / 2)
+        ax.set_xticklabels(self.quantiles.cpu().numpy())
+        ax.legend()
+
+        task.get_logger().report_matplotlib_figure(title='Quantile Performance Comparison', series=series_name, report_image=True, figure=plt, iteration=iteration)
+        plt.close()
+                
+class NonAutoRegressiveQuantileRegression(Trainer):
+    def __init__(self, model: torch.nn.Module, optimizer: torch.optim.Optimizer, data_processor: DataProcessor, quantiles: list, device: torch.device, clearml_helper: ClearMLHelper = None, debug: bool = True):
+        quantiles_tensor = torch.tensor(quantiles)
+        quantiles_tensor = quantiles_tensor.to(device)
+        self.quantiles = quantiles
+
+        criterion = NonAutoRegressivePinballLoss(quantiles=quantiles_tensor)
+        super().__init__(model=model, optimizer=optimizer, criterion=criterion, data_processor=data_processor, device=device, clearml_helper=clearml_helper, debug=debug)
+
+    @staticmethod
+    def sample_from_dist(quantiles, output_values):
+        reshaped_values = output_values.reshape(-1, len(quantiles))
+        samples = []
+        for row in reshaped_values:
+            inverse_cdf = interp1d(quantiles, row, kind='linear', bounds_error=False, fill_value="extrapolate")
+            uniform_random_numbers = np.random.uniform(0, 1, 1000)
+            new_samples = inverse_cdf(uniform_random_numbers)
+            samples.append(np.mean(new_samples))
+        return np.array(samples)
+        
+    def log_final_metrics(self, task, dataloader, train: bool = True):
+        metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }
+        transformed_metrics = { metric.__class__.__name__: 0.0 for metric in self.metrics_to_track }
+
+        with torch.no_grad():
+            for inputs, targets in dataloader:
+                inputs, targets = inputs.to(self.device), targets
+
+                
+                outputs = self.model(inputs)
+                outputted_samples = [self.sample_from_dist(self.quantiles.cpu(), output.cpu().numpy()) for output in outputs]
+                # to tensor
+                outputted_samples = torch.tensor(outputted_samples)
+
+                inversed_outputs = torch.tensor(self.data_processor.inverse_transform(outputted_samples))
+                inversed_inputs = torch.tensor(self.data_processor.inverse_transform(targets))
+
+                # set on same device
+                inversed_outputs = inversed_outputs.to(self.device)
+                inversed_inputs = inversed_inputs.to(self.device)
+                outputted_samples = outputted_samples.to(self.device)
+
+                for metric in self.metrics_to_track:
+                    transformed_metrics[metric.__class__.__name__] += metric(outputted_samples, targets.to(self.device))
+                    metrics[metric.__class__.__name__] += metric(inversed_outputs, inversed_inputs)
+
+            for metric in self.metrics_to_track:
+                metrics[metric.__class__.__name__] /= len(dataloader)
+                transformed_metrics[metric.__class__.__name__] /= len(dataloader)
+
+            for metric_name, metric_value in metrics.items():
+                if train:
+                    metric_name = f'train_{metric_name}'
+                else:
+                    metric_name = f'test_{metric_name}'
+                
+                task.get_logger().report_single_value(name=metric_name, value=metric_value)
+
+            for metric_name, metric_value in transformed_metrics.items():
+                if train:
+                    metric_name = f'train_transformed_{metric_name}'
+                else:
+                    metric_name = f'test_transformed_{metric_name}'
+
+                task.get_logger().report_single_value(name=metric_name, value=metric_value)
+
+    
+    def get_plot(self, current_day, next_day, predictions, show_legend: bool = True):
+        fig = go.Figure()
+
+        # Convert to numpy for plotting
+        current_day_np = current_day.view(-1).cpu().numpy()
+        next_day_np = next_day.view(-1).cpu().numpy()
+        
+        # reshape predictions to (n, len(quantiles))$
+        predictions_np = predictions.cpu().numpy().reshape(-1, len(self.quantiles))
+
+        # Add traces for current and next day
+        fig.add_trace(go.Scatter(x=np.arange(96), y=current_day_np, name="Current Day"))
+        fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))
+
+        for i, q in enumerate(self.quantiles):
+            fig.add_trace(go.Scatter(x=96 + np.arange(96), y=predictions_np[:, i], 
+                                    name=f"Prediction (Q={q})", line=dict(dash='dash')))
+
+        # Update the layout
+        fig.update_layout(title="Predictions and Quantiles", showlegend=show_legend)
+        
+        return fig
+

src/trainers/trainer.py

@@ -67,6 +67,8 @@ class Trainer:
         task.connect(self.criterion, name="criterion")
         task.connect(self.data_processor, name="data_processor")
 
+        task.connect(self.data_processor.data_config, name="data_features")
+
         return task
     
     def random_samples(self, train: bool = True, num_samples: int = 10):
@@ -82,58 +84,68 @@ class Trainer:
 
 
     def train(self, epochs: int):
-        train_loader, test_loader = self.data_processor.get_dataloaders(predict_sequence_length=self.model.output_size)
+        try:
+            train_loader, test_loader = self.data_processor.get_dataloaders(predict_sequence_length=self.model.output_size)
 
-        train_samples = self.random_samples(train=True)
-        test_samples = self.random_samples(train=False)
+            train_samples = self.random_samples(train=True)
+            test_samples = self.random_samples(train=False)
 
-        task = self.init_clearml_task()
+            task = self.init_clearml_task()
 
-        self.best_score = None
-        counter = 0
+            self.best_score = None
+            counter = 0
 
-        for epoch in range(1, epochs + 1):
-            self.model.train()
-            running_loss = 0.0
+            for epoch in range(1, epochs + 1):
+                self.model.train()
+                running_loss = 0.0
 
-            for inputs, targets in train_loader:
-                inputs, targets = inputs.to(self.device), targets.to(self.device)
+                for inputs, targets in train_loader:
+                    inputs, targets = inputs.to(self.device), targets.to(self.device)
 
-                self.optimizer.zero_grad()
-                output = self.model(inputs)
+                    self.optimizer.zero_grad()
+                    output = self.model(inputs)
 
-                loss = self.criterion(output, targets)
-                loss.backward()
-                self.optimizer.step()
+                    loss = self.criterion(output, targets)
+                    loss.backward()
+                    self.optimizer.step()
 
-                running_loss += loss.item()
-                
+                    running_loss += loss.item()
+                    
 
-            running_loss /= len(train_loader.dataset)
-            test_loss = self.test(test_loader)
+                running_loss /= len(train_loader.dataset)
+                test_loss = self.test(test_loader)
 
-            if self.patience is not None:
-                if self.best_score is None or test_loss < self.best_score + self.delta:
-                    self.save_checkpoint(test_loss, task, epoch)
-                    counter = 0
-                else:
-                    counter += 1
-                    if counter >= self.patience:
-                        print('Early stopping triggered')
-                        break
+                if self.patience is not None:
+                    if self.best_score is None or test_loss < self.best_score + self.delta:
+                        self.save_checkpoint(test_loss, task, epoch)
+                        counter = 0
+                    else:
+                        counter += 1
+                        if counter >= self.patience:
+                            print('Early stopping triggered')
+                            break
+
+                if task:
+                    task.get_logger().report_scalar(title=self.criterion.__class__.__name__, series="train", value=running_loss, iteration=epoch)
+                    task.get_logger().report_scalar(title=self.criterion.__class__.__name__, series="test", value=test_loss, iteration=epoch)
+                    
+
+                    if epoch % self.plot_every_n_epochs == 0:
+                        self.debug_plots(task, True, train_loader, train_samples, epoch)
+                        self.debug_plots(task, False, test_loader, test_samples, epoch)
+
+                        if hasattr(self, 'plot_quantile_percentages'):
+                            self.plot_quantile_percentages(task, train_loader, True, epoch)
+                            self.plot_quantile_percentages(task, test_loader, False, epoch)
 
             if task:
-                task.get_logger().report_scalar(title=self.criterion.__class__.__name__, series="train", value=running_loss, iteration=epoch)
-                task.get_logger().report_scalar(title=self.criterion.__class__.__name__, series="test", value=test_loss, iteration=epoch)
-                
-
-                if epoch % self.plot_every_n_epochs == 0:
-                    self.debug_plots(task, True, train_loader, train_samples, epoch)
-                    self.debug_plots(task, False, test_loader, test_samples, epoch)
-
-        if task:
-            self.finish_training(task=task)
-            task.close()
+                self.finish_training(task=task)
+                task.close()
+        except Exception:
+            if task:
+                task.close()
+                task.set_archived(True)
+            raise
 
 
     def log_final_metrics(self, task, dataloader, train: bool = True):
@@ -178,10 +190,12 @@ class Trainer:
             self.model.load_state_dict(torch.load('checkpoint.pt'))
         self.model.eval()
 
-        transformed_train_loader, transformed_test_loader = self.data_processor.get_dataloaders(predict_sequence_length=self.model.output_size)
+        train_loader, test_loader = self.data_processor.get_dataloaders(predict_sequence_length=self.model.output_size)
 
-        # self.log_final_metrics(task, transformed_train_loader, train=True)
-        self.log_final_metrics(task, transformed_test_loader, train=False)
+        if not hasattr(self, 'plot_quantile_percentages'):
+            self.log_final_metrics(task, train_loader, train=True)
+
+        self.log_final_metrics(task, test_loader, train=False)
 
 
     def test(self, test_loader: torch.utils.data.DataLoader):
@@ -242,9 +256,9 @@ class Trainer:
                 fig.add_trace(trace, row=row, col=col)
 
 
-            loss = self.criterion(predictions.to(self.device), target.squeeze(-1).to(self.device)).item()
+            # 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}")
+            # 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)