Merge branch 'February-Report' into main

src/data/dataset.py

@@ -25,12 +25,19 @@ class NrvDataset(Dataset):
         self.sequence_length = sequence_length
         self.predict_sequence_length = predict_sequence_length
 
-        self.samples_to_skip = self.skip_samples(dataframe=dataframe)
+        self.samples_to_skip = self.skip_samples(dataframe=dataframe, full_day_skip=self.full_day_skip)
         total_indices = set(
             range(len(dataframe) - self.sequence_length - self.predict_sequence_length)
         )
         self.valid_indices = sorted(list(total_indices - set(self.samples_to_skip)))
 
+        # full day indices
+        full_day_skipped_samples = self.skip_samples(dataframe=dataframe, full_day_skip=True)
+        full_day_total_indices = set(
+            range(len(dataframe) - self.sequence_length - self.predict_sequence_length)
+        )
+        self.full_day_valid_indices = sorted(list(full_day_total_indices - set(full_day_skipped_samples)))
+
         self.history_features = []
         if self.data_config.LOAD_HISTORY:
             self.history_features.append("total_load")
@@ -73,7 +80,7 @@ class NrvDataset(Dataset):
 
         self.history_features, self.forecast_features = self.preprocess_data(dataframe)
 
-    def skip_samples(self, dataframe):
+    def skip_samples(self, dataframe, full_day_skip):
         nan_rows = dataframe[dataframe.isnull().any(axis=1)]
         nan_indices = nan_rows.index
         skip_indices = [
@@ -91,7 +98,7 @@ class NrvDataset(Dataset):
 
         # add indices that are not the start of a day (00:15) to the skip indices (use datetime column)
         # get indices of all 00:15 timestamps
-        if self.full_day_skip:
+        if full_day_skip:
             start_of_day_indices = dataframe[
                 dataframe["datetime"].dt.time != pd.Timestamp("00:00:00").time()
             ].index

src/policies/policy_executer.py

@@ -8,7 +8,8 @@ import pandas as pd
 import datetime
 from tqdm import tqdm
 from src.utils.imbalance_price_calculator import ImbalancePriceCalculator
-import time
+import seaborn as sns
+import matplotlib.pyplot as plt
 import plotly.express as px
 
 ### import functions ###
@@ -16,7 +17,7 @@ from src.trainers.quantile_trainer import auto_regressive as quantile_auto_regre
 from src.trainers.diffusion_trainer import sample_diffusion
 from src.utils.clearml import ClearMLHelper
 
-# argparse to parse task id and model type
+### Arguments ###
 parser = argparse.ArgumentParser()
 parser.add_argument('--task_id', type=str, default=None)
 parser.add_argument('--model_type', type=str, default=None)
@@ -27,6 +28,7 @@ assert args.task_id is not None, "Please specify task id"
 assert args.model_type is not None, "Please specify model type"
 assert args.model_name is not None, "Please specify model name"
 
+### Baseline Policy ###
 battery = Battery(2, 1)
 baseline_policy = BaselinePolicy(battery, data_path="")
 
@@ -163,20 +165,17 @@ def get_next_day_profits_for_date(model, data_processor, test_loader, date, ipc,
     return predicted_nrv_profits_cycles, baseline_profits_cycles, _charge_thresholds, _discharge_thresholds
 
 def next_day_test_set(model, data_processor, test_loader, ipc, predict_NRV: callable):
-    penalties = [0, 10, 50, 150, 300, 500, 600, 800, 1000, 1500, 2000, 2500]
+    penalties = [0, 50, 250, 500, 1000, 1500]
     predicted_nrv_profits_cycles = {i: [0, 0] for i in penalties}
     baseline_profits_cycles = {i: [0, 0] for i in penalties}
 
     charge_thresholds = {}
     discharge_thresholds = {}
 
-    # get all dates in test set
     dates = baseline_policy.test_data["DateTime"].dt.date.unique()
-
-    # dates back to datetime
     dates = pd.to_datetime(dates)
 
-    for date in tqdm(dates[:10]):
+    for date in tqdm(dates):
         try:
             new_predicted_nrv_profits_cycles, new_baseline_profits_cycles, new_charge_thresholds, new_discharge_thresholds = get_next_day_profits_for_date(model, data_processor, test_loader, date, ipc, predict_NRV, penalties)
             
@@ -191,8 +190,7 @@ def next_day_test_set(model, data_processor, test_loader, ipc, predict_NRV: call
                 baseline_profits_cycles[penalty][1] += new_baseline_profits_cycles[penalty][1]
 
         except Exception as e:
-            # print(f"Error for date {date}")
+            print(f"Error for date {date}")
-            raise e
 
     return predicted_nrv_profits_cycles, baseline_profits_cycles, charge_thresholds, discharge_thresholds
 
@@ -222,9 +220,6 @@ def main():
     # the charge_thresholds is a dictionary with date as key. The values of the dictionary is another dictionary with keys as penalties and values as the charge thresholds
     # create density plot that shows a density plot of the charge thresholds for each penalty (use seaborn displot) (One plot with a different color for each penalty)
 
-    import seaborn as sns
-    import matplotlib.pyplot as plt
-
     charge_thresholds_for_penalty = {}
     for d in charge_thresholds.values():
         for penalty, thresholds in d.items():
@@ -239,47 +234,73 @@ def main():
                 discharge_thresholds_for_penalty[penalty] = []
             discharge_thresholds_for_penalty[penalty].extend(thresholds)
 
+    def plot_threshold_distribution(thresholds: dict, title: str):
+        data_to_plot = []
+        for penalty, values in thresholds.items():
+            for value in values:
+                data_to_plot.append({'Penalty': penalty, 'Value': value.item()})
+        df = pd.DataFrame(data_to_plot)
+        palette = sns.color_palette("bright", len(thresholds.keys()))
+        fig = sns.displot(data=df, x="Value", hue="Penalty", kind="kde", palette=palette)
+        plt.title('Density of Charge Thresholds by Penalty')
+        plt.xlabel('Charge Threshold')
+        plt.ylabel('Density')
+        plt.legend(title='Penalty')
+        task.get_logger().report_matplotlib_figure(
+            "Policy Results", 
+            title, 
+            iteration=0, 
+            figure=fig
+        )
+        plt.close()
+
     ### Plot charge thresholds distribution ###
-    data_to_plot = []
+    plot_threshold_distribution(charge_thresholds_for_penalty, "Charge Thresholds")
-    for penalty, values in charge_thresholds_for_penalty.items():
-        for value in values:
-            data_to_plot.append({'Penalty': penalty, 'Value': value.item()})
-    df = pd.DataFrame(data_to_plot)
-    print(df.head())
-    palette = sns.color_palette("bright", len(charge_thresholds.keys()))
-    fig = sns.displot(data=df, x="Value", hue="Penalty", kind="kde", palette=palette)
-    plt.title('Density of Charge Thresholds by Penalty')
-    plt.xlabel('Charge Threshold')
-    plt.ylabel('Density')
-    plt.legend(title='Penalty')
-    task.get_logger().report_matplotlib_figure(
-        "Policy Results", 
-        "Charge Thresholds", 
-        iteration=0, 
-        figure=fig
-    )
-    plt.close()
 
     ### Plot discharge thresholds distribution ###
-    data_to_plot = []
+    plot_threshold_distribution(discharge_thresholds_for_penalty, "Discharge Thresholds")
-    for penalty, values in discharge_thresholds_for_penalty.items():
-        for value in values:
-            data_to_plot.append({'Penalty': penalty, 'Value': value.item()})
-    df = pd.DataFrame(data_to_plot)
-    palette = sns.color_palette("bright", len(discharge_thresholds.keys()))
-    fig = sns.displot(data=df, x="Value", hue="Penalty", kind="kde", palette=palette)
-    plt.title('Density of Charge Thresholds by Penalty')
-    plt.xlabel('Charge Threshold')
-    plt.ylabel('Density')
-    plt.legend(title='Penalty')
-    task.get_logger().report_matplotlib_figure(
-        "Policy Results", 
-        "Discharge Thresholds", 
-        iteration=0, 
-        figure=fig
-    )
-    plt.close()
 
+    def plot_thresholds_per_day(thresholds: dict, title: str):
+        # plot mean charge threshold per day (per penalty (other color))
+        data_to_plot = []
+        for date, values in thresholds.items():
+            for penalty, value in values.items():
+                mean_val = value.mean().item()
+                std_val = value.std().item()  # Calculate standard deviation
+                data_to_plot.append({'Date': date, 'Penalty': penalty, 'Mean': mean_val, 'StdDev': std_val})
+                print(f"Date: {date}, Penalty: {penalty}, Mean: {mean_val}, StdDev: {std_val}")
+        df = pd.DataFrame(data_to_plot)
+        df["Date"] = pd.to_datetime(df["Date"])
+
+        fig = px.line(
+            df,
+            x="Date",
+            y="Mean",
+            color="Penalty",
+            title=title,
+            labels={"Mean": "Threshold", "Date": "Date"},
+            markers=True,  # Adds markers to the lines
+            hover_data=["Penalty"],  # Adds additional hover information
+        )
+
+        fig.update_layout(
+            width=1000,  # Set the width of the figure
+            height=600,  # Set the height of the figure
+            title_x=0.5,  # Center the title horizontally
+        )
+
+        task.get_logger().report_plotly(
+            "Thresholds per Day", 
+            title, 
+            iteration=0, 
+            figure=fig
+        )
+
+    ### Plot mean charge thresholds per day ###
+    plot_thresholds_per_day(charge_thresholds, "Mean Charge Thresholds per Day")
+
+    ### Plot mean discharge thresholds per day ###
+    plot_thresholds_per_day(discharge_thresholds, "Mean Discharge Thresholds per Day")
 
 
     # create dataframe with columns "name", "penalty", "profit", "cycles"

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)
+        for actual_idx, idx in sample_indices.items():
-        rows = num_samples  # One row per sample since we only want one column
+            auto_regressive_output = self.auto_regressive(data_loader.dataset, [idx]*1000)
-
-        #  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])
             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)
+            # keep one initial
-            target = target.squeeze(0)
+            initial = initial[0]
+            target = target[0]
 
-            sub_fig = self.get_plot(initial, target, predictions, show_legend=(i == 0))
+            predictions = predictions
 
-            row = i + 1
+            fig = self.get_plot(initial, target, predictions, show_legend=(0 == 0))
-            col = 1
 
-            for trace in sub_fig.data:
+            task.get_logger().report_matplotlib_figure(
-                fig.add_trace(trace, row=row, col=col)
+                title="Training" if train else "Testing",
-
+                series=f'Sample {actual_idx}',
-            if cols == 2:
+                iteration=epoch,
-                error_sub_fig = self.get_plot_error(
+                figure=fig,
-                    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}"
             )
 
-        # 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()

src/trainers/diffusion_trainer.py

@@ -19,8 +19,6 @@ def sample_diffusion(model: DiffusionModel, n: int, inputs: torch.tensor, noise_
     alpha = 1. - beta
     alpha_hat = torch.cumprod(alpha, dim=0)
 
-    # inputs: (num_features) -> (batch_size, num_features)
-    # inputs: (time_steps, num_features) -> (batch_size, time_steps, num_features)
     if len(inputs.shape) == 2:
         inputs = inputs.repeat(n, 1)
     elif len(inputs.shape) == 3:
@@ -42,17 +40,17 @@ def sample_diffusion(model: DiffusionModel, n: int, inputs: torch.tensor, noise_
                 noise = torch.zeros_like(x)
 
             x = 1/torch.sqrt(_alpha) * (x-((1-_alpha) / (torch.sqrt(1 - _alpha_hat))) * predicted_noise) + torch.sqrt(_beta) * noise
+    x = torch.clamp(x, -1.0, 1.0)
     return x
 
 
-
 class DiffusionTrainer:
     def __init__(self, model: nn.Module, data_processor: DataProcessor, device: torch.device):
         self.model = model
         self.device = device
 
-        self.noise_steps = 20
+        self.noise_steps = 30
-        self.beta_start = 1e-4
+        self.beta_start = 0.0001
         self.beta_end = 0.02
         self.ts_length = 96
         
@@ -98,7 +96,16 @@ class DiffusionTrainer:
         else:
             loader = test_loader
 
-        indices = np.random.randint(0, len(loader.dataset) - 1, size=num_samples)
+        # set seed
+        np.random.seed(42)
+
+        actual_indices = np.random.choice(loader.dataset.full_day_valid_indices, num_samples, replace=False)
+        indices = {}
+        for i in actual_indices:
+            indices[i] = loader.dataset.valid_indices.index(i)
+
+        print(actual_indices)
+        
         return indices
     
     def init_clearml_task(self, task):
@@ -173,7 +180,7 @@ class DiffusionTrainer:
 
         
     def debug_plots(self, task, training: bool, data_loader, sample_indices, epoch):
-        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)
@@ -182,18 +189,21 @@ class DiffusionTrainer:
             self.model.eval()
             with torch.no_grad():
                 samples = self.sample(self.model, 100, features).cpu().numpy()
+                samples = self.data_processor.inverse_transform(samples)
+                target = self.data_processor.inverse_transform(target)
                
-            ci_99_upper = np.quantile(samples, 0.99, axis=0)
+            ci_99_upper = np.quantile(samples, 0.995, axis=0)
-            ci_99_lower = np.quantile(samples, 0.01, axis=0)
+            ci_99_lower = np.quantile(samples, 0.005, axis=0)
 
-            ci_95_upper = np.quantile(samples, 0.95, axis=0)
+            ci_95_upper = np.quantile(samples, 0.975, axis=0)
-            ci_95_lower = np.quantile(samples, 0.05, axis=0)
+            ci_95_lower = np.quantile(samples, 0.025, axis=0)
 
-            ci_90_upper = np.quantile(samples, 0.9, axis=0)
+            ci_90_upper = np.quantile(samples, 0.95, axis=0)
-            ci_90_lower = np.quantile(samples, 0.1, axis=0)
+            ci_90_lower = np.quantile(samples, 0.05, axis=0)
+
+            ci_50_lower = np.quantile(samples, 0.25, axis=0)
+            ci_50_upper = np.quantile(samples, 0.75, axis=0)
 
-            ci_50_upper = np.quantile(samples, 0.5, axis=0)
-            ci_50_lower = np.quantile(samples, 0.5, axis=0)
 
             sns.set_theme()
             time_steps = np.arange(0, 96)
@@ -219,7 +229,7 @@ class DiffusionTrainer:
 
             task.get_logger().report_matplotlib_figure(
                 title="Training" if training else "Testing",
-                series=f'Sample {i}',
+                series=f'Sample {actual_idx}',
                 iteration=epoch,
                 figure=fig,
             )

src/trainers/quantile_trainer.py

@@ -10,7 +10,9 @@ import plotly.graph_objects as go
 import numpy as np
 import matplotlib.pyplot as plt
 from scipy.interpolate import CubicSpline
-
+import matplotlib.pyplot as plt
+import seaborn as sns
+import matplotlib.patches as mpatches
 
 def sample_from_dist(quantiles, preds):
     if isinstance(preds, torch.Tensor):
@@ -261,35 +263,35 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
                     name="test_CRPS_from_samples_transformed", value=np.mean(crps_from_samples_metric)
                 )
 
-    def get_plot_error(
+    # def get_plot_error(
-        self,
+    #     self,
-        next_day,
+    #     next_day,
-        predictions,
+    #     predictions,
-    ):
+    # ):
-        metric = PinballLoss(quantiles=self.quantiles)
+    #     metric = PinballLoss(quantiles=self.quantiles)
-        fig = go.Figure()
+    #     fig = go.Figure()
 
-        next_day_np = next_day.view(-1).cpu().numpy()
+    #     next_day_np = next_day.view(-1).cpu().numpy()
-        predictions_np = predictions.cpu().numpy()
+    #     predictions_np = predictions.cpu().numpy()
 
-        if True:
+    #     if True:
-            next_day_np = self.data_processor.inverse_transform(next_day_np)
+    #         next_day_np = self.data_processor.inverse_transform(next_day_np)
-            predictions_np = self.data_processor.inverse_transform(predictions_np)
+    #         predictions_np = self.data_processor.inverse_transform(predictions_np)
 
-        # for each time step, calculate the error using the metric
+    #     # for each time step, calculate the error using the metric
-        errors = []
+    #     errors = []
-        for i in range(96):
+    #     for i in range(96):
 
-            target_tensor = torch.tensor(next_day_np[i]).unsqueeze(0)
+    #         target_tensor = torch.tensor(next_day_np[i]).unsqueeze(0)
-            prediction_tensor = torch.tensor(predictions_np[i]).unsqueeze(0)
+    #         prediction_tensor = torch.tensor(predictions_np[i]).unsqueeze(0)
 
-            errors.append(metric(prediction_tensor, target_tensor))
+    #         errors.append(metric(prediction_tensor, target_tensor))
 
-        # plot the error 
+    #     # plot the error 
-        fig.add_trace(go.Scatter(x=np.arange(96), y=errors, name=metric.__class__.__name__))
+    #     fig.add_trace(go.Scatter(x=np.arange(96), y=errors, name=metric.__class__.__name__))
-        fig.update_layout(title=f"Error of {metric.__class__.__name__} for each time step")
+    #     fig.update_layout(title=f"Error of {metric.__class__.__name__} for each time step")
 
-        return fig
+    #     return fig
 
 
     def get_plot(
@@ -312,26 +314,59 @@ class AutoRegressiveQuantileTrainer(AutoRegressiveTrainer):
             next_day_np = self.data_processor.inverse_transform(next_day_np)
             predictions_np = self.data_processor.inverse_transform(predictions_np)
 
+        ci_99_upper = np.quantile(predictions_np, 0.995, axis=0)
+        ci_99_lower = np.quantile(predictions_np, 0.005, axis=0)
+
+        ci_95_upper = np.quantile(predictions_np, 0.975, axis=0)
+        ci_95_lower = np.quantile(predictions_np, 0.025, axis=0)
+
+        ci_90_upper = np.quantile(predictions_np, 0.95, axis=0)
+        ci_90_lower = np.quantile(predictions_np, 0.05, axis=0)
+
+        ci_50_lower = np.quantile(predictions_np, 0.25, axis=0)
+        ci_50_upper = np.quantile(predictions_np, 0.75, axis=0)
+
         # 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=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"))
+        # fig.add_trace(go.Scatter(x=96 + np.arange(96), y=next_day_np, name="Next Day"))
 
-        for i, q in enumerate(self.quantiles):
+        # for i, q in enumerate(self.quantiles):
-            fig.add_trace(
+        #     fig.add_trace(
-                go.Scatter(
+        #         go.Scatter(
-                    x=96 + np.arange(96),
+        #             x=96 + np.arange(96),
-                    y=predictions_np[:, i],
+        #             y=predictions_np[:, i],
-                    name=f"Prediction (Q={q})",
+        #             name=f"Prediction (Q={q})",
-                    line=dict(dash="dash"),
+        #             line=dict(dash="dash"),
-                )
+        #         )
-            )
+        #     )
 
-        # Update the layout
+        # # Update the layout
-        fig.update_layout(
+        # fig.update_layout(
-            title="Predictions and Quantiles of the Linear Model",
+        #     title="Predictions and Quantiles of the Linear Model",
-            showlegend=show_legend,
+        #     showlegend=show_legend,
-        )
+        # )
 
+        sns.set_theme()
+        time_steps = np.arange(0, 96)
+
+        fig, ax = plt.subplots(figsize=(20, 10))
+        ax.plot(time_steps, predictions_np.mean(axis=0), label="Mean of NRV samples", linewidth=3)
+        # ax.fill_between(time_steps, ci_lower, ci_upper, color='b', alpha=0.2, label='Full Interval')
+
+        ax.fill_between(time_steps, ci_99_lower, ci_99_upper, color='b', alpha=0.2, label='99% Interval')
+        ax.fill_between(time_steps, ci_95_lower, ci_95_upper, color='b', alpha=0.2, label='95% Interval')
+        ax.fill_between(time_steps, ci_90_lower, ci_90_upper, color='b', alpha=0.2, label='90% Interval')
+        ax.fill_between(time_steps, ci_50_lower, ci_50_upper, color='b', alpha=0.2, label='50% Interval')
+
+        ax.plot(next_day_np, label="Real NRV", linewidth=3)
+        # full_interval_patch = mpatches.Patch(color='b', alpha=0.2, label='Full Interval')
+        ci_99_patch = mpatches.Patch(color='b', alpha=0.3, label='99% Interval')
+        ci_95_patch = mpatches.Patch(color='b', alpha=0.4, label='95% Interval')
+        ci_90_patch = mpatches.Patch(color='b', alpha=0.5, label='90% Interval')
+        ci_50_patch = mpatches.Patch(color='b', alpha=0.6, label='50% Interval')
+
+
+        ax.legend(handles=[ci_99_patch, ci_95_patch, ci_90_patch, ci_50_patch, ax.lines[0], ax.lines[1]])
         return fig
 
     def auto_regressive(self, dataset, idx_batch, sequence_length: int = 96):

src/trainers/trainer.py

@@ -86,7 +86,7 @@ class Trainer:
 
     def random_samples(self, train: bool = True, num_samples: int = 10):
         train_loader, test_loader = self.data_processor.get_dataloaders(
-            predict_sequence_length=self.model.output_size
+            predict_sequence_length=96
         )
 
         if train:
@@ -94,7 +94,14 @@ class Trainer:
         else:
             loader = test_loader
 
-        indices = np.random.randint(0, len(loader.dataset) - 1, size=num_samples)
+        np.random.seed(42)
+        actual_indices = np.random.choice(loader.dataset.full_day_valid_indices, num_samples, replace=False)
+        indices = {}
+        for i in actual_indices:
+            indices[i] = loader.dataset.valid_indices.index(i)
+
+        print(actual_indices)
+        
         return indices
 
     def train(self, epochs: int, remotely: bool = False, task: Task = None):
@@ -107,8 +114,8 @@ class Trainer:
                 predict_sequence_length=self.model.output_size
             )
 
-            train_samples = self.random_samples(train=True)
+            train_samples = self.random_samples(train=True, num_samples=5)
-            test_samples = self.random_samples(train=False)
+            test_samples = self.random_samples(train=False, num_samples=5)
 
             self.init_clearml_task(task)
 

src/training_scripts/diffusion_training.py

@@ -38,7 +38,7 @@ data_config.NOMINAL_NET_POSITION = True
 data_config = task.connect(data_config, name="data_features")
 
 data_processor = DataProcessor(data_config, path="", lstm=False)
-data_processor.set_batch_size(128)
+data_processor.set_batch_size(64)
 data_processor.set_full_day_skip(True)
 
 inputDim = data_processor.get_input_size()
@@ -47,15 +47,15 @@ print("Input dim: ", inputDim)
 model_parameters = {
     "epochs": 5000,
     "learning_rate": 0.0001,
-    "hidden_sizes": [512, 512, 512],
+    "hidden_sizes": [128, 128],
-    "time_dim": 64,
+    "time_dim": 8,
 }
 
 model_parameters = task.connect(model_parameters, name="model_parameters")
 
 #### Model ####
-# model = SimpleDiffusionModel(96, model_parameters["hidden_sizes"], other_inputs_dim=inputDim[1], time_dim=model_parameters["time_dim"])
+model = SimpleDiffusionModel(96, model_parameters["hidden_sizes"], other_inputs_dim=inputDim[1], time_dim=model_parameters["time_dim"])
-model = GRUDiffusionModel(96, model_parameters["hidden_sizes"], other_inputs_dim=inputDim[2], time_dim=model_parameters["time_dim"], gru_hidden_size=256)
+# model = GRUDiffusionModel(96, model_parameters["hidden_sizes"], other_inputs_dim=inputDim[2], time_dim=model_parameters["time_dim"], gru_hidden_size=128)
 
 print("Starting training ...")