sampling.py

"""Sampling methods for diffusion models.

This module implements different sampling strategies for denoising diffusion models:
- DDPM: Original sampling from Ho et al. (2020)
- DDIM: Deterministic sampling from Song et al. (2021) - faster inference
"""

from abc import ABC, abstractmethod

import torch
from tqdm import tqdm

from ddpm.config import DiffusionConfig
from ddpm.models import ConditionalUNet


class DiffusionSampler(ABC):
    """Abstract base class for diffusion samplers.

    Args:
        config: Diffusion configuration
        model: Conditional U-Net for denoising
        betas: Noise schedule
        alphas_cumprod: Cumulative product of alphas
    """

    def __init__(
        self,
        config: DiffusionConfig,
        model: ConditionalUNet,
        betas: torch.Tensor,
        alphas_cumprod: torch.Tensor,
    ):
        self.config = config
        self.model = model
        self.device = config.device
        self.betas = betas
        self.alphas_cumprod = alphas_cumprod
        self.alphas = 1.0 - betas

    @abstractmethod
    def sample(
        self,
        conditioning: torch.Tensor,
        n_samples: int = 1,
        require_grad: bool = False,
        initial_noise: torch.Tensor | None = None,
        show_progress: bool = True,
    ) -> torch.Tensor:
        """Generate samples given conditioning data.

        Args:
            conditioning: Conditioning data
            n_samples: Number of samples to generate
            require_grad: If True, enable gradient computation
            initial_noise: Optional fixed initial noise tensor
            show_progress: If True, show tqdm progress bar

        Returns:
            Generated samples
        """
        pass

    def _prepare_sampling(
        self,
        conditioning: torch.Tensor,
        n_samples: int,
        require_grad: bool,
        initial_noise: torch.Tensor | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Prepare conditioning and initial noise for sampling.

        Handles model eval mode, conditioning expansion, device transfer,
        gradient setup, and noise initialisation.

        Args:
            conditioning: Raw conditioning data
            n_samples: Number of samples per conditioning
            require_grad: Whether to enable gradient computation
            initial_noise: Optional fixed initial noise

        Returns:
            Tuple of (prepared conditioning, initial noise x_t)
        """
        if not require_grad:
            self.model.eval()

        # Ensure conditioning has batch dimension
        if conditioning.dim() == 1:
            conditioning = conditioning.unsqueeze(0).repeat(n_samples, 1)
        elif conditioning.dim() == 2:
            if self.config.n_conditioning_channels > 1:
                conditioning = conditioning.unsqueeze(0).repeat(n_samples, 1, 1)
            else:
                conditioning = conditioning.repeat_interleave(n_samples, dim=0)
        elif conditioning.dim() == 3:
            conditioning = conditioning.repeat_interleave(n_samples, dim=0)

        conditioning = conditioning.to(self.device)
        if require_grad and not conditioning.requires_grad:
            conditioning.requires_grad_(True)

        if initial_noise is not None:
            x_t = initial_noise.to(self.device)
        else:
            x_t = torch.randn(
                conditioning.shape[0], self.config.trajectory_length, device=self.device
            )

        return conditioning, x_t


class DDPMSampler(DiffusionSampler):
    """DDPM sampling (Ho et al., 2020).

    Uses the full reverse diffusion process with stochastic sampling.
    Requires all T timesteps for generation.
    """

    def __init__(
        self,
        config: DiffusionConfig,
        model: ConditionalUNet,
        betas: torch.Tensor,
        alphas_cumprod: torch.Tensor,
    ):
        super().__init__(config, model, betas, alphas_cumprod)

        # Precompute posterior variance for efficiency
        alphas_cumprod_prev = torch.nn.functional.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
        self.posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)

    def reverse_step(
        self, x_t: torch.Tensor, t: torch.Tensor, conditioning: torch.Tensor, add_noise: bool = True
    ) -> torch.Tensor:
        """Single DDPM reverse diffusion step.

        Args:
            x_t: Noisy data at timestep t
            t: Current timestep
            conditioning: Conditioning data
            add_noise: Whether to add noise (False for t=0)

        Returns:
            x_{t-1}: Less noisy data
        """
        # Predict noise
        if add_noise or x_t.requires_grad or conditioning.requires_grad:
            predicted_noise = self.model(x_t, t, conditioning)
        else:
            with torch.no_grad():
                predicted_noise = self.model(x_t, t, conditioning)

        # Extract coefficients
        alpha = self.alphas[t][:, None]
        alpha_cumprod = self.alphas_cumprod[t][:, None]
        beta = self.betas[t][:, None]

        # Compute mean
        mean = (x_t - beta * predicted_noise / torch.sqrt(1.0 - alpha_cumprod)) / torch.sqrt(alpha)

        if add_noise and t[0] > 0:
            # Add noise except for the last step
            posterior_variance = self.posterior_variance[t][:, None]
            noise = torch.randn_like(x_t)
            x_t_minus_1 = mean + torch.sqrt(posterior_variance) * noise
        else:
            x_t_minus_1 = mean

        return x_t_minus_1

    def sample(
        self,
        conditioning: torch.Tensor,
        n_samples: int = 1,
        require_grad: bool = False,
        initial_noise: torch.Tensor | None = None,
        show_progress: bool = True,
    ) -> torch.Tensor:
        """Generate samples using DDPM sampling.

        Args:
            conditioning: Conditioning data of shape [trajectory_length] or [batch_size, trajectory_length]
            n_samples: Number of samples to generate per conditioning
            require_grad: If True, enable gradient computation
            initial_noise: Optional fixed initial noise tensor. If provided, must have shape
                          [batch_size * n_samples, trajectory_length]. Useful for consistent
                          gradients in optimisation.
            show_progress: If True, show tqdm progress bar. Disable for optimisation loops.

        Returns:
            Generated samples of shape [n_samples, trajectory_length]
        """
        conditioning, x_t = self._prepare_sampling(
            conditioning, n_samples, require_grad, initial_noise
        )

        # Reverse diffusion
        timestep_iterator = range(self.config.n_timesteps - 1, -1, -1)
        if show_progress:
            timestep_iterator = tqdm(timestep_iterator, desc="DDPM Sampling")

        for t_idx in timestep_iterator:
            t = torch.full((conditioning.shape[0],), t_idx, dtype=torch.long).to(self.device)
            x_t = self.reverse_step(x_t, t, conditioning, add_noise=(t_idx > 0))

        return x_t


class DDIMSampler(DiffusionSampler):
    """DDIM sampling (Song et al., 2021).

    Deterministic (or semi-deterministic) sampling that allows for faster inference
    by skipping timesteps. Can generate samples in fewer steps than DDPM.

    Args:
        config: Diffusion configuration
        model: Conditional U-Net for denoising
        betas: Noise schedule
        alphas_cumprod: Cumulative product of alphas
        eta: Controls stochasticity (0 = deterministic, 1 = stochastic like DDPM)
        num_inference_steps: Number of steps to use (can be less than training steps)
    """

    def __init__(
        self,
        config: DiffusionConfig,
        model: ConditionalUNet,
        betas: torch.Tensor,
        alphas_cumprod: torch.Tensor,
        eta: float = 0.0,
        num_inference_steps: int | None = None,
    ):
        super().__init__(config, model, betas, alphas_cumprod)
        self.eta = eta

        # Use fewer steps for faster inference
        self.num_inference_steps = num_inference_steps or config.n_timesteps

        # Create timestep schedule (evenly spaced) while honouring the
        # requested step count exactly.  The schedule is stored in ascending
        # order and reversed at sampling time.
        if self.num_inference_steps < config.n_timesteps:
            self.timesteps = (
                torch.linspace(
                    config.n_timesteps - 1,
                    0,
                    steps=self.num_inference_steps,
                )
                .round()
                .long()
                .flip(0)
            )
            self.timesteps = torch.unique_consecutive(self.timesteps)
        else:
            self.timesteps = torch.arange(config.n_timesteps).long()

    def reverse_step(
        self, x_t: torch.Tensor, t: int, t_prev: int, conditioning: torch.Tensor
    ) -> torch.Tensor:
        """Single DDIM reverse step.

        Uses the deterministic (or semi-deterministic) DDIM update rule:
        x_{t-1} = sqrt(alpha_{t-1}) * pred_x0 + sqrt(1 - alpha_{t-1} - sigma_t^2) * epsilon + sigma_t * noise

        where pred_x0 = (x_t - sqrt(1 - alpha_t) * epsilon) / sqrt(alpha_t)

        Args:
            x_t: Noisy data at timestep t
            t: Current timestep index
            t_prev: Previous timestep index (can skip timesteps)
            conditioning: Conditioning data

        Returns:
            x_{t_prev}: Less noisy data at previous timestep
        """
        # Predict noise
        t_tensor = torch.full((x_t.shape[0],), t, dtype=torch.long, device=self.device)

        if x_t.requires_grad or conditioning.requires_grad:
            predicted_noise = self.model(x_t, t_tensor, conditioning)
        else:
            with torch.no_grad():
                predicted_noise = self.model(x_t, t_tensor, conditioning)

        # Get alpha values
        alpha_t = self.alphas_cumprod[t]
        alpha_t_prev = self.alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0)

        # Predict x_0 from x_t and predicted noise
        pred_x0 = (x_t - torch.sqrt(1.0 - alpha_t) * predicted_noise) / torch.sqrt(alpha_t)

        # Compute variance (eta controls stochasticity)
        sigma_t = self.eta * torch.sqrt(
            (1.0 - alpha_t_prev) / (1.0 - alpha_t) * (1.0 - alpha_t / alpha_t_prev)
        )

        # Compute direction pointing to x_t
        dir_xt = torch.sqrt(1.0 - alpha_t_prev - sigma_t**2) * predicted_noise

        # Compute x_{t-1}
        x_t_prev = torch.sqrt(alpha_t_prev) * pred_x0 + dir_xt

        # Add noise if eta > 0 and not at the last step
        if self.eta > 0 and t_prev >= 0:
            noise = torch.randn_like(x_t)
            x_t_prev = x_t_prev + sigma_t * noise

        return x_t_prev

    def sample(
        self,
        conditioning: torch.Tensor,
        n_samples: int = 1,
        require_grad: bool = False,
        initial_noise: torch.Tensor | None = None,
        show_progress: bool = True,
        grad_steps: int | None = None,
        grad_steps_spread: bool = False,
    ) -> torch.Tensor:
        """Generate samples using DDIM sampling.

        Args:
            conditioning: Conditioning data of shape [trajectory_length] or [batch_size, trajectory_length]
            n_samples: Number of samples to generate per conditioning
            require_grad: If True, enable gradient computation
            initial_noise: Optional fixed initial noise tensor. If provided, must have shape
                          [batch_size * n_samples, trajectory_length]. Useful for consistent
                          gradients in optimisation.
            show_progress: If True, show tqdm progress bar. Disable for optimisation loops.
            grad_steps: If set, restricts autograd to a subset of denoising steps.  By
                default (``grad_steps_spread=False``) these are the *final* ``grad_steps``
                steps (largest denoising effect, lowest noise).  If ``grad_steps_spread=True``
                the steps are chosen *linearly across the full chain*, capturing gradient
                signal at the start, middle, and end of denoising.
            grad_steps_spread: If ``True``, spread the ``grad_steps`` gradient-carrying
                steps uniformly across the chain instead of concentrating them at the end.

        Returns:
            Generated samples of shape [n_samples, trajectory_length]
        """
        conditioning, x_t = self._prepare_sampling(
            conditioning, n_samples, require_grad, initial_noise
        )

        if x_t.shape[0] != conditioning.shape[0]:
            raise ValueError(
                f"initial_noise batch size ({x_t.shape[0]}) must match "
                f"conditioning batch size ({conditioning.shape[0]})"
            )

        # Reverse diffusion with potentially fewer steps
        timesteps_reversed = self.timesteps.flip(0)
        n_total = len(timesteps_reversed)

        # Determine which step indices carry gradients
        if grad_steps is None:
            grad_indices: set[int] | None = None
        elif grad_steps_spread:
            # Linearly spaced across the full chain (inclusive of first and last)
            grad_indices = set(
                int(round(i))
                for i in torch.linspace(0, n_total - 1, min(grad_steps, n_total)).tolist()
            )
        else:
            # Final grad_steps steps only
            grad_indices = set(range(n_total - grad_steps, n_total))

        if show_progress:
            desc = f"DDIM Sampling ({self.num_inference_steps} steps)"
            timestep_iterator = enumerate(tqdm(timesteps_reversed, desc=desc))
        else:
            timestep_iterator = enumerate(timesteps_reversed)

        for i, t_idx in timestep_iterator:
            t = int(t_idx.item())
            # Get previous timestep (or -1 for the last step)
            t_prev = (
                int(timesteps_reversed[i + 1].item()) if i + 1 < len(timesteps_reversed) else -1
            )

            # Optionally restrict autograd graph to selected steps
            if grad_indices is not None and i not in grad_indices:
                with torch.no_grad():
                    x_t = self.reverse_step(x_t.detach(), t, t_prev, conditioning.detach())
            else:
                x_t = self.reverse_step(x_t, t, t_prev, conditioning)

        return x_t


def create_sampler(
    sampler_type: str,
    config: DiffusionConfig,
    model: ConditionalUNet,
    betas: torch.Tensor,
    alphas_cumprod: torch.Tensor,
    **kwargs,
) -> DiffusionSampler:
    """Factory function to create a sampler.

    Args:
        sampler_type: Type of sampler ('ddpm' or 'ddim')
        config: Diffusion configuration
        model: Conditional U-Net
        betas: Noise schedule
        alphas_cumprod: Cumulative product of alphas
        **kwargs: Additional arguments for specific samplers
            - For DDIM: eta (float), num_inference_steps (int)

    Returns:
        Sampler instance
    """
    sampler_type = sampler_type.lower()

    if sampler_type == "ddpm":
        return DDPMSampler(config, model, betas, alphas_cumprod)
    elif sampler_type == "ddim":
        return DDIMSampler(
            config,
            model,
            betas,
            alphas_cumprod,
            eta=kwargs.get("eta", 0.0),
            num_inference_steps=kwargs.get("num_inference_steps", None),
        )
    else:
        raise ValueError(f"Unknown sampler type: {sampler_type}. Choose 'ddpm' or 'ddim'.")