flow_matching.py

"""Conditional flow-matching generative model (Appendix B.2).

The generator learns a velocity field that transports Gaussian noise to the target
conditional distribution P(Y | X).  Training pairs each observation y with noise
along the Gaussian probability path

    y_t = alpha(t) * y + beta(t) * noise,     alpha(t) = t,  beta(t) = sqrt(1 - t),

so that y_0 is pure noise and y_1 is the data.  The network regresses on the
conditional vector field of that path; at sampling time we start from noise and
integrate the learned field forward with Euler steps.

For this alpha/beta pair the conditional vector field reduces to the expression
quoted in the paper, v*(y_t, t) = y - noise / (2 sqrt(1 - t)); `conditional_vector_field`
implements the general form, which agrees with it.

The conformal methods only ever consume samples, so nothing downstream depends on
this particular generator -- any conditional sampler could be substituted.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F

import numpy as np
from scipy.stats import multivariate_normal


class GaussianPath:
    """Gaussian probability path y_t = alpha(t) * y + beta(t) * noise."""

    def __init__(self, alpha_fn, beta_fn):
        self.alpha = alpha_fn  # A function defining sigma(t)
        self.beta = beta_fn

    def sample_conditional_path(self, z, t):
        """Draw y_t on the path from a data point z at time t."""

        alpha_t = self.alpha(t)
        beta_t = self.beta(t)
        x_t = alpha_t * z + beta_t * torch.randn_like(z)
        return x_t

    def conditional_vector_field(self, x, z, t):
        """Velocity that transports the path, the regression target for the network.

        x is the noised point y_t, z the clean data point.  For alpha(t) = t and
        beta(t) = sqrt(1 - t) this equals z - noise / (2 sqrt(1 - t)).
        """

        cvf = (self.alpha.dt(t) - (self.beta.dt(t) / self.beta(t)) * self.alpha(t)) * z + (self.beta.dt(t) / self.beta(t)) * x
        return cvf


class LinearAlpha():
    """Data schedule alpha(t) = t, ramping the signal in linearly."""

    def __init__(self):
        pass

    def __call__(self, t):
        alpha_t = t
        return alpha_t

    def dt(self, t):
        return torch.ones_like(t)


class SquareRootBeta():
    """Noise schedule beta(t) = sqrt(1 - t), vanishing as the path reaches the data."""

    def __init__(self):
        pass

    def __call__(self, t):
        beta_t = torch.sqrt(1 - t)
        return beta_t

    def dt(self, t):
        return - 0.5 / (torch.sqrt(1 - t) + 1e-4)


class FlowMatchingNet(nn.Module):
    def __init__(self, input_dim, condition_dim, hidden_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim + condition_dim + 1, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, hidden_dim)
        self.fc5 = nn.Linear(hidden_dim, input_dim)

    def forward(self, y, condition, t):
        """Predict the velocity at (y, t) given the condition.

        y: (B, input_dim), condition: (B, condition_dim), t: (B, 1) or (B,).
        Returns (B, input_dim).
        """

        t = t.view(-1, 1)  # Ensure t has the right shape
        x = torch.cat([y, condition, t], dim=-1)  # Concatenate y, condition, and time
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        return self.fc5(x)


def flow_matching_loss(model, gaussian_path, z, condition, t, device):
    """Mean squared error between the predicted and true conditional velocity."""


    x_t = gaussian_path.sample_conditional_path(z, t)
    cvf = gaussian_path.conditional_vector_field(x_t, z, t)
    vf_pred = model(x_t, condition, t)

    # MSE loss
    return F.mse_loss(vf_pred, cvf)


def train_flow_matching(model, gaussian_path, dataloader, optimizer, num_epochs, device):
    """Train the velocity network, sampling a fresh t ~ Uniform[0, 1] per batch."""

    model.train()
    for epoch in range(num_epochs):
        Loss = 0  # Initialize loss for this epoch

        for batch in dataloader:
            condition, z = batch  # Original data and condition
            condition, z = condition.to(device), z.to(device)

            # Sample random time t
            t = torch.rand((z.size(0), 1), device=device)  # Uniformly sample t in [0, 1]

            # Compute loss
            loss = flow_matching_loss(model, gaussian_path, z, condition, t, device)

            # Backpropagation
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            Loss += loss.item()

        # print(f"Epoch {epoch + 1}, Loss: {Loss:.4f}", flush=True)


def generate_data(model, gaussian_path, condition, timesteps, device):
    """Integrate noise forward into samples for a batch of (repeated) conditions.

    Returns the generated samples and, for each, the log-density of the Gaussian
    noise it started from.  Those log-densities are the per-sample confidence
    scores HD-PCP filters on: the flow is a deterministic map out of the noise, so
    a sample seeded nearer the noise distribution's mode is one the model places in
    its high-density region.

    Note the loop takes one Euler step at each of `timesteps` grid points including
    t = 1, so it advances slightly past t = 1.  This is kept as-is because it is
    what produced the published results.
    """
    # Start from Gaussian noise
    input_dim = model.fc5.out_features
    x = torch.randn(condition.size(0), input_dim).to(device)
    scores = multivariate_normal.logpdf(x.cpu(), mean=np.zeros(input_dim), cov=np.eye(input_dim)).reshape(-1, 1)
    # (n_samples, 1)
    t_space = torch.linspace(0, 1, timesteps, device=device)

    for t in t_space:
        with torch.no_grad():
            dt = 1 / (timesteps - 1)
            u_t = model(x, condition, t * torch.ones(condition.size(0), 1).to(device))
            x = x + u_t * dt  # Euler step

    return x, scores


def generate_samples_for_dataset(model, gaussian_path, data_loader, num_samples, timesteps, device):
    """Generate an ensemble of `num_samples` responses for every input in a loader.

    Returns samples (N, num_samples, dim_y), their noise-seed log-densities
    (N, num_samples, 1), and the conditions (N, num_samples, dim_x).
    """
    model.eval()  # Set the model to evaluation mode
    all_generated_samples = []  # To store all generated samples
    all_generated_scores = [] # To store all generated scores
    conditions = []  # To store corresponding conditions

    for condition_batch in data_loader:
        X = condition_batch[0].to(device)  # Get condition inputs from calib_loader

        # Generate num_samples_per_condition for each condition in the batch
        for condition in X:
            condition = condition.unsqueeze(0).repeat(num_samples, 1)  # Repeat condition
            generated_samples, scores = generate_data(model, gaussian_path, condition, timesteps, device)
            all_generated_samples.append(generated_samples.cpu().numpy())
            conditions.append(condition.cpu().numpy())
            all_generated_scores.append(scores)

    # Convert lists to arrays for easier handling
    all_generated_samples = np.array(all_generated_samples)
    all_generated_scores = np.array(all_generated_scores)
    conditions = np.array(conditions)
    return all_generated_samples, all_generated_scores, conditions