sampling.py

#!/usr/bin/env python3
"""
Langevin Dynamics Sampling
==========================

Implements flexible Langevin dynamics sampling with comprehensive trajectory tracking.
Works with learned networks, true score functions, and annealed sampling.
"""

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from tqdm import tqdm
from typing import Tuple, Optional, Union, List, Callable
import warnings
from abc import ABC, abstractmethod


class BaseLangevinSampler(ABC):
    """
    Base class for Langevin dynamics samplers containing common functionality.
    """
    
    def __init__(self, 
                 score_function: Callable[[torch.Tensor], torch.Tensor],
                 bounds: Tuple[float, float, float, float],
                 device: torch.device = None,
                 dataset = None):
        """
        Initialize the base Langevin sampler.
        
        Args:
            score_function: Function that takes x and returns score (gradient of log probability)
            bounds: Domain bounds (x_min, x_max, y_min, y_max)
            device: Device to run on
            dataset: Optional dataset for energy computation
        """
        self.score_function = score_function
        self.bounds = bounds
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.dataset = dataset
        
        # Initialize trajectory storage
        self.trajectory = []
        self.energies = []
        self.step_sizes = []
        self.score_magnitudes = []
        
        # Sampling statistics
        self.total_steps = 0
        self.acceptance_rate = 0.0
        self.final_samples = None
        
        # Bounds for sampling
        self.x_min, self.x_max, self.y_min, self.y_max = bounds
    
    def _validate_step_size(self, step_size: float) -> float:
        """Validate and potentially adjust step size."""
        if step_size <= 0:
            raise ValueError(f"Step size must be positive, got {step_size}")
        
        # Warn if step size seems too large
        domain_size = min(self.x_max - self.x_min, self.y_max - self.y_min)
        if step_size > domain_size * 0.1:
            warnings.warn(f"Large step size {step_size:.4f} relative to domain size {domain_size:.4f}")
        
        return step_size
    
    def _validate_samples(self, samples: torch.Tensor, clip_to_bounds: bool = True) -> torch.Tensor:
        """Validate that samples are within bounds and optionally clip them."""
        if samples.dim() != 2 or samples.shape[1] != 2:
            raise ValueError(f"Samples must be 2D with shape (n, 2), got {samples.shape}")
        
        # Check bounds
        x_vals, y_vals = samples[:, 0], samples[:, 1]
        x_in_bounds = (x_vals >= self.x_min) & (x_vals <= self.x_max)
        y_in_bounds = (y_vals >= self.y_min) & (y_vals <= self.y_max)
        
        if not torch.all(x_in_bounds & y_in_bounds):
            if clip_to_bounds:
                # Clip samples to domain bounds
                samples = samples.clone()
                samples[:, 0] = torch.clamp(samples[:, 0], self.x_min, self.x_max)
                samples[:, 1] = torch.clamp(samples[:, 1], self.y_min, self.y_max)
            else:
                warnings.warn("Some samples are outside domain bounds")
        
        return samples
    
    def _compute_energy(self, samples: torch.Tensor) -> torch.Tensor:
        """Compute energy (negative log probability) for samples."""
        if self.dataset is None:
            return torch.zeros(samples.shape[0], device=self.device)
        
        try:
            if hasattr(self.dataset, 'energy'):
                return self.dataset.energy(samples)
            elif hasattr(self.dataset, 'pdf'):
                # Convert to numpy for pdf computation
                samples_np = samples.cpu().numpy()
                pdf_values = self.dataset.pdf(samples_np[:, 0], samples_np[:, 1])
                # Avoid log(0) by adding small epsilon
                log_pdf = np.log(np.maximum(pdf_values, 1e-10))
                return -torch.tensor(log_pdf, dtype=torch.float32, device=self.device)
            else:
                return torch.zeros(samples.shape[0], device=self.device)
        except Exception as e:
            warnings.warn(f"Energy computation failed: {e}")
            return torch.zeros(samples.shape[0], device=self.device)
    
    def _store_trajectory_step(self, samples: torch.Tensor, step_size: float, 
                              scores: torch.Tensor, store_every: int = 1):
        """Store trajectory information."""
        if self.total_steps % store_every == 0:
            self.trajectory.append(samples.clone().detach().cpu())
            self.energies.append(self._compute_energy(samples).detach().cpu())
            self.step_sizes.append(step_size)
            
            # Store score magnitudes
            score_mags = torch.norm(scores, dim=1)
            self.score_magnitudes.append(score_mags.detach().cpu())
    
    def get_trajectory_info(self) -> dict:
        """Get comprehensive trajectory information."""
        if not self.trajectory:
            return {"message": "No trajectory data available"}
        
        final_samples = self.trajectory[-1]
        final_energies = self.energies[-1]
        
        # Convert lists to tensors for analysis
        all_energies = torch.cat(self.energies, dim=0) if self.energies else torch.tensor([])
        all_score_mags = torch.cat(self.score_magnitudes, dim=0) if self.score_magnitudes else torch.tensor([])
        
        return {
            'total_steps': self.total_steps,
            'trajectory_length': len(self.trajectory),
            'trajectory': self.trajectory,  # Include the actual trajectory data
            'final_samples': final_samples,
            'final_energies': final_energies,
            'step_sizes': self.step_sizes,
            'energy_stats': {
                'mean': torch.mean(all_energies).item() if len(all_energies) > 0 else 0,
                'std': torch.std(all_energies).item() if len(all_energies) > 0 else 0,
                'min': torch.min(all_energies).item() if len(all_energies) > 0 else 0,
                'max': torch.max(all_energies).item() if len(all_energies) > 0 else 0
            },
            'score_magnitude_stats': {
                'mean': torch.mean(all_score_mags).item() if len(all_score_mags) > 0 else 0,
                'std': torch.std(all_score_mags).item() if len(all_score_mags) > 0 else 0,
                'min': torch.min(all_score_mags).item() if len(all_score_mags) > 0 else 0,
                'max': torch.max(all_score_mags).item() if len(all_score_mags) > 0 else 0
            }
        }
    
    def visualize_trajectory(self, subsample: int = 1, show_energy: bool = True, 
                           show_scores: bool = True, save_name: Optional[str] = None):
        """Visualize the sampling trajectory."""
        if not self.trajectory:
            print("No trajectory data to visualize")
            return
        
        # Subsample trajectory for visualization
        traj_indices = range(0, len(self.trajectory), subsample)
        traj_subset = [self.trajectory[i] for i in traj_indices]
        
        n_plots = 2 + (1 if show_energy else 0) + (1 if show_scores else 0)
        fig, axes = plt.subplots(2, (n_plots + 1) // 2, figsize=(6 * ((n_plots + 1) // 2), 12))
        if n_plots <= 2:
            axes = axes.reshape(2, 1)
        
        # Plot 1: Trajectory evolution
        ax1 = axes[0, 0]
        for i, samples in enumerate(traj_subset):
            alpha = 0.1 + 0.9 * i / max(1, len(traj_subset) - 1)
            color = plt.cm.viridis(i / max(1, len(traj_subset) - 1))
            ax1.scatter(samples[:, 0], samples[:, 1], alpha=alpha, s=10, c=[color])
        
        ax1.set_xlim(self.x_min, self.x_max)
        ax1.set_ylim(self.y_min, self.y_max)
        ax1.set_title('Trajectory Evolution')
        ax1.set_xlabel('X')
        ax1.set_ylabel('Y')
        ax1.grid(True, alpha=0.3)
        
        # Plot 2: Final samples
        ax2 = axes[0, 1]
        final_samples = self.trajectory[-1]
        ax2.scatter(final_samples[:, 0], final_samples[:, 1], alpha=0.6, s=20, c='red')
        ax2.set_xlim(self.x_min, self.x_max)
        ax2.set_ylim(self.y_min, self.y_max)
        ax2.set_title('Final Samples')
        ax2.set_xlabel('X')
        ax2.set_ylabel('Y')
        ax2.grid(True, alpha=0.3)
        
        plot_idx = 2
        
        # Plot 3: Energy evolution (if requested)
        if show_energy and self.energies:
            ax3 = axes[1, 0]
            # Plot energy statistics over time
            energy_means = [torch.mean(e).item() for e in self.energies[::subsample]]
            energy_stds = [torch.std(e).item() for e in self.energies[::subsample]]
            
            steps = range(0, len(self.energies), subsample)
            ax3.plot(steps, energy_means, 'b-', label='Mean Energy')
            ax3.fill_between(steps, 
                           [m - s for m, s in zip(energy_means, energy_stds)],
                           [m + s for m, s in zip(energy_means, energy_stds)],
                           alpha=0.3, color='blue')
            ax3.set_xlabel('Step')
            ax3.set_ylabel('Energy')
            ax3.set_title('Energy Evolution')
            ax3.legend()
            ax3.grid(True, alpha=0.3)
            plot_idx += 1
        
        # Plot 4: Score magnitudes (if requested)
        if show_scores and self.score_magnitudes:
            ax4 = axes[1, 1] if show_energy else axes[1, 0]
            # Plot score magnitude statistics over time
            score_means = [torch.mean(s).item() for s in self.score_magnitudes[::subsample]]
            score_stds = [torch.std(s).item() for s in self.score_magnitudes[::subsample]]
            
            steps = range(0, len(self.score_magnitudes), subsample)
            ax4.plot(steps, score_means, 'g-', label='Mean Score Magnitude')
            ax4.fill_between(steps,
                           [m - s for m, s in zip(score_means, score_stds)],
                           [m + s for m, s in zip(score_means, score_stds)],
                           alpha=0.3, color='green')
            ax4.set_xlabel('Step')
            ax4.set_ylabel('Score Magnitude')
            ax4.set_title('Score Magnitude Evolution')
            ax4.legend()
            ax4.grid(True, alpha=0.3)
        
        plt.tight_layout()
        
        if save_name:
            plt.savefig(save_name, dpi=150, bbox_inches='tight')
        
        plt.show()
    
    def create_animation(self, save_name: Optional[str] = None, interval: int = 100, 
                        subsample: int = 1) -> Optional[FuncAnimation]:
        """Create an animation of the sampling process."""
        if not self.trajectory:
            print("No trajectory data to animate")
            return None
        
        # Subsample trajectory for animation
        traj_indices = range(0, len(self.trajectory), subsample)
        traj_subset = [self.trajectory[i] for i in traj_indices]
        
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
        
        # Setup axes
        ax1.set_xlim(self.x_min, self.x_max)
        ax1.set_ylim(self.y_min, self.y_max)
        ax1.set_title('Langevin Dynamics Sampling')
        ax1.set_xlabel('X')
        ax1.set_ylabel('Y')
        ax1.grid(True, alpha=0.3)
        
        ax2.set_xlim(self.x_min, self.x_max)
        ax2.set_ylim(self.y_min, self.y_max)
        ax2.set_title('Accumulated Samples')
        ax2.set_xlabel('X')
        ax2.set_ylabel('Y')
        ax2.grid(True, alpha=0.3)
        
        # Initialize plots
        scat1 = ax1.scatter([], [], s=30, c='red', alpha=0.7)
        scat2 = ax2.scatter([], [], s=10, c='blue', alpha=0.3)
        
        # Accumulate samples for ax2
        accumulated_samples = []
        
        def animate(frame):
            nonlocal accumulated_samples
            
            if frame < len(traj_subset):
                current_samples = traj_subset[frame]
                
                # Update current samples (ax1)
                scat1.set_offsets(current_samples.numpy())
                
                # Accumulate samples (ax2)
                accumulated_samples.append(current_samples)
                if len(accumulated_samples) > 0:
                    all_samples = torch.cat(accumulated_samples, dim=0)
                    scat2.set_offsets(all_samples.numpy())
                
                # Update titles with step information
                ax1.set_title(f'Current Samples (Step {frame * subsample})')
                ax2.set_title(f'Accumulated Samples (Total: {len(all_samples) if len(accumulated_samples) > 0 else 0})')
            
            return scat1, scat2
        
        anim = FuncAnimation(fig, animate, frames=len(traj_subset), 
                           interval=interval, blit=False, repeat=True)
        
        if save_name:
            anim.save(save_name, writer='pillow', fps=10)
            print(f"Animation saved as {save_name}")
        
        plt.show()
        return anim
    
    def clear_trajectory(self):
        """Clear stored trajectory data."""
        self.trajectory = []
        self.energies = []
        self.step_sizes = []
        self.score_magnitudes = []
        self.total_steps = 0
        self.acceptance_rate = 0.0
        self.final_samples = None
    
    @abstractmethod
    def sample(self, *args, **kwargs):
        """Abstract method for sampling - must be implemented by subclasses."""
        pass


class LangevinSampler(BaseLangevinSampler):
    """
    Standard Langevin dynamics sampler.
    """
    
    def __init__(self, 
                 score_function: Callable[[torch.Tensor], torch.Tensor],
                 bounds: Tuple[float, float, float, float],
                 device: torch.device = None,
                 dataset = None):
        """
        Initialize the Langevin sampler.
        
        Args:
            score_function: Function that takes x and returns score (gradient of log probability)
            bounds: Domain bounds (x_min, x_max, y_min, y_max)
            device: Device to run on
            dataset: Optional dataset for energy computation
        """
        super().__init__(score_function, bounds, device, dataset)
    
    def sample(self, num_samples: int, num_steps: int, step_size: float = 0.01,
               initialization: str = 'uniform', store_trajectory: bool = True,
               store_every: int = 1, verbose: bool = True) -> Tuple[torch.Tensor, dict]:
        """
        Generate samples using Langevin dynamics.
        
        Args:
            num_samples: Number of samples to generate
            num_steps: Number of Langevin steps
            step_size: Step size for Langevin dynamics
            initialization: How to initialize samples ('uniform', 'normal', 'center')
            store_trajectory: Whether to store the full trajectory
            store_every: Store trajectory every N steps
            verbose: Whether to show progress
            
        Returns:
            Tuple of (final_samples, info_dict)
        """
        step_size = self._validate_step_size(step_size)
        
        # Initialize samples
        if initialization == 'uniform':
            samples = torch.rand(num_samples, 2, device=self.device)
            samples[:, 0] = samples[:, 0] * (self.x_max - self.x_min) + self.x_min
            samples[:, 1] = samples[:, 1] * (self.y_max - self.y_min) + self.y_min
        elif initialization == 'normal':
            center_x = (self.x_min + self.x_max) / 2
            center_y = (self.y_min + self.y_max) / 2
            scale = min(self.x_max - self.x_min, self.y_max - self.y_min) / 4
            samples = torch.randn(num_samples, 2, device=self.device) * scale
            samples[:, 0] += center_x
            samples[:, 1] += center_y
        elif initialization == 'center':
            center_x = (self.x_min + self.x_max) / 2
            center_y = (self.y_min + self.y_max) / 2
            samples = torch.full((num_samples, 2), fill_value=0.0, device=self.device)
            samples[:, 0] = center_x
            samples[:, 1] = center_y
        else:
            raise ValueError(f"Unknown initialization: {initialization}")
        
        samples = self._validate_samples(samples)
        
        # Clear previous trajectory
        if store_trajectory:
            self.clear_trajectory()
        
        # Langevin dynamics
        pbar = tqdm(range(num_steps), desc="Langevin Sampling") if verbose else range(num_steps)
        
        for step in pbar:
            # Compute scores
            scores = self.score_function(samples)
            
            # Langevin update: x_{t+1} = x_t + (step_size/2) * score + sqrt(step_size) * noise
            noise = torch.randn_like(samples)
            samples = samples + (step_size / 2) * scores + torch.sqrt(torch.tensor(step_size)) * noise
            
            # Don't clip samples to domain bounds (let them flow freely)
            samples = self._validate_samples(samples, clip_to_bounds=False)
            
            # Store trajectory
            if store_trajectory:
                self._store_trajectory_step(samples, step_size, scores, store_every)
            
            self.total_steps += 1
            
            # Update progress bar
            if verbose and isinstance(pbar, tqdm):
                if step % 10 == 0:
                    score_mag = torch.mean(torch.norm(scores, dim=1)).item()
                    pbar.set_postfix({'score_mag': f'{score_mag:.4f}'})
        
        samples = self._validate_samples(samples, clip_to_bounds=False)
        self.final_samples = samples
        
        # Compute final info
        info = self.get_trajectory_info()
        info.update({
            'num_samples': num_samples,
            'num_steps': num_steps,
            'step_size': step_size,
            'initialization': initialization,
            'final_score_magnitude': torch.mean(torch.norm(self.score_function(samples), dim=1)).item()
        })
        
        return samples, info


class AnnealedLangevinSampler(BaseLangevinSampler):
    """
    Annealed Langevin dynamics sampler for noise conditional score networks.
    """
    
    def __init__(self, 
                 score_network,
                 noise_levels: List[float],
                 bounds: Tuple[float, float, float, float],
                 device: torch.device = None,
                 dataset = None):
        """
        Initialize the annealed Langevin sampler.
        
        Args:
            score_network: Noise conditional score network s_θ(x, σ)
            noise_levels: List of noise levels σ_L > ... > σ_1
            bounds: Domain bounds (x_min, x_max, y_min, y_max)
            device: Device to run on
            dataset: Optional dataset for energy computation
        """
        # Score function wrapper for noise conditional network
        self.score_network = score_network
        self.current_sigma = None
        
        def score_function(x):
            if self.current_sigma is None:
                raise ValueError("Noise level not set")
            sigma_tensor = torch.full((x.shape[0],), self.current_sigma, device=x.device)
            return self.score_network(x, sigma_tensor)
        
        super().__init__(score_function, bounds, device, dataset)
        
        # Sort noise levels in descending order for annealing
        self.noise_levels = sorted(noise_levels, reverse=True)
        self.num_noise_levels = len(noise_levels)
        
        # Annealing-specific storage
        self.noise_level_history = []
        self.per_noise_trajectories = []
    
    def _compute_adaptive_step_size(self, sigma: float, base_step_size: float = 0.01) -> float:
        """
        Compute adaptive step size based on current noise level.
        
        Uses formula: ε_i = ε * (σ_i / σ_max)²
        """
        sigma_max = max(self.noise_levels)
        return base_step_size * (sigma / sigma_max) ** 2
    
    def sample(self, num_samples: int, steps_per_noise: int, 
               base_step_size: float = 0.01, initialization: str = 'uniform',
               store_trajectory: bool = True, store_every: int = 1,
               verbose: bool = True) -> Tuple[torch.Tensor, dict]:
        """
        Generate samples using annealed Langevin dynamics.
        
        Args:
            num_samples: Number of samples to generate
            steps_per_noise: Number of Langevin steps per noise level
            base_step_size: Base step size (will be adapted per noise level)
            initialization: How to initialize samples ('uniform', 'normal', 'prior')
            store_trajectory: Whether to store the full trajectory
            store_every: Store trajectory every N steps
            verbose: Whether to show progress
            
        Returns:
            Tuple of (final_samples, info_dict)
        """
        base_step_size = self._validate_step_size(base_step_size)
        
        # Initialize samples
        if initialization == 'uniform':
            samples = torch.rand(num_samples, 2, device=self.device)
            samples[:, 0] = samples[:, 0] * (self.x_max - self.x_min) + self.x_min
            samples[:, 1] = samples[:, 1] * (self.y_max - self.y_min) + self.y_min
        elif initialization == 'normal':
            center_x = (self.x_min + self.x_max) / 2
            center_y = (self.y_min + self.y_max) / 2
            scale = min(self.x_max - self.x_min, self.y_max - self.y_min) / 4
            samples = torch.randn(num_samples, 2, device=self.device) * scale
            samples[:, 0] += center_x
            samples[:, 1] += center_y
        elif initialization == 'prior':
            # Initialize from Gaussian noise (as in the literature)
            samples = torch.randn(num_samples, 2, device=self.device)
            # Scale to fit within bounds
            samples = samples * 0.5  # Make it more reasonable
            center_x = (self.x_min + self.x_max) / 2
            center_y = (self.y_min + self.y_max) / 2
            samples[:, 0] += center_x
            samples[:, 1] += center_y
        else:
            raise ValueError(f"Unknown initialization: {initialization}")
        
        samples = self._validate_samples(samples)
        
        # Clear previous trajectory
        if store_trajectory:
            self.clear_trajectory()
            self.per_noise_trajectories = []
            self.noise_level_history = []
        
        # Annealed Langevin dynamics
        total_steps = len(self.noise_levels) * steps_per_noise
        pbar = tqdm(total=total_steps, desc="Annealed Langevin") if verbose else None
        
        if verbose:
            print(f"Starting Annealed Langevin Dynamics:")
            print(f"  📊 Noise levels: {len(self.noise_levels)} (σ: {self.noise_levels[-1]:.3f} → {self.noise_levels[0]:.3f})")
            print(f"  🎯 Samples: {num_samples}")
            print(f"  🚀 Steps per noise level: {steps_per_noise}")
            print(f"  📈 Total steps: {total_steps}")
        
        for level_idx, sigma in enumerate(self.noise_levels):
            self.current_sigma = sigma
            step_size = self._compute_adaptive_step_size(sigma, base_step_size)
            
            if verbose:
                print(f"  🔄 Noise level {level_idx + 1}/{len(self.noise_levels)}: σ={sigma:.4f}, ε={step_size:.4f}")
            
            level_trajectory = []
            
            # Run Langevin dynamics for this noise level
            for step in range(steps_per_noise):
                # Compute scores with current noise level
                scores = self.score_function(samples)
                
                # Langevin update
                noise = torch.randn_like(samples)
                samples = samples + (step_size / 2) * scores + torch.sqrt(torch.tensor(step_size)) * noise
                
                # Don't clip samples to domain bounds (let them flow freely)
                samples = self._validate_samples(samples, clip_to_bounds=False)
                
                # Store trajectory
                if store_trajectory:
                    self._store_trajectory_step(samples, step_size, scores, store_every)
                    if step % store_every == 0:
                        level_trajectory.append(samples.clone().detach().cpu())
                
                self.total_steps += 1
                
                # Update progress bar
                if verbose and pbar:
                    pbar.update(1)
                    if step % 10 == 0:
                        score_mag = torch.mean(torch.norm(scores, dim=1)).item()
                        pbar.set_postfix({
                            'noise_level': f'{sigma:.4f}',
                            'step_size': f'{step_size:.4f}',
                            'score_mag': f'{score_mag:.4f}'
                        })
            
            # Store per-noise-level information
            if store_trajectory:
                self.per_noise_trajectories.append(level_trajectory)
                self.noise_level_history.extend([sigma] * steps_per_noise)
        
        if verbose and pbar:
            pbar.close()
        
        if verbose:
            print(f"  ✅ Annealed Langevin Dynamics completed! Final samples ready.")
        
        samples = self._validate_samples(samples, clip_to_bounds=False)
        self.final_samples = samples
        
        # Compute final info
        info = self.get_trajectory_info()
        info.update({
            'num_samples': num_samples,
            'steps_per_noise': steps_per_noise,
            'base_step_size': base_step_size,
            'noise_levels': self.noise_levels,
            'total_noise_levels': len(self.noise_levels),
            'initialization': initialization,
            'final_score_magnitude': torch.mean(torch.norm(self.score_function(samples), dim=1)).item(),
            'adaptive_step_sizes': [self._compute_adaptive_step_size(sigma, base_step_size) for sigma in self.noise_levels]
        })
        
        return samples, info
    
    def visualize_annealing_process(self, subsample: int = 1, save_name: Optional[str] = None):
        """
        Visualize the annealing process across different noise levels.
        """
        if not self.per_noise_trajectories:
            print("No per-noise-level trajectory data to visualize")
            return
        
        num_levels = len(self.per_noise_trajectories)
        cols = min(4, num_levels)
        rows = (num_levels + cols - 1) // cols
        
        fig, axes = plt.subplots(rows, cols, figsize=(5 * cols, 4 * rows))
        if rows == 1:
            axes = axes.reshape(1, -1)
        elif cols == 1:
            axes = axes.reshape(-1, 1)
        
        for level_idx, (sigma, level_traj) in enumerate(zip(self.noise_levels, self.per_noise_trajectories)):
            row = level_idx // cols
            col = level_idx % cols
            ax = axes[row, col]
            
            # Plot trajectory evolution for this noise level
            for i, samples in enumerate(level_traj[::subsample]):
                alpha = 0.3 + 0.7 * i / max(1, len(level_traj) - 1)
                size = 5 + 15 * i / max(1, len(level_traj) - 1)
                color = plt.cm.plasma(i / max(1, len(level_traj) - 1))
                ax.scatter(samples[:, 0], samples[:, 1], alpha=alpha, s=size, c=[color])
            
            ax.set_xlim(self.x_min, self.x_max)
            ax.set_ylim(self.y_min, self.y_max)
            ax.set_title(f'Noise Level σ = {sigma:.4f}')
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            ax.grid(True, alpha=0.3)
        
        # Hide unused subplots
        for level_idx in range(num_levels, rows * cols):
            row = level_idx // cols
            col = level_idx % cols
            axes[row, col].set_visible(False)
        
        plt.tight_layout()
        
        if save_name:
            plt.savefig(save_name, dpi=150, bbox_inches='tight')
        
        plt.show()


class DDPMSampler:
    """
    DDPM (Denoising Diffusion Probabilistic Model) sampler.
    Implements the reverse diffusion process for sampling.
    """
    
    def __init__(self, diffusion_model, device: torch.device = None):
        """
        Initialize the DDPM sampler.
        
        Args:
            diffusion_model: Trained DiffusionModel with network and scheduler
            device: Device to run on
        """
        self.model = diffusion_model
        self.network = diffusion_model.network
        self.scheduler = diffusion_model.scheduler
        self.device = device or next(self.network.parameters()).device
        
        # Move model to device
        self.model.to(self.device)
        self.scheduler.to(self.device)
        
        # Sampling statistics
        self.total_steps = 0
        self.final_samples = None
    
    def sample_step(self, x: torch.Tensor, t: int) -> torch.Tensor:
        """
        Perform a single reverse diffusion step.
        
        Args:
            x: Current noisy data (batch_size, 2)
            t: Current timestep
            
        Returns:
            Denoised data (batch_size, 2)
        """
        batch_size = x.shape[0]
        device = x.device
        
        # Create timestep tensor
        timesteps = torch.full((batch_size,), t, device=device, dtype=torch.long)
        
        # Predict noise
        predicted_noise = self.network(x, timesteps)
        
        # Compute denoised prediction
        alpha = self.scheduler.alphas[t]
        alpha_cumprod = self.scheduler.alphas_cumprod[t]
        beta = self.scheduler.betas[t]
        
        # Compute x_{t-1}
        x_prev = (1 / torch.sqrt(alpha)) * (
            x - (beta / torch.sqrt(1 - alpha_cumprod)) * predicted_noise
        )
        
        # Add noise if not the last step
        if t > 0:
            posterior_variance = self.scheduler.posterior_variance[t]
            noise = torch.randn_like(x)
            x_prev = x_prev + torch.sqrt(posterior_variance) * noise
        
        return x_prev
    
    def sample(self, num_samples: int, store_trajectory: bool = False, 
               store_every: int = 100, verbose: bool = True) -> Tuple[torch.Tensor, dict]:
        """
        Generate samples using the reverse diffusion process.
        
        Args:
            num_samples: Number of samples to generate
            store_trajectory: Whether to store the trajectory (optional for consistency)
            store_every: Store trajectory every N steps (optional for consistency)
            verbose: Whether to show progress
            
        Returns:
            Tuple of (generated_samples, info_dict)
        """
        self.model.eval()
        
        with torch.no_grad():
            # Start from random noise
            x = torch.randn(num_samples, self.network.dim, device=self.device)
            
            # Reverse diffusion process  
            num_timesteps = self.scheduler.num_timesteps
            
            # Store trajectory if requested
            trajectory = []
            timesteps = []
            if store_trajectory:
                trajectory.append(x.clone().cpu())
                timesteps.append(num_timesteps)  # Initial timestep
            pbar = tqdm(reversed(range(num_timesteps)), desc="DDPM Sampling") if verbose else reversed(range(num_timesteps))
            
            for step_idx, t in enumerate(pbar):
                x = self.sample_step(x, t)
                
                # Store trajectory
                if store_trajectory and step_idx % store_every == 0:
                    trajectory.append(x.clone().cpu())
                    timesteps.append(t)
                
                self.total_steps += 1
                
                # Update progress bar
                if verbose and isinstance(pbar, tqdm):
                    if step_idx % 50 == 0:
                        pbar.set_postfix({'timestep': f'{t}'})
        
        self.final_samples = x
        
        # Compute info
        info = {
            'num_samples': num_samples,
            'num_timesteps': num_timesteps,
            'total_steps': self.total_steps,
            'final_samples': x.cpu(),
            'sampling_method': 'ddpm_reverse_diffusion'
        }
        
        if store_trajectory:
            info['trajectory'] = trajectory
            info['timesteps'] = timesteps
            info['trajectory_length'] = len(trajectory)
        
        return x, info
    
    def clear_trajectory(self):
        """Clear stored trajectory data (for consistency with other samplers)."""
        self.total_steps = 0
        self.final_samples = None


def create_langevin_sampler(score_fn: Callable[[torch.Tensor], torch.Tensor],
                           bounds: Tuple[float, float, float, float],
                           device: torch.device = torch.device('cpu'),
                           dataset = None) -> LangevinSampler:
    """
    Convenience function to create a LangevinSampler.
    
    Args:
        score_fn: Score function (neural network, exact function, etc.)
        bounds: Domain bounds (x_min, x_max, y_min, y_max)
        device: Device to use
        dataset: Optional dataset for energy tracking
        
    Returns:
        LangevinSampler instance
    """
    return LangevinSampler(
        score_function=score_fn,
        bounds=bounds,
        device=device,
        dataset=dataset
    )


def create_exact_sampler(dataset, bounds: Tuple[float, float, float, float], 
                        device: torch.device = None) -> LangevinSampler:
    """
    Convenience function to create a LangevinSampler with exact score function.
    
    Args:
        dataset: Dataset object with true_score method
        bounds: Domain bounds (x_min, x_max, y_min, y_max)  
        device: Device to use (inferred from dataset if None)
        
    Returns:
        LangevinSampler instance
    """
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    def exact_score_function(x):
        x_np = x.cpu().numpy()
        scores = dataset.true_score(x_np[:, 0], x_np[:, 1])
        if isinstance(scores, torch.Tensor):
            return scores.to(device)
        else:
            return torch.tensor(scores, dtype=torch.float32, device=device).view(x.shape[0], 2)
    
    return LangevinSampler(exact_score_function, bounds, device, dataset)


def create_learned_sampler(score_network, bounds: Tuple[float, float, float, float],
                          device: torch.device = None, dataset = None) -> LangevinSampler:
    """
    Convenience function to create a LangevinSampler with learned score function.
    
    Args:
        score_network: Trained score network
        bounds: Domain bounds (x_min, x_max, y_min, y_max)
        device: Device to use (inferred from network if None)
        dataset: Optional dataset for energy tracking
        
    Returns:
        LangevinSampler instance
    """
    if device is None:
        device = next(score_network.parameters()).device
    
    def learned_score_function(x):
        score_network.eval()
        with torch.no_grad():
            return score_network(x.to(device))
    
    return LangevinSampler(learned_score_function, bounds, device, dataset)


def create_annealed_sampler_from_ncsn(score_network, noise_levels: List[float],
                                     dataset, device: torch.device = None) -> AnnealedLangevinSampler:
    """
    Convenience function to create an AnnealedLangevinSampler from NCSN results.
    
    Args:
        score_network: Trained noise conditional score network
        noise_levels: List of noise levels used in training
        dataset: Dataset object (for bounds and energy computation)
        device: Device to use (inferred from network if None)
        
    Returns:
        AnnealedLangevinSampler instance
    """
    if device is None:
        device = next(score_network.parameters()).device
    
    bounds = dataset.get_bounds()
    
    return AnnealedLangevinSampler(
        score_network=score_network,
        noise_levels=noise_levels,
        bounds=bounds,
        device=device,
        dataset=dataset
    )


def create_annealed_langevin_sampler(score_network,
                                   noise_levels: List[float],
                                   bounds: Tuple[float, float, float, float],
                                   device: torch.device = None,
                                   dataset = None) -> AnnealedLangevinSampler:
    """
    Convenience function to create an AnnealedLangevinSampler.
    
    Args:
        score_network: Noise conditional score network
        noise_levels: List of noise levels
        bounds: Domain bounds
        device: Device to use (inferred from network if None)
        dataset: Optional dataset for energy tracking
        
    Returns:
        AnnealedLangevinSampler instance
    """
    if device is None:
        device = next(score_network.parameters()).device
    
    return AnnealedLangevinSampler(
        score_network=score_network,
        noise_levels=noise_levels,
        bounds=bounds,
        device=device,
        dataset=dataset
    )


def create_ddpm_sampler(diffusion_model, device: torch.device = None) -> DDPMSampler:
    """
    Create a DDPM sampler.
    
    Args:
        diffusion_model: Trained diffusion model
        device: Device to run on
        
    Returns:
        DDPM sampler
    """
    if device is None:
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
    return DDPMSampler(diffusion_model, device)