Upload 10 files

Data_generation_tool_kit/Hier_diffusion_energy/global_scaler.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b56b496f31ec90c863f0841bc8dd9c60d990662d93093a1dad427c012a721c2f
+size 477

Data_generation_tool_kit/Hier_diffusion_energy/hierarchial_diffusion_model.py

@@ -0,0 +1,384 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import List, Optional, Dict
+from tqdm import tqdm
+
+
+class SinusoidalPositionEmbeddings(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, time: torch.Tensor) -> torch.Tensor:
+        device = time.device
+        half_dim = self.dim // 2
+        embeddings = math.log(10000) / (half_dim - 1)
+        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
+        embeddings = time[:, None] * embeddings[None, :]
+        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
+        return embeddings
+
+
+class ResnetBlock1D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, *, time_emb_dim: int = None, dropout: float = 0.1):
+        super().__init__()
+        self.time_mlp = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(time_emb_dim, out_channels * 2)
+        ) if time_emb_dim is not None else None
+
+        self.block1_conv = nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.block1_norm = nn.GroupNorm(8, out_channels, affine=False)
+        self.block1_act = nn.SiLU()
+
+        self.block2_conv = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)
+        self.block2_norm = nn.GroupNorm(8, out_channels)
+        self.block2_act = nn.SiLU()
+        self.block2_dropout = nn.Dropout(dropout)
+
+        self.res_conv = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
+
+    def forward(self, x: torch.Tensor, time_emb: torch.Tensor = None) -> torch.Tensor:
+        h = self.block1_conv(x)
+        h = self.block1_norm(h)
+
+        if self.time_mlp is not None and time_emb is not None:
+            scale_shift = self.time_mlp(time_emb)
+            scale, shift = scale_shift.chunk(2, dim=1)
+            h = h * (scale.unsqueeze(-1) + 1) + shift.unsqueeze(-1)
+
+        h = self.block1_act(h)
+
+        h = self.block2_act(self.block2_norm(self.block2_conv(h)))
+        h = self.block2_dropout(h)
+        return h + self.res_conv(x)
+
+
+class AttentionBlock1D(nn.Module):
+    def __init__(self, channels: int, num_heads: int = 8):
+        super().__init__()
+        self.channels = channels
+        self.num_heads = num_heads
+        assert channels % num_heads == 0, "channels must be divisible by num_heads"
+        self.head_dim = channels // num_heads
+        
+        self.norm = nn.GroupNorm(8, channels)
+        self.qkv = nn.Conv1d(channels, channels * 3, 1)
+        self.proj = nn.Conv1d(channels, channels, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, C, L = x.shape
+        h = self.norm(x)
+        
+        qkv = self.qkv(h)
+        qkv = qkv.view(B, 3, self.num_heads, self.head_dim, L)
+        qkv = qkv.permute(1, 0, 2, 4, 3)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0)
+        
+        out = out.permute(0, 1, 3, 2)
+        out = out.contiguous().view(B, C, L)
+        
+        return x + self.proj(out)
+
+
+class DownBlock1D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, time_emb_dim: int, dropout: float, use_attention: bool, num_blocks: int = 2):
+        super().__init__()
+        self.resnets = nn.ModuleList([
+            ResnetBlock1D(in_channels if i == 0 else out_channels, out_channels, time_emb_dim=time_emb_dim, dropout=dropout)
+            for i in range(num_blocks)
+        ])
+        self.attn = AttentionBlock1D(out_channels) if use_attention else nn.Identity()
+        self.downsampler = nn.Conv1d(out_channels, out_channels, kernel_size=4, stride=2, padding=1)
+
+    def forward(self, x, time_emb):
+        for resnet in self.resnets:
+            x = resnet(x, time_emb)
+        x = self.attn(x)
+        skip = x
+        x = self.downsampler(x)
+        return x, skip
+
+
+class UpBlock1D(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, time_emb_dim: int, dropout: float, use_attention: bool, num_blocks: int = 2):
+        super().__init__()
+        self.resnets = nn.ModuleList()
+        self.resnets.append(ResnetBlock1D(in_channels * 2, out_channels, time_emb_dim=time_emb_dim, dropout=dropout))
+        for _ in range(num_blocks - 1):
+            self.resnets.append(ResnetBlock1D(out_channels, out_channels, time_emb_dim=time_emb_dim, dropout=dropout))
+        self.attn = AttentionBlock1D(out_channels) if use_attention else nn.Identity()
+        self.upsampler = nn.ConvTranspose1d(in_channels, in_channels, kernel_size=4, stride=2, padding=1)
+
+    def forward(self, x, skip_x, time_emb):
+        x = self.upsampler(x)
+        
+        if x.size(-1) != skip_x.size(-1):
+            diff_L = skip_x.size(-1) - x.size(-1)
+            if diff_L > 0:
+                x = F.pad(x, [diff_L // 2, diff_L - diff_L // 2])
+            elif diff_L < 0:
+                x = x[:, :, :skip_x.size(-1)]
+        
+        x = torch.cat([skip_x, x], dim=1)
+        
+        for resnet in self.resnets:
+            x = resnet(x, time_emb)
+        return self.attn(x)
+
+
+class ConditionalUnet(nn.Module):
+    def __init__(self, in_channels: int, num_houses: int, embedding_dim: int = 64, 
+                 hidden_dims: List[int] = [64, 128, 256], 
+                 dropout: float = 0.1, use_attention: bool = True, 
+                 cond_channels: int = 0, blocks_per_level: int = 2):
+        super().__init__()
+        time_emb_dim = hidden_dims[0] * 4
+
+        self.time_mlp = nn.Sequential(
+            SinusoidalPositionEmbeddings(hidden_dims[0]), 
+            nn.Linear(hidden_dims[0], time_emb_dim), 
+            nn.SiLU(), 
+            nn.Linear(time_emb_dim, time_emb_dim)
+        )
+        
+        self.house_embedding = nn.Embedding(num_houses, embedding_dim)
+        self.house_proj = nn.Linear(embedding_dim, time_emb_dim)
+
+        self.day_of_week_embedding = nn.Embedding(7, embedding_dim)
+        self.day_of_year_embedding = nn.Embedding(366, embedding_dim)
+        
+        self.day_of_week_proj = nn.Linear(embedding_dim, time_emb_dim)
+        self.day_of_year_proj = nn.Linear(embedding_dim, time_emb_dim)
+
+        self.init_conv = nn.Conv1d(in_channels + cond_channels, hidden_dims[0], kernel_size=7, padding=3)
+        
+        num_resolutions = len(hidden_dims)
+        self.down_blocks = nn.ModuleList([
+            DownBlock1D(hidden_dims[i], hidden_dims[i+1], time_emb_dim, dropout, use_attention, blocks_per_level)
+            for i in range(num_resolutions - 1)
+        ])
+        
+        self.mid_block1 = ResnetBlock1D(hidden_dims[-1], hidden_dims[-1], time_emb_dim=time_emb_dim, dropout=dropout)
+        self.mid_attn = AttentionBlock1D(hidden_dims[-1])
+        self.mid_block2 = ResnetBlock1D(hidden_dims[-1], hidden_dims[-1], time_emb_dim=time_emb_dim, dropout=dropout)
+        
+        self.up_blocks = nn.ModuleList([
+            UpBlock1D(hidden_dims[i+1], hidden_dims[i], time_emb_dim, dropout, use_attention, blocks_per_level)
+            for i in reversed(range(num_resolutions - 1))
+        ])
+        
+        self.final_conv = nn.Sequential(
+            ResnetBlock1D(hidden_dims[0], hidden_dims[0], time_emb_dim=time_emb_dim, dropout=dropout), 
+            nn.Conv1d(hidden_dims[0], in_channels, 1)
+        )
+
+    def forward(self, x: torch.Tensor, timestep: torch.Tensor, conditions: Dict[str, torch.Tensor], 
+                conditioning_signal: Optional[torch.Tensor] = None) -> torch.Tensor:
+        time_emb = self.time_mlp(timestep)
+        
+        house_id = conditions["house_id"]
+        day_of_week = conditions["day_of_week"]
+        day_of_year = conditions["day_of_year"]
+
+        house_emb = self.house_proj(self.house_embedding(house_id))
+        dow_emb = self.day_of_week_proj(self.day_of_week_embedding(day_of_week))
+        doy_emb = self.day_of_year_proj(self.day_of_year_embedding(day_of_year))
+        
+        emb = time_emb + house_emb + dow_emb + doy_emb
+
+        x = x.permute(0, 2, 1)
+        if conditioning_signal is not None:
+            x = torch.cat([x, conditioning_signal.permute(0, 2, 1)], dim=1)
+        
+        x = self.init_conv(x)
+        
+        skip_connections = []
+        for down_block in self.down_blocks:
+            x, skip_x = down_block(x, emb)
+            skip_connections.append(skip_x)
+        
+        x = self.mid_block1(x, emb)
+        x = self.mid_attn(x)
+        x = self.mid_block2(x, emb)
+        
+        for up_block in self.up_blocks:
+            x = up_block(x, skip_connections.pop(), emb)
+            
+        return self.final_conv(x).permute(0, 2, 1)
+
+
+class ImprovedDiffusionModel(nn.Module):
+    def __init__(self, base_model: ConditionalUnet, num_timesteps: int, channel_weights: torch.Tensor = None):
+        super().__init__()
+        self.model = base_model
+        self.num_timesteps = num_timesteps
+        self.channel_weights = channel_weights
+        
+        betas = self._cosine_beta_schedule(num_timesteps)
+        alphas = 1.0 - betas
+        alphas_cumprod = torch.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+        
+        self.register_buffer('betas', betas)
+        self.register_buffer('alphas', alphas)
+        self.register_buffer('alphas_cumprod', alphas_cumprod)
+        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev)
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1.0 - alphas_cumprod))
+        
+        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        posterior_variance = torch.clamp(posterior_variance, min=1e-20)
+        self.register_buffer('posterior_variance', posterior_variance)
+
+    def _cosine_beta_schedule(self, timesteps, s=0.008):
+        steps = timesteps + 1
+        x = torch.linspace(0, timesteps, steps, dtype=torch.float64)
+        alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi * 0.5) ** 2
+        alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+        betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+        return torch.clip(betas, 0.0001, 0.9999).float()
+
+    def q_sample(self, x_start, t, noise=None):
+        if noise is None: noise = torch.randn_like(x_start)
+        sqrt_alphas_cumprod_t = self.sqrt_alphas_cumprod[t].view(-1, 1, 1)
+        sqrt_one_minus_alphas_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1)
+        return sqrt_alphas_cumprod_t * x_start + sqrt_one_minus_alphas_cumprod_t * noise
+
+    def forward(self, x_0: torch.Tensor, conditions: Dict[str, torch.Tensor], 
+            conditioning_signal: Optional[torch.Tensor] = None) -> torch.Tensor:
+        t = torch.randint(0, self.num_timesteps, (x_0.shape[0],), device=x_0.device).long()
+        noise = torch.randn_like(x_0)
+        x_t = self.q_sample(x_0, t, noise)
+        predicted_noise = self.model(x_t, t, conditions, conditioning_signal)
+
+        # --- START: MODIFIED LOSS CALCULATION ---
+        loss = F.huber_loss(noise, predicted_noise, reduction='none')
+
+        if self.channel_weights is not None:
+            # Apply weights [B, L, C] * [1, 1, C]
+            weights = self.channel_weights.to(loss.device).view(1, 1, -1)
+            loss = (loss * weights).mean()
+        else:
+            loss = loss.mean()
+
+        return loss
+    # --- END: MODIFIED LOSS CALCULATION ---
+
+    @torch.no_grad()
+    def sample(self, num_samples: int, conditions: Dict[str, torch.Tensor], shape: tuple, 
+               conditioning_signal: Optional[torch.Tensor] = None) -> torch.Tensor:
+        device = next(self.model.parameters()).device
+        x = torch.randn(num_samples, *shape, device=device)
+        
+        for t in tqdm(reversed(range(self.num_timesteps)), desc="Sampling", total=self.num_timesteps, leave=False):
+            t_batch = torch.full((num_samples,), t, device=device, dtype=torch.long)
+            predicted_noise = self.model(x, t_batch, conditions, conditioning_signal)
+            
+            alpha_t = self.alphas[t]
+            sqrt_one_minus_alpha_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t]
+            
+            mean = (1 / torch.sqrt(alpha_t)) * (x - ((1 - alpha_t) / sqrt_one_minus_alpha_cumprod_t) * predicted_noise)
+            
+            if t > 0:
+                noise = torch.randn_like(x)
+                variance = self.posterior_variance[t]
+                x = mean + torch.sqrt(variance) * noise
+            else:
+                x = mean
+                
+        return x
+
+
+class HierarchicalDiffusionModel(nn.Module):
+    def __init__(self, in_channels: int, num_houses: int, downscale_factor: int, channel_weights: Optional[torch.Tensor] = None, **model_kwargs):
+        super().__init__()
+        self.downscale_factor = downscale_factor
+        self.fine_chunk_size = 2 * 96
+
+        # Pop num_timesteps *only once* at the top
+        num_timesteps = model_kwargs.pop("num_timesteps")
+        
+        self.downsampler = nn.Conv1d(in_channels, in_channels, kernel_size=downscale_factor, stride=downscale_factor)
+        self.upsampler = nn.ConvTranspose1d(in_channels, in_channels, kernel_size=downscale_factor, stride=downscale_factor)
+        
+        # Now num_timesteps can be passed to both models without error
+        self.coarse_model = ImprovedDiffusionModel(
+            ConditionalUnet(in_channels=in_channels, num_houses=num_houses, **model_kwargs), 
+            num_timesteps,
+            channel_weights=channel_weights
+        )
+        self.fine_model = ImprovedDiffusionModel(
+            ConditionalUnet(in_channels=in_channels, num_houses=num_houses, 
+                          cond_channels=in_channels, **model_kwargs), 
+            num_timesteps,
+            channel_weights=channel_weights
+        )
+        
+    def forward(self, x_0: torch.Tensor, conditions: Dict[str, torch.Tensor]) -> torch.Tensor:
+        x_0_coarse = self.downsampler(x_0.permute(0, 2, 1)).permute(0, 2, 1)
+        coarse_loss = self.coarse_model(x_0_coarse, conditions)
+        
+        with torch.no_grad():
+            x_0_coarse_upsampled = self.upsampler(x_0_coarse.detach().permute(0, 2, 1)).permute(0, 2, 1)
+            
+            if x_0_coarse_upsampled.shape[1] != x_0.shape[1]:
+                diff = x_0.shape[1] - x_0_coarse_upsampled.shape[1]
+                if diff > 0: x_0_coarse_upsampled = F.pad(x_0_coarse_upsampled, [0, 0, 0, diff])
+                else: x_0_coarse_upsampled = x_0_coarse_upsampled[:, :x_0.shape[1], :]
+            x_0_fine_residual = x_0 - x_0_coarse_upsampled
+        
+        full_length = x_0.shape[1]
+        if full_length > self.fine_chunk_size:
+            start_index = torch.randint(0, full_length - self.fine_chunk_size + 1, (1,)).item()
+        else:
+            start_index = 0
+            self.fine_chunk_size = full_length
+        
+        residual_chunk = x_0_fine_residual[:, start_index:start_index + self.fine_chunk_size, :]
+        conditioning_chunk = x_0_coarse_upsampled[:, start_index:start_index + self.fine_chunk_size, :]
+        
+        fine_loss = self.fine_model(residual_chunk, conditions, conditioning_signal=conditioning_chunk)
+        
+        fine_loss_weight = 1.5
+        return coarse_loss + (fine_loss * fine_loss_weight)
+
+    @torch.no_grad()
+    def sample(self, num_samples: int, conditions: Dict[str, torch.Tensor], shape: tuple) -> torch.Tensor:
+        full_length, num_features = shape
+        device = next(self.parameters()).device
+        
+        conditions = {k: v.to(device) for k, v in conditions.items()}
+        
+        print("--- Stage 1: Sampling Coarse Structure ---")
+        coarse_shape = (full_length // self.downscale_factor, num_features)
+        generated_coarse = self.coarse_model.sample(num_samples, conditions, shape=coarse_shape)
+        upsampled_coarse = self.upsampler(generated_coarse.permute(0, 2, 1)).permute(0, 2, 1)
+        
+        if upsampled_coarse.shape[1] != full_length:
+            diff = full_length - upsampled_coarse.shape[1]
+            if diff > 0: upsampled_coarse = F.pad(upsampled_coarse, [0, 0, 0, diff])
+            else: upsampled_coarse = upsampled_coarse[:, :full_length, :]
+        
+        print("--- Stage 2: Sampling Fine Details ---")
+        stitched_fine_residual = torch.zeros_like(upsampled_coarse)
+        
+        for start_index in tqdm(range(0, full_length, self.fine_chunk_size), desc="Fine chunks"):
+            end_index = min(start_index + self.fine_chunk_size, full_length)
+            chunk_length = end_index - start_index
+            fine_shape = (chunk_length, num_features)
+            conditioning_chunk = upsampled_coarse[:, start_index:end_index, :]
+            
+            generated_fine_chunk = self.fine_model.sample(
+                num_samples, conditions, shape=fine_shape, 
+                conditioning_signal=conditioning_chunk
+            )
+            
+            stitched_fine_residual[:, start_index:end_index, :] = generated_fine_chunk
+        
+        final_sample = upsampled_coarse + stitched_fine_residual
+        return final_sample

Data_generation_tool_kit/dataloader.py

@@ -0,0 +1,179 @@
+import torch
+import pandas as pd
+import numpy as np
+from sklearn.preprocessing import MinMaxScaler
+import joblib
+import os
+from typing import Tuple, Dict
+import warnings
+warnings.filterwarnings('ignore')
+
+class MultiHouseDataset(torch.utils.data.Dataset):
+  
+    def __init__(self, data_dir: str, window_size: int = 96, step_size: int = 1,
+                 scaler_path: str = 'global_scaler.gz', cache_in_memory: bool = True,
+                 dtype: torch.dtype = torch.float32, limit_to_one_year: bool = True):
+        self.window_size = window_size
+        self.step_size = step_size
+        self.cache_in_memory = cache_in_memory
+        self.dtype = dtype
+        self.limit_to_one_year = limit_to_one_year
+        
+        all_files = sorted([f for f in os.listdir(data_dir) if f.endswith('.csv')])
+        print(f"Found {len(all_files)} house files in '{data_dir}'.")
+        
+        self.num_houses = len(all_files)
+        
+        print("Reading house data...")
+        if self.limit_to_one_year:
+            print("INFO: Limiting data to the first year (17,520 samples) for each house.")
+            
+        data_per_house = []
+        timestamps_per_house = []
+    
+        SAMPLES_PER_YEAR = 17520 
+        
+        for filename in all_files:
+            df = pd.read_csv(os.path.join(data_dir, filename), parse_dates=['timestamp'])
+            timestamps_per_house.append(df['timestamp'].values)
+            time_series_values = df[['grid_usage', 'solar_generation']].values.astype(np.float32)
+
+            if self.limit_to_one_year:
+                time_series_values = time_series_values[:SAMPLES_PER_YEAR]
+
+            num_timesteps = len(time_series_values)
+            timesteps_of_day = np.arange(num_timesteps) % 48 
+
+            sin_time = np.sin(2 * np.pi * timesteps_of_day / 48.0).astype(np.float32)
+            cos_time = np.cos(2 * np.pi * timesteps_of_day / 48.0).astype(np.float32)
+
+            time_series_values = np.concatenate([
+                time_series_values, 
+                sin_time[:, np.newaxis],
+                cos_time[:, np.newaxis]
+            ], axis=1)
+
+            data_per_house.append(time_series_values)
+
+        if os.path.exists(scaler_path):
+            scaler = joblib.load(scaler_path)
+            print(f"Scaler loaded from {scaler_path}")
+        else:
+            print("Fitting global scaler...")
+            combined_data = np.vstack(data_per_house)
+            scaler = MinMaxScaler(feature_range=(-1, 1))
+            scaler.fit(combined_data)
+            joblib.dump(scaler, scaler_path)
+            print(f"Scaler saved to {scaler_path}")
+        
+        if self.cache_in_memory:
+            print("Caching normalized data...")
+            self.normalized_data_per_house = []
+            for series in data_per_house:
+                normalized = scaler.transform(series)
+                tensor_data = torch.from_numpy(normalized).to(dtype=self.dtype)
+                self.normalized_data_per_house.append(tensor_data)
+        else:
+            self.normalized_data_per_house = []
+            for series in data_per_house:
+                self.normalized_data_per_house.append(scaler.transform(series))
+        
+        del data_per_house
+            
+        print("Pre-computing mappings...")
+        
+        self.windows_per_house = [(len(d) - self.window_size) // self.step_size + 1 for d in self.normalized_data_per_house]
+        self.cumulative_windows = np.cumsum([0] + self.windows_per_house)
+        self.total_windows = self.cumulative_windows[-1]
+
+        self.sample_to_house = np.empty(self.total_windows, dtype=np.int32)
+        self.sample_to_local_idx = np.empty(self.total_windows, dtype=np.int32)
+        self.sample_to_day_of_week = np.empty(self.total_windows, dtype=np.int32)
+        self.sample_to_day_of_year = np.empty(self.total_windows, dtype=np.int32)
+
+        for house_idx in range(self.num_houses):
+            start_global_idx = self.cumulative_windows[house_idx]
+            end_global_idx = self.cumulative_windows[house_idx + 1]
+            num_windows_for_this_house = self.windows_per_house[house_idx]
+
+            self.sample_to_house[start_global_idx:end_global_idx] = house_idx
+            
+            local_indices = np.arange(num_windows_for_this_house) * self.step_size
+            self.sample_to_local_idx[start_global_idx:end_global_idx] = local_indices
+
+            house_timestamps = pd.Series(timestamps_per_house[house_idx][local_indices])
+            self.sample_to_day_of_week[start_global_idx:end_global_idx] = house_timestamps.dt.dayofweek
+            self.sample_to_day_of_year[start_global_idx:end_global_idx] = house_timestamps.dt.dayofyear - 1 
+
+        print(f"Dataset initialized. Total windows: {self.total_windows} from {self.num_houses} houses.")
+        memory_usage = sum(data.numel() * data.element_size() for data in self.normalized_data_per_house) / 1e6 if self.cache_in_memory else 0
+        print(f"Memory usage for cached tensors: {memory_usage:.1f} MB")
+
+    def __len__(self) -> int:
+        return self.total_windows
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:
+        if idx < 0 or idx >= self.total_windows:
+            raise IndexError("Index out of range")
+
+        house_index = self.sample_to_house[idx]
+        local_start_pos = self.sample_to_local_idx[idx]
+        
+        window_data = self.normalized_data_per_house[house_index][local_start_pos : local_start_pos + self.window_size]
+        
+        conditions = {
+            "house_id": torch.tensor(house_index, dtype=torch.long),
+            "day_of_week": torch.tensor(self.sample_to_day_of_week[idx], dtype=torch.long),
+            "day_of_year": torch.tensor(self.sample_to_day_of_year[idx], dtype=torch.long),
+        }
+        
+        return window_data, conditions
+
+    def get_memory_usage(self) -> dict:
+        if self.cache_in_memory:
+            tensor_memory = sum(data.numel() * data.element_size() for data in self.normalized_data_per_house) / 1e6
+        else:
+            tensor_memory = 0
+        
+        mapping_memory = (self.sample_to_house.nbytes + self.sample_to_local_idx.nbytes) / 1e6
+        
+        return {
+            'tensor_cache_mb': tensor_memory,
+            'mapping_arrays_mb': mapping_memory,
+            'total_mb': tensor_memory + mapping_memory
+        }
+
+class LatentDataset(torch.utils.data.Dataset):
+    def __init__(self, latent_vectors: torch.Tensor, house_ids: torch.Tensor):
+        assert len(latent_vectors) == len(house_ids), "Latent vectors and house IDs must have same length"
+        self.latent_vectors = latent_vectors.contiguous()
+        self.house_ids = house_ids.contiguous()
+        
+    def __len__(self) -> int:
+        return len(self.latent_vectors)
+    
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.latent_vectors[idx], self.house_ids[idx]
+
+if __name__ == "__main__":
+    import time
+    
+    DATA_DIRECTORY = './data/per_house/'
+    
+    if os.path.exists(DATA_DIRECTORY):
+        print("--- Testing Dataset Setup ---")
+        
+        start_time = time.time()
+        dataset = MultiHouseDataset(data_dir=DATA_DIRECTORY, window_size=96, step_size=96)
+        init_time = time.time() - start_time
+        print(f"Dataset initialization: {init_time:.2f}s")
+        print(f"Memory usage: {dataset.get_memory_usage()}")
+        
+        if len(dataset) > 0:
+            first_sample, first_conditions = dataset[0]
+            
+            print(f"\nSample data shape: {first_sample.shape}")
+            print(f"Sample conditions: {first_conditions}")
+            print(f"Total houses: {dataset.num_houses}")
+    else:
+        print(f"ERROR: Data directory not found at '{DATA_DIRECTORY}'. Please create and populate this directory.")

Data_generation_tool_kit/generate.py

@@ -0,0 +1,291 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import pandas as pd
+import os
+import joblib
+import math
+import datetime
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+import matplotlib.dates as mdates
+
+# =============================================================================
+# 1. MODEL CLASS DEFINITIONS
+# =============================================================================
+
+try:
+    from hierarchical_diffusion_model import (
+        HierarchicalDiffusionModel, ConditionalUnet, ResnetBlock1D, 
+        AttentionBlock1D, DownBlock1D, UpBlock1D, 
+        SinusoidalPositionEmbeddings, ImprovedDiffusionModel
+    )
+    print("Diffusion model classes imported.")
+except ImportError:
+    print("="*50)
+    print("ERROR: Could not import model classes from 'hierarchical_diffusion_model.py'.")
+    print("="*50)
+    exit()
+
+
+# =============================================================================
+# 2. HELPER FUNCTIONS
+# =============================================================================
+
+def add_amplitude_jitter(series, daily_samples=48, scale=0.05):
+    series = series.copy()
+    num_days = len(series) // daily_samples
+    if num_days == 0: return series
+    factors = np.random.normal(1.0, scale, size=num_days)
+    for d in range(num_days):
+        start, end = d * daily_samples, (d + 1) * daily_samples
+        series[start:end] *= factors[d]
+    return series
+
+def add_cloud_variability(pv, timestamps, base_sigma=0.25):
+    pv = pv.copy()
+    if len(pv) == 0: return pv
+    days = pd.Series(pv, index=timestamps).groupby(timestamps.date)
+    adjusted = []
+    for day, vals in days:
+        cloud_factor = np.random.lognormal(mean=-0.02, sigma=base_sigma)
+        hour = vals.index.hour
+        day_pv = np.where((hour >= 6) & (hour <= 18), vals * cloud_factor, 0.0)
+        adjusted.append(day_pv)
+    if not adjusted: return np.array([])
+    return np.concatenate(adjusted)
+
+def enforce_physics(df: pd.DataFrame, pv_cap_kw: float | None = None) -> pd.DataFrame:
+    df = df.copy()
+    df['solar_generation'] = np.clip(df['solar_generation'], 0.0, None)
+    hour = df.index.hour
+    night = (hour < 7) | (hour > 18)
+    df.loc[night, 'solar_generation'] = 0.0
+    export_mask = df['grid_usage'] < 0
+    if export_mask.any():
+        limited_export = -np.minimum(-df.loc[export_mask, 'grid_usage'], df.loc[export_mask, 'solar_generation'])
+        df.loc[export_mask, 'grid_usage'] = limited_export
+        zero_pv_neg_grid = export_mask & (df['solar_generation'] <= 1e-6)
+        df.loc[zero_pv_neg_grid, 'grid_usage'] = 0.0
+    if pv_cap_kw is not None:
+        df['solar_generation'] = np.clip(df['solar_generation'], 0.0, pv_cap_kw)
+    return df
+
+def calculate_generation_length(duration: str, samples_per_day: int) -> int:
+    """Calculate samples needed."""
+    if duration == '1_year':
+        return 365 * samples_per_day
+    elif duration == '6_months':
+        return 182 * samples_per_day
+    elif duration == '2_months':
+        return 60 * samples_per_day
+    elif duration == '1_month':
+        return 30 * samples_per_day
+    elif duration == '14_days':
+        return 14 * samples_per_day
+    elif duration == '7_days':
+        return 7 * samples_per_day
+    elif duration == '2_days':
+        return 2 * samples_per_day
+    else:
+        print(f"Warning: Unknown duration '{duration}'. Defaulting to 1 year.")
+        return 365 * samples_per_day
+
+# =============================================================================
+# 3. HARDCODED CONFIGURATION
+# =============================================================================
+
+class Config:
+    # --- Paths and Directories ---
+    MODEL_PATH = './trained_model/best_hierarchical_model.pth' 
+    SCALER_PATH = './data/global_scaler.gz'
+    ORIGINAL_DATA_DIR = './data/per_house'
+    OUTPUT_DIR = './generated_data' 
+
+    # --- Generation Parameters ---
+    GENERATION_DURATION = '1_year' 
+    NUM_PROFILES_TO_GENERATE = 2000
+    PLOTS_TO_GENERATE = 20      
+    GENERATION_BATCH_SIZE = 128
+
+    # --- Model & Training Parameters ---
+    TRAINING_WINDOW_DAYS = 14 
+    
+    NUM_HOUSES_TRAINED_ON = 300 
+    SAMPLES_PER_DAY = 48
+    NUM_FEATURES = 4            
+    DOWNSCALE_FACTOR = 4
+    EMBEDDING_DIM = 64
+    HIDDEN_SIZE = 512
+    HIDDEN_DIMS = [HIDDEN_SIZE // 4, HIDDEN_SIZE // 2, HIDDEN_SIZE] 
+    DROPOUT = 0.1
+    USE_ATTENTION = True
+    DIFFUSION_TIMESTEPS = 500
+    BLOCKS_PER_LEVEL = 3
+
+
+# =============================================================================
+# 4. MAIN GENERATION LOGIC
+# =============================================================================
+
+def main(cfg, run_output_dir):
+    """Main generation logic."""
+    DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+    print(f"Using device: {DEVICE}")
+
+    csv_output_dir = os.path.join(run_output_dir, 'csv')
+    plot_output_dir = os.path.join(run_output_dir, 'plots')
+    os.makedirs(csv_output_dir, exist_ok=True)
+    os.makedirs(plot_output_dir, exist_ok=True)
+    
+    print("Loading resources...")
+    try:
+        scaler = joblib.load(cfg.SCALER_PATH)
+        if scaler.n_features_in_ != cfg.NUM_FEATURES:
+            print(f"WARNING: Scaler was fit on {scaler.n_features_in_} features, but model expects {cfg.NUM_FEATURES}.")
+            
+        original_files = sorted([f for f in os.listdir(cfg.ORIGINAL_DATA_DIR) if f.endswith('.csv')])
+        if not original_files:
+            raise FileNotFoundError("No original data files found to extract timestamps.")
+        
+        sample_original_df = pd.read_csv(os.path.join(cfg.ORIGINAL_DATA_DIR, original_files[0]), index_col='timestamp', parse_dates=True)
+        
+        # Load 1 year timestamps
+        full_timestamps = sample_original_df.index[:(365 * cfg.SAMPLES_PER_DAY)]
+        
+        # Goal length
+        total_samples_needed = calculate_generation_length(cfg.GENERATION_DURATION, cfg.SAMPLES_PER_DAY)
+        
+        # Training window length
+        TRAINING_WINDOW_SAMPLES = cfg.TRAINING_WINDOW_DAYS * cfg.SAMPLES_PER_DAY
+        
+        # Clamping to max
+        if total_samples_needed > len(full_timestamps):
+            print(f"Warning: Requested {total_samples_needed} samples, but file has {len(full_timestamps)}. Clamping to max.")
+            total_samples_needed = len(full_timestamps)
+            
+        print(f"Goal: Generate {total_samples_needed} samples ({cfg.GENERATION_DURATION}) per profile.")
+        print(f"Strategy: Stitching {TRAINING_WINDOW_SAMPLES}-sample chunks.")
+        
+        model = HierarchicalDiffusionModel(
+            in_channels=cfg.NUM_FEATURES,
+            num_houses=cfg.NUM_HOUSES_TRAINED_ON,
+            downscale_factor=cfg.DOWNSCALE_FACTOR,
+            embedding_dim=cfg.EMBEDDING_DIM,
+            hidden_dims=cfg.HIDDEN_DIMS,
+            dropout=cfg.DROPOUT, 
+            use_attention=cfg.USE_ATTENTION,
+            num_timesteps=cfg.DIFFUSION_TIMESTEPS,
+            blocks_per_level=cfg.BLOCKS_PER_LEVEL
+        )
+        
+        model.load_state_dict(torch.load(cfg.MODEL_PATH, map_location=DEVICE))
+        model.to(DEVICE)
+        model.eval()
+        print("Model, scaler, timestamps ready.")
+
+    except FileNotFoundError as e:
+        print(f"ERROR: A required file was not found. Details: {e}")
+        return
+    except Exception as e:
+        print(f"An error occurred during setup: {e}")
+        return
+
+    num_batches = math.ceil(cfg.NUM_PROFILES_TO_GENERATE / cfg.GENERATION_BATCH_SIZE)
+    house_counter = 0
+
+    pbar = tqdm(range(num_batches), desc="Generating Batches")
+    for i in pbar:
+        current_batch_size = min(cfg.GENERATION_BATCH_SIZE, cfg.NUM_PROFILES_TO_GENERATE - house_counter)
+        if current_batch_size <= 0: break
+        pbar.set_postfix({'batch_size': current_batch_size})
+        
+        # --- STITCHING LOGIC ---
+        num_chunks_needed = math.ceil(total_samples_needed / TRAINING_WINDOW_SAMPLES)
+        batch_chunks_list = []
+
+        for chunk_idx in range(num_chunks_needed):
+            # Calculate chunk length
+            samples_remaining = total_samples_needed - (chunk_idx * TRAINING_WINDOW_SAMPLES)
+            current_chunk_length = min(TRAINING_WINDOW_SAMPLES, samples_remaining)
+            
+            shape_to_generate = (current_chunk_length, cfg.NUM_FEATURES)
+
+            # Generate random conditions
+            sample_conditions = {
+                "house_id": torch.randint(0, cfg.NUM_HOUSES_TRAINED_ON, (current_batch_size,), device=DEVICE),
+                "day_of_week": torch.randint(0, 7, (current_batch_size,), device=DEVICE),
+                "day_of_year": torch.randint(0, 365, (current_batch_size,), device=DEVICE)
+            }
+            
+            with torch.no_grad():
+                # Generate one chunk
+                generated_chunk_data = model.sample(current_batch_size, sample_conditions, shape=shape_to_generate)
+            
+            batch_chunks_list.append(generated_chunk_data.cpu().numpy())
+        
+        # Stitch chunks together
+        generated_data_np = np.concatenate(batch_chunks_list, axis=1)
+        # --- END OF STITCHING LOGIC ---
+
+        # --- Post-processing loop ---
+        for j in range(current_batch_size):
+            current_house_num = house_counter + 1
+            # Select timestamps
+            profile_timestamps = full_timestamps[:total_samples_needed]
+            normalized_series = generated_data_np[j]
+            
+            unscaled_series = scaler.inverse_transform(normalized_series)
+            
+            df = pd.DataFrame(
+                unscaled_series, 
+                columns=['grid_usage', 'solar_generation', 'sin_time', 'cos_time'], 
+                index=profile_timestamps
+            )
+
+            df = enforce_physics(df)
+            df['grid_usage'] = add_amplitude_jitter(df['grid_usage'].values, scale=0.08, daily_samples=cfg.SAMPLES_PER_DAY)
+            df['solar_generation'] = add_cloud_variability(df['solar_generation'].values, df.index, base_sigma=0.3)
+            df = enforce_physics(df) 
+            
+            df_to_save = df[['grid_usage', 'solar_generation']]
+            df_to_save.to_csv(os.path.join(csv_output_dir, f'generated_house_{current_house_num}.csv'))
+
+            if house_counter < cfg.PLOTS_TO_GENERATE:
+                plot_df = df_to_save.head(cfg.SAMPLES_PER_DAY * 14) 
+                plt.figure(figsize=(15, 6))
+                plt.plot(plot_df.index, plot_df['grid_usage'], label='Grid Usage', color='dodgerblue', alpha=0.9)
+                plt.plot(plot_df.index, plot_df['solar_generation'], label='Solar Generation', color='darkorange', alpha=0.9)
+                plt.title(f'Generated Data for Profile {current_house_num} (First 14 Days)')
+                plt.xlabel('Timestamp'); plt.ylabel('Power (kW)'); plt.legend(); plt.grid(True, which='both', linestyle='--', linewidth=0.5)
+                plt.tight_layout()
+                plt.savefig(os.path.join(plot_output_dir, f'generated_profile_{current_house_num}_plot.png'))
+                plt.close()
+            
+            house_counter += 1
+
+    print(f"\nSuccessfully generated and saved {house_counter} house profiles.")
+    if cfg.PLOTS_TO_GENERATE > 0:
+        print(f"Plots saved to '{plot_output_dir}'.")
+
+
+# =============================================================================
+# 5. --- Main execution block ---
+# =============================================================================
+
+if __name__ == '__main__':
+    config = Config()
+    
+    # Create unique output directory
+    run_timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    run_name = f"generation_run_{config.GENERATION_DURATION}_{run_timestamp}"
+    run_output_dir = os.path.join(config.OUTPUT_DIR, run_name)
+    os.makedirs(run_output_dir, exist_ok=True)
+    
+    print(f"Starting new generation run: {run_name}")
+    print(f"All outputs will be saved to: {run_output_dir}")
+    
+    # Run generation
+    main(config, run_output_dir)
+    
+    print("\nGeneration process complete.")

Data_generation_tool_kit/train.py

@@ -0,0 +1,209 @@
+import torch
+from torch.utils.data import DataLoader, random_split, Subset
+from torch.cuda.amp import autocast, GradScaler
+from tqdm import tqdm
+import numpy as np
+import os
+import datetime
+import pandas as pd
+import matplotlib.pyplot as plt
+import math
+import joblib  
+
+from dataloader import MultiHouseDataset
+from hierarchical_diffusion_model import HierarchicalDiffusionModel
+
+if torch.cuda.is_available():
+    DEVICE = "cuda"
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+    print("Using NVIDIA CUDA backend.")
+elif torch.backends.mps.is_available():
+    DEVICE = "mps"
+    print("Using Apple MPS backend.")
+else:
+    DEVICE = "cpu"
+    print("Using CPU.")
+
+EPOCHS = 200
+LEARNING_RATE = 1e-4
+BATCH_SIZE = 512
+USE_AMP = True
+GRADIENT_CLIP_VAL = 0.1
+
+WINDOW_DURATION = '14_days' 
+
+DATA_DIRECTORY = './data/per_house'
+NUM_WORKERS = os.cpu_count() // 2
+PIN_MEMORY = True
+USE_ATTENTION = True
+DROPOUT = 0.1
+HIDDEN_SIZE = 512
+EMBEDDING_DIM = 64
+DIFFUSION_TIMESTEPS = 500
+DOWNSCALE_FACTOR = 4
+
+def calculate_window_size(duration: str) -> int:
+    SAMPLES_PER_DAY = 48
+    mapping = {
+        '2_days': 2 * SAMPLES_PER_DAY,
+        '7_days': 7 * SAMPLES_PER_DAY,
+        '14_days': 14 * SAMPLES_PER_DAY,
+        '15_days': 15 * SAMPLES_PER_DAY,
+        '30_days': 30 * SAMPLES_PER_DAY
+    }
+    if duration not in mapping:
+        raise ValueError(f"Invalid WINDOW_DURATION: {duration}")
+    return mapping[duration]
+
+def denormalize_data(normalized_data, scaler_path='global_scaler.gz'):
+    scaler = joblib.load(scaler_path)
+    original_shape = normalized_data.shape
+    if len(original_shape) == 3:
+        batch_size, seq_len, features = original_shape
+        normalized_flat = normalized_data.reshape(-1, features)
+        denormalized_flat = scaler.inverse_transform(normalized_flat)
+        return denormalized_flat.reshape(original_shape)
+    else:
+        return scaler.inverse_transform(normalized_data)
+
+def moving_average(data, window_size):
+    return np.convolve(data, np.ones(window_size), 'valid') / window_size
+
+def save_and_plot_loss(loss_dict, title, filepath, window_size=10):
+    plt.figure(figsize=(12, 6))
+    for label, losses in loss_dict.items():
+        pd.DataFrame({label: losses}).to_csv(f"{filepath}_{label.lower().replace(' ', '_')}.csv", index=False)
+        plt.plot(losses, label=f'Raw {label}', alpha=0.3)
+        if len(losses) > window_size:
+            smoothed_losses = moving_average(losses, window_size)
+            plt.plot(np.arange(window_size - 1, len(losses)), smoothed_losses, label=f'Smoothed {label}')
+    plt.title(title)
+    plt.xlabel('Epoch'); plt.ylabel('Loss')
+    plt.legend(); plt.grid(True)
+    plt.savefig(f"{filepath}.png"); plt.close()
+    print(f" Loss plot saved to {filepath}.png")
+
+def train_diffusion(log_dir, model_save_path):
+    print("--- Starting Hierarchical Diffusion Training ---")
+    window_size = calculate_window_size(WINDOW_DURATION)
+    print(f"Using window duration: {WINDOW_DURATION} ({window_size} samples)")
+
+    dataset = MultiHouseDataset(
+        data_dir=DATA_DIRECTORY, 
+        window_size=window_size, 
+        step_size=window_size//2,  
+        limit_to_one_year=False
+    )
+    print(f"Dataset loaded: {len(dataset)} samples, {dataset.num_houses} houses, {dataset[0][0].shape[1]} features.")
+
+    val_split = 0.1
+    val_size = int(len(dataset) * val_split)
+    train_size = len(dataset) - val_size
+    train_dataset, val_dataset = random_split(dataset, [train_size, val_size])
+    print(f"Train size: {train_size}, Validation size: {val_size}")
+
+    train_dataloader = DataLoader(
+        train_dataset, batch_size=BATCH_SIZE, shuffle=True, 
+        num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY, drop_last=True
+    )
+    val_dataloader = DataLoader(
+        val_dataset, batch_size=BATCH_SIZE*2, shuffle=False, 
+        num_workers=NUM_WORKERS, pin_memory=PIN_MEMORY
+    )
+    
+    channel_weights = torch.tensor([1.0, 8.0, 1.0, 1.0], device=DEVICE)
+    print(f"Using channel weights: {channel_weights}")
+
+    model = HierarchicalDiffusionModel(
+        in_channels=dataset[0][0].shape[1],
+        num_houses=dataset.num_houses,
+        downscale_factor=DOWNSCALE_FACTOR,
+        channel_weights=channel_weights,
+        embedding_dim=EMBEDDING_DIM,
+        hidden_dims=[HIDDEN_SIZE // 4, HIDDEN_SIZE // 2, HIDDEN_SIZE],
+        dropout=DROPOUT, 
+        use_attention=USE_ATTENTION,
+        num_timesteps=DIFFUSION_TIMESTEPS,
+        blocks_per_level=3
+    ).to(DEVICE)
+    
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=EPOCHS)
+    scaler = GradScaler(enabled=(USE_AMP and DEVICE == "cuda"))
+    
+    train_losses, val_losses = [], []
+    best_val_loss = float('inf')
+
+    print(f"Starting training for {EPOCHS} epochs...")
+    for epoch in range(EPOCHS):
+        model.train()
+        total_train_loss = 0.0
+        pbar = tqdm(train_dataloader, desc=f"Epoch {epoch+1}/{EPOCHS} (Train)")
+        
+        for clean_data, conditions in pbar:
+            clean_data = clean_data.to(DEVICE, non_blocking=PIN_MEMORY)
+            conditions = {k: v.to(DEVICE, non_blocking=PIN_MEMORY) for k, v in conditions.items()}
+            
+            optimizer.zero_grad(set_to_none=True)
+            with autocast(enabled=(USE_AMP and DEVICE == "cuda")):
+                loss = model(clean_data, conditions)
+
+            scaler.scale(loss).backward()
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), GRADIENT_CLIP_VAL)
+            scaler.step(optimizer)
+            scaler.update()
+            
+            total_train_loss += loss.item()
+            pbar.set_postfix({'loss': f'{loss.item():.6f}', 'lr': f'{scheduler.get_last_lr()[0]:.2e}'})
+        
+        avg_train_loss = total_train_loss / len(train_dataloader)
+        train_losses.append(avg_train_loss)
+        
+        model.eval()
+        total_val_loss = 0.0
+        with torch.no_grad():
+            for clean_data, conditions in tqdm(val_dataloader, desc="Validating"):
+                clean_data = clean_data.to(DEVICE, non_blocking=PIN_MEMORY)
+                conditions = {k: v.to(DEVICE, non_blocking=PIN_MEMORY) for k, v in conditions.items()}
+                with autocast(enabled=(USE_AMP and DEVICE == "cuda")):
+                    loss = model(clean_data, conditions)
+                total_val_loss += loss.item()
+        
+        avg_val_loss = total_val_loss / len(val_dataloader)
+        val_losses.append(avg_val_loss)
+        
+        print(f"Epoch {epoch+1}/{EPOCHS} | Train Loss: {avg_train_loss:.6f} | Val Loss: {avg_val_loss:.6f}")
+        
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            torch.save(model.state_dict(), model_save_path)
+            print(f"New best model saved to {model_save_path} (Val Loss: {best_val_loss:.6f})")
+        
+        scheduler.step()
+
+    print("--- Training complete ---")
+    save_and_plot_loss(
+        {'Train Loss': train_losses, 'Validation Loss': val_losses},
+        'Hierarchical Diffusion Model Training & Validation Loss',
+        os.path.join(log_dir, 'diffusion_loss_curves')
+    )
+    
+    return dataset
+
+if __name__ == "__main__":
+    timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+    run_name = f"hierarchical_diffusion_{WINDOW_DURATION}_{timestamp}"
+    log_dir = os.path.join("./training_logs", run_name)
+    os.makedirs(log_dir, exist_ok=True)
+    model_path = os.path.join(log_dir, 'best_hierarchical_model.pth')
+
+    print(f"Starting new run: {run_name}")
+    print(f"Logs and models will be saved to: {log_dir}")
+
+    full_dataset = train_diffusion(log_dir=log_dir, model_save_path=model_path)
+    
+    print("\nTraining and best model saving complete.")
+    print(f"Model saved to: {model_path}")
+    print(f"Loss curves saved to: {os.path.join(log_dir, 'diffusion_loss_curves.png')}")