generative_models.py

"""Training and sampling of the conditional generative model (Appendix B.2).

Wraps the conditional flow-matching generator so `main.py` deals only in
ensembles.  The model is trained on the 60% training split, then sampled on the
calibration and test splits to produce, for every input, an ensemble of candidate
responses that the conformal methods operate on.

Standardisation matters here: the network is trained on standardised X and Y, and
the generated ensembles are mapped back to the original units before returning, so
that prediction set volumes are reported in the response's own units.
"""

import os

import numpy as np
import torch
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader

import flow_matching
from dataset import DatasetTensor, get_dataset


class GenerativeModel:
    """Conditional flow-matching generator over one dataset's train/calib/test split."""

    def __init__(self, args):
        self.args = args
        self.model_type = args.model_type
        self.X, self.Y = get_dataset(args.dataset, data_path=args.data_path)

        if self.model_type != 'flow-matching':
            raise ValueError(f'Unsupported model_type {self.model_type!r}; '
                             'this repository implements "flow-matching".')

        self.model = flow_matching.FlowMatchingNet(
            input_dim=self.Y.shape[1],
            condition_dim=self.X.shape[1],
            hidden_dim=args.hidden_dim,
        ).to(args.device)
        self.optimizer = torch.optim.Adam(self.model.parameters(), lr=args.lr)

        # Gaussian probability path y_t = t * y + sqrt(1 - t) * noise, whose
        # conditional vector field is the regression target during training.
        self.gaussian_path = flow_matching.GaussianPath(
            flow_matching.LinearAlpha(), flow_matching.SquareRootBeta())

    def prep_data(self):
        """Split 60/20/20, standardise, and build the three dataloaders."""
        # 60/20/20 split. Rows were already shuffled by get_dataset, so slicing in
        # order is a random split; the generator sees only the first 60%.
        N = self.X.shape[0]
        train, calib, test = np.split(range(N), [int(.6 * N), int(.8 * N)])
        print(f'train size: {len(train)}, calib size: {len(calib)}, '
              f'test size: {len(test)}', flush=True)

        # Fit the scalers on training data only; calibration and test inherit them.
        self.x_scaler = StandardScaler().fit(self.X[train])
        self.y_scaler = StandardScaler().fit(self.Y[train])

        # Only the training loader shuffles; calibration and test must keep their
        # order so ensembles line up with the observations they belong to.
        loaders = {}
        for name, idx, shuffle in [('train', train, True), ('calib', calib, False),
                                   ('test', test, False)]:
            dataset = DatasetTensor(self.x_scaler.transform(self.X[idx]),
                                    self.y_scaler.transform(self.Y[idx]))
            loaders[name] = DataLoader(dataset, batch_size=self.args.batch_size,
                                       shuffle=shuffle)

        self.train_loader = loaders['train']
        self.calib_loader = loaders['calib']
        self.test_loader = loaders['test']

    def train(self):
        """Fit the vector field by regressing on the conditional flow."""
        flow_matching.train_flow_matching(
            self.model, self.gaussian_path, self.train_loader,
            self.optimizer, self.args.n_epochs, self.args.device)

    def sample(self):
        """Draw ensembles for the calibration and test splits, in original units.

        Alongside the samples, each member carries the log-density of the Gaussian
        noise it was integrated from; HD-PCP uses these as per-sample confidence.
        """
        # One ensemble of n_samples responses per input, for both splits.
        calib_samples, calib_scores, calib_conditions = flow_matching.generate_samples_for_dataset(
            self.model, self.gaussian_path, self.calib_loader,
            self.args.n_samples, self.args.timesteps, self.args.device)
        test_samples, test_scores, test_conditions = flow_matching.generate_samples_for_dataset(
            self.model, self.gaussian_path, self.test_loader,
            self.args.n_samples, self.args.timesteps, self.args.device)

        # Undo standardisation so volumes are reported in the response's own units.
        unscale = lambda array, scaler: scaler.inverse_transform(
            array.reshape(-1, array.shape[-1])).reshape(array.shape)

        calib_samples = unscale(calib_samples, self.y_scaler)
        test_samples = unscale(test_samples, self.y_scaler)
        calib_conditions = unscale(calib_conditions, self.x_scaler)
        test_conditions = unscale(test_conditions, self.x_scaler)

        return (calib_samples, calib_scores, calib_conditions,
                test_samples, test_scores, test_conditions)

    def get_ground_truth(self):
        """Observed responses for the calibration and test splits, in original units."""
        N = self.X.shape[0]
        _, calib, test = np.split(range(N), [int(.6 * N), int(.8 * N)])
        return self.Y[calib], self.Y[test]

    def save(self):
        """Persist the trained network so later CP runs can skip training."""
        torch.save(self.model.state_dict(), os.path.join(self.args.model_path, 'model.pth'))

    def load(self):
        """Restore a previously trained network."""
        self.model.load_state_dict(
            torch.load(os.path.join(self.args.model_path, 'model.pth')))