Added trainer for Diffusion model

src/policies/simple_baseline.py

@@ -2,6 +2,7 @@ import pandas as pd
 import numpy as np
 import itertools
 from tqdm import tqdm
+import torch
 
 imbalance_prices = "data/imbalance_prices.csv"
 
@@ -126,98 +127,65 @@ class BaselinePolicy():
 
         return df, df_test
     
-    def get_optimal_thresholds(self, imbalance_prices, charge_thresholds, discharge_thresholds, tqdm=True):
-        df = None
-
+    def get_optimal_thresholds(self, imbalance_prices, charge_thresholds, discharge_thresholds):
         threshold_pairs = itertools.product(charge_thresholds, discharge_thresholds)
         threshold_pairs = filter(lambda x: x[0] < x[1], threshold_pairs)
 
-        # disable tqdm if necessary
-        if tqdm:
-            threshold_pairs = tqdm(threshold_pairs)
+        threshold_pairs = list(threshold_pairs)
 
-        for charge_threshold, discharge_threshold in threshold_pairs:
-            self.battery.reset()
-            total_charging_cost = 0
-            total_discharging_profit = 0
-            number_of_charges = 0
-            number_of_discharges = 0
-            mean_charging_price = 0
-            mean_discharging_price = 0
+        charge_thresholds = torch.tensor([x[0] for x in threshold_pairs])
+        discharge_thresholds = torch.tensor([x[1] for x in threshold_pairs])
 
-            for index, price in enumerate(imbalance_prices):
-                if price < charge_threshold:
-                    total_charging_cost += self.battery.charge() * price
-                    mean_charging_price += price
-                    number_of_charges += 1
-                elif price > discharge_threshold:
-                    total_discharging_profit += self.battery.discharge() * price
-                    mean_discharging_price += price
-                    number_of_discharges += 1
+        next_day_charge_thresholds, next_day_discharge_thresholds = [], []
 
-            if number_of_charges == 0:
-                mean_charging_price = 0
-            else:
-                mean_charging_price /= number_of_charges
+        for ip in imbalance_prices:
+            new_imbalance_prices = ip.repeat(len(charge_thresholds), 1)
 
-            if number_of_discharges == 0:
-                mean_discharging_price = 0
-            else:
-                mean_discharging_price /= number_of_discharges
+            profits, charge_cycles = self.simulate(new_imbalance_prices, charge_thresholds, discharge_thresholds)
 
-            new_df = pd.DataFrame([[charge_threshold, discharge_threshold, total_charging_cost, total_discharging_profit, total_discharging_profit - total_charging_cost, self.battery.charge_cycles, mean_charging_price, mean_discharging_price]], columns=["Charge threshold", "Discharge threshold", "Charging Cost", "Discharging Profit", "Total Profit", "Charge cycles", "Mean charging price", "Mean discharging price"])
-            if df is None:
-                df = new_df
-            else:
-                df = pd.concat([df, new_df])
+            sorted_profits, sorted_indices = torch.sort(profits, descending=True)
 
-        df = df.sort_values(by=['Total Profit'], ascending=False)
+            # Reorder other tensors based on sorted indices
+            sorted_charge_thresholds = charge_thresholds[sorted_indices]
+            sorted_discharge_thresholds = discharge_thresholds[sorted_indices]
 
-        return df
-    
-    def simulate(self, imbalance_prices, charge_threshold, discharge_threshold):
-        self.battery.reset()
-        total_charging_cost = 0
-        total_discharging_profit = 0
-        number_of_charges = 0
-        number_of_discharges = 0
-        mean_charging_price = 0
-        mean_discharging_price = 0
+            # get the optimal thresholds
+            next_day_charge_threshold = sorted_charge_thresholds[0]
+            next_day_discharge_threshold = sorted_discharge_thresholds[0]
 
-        # for each timestep also print what is happening
-        for i, price in enumerate(imbalance_prices):
-            if price < charge_threshold:
-                charge = self.battery.charge()
-                total_charging_cost += charge * price
-                mean_charging_price += price
-                number_of_charges += 1
-                # print(f"{i}: Charging battery with {charge}MWh at {price}€/MWh")
-                
-            elif price > discharge_threshold:
-                discharge = self.battery.discharge()
-                total_discharging_profit += discharge * price
-                mean_discharging_price += price
-                number_of_discharges += 1
-                # print(f"{i}: Discharging battery with {discharge}MWh at {price}€/MWh")
+            next_day_charge_thresholds.append(next_day_charge_threshold)
+            next_day_discharge_thresholds.append(next_day_discharge_threshold)
 
-        if number_of_charges == 0:
-            mean_charging_price = 0
-        else:
-            mean_charging_price /= number_of_charges
+        # return float tensors
+        return torch.tensor(next_day_charge_thresholds, dtype=torch.float), torch.tensor(next_day_discharge_thresholds, dtype=torch.float)
 
-        if number_of_discharges == 0:
-            mean_discharging_price = 0
-        else:
-            mean_discharging_price /= number_of_discharges
+    def simulate(self, price_matrix, charge_thresholds: torch.tensor, discharge_thresholds: torch.tensor):
+        batch_size = price_matrix.shape[0]
 
-        # print(f"Total charging cost: {total_charging_cost}")
-        # print(f"Total discharging profit: {total_discharging_profit}")
-        # print(f"Total profit: {total_discharging_profit - total_charging_cost}")
-        # print(f"Charge cycles: {self.battery.charge_cycles}")
-        # print(f"Mean charging price: {mean_charging_price}")
-        # print(f"Mean discharging price: {mean_discharging_price}")
+        charge_matrix = torch.zeros(price_matrix.shape)
 
-        return total_discharging_profit - total_charging_cost, self.battery.charge_cycles
+        charge_thresholds_reshaped = charge_thresholds.repeat(price_matrix.shape[1], 1).T
+        discharge_thresholds_reshaped = discharge_thresholds.repeat(price_matrix.shape[1], 1).T
+
+        charge_matrix[price_matrix < charge_thresholds_reshaped] = 1
+        
+        charge_matrix[price_matrix > discharge_thresholds_reshaped] = -1
+
+        battery_states = torch.zeros(batch_size)
+        profits = torch.zeros(batch_size)
+        charge_cycles = torch.zeros(batch_size)
+
+        for i in range(price_matrix.shape[1]):
+            discharge_mask = ~((charge_matrix[:, i] == -1) & (battery_states == 0))
+            charge_mask = ~((charge_matrix[:, i] == 1) & (battery_states == self.battery.capacity))
+
+            mask = discharge_mask & charge_mask
+
+            battery_states[mask] += charge_matrix[:, i][mask] * self.battery.power / 4
+            profits[mask] += -charge_matrix[:, i][mask] * price_matrix[:, i][mask] * self.battery.power / 4
+            charge_cycles[mask] += torch.abs(charge_matrix[:, i][mask]) * (self.battery.power / 4) / self.battery.capacity
+
+        return profits, charge_cycles