From 21bb7e850fbb67bb2d314bdf1a7b522cd1477bef Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Thu, 16 Jul 2026 12:43:18 -0400
Subject: [PATCH 01/28] wip

---
 sdgym/synthesizers/tabddpm.py | 1224 +++++++++++++++++++++++++++++++++
 1 file changed, 1224 insertions(+)
 create mode 100644 sdgym/synthesizers/tabddpm.py

diff --git a/sdgym/synthesizers/tabddpm.py b/sdgym/synthesizers/tabddpm.py
new file mode 100644
index 00000000..6b353fe9
--- /dev/null
+++ b/sdgym/synthesizers/tabddpm.py
@@ -0,0 +1,1224 @@
+"""TabDDPMSynthesizer -- SDGym synthesizer built on TabDDPM.
+Paper: "TabDDPM: Modelling Tabular Data with Diffusion Models (2022)"
+https://arxiv.org/abs/2209.15421
+
+Original implementation is provided:
+https://github.com/yandex-research/tab-ddpm/tree/main/tab_ddpm
+"""
+
+import math
+from typing import Callable, Dict, List, Optional, Union
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from sklearn.preprocessing import OrdinalEncoder, QuantileTransformer
+from torch import Tensor
+
+
+CAT_MISSING_VALUE = '__nan__'
+
+
+class FoundNANsError(BaseException):
+    """Found NANs during sampling."""
+
+    def __init__(self, message='Found NANs during sampling.'):
+        super(FoundNANsError, self).__init__(message)
+
+
+def sum_except_batch(x, num_dims=1):
+    """Sum all dimensions except the first ``num_dims`` batch dimensions."""
+    return x.reshape(*x.shape[:num_dims], -1).sum(-1)
+
+
+def mean_flat(tensor):
+    """Take the mean over all non-batch dimensions."""
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def ohe_to_categories(ohe, K):
+    K = torch.from_numpy(K)
+    indices = torch.cat([torch.zeros((1,)), K.cumsum(dim=0)], dim=0).int().tolist()
+    res = []
+    for i in range(len(indices) - 1):
+        res.append(ohe[:, indices[i]:indices[i + 1]].argmax(dim=1))
+    return torch.stack(res, dim=1)
+
+
+def log_1_min_a(a):
+    return torch.log(1 - a.exp() + 1e-40)
+
+
+def log_add_exp(a, b):
+    maximum = torch.max(a, b)
+    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))
+
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    t = t.to(a.device)
+    out = a.gather(-1, t)
+    while len(out.shape) < len(x_shape):
+        out = out[..., None]
+    return out.expand(x_shape)
+
+
+def log_categorical(log_x_start, log_prob):
+    return (log_x_start.exp() * log_prob).sum(dim=1)
+
+
+def index_to_log_onehot(x, num_classes):
+    onehots = []
+    for i in range(len(num_classes)):
+        onehots.append(F.one_hot(x[:, i], num_classes[i]))
+
+    x_onehot = torch.cat(onehots, dim=1)
+    log_onehot = torch.log(x_onehot.float().clamp(min=1e-30))
+    return log_onehot
+
+
+@torch.jit.script
+def log_sub_exp(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    m = torch.maximum(a, b)
+    return torch.log(torch.exp(a - m) - torch.exp(b - m)) + m
+
+
+@torch.jit.script
+def sliced_logsumexp(x, slices):
+    lse = torch.logcumsumexp(
+        torch.nn.functional.pad(x, [1, 0, 0, 0], value=-float('inf')),
+        dim=-1)
+
+    slice_starts = slices[:-1]
+    slice_ends = slices[1:]
+
+    slice_lse = log_sub_exp(lse[:, slice_ends], lse[:, slice_starts])
+    slice_lse_repeated = torch.repeat_interleave(
+        slice_lse,
+        slice_ends - slice_starts,
+        dim=-1
+    )
+    return slice_lse_repeated
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """Get a pre-defined beta schedule for the given name."""
+    if schedule_name == 'linear':
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif schedule_name == 'cosine':
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f'unknown beta schedule: {schedule_name}')
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """Create sinusoidal timestep embeddings."""
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+class Block(nn.Module):
+    """The main building block of `MLP`."""
+
+    def __init__(
+        self,
+        *,
+        d_in: int,
+        d_out: int,
+        bias: bool,
+        activation: Union[str, Callable[..., nn.Module]],
+        dropout: float,
+    ) -> None:
+        super().__init__()
+        self.linear = nn.Linear(d_in, d_out, bias)
+        self.activation = getattr(nn, activation)() if isinstance(activation, str) else activation()
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.dropout(self.activation(self.linear(x)))
+
+class MLP(nn.Module):
+    """The MLP model from "Revisiting Deep Learning Models for Tabular Data".
+
+    MLP: (in) -> Block -> ... -> Block -> Linear -> (out)
+    Block: (in) -> Linear -> Activation -> Dropout -> (out)
+    """
+    def __init__(
+        self,
+        *,
+        d_in: int,
+        d_layers: List[int],
+        dropouts: Union[float, List[float]],
+        activation: Union[str, Callable[[], nn.Module]],
+        d_out: int,
+    ) -> None:
+        super().__init__()
+        if isinstance(dropouts, float):
+            dropouts = [dropouts] * len(d_layers)
+        assert len(d_layers) == len(dropouts)
+
+        self.blocks = nn.ModuleList(
+            [
+                MLP.Block(
+                    d_in=d_layers[i - 1] if i else d_in,
+                    d_out=d,
+                    bias=True,
+                    activation=activation,
+                    dropout=dropout,
+                )
+                for i, (d, dropout) in enumerate(zip(d_layers, dropouts))
+            ]
+        )
+        self.head = nn.Linear(d_layers[-1] if d_layers else d_in, d_out)
+
+    @classmethod
+    def make_baseline(
+        cls,
+        d_in: int,
+        d_layers: List[int],
+        dropout: float,
+        d_out: int,
+    ) -> 'MLP':
+        """Create a "baseline" MLP: ReLU activations, uniform dropout."""
+        assert isinstance(dropout, float)
+        if len(d_layers) > 2:
+            assert len(set(d_layers[1:-1])) == 1, (
+                'if d_layers contains more than two elements, then'
+                ' all elements except for the first and the last ones must be equal.'
+            )
+        return MLP(
+            d_in=d_in,
+            d_layers=d_layers,
+            dropouts=dropout,
+            activation='ReLU',
+            d_out=d_out,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x.float()
+        for block in self.blocks:
+            x = block(x)
+        x = self.head(x)
+        return x
+
+
+class MLPDiffusion(nn.Module):
+    """MLP denoiser with sinusoidal timestep embedding and optional label conditioning."""
+
+    def __init__(self, d_in, num_classes, is_y_cond, rtdl_params, dim_t=128):
+        super().__init__()
+        self.dim_t = dim_t
+        self.num_classes = num_classes
+        self.is_y_cond = is_y_cond
+
+        rtdl_params = dict(rtdl_params)
+        rtdl_params['d_in'] = dim_t
+        rtdl_params['d_out'] = d_in
+
+        self.mlp = MLP.make_baseline(**rtdl_params)
+
+        if self.num_classes > 0 and is_y_cond:
+            self.label_emb = nn.Embedding(self.num_classes, dim_t)
+        elif self.num_classes == 0 and is_y_cond:
+            self.label_emb = nn.Linear(1, dim_t)
+
+        self.proj = nn.Linear(d_in, dim_t)
+        self.time_embed = nn.Sequential(
+            nn.Linear(dim_t, dim_t),
+            nn.SiLU(),
+            nn.Linear(dim_t, dim_t)
+        )
+
+    def forward(self, x, timesteps, y=None):
+        emb = self.time_embed(timestep_embedding(timesteps, self.dim_t))
+        if self.is_y_cond and y is not None:
+            if self.num_classes > 0:
+                y = y.squeeze()
+            else:
+                y = y.resize(y.size(0), 1).float()
+            emb += F.silu(self.label_emb(y))
+        x = self.proj(x) + emb
+        return self.mlp(x)
+
+
+# =====================================================================
+# Gaussian + multinomial diffusion
+# (from tab_ddpm/gaussian_multinomial_diffsuion.py, trimmed to the paths
+# actually used by the paper's pipeline: 'mse' Gaussian loss, 'eps'
+# parametrization, 'vb_stochastic' multinomial loss, uniform time sampling
+# and ancestral sampling)
+# =====================================================================
+
+class GaussianMultinomialDiffusion(torch.nn.Module):
+    """Joint diffusion: Gaussian over numerical features, multinomial over categorical."""
+
+    def __init__(
+            self,
+            num_classes: np.ndarray,
+            num_numerical_features: int,
+            denoise_fn,
+            num_timesteps=1000,
+            scheduler='cosine',
+            device=torch.device('cpu')
+    ):
+        super(GaussianMultinomialDiffusion, self).__init__()
+
+        self.num_numerical_features = num_numerical_features
+        self.num_classes = num_classes  # it is a vector [K1, K2, ..., Km]
+        self.num_classes_expanded = torch.from_numpy(
+            np.concatenate([num_classes[i].repeat(num_classes[i]) for i in range(len(num_classes))])
+        ).to(device)
+
+        self.slices_for_classes = [np.arange(self.num_classes[0])]
+        offsets = np.cumsum(self.num_classes)
+        for i in range(1, len(offsets)):
+            self.slices_for_classes.append(np.arange(offsets[i - 1], offsets[i]))
+        self.offsets = torch.from_numpy(np.append([0], offsets)).to(device)
+
+        self._denoise_fn = denoise_fn
+        self.num_timesteps = num_timesteps
+        self.scheduler = scheduler
+
+        alphas = 1. - get_named_beta_schedule(scheduler, num_timesteps)
+        alphas = torch.tensor(alphas.astype('float64'))
+        betas = 1. - alphas
+
+        log_alpha = np.log(alphas)
+        log_cumprod_alpha = np.cumsum(log_alpha)
+
+        log_1_min_alpha = log_1_min_a(log_alpha)
+        log_1_min_cumprod_alpha = log_1_min_a(log_cumprod_alpha)
+
+        alphas_cumprod = np.cumprod(alphas, axis=0)
+        alphas_cumprod_prev = torch.tensor(np.append(1.0, alphas_cumprod[:-1]))
+        alphas_cumprod_next = torch.tensor(np.append(alphas_cumprod[1:], 0.0))
+        sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
+        sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - alphas_cumprod)
+        sqrt_recip_alphas_cumprod = np.sqrt(1.0 / alphas_cumprod)
+        sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / alphas_cumprod - 1)
+
+        # Gaussian diffusion
+        self.posterior_variance = (
+            betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        )
+        self.posterior_log_variance_clipped = torch.from_numpy(
+            np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:]))
+        ).float().to(device)
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        ).float().to(device)
+        self.posterior_mean_coef2 = (
+            (1.0 - alphas_cumprod_prev)
+            * np.sqrt(alphas.numpy())
+            / (1.0 - alphas_cumprod)
+        ).float().to(device)
+
+        assert log_add_exp(log_alpha, log_1_min_alpha).abs().sum().item() < 1.e-5
+        assert log_add_exp(log_cumprod_alpha, log_1_min_cumprod_alpha).abs().sum().item() < 1e-5
+        assert (np.cumsum(log_alpha) - log_cumprod_alpha).abs().sum().item() < 1.e-5
+
+        # Convert to float32 and register buffers.
+        self.register_buffer('alphas', alphas.float().to(device))
+        self.register_buffer('log_alpha', log_alpha.float().to(device))
+        self.register_buffer('log_1_min_alpha', log_1_min_alpha.float().to(device))
+        self.register_buffer('log_1_min_cumprod_alpha', log_1_min_cumprod_alpha.float().to(device))
+        self.register_buffer('log_cumprod_alpha', log_cumprod_alpha.float().to(device))
+        self.register_buffer('alphas_cumprod', alphas_cumprod.float().to(device))
+        self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float().to(device))
+        self.register_buffer('alphas_cumprod_next', alphas_cumprod_next.float().to(device))
+        self.register_buffer('sqrt_alphas_cumprod', sqrt_alphas_cumprod.float().to(device))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', sqrt_one_minus_alphas_cumprod.float().to(device))
+        self.register_buffer('sqrt_recip_alphas_cumprod', sqrt_recip_alphas_cumprod.float().to(device))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', sqrt_recipm1_alphas_cumprod.float().to(device))
+
+
+    def gaussian_q_sample(self, x_start, t, noise=None):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+            * noise
+        )
+
+    def gaussian_q_posterior_mean_variance(self, x_start, x_t, t):
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = extract(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def gaussian_p_mean_variance(self, model_output, x, t):
+        B = x.shape[0]
+        assert t.shape == (B,)
+
+        model_variance = torch.cat(
+            [self.posterior_variance[1].unsqueeze(0).to(x.device), (1. - self.alphas)[1:]], dim=0
+        )
+        model_log_variance = torch.log(model_variance)
+
+        model_variance = extract(model_variance, t, x.shape)
+        model_log_variance = extract(model_log_variance, t, x.shape)
+
+        # 'eps' parametrization: the network predicts the noise
+        pred_xstart = self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+
+        model_mean, _, _ = self.gaussian_q_posterior_mean_variance(
+            x_start=pred_xstart, x_t=x, t=t
+        )
+
+        return {
+            'mean': model_mean,
+            'variance': model_variance,
+            'log_variance': model_log_variance,
+            'pred_xstart': pred_xstart,
+        }
+
+    def _gaussian_loss(self, model_out, noise):
+        # 'mse' loss: the network predicts the added noise
+        return mean_flat((noise - model_out) ** 2)
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def gaussian_p_sample(self, model_out, x, t):
+        out = self.gaussian_p_mean_variance(model_out, x, t)
+        noise = torch.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+
+        sample = out['mean'] + nonzero_mask * torch.exp(0.5 * out['log_variance']) * noise
+        return {'sample': sample, 'pred_xstart': out['pred_xstart']}
+
+
+    def multinomial_kl(self, log_prob1, log_prob2):
+        kl = (log_prob1.exp() * (log_prob1 - log_prob2)).sum(dim=1)
+        return kl
+
+    def q_pred_one_timestep(self, log_x_t, t):
+        log_alpha_t = extract(self.log_alpha, t, log_x_t.shape)
+        log_1_min_alpha_t = extract(self.log_1_min_alpha, t, log_x_t.shape)
+
+        # alpha_t * E[xt] + (1 - alpha_t) 1 / K
+        log_probs = log_add_exp(
+            log_x_t + log_alpha_t,
+            log_1_min_alpha_t - torch.log(self.num_classes_expanded)
+        )
+
+        return log_probs
+
+    def q_pred(self, log_x_start, t):
+        log_cumprod_alpha_t = extract(self.log_cumprod_alpha, t, log_x_start.shape)
+        log_1_min_cumprod_alpha = extract(self.log_1_min_cumprod_alpha, t, log_x_start.shape)
+
+        log_probs = log_add_exp(
+            log_x_start + log_cumprod_alpha_t,
+            log_1_min_cumprod_alpha - torch.log(self.num_classes_expanded)
+        )
+
+        return log_probs
+
+    def predict_start(self, model_out, log_x_t):
+        assert model_out.size(0) == log_x_t.size(0)
+        assert model_out.size(1) == self.num_classes.sum(), f'{model_out.size()}'
+
+        log_pred = torch.empty_like(model_out)
+        for ix in self.slices_for_classes:
+            log_pred[:, ix] = F.log_softmax(model_out[:, ix], dim=1)
+        return log_pred
+
+    def q_posterior(self, log_x_start, log_x_t, t):
+        # q(xt-1 | xt, x0) = q(xt | xt-1, x0) * q(xt-1 | x0) / q(xt | x0)
+        # where q(xt | xt-1, x0) = q(xt | xt-1).
+        
+        t_minus_1 = t - 1
+        # Remove negative values, will not be used anyway for final decoder
+        t_minus_1 = torch.where(t_minus_1 < 0, torch.zeros_like(t_minus_1), t_minus_1)
+        log_EV_qxtmin_x0 = self.q_pred(log_x_start, t_minus_1)
+
+        num_axes = (1,) * (len(log_x_start.size()) - 1)
+        t_broadcast = t.to(log_x_start.device).view(-1, *num_axes) * torch.ones_like(log_x_start)
+        log_EV_qxtmin_x0 = torch.where(t_broadcast == 0, log_x_start, log_EV_qxtmin_x0.to(torch.float32))
+
+        # Note: the formula uses log q_pred_one_timestep(x_t, t), _NOT_ x_tmin1.
+        unnormed_logprobs = log_EV_qxtmin_x0 + self.q_pred_one_timestep(log_x_t, t)
+
+        log_EV_xtmin_given_xt_given_xstart = \
+            unnormed_logprobs \
+            - sliced_logsumexp(unnormed_logprobs, self.offsets)
+
+        return log_EV_xtmin_given_xt_given_xstart
+
+    def p_pred(self, model_out, log_x, t):
+        # 'x0' parametrization
+        log_x_recon = self.predict_start(model_out, log_x)
+        log_model_pred = self.q_posterior(
+            log_x_start=log_x_recon, log_x_t=log_x, t=t)
+        return log_model_pred
+
+    @torch.no_grad()
+    def p_sample(self, model_out, log_x, t):
+        model_log_prob = self.p_pred(model_out, log_x=log_x, t=t)
+        out = self.log_sample_categorical(model_log_prob)
+        return out
+
+    def log_sample_categorical(self, logits):
+        full_sample = []
+        for i in range(len(self.num_classes)):
+            one_class_logits = logits[:, self.slices_for_classes[i]]
+            uniform = torch.rand_like(one_class_logits)
+            gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
+            sample = (gumbel_noise + one_class_logits).argmax(dim=1)
+            full_sample.append(sample.unsqueeze(1))
+        full_sample = torch.cat(full_sample, dim=1)
+        log_sample = index_to_log_onehot(full_sample, self.num_classes)
+        return log_sample
+
+    def q_sample(self, log_x_start, t):
+        log_EV_qxt_x0 = self.q_pred(log_x_start, t)
+        log_sample = self.log_sample_categorical(log_EV_qxt_x0)
+        return log_sample
+
+    def kl_prior(self, log_x_start):
+        b = log_x_start.size(0)
+        device = log_x_start.device
+        ones = torch.ones(b, device=device).long()
+
+        log_qxT_prob = self.q_pred(log_x_start, t=(self.num_timesteps - 1) * ones)
+        log_half_prob = -torch.log(self.num_classes_expanded * torch.ones_like(log_qxT_prob))
+
+        kl_prior = self.multinomial_kl(log_qxT_prob, log_half_prob)
+        return sum_except_batch(kl_prior)
+
+    def compute_Lt(self, model_out, log_x_start, log_x_t, t, detach_mean=False):
+        log_true_prob = self.q_posterior(
+            log_x_start=log_x_start, log_x_t=log_x_t, t=t)
+        log_model_prob = self.p_pred(model_out, log_x=log_x_t, t=t)
+
+        if detach_mean:
+            log_model_prob = log_model_prob.detach()
+
+        kl = self.multinomial_kl(log_true_prob, log_model_prob)
+        kl = sum_except_batch(kl)
+
+        decoder_nll = -log_categorical(log_x_start, log_model_prob)
+        decoder_nll = sum_except_batch(decoder_nll)
+
+        mask = (t == torch.zeros_like(t)).float()
+        loss = mask * decoder_nll + (1. - mask) * kl
+
+        return loss
+
+    def sample_time(self, b, device):
+        # uniform time sampling
+        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
+        pt = torch.ones_like(t).float() / self.num_timesteps
+        return t, pt
+
+    def _multinomial_loss(self, model_out, log_x_start, log_x_t, t, pt):
+        # 'vb_stochastic' loss
+        kl = self.compute_Lt(model_out, log_x_start, log_x_t, t)
+        kl_prior = self.kl_prior(log_x_start)
+        # Upweigh loss term of the kl
+        vb_loss = kl / pt + kl_prior
+        return vb_loss
+
+
+    def mixed_loss(self, x, out_dict):
+        b = x.shape[0]
+        device = x.device
+        t, pt = self.sample_time(b, device)
+
+        x_num = x[:, :self.num_numerical_features]
+        x_cat = x[:, self.num_numerical_features:]
+
+        x_num_t = x_num
+        log_x_cat_t = x_cat
+        noise = None
+        if x_num.shape[1] > 0:
+            noise = torch.randn_like(x_num)
+            x_num_t = self.gaussian_q_sample(x_num, t, noise=noise)
+        if x_cat.shape[1] > 0:
+            log_x_cat = index_to_log_onehot(x_cat.long(), self.num_classes)
+            log_x_cat_t = self.q_sample(log_x_start=log_x_cat, t=t)
+
+        x_in = torch.cat([x_num_t, log_x_cat_t], dim=1)
+
+        model_out = self._denoise_fn(x_in, t, **out_dict)
+
+        model_out_num = model_out[:, :self.num_numerical_features]
+        model_out_cat = model_out[:, self.num_numerical_features:]
+
+        loss_multi = torch.zeros((1,), device=device).float()
+        loss_gauss = torch.zeros((1,), device=device).float()
+        if x_cat.shape[1] > 0:
+            loss_multi = self._multinomial_loss(
+                model_out_cat, log_x_cat, log_x_cat_t, t, pt
+            ) / len(self.num_classes)
+
+        if x_num.shape[1] > 0:
+            loss_gauss = self._gaussian_loss(model_out_num, noise)
+
+        return loss_multi.mean(), loss_gauss.mean()
+
+    @torch.no_grad()
+    def sample(self, num_samples, y_dist):
+        b = num_samples
+        device = self.log_alpha.device
+        z_norm = torch.randn((b, self.num_numerical_features), device=device)
+
+        has_cat = self.num_classes[0] != 0
+        log_z = torch.zeros((b, 0), device=device).float()
+        if has_cat:
+            uniform_logits = torch.zeros((b, len(self.num_classes_expanded)), device=device)
+            log_z = self.log_sample_categorical(uniform_logits)
+
+        y = torch.multinomial(
+            y_dist,
+            num_samples=b,
+            replacement=True
+        )
+        out_dict = {'y': y.long().to(device)}
+        for i in reversed(range(0, self.num_timesteps)):
+            t = torch.full((b,), i, device=device, dtype=torch.long)
+            model_out = self._denoise_fn(
+                torch.cat([z_norm, log_z], dim=1).float(),
+                t,
+                **out_dict
+            )
+            model_out_num = model_out[:, :self.num_numerical_features]
+            model_out_cat = model_out[:, self.num_numerical_features:]
+            if self.num_numerical_features > 0:
+                z_norm = self.gaussian_p_sample(model_out_num, z_norm, t)['sample']
+            if has_cat:
+                log_z = self.p_sample(model_out_cat, log_z, t)
+
+        z_ohe = torch.exp(log_z).round()
+        z_cat = log_z
+        if has_cat:
+            z_cat = ohe_to_categories(z_ohe, self.num_classes)
+        sample = torch.cat([z_norm, z_cat], dim=1).cpu()
+        return sample, out_dict
+
+    def sample_all(self, num_samples, batch_size, y_dist, verbose=False):
+        all_y = []
+        all_samples = []
+        num_generated = 0
+        max_attempts = 10 * math.ceil(num_samples / batch_size) + 10
+        attempts = 0
+        while num_generated < num_samples:
+            if attempts >= max_attempts:
+                raise FoundNANsError(
+                    'Sampling keeps producing NaNs; the model may be undertrained '
+                    'or the learning rate too high.'
+                )
+            attempts += 1
+
+            b = min(batch_size, num_samples - num_generated)
+            sample, out_dict = self.sample(b, y_dist)
+            mask_nan = torch.any(sample.isnan(), dim=1)
+            sample = sample[~mask_nan]
+            y = out_dict['y'][~mask_nan]
+
+            all_samples.append(sample)
+            all_y.append(y.cpu())
+            num_generated += sample.shape[0]
+            if verbose:
+                print(f'Sampled {min(num_generated, num_samples)}/{num_samples} rows', end='\r')
+
+        if verbose:
+            print()
+        x_gen = torch.cat(all_samples, dim=0)[:num_samples]
+        y_gen = torch.cat(all_y, dim=0)[:num_samples]
+
+        return x_gen, y_gen
+
+
+# =====================================================================
+# Data transformer: SDV metadata-driven DataFrame <-> model matrix
+# (replaces lib/data.py's Dataset/Transformations for the DataFrame API;
+# uses the paper's preprocessing: quantile normalization for numerical
+# features and ordinal encoding for categorical features)
+# =====================================================================
+
+class _DataTransformer:
+    """Converts a DataFrame into (numerical block, categorical index block) and back."""
+
+    _NUMERICAL_ROLES = ('numerical', 'datetime')
+
+    def __init__(self, columns_metadata, normalization='quantile', seed=0):
+        self._columns_metadata = columns_metadata
+        self._normalization = normalization
+        self._seed = seed
+
+        self._roles = {}
+        for column, spec in columns_metadata.items():
+            sdtype = spec.get('sdtype', 'categorical')
+            role = sdtype
+            if role not in ('numerical', 'datetime', 'boolean', 'id'):
+                # 'categorical' and any unrecognized/PII sdtype
+                role = 'categorical'
+            self._roles[column] = role
+
+        self.num_columns: List[str] = []
+        self.cat_columns: List[str] = []
+        self.id_columns: List[str] = []
+
+    # -- fitting -------------------------------------------------------
+
+    def fit(self, df: pd.DataFrame) -> None:
+        for column in df.columns:
+            role = self._roles[column]
+            if role in self._NUMERICAL_ROLES:
+                self.num_columns.append(column)
+            elif role == 'id':
+                self.id_columns.append(column)
+            else:
+                self.cat_columns.append(column)
+
+        self._dtypes = {column: df[column].dtype for column in df.columns}
+        self._id_is_numeric = {
+            column: pd.api.types.is_numeric_dtype(df[column]) for column in self.id_columns
+        }
+
+        # Numerical block: raw floats (datetimes as epoch nanoseconds)
+        X_num = self._to_numeric_block(df)
+
+        self._num_means = None
+        self._num_transform = None
+        self._disc_uniques = {}
+        if X_num.shape[1] > 0:
+            col_means = np.nanmean(X_num, axis=0)
+            col_means = np.where(np.isnan(col_means), 0.0, col_means)
+            self._num_means = col_means
+            inds = np.where(np.isnan(X_num))
+            X_num[inds] = np.take(col_means, inds[1])
+
+            # Columns holding few unique integer values are snapped back to the
+            # observed values after sampling (same heuristic as scripts/sample.py).
+            for j, column in enumerate(self.num_columns):
+                if self._roles[column] != 'numerical':
+                    continue
+                uniq = np.unique(X_num[:, j])
+                if len(uniq) <= 32 and np.allclose(uniq, np.round(uniq)):
+                    self._disc_uniques[j] = uniq
+
+            if self._normalization is not None:
+                self._num_transform = QuantileTransformer(
+                    output_distribution='normal',
+                    n_quantiles=max(min(X_num.shape[0] // 30, 1000), 10),
+                    subsample=int(1e9),
+                    random_state=self._seed,
+                )
+                self._num_transform.fit(X_num)
+
+        # Categorical block: ordinal codes ('__nan__' is its own category)
+        self._cat_transform = None
+        self.category_sizes = np.array([0])
+        if self.cat_columns:
+            X_cat = self._to_categorical_block(df)
+            self._cat_transform = OrdinalEncoder(dtype='int64')
+            self._cat_transform.fit(X_cat)
+            self.category_sizes = np.array(
+                [len(categories) for categories in self._cat_transform.categories_]
+            )
+
+    def transform(self, df: pd.DataFrame):
+        X_num = self._to_numeric_block(df)
+        if X_num.shape[1] > 0:
+            inds = np.where(np.isnan(X_num))
+            X_num[inds] = np.take(self._num_means, inds[1])
+            if self._num_transform is not None:
+                X_num = self._num_transform.transform(X_num)
+
+        X_cat = np.empty((len(df), 0), dtype='int64')
+        if self.cat_columns:
+            X_cat = self._cat_transform.transform(self._to_categorical_block(df))
+
+        return X_num.astype('float32'), X_cat
+
+    # -- inverting -----------------------------------------------------
+
+    def inverse_transform(self, X_num: np.ndarray, X_cat: np.ndarray) -> pd.DataFrame:
+        n_rows = max(X_num.shape[0], X_cat.shape[0])
+        columns = {}
+
+        if self.num_columns:
+            if self._num_transform is not None:
+                X_num = self._num_transform.inverse_transform(X_num)
+            for j, uniq in self._disc_uniques.items():
+                dist = np.abs(X_num[:, j][:, None] - uniq[None, :])
+                X_num[:, j] = uniq[dist.argmin(axis=1)]
+            for j, column in enumerate(self.num_columns):
+                columns[column] = self._from_numeric_column(column, X_num[:, j])
+
+        if self.cat_columns:
+            decoded = self._cat_transform.inverse_transform(
+                np.round(X_cat).astype('int64')
+            )
+            for j, column in enumerate(self.cat_columns):
+                columns[column] = self._from_categorical_column(column, decoded[:, j])
+
+        for column in self.id_columns:
+            if self._id_is_numeric[column]:
+                columns[column] = np.arange(n_rows)
+            else:
+                columns[column] = np.array([f'{column}_{i}' for i in range(n_rows)])
+
+        return pd.DataFrame(columns)
+
+    # -- per-column helpers ---------------------------------------------
+
+    def _to_numeric_block(self, df: pd.DataFrame) -> np.ndarray:
+        parts = []
+        for column in self.num_columns:
+            if self._roles[column] == 'datetime':
+                fmt = self._columns_metadata[column].get('datetime_format')
+                series = pd.to_datetime(df[column], format=fmt, errors='coerce')
+                values = series.values.astype('int64').astype('float64')
+                values[series.isna().values] = np.nan
+            else:
+                values = pd.to_numeric(df[column], errors='coerce').astype('float64').values
+            parts.append(values)
+        if not parts:
+            return np.empty((len(df), 0), dtype='float64')
+        return np.column_stack(parts)
+
+    def _to_categorical_block(self, df: pd.DataFrame) -> np.ndarray:
+        parts = []
+        for column in self.cat_columns:
+            series = df[column]
+            values = series.astype(object).where(series.notna(), CAT_MISSING_VALUE).astype(str)
+            parts.append(values.values)
+        return np.column_stack(parts)
+
+    def _from_numeric_column(self, column: str, values: np.ndarray):
+        if self._roles[column] == 'datetime':
+            return pd.to_datetime(np.round(values).astype('int64'))
+        dtype = self._dtypes[column]
+        if pd.api.types.is_integer_dtype(dtype):
+            return np.round(values).astype(dtype)
+        return values.astype('float64')
+
+    def _from_categorical_column(self, column: str, values: np.ndarray):
+        series = pd.Series(values, dtype=object)
+        series = series.where(series != CAT_MISSING_VALUE, np.nan)
+        if self._roles[column] == 'boolean':
+            return series.map({'True': True, 'False': False})
+        try:
+            return series.astype(self._dtypes[column])
+        except (ValueError, TypeError):
+            return series
+
+
+# =====================================================================
+# Trainer (from scripts/train.py, without EMA -- the original pipeline
+# samples from the non-EMA weights)
+# =====================================================================
+
+class _FastTensorDataLoader:
+    """Infinite iterator over shuffled (X, y) batches; faster than DataLoader for tensors."""
+
+    def __init__(self, X, y, batch_size):
+        self.X = X
+        self.y = y
+        self.batch_size = min(batch_size, X.shape[0])
+
+    def __iter__(self):
+        while True:
+            perm = torch.randperm(self.X.shape[0])
+            X, y = self.X[perm], self.y[perm]
+            for i in range(0, X.shape[0], self.batch_size):
+                yield X[i:i + self.batch_size], y[i:i + self.batch_size]
+
+
+class _Trainer:
+    def __init__(self, diffusion, train_iter, lr, weight_decay, steps, device, verbose=True):
+        self.diffusion = diffusion
+        self.train_iter = iter(train_iter)
+        self.steps = steps
+        self.init_lr = lr
+        self.optimizer = torch.optim.AdamW(
+            self.diffusion.parameters(), lr=lr, weight_decay=weight_decay
+        )
+        self.device = device
+        self.verbose = verbose
+        self.log_every = 100
+        self.loss_history = []
+
+    def _anneal_lr(self, step):
+        frac_done = step / self.steps
+        lr = self.init_lr * (1 - frac_done)
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
+
+    def _run_step(self, x, out_dict):
+        x = x.to(self.device)
+        for k in out_dict:
+            out_dict[k] = out_dict[k].long().to(self.device)
+        self.optimizer.zero_grad()
+        loss_multi, loss_gauss = self.diffusion.mixed_loss(x, out_dict)
+        loss = loss_multi + loss_gauss
+        loss.backward()
+        self.optimizer.step()
+        return loss_multi, loss_gauss
+
+    def run_loop(self):
+        curr_loss_multi = 0.0
+        curr_loss_gauss = 0.0
+        curr_count = 0
+
+        for step in range(self.steps):
+            x, y = next(self.train_iter)
+            batch_loss_multi, batch_loss_gauss = self._run_step(x, {'y': y})
+
+            self._anneal_lr(step)
+
+            curr_count += len(x)
+            curr_loss_multi += batch_loss_multi.item() * len(x)
+            curr_loss_gauss += batch_loss_gauss.item() * len(x)
+
+            if (step + 1) % self.log_every == 0:
+                mloss = np.around(curr_loss_multi / curr_count, 4)
+                gloss = np.around(curr_loss_gauss / curr_count, 4)
+                self.loss_history.append(
+                    {'step': step + 1, 'mloss': mloss, 'gloss': gloss, 'loss': mloss + gloss}
+                )
+                if self.verbose:
+                    print(f'Step {step + 1}/{self.steps} MLoss: {mloss} GLoss: {gloss} '
+                          f'Sum: {np.around(mloss + gloss, 4)}')
+                curr_count = 0
+                curr_loss_multi = 0.0
+                curr_loss_gauss = 0.0
+
+
+# =====================================================================
+# Public synthesizer
+# =====================================================================
+
+class TabDDPMSynthesizer:
+    """SDV-style single-table synthesizer based on TabDDPM (arXiv:2209.15421).
+
+    Args:
+        metadata:
+            SDV metadata describing the table -- either an SDV metadata object
+            (``SingleTableMetadata`` / ``Metadata``, anything exposing
+            ``to_dict()``) or the equivalent dictionary with a ``'columns'``
+            key (single-table format) or a ``'tables'`` key (multi-table
+            format holding exactly one table).
+        hyperparameters:
+            Optional dict overriding any of ``DEFAULT_HYPERPARAMETERS``:
+
+            - ``target_column`` (str or None): categorical/boolean column to
+              condition the diffusion on (the paper's class-conditional setup
+              for classification datasets). ``None`` trains unconditionally.
+            - ``d_layers`` (list of int): hidden layer sizes of the MLP denoiser.
+            - ``dropout`` (float): dropout of the MLP denoiser.
+            - ``dim_t`` (int): timestep/label embedding dimension.
+            - ``num_timesteps`` (int): diffusion timesteps T.
+            - ``scheduler`` (str): ``'cosine'`` or ``'linear'`` beta schedule.
+            - ``steps`` (int): training iterations.
+            - ``lr``, ``weight_decay``, ``batch_size``: AdamW optimizer settings.
+            - ``normalization`` (str or None): ``'quantile'`` (paper default)
+              or ``None`` to skip normalizing numerical features.
+            - ``sample_batch_size`` (int): batch size used during sampling.
+            - ``device`` (str or None): e.g. ``'cuda'``/``'cpu'``;
+              ``None`` auto-selects.
+            - ``seed`` (int): random seed used for fitting.
+            - ``verbose`` (bool): print training/sampling progress.
+    """
+
+    DEFAULT_HYPERPARAMETERS = {
+        'target_column': None,
+        'd_layers': [256, 256],
+        'dropout': 0.0,
+        'dim_t': 128,
+        'num_timesteps': 1000,
+        'scheduler': 'cosine',
+        'steps': 1000,
+        'lr': 0.001,
+        'weight_decay': 1e-5,
+        'batch_size': 4096,
+        'normalization': 'quantile',
+        'sample_batch_size': 10000,
+        'device': None,
+        'seed': 0,
+        'verbose': True,
+    }
+
+    def __init__(self, metadata, hyperparameters: Optional[dict] = None):
+        self._table_name, self._table_metadata = self._parse_metadata(metadata)
+
+        hyperparameters = dict(hyperparameters or {})
+        unknown = set(hyperparameters) - set(self.DEFAULT_HYPERPARAMETERS)
+        if unknown:
+            raise ValueError(
+                f'Unknown hyperparameters: {sorted(unknown)}. '
+                f'Valid keys are: {sorted(self.DEFAULT_HYPERPARAMETERS)}.'
+            )
+        self.hyperparameters = {**self.DEFAULT_HYPERPARAMETERS, **hyperparameters}
+
+        device = self.hyperparameters['device']
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self._device = torch.device(device)
+
+        self._fitted = False
+
+    # -- public API ------------------------------------------------------
+
+    def fit(self, data: Dict[str, pd.DataFrame]) -> None:
+        """Fit the synthesizer on real data.
+
+        Args:
+            data: dictionary mapping table names to pandas DataFrames.
+                Must contain exactly one table matching the metadata.
+        """
+        table_name, df = self._validate_data(data)
+        self._table_name = table_name
+
+        hp = self.hyperparameters
+        torch.manual_seed(hp['seed'])
+        np.random.seed(hp['seed'])
+
+        # Target column for class-conditional training (paper's setup for
+        # classification datasets); everything else goes through the transformer.
+        target_column = hp['target_column']
+        columns_metadata = dict(self._table_metadata['columns'])
+        self._n_target_classes = 0
+        self._target_encoder = None
+        if target_column is not None:
+            if target_column not in df.columns:
+                raise ValueError(f"target_column '{target_column}' not found in the data.")
+            sdtype = columns_metadata[target_column].get('sdtype', 'categorical')
+            if sdtype not in ('categorical', 'boolean'):
+                raise ValueError(
+                    f"target_column '{target_column}' must be categorical or boolean "
+                    f'(got sdtype={sdtype!r}). Numerical columns are modelled '
+                    'unconditionally; leave target_column unset.'
+                )
+            target_series = df[target_column]
+            target_values = (
+                target_series.astype(object)
+                .where(target_series.notna(), CAT_MISSING_VALUE)
+                .astype(str)
+                .values.reshape(-1, 1)
+            )
+            self._target_encoder = OrdinalEncoder(dtype='int64')
+            y = self._target_encoder.fit_transform(target_values).reshape(-1)
+            self._n_target_classes = len(self._target_encoder.categories_[0])
+            self._target_dtype = target_series.dtype
+            self._target_is_boolean = sdtype == 'boolean'
+            columns_metadata.pop(target_column)
+            df = df.drop(columns=[target_column])
+        else:
+            y = np.zeros(len(df), dtype='int64')
+
+        # Preprocess: quantile-normalized numerical block + ordinal categorical block
+        self._transformer = _DataTransformer(
+            columns_metadata, normalization=hp['normalization'], seed=hp['seed']
+        )
+        self._transformer.fit(df)
+        X_num, X_cat = self._transformer.transform(df)
+
+        n_num = X_num.shape[1]
+        K = self._transformer.category_sizes
+        d_in = int(n_num + K.sum())
+        if d_in == 0:
+            raise ValueError('The table has no modelable columns (only id columns?).')
+
+        # Empirical distribution of the conditioning label (uniform over the
+        # single dummy class when training unconditionally).
+        y_tensor = torch.from_numpy(y)
+        self._y_dist = torch.bincount(y_tensor).float()
+
+        # Build the denoiser and the joint Gaussian/multinomial diffusion
+        model = MLPDiffusion(
+            d_in=d_in,
+            num_classes=self._n_target_classes,
+            is_y_cond=target_column is not None,
+            rtdl_params={'d_layers': list(hp['d_layers']), 'dropout': hp['dropout']},
+            dim_t=hp['dim_t'],
+        ).to(self._device)
+
+        self._diffusion = GaussianMultinomialDiffusion(
+            num_classes=K,
+            num_numerical_features=n_num,
+            denoise_fn=model,
+            num_timesteps=hp['num_timesteps'],
+            scheduler=hp['scheduler'],
+            device=self._device,
+        ).to(self._device)
+        self._diffusion.train()
+
+        X = torch.from_numpy(
+            np.concatenate([X_num, X_cat.astype('float32')], axis=1)
+        ).float()
+        train_loader = _FastTensorDataLoader(X, y_tensor, batch_size=hp['batch_size'])
+
+        trainer = _Trainer(
+            self._diffusion,
+            train_loader,
+            lr=hp['lr'],
+            weight_decay=hp['weight_decay'],
+            steps=hp['steps'],
+            device=self._device,
+            verbose=hp['verbose'],
+        )
+        trainer.run_loop()
+        self.loss_history = pd.DataFrame(trainer.loss_history)
+
+        self._column_order = list(self._table_metadata['columns'])
+        self._fitted = True
+
+    def sample(self, num_rows: int) -> Dict[str, pd.DataFrame]:
+        """Sample synthetic rows from the fitted synthesizer.
+
+        Args:
+            num_rows: number of rows to generate.
+
+        Returns:
+            Dictionary mapping the table name to a synthetic DataFrame with
+            the same columns as the training data.
+        """
+        if not self._fitted:
+            raise RuntimeError('The synthesizer has not been fitted; call fit() first.')
+        if num_rows <= 0:
+            raise ValueError('num_rows must be a positive integer.')
+
+        hp = self.hyperparameters
+        self._diffusion.eval()
+        x_gen, y_gen = self._diffusion.sample_all(
+            num_rows,
+            batch_size=hp['sample_batch_size'],
+            y_dist=self._y_dist,
+            verbose=hp['verbose'],
+        )
+
+        X_gen = x_gen.numpy()
+        n_num = self._diffusion.num_numerical_features
+        X_num = X_gen[:, :n_num]
+        has_cat = self._transformer.category_sizes[0] != 0
+        X_cat = X_gen[:, n_num:] if has_cat else np.empty((num_rows, 0), dtype='int64')
+
+        df = self._transformer.inverse_transform(X_num, X_cat)
+
+        if self._target_encoder is not None:
+            decoded = self._target_encoder.inverse_transform(
+                y_gen.numpy().reshape(-1, 1)
+            ).reshape(-1)
+            series = pd.Series(decoded, dtype=object)
+            series = series.where(series != CAT_MISSING_VALUE, np.nan)
+            if self._target_is_boolean:
+                series = series.map({'True': True, 'False': False})
+            else:
+                try:
+                    series = series.astype(self._target_dtype)
+                except (ValueError, TypeError):
+                    pass
+            df[self.hyperparameters['target_column']] = series
+
+        df = df[[column for column in self._column_order if column in df.columns]]
+        return {self._table_name: df}
+
+    # -- internal helpers --------------------------------------------------
+
+    @staticmethod
+    def _parse_metadata(metadata):
+        if hasattr(metadata, 'to_dict'):
+            metadata = metadata.to_dict()
+        if not isinstance(metadata, dict):
+            raise TypeError(
+                'metadata must be an SDV metadata object or its dict representation.'
+            )
+
+        if 'tables' in metadata:
+            tables = metadata['tables']
+            if len(tables) != 1:
+                raise ValueError(
+                    'TabDDPMSynthesizer is a single-table synthesizer; the metadata '
+                    f'describes {len(tables)} tables.'
+                )
+            table_name, table_metadata = next(iter(tables.items()))
+        elif 'columns' in metadata:
+            table_name, table_metadata = None, metadata
+        else:
+            raise ValueError(
+                "metadata dict must contain either a 'columns' key (single-table "
+                "format) or a 'tables' key (multi-table format)."
+            )
+
+        if not table_metadata.get('columns'):
+            raise ValueError('The table metadata does not define any columns.')
+        return table_name, table_metadata
+
+    def _validate_data(self, data):
+        if not isinstance(data, dict) or not all(
+            isinstance(df, pd.DataFrame) for df in data.values()
+        ):
+            raise TypeError('data must be a dictionary mapping table names to DataFrames.')
+        if len(data) != 1:
+            raise ValueError(
+                'TabDDPMSynthesizer is a single-table synthesizer; got '
+                f'{len(data)} tables: {sorted(data)}.'
+            )
+
+        table_name, df = next(iter(data.items()))
+        if self._table_name is not None and table_name != self._table_name:
+            raise ValueError(
+                f"The metadata describes table '{self._table_name}' but the data "
+                f"contains table '{table_name}'."
+            )
+
+        metadata_columns = set(self._table_metadata['columns'])
+        data_columns = set(df.columns)
+        missing = metadata_columns - data_columns
+        extra = data_columns - metadata_columns
+        if missing or extra:
+            raise ValueError(
+                'The data does not match the metadata. '
+                f'Missing columns: {sorted(missing)}. Unexpected columns: {sorted(extra)}.'
+            )
+        return table_name, df

From 4be6278c69d18c0874589539a02cd3056f87ca48 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Fri, 17 Jul 2026 10:30:54 -0400
Subject: [PATCH 02/28] add tabddpm for single-table synthesis

---
 sdgym/synthesizers/__init__.py                |   2 +
 sdgym/synthesizers/tabddpm.py                 | 697 +++++++++---------
 .../integration/synthesizers/test_tabddpm.py  |  28 +
 tests/unit/synthesizers/test_tabddpm.py       | 121 +++
 4 files changed, 498 insertions(+), 350 deletions(-)
 create mode 100644 tests/integration/synthesizers/test_tabddpm.py
 create mode 100644 tests/unit/synthesizers/test_tabddpm.py

diff --git a/sdgym/synthesizers/__init__.py b/sdgym/synthesizers/__init__.py
index 67368dc7..e17d3e77 100644
--- a/sdgym/synthesizers/__init__.py
+++ b/sdgym/synthesizers/__init__.py
@@ -8,6 +8,7 @@
 from sdgym.synthesizers.identity import DataIdentity
 from sdgym.synthesizers.column import ColumnSynthesizer
 from sdgym.synthesizers.realtabformer import RealTabFormerSynthesizer
+from sdgym.synthesizers.tabddpm import TabDDPMSynthesizer
 from sdgym.synthesizers.uniform import UniformSynthesizer, MultiTableUniformSynthesizer
 from sdgym.synthesizers.utils import (
     get_available_single_table_synthesizers,
@@ -21,6 +22,7 @@
     'ColumnSynthesizer',
     'UniformSynthesizer',
     'RealTabFormerSynthesizer',
+    'TabDDPMSynthesizer',
     'create_single_table_synthesizer',
     'create_multi_table_synthesizer',
     'create_synthesizer_variant',
diff --git a/sdgym/synthesizers/tabddpm.py b/sdgym/synthesizers/tabddpm.py
index 6b353fe9..bdeb9db5 100644
--- a/sdgym/synthesizers/tabddpm.py
+++ b/sdgym/synthesizers/tabddpm.py
@@ -1,13 +1,16 @@
-"""TabDDPMSynthesizer -- SDGym synthesizer built on TabDDPM.
+"""TabDDPMSynthesizer -- SDGym synthesizer for TabDDPM.
+
 Paper: "TabDDPM: Modelling Tabular Data with Diffusion Models (2022)"
 https://arxiv.org/abs/2209.15421
 
 Original implementation is provided:
-https://github.com/yandex-research/tab-ddpm/tree/main/tab_ddpm
+https://github.com/yandex-research/tab-ddpm/tree/main/tab_ddpm.
 """
 
+import logging
 import math
-from typing import Callable, Dict, List, Optional, Union
+import sys
+from typing import Callable, List, Union
 
 import numpy as np
 import pandas as pd
@@ -17,6 +20,7 @@
 from sklearn.preprocessing import OrdinalEncoder, QuantileTransformer
 from torch import Tensor
 
+from sdgym.synthesizers.base import BaselineSynthesizer
 
 CAT_MISSING_VALUE = '__nan__'
 
@@ -39,24 +43,25 @@ def mean_flat(tensor):
 
 
 def ohe_to_categories(ohe, K):
+    """Apply one hot encoding."""
     K = torch.from_numpy(K)
     indices = torch.cat([torch.zeros((1,)), K.cumsum(dim=0)], dim=0).int().tolist()
     res = []
     for i in range(len(indices) - 1):
-        res.append(ohe[:, indices[i]:indices[i + 1]].argmax(dim=1))
+        res.append(ohe[:, indices[i] : indices[i + 1]].argmax(dim=1))
     return torch.stack(res, dim=1)
 
 
-def log_1_min_a(a):
+def _log_1_min_a(a):
     return torch.log(1 - a.exp() + 1e-40)
 
 
-def log_add_exp(a, b):
+def _log_add_exp(a, b):
     maximum = torch.max(a, b)
     return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))
 
 
-def extract(a, t, x_shape):
+def _extract(a, t, x_shape):
     b, *_ = t.shape
     t = t.to(a.device)
     out = a.gather(-1, t)
@@ -65,11 +70,11 @@ def extract(a, t, x_shape):
     return out.expand(x_shape)
 
 
-def log_categorical(log_x_start, log_prob):
+def _log_categorical(log_x_start, log_prob):
     return (log_x_start.exp() * log_prob).sum(dim=1)
 
 
-def index_to_log_onehot(x, num_classes):
+def _index_to_log_onehot(x, num_classes):
     onehots = []
     for i in range(len(num_classes)):
         onehots.append(F.one_hot(x[:, i], num_classes[i]))
@@ -79,27 +84,19 @@ def index_to_log_onehot(x, num_classes):
     return log_onehot
 
 
-@torch.jit.script
-def log_sub_exp(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+def _log_sub_exp(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
     m = torch.maximum(a, b)
     return torch.log(torch.exp(a - m) - torch.exp(b - m)) + m
 
 
-@torch.jit.script
-def sliced_logsumexp(x, slices):
-    lse = torch.logcumsumexp(
-        torch.nn.functional.pad(x, [1, 0, 0, 0], value=-float('inf')),
-        dim=-1)
+def _sliced_logsumexp(x, slices):
+    lse = torch.logcumsumexp(torch.nn.functional.pad(x, [1, 0, 0, 0], value=-float('inf')), dim=-1)
 
     slice_starts = slices[:-1]
     slice_ends = slices[1:]
 
-    slice_lse = log_sub_exp(lse[:, slice_ends], lse[:, slice_starts])
-    slice_lse_repeated = torch.repeat_interleave(
-        slice_lse,
-        slice_ends - slice_starts,
-        dim=-1
-    )
+    slice_lse = _log_sub_exp(lse[:, slice_ends], lse[:, slice_starts])
+    slice_lse_repeated = torch.repeat_interleave(slice_lse, slice_ends - slice_starts, dim=-1)
     return slice_lse_repeated
 
 
@@ -111,9 +108,7 @@ def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
         scale = 1000 / num_diffusion_timesteps
         beta_start = scale * 0.0001
         beta_end = scale * 0.02
-        return np.linspace(
-            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
-        )
+        return np.linspace(beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64)
     elif schedule_name == 'cosine':
         return betas_for_alpha_bar(
             num_diffusion_timesteps,
@@ -124,7 +119,9 @@ def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
 
 
 def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
-    """Create a beta schedule that discretizes the given alpha_t_bar function,
+    """Beta schedule.
+
+    Create a beta schedule that discretizes the given alpha_t_bar function,
     which defines the cumulative product of (1-beta) over time from t = [0,1].
     """
     betas = []
@@ -148,32 +145,36 @@ def timestep_embedding(timesteps, dim, max_period=10000):
     return embedding
 
 
-class Block(nn.Module):
-    """The main building block of `MLP`."""
-
-    def __init__(
-        self,
-        *,
-        d_in: int,
-        d_out: int,
-        bias: bool,
-        activation: Union[str, Callable[..., nn.Module]],
-        dropout: float,
-    ) -> None:
-        super().__init__()
-        self.linear = nn.Linear(d_in, d_out, bias)
-        self.activation = getattr(nn, activation)() if isinstance(activation, str) else activation()
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.dropout(self.activation(self.linear(x)))
-
 class MLP(nn.Module):
     """The MLP model from "Revisiting Deep Learning Models for Tabular Data".
 
     MLP: (in) -> Block -> ... -> Block -> Linear -> (out)
     Block: (in) -> Linear -> Activation -> Dropout -> (out)
     """
+
+    class Block(nn.Module):
+        """The main building block of `MLP`."""
+
+        def __init__(
+            self,
+            *,
+            d_in: int,
+            d_out: int,
+            bias: bool,
+            activation: Union[str, Callable[..., nn.Module]],
+            dropout: float,
+        ) -> None:
+            super().__init__()
+            self.linear = nn.Linear(d_in, d_out, bias)
+            self.activation = (
+                getattr(nn, activation)() if isinstance(activation, str) else activation()
+            )
+            self.dropout = nn.Dropout(dropout)
+
+        def forward(self, x: Tensor) -> Tensor:
+            """Forward pass."""
+            return self.dropout(self.activation(self.linear(x)))
+
     def __init__(
         self,
         *,
@@ -188,18 +189,16 @@ def __init__(
             dropouts = [dropouts] * len(d_layers)
         assert len(d_layers) == len(dropouts)
 
-        self.blocks = nn.ModuleList(
-            [
-                MLP.Block(
-                    d_in=d_layers[i - 1] if i else d_in,
-                    d_out=d,
-                    bias=True,
-                    activation=activation,
-                    dropout=dropout,
-                )
-                for i, (d, dropout) in enumerate(zip(d_layers, dropouts))
-            ]
-        )
+        self.blocks = nn.ModuleList([
+            MLP.Block(
+                d_in=d_layers[i - 1] if i else d_in,
+                d_out=d,
+                bias=True,
+                activation=activation,
+                dropout=dropout,
+            )
+            for i, (d, dropout) in enumerate(zip(d_layers, dropouts))
+        ])
         self.head = nn.Linear(d_layers[-1] if d_layers else d_in, d_out)
 
     @classmethod
@@ -226,6 +225,7 @@ def make_baseline(
         )
 
     def forward(self, x: Tensor) -> Tensor:
+        """Forward pass."""
         x = x.float()
         for block in self.blocks:
             x = block(x)
@@ -254,13 +254,10 @@ def __init__(self, d_in, num_classes, is_y_cond, rtdl_params, dim_t=128):
             self.label_emb = nn.Linear(1, dim_t)
 
         self.proj = nn.Linear(d_in, dim_t)
-        self.time_embed = nn.Sequential(
-            nn.Linear(dim_t, dim_t),
-            nn.SiLU(),
-            nn.Linear(dim_t, dim_t)
-        )
+        self.time_embed = nn.Sequential(nn.Linear(dim_t, dim_t), nn.SiLU(), nn.Linear(dim_t, dim_t))
 
     def forward(self, x, timesteps, y=None):
+        """Forward pass."""
         emb = self.time_embed(timestep_embedding(timesteps, self.dim_t))
         if self.is_y_cond and y is not None:
             if self.num_classes > 0:
@@ -272,25 +269,24 @@ def forward(self, x, timesteps, y=None):
         return self.mlp(x)
 
 
-# =====================================================================
-# Gaussian + multinomial diffusion
-# (from tab_ddpm/gaussian_multinomial_diffsuion.py, trimmed to the paths
-# actually used by the paper's pipeline: 'mse' Gaussian loss, 'eps'
-# parametrization, 'vb_stochastic' multinomial loss, uniform time sampling
-# and ancestral sampling)
-# =====================================================================
-
 class GaussianMultinomialDiffusion(torch.nn.Module):
-    """Joint diffusion: Gaussian over numerical features, multinomial over categorical."""
+    """Joint diffusion: Gaussian over numerical features, multinomial over categorical.
+
+    Simplified code to use the paper's pipeline:
+    - 'mse' Gaussian loss
+    - 'eps' parametrization
+    - 'vb_stochastic' multinomial loss
+    - uniform time sampling and ancestral sampling
+    """
 
     def __init__(
-            self,
-            num_classes: np.ndarray,
-            num_numerical_features: int,
-            denoise_fn,
-            num_timesteps=1000,
-            scheduler='cosine',
-            device=torch.device('cpu')
+        self,
+        num_classes: np.ndarray,
+        num_numerical_features: int,
+        denoise_fn,
+        num_timesteps=1000,
+        scheduler='cosine',
+        device=torch.device('cpu'),
     ):
         super(GaussianMultinomialDiffusion, self).__init__()
 
@@ -310,15 +306,15 @@ def __init__(
         self.num_timesteps = num_timesteps
         self.scheduler = scheduler
 
-        alphas = 1. - get_named_beta_schedule(scheduler, num_timesteps)
+        alphas = 1.0 - get_named_beta_schedule(scheduler, num_timesteps)
         alphas = torch.tensor(alphas.astype('float64'))
-        betas = 1. - alphas
+        betas = 1.0 - alphas
 
         log_alpha = np.log(alphas)
         log_cumprod_alpha = np.cumsum(log_alpha)
 
-        log_1_min_alpha = log_1_min_a(log_alpha)
-        log_1_min_cumprod_alpha = log_1_min_a(log_cumprod_alpha)
+        log_1_min_alpha = _log_1_min_a(log_alpha)
+        log_1_min_cumprod_alpha = _log_1_min_a(log_cumprod_alpha)
 
         alphas_cumprod = np.cumprod(alphas, axis=0)
         alphas_cumprod_prev = torch.tensor(np.append(1.0, alphas_cumprod[:-1]))
@@ -329,24 +325,25 @@ def __init__(
         sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / alphas_cumprod - 1)
 
         # Gaussian diffusion
-        self.posterior_variance = (
-            betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        self.posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+        self.posterior_log_variance_clipped = (
+            torch
+            .from_numpy(np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:])))
+            .float()
+            .to(device)
         )
-        self.posterior_log_variance_clipped = torch.from_numpy(
-            np.log(np.append(self.posterior_variance[1], self.posterior_variance[1:]))
-        ).float().to(device)
         self.posterior_mean_coef1 = (
-            betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)
-        ).float().to(device)
+            (betas * np.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)).float().to(device)
+        )
         self.posterior_mean_coef2 = (
-            (1.0 - alphas_cumprod_prev)
-            * np.sqrt(alphas.numpy())
-            / (1.0 - alphas_cumprod)
-        ).float().to(device)
+            ((1.0 - alphas_cumprod_prev) * np.sqrt(alphas.numpy()) / (1.0 - alphas_cumprod))
+            .float()
+            .to(device)
+        )
 
-        assert log_add_exp(log_alpha, log_1_min_alpha).abs().sum().item() < 1.e-5
-        assert log_add_exp(log_cumprod_alpha, log_1_min_cumprod_alpha).abs().sum().item() < 1e-5
-        assert (np.cumsum(log_alpha) - log_cumprod_alpha).abs().sum().item() < 1.e-5
+        assert _log_add_exp(log_alpha, log_1_min_alpha).abs().sum().item() < 1.0e-5
+        assert _log_add_exp(log_cumprod_alpha, log_1_min_cumprod_alpha).abs().sum().item() < 1e-5
+        assert (np.cumsum(log_alpha) - log_cumprod_alpha).abs().sum().item() < 1.0e-5
 
         # Convert to float32 and register buffers.
         self.register_buffer('alphas', alphas.float().to(device))
@@ -358,51 +355,51 @@ def __init__(
         self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float().to(device))
         self.register_buffer('alphas_cumprod_next', alphas_cumprod_next.float().to(device))
         self.register_buffer('sqrt_alphas_cumprod', sqrt_alphas_cumprod.float().to(device))
-        self.register_buffer('sqrt_one_minus_alphas_cumprod', sqrt_one_minus_alphas_cumprod.float().to(device))
-        self.register_buffer('sqrt_recip_alphas_cumprod', sqrt_recip_alphas_cumprod.float().to(device))
-        self.register_buffer('sqrt_recipm1_alphas_cumprod', sqrt_recipm1_alphas_cumprod.float().to(device))
-
+        self.register_buffer(
+            'sqrt_one_minus_alphas_cumprod', sqrt_one_minus_alphas_cumprod.float().to(device)
+        )
+        self.register_buffer(
+            'sqrt_recip_alphas_cumprod', sqrt_recip_alphas_cumprod.float().to(device)
+        )
+        self.register_buffer(
+            'sqrt_recipm1_alphas_cumprod', sqrt_recipm1_alphas_cumprod.float().to(device)
+        )
 
-    def gaussian_q_sample(self, x_start, t, noise=None):
+    def _gaussian_q_sample(self, x_start, t, noise=None):
         if noise is None:
             noise = torch.randn_like(x_start)
         assert noise.shape == x_start.shape
         return (
-            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
-            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
-            * noise
+            _extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
-    def gaussian_q_posterior_mean_variance(self, x_start, x_t, t):
+    def _gaussian_q_posterior_mean_variance(self, x_start, x_t, t):
         assert x_start.shape == x_t.shape
         posterior_mean = (
-            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
-            + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
-        )
-        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
-        posterior_log_variance_clipped = extract(
-            self.posterior_log_variance_clipped, t, x_t.shape
+            _extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
         )
+        posterior_variance = _extract(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract(self.posterior_log_variance_clipped, t, x_t.shape)
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
-    def gaussian_p_mean_variance(self, model_output, x, t):
+    def _gaussian_p_mean_variance(self, model_output, x, t):
         B = x.shape[0]
         assert t.shape == (B,)
 
         model_variance = torch.cat(
-            [self.posterior_variance[1].unsqueeze(0).to(x.device), (1. - self.alphas)[1:]], dim=0
+            [self.posterior_variance[1].unsqueeze(0).to(x.device), (1.0 - self.alphas)[1:]], dim=0
         )
         model_log_variance = torch.log(model_variance)
 
-        model_variance = extract(model_variance, t, x.shape)
-        model_log_variance = extract(model_log_variance, t, x.shape)
+        model_variance = _extract(model_variance, t, x.shape)
+        model_log_variance = _extract(model_log_variance, t, x.shape)
 
         # 'eps' parametrization: the network predicts the noise
         pred_xstart = self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
 
-        model_mean, _, _ = self.gaussian_q_posterior_mean_variance(
-            x_start=pred_xstart, x_t=x, t=t
-        )
+        model_mean, _, _ = self._gaussian_q_posterior_mean_variance(x_start=pred_xstart, x_t=x, t=t)
 
         return {
             'mean': model_mean,
@@ -418,12 +415,12 @@ def _gaussian_loss(self, model_out, noise):
     def _predict_xstart_from_eps(self, x_t, t, eps):
         assert x_t.shape == eps.shape
         return (
-            extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
-            - extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+            _extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
         )
 
-    def gaussian_p_sample(self, model_out, x, t):
-        out = self.gaussian_p_mean_variance(model_out, x, t)
+    def _gaussian_p_sample(self, model_out, x, t):
+        out = self._gaussian_p_mean_variance(model_out, x, t)
         noise = torch.randn_like(x)
         nonzero_mask = (
             (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
@@ -432,35 +429,33 @@ def gaussian_p_sample(self, model_out, x, t):
         sample = out['mean'] + nonzero_mask * torch.exp(0.5 * out['log_variance']) * noise
         return {'sample': sample, 'pred_xstart': out['pred_xstart']}
 
-
-    def multinomial_kl(self, log_prob1, log_prob2):
+    def _multinomial_kl(self, log_prob1, log_prob2):
         kl = (log_prob1.exp() * (log_prob1 - log_prob2)).sum(dim=1)
         return kl
 
-    def q_pred_one_timestep(self, log_x_t, t):
-        log_alpha_t = extract(self.log_alpha, t, log_x_t.shape)
-        log_1_min_alpha_t = extract(self.log_1_min_alpha, t, log_x_t.shape)
+    def _q_pred_one_timestep(self, log_x_t, t):
+        log_alpha_t = _extract(self.log_alpha, t, log_x_t.shape)
+        log_1_min_alpha_t = _extract(self.log_1_min_alpha, t, log_x_t.shape)
 
         # alpha_t * E[xt] + (1 - alpha_t) 1 / K
-        log_probs = log_add_exp(
-            log_x_t + log_alpha_t,
-            log_1_min_alpha_t - torch.log(self.num_classes_expanded)
+        log_probs = _log_add_exp(
+            log_x_t + log_alpha_t, log_1_min_alpha_t - torch.log(self.num_classes_expanded)
         )
 
         return log_probs
 
-    def q_pred(self, log_x_start, t):
-        log_cumprod_alpha_t = extract(self.log_cumprod_alpha, t, log_x_start.shape)
-        log_1_min_cumprod_alpha = extract(self.log_1_min_cumprod_alpha, t, log_x_start.shape)
+    def _q_pred(self, log_x_start, t):
+        log_cumprod_alpha_t = _extract(self.log_cumprod_alpha, t, log_x_start.shape)
+        log_1_min_cumprod_alpha = _extract(self.log_1_min_cumprod_alpha, t, log_x_start.shape)
 
-        log_probs = log_add_exp(
+        log_probs = _log_add_exp(
             log_x_start + log_cumprod_alpha_t,
-            log_1_min_cumprod_alpha - torch.log(self.num_classes_expanded)
+            log_1_min_cumprod_alpha - torch.log(self.num_classes_expanded),
         )
 
         return log_probs
 
-    def predict_start(self, model_out, log_x_t):
+    def _predict_start(self, model_out, log_x_t):
         assert model_out.size(0) == log_x_t.size(0)
         assert model_out.size(1) == self.num_classes.sum(), f'{model_out.size()}'
 
@@ -469,42 +464,43 @@ def predict_start(self, model_out, log_x_t):
             log_pred[:, ix] = F.log_softmax(model_out[:, ix], dim=1)
         return log_pred
 
-    def q_posterior(self, log_x_start, log_x_t, t):
+    def _q_posterior(self, log_x_start, log_x_t, t):
         # q(xt-1 | xt, x0) = q(xt | xt-1, x0) * q(xt-1 | x0) / q(xt | x0)
         # where q(xt | xt-1, x0) = q(xt | xt-1).
-        
+
         t_minus_1 = t - 1
         # Remove negative values, will not be used anyway for final decoder
         t_minus_1 = torch.where(t_minus_1 < 0, torch.zeros_like(t_minus_1), t_minus_1)
-        log_EV_qxtmin_x0 = self.q_pred(log_x_start, t_minus_1)
+        log_EV_qxtmin_x0 = self._q_pred(log_x_start, t_minus_1)
 
         num_axes = (1,) * (len(log_x_start.size()) - 1)
         t_broadcast = t.to(log_x_start.device).view(-1, *num_axes) * torch.ones_like(log_x_start)
-        log_EV_qxtmin_x0 = torch.where(t_broadcast == 0, log_x_start, log_EV_qxtmin_x0.to(torch.float32))
+        log_EV_qxtmin_x0 = torch.where(
+            t_broadcast == 0, log_x_start, log_EV_qxtmin_x0.to(torch.float32)
+        )
 
         # Note: the formula uses log q_pred_one_timestep(x_t, t), _NOT_ x_tmin1.
-        unnormed_logprobs = log_EV_qxtmin_x0 + self.q_pred_one_timestep(log_x_t, t)
+        unnormed_logprobs = log_EV_qxtmin_x0 + self._q_pred_one_timestep(log_x_t, t)
 
-        log_EV_xtmin_given_xt_given_xstart = \
-            unnormed_logprobs \
-            - sliced_logsumexp(unnormed_logprobs, self.offsets)
+        log_EV_xtmin_given_xt_given_xstart = unnormed_logprobs - _sliced_logsumexp(
+            unnormed_logprobs, self.offsets
+        )
 
         return log_EV_xtmin_given_xt_given_xstart
 
-    def p_pred(self, model_out, log_x, t):
+    def _p_pred(self, model_out, log_x, t):
         # 'x0' parametrization
-        log_x_recon = self.predict_start(model_out, log_x)
-        log_model_pred = self.q_posterior(
-            log_x_start=log_x_recon, log_x_t=log_x, t=t)
+        log_x_recon = self._predict_start(model_out, log_x)
+        log_model_pred = self._q_posterior(log_x_start=log_x_recon, log_x_t=log_x, t=t)
         return log_model_pred
 
     @torch.no_grad()
-    def p_sample(self, model_out, log_x, t):
-        model_log_prob = self.p_pred(model_out, log_x=log_x, t=t)
-        out = self.log_sample_categorical(model_log_prob)
+    def _p_sample(self, model_out, log_x, t):
+        model_log_prob = self._p_pred(model_out, log_x=log_x, t=t)
+        out = self._log_sample_categorical(model_log_prob)
         return out
 
-    def log_sample_categorical(self, logits):
+    def _log_sample_categorical(self, logits):
         full_sample = []
         for i in range(len(self.num_classes)):
             one_class_logits = logits[:, self.slices_for_classes[i]]
@@ -513,45 +509,44 @@ def log_sample_categorical(self, logits):
             sample = (gumbel_noise + one_class_logits).argmax(dim=1)
             full_sample.append(sample.unsqueeze(1))
         full_sample = torch.cat(full_sample, dim=1)
-        log_sample = index_to_log_onehot(full_sample, self.num_classes)
+        log_sample = _index_to_log_onehot(full_sample, self.num_classes)
         return log_sample
 
-    def q_sample(self, log_x_start, t):
-        log_EV_qxt_x0 = self.q_pred(log_x_start, t)
-        log_sample = self.log_sample_categorical(log_EV_qxt_x0)
+    def _q_sample(self, log_x_start, t):
+        log_EV_qxt_x0 = self._q_pred(log_x_start, t)
+        log_sample = self._log_sample_categorical(log_EV_qxt_x0)
         return log_sample
 
-    def kl_prior(self, log_x_start):
+    def _kl_prior(self, log_x_start):
         b = log_x_start.size(0)
         device = log_x_start.device
         ones = torch.ones(b, device=device).long()
 
-        log_qxT_prob = self.q_pred(log_x_start, t=(self.num_timesteps - 1) * ones)
+        log_qxT_prob = self._q_pred(log_x_start, t=(self.num_timesteps - 1) * ones)
         log_half_prob = -torch.log(self.num_classes_expanded * torch.ones_like(log_qxT_prob))
 
-        kl_prior = self.multinomial_kl(log_qxT_prob, log_half_prob)
+        kl_prior = self._multinomial_kl(log_qxT_prob, log_half_prob)
         return sum_except_batch(kl_prior)
 
-    def compute_Lt(self, model_out, log_x_start, log_x_t, t, detach_mean=False):
-        log_true_prob = self.q_posterior(
-            log_x_start=log_x_start, log_x_t=log_x_t, t=t)
-        log_model_prob = self.p_pred(model_out, log_x=log_x_t, t=t)
+    def _compute_Lt(self, model_out, log_x_start, log_x_t, t, detach_mean=False):
+        log_true_prob = self._q_posterior(log_x_start=log_x_start, log_x_t=log_x_t, t=t)
+        log_model_prob = self._p_pred(model_out, log_x=log_x_t, t=t)
 
         if detach_mean:
             log_model_prob = log_model_prob.detach()
 
-        kl = self.multinomial_kl(log_true_prob, log_model_prob)
+        kl = self._multinomial_kl(log_true_prob, log_model_prob)
         kl = sum_except_batch(kl)
 
-        decoder_nll = -log_categorical(log_x_start, log_model_prob)
+        decoder_nll = -_log_categorical(log_x_start, log_model_prob)
         decoder_nll = sum_except_batch(decoder_nll)
 
         mask = (t == torch.zeros_like(t)).float()
-        loss = mask * decoder_nll + (1. - mask) * kl
+        loss = mask * decoder_nll + (1.0 - mask) * kl
 
         return loss
 
-    def sample_time(self, b, device):
+    def _sample_time(self, b, device):
         # uniform time sampling
         t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
         pt = torch.ones_like(t).float() / self.num_timesteps
@@ -559,44 +554,43 @@ def sample_time(self, b, device):
 
     def _multinomial_loss(self, model_out, log_x_start, log_x_t, t, pt):
         # 'vb_stochastic' loss
-        kl = self.compute_Lt(model_out, log_x_start, log_x_t, t)
-        kl_prior = self.kl_prior(log_x_start)
+        kl = self._compute_Lt(model_out, log_x_start, log_x_t, t)
+        kl_prior = self._kl_prior(log_x_start)
         # Upweigh loss term of the kl
         vb_loss = kl / pt + kl_prior
         return vb_loss
 
-
-    def mixed_loss(self, x, out_dict):
+    def _mixed_loss(self, x, out_dict):
         b = x.shape[0]
         device = x.device
-        t, pt = self.sample_time(b, device)
+        t, pt = self._sample_time(b, device)
 
-        x_num = x[:, :self.num_numerical_features]
-        x_cat = x[:, self.num_numerical_features:]
+        x_num = x[:, : self.num_numerical_features]
+        x_cat = x[:, self.num_numerical_features :]
 
         x_num_t = x_num
         log_x_cat_t = x_cat
         noise = None
         if x_num.shape[1] > 0:
             noise = torch.randn_like(x_num)
-            x_num_t = self.gaussian_q_sample(x_num, t, noise=noise)
+            x_num_t = self._gaussian_q_sample(x_num, t, noise=noise)
         if x_cat.shape[1] > 0:
-            log_x_cat = index_to_log_onehot(x_cat.long(), self.num_classes)
-            log_x_cat_t = self.q_sample(log_x_start=log_x_cat, t=t)
+            log_x_cat = _index_to_log_onehot(x_cat.long(), self.num_classes)
+            log_x_cat_t = self._q_sample(log_x_start=log_x_cat, t=t)
 
         x_in = torch.cat([x_num_t, log_x_cat_t], dim=1)
 
         model_out = self._denoise_fn(x_in, t, **out_dict)
 
-        model_out_num = model_out[:, :self.num_numerical_features]
-        model_out_cat = model_out[:, self.num_numerical_features:]
+        model_out_num = model_out[:, : self.num_numerical_features]
+        model_out_cat = model_out[:, self.num_numerical_features :]
 
         loss_multi = torch.zeros((1,), device=device).float()
         loss_gauss = torch.zeros((1,), device=device).float()
         if x_cat.shape[1] > 0:
-            loss_multi = self._multinomial_loss(
-                model_out_cat, log_x_cat, log_x_cat_t, t, pt
-            ) / len(self.num_classes)
+            loss_multi = self._multinomial_loss(model_out_cat, log_x_cat, log_x_cat_t, t, pt) / len(
+                self.num_classes
+            )
 
         if x_num.shape[1] > 0:
             loss_gauss = self._gaussian_loss(model_out_num, noise)
@@ -604,7 +598,7 @@ def mixed_loss(self, x, out_dict):
         return loss_multi.mean(), loss_gauss.mean()
 
     @torch.no_grad()
-    def sample(self, num_samples, y_dist):
+    def _sample(self, num_samples, y_dist):
         b = num_samples
         device = self.log_alpha.device
         z_norm = torch.randn((b, self.num_numerical_features), device=device)
@@ -613,27 +607,19 @@ def sample(self, num_samples, y_dist):
         log_z = torch.zeros((b, 0), device=device).float()
         if has_cat:
             uniform_logits = torch.zeros((b, len(self.num_classes_expanded)), device=device)
-            log_z = self.log_sample_categorical(uniform_logits)
+            log_z = self._log_sample_categorical(uniform_logits)
 
-        y = torch.multinomial(
-            y_dist,
-            num_samples=b,
-            replacement=True
-        )
+        y = torch.multinomial(y_dist, num_samples=b, replacement=True)
         out_dict = {'y': y.long().to(device)}
         for i in reversed(range(0, self.num_timesteps)):
             t = torch.full((b,), i, device=device, dtype=torch.long)
-            model_out = self._denoise_fn(
-                torch.cat([z_norm, log_z], dim=1).float(),
-                t,
-                **out_dict
-            )
-            model_out_num = model_out[:, :self.num_numerical_features]
-            model_out_cat = model_out[:, self.num_numerical_features:]
+            model_out = self._denoise_fn(torch.cat([z_norm, log_z], dim=1).float(), t, **out_dict)
+            model_out_num = model_out[:, : self.num_numerical_features]
+            model_out_cat = model_out[:, self.num_numerical_features :]
             if self.num_numerical_features > 0:
-                z_norm = self.gaussian_p_sample(model_out_num, z_norm, t)['sample']
+                z_norm = self._gaussian_p_sample(model_out_num, z_norm, t)['sample']
             if has_cat:
-                log_z = self.p_sample(model_out_cat, log_z, t)
+                log_z = self._p_sample(model_out_cat, log_z, t)
 
         z_ohe = torch.exp(log_z).round()
         z_cat = log_z
@@ -642,7 +628,7 @@ def sample(self, num_samples, y_dist):
         sample = torch.cat([z_norm, z_cat], dim=1).cpu()
         return sample, out_dict
 
-    def sample_all(self, num_samples, batch_size, y_dist, verbose=False):
+    def _sample_all(self, num_samples, batch_size, y_dist, verbose=False):
         all_y = []
         all_samples = []
         num_generated = 0
@@ -657,7 +643,7 @@ def sample_all(self, num_samples, batch_size, y_dist, verbose=False):
             attempts += 1
 
             b = min(batch_size, num_samples - num_generated)
-            sample, out_dict = self.sample(b, y_dist)
+            sample, out_dict = self._sample(b, y_dist)
             mask_nan = torch.any(sample.isnan(), dim=1)
             sample = sample[~mask_nan]
             y = out_dict['y'][~mask_nan]
@@ -666,23 +652,14 @@ def sample_all(self, num_samples, batch_size, y_dist, verbose=False):
             all_y.append(y.cpu())
             num_generated += sample.shape[0]
             if verbose:
-                print(f'Sampled {min(num_generated, num_samples)}/{num_samples} rows', end='\r')
+                sys.stdout.write(f'Sampled {min(num_generated, num_samples)}/{num_samples} rows\n')
 
-        if verbose:
-            print()
         x_gen = torch.cat(all_samples, dim=0)[:num_samples]
         y_gen = torch.cat(all_y, dim=0)[:num_samples]
 
         return x_gen, y_gen
 
 
-# =====================================================================
-# Data transformer: SDV metadata-driven DataFrame <-> model matrix
-# (replaces lib/data.py's Dataset/Transformations for the DataFrame API;
-# uses the paper's preprocessing: quantile normalization for numerical
-# features and ordinal encoding for categorical features)
-# =====================================================================
-
 class _DataTransformer:
     """Converts a DataFrame into (numerical block, categorical index block) and back."""
 
@@ -697,8 +674,12 @@ def __init__(self, columns_metadata, normalization='quantile', seed=0):
         for column, spec in columns_metadata.items():
             sdtype = spec.get('sdtype', 'categorical')
             role = sdtype
-            if role not in ('numerical', 'datetime', 'boolean', 'id'):
-                # 'categorical' and any unrecognized/PII sdtype
+            if spec.get('pii', False):
+                # PII columns are not modelled; placeholders are regenerated
+                # at sampling time (like SDV's anonymization behavior).
+                role = 'id'
+            elif role not in ('numerical', 'datetime', 'boolean', 'id'):
+                # 'categorical' and any unrecognized sdtype
                 role = 'categorical'
             self._roles[column] = role
 
@@ -706,8 +687,6 @@ def __init__(self, columns_metadata, normalization='quantile', seed=0):
         self.cat_columns: List[str] = []
         self.id_columns: List[str] = []
 
-    # -- fitting -------------------------------------------------------
-
     def fit(self, df: pd.DataFrame) -> None:
         for column in df.columns:
             role = self._roles[column]
@@ -737,7 +716,7 @@ def fit(self, df: pd.DataFrame) -> None:
             X_num[inds] = np.take(col_means, inds[1])
 
             # Columns holding few unique integer values are snapped back to the
-            # observed values after sampling (same heuristic as scripts/sample.py).
+            # observed values after sampling.
             for j, column in enumerate(self.num_columns):
                 if self._roles[column] != 'numerical':
                     continue
@@ -761,9 +740,9 @@ def fit(self, df: pd.DataFrame) -> None:
             X_cat = self._to_categorical_block(df)
             self._cat_transform = OrdinalEncoder(dtype='int64')
             self._cat_transform.fit(X_cat)
-            self.category_sizes = np.array(
-                [len(categories) for categories in self._cat_transform.categories_]
-            )
+            self.category_sizes = np.array([
+                len(categories) for categories in self._cat_transform.categories_
+            ])
 
     def transform(self, df: pd.DataFrame):
         X_num = self._to_numeric_block(df)
@@ -779,8 +758,6 @@ def transform(self, df: pd.DataFrame):
 
         return X_num.astype('float32'), X_cat
 
-    # -- inverting -----------------------------------------------------
-
     def inverse_transform(self, X_num: np.ndarray, X_cat: np.ndarray) -> pd.DataFrame:
         n_rows = max(X_num.shape[0], X_cat.shape[0])
         columns = {}
@@ -795,9 +772,7 @@ def inverse_transform(self, X_num: np.ndarray, X_cat: np.ndarray) -> pd.DataFram
                 columns[column] = self._from_numeric_column(column, X_num[:, j])
 
         if self.cat_columns:
-            decoded = self._cat_transform.inverse_transform(
-                np.round(X_cat).astype('int64')
-            )
+            decoded = self._cat_transform.inverse_transform(np.round(X_cat).astype('int64'))
             for j, column in enumerate(self.cat_columns):
                 columns[column] = self._from_categorical_column(column, decoded[:, j])
 
@@ -809,18 +784,16 @@ def inverse_transform(self, X_num: np.ndarray, X_cat: np.ndarray) -> pd.DataFram
 
         return pd.DataFrame(columns)
 
-    # -- per-column helpers ---------------------------------------------
-
     def _to_numeric_block(self, df: pd.DataFrame) -> np.ndarray:
         parts = []
         for column in self.num_columns:
             if self._roles[column] == 'datetime':
                 fmt = self._columns_metadata[column].get('datetime_format')
                 series = pd.to_datetime(df[column], format=fmt, errors='coerce')
-                values = series.values.astype('int64').astype('float64')
-                values[series.isna().values] = np.nan
+                values = series.to_numpy().astype('int64').astype('float64')
+                values[series.isna().to_numpy()] = np.nan
             else:
-                values = pd.to_numeric(df[column], errors='coerce').astype('float64').values
+                values = pd.to_numeric(df[column], errors='coerce').astype('float64').to_numpy()
             parts.append(values)
         if not parts:
             return np.empty((len(df), 0), dtype='float64')
@@ -831,12 +804,16 @@ def _to_categorical_block(self, df: pd.DataFrame) -> np.ndarray:
         for column in self.cat_columns:
             series = df[column]
             values = series.astype(object).where(series.notna(), CAT_MISSING_VALUE).astype(str)
-            parts.append(values.values)
+            parts.append(values.to_numpy())
         return np.column_stack(parts)
 
     def _from_numeric_column(self, column: str, values: np.ndarray):
         if self._roles[column] == 'datetime':
-            return pd.to_datetime(np.round(values).astype('int64'))
+            stamps = pd.to_datetime(np.round(values).astype('int64'))
+            fmt = self._columns_metadata[column].get('datetime_format')
+            if fmt is not None:
+                return stamps.strftime(fmt)
+            return stamps
         dtype = self._dtypes[column]
         if pd.api.types.is_integer_dtype(dtype):
             return np.round(values).astype(dtype)
@@ -853,14 +830,7 @@ def _from_categorical_column(self, column: str, values: np.ndarray):
             return series
 
 
-# =====================================================================
-# Trainer (from scripts/train.py, without EMA -- the original pipeline
-# samples from the non-EMA weights)
-# =====================================================================
-
 class _FastTensorDataLoader:
-    """Infinite iterator over shuffled (X, y) batches; faster than DataLoader for tensors."""
-
     def __init__(self, X, y, batch_size):
         self.X = X
         self.y = y
@@ -871,7 +841,7 @@ def __iter__(self):
             perm = torch.randperm(self.X.shape[0])
             X, y = self.X[perm], self.y[perm]
             for i in range(0, X.shape[0], self.batch_size):
-                yield X[i:i + self.batch_size], y[i:i + self.batch_size]
+                yield X[i : i + self.batch_size], y[i : i + self.batch_size]
 
 
 class _Trainer:
@@ -899,7 +869,7 @@ def _run_step(self, x, out_dict):
         for k in out_dict:
             out_dict[k] = out_dict[k].long().to(self.device)
         self.optimizer.zero_grad()
-        loss_multi, loss_gauss = self.diffusion.mixed_loss(x, out_dict)
+        loss_multi, loss_gauss = self.diffusion._mixed_loss(x, out_dict)
         loss = loss_multi + loss_gauss
         loss.backward()
         self.optimizer.step()
@@ -923,109 +893,121 @@ def run_loop(self):
             if (step + 1) % self.log_every == 0:
                 mloss = np.around(curr_loss_multi / curr_count, 4)
                 gloss = np.around(curr_loss_gauss / curr_count, 4)
-                self.loss_history.append(
-                    {'step': step + 1, 'mloss': mloss, 'gloss': gloss, 'loss': mloss + gloss}
-                )
+                self.loss_history.append({
+                    'step': step + 1,
+                    'mloss': mloss,
+                    'gloss': gloss,
+                    'loss': mloss + gloss,
+                })
                 if self.verbose:
-                    print(f'Step {step + 1}/{self.steps} MLoss: {mloss} GLoss: {gloss} '
-                          f'Sum: {np.around(mloss + gloss, 4)}')
+                    sys.stdout.write(
+                        f'Step {step + 1}/{self.steps} MLoss: {mloss} GLoss: {gloss} '
+                        f'Sum: {np.around(mloss + gloss, 4)}\n'
+                    )
                 curr_count = 0
                 curr_loss_multi = 0.0
                 curr_loss_gauss = 0.0
 
 
-# =====================================================================
-# Public synthesizer
-# =====================================================================
-
-class TabDDPMSynthesizer:
-    """SDV-style single-table synthesizer based on TabDDPM (arXiv:2209.15421).
+class TabDDPM:
+    """Single table synthesizer based on TabDDPM.
 
     Args:
-        metadata:
-            SDV metadata describing the table -- either an SDV metadata object
-            (``SingleTableMetadata`` / ``Metadata``, anything exposing
-            ``to_dict()``) or the equivalent dictionary with a ``'columns'``
-            key (single-table format) or a ``'tables'`` key (multi-table
-            format holding exactly one table).
-        hyperparameters:
-            Optional dict overriding any of ``DEFAULT_HYPERPARAMETERS``:
-
-            - ``target_column`` (str or None): categorical/boolean column to
-              condition the diffusion on (the paper's class-conditional setup
-              for classification datasets). ``None`` trains unconditionally.
-            - ``d_layers`` (list of int): hidden layer sizes of the MLP denoiser.
-            - ``dropout`` (float): dropout of the MLP denoiser.
-            - ``dim_t`` (int): timestep/label embedding dimension.
-            - ``num_timesteps`` (int): diffusion timesteps T.
-            - ``scheduler`` (str): ``'cosine'`` or ``'linear'`` beta schedule.
-            - ``steps`` (int): training iterations.
-            - ``lr``, ``weight_decay``, ``batch_size``: AdamW optimizer settings.
-            - ``normalization`` (str or None): ``'quantile'`` (paper default)
-              or ``None`` to skip normalizing numerical features.
-            - ``sample_batch_size`` (int): batch size used during sampling.
-            - ``device`` (str or None): e.g. ``'cuda'``/``'cpu'``;
-              ``None`` auto-selects.
-            - ``seed`` (int): random seed used for fitting.
-            - ``verbose`` (bool): print training/sampling progress.
+        metadata (sdv.metadata.Metadata):
+            The metadata describing the data.
+        target_column (str or None):
+            Categorical/boolean column to condition the diffusion on (the paper's
+            class-conditional setup for classification datasets). ``None`` trains
+            unconditionally.
+        d_layers (List[int]):
+            Hidden layer sizes of the MLP denoiser.
+        dropout (float):
+            Dropout of the MLP denoiser.
+        dim_t (int):
+            Timestep/label embedding dimension.
+        num_timesteps (int):
+            Diffusion timesteps T.
+        scheduler (str):
+            ``'cosine'`` or ``'linear'`` beta schedule.
+        steps (int):
+            Training iterations.
+        lr (float):
+            Learning rate for the optimizer.
+        weight_decay (float):
+            Weight decay for the optimizer.
+        batch_size (int):
+            Batch size.
+        normalization (str or None):
+            ``'quantile'`` (paper default) or ``None`` to skip normalizing numerical features.
+        sample_batch_size (int):
+            Batch size used during sampling.
+        device (str or None):
+            Whether to use ``'cuda'`` or ``'cpu'``. If None, auto-select is used.
+        seed (int):
+            Random seed used for fitting.
+        verbose (bool):
+            Print training / sampling progress.
     """
 
-    DEFAULT_HYPERPARAMETERS = {
-        'target_column': None,
-        'd_layers': [256, 256],
-        'dropout': 0.0,
-        'dim_t': 128,
-        'num_timesteps': 1000,
-        'scheduler': 'cosine',
-        'steps': 1000,
-        'lr': 0.001,
-        'weight_decay': 1e-5,
-        'batch_size': 4096,
-        'normalization': 'quantile',
-        'sample_batch_size': 10000,
-        'device': None,
-        'seed': 0,
-        'verbose': True,
-    }
-
-    def __init__(self, metadata, hyperparameters: Optional[dict] = None):
-        self._table_name, self._table_metadata = self._parse_metadata(metadata)
-
-        hyperparameters = dict(hyperparameters or {})
-        unknown = set(hyperparameters) - set(self.DEFAULT_HYPERPARAMETERS)
-        if unknown:
-            raise ValueError(
-                f'Unknown hyperparameters: {sorted(unknown)}. '
-                f'Valid keys are: {sorted(self.DEFAULT_HYPERPARAMETERS)}.'
-            )
-        self.hyperparameters = {**self.DEFAULT_HYPERPARAMETERS, **hyperparameters}
+    def __init__(
+        self,
+        metadata,
+        target_column=None,
+        d_layers=None,
+        dropout=0.0,
+        dim_t=128,
+        num_timesteps=1000,
+        scheduler='cosine',
+        steps=1000,
+        lr=0.001,
+        weight_decay=1e-5,
+        batch_size=4096,
+        normalization='quantile',
+        sample_batch_size=10000,
+        device=None,
+        seed=0,
+        verbose=False,
+    ):
+        self.target_column = target_column
+        self.d_layers = d_layers or [256, 256]
+        self.dropout = dropout
+        self.dim_t = dim_t
+        self.num_timesteps = num_timesteps
+        self.scheduler = scheduler
+        self.steps = steps
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.batch_size = batch_size
+        self.normalization = normalization
+        self.sample_batch_size = sample_batch_size
+        self.seed = seed
+        self.verbose = verbose
 
-        device = self.hyperparameters['device']
+        self._table_name, self._table_metadata = self._parse_metadata(metadata)
         if device is None:
             device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self._device = torch.device(device)
-
         self._fitted = False
 
-    # -- public API ------------------------------------------------------
-
-    def fit(self, data: Dict[str, pd.DataFrame]) -> None:
+    def fit(self, data):
         """Fit the synthesizer on real data.
 
         Args:
-            data: dictionary mapping table names to pandas DataFrames.
-                Must contain exactly one table matching the metadata.
+            data (pandas.DataFrame):
+                The data to fit the synthesizer to.
         """
+        if isinstance(data, pd.DataFrame):
+            data = {self._table_name: data}
+
         table_name, df = self._validate_data(data)
         self._table_name = table_name
 
-        hp = self.hyperparameters
-        torch.manual_seed(hp['seed'])
-        np.random.seed(hp['seed'])
+        torch.manual_seed(self.seed)
+        np.random.seed(self.seed)
 
         # Target column for class-conditional training (paper's setup for
         # classification datasets); everything else goes through the transformer.
-        target_column = hp['target_column']
+        target_column = self.target_column
         columns_metadata = dict(self._table_metadata['columns'])
         self._n_target_classes = 0
         self._target_encoder = None
@@ -1041,10 +1023,12 @@ def fit(self, data: Dict[str, pd.DataFrame]) -> None:
                 )
             target_series = df[target_column]
             target_values = (
-                target_series.astype(object)
+                target_series
+                .astype(object)
                 .where(target_series.notna(), CAT_MISSING_VALUE)
                 .astype(str)
-                .values.reshape(-1, 1)
+                .to_numpy()
+                .reshape(-1, 1)
             )
             self._target_encoder = OrdinalEncoder(dtype='int64')
             y = self._target_encoder.fit_transform(target_values).reshape(-1)
@@ -1058,7 +1042,7 @@ def fit(self, data: Dict[str, pd.DataFrame]) -> None:
 
         # Preprocess: quantile-normalized numerical block + ordinal categorical block
         self._transformer = _DataTransformer(
-            columns_metadata, normalization=hp['normalization'], seed=hp['seed']
+            columns_metadata, normalization=self.normalization, seed=self.seed
         )
         self._transformer.fit(df)
         X_num, X_cat = self._transformer.transform(df)
@@ -1079,33 +1063,31 @@ def fit(self, data: Dict[str, pd.DataFrame]) -> None:
             d_in=d_in,
             num_classes=self._n_target_classes,
             is_y_cond=target_column is not None,
-            rtdl_params={'d_layers': list(hp['d_layers']), 'dropout': hp['dropout']},
-            dim_t=hp['dim_t'],
+            rtdl_params={'d_layers': list(self.d_layers), 'dropout': self.dropout},
+            dim_t=self.dim_t,
         ).to(self._device)
 
         self._diffusion = GaussianMultinomialDiffusion(
             num_classes=K,
             num_numerical_features=n_num,
             denoise_fn=model,
-            num_timesteps=hp['num_timesteps'],
-            scheduler=hp['scheduler'],
+            num_timesteps=self.num_timesteps,
+            scheduler=self.scheduler,
             device=self._device,
         ).to(self._device)
         self._diffusion.train()
 
-        X = torch.from_numpy(
-            np.concatenate([X_num, X_cat.astype('float32')], axis=1)
-        ).float()
-        train_loader = _FastTensorDataLoader(X, y_tensor, batch_size=hp['batch_size'])
+        X = torch.from_numpy(np.concatenate([X_num, X_cat.astype('float32')], axis=1)).float()
+        train_loader = _FastTensorDataLoader(X, y_tensor, batch_size=self.batch_size)
 
         trainer = _Trainer(
             self._diffusion,
             train_loader,
-            lr=hp['lr'],
-            weight_decay=hp['weight_decay'],
-            steps=hp['steps'],
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+            steps=self.steps,
             device=self._device,
-            verbose=hp['verbose'],
+            verbose=self.verbose,
         )
         trainer.run_loop()
         self.loss_history = pd.DataFrame(trainer.loss_history)
@@ -1113,28 +1095,28 @@ def fit(self, data: Dict[str, pd.DataFrame]) -> None:
         self._column_order = list(self._table_metadata['columns'])
         self._fitted = True
 
-    def sample(self, num_rows: int) -> Dict[str, pd.DataFrame]:
+    def sample(self, num_rows):
         """Sample synthetic rows from the fitted synthesizer.
 
         Args:
-            num_rows: number of rows to generate.
+            num_rows (int):
+                Amount of rows to sample.
 
         Returns:
-            Dictionary mapping the table name to a synthetic DataFrame with
-            the same columns as the training data.
+            pandas.DataFrame:
+                Sampled data.
         """
         if not self._fitted:
             raise RuntimeError('The synthesizer has not been fitted; call fit() first.')
         if num_rows <= 0:
             raise ValueError('num_rows must be a positive integer.')
 
-        hp = self.hyperparameters
         self._diffusion.eval()
-        x_gen, y_gen = self._diffusion.sample_all(
+        x_gen, y_gen = self._diffusion._sample_all(
             num_rows,
-            batch_size=hp['sample_batch_size'],
+            batch_size=self.sample_batch_size,
             y_dist=self._y_dist,
-            verbose=hp['verbose'],
+            verbose=self.verbose,
         )
 
         X_gen = x_gen.numpy()
@@ -1146,9 +1128,9 @@ def sample(self, num_rows: int) -> Dict[str, pd.DataFrame]:
         df = self._transformer.inverse_transform(X_num, X_cat)
 
         if self._target_encoder is not None:
-            decoded = self._target_encoder.inverse_transform(
-                y_gen.numpy().reshape(-1, 1)
-            ).reshape(-1)
+            decoded = self._target_encoder.inverse_transform(y_gen.numpy().reshape(-1, 1)).reshape(
+                -1
+            )
             series = pd.Series(decoded, dtype=object)
             series = series.where(series != CAT_MISSING_VALUE, np.nan)
             if self._target_is_boolean:
@@ -1158,21 +1140,17 @@ def sample(self, num_rows: int) -> Dict[str, pd.DataFrame]:
                     series = series.astype(self._target_dtype)
                 except (ValueError, TypeError):
                     pass
-            df[self.hyperparameters['target_column']] = series
+            df[self.target_column] = series
 
         df = df[[column for column in self._column_order if column in df.columns]]
-        return {self._table_name: df}
-
-    # -- internal helpers --------------------------------------------------
+        return df
 
     @staticmethod
     def _parse_metadata(metadata):
         if hasattr(metadata, 'to_dict'):
             metadata = metadata.to_dict()
         if not isinstance(metadata, dict):
-            raise TypeError(
-                'metadata must be an SDV metadata object or its dict representation.'
-            )
+            raise TypeError('metadata must be an SDV metadata object or its dict representation.')
 
         if 'tables' in metadata:
             tables = metadata['tables']
@@ -1222,3 +1200,22 @@ def _validate_data(self, data):
                 f'Missing columns: {sorted(missing)}. Unexpected columns: {sorted(extra)}.'
             )
         return table_name, df
+
+
+class TabDDPMSynthesizer(BaselineSynthesizer):
+    """Custom wrapper for the TabDDPM synthesizer to make it work with SDGym."""
+
+    LOGGER = logging.getLogger(__name__)
+    _MODEL_KWARGS = None
+    _MODALITY_FLAG = 'single_table'
+
+    def _fit(self, data, metadata):
+        model_kwargs = self._MODEL_KWARGS.copy() if self._MODEL_KWARGS else {}
+        model = TabDDPM(metadata, **model_kwargs)
+        model.fit(data)
+
+        self._internal_synthesizer = model
+
+    def _sample_from_synthesizer(self, synthesizer, n_sample):
+        """Sample synthetic data with specified sample count."""
+        return synthesizer._internal_synthesizer.sample(n_sample)
diff --git a/tests/integration/synthesizers/test_tabddpm.py b/tests/integration/synthesizers/test_tabddpm.py
new file mode 100644
index 00000000..5c977404
--- /dev/null
+++ b/tests/integration/synthesizers/test_tabddpm.py
@@ -0,0 +1,28 @@
+import sys
+
+import pytest
+
+from sdgym import load_dataset
+from sdgym.synthesizers import TabDDPMSynthesizer
+
+
+@pytest.mark.skipif(sys.platform.startswith('darwin'), reason='Test not supported on github MacOS')
+def test_tabddpm_end_to_end():
+    """Test it without metrics."""
+    # Setup
+    data, metadata_dict = load_dataset(
+        'single_table', 'student_placements', limit_dataset_size=False
+    )
+    tabddpm_instance = TabDDPMSynthesizer()
+    tabddpm_instance._MODEL_KWARGS = {'steps': 5000, 'num_timesteps': 1000}
+
+    # Run
+    trained_synthesizer = tabddpm_instance.get_trained_synthesizer(data, metadata_dict)
+    sampled_data = tabddpm_instance.sample_from_synthesizer(trained_synthesizer, n_samples=10)
+
+    # Assert
+    assert sampled_data.shape[1] == data.shape[1], (
+        f'Sampled data shape {sampled_data.shape} does not match original data shape {data.shape}'
+    )
+
+    assert set(sampled_data.columns) == set(data.columns)
diff --git a/tests/unit/synthesizers/test_tabddpm.py b/tests/unit/synthesizers/test_tabddpm.py
new file mode 100644
index 00000000..05fc3b43
--- /dev/null
+++ b/tests/unit/synthesizers/test_tabddpm.py
@@ -0,0 +1,121 @@
+"""Tests for the tabddpm module."""
+
+from unittest.mock import MagicMock
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from sdgym.synthesizers.tabddpm import TabDDPM, TabDDPMSynthesizer
+
+
+@pytest.fixture
+def sample_data():
+    """Provide sample data for testing."""
+    n_samples = 10
+    num_values = np.random.normal(size=n_samples)
+
+    return pd.DataFrame({
+        'num': num_values,
+    })
+
+
+class TestTabDDPMSynthesizer:
+    """Unit tests for TabDDPMSynthesizer integration with SDGym."""
+
+    def test__get_trained_synthesizer(self):
+        """Test _get_trained_synthesizer
+
+        Initializes TabDDPM and fits TabDDPM with correct parameters.
+        """
+        # Setup
+        synthesizer = TabDDPMSynthesizer()
+        metadata = {
+            'primary_key': 'guest_email',
+            'METADATA_SPEC_VERSION': 'SINGLE_TABLE_V1',
+            'columns': {
+                'guest_email': {'sdtype': 'email', 'pii': True},
+                'has_rewards': {'sdtype': 'boolean'},
+                'room_type': {'sdtype': 'categorical'},
+                'amenities_fee': {'sdtype': 'numerical', 'computer_representation': 'Float'},
+                'checkin_date': {'sdtype': 'datetime', 'datetime_format': '%d %b %Y'},
+                'checkout_date': {'sdtype': 'datetime', 'datetime_format': '%d %b %Y'},
+                'room_rate': {'sdtype': 'numerical', 'computer_representation': 'Float'},
+                'billing_address': {'sdtype': 'address', 'pii': True},
+                'credit_card_number': {'sdtype': 'credit_card_number', 'pii': True},
+            },
+        }
+
+        data = {
+            'guest_email': {
+                0: 'michaelsanders@shaw.net',
+                1: 'randy49@brown.biz',
+                2: 'webermelissa@neal.com',
+                3: 'gsims@terry.com',
+                4: 'misty33@smith.biz',
+            },
+            'has_rewards': {0: False, 1: False, 2: True, 3: False, 4: False},
+            'room_type': {0: 'BASIC', 1: 'BASIC', 2: 'DELUXE', 3: 'BASIC', 4: 'BASIC'},
+            'amenities_fee': {0: 37.89, 1: 24.37, 2: 0.0, 3: np.nan, 4: 16.45},
+            'checkin_date': {
+                0: '27 Dec 2020',
+                1: '30 Dec 2020',
+                2: '17 Sep 2020',
+                3: '28 Dec 2020',
+                4: '05 Apr 2020',
+            },
+            'checkout_date': {
+                0: '29 Dec 2020',
+                1: '02 Jan 2021',
+                2: '18 Sep 2020',
+                3: '31 Dec 2020',
+                4: np.nan,
+            },
+            'room_rate': {0: 131.23, 1: 114.43, 2: 368.33, 3: 115.61, 4: 122.41},
+            'billing_address': {
+                0: '49380 Rivers Street\nSpencerville, AK 68265',
+                1: '88394 Boyle Meadows\nConleyberg, TN 22063',
+                2: '0323 Lisa Station Apt. 208\nPort Thomas, LA 82585',
+                3: '77 Massachusetts Ave\nCambridge, MA 02139',
+                4: '1234 Corporate Drive\nBoston, MA 02116',
+            },
+            'credit_card_number': {
+                0: 4075084747483975747,
+                1: 180072822063468,
+                2: 38983476971380,
+                3: 4969551998845740,
+                4: 3558512986488983,
+            },
+        }
+
+        real_data = pd.DataFrame(data)
+        n_sample = 5
+
+        # Run
+        result = synthesizer._get_trained_synthesizer(real_data, metadata)
+        sampled_data = synthesizer._sample_from_synthesizer(result, n_sample)
+
+        # Assert
+        assert isinstance(result._internal_synthesizer, TabDDPM)
+        assert isinstance(result, TabDDPMSynthesizer)
+        assert len(sampled_data) == n_sample
+
+    def test__sample_from_synthesizer(self):
+        """Test _sample_from_synthesizer generates data with the specified sample size."""
+        # Setup
+        trained_model = MagicMock()
+        trained_model._internal_synthesizer = MagicMock()
+        trained_model._internal_synthesizer.sample.return_value = MagicMock(
+            shape=(10, 5)
+        )  # Mock sample data shape
+        n_sample = 10
+        synthesizer = TabDDPMSynthesizer()
+
+        # Run
+        synthetic_data = synthesizer._sample_from_synthesizer(trained_model, n_sample)
+
+        # Assert
+        trained_model._internal_synthesizer.sample.assert_called_once_with(n_sample)
+        assert synthetic_data.shape[0] == n_sample, (
+            f'Expected {n_sample} rows, but got {synthetic_data.shape[0]}'
+        )

From f18a5ca70e6a915d225b7abd98ccff44ed9f0767 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Fri, 17 Jul 2026 11:09:24 -0400
Subject: [PATCH 03/28] update synthesizer list test

---
 tests/integration/synthesizers/test_utils.py | 1 +
 tests/unit/synthesizers/test_utils.py        | 1 +
 2 files changed, 2 insertions(+)

diff --git a/tests/integration/synthesizers/test_utils.py b/tests/integration/synthesizers/test_utils.py
index ad78c519..2ca3f895 100644
--- a/tests/integration/synthesizers/test_utils.py
+++ b/tests/integration/synthesizers/test_utils.py
@@ -15,6 +15,7 @@ def test_get_available_single_table_synthesizers():
         'GaussianCopulaSynthesizer',
         'RealTabFormerSynthesizer',
         'TVAESynthesizer',
+        'TabDDPMSynthesizer',
         'UniformSynthesizer',
     ]
 
diff --git a/tests/unit/synthesizers/test_utils.py b/tests/unit/synthesizers/test_utils.py
index 0881994a..0e7d5f5b 100644
--- a/tests/unit/synthesizers/test_utils.py
+++ b/tests/unit/synthesizers/test_utils.py
@@ -14,6 +14,7 @@ def test__get_supported_synthesizers():
         'MultiTableUniformSynthesizer',
         'RealTabFormerSynthesizer',
         'TVAESynthesizer',
+        'TabDDPMSynthesizer',
         'UniformSynthesizer',
     ]
 

From 431e05325870898930c551f4f18f15c8ab051fe8 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Fri, 17 Jul 2026 16:40:31 -0400
Subject: [PATCH 04/28] update code

---
 sdgym/synthesizers/tabddpm.py | 176 ++++++++++++++++------------------
 1 file changed, 82 insertions(+), 94 deletions(-)

diff --git a/sdgym/synthesizers/tabddpm.py b/sdgym/synthesizers/tabddpm.py
index bdeb9db5..33d0e0c7 100644
--- a/sdgym/synthesizers/tabddpm.py
+++ b/sdgym/synthesizers/tabddpm.py
@@ -10,6 +10,7 @@
 import logging
 import math
 import sys
+from copy import deepcopy
 from typing import Callable, List, Union
 
 import numpy as np
@@ -17,6 +18,7 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from sdv.metadata import Metadata
 from sklearn.preprocessing import OrdinalEncoder, QuantileTransformer
 from torch import Tensor
 
@@ -132,6 +134,16 @@ def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
     return np.array(betas)
 
 
+def update_ema(target_params, source_params, rate=0.999):
+    """Update parametes using EMA.
+
+    Update target parameters to be closer to those of source parameters using
+    an exponential moving average.
+    """
+    for targ, src in zip(target_params, source_params):
+        targ.detach().mul_(rate).add_(src.detach(), alpha=1 - rate)
+
+
 def timestep_embedding(timesteps, dim, max_period=10000):
     """Create sinusoidal timestep embeddings."""
     half = dim // 2
@@ -270,28 +282,33 @@ def forward(self, x, timesteps, y=None):
 
 
 class GaussianMultinomialDiffusion(torch.nn.Module):
-    """Joint diffusion: Gaussian over numerical features, multinomial over categorical.
-
-    Simplified code to use the paper's pipeline:
-    - 'mse' Gaussian loss
-    - 'eps' parametrization
-    - 'vb_stochastic' multinomial loss
-    - uniform time sampling and ancestral sampling
-    """
+    """Joint diffusion: Gaussian over numerical features, multinomial over categorical."""
 
     def __init__(
         self,
-        num_classes: np.ndarray,
+        num_classes: np.array,
         num_numerical_features: int,
         denoise_fn,
         num_timesteps=1000,
+        gaussian_loss_type='mse',
+        gaussian_parametrization='eps',
+        multinomial_loss_type='vb_stochastic',
+        parametrization='x0',
         scheduler='cosine',
         device=torch.device('cpu'),
     ):
         super(GaussianMultinomialDiffusion, self).__init__()
+        assert multinomial_loss_type in ('vb_stochastic', 'vb_all')
+        assert parametrization in ('x0', 'direct')
+
+        if multinomial_loss_type == 'vb_all':
+            sys.stdout.write(
+                'Computing the loss using the bound on _all_ timesteps. '
+                'This is expensive both in terms of memory and computation.\n'
+            )
 
         self.num_numerical_features = num_numerical_features
-        self.num_classes = num_classes  # it is a vector [K1, K2, ..., Km]
+        self.num_classes = num_classes  # it as a vector [K1, K2, ..., Km]
         self.num_classes_expanded = torch.from_numpy(
             np.concatenate([num_classes[i].repeat(num_classes[i]) for i in range(len(num_classes))])
         ).to(device)
@@ -303,7 +320,11 @@ def __init__(
         self.offsets = torch.from_numpy(np.append([0], offsets)).to(device)
 
         self._denoise_fn = denoise_fn
+        self.gaussian_loss_type = gaussian_loss_type
+        self.gaussian_parametrization = gaussian_parametrization
+        self.multinomial_loss_type = multinomial_loss_type
         self.num_timesteps = num_timesteps
+        self.parametrization = parametrization
         self.scheduler = scheduler
 
         alphas = 1.0 - get_named_beta_schedule(scheduler, num_timesteps)
@@ -365,6 +386,9 @@ def __init__(
             'sqrt_recipm1_alphas_cumprod', sqrt_recipm1_alphas_cumprod.float().to(device)
         )
 
+        self.register_buffer('Lt_history', torch.zeros(num_timesteps))
+        self.register_buffer('Lt_count', torch.zeros(num_timesteps))
+
     def _gaussian_q_sample(self, x_start, t, noise=None):
         if noise is None:
             noise = torch.randn_like(x_start)
@@ -844,9 +868,15 @@ def __iter__(self):
                 yield X[i : i + self.batch_size], y[i : i + self.batch_size]
 
 
-class _Trainer:
+class Trainer:
+    """Diffusion trainer."""
+
     def __init__(self, diffusion, train_iter, lr, weight_decay, steps, device, verbose=True):
         self.diffusion = diffusion
+        self.ema_model = deepcopy(self.diffusion._denoise_fn)
+        for param in self.ema_model.parameters():
+            param.detach_()
+
         self.train_iter = iter(train_iter)
         self.steps = steps
         self.init_lr = lr
@@ -854,9 +884,10 @@ def __init__(self, diffusion, train_iter, lr, weight_decay, steps, device, verbo
             self.diffusion.parameters(), lr=lr, weight_decay=weight_decay
         )
         self.device = device
-        self.verbose = verbose
+        self.loss_history = pd.DataFrame(columns=['step', 'mloss', 'gloss', 'loss'])
         self.log_every = 100
-        self.loss_history = []
+        self.print_every = 500
+        self.ema_every = 1000
 
     def _anneal_lr(self, step):
         frac_done = step / self.steps
@@ -876,13 +907,16 @@ def _run_step(self, x, out_dict):
         return loss_multi, loss_gauss
 
     def run_loop(self):
+        """Training loop."""
+        step = 0
         curr_loss_multi = 0.0
         curr_loss_gauss = 0.0
-        curr_count = 0
 
-        for step in range(self.steps):
-            x, y = next(self.train_iter)
-            batch_loss_multi, batch_loss_gauss = self._run_step(x, {'y': y})
+        curr_count = 0
+        while step < self.steps:
+            x, out_dict = next(self.train_iter)
+            out_dict = {'y': out_dict}
+            batch_loss_multi, batch_loss_gauss = self._run_step(x, out_dict)
 
             self._anneal_lr(step)
 
@@ -893,20 +927,24 @@ def run_loop(self):
             if (step + 1) % self.log_every == 0:
                 mloss = np.around(curr_loss_multi / curr_count, 4)
                 gloss = np.around(curr_loss_gauss / curr_count, 4)
-                self.loss_history.append({
-                    'step': step + 1,
-                    'mloss': mloss,
-                    'gloss': gloss,
-                    'loss': mloss + gloss,
-                })
-                if self.verbose:
+                if (step + 1) % self.print_every == 0:
                     sys.stdout.write(
-                        f'Step {step + 1}/{self.steps} MLoss: {mloss} GLoss: {gloss} '
-                        f'Sum: {np.around(mloss + gloss, 4)}\n'
+                        f'Step {(step + 1)}/{self.steps} '
+                        f'MLoss: {mloss} GLoss: {gloss} Sum: {mloss + gloss}\n'
                     )
+                self.loss_history.loc[len(self.loss_history)] = [
+                    step + 1,
+                    mloss,
+                    gloss,
+                    mloss + gloss,
+                ]
                 curr_count = 0
-                curr_loss_multi = 0.0
                 curr_loss_gauss = 0.0
+                curr_loss_multi = 0.0
+
+            update_ema(self.ema_model.parameters(), self.diffusion._denoise_fn.parameters())
+
+            step += 1
 
 
 class TabDDPM:
@@ -983,7 +1021,13 @@ def __init__(
         self.seed = seed
         self.verbose = verbose
 
-        self._table_name, self._table_metadata = self._parse_metadata(metadata)
+        if isinstance(metadata, dict):
+            metadata = Metadata.load_from_dict(metadata)
+
+        metadata.validate()
+        self._metadata = metadata
+        self._table_name = list(metadata.tables)[0]
+        self._table_metadata = metadata.tables[self._table_name].to_dict()
         if device is None:
             device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self._device = torch.device(device)
@@ -996,11 +1040,11 @@ def fit(self, data):
             data (pandas.DataFrame):
                 The data to fit the synthesizer to.
         """
+        data = data.copy()
         if isinstance(data, pd.DataFrame):
-            data = {self._table_name: data}
+            data_dict = {self._table_name: data}
 
-        table_name, df = self._validate_data(data)
-        self._table_name = table_name
+        self._metadata.validate_data(data_dict)
 
         torch.manual_seed(self.seed)
         np.random.seed(self.seed)
@@ -1012,7 +1056,7 @@ def fit(self, data):
         self._n_target_classes = 0
         self._target_encoder = None
         if target_column is not None:
-            if target_column not in df.columns:
+            if target_column not in data.columns:
                 raise ValueError(f"target_column '{target_column}' not found in the data.")
             sdtype = columns_metadata[target_column].get('sdtype', 'categorical')
             if sdtype not in ('categorical', 'boolean'):
@@ -1021,7 +1065,7 @@ def fit(self, data):
                     f'(got sdtype={sdtype!r}). Numerical columns are modelled '
                     'unconditionally; leave target_column unset.'
                 )
-            target_series = df[target_column]
+            target_series = data[target_column]
             target_values = (
                 target_series
                 .astype(object)
@@ -1036,16 +1080,16 @@ def fit(self, data):
             self._target_dtype = target_series.dtype
             self._target_is_boolean = sdtype == 'boolean'
             columns_metadata.pop(target_column)
-            df = df.drop(columns=[target_column])
+            data = data.drop(columns=[target_column])
         else:
-            y = np.zeros(len(df), dtype='int64')
+            y = np.zeros(len(data), dtype='int64')
 
         # Preprocess: quantile-normalized numerical block + ordinal categorical block
         self._transformer = _DataTransformer(
             columns_metadata, normalization=self.normalization, seed=self.seed
         )
-        self._transformer.fit(df)
-        X_num, X_cat = self._transformer.transform(df)
+        self._transformer.fit(data)
+        X_num, X_cat = self._transformer.transform(data)
 
         n_num = X_num.shape[1]
         K = self._transformer.category_sizes
@@ -1080,7 +1124,7 @@ def fit(self, data):
         X = torch.from_numpy(np.concatenate([X_num, X_cat.astype('float32')], axis=1)).float()
         train_loader = _FastTensorDataLoader(X, y_tensor, batch_size=self.batch_size)
 
-        trainer = _Trainer(
+        trainer = Trainer(
             self._diffusion,
             train_loader,
             lr=self.lr,
@@ -1145,62 +1189,6 @@ def sample(self, num_rows):
         df = df[[column for column in self._column_order if column in df.columns]]
         return df
 
-    @staticmethod
-    def _parse_metadata(metadata):
-        if hasattr(metadata, 'to_dict'):
-            metadata = metadata.to_dict()
-        if not isinstance(metadata, dict):
-            raise TypeError('metadata must be an SDV metadata object or its dict representation.')
-
-        if 'tables' in metadata:
-            tables = metadata['tables']
-            if len(tables) != 1:
-                raise ValueError(
-                    'TabDDPMSynthesizer is a single-table synthesizer; the metadata '
-                    f'describes {len(tables)} tables.'
-                )
-            table_name, table_metadata = next(iter(tables.items()))
-        elif 'columns' in metadata:
-            table_name, table_metadata = None, metadata
-        else:
-            raise ValueError(
-                "metadata dict must contain either a 'columns' key (single-table "
-                "format) or a 'tables' key (multi-table format)."
-            )
-
-        if not table_metadata.get('columns'):
-            raise ValueError('The table metadata does not define any columns.')
-        return table_name, table_metadata
-
-    def _validate_data(self, data):
-        if not isinstance(data, dict) or not all(
-            isinstance(df, pd.DataFrame) for df in data.values()
-        ):
-            raise TypeError('data must be a dictionary mapping table names to DataFrames.')
-        if len(data) != 1:
-            raise ValueError(
-                'TabDDPMSynthesizer is a single-table synthesizer; got '
-                f'{len(data)} tables: {sorted(data)}.'
-            )
-
-        table_name, df = next(iter(data.items()))
-        if self._table_name is not None and table_name != self._table_name:
-            raise ValueError(
-                f"The metadata describes table '{self._table_name}' but the data "
-                f"contains table '{table_name}'."
-            )
-
-        metadata_columns = set(self._table_metadata['columns'])
-        data_columns = set(df.columns)
-        missing = metadata_columns - data_columns
-        extra = data_columns - metadata_columns
-        if missing or extra:
-            raise ValueError(
-                'The data does not match the metadata. '
-                f'Missing columns: {sorted(missing)}. Unexpected columns: {sorted(extra)}.'
-            )
-        return table_name, df
-
 
 class TabDDPMSynthesizer(BaselineSynthesizer):
     """Custom wrapper for the TabDDPM synthesizer to make it work with SDGym."""

From 68feb7fb8652be8f7824cf80dd98ede814d870a0 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Wed, 22 Jul 2026 17:12:27 -0400
Subject: [PATCH 05/28] add clavaddpm and tests

---
 sdgym/synthesizers/__init__.py                |    2 +
 sdgym/synthesizers/clavaddpm.py               | 1022 +++++++++++++++++
 .../synthesizers/test_clavaddpm.py            |   32 +
 tests/unit/synthesizers/test_clavaddpm.py     |   98 ++
 4 files changed, 1154 insertions(+)
 create mode 100644 sdgym/synthesizers/clavaddpm.py
 create mode 100644 tests/integration/synthesizers/test_clavaddpm.py
 create mode 100644 tests/unit/synthesizers/test_clavaddpm.py

diff --git a/sdgym/synthesizers/__init__.py b/sdgym/synthesizers/__init__.py
index e17d3e77..1a55c32b 100644
--- a/sdgym/synthesizers/__init__.py
+++ b/sdgym/synthesizers/__init__.py
@@ -9,6 +9,7 @@
 from sdgym.synthesizers.column import ColumnSynthesizer
 from sdgym.synthesizers.realtabformer import RealTabFormerSynthesizer
 from sdgym.synthesizers.tabddpm import TabDDPMSynthesizer
+from sdgym.synthesizers.clavaddpm import ClavaDDPMSynthesizer
 from sdgym.synthesizers.uniform import UniformSynthesizer, MultiTableUniformSynthesizer
 from sdgym.synthesizers.utils import (
     get_available_single_table_synthesizers,
@@ -29,6 +30,7 @@
     'get_available_single_table_synthesizers',
     'get_available_multi_table_synthesizers',
     'MultiTableUniformSynthesizer',
+    'ClavaDDPMSynthesizer',
 ]
 
 for sdv_name in _get_all_sdv_synthesizers():
diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
new file mode 100644
index 00000000..ea59717f
--- /dev/null
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -0,0 +1,1022 @@
+"""ClavaDDPMSynthesizer -- SDGym synthesizer for ClavaDDPM.
+
+Paper: "ClavaDDPM: Multi-relational Data Synthesis with Cluster-guided Diffusion Models (2024)"
+https://arxiv.org/abs/2405.17724
+
+Original implementation is provided:
+https://github.com/weipang142857/ClavaDDPM/tree/main.
+"""
+
+import logging
+import random
+import sys
+from collections import defaultdict
+from datetime import datetime, timedelta
+
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.cluster import KMeans
+from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
+from sklearn.neighbors import NearestNeighbors
+from sklearn.preprocessing import LabelEncoder, MinMaxScaler, OneHotEncoder
+
+from sdgym.synthesizers.base import MultiTableBaselineSynthesizer
+from sdgym.synthesizers.tabddpm import CAT_MISSING_VALUE, TabDDPM, ohe_to_categories
+
+
+def get_group_data_dict(np_data, group_id_attrs=[0]):
+    """Grouping dictionary from pipeline_utils.py."""
+    group_data_dict = {}
+    data_len = len(np_data)
+    for i in range(data_len):
+        row_id = tuple(np_data[i, group_id_attrs])
+        if row_id not in group_data_dict:
+            group_data_dict[row_id] = []
+        group_data_dict[row_id].append(np_data[i])
+    return group_data_dict
+
+
+def get_group_data(np_data, group_id_attrs=[0]):
+    """Grouping list from pipeline_utils.py."""
+    group_data_list = []
+    data_len = len(np_data)
+    i = 0
+    while i < data_len:
+        group = []
+        row_id = np_data[i, group_id_attrs]
+        while (np_data[i, group_id_attrs] == row_id).all():
+            group.append(np_data[i])
+            i += 1
+            if i >= data_len:
+                break
+        group_data_list.append(np.array(group))
+    return np.array(group_data_list, dtype=object)
+
+
+def min_max_normalize_sklearn(matrix):
+    """Apply MinMaxScaler to each column from pipeline_utils.py."""
+    scaler = MinMaxScaler(feature_range=(-1, 1))
+    normalized_data = np.empty((matrix.shape[0], 0))
+    for col in range(matrix.shape[1]):
+        column = matrix[:, col].reshape(-1, 1)
+        transformed_column = scaler.fit_transform(column)
+        normalized_data = np.concatenate((normalized_data, transformed_column), axis=1)
+
+    return normalized_data
+
+
+def aggregate_and_sample(cluster_probabilities, child_group_lengths):
+    """Aggregate the distribution and sample from pipeline_modules.py."""
+    group_cluster_labels = []
+    curr_index = 0
+    agree_rates = []
+    for group_length in child_group_lengths:
+        group_probability_distribution = np.mean(
+            cluster_probabilities[curr_index : curr_index + group_length], axis=0
+        )
+        group_cluster_label = np.random.choice(
+            range(len(group_probability_distribution)), p=group_probability_distribution
+        )
+        group_cluster_labels.append(group_cluster_label)
+        agree_rates.append(np.max(group_probability_distribution))
+        curr_index += group_length
+    return group_cluster_labels, agree_rates
+
+
+def freq_to_prob(freq_dict):
+    """Converts a dict of frequencies to a dict of probabilities from pipeline_utils.py."""
+    prob_dict = {}
+    for key in freq_dict:
+        prob_dict[key] = freq_dict[key] / sum(list(freq_dict.values()))
+    return prob_dict
+
+
+def sample_from_dict(probabilities):
+    """Sample using a dict of probabilities from pipeline_utils.py."""
+    random_number = random.random()
+    cumulative_sum = 0
+    selected_key = None
+    for key, probability in probabilities.items():
+        cumulative_sum += probability
+        if cumulative_sum >= random_number:
+            selected_key = key
+            break
+    return selected_key
+
+
+def get_df_without_id(df, id_cols):
+    """Drop id columns based on `id_cols` from pipeline_utils.py."""
+    return df.drop(columns=[col for col in id_cols if col in df.columns])
+
+
+def calculate_days_since_earliest_date(dates, date_format='%y%m%d'):
+    """Encode a date column as integer days since its earliest date from preprocess_utils.py."""
+    parsed = [
+        datetime.strptime(str(date), date_format) if pd.notna(date) else None for date in dates
+    ]
+    earliest_date = min(date for date in parsed if date is not None)
+    days_since = [(date - earliest_date).days if date is not None else np.nan for date in parsed]
+    return days_since, earliest_date.strftime(date_format)
+
+
+def reconstruct_dates(days_since, earliest_date_str, date_format='%y%m%d'):
+    """Inverse of `calculate_days_since_earliest_date` from preprocess_utils.py."""
+    earliest_date = datetime.strptime(earliest_date_str, date_format)
+    return [
+        (earliest_date + timedelta(days=int(round(float(days))))).strftime(date_format)
+        if pd.notna(days)
+        else None
+        for days in days_since
+    ]
+
+
+def table_label_encode(df, discrete_cols):
+    """Label-encode the discrete columns of a table from preprocess_utils.py."""
+    df = df.copy()
+    label_encoders = {}
+    for col in discrete_cols:
+        le = LabelEncoder()
+        df[col] = le.fit_transform(df[col])
+        label_encoders[col] = le
+    return df, label_encoders
+
+
+def table_label_decode(df, label_encoders):
+    """Inverse of :func:`table_label_encode` from preprocess_utils.py."""
+    df = df.copy()
+    for col, le in label_encoders.items():
+        df[col] = le.inverse_transform(df[col])
+    return df
+
+
+def get_domain(df, id_cols, discrete_cols):
+    """Build the ``{col: {'size', 'type'}}`` domain of a table from preprocess_utils.py."""
+    domain = {}
+    for col in df.columns:
+        if col in discrete_cols:
+            domain[col] = {'size': len(df[col].unique()), 'type': 'discrete'}
+        elif col not in id_cols:
+            domain[col] = {'size': len(df[col].unique()), 'type': 'continuous'}
+    return domain
+
+
+def topological_sort(graph):
+    """Order tables into ``[parent, child]`` relations from preprocess_utils.py."""
+    in_degree = {node: 0 for node in graph}
+    for node in graph:
+        for child in graph[node]['children']:
+            in_degree[child] += 1
+
+    zero_in_degree = [node for node, degree in in_degree.items() if degree == 0]
+
+    sorted_order = []
+    for node in zero_in_degree:
+        sorted_order.append([None, node])
+
+    queue = zero_in_degree[:]
+    while queue:
+        current = queue.pop(0)
+        for child in graph[current]['children']:
+            in_degree[child] -= 1
+            if in_degree[child] == 0:
+                queue.append(child)
+            sorted_order.append([current, child])
+
+    return sorted_order
+
+
+def pair_clustering_keep_id(
+    child_df,
+    child_domain_dict,
+    parent_df,
+    parent_domain_dict,
+    child_primary_key,
+    parent_primary_key,
+    foreign_key,
+    num_clusters,
+    parent_scale,
+    key_scale,
+    parent_name,
+    child_name,
+    clustering_method='kmeans',
+    seed=0,
+):
+    """Cluster child rows augmented with their parent's features from pipeline_modules.py."""
+    original_child_cols = list(child_df.columns)
+    original_parent_cols = list(parent_df.columns)
+
+    relation_cluster_name = f'{parent_name}_{child_name}_cluster'
+
+    child_data = child_df.to_numpy()
+    parent_data = parent_df.to_numpy()
+
+    child_num_cols = []
+    child_cat_cols = []
+
+    parent_num_cols = []
+    parent_cat_cols = []
+
+    for col_index, col in enumerate(original_child_cols):
+        if col in child_domain_dict:
+            if child_domain_dict[col]['type'] == 'discrete':
+                child_cat_cols.append((col_index, col))
+            else:
+                child_num_cols.append((col_index, col))
+
+    for col_index, col in enumerate(original_parent_cols):
+        if col in parent_domain_dict:
+            if parent_domain_dict[col]['type'] == 'discrete':
+                parent_cat_cols.append((col_index, col))
+            else:
+                parent_num_cols.append((col_index, col))
+
+    child_primary_key_index = original_child_cols.index(child_primary_key)
+    parent_primary_key_index = original_parent_cols.index(parent_primary_key)
+    foreing_key_index = original_child_cols.index(foreign_key)
+
+    # sort child data by foreign key
+    sorted_child_data = child_data[np.argsort(child_data[:, foreing_key_index])]
+    child_group_data_dict = get_group_data_dict(sorted_child_data, [foreing_key_index])
+
+    # sort parent data by primary key
+    sorted_parent_data = parent_data[np.argsort(parent_data[:, parent_primary_key_index])]
+
+    group_lengths = []
+    unique_group_ids = sorted_parent_data[:, parent_primary_key_index]
+    for group_id in unique_group_ids:
+        group_id = tuple([group_id])
+        if group_id not in child_group_data_dict:
+            group_lengths.append(0)
+        else:
+            group_lengths.append(len(child_group_data_dict[group_id]))
+
+    group_lengths = np.array(group_lengths, dtype=int)
+
+    sorted_parent_data_repeated = np.repeat(sorted_parent_data, group_lengths, axis=0)
+    assert (
+        sorted_parent_data_repeated[:, parent_primary_key_index]
+        == sorted_child_data[:, foreing_key_index]
+    ).all()
+
+    child_group_data = get_group_data(sorted_child_data, [foreing_key_index])
+
+    sorted_child_num_data = sorted_child_data[:, [ci for ci, _ in child_num_cols]]
+    sorted_child_cat_data = sorted_child_data[:, [ci for ci, _ in child_cat_cols]]
+    sorted_parent_num_data = sorted_parent_data_repeated[:, [ci for ci, _ in parent_num_cols]]
+    sorted_parent_cat_data = sorted_parent_data_repeated[:, [ci for ci, _ in parent_cat_cols]]
+
+    joint_num_matrix = np.concatenate([sorted_child_num_data, sorted_parent_num_data], axis=1)
+    joint_cat_matrix = np.concatenate([sorted_child_cat_data, sorted_parent_cat_data], axis=1)
+
+    # Impute missing numerical values with their column mean before clustering.
+    # The reference assumed pre-imputed data; this mirrors the mean-fill that
+    # TabDDPM's own transformer applies, so clustering (which cannot take NaNs)
+    # is robust to real-world tables with missing values.
+    if joint_num_matrix.shape[1] > 0:
+        joint_num_matrix = joint_num_matrix.astype(float)
+        col_means = np.nanmean(joint_num_matrix, axis=0)
+        col_means = np.where(np.isnan(col_means), 0.0, col_means)
+        nan_positions = np.where(np.isnan(joint_num_matrix))
+        joint_num_matrix[nan_positions] = np.take(col_means, nan_positions[1])
+
+    joint_num_matrix_p_index = sorted_child_num_data.shape[1]
+    cat_one_hot = None
+    if joint_cat_matrix.shape[1] > 0:
+        joint_cat_matrix_p_index = sorted_child_cat_data.shape[1]
+
+        cat_converted = []
+        for i in range(joint_cat_matrix.shape[1]):
+            # skip huge categoricals to avoid an explosive one-hot encoding
+            if len(np.unique(joint_cat_matrix[:, i])) > 1000:
+                continue
+            label_encoder = LabelEncoder()
+            cat_converted.append(label_encoder.fit_transform(joint_cat_matrix[:, i]).astype(float))
+        cat_converted = np.vstack(cat_converted).T
+
+        cat_one_hot = np.empty((cat_converted.shape[0], 0))
+        for col in range(cat_converted.shape[1]):
+            encoder = OneHotEncoder(sparse_output=False)
+            column = cat_converted[:, col].reshape(-1, 1)
+            cat_one_hot = np.concatenate((cat_one_hot, encoder.fit_transform(column)), axis=1)
+
+        cat_one_hot[:, joint_cat_matrix_p_index:] = (
+            parent_scale * cat_one_hot[:, joint_cat_matrix_p_index:]
+        )
+
+    num_min_max = min_max_normalize_sklearn(joint_num_matrix)
+
+    # key channel: parent identity, factorized so arbitrary key types normalize
+    key_factorized = pd.factorize(sorted_parent_data_repeated[:, parent_primary_key_index])[0]
+    key_min_max = min_max_normalize_sklearn(key_factorized.astype(float).reshape(-1, 1))
+    key_scaled = key_scale * key_min_max
+
+    num_min_max[:, joint_num_matrix_p_index:] = (
+        parent_scale * num_min_max[:, joint_num_matrix_p_index:]
+    )
+
+    if joint_cat_matrix.shape[1] > 0:
+        cluster_data = np.concatenate((num_min_max, cat_one_hot, key_scaled), axis=1)
+    else:
+        cluster_data = np.concatenate((num_min_max, key_scaled), axis=1)
+
+    child_group_lengths = np.array([len(group) for group in child_group_data], dtype=int)
+    num_clusters = min(num_clusters, len(cluster_data))
+
+    if clustering_method == 'kmeans':
+        kmeans = KMeans(n_clusters=num_clusters, n_init='auto', init='k-means++', random_state=seed)
+        kmeans.fit(cluster_data)
+        cluster_labels = kmeans.labels_
+    elif clustering_method == 'both':
+        gmm = GaussianMixture(
+            n_components=num_clusters,
+            covariance_type='diag',
+            init_params='k-means++',
+            tol=0.0001,
+            random_state=seed,
+        )
+        gmm.fit(cluster_data)
+        cluster_labels = gmm.predict(cluster_data)
+    elif clustering_method == 'variational':
+        gmm = BayesianGaussianMixture(
+            n_components=num_clusters,
+            covariance_type='diag',
+            init_params='k-means++',
+            tol=0.0001,
+            random_state=seed,
+        )
+        gmm.fit(cluster_data)
+        cluster_labels = gmm.predict_proba(cluster_data)
+    elif clustering_method == 'gmm':
+        gmm = GaussianMixture(n_components=num_clusters, covariance_type='diag', random_state=seed)
+        gmm.fit(cluster_data)
+        cluster_labels = gmm.predict(cluster_data)
+
+    if clustering_method == 'variational':
+        group_cluster_labels, agree_rates = aggregate_and_sample(
+            cluster_labels, child_group_lengths
+        )
+    else:
+        # voting to determine the cluster label for each parent
+        group_cluster_labels = []
+        curr_index = 0
+        agree_rates = []
+        for group_length in child_group_lengths:
+            most_common_label_count = np.max(
+                np.bincount(cluster_labels[curr_index : curr_index + group_length])
+            )
+            group_cluster_label = np.argmax(
+                np.bincount(cluster_labels[curr_index : curr_index + group_length])
+            )
+            group_cluster_labels.append(group_cluster_label)
+            agree_rates.append(most_common_label_count / group_length)
+            curr_index += group_length
+
+    group_assignment = np.repeat(group_cluster_labels, child_group_lengths, axis=0).reshape((-1, 1))
+    sorted_child_data_with_cluster = np.concatenate([sorted_child_data, group_assignment], axis=1)
+
+    # per-cluster distribution of "how many children a parent has"
+    group_labels_list = group_cluster_labels
+    group_lengths_list = child_group_lengths.tolist()
+    group_lengths_dict = {}
+    for i in range(len(group_labels_list)):
+        group_label = group_labels_list[i]
+        if group_label not in group_lengths_dict:
+            group_lengths_dict[group_label] = defaultdict(int)
+        group_lengths_dict[group_label][group_lengths_list[i]] += 1
+
+    group_lengths_prob_dicts = {
+        group_label: freq_to_prob(freq_dict)
+        for group_label, freq_dict in group_lengths_dict.items()
+    }
+
+    # attach cluster label to the child in its original order
+    sorted_child_ids = sorted_child_data[:, child_primary_key_index]
+    child_id_to_cluster = dict(zip(sorted_child_ids, group_assignment.flatten()))
+    child_df_with_cluster = child_df.copy()
+    child_df_with_cluster[relation_cluster_name] = (
+        child_df[child_primary_key].map(child_id_to_cluster).astype(int)
+    )
+
+    # a parent inherits its children's cluster; childless parents get a fresh id
+    parent_id_to_cluster = {}
+    for i in range(len(sorted_child_data)):
+        parent_id = sorted_child_data[i, foreing_key_index]
+        if parent_id in parent_id_to_cluster:
+            continue
+        parent_id_to_cluster[parent_id] = sorted_child_data_with_cluster[i, -1]
+
+    max_cluster_label = max(parent_id_to_cluster.values())
+    parent_df_with_cluster = parent_df.copy()
+    parent_clusters = parent_df[parent_primary_key].map(parent_id_to_cluster)
+    parent_df_with_cluster[relation_cluster_name] = parent_clusters.fillna(
+        max_cluster_label + 1
+    ).astype(int)
+
+    new_col_entry = {
+        'type': 'discrete',
+        'size': len(set(parent_df_with_cluster[relation_cluster_name])),
+    }
+    parent_domain_dict[relation_cluster_name] = new_col_entry.copy()
+    child_domain_dict[relation_cluster_name] = new_col_entry.copy()
+
+    return parent_df_with_cluster, child_df_with_cluster, group_lengths_prob_dicts
+
+
+def match_tables(A, B, n_clusters=25, unique_matching=True, batch_size=100):
+    """Nearest-neighbour match of every row of ``A`` to a row of ``B``.
+
+    n_clusters is used by the original implementation for faiss.IndexIVFFlat.
+    """
+    A = np.ascontiguousarray(A, dtype=np.float32)
+    B = np.ascontiguousarray(B, dtype=np.float32)
+
+    if not unique_matching:
+        distances, indices = NearestNeighbors(n_neighbors=1).fit(B).kneighbors(A)
+        return indices.flatten().tolist(), distances.flatten().tolist()
+
+    k = min(len(B), 50)
+    distances, indices = NearestNeighbors(n_neighbors=k).fit(B).kneighbors(A)
+    used = set()
+    matched_indices = []
+    matched_distances = []
+    for row_indices, row_distances in zip(indices, distances):
+        chosen, chosen_distance = None, 0.0
+        for candidate, distance in zip(row_indices, row_distances):
+            if int(candidate) not in used:
+                chosen, chosen_distance = int(candidate), float(distance)
+                break
+        if chosen is None:
+            chosen = next(j for j in range(len(B)) if j not in used)
+        used.add(chosen)
+        matched_indices.append(chosen)
+        matched_distances.append(chosen_distance)
+
+    return matched_indices, matched_distances
+
+
+def handle_multi_parent(
+    child,
+    parents,
+    synthetic_tables,
+    id_cols,
+    n_clusters=1,
+    unique_matching=True,
+    batch_size=100,
+    no_matching=False,
+):
+    """Reconcile a child generated once per parent into a single table."""
+    synthetic_child_dfs = [
+        (synthetic_tables[(parent, child)]['df'].copy(), fk) for parent, fk, _ppk in parents
+    ]
+    anchor_index = int(np.argmin([len(df) for df, _ in synthetic_child_dfs]))
+    anchor = synthetic_child_dfs[anchor_index]
+    synthetic_child_dfs.pop(anchor_index)
+    for df, fk in synthetic_child_dfs:
+        df_without_ids = get_df_without_id(df, id_cols)
+        anchor_df_without_ids = get_df_without_id(anchor[0], id_cols)
+
+        df_val = df_without_ids.to_numpy().astype(float)
+        anchor_val = anchor_df_without_ids.to_numpy().astype(float)
+        if len(df_val.shape) == 1:
+            df_val = df_val.reshape(-1, 1)
+            anchor_val = anchor_val.reshape(-1, 1)
+
+        indices, _ = match_tables(
+            anchor_val,
+            df_val,
+            n_clusters=n_clusters,
+            unique_matching=unique_matching,
+            batch_size=batch_size,
+        )
+        if no_matching:
+            indices = np.random.permutation(indices)
+        df = df.iloc[indices]
+        anchor[0][fk] = df[fk].to_numpy()
+    return anchor[0]
+
+
+@torch.no_grad()
+def _sample_step(diffusion, y):
+    """One reverse-diffusion pass conditioned on the label tensor ``y``.
+
+    Mirrors ``GaussianMultinomialDiffusion._sample`` but takes explicit
+    per-row labels instead of drawing them from the empirical distribution.
+    Similar to the original classifier-guided ``conditional_sample`` implemtation.
+    """
+    device = diffusion.log_alpha.device
+    b = y.shape[0]
+    z_norm = torch.randn((b, diffusion.num_numerical_features), device=device)
+
+    has_cat = diffusion.num_classes[0] != 0
+    log_z = torch.zeros((b, 0), device=device).float()
+    if has_cat:
+        uniform_logits = torch.zeros((b, len(diffusion.num_classes_expanded)), device=device)
+        log_z = diffusion._log_sample_categorical(uniform_logits)
+
+    out_dict = {'y': y.long().to(device)}
+    for i in reversed(range(diffusion.num_timesteps)):
+        t = torch.full((b,), i, device=device, dtype=torch.long)
+        model_out = diffusion._denoise_fn(torch.cat([z_norm, log_z], dim=1).float(), t, **out_dict)
+        model_out_num = model_out[:, : diffusion.num_numerical_features]
+        model_out_cat = model_out[:, diffusion.num_numerical_features :]
+        if diffusion.num_numerical_features > 0:
+            z_norm = diffusion._gaussian_p_sample(model_out_num, z_norm, t)['sample']
+        if has_cat:
+            log_z = diffusion._p_sample(model_out_cat, log_z, t)
+
+    z_cat = log_z
+    if has_cat:
+        z_cat = ohe_to_categories(torch.exp(log_z).round(), diffusion.num_classes)
+    return torch.cat([z_norm, z_cat], dim=1).cpu()
+
+
+def _batched_conditional_sample(diffusion, y_codes, batch_size):
+    """Generate one row per entry of ``y_codes``, retrying NaN rows."""
+    n = len(y_codes)
+    n_num = diffusion.num_numerical_features
+    n_cat = len(diffusion.num_classes) if diffusion.num_classes[0] != 0 else 0
+    out = np.zeros((n, n_num + n_cat), dtype='float32')
+
+    pending = np.arange(n)
+    for _ in range(20):
+        if len(pending) == 0:
+            break
+        still = []
+        for start in range(0, len(pending), batch_size):
+            idx = pending[start : start + batch_size]
+            sample = _sample_step(diffusion, torch.from_numpy(y_codes[idx])).numpy()
+            bad = np.isnan(sample).any(axis=1)
+            out[idx[~bad]] = sample[~bad]
+            still.extend(idx[bad].tolist())
+        pending = np.array(still, dtype=int)
+
+    return out
+
+
+def sample_from_diffusion(model, labels, sample_batch_size):
+    """Decode a fitted ``TabDDPM`` conditioned on cluster ``labels``."""
+    if len(labels) == 0:
+        return pd.DataFrame(columns=list(model._column_order))
+
+    y_str = np.array([str(int(v)) for v in labels]).reshape(-1, 1)
+    y_codes = model._target_encoder.transform(y_str).reshape(-1).astype('int64')
+
+    diffusion = model._diffusion
+    diffusion.eval()
+    x_gen = _batched_conditional_sample(diffusion, y_codes, sample_batch_size)
+
+    n_num = diffusion.num_numerical_features
+    has_cat = model._transformer.category_sizes[0] != 0
+    x_num = x_gen[:, :n_num]
+    x_cat = x_gen[:, n_num:] if has_cat else np.empty((len(x_gen), 0), dtype='int64')
+    df = model._transformer.inverse_transform(x_num, x_cat)
+
+    decoded = model._target_encoder.inverse_transform(y_codes.reshape(-1, 1)).reshape(-1)
+    series = pd.Series(decoded, dtype=object).where(lambda s: s != CAT_MISSING_VALUE, np.nan)
+    if model._target_is_boolean:
+        series = series.map({'True': True, 'False': False})
+    else:
+        try:
+            series = series.astype(model._target_dtype)
+        except (ValueError, TypeError):
+            pass
+    df[model.target_column] = series
+
+    return df[[column for column in model._column_order if column in df.columns]]
+
+
+class ClavaDDPM:
+    """Cluster-guided multi-table diffusion synthesizer.
+
+    Args:
+        metadata (sdv.metadata.Metadata):
+            The metadata describing the data.
+        num_clusters (int or dict):
+            Latent clusters per relationship; an int applies to all, a dict maps
+            a child table name to its own count.
+        parent_scale (float):
+            Weight of the parent's features when clustering.
+        key_scale (float):
+            Weight of the parent-identity (key) channel when clustering.
+        clustering_method (str):
+            ``'both'`` / ``'gmm'`` (Gaussian mixture), ``'kmeans'`` or
+            ``'variational'`` (Bayesian GMM).
+        d_layers (List[int]):
+            Hidden layer sizes of the MLP denoiser.
+        dropout (float):
+            Dropout of the MLP denoiser.
+        dim_t (int):
+            Timestep/label embedding dimension.
+        num_timesteps (int):
+            Diffusion timesteps T.
+        scheduler (str):
+            ``'cosine'`` or ``'linear'`` beta schedule.
+        diffusion_iterations (int):
+            Training iterations.
+        batch_size (int):
+            Batch size.
+        lr (float):
+            Learning rate for the optimizer.
+        weight_decay (float):
+            Weight decay for the optimizer.
+        sampling_batch_size (int):
+            Batch size used during sampling.
+        num_matching_clusters (int):
+            Number of clusters to match. Default 1.
+        matching_batch_size (int):
+            Number of samples in the batch. Default 1000.
+        unique_matching (bool):
+            Multi-parent matching controls for unique matching.
+        no_matching (bool):
+            Multi-parent matching controls for no matching.
+        device (str or None):
+            Whether to use ``'cuda'`` or ``'cpu'``. If None, auto-select is used.
+        seed (int):
+            Random seed.
+        verbose (bool):
+            Show progress.
+    """
+
+    def __init__(
+        self,
+        metadata,
+        num_clusters=25,
+        parent_scale=1.0,
+        key_scale=1.0,
+        clustering_method='both',
+        d_layers=(512, 1024, 1024, 1024, 1024, 512),
+        dropout=0.0,
+        dim_t=128,
+        num_timesteps=2000,
+        scheduler='cosine',
+        diffusion_iterations=100000,
+        batch_size=4096,
+        lr=6e-4,
+        weight_decay=1e-5,
+        sampling_batch_size=20000,
+        num_matching_clusters=1,
+        matching_batch_size=1000,
+        unique_matching=True,
+        no_matching=False,
+        device=None,
+        seed=0,
+        verbose=False,
+    ):
+        self.num_clusters = num_clusters
+        self.parent_scale = parent_scale
+        self.key_scale = key_scale
+        self.clustering_method = clustering_method
+        self.d_layers = list(d_layers)
+        self.dropout = dropout
+        self.dim_t = dim_t
+        self.num_timesteps = num_timesteps
+        self.scheduler = scheduler
+        self.diffusion_iterations = diffusion_iterations
+        self.batch_size = batch_size
+        self.lr = lr
+        self.weight_decay = weight_decay
+        self.sampling_batch_size = sampling_batch_size
+        self.num_matching_clusters = num_matching_clusters
+        self.matching_batch_size = matching_batch_size
+        self.unique_matching = unique_matching
+        self.no_matching = no_matching
+        self.device = device
+        self.seed = seed
+        self.verbose = verbose
+
+        self._parse_metadata(metadata)
+        self._fitted = False
+
+    def _parse_metadata(self, metadata):
+        meta = metadata if isinstance(metadata, dict) else metadata.to_dict()
+
+        self._table_names = list(meta['tables'])
+        self._primary_key = {}
+        # child -> list of (parent_table, foreign_key_col, parent_primary_key)
+        self._parents = {name: [] for name in self._table_names}
+        self._children = {name: [] for name in self._table_names}
+
+        for name, table_meta in meta['tables'].items():
+            self._primary_key[name] = table_meta.get('primary_key')
+
+        for relationship in meta.get('relationships', []):
+            parent = relationship['parent_table_name']
+            child = relationship['child_table_name']
+            self._parents[child].append((
+                parent,
+                relationship['child_foreign_key'],
+                relationship['parent_primary_key'],
+            ))
+            self._children[parent].append(child)
+
+        # Preprocessing
+        self._discrete_cols = {}
+        self._datetime_cols = {}  # name -> {col: datetime_format}
+        for name, table_meta in meta['tables'].items():
+            id_cols = self._id_cols(name)
+            discrete, datetimes = [], {}
+            for column, spec in table_meta['columns'].items():
+                sdtype = spec.get('sdtype', 'categorical')
+                if column in id_cols or sdtype == 'id':
+                    continue
+                if sdtype == 'datetime':
+                    datetimes[column] = spec.get('datetime_format')
+                elif sdtype != 'numerical':
+                    discrete.append(column)
+            self._discrete_cols[name] = discrete
+            self._datetime_cols[name] = datetimes
+
+        graph = {name: {'children': self._children[name]} for name in self._table_names}
+        self._relation_order = [tuple(edge) for edge in topological_sort(graph)]
+
+    def _num_clusters_for(self, child):
+        if isinstance(self.num_clusters, dict):
+            return self.num_clusters[child]
+        return self.num_clusters
+
+    def _id_cols(self, table):
+        return [self._primary_key[table]] + [fk for _p, fk, _ppk in self._parents[table]]
+
+    def fit(self, data):
+        """Fit this model to the original data.
+
+        Args:
+            data (dict):
+                Dictionary mapping each table name to a ``pandas.DataFrame``.
+        """
+        missing = set(self._table_names) - set(data)
+        if missing:
+            raise ValueError(f'Missing tables in data: {sorted(missing)}')
+
+        random.seed(self.seed)
+        np.random.seed(self.seed)
+        torch.manual_seed(self.seed)
+
+        # Preprocess every table up front: turn datetimes into numeric
+        # day-counts, label-encode the discrete columns, and derive the
+        # ``domain``. Encoders / date anchors are kept so ``sample`` can invert
+        # them.
+        self._tables = {}
+        for name in self._table_names:
+            output_cols = list(data[name].columns)
+            df = data[name].copy().reset_index(drop=True)
+
+            date_info = {}
+            for col, date_format in self._datetime_cols[name].items():
+                if date_format is None:
+                    df[col] = pd.to_datetime(df[col], errors='coerce').dt.strftime('%Y-%m-%d')
+                    date_format = '%Y-%m-%d'
+                days_since, earliest = calculate_days_since_earliest_date(df[col], date_format)
+                df[col] = np.asarray(days_since, dtype=float)
+                date_info[col] = (earliest, date_format)
+
+            df, label_encoders = table_label_encode(df, self._discrete_cols[name])
+            domain = get_domain(df, self._id_cols(name), self._discrete_cols[name])
+
+            self._tables[name] = {
+                'df': df,
+                'domain': domain,
+                'output_cols': output_cols,
+                'original_len': len(data[name]),
+                'pk': self._primary_key[name],
+                'parents': list(self._parents[name]),
+                'label_encoders': label_encoders,
+                'date_info': date_info,
+            }
+
+        self._clustering()
+        self._training()
+
+        self._fitted = True
+        return self
+
+    def _clustering(self):
+        """clava_clustering: cluster every parent->child relationship."""
+        relation_order_reversed = self._relation_order[::-1]
+        self._all_group_lengths_prob_dicts = {}
+
+        for parent, child in relation_order_reversed:
+            if parent is None:
+                continue
+            if self.verbose:
+                sys.stdout.write(f'Clustering {parent} -> {child}')
+
+            fk = next(f for p, f, _ppk in self._tables[child]['parents'] if p == parent)
+
+            parent_df, child_df, group_lengths_prob_dicts = pair_clustering_keep_id(
+                self._tables[child]['df'],
+                self._tables[child]['domain'],
+                self._tables[parent]['df'],
+                self._tables[parent]['domain'],
+                self._tables[child]['pk'],
+                self._tables[parent]['pk'],
+                fk,
+                self._num_clusters_for(child),
+                self.parent_scale,
+                self.key_scale,
+                parent,
+                child,
+                clustering_method=self.clustering_method,
+                seed=self.seed,
+            )
+            self._tables[parent]['df'] = parent_df
+            self._tables[child]['df'] = child_df
+            self._all_group_lengths_prob_dicts[(parent, child)] = group_lengths_prob_dicts
+
+    def _training(self):
+        """clava_training: train one TabDDPM per relationship (child_training)."""
+        self._models = {}
+        for parent, child in self._relation_order:
+            if self.verbose:
+                sys.stdout.write(f'Training {parent} -> {child}')
+            df_with_cluster = self._tables[child]['df']
+            df_without_id = get_df_without_id(df_with_cluster, self._id_cols(child))
+            self._models[(parent, child)] = self._child_training(
+                df_without_id, self._tables[child]['domain'], parent, child
+            )
+
+    def _child_training(self, df_without_id, domain, parent, child):
+        """child_training: fit a (conditional) TabDDPM for one relationship.
+
+        Root relationships (``parent is None``) train unconditionally; child
+        relationships condition on the ``{parent}_{child}_cluster`` column via
+        TabDDPM's ``target_column`` label embedding.
+        """
+        y_col = None if parent is None else f'{parent}_{child}_cluster'
+
+        columns_meta = {}
+        for column in df_without_id.columns:
+            if column not in domain:
+                continue
+            sdtype = 'categorical' if domain[column]['type'] == 'discrete' else 'numerical'
+            columns_meta[column] = {'sdtype': sdtype}
+        table_metadata = {
+            'METADATA_SPEC_VERSION': 'V1',
+            'tables': {child: {'columns': columns_meta}},
+        }
+
+        model = TabDDPM(
+            table_metadata,
+            target_column=y_col,
+            d_layers=self.d_layers,
+            dropout=self.dropout,
+            dim_t=self.dim_t,
+            num_timesteps=self.num_timesteps,
+            scheduler=self.scheduler,
+            steps=self.diffusion_iterations,
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+            batch_size=self.batch_size,
+            sample_batch_size=self.sampling_batch_size,
+            device=self.device,
+            seed=self.seed,
+            verbose=self.verbose,
+        )
+        model.fit(df_without_id[list(columns_meta)])
+        return model
+
+    def sample(self, scale=1.0):
+        """Generate synthetic data for the entire dataset.
+
+        Args:
+            scale (float):
+                A float representing how much to scale the data by. If scale is set to ``1.0``,
+                this does not scale the sizes of the tables. If ``scale`` is greater than ``1.0``
+                create more rows than the original data by a factor of ``scale``.
+                If ``scale`` is lower than ``1.0`` create fewer rows by the factor of ``scale``
+                than the original tables. Defaults to ``1.0``.
+        """
+        if not self._fitted:
+            raise RuntimeError('The synthesizer has not been fitted; call fit() first.')
+
+        synthetic_tables = {}
+
+        for parent, child in self._relation_order:
+            if self.verbose:
+                sys.stdout.write(f'Generating {parent} -> {child}')
+            result = self._models[(parent, child)]
+            df_with_cluster = self._tables[child]['df']
+            df_without_id = get_df_without_id(df_with_cluster, self._id_cols(child))
+
+            if parent is None:
+                child_generated = result.sample(round(scale * len(df_without_id)))
+                child_keys = list(range(len(child_generated)))
+                generated = child_generated.copy()
+                generated.insert(0, self._tables[child]['pk'], child_keys)
+                synthetic_tables[(parent, child)] = {'df': generated, 'keys': child_keys}
+            else:
+                # any already-generated table for this parent carries its cluster
+                for (_p, cname), value in synthetic_tables.items():
+                    if cname == parent:
+                        parent_synthetic_df = value['df']
+                        parent_keys = value['keys']
+                        break
+
+                cluster_col = f'{parent}_{child}_cluster'
+                group_labels = (
+                    parent_synthetic_df[cluster_col].astype(float).round().astype(int).tolist()
+                )
+                group_lengths_prob_dicts = self._all_group_lengths_prob_dicts[(parent, child)]
+
+                sampled_group_sizes = []
+                ys = []
+                for group_label in group_labels:
+                    if group_label not in group_lengths_prob_dicts:
+                        sampled_group_sizes.append(0)
+                        continue
+                    sampled_group_size = sample_from_dict(group_lengths_prob_dicts[group_label])
+                    sampled_group_sizes.append(sampled_group_size)
+                    ys.extend([group_label] * sampled_group_size)
+
+                child_generated = sample_from_diffusion(
+                    result, np.array(ys), self.sampling_batch_size
+                )
+
+                child_foreign_keys = np.repeat(parent_keys, sampled_group_sizes, axis=0)
+                child_primary_keys = np.arange(len(child_generated))
+                fk = next(f for p, f, _ppk in self._tables[child]['parents'] if p == parent)
+
+                generated = child_generated.copy().reset_index(drop=True)
+                generated[self._tables[child]['pk']] = child_primary_keys
+                generated[fk] = child_foreign_keys
+                synthetic_tables[(parent, child)] = {
+                    'df': generated,
+                    'keys': list(child_primary_keys),
+                }
+
+        final_tables = {}
+        for parent, child in self._relation_order:
+            if child in final_tables:
+                continue
+            if len(self._tables[child]['parents']) > 1:
+                final_tables[child] = handle_multi_parent(
+                    child,
+                    self._tables[child]['parents'],
+                    synthetic_tables,
+                    self._id_cols(child),
+                    n_clusters=self.num_matching_clusters,
+                    unique_matching=self.unique_matching,
+                    batch_size=self.matching_batch_size,
+                    no_matching=self.no_matching,
+                )
+            else:
+                final_tables[child] = synthetic_tables[(parent, child)]['df']
+
+        return {child: self._postprocess(child, table) for child, table in final_tables.items()}
+
+    def _postprocess(self, table, df):
+        """Invert fit-time preprocessing: decode discretes, rebuild dates, order cols."""
+        df = df.copy()
+        state = self._tables[table]
+
+        # decode label-encoded discrete columns back to their original values
+        encoders = {col: le for col, le in state['label_encoders'].items() if col in df.columns}
+        for col in encoders:
+            df[col] = df[col].astype(int)
+        df = table_label_decode(df, encoders)
+
+        # turn numeric day-counts back into formatted date strings
+        for col, (earliest, date_format) in state['date_info'].items():
+            if col in df.columns:
+                df[col] = reconstruct_dates(df[col].to_numpy(), earliest, date_format)
+
+        return df[state['output_cols']].reset_index(drop=True)
+
+
+class ClavaDDPMSynthesizer(MultiTableBaselineSynthesizer):
+    """Custom wrapper for the ClavaDDPM synthesizer to make it work with SDGym."""
+
+    LOGGER = logging.getLogger(__name__)
+    _MODEL_KWARGS = None
+    _MODALITY_FLAG = 'multi_table'
+
+    def _fit(self, data, metadata):
+        """Fit the synthesizer to the multi-table data.
+
+        Args:
+            data (dict):
+                A dict mapping table name to table data.
+            metadata (sdv.metadata.MultiTableMetadata):
+                The multi-table metadata describing the data.
+        """
+        model_kwargs = self._MODEL_KWARGS.copy() if self._MODEL_KWARGS else {}
+        model = ClavaDDPM(metadata, **model_kwargs)
+        model.fit(data)
+
+        self._internal_synthesizer = model
+
+    def _sample_from_synthesizer(self, synthesizer, scale):
+        """Sample data from the provided synthesizer.
+
+        Args:
+            synthesizer (SDGym synthesizer):
+                The synthesizer object to sample data from.
+            scale (float):
+                The scale of data to sample.
+                Defaults to 1.0.
+
+        Returns:
+            dict:  A dict mapping table name to the sampled data.
+        """
+        return synthesizer._internal_synthesizer.sample(scale)
diff --git a/tests/integration/synthesizers/test_clavaddpm.py b/tests/integration/synthesizers/test_clavaddpm.py
new file mode 100644
index 00000000..426acd20
--- /dev/null
+++ b/tests/integration/synthesizers/test_clavaddpm.py
@@ -0,0 +1,32 @@
+from sdgym import load_dataset
+from sdgym.synthesizers import ClavaDDPMSynthesizer
+
+
+def test_clavaddpm_end_to_end():
+    """Test it without metrics."""
+    # Setup
+    data, metadata_dict = load_dataset('multi_table', 'fake_hotels', limit_dataset_size=False)
+    clavaddpm_instance = ClavaDDPMSynthesizer()
+    ClavaDDPMSynthesizer._MODEL_KWARGS = {
+        'num_clusters': 5,
+        'clustering_method': 'both',
+        'd_layers': (128, 128),
+        'diffusion_iterations': 300,
+        'batch_size': 512,
+        'sampling_batch_size': 4096,
+        'num_timesteps': 100,
+        'verbose': True,
+    }
+
+    # Run
+    trained_synthesizer = clavaddpm_instance.get_trained_synthesizer(data, metadata_dict)
+    sampled_data = clavaddpm_instance.sample_from_synthesizer(trained_synthesizer, scale=1)
+
+    # Assert
+    for table_name, table in sampled_data.items():
+        assert table.shape[1] == data[table_name].shape[1], (
+            f'Sampled data shape {sampled_data.shape} does '
+            f'not match original data shape {data.shape}'
+        )
+
+        assert set(table.columns) == set(data[table_name].columns)
diff --git a/tests/unit/synthesizers/test_clavaddpm.py b/tests/unit/synthesizers/test_clavaddpm.py
new file mode 100644
index 00000000..dfd55b23
--- /dev/null
+++ b/tests/unit/synthesizers/test_clavaddpm.py
@@ -0,0 +1,98 @@
+"""Tests for the clavaddmp module."""
+
+from unittest.mock import MagicMock, patch
+
+import pandas as pd
+from sdv.metadata import Metadata
+
+from sdgym.synthesizers.clavaddpm import ClavaDDPM, ClavaDDPMSynthesizer
+
+
+class TestClavaDDPMSynthesizer:
+    @patch('sdgym.synthesizers.clavaddpm.ClavaDDPM')
+    def test__fit_mock(self, mock_clavaddpm):
+        """Test `_fit` builds and fits mock ClavaDDPM."""
+        # Setup
+        synthesizer = ClavaDDPMSynthesizer()
+        data = {'table1': pd.DataFrame({'col1': [1, 2, 3]})}
+        metadata = MagicMock()
+
+        # Run
+        synthesizer._fit(data, metadata)
+
+        # Assert
+        mock_clavaddpm.assert_called_once_with(metadata)
+        mock_clavaddpm.return_value.fit.assert_called_once_with(data)
+        assert synthesizer._internal_synthesizer is mock_clavaddpm.return_value
+
+    @patch('sdgym.synthesizers.clavaddpm.ClavaDDPM')
+    def test__fit_passes_model_kwargs(self, mock_clavaddpm):
+        """Test `_fit` forwards `_MODEL_KWARGS` to ClavaDDPM."""
+        # Setup
+        synthesizer = ClavaDDPMSynthesizer()
+        synthesizer._MODEL_KWARGS = {'num_clusters': 5, 'num_timesteps': 100}
+        data = {'table1': pd.DataFrame({'col1': [1, 2, 3]})}
+        metadata = MagicMock()
+
+        # Run
+        synthesizer._fit(data, metadata)
+
+        # Assert
+        mock_clavaddpm.assert_called_once_with(metadata, num_clusters=5, num_timesteps=100)
+
+    def test__fit(self):
+        """Test the `_fit` method."""
+        # Setup
+        synthesizer = ClavaDDPMSynthesizer()
+        synthesizer._MODEL_KWARGS = {'diffusion_iterations': 300, 'num_timesteps': 100}
+        data = {
+            'table1': pd.DataFrame({
+                'col1': [1, 2, 3],
+                'col2': ['A', 'B', 'C'],
+            }),
+            'table2': pd.DataFrame({
+                'col3': [10.0, 20.0, 30.0],
+                'col4': [True, False, True],
+            }),
+        }
+
+        metadata = Metadata.load_from_dict({
+            'tables': {
+                'table1': {
+                    'columns': {
+                        'col1': {'sdtype': 'numerical'},
+                        'col2': {'sdtype': 'categorical'},
+                    },
+                    'primary_key': 'col1',
+                },
+                'table2': {
+                    'columns': {
+                        'col3': {'sdtype': 'numerical'},
+                        'col4': {'sdtype': 'boolean'},
+                    },
+                    'primary_key': 'col3',
+                },
+            },
+            'relationships': [],
+        })
+
+        # Run
+        synthesizer._fit(data, metadata)
+
+        # Assert
+        assert isinstance(synthesizer._internal_synthesizer, ClavaDDPM)
+        assert isinstance(synthesizer, ClavaDDPMSynthesizer)
+
+    @patch('sdgym.synthesizers.clavaddpm.ClavaDDPM')
+    def test__get_trained_synthesizer(self, mock_clavaddpm):
+        """Test `_get_trained_synthesizer` with mock."""
+        # Setup
+        synthesizer = ClavaDDPMSynthesizer()
+        data = {'table1': pd.DataFrame({'col1': [1, 2, 3]})}
+        metadata = MagicMock()
+        # Run
+        trained_synthesizer = synthesizer._get_trained_synthesizer(data, metadata)
+        # Assert
+        assert isinstance(trained_synthesizer, ClavaDDPMSynthesizer)
+        assert trained_synthesizer._internal_synthesizer is mock_clavaddpm.return_value
+        mock_clavaddpm.return_value.fit.assert_called_once_with(data)

From a1853580808bbf22383a3de36eb63117ec426b51 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Wed, 22 Jul 2026 18:08:20 -0400
Subject: [PATCH 06/28] add benchmark_multi_table test

---
 tests/integration/synthesizers/test_utils.py |  6 ++++-
 tests/integration/test_benchmark.py          | 25 ++++++++++++++++++++
 tests/unit/synthesizers/test_utils.py        |  3 ++-
 3 files changed, 32 insertions(+), 2 deletions(-)

diff --git a/tests/integration/synthesizers/test_utils.py b/tests/integration/synthesizers/test_utils.py
index 2ca3f895..4a1e4d9d 100644
--- a/tests/integration/synthesizers/test_utils.py
+++ b/tests/integration/synthesizers/test_utils.py
@@ -29,7 +29,11 @@ def test_get_available_single_table_synthesizers():
 def test_get_available_multi_table_synthesizers():
     """Test the `get_available_multi_table_synthesizers` method"""
     # Setup
-    expected_synthesizers = ['HMASynthesizer', 'MultiTableUniformSynthesizer']
+    expected_synthesizers = [
+        'ClavaDDPMSynthesizer',
+        'HMASynthesizer',
+        'MultiTableUniformSynthesizer'
+    ]
 
     # Run
     synthesizers = get_available_multi_table_synthesizers()
diff --git a/tests/integration/test_benchmark.py b/tests/integration/test_benchmark.py
index cdf87685..39310411 100644
--- a/tests/integration/test_benchmark.py
+++ b/tests/integration/test_benchmark.py
@@ -929,6 +929,31 @@ def test_benchmark_multi_table_basic_synthesizers():
     ]
 
 
+def test_benchmark_multi_table_clavaddpm_no_metrics():
+    """Test it without metrics."""
+    # Run
+    custom_synthesizer = create_synthesizer_variant(
+        display_name='ClavaDDPMSynthesizer',
+        synthesizer_class='ClavaDDPMSynthesizer',
+        synthesizer_parameters={'diffusion_iterations': 300, 'num_timesteps': 100},
+    )
+    output = sdgym.benchmark_multi_table(
+        synthesizers=[],
+        custom_synthesizers=[custom_synthesizer],
+        sdv_datasets=['fake_hotels'],
+        compute_quality_score=False,
+        compute_diagnostic_score=False,
+    )
+
+    # Assert
+    train_time = output['Train_Time'][0]
+    sample_time = output['Sample_Time'][0]
+    assert isinstance(train_time, (int, float, complex)), 'Train_Time is not numerical'
+    assert isinstance(sample_time, (int, float, complex)), 'Sample_Time is not numerical'
+    assert train_time >= 0
+    assert sample_time >= 0
+
+
 @pytest.mark.skipif(
     not os.getenv('AWS_ACCESS_KEY_ID') or not os.getenv('AWS_SECRET_ACCESS_KEY'),
     reason='MovieLens benchmark requires AWS credentials for private dataset access.',
diff --git a/tests/unit/synthesizers/test_utils.py b/tests/unit/synthesizers/test_utils.py
index 0e7d5f5b..83558a9c 100644
--- a/tests/unit/synthesizers/test_utils.py
+++ b/tests/unit/synthesizers/test_utils.py
@@ -6,6 +6,7 @@ def test__get_supported_synthesizers():
     # Setup
     expected_synthesizers = [
         'CTGANSynthesizer',
+        'ClavaDDPMSynthesizer',
         'ColumnSynthesizer',
         'CopulaGANSynthesizer',
         'DataIdentity',
@@ -15,7 +16,7 @@ def test__get_supported_synthesizers():
         'RealTabFormerSynthesizer',
         'TVAESynthesizer',
         'TabDDPMSynthesizer',
-        'UniformSynthesizer',
+        'UniformSynthesizer'
     ]
 
     # Run

From f71f40a55a99bc5f77b7c6fc999ce7c8a16f5f2d Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Wed, 22 Jul 2026 18:12:04 -0400
Subject: [PATCH 07/28] fix lint

---
 tests/integration/synthesizers/test_utils.py | 2 +-
 tests/unit/synthesizers/test_utils.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/integration/synthesizers/test_utils.py b/tests/integration/synthesizers/test_utils.py
index 4a1e4d9d..05b21353 100644
--- a/tests/integration/synthesizers/test_utils.py
+++ b/tests/integration/synthesizers/test_utils.py
@@ -32,7 +32,7 @@ def test_get_available_multi_table_synthesizers():
     expected_synthesizers = [
         'ClavaDDPMSynthesizer',
         'HMASynthesizer',
-        'MultiTableUniformSynthesizer'
+        'MultiTableUniformSynthesizer',
     ]
 
     # Run
diff --git a/tests/unit/synthesizers/test_utils.py b/tests/unit/synthesizers/test_utils.py
index 83558a9c..f357aa64 100644
--- a/tests/unit/synthesizers/test_utils.py
+++ b/tests/unit/synthesizers/test_utils.py
@@ -16,7 +16,7 @@ def test__get_supported_synthesizers():
         'RealTabFormerSynthesizer',
         'TVAESynthesizer',
         'TabDDPMSynthesizer',
-        'UniformSynthesizer'
+        'UniformSynthesizer',
     ]
 
     # Run

From 21d009441ae2cfde8e9583942f0d16f8622a6aac Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Mon, 27 Jul 2026 10:48:00 -0400
Subject: [PATCH 08/28] fix sklearn onhotencoding version

---
 sdgym/synthesizers/clavaddpm.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index ea59717f..2fed139d 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -15,7 +15,9 @@
 
 import numpy as np
 import pandas as pd
+import sklearn
 import torch
+from packaging.version import Version
 from sklearn.cluster import KMeans
 from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
 from sklearn.neighbors import NearestNeighbors
@@ -296,7 +298,12 @@ def pair_clustering_keep_id(
 
         cat_one_hot = np.empty((cat_converted.shape[0], 0))
         for col in range(cat_converted.shape[1]):
-            encoder = OneHotEncoder(sparse_output=False)
+            sklearn_version = Version(sklearn.__version__).release
+            if sklearn_version[:2] >= (1, 2):
+                encoder = OneHotEncoder(sparse_output=False)
+            else:
+                encoder = OneHotEncoder(sparse=False)
+
             column = cat_converted[:, col].reshape(-1, 1)
             cat_one_hot = np.concatenate((cat_one_hot, encoder.fit_transform(column)), axis=1)
 

From 6ef5fa480f91345551aa60a1c08913ea2215a476 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Mon, 27 Jul 2026 11:40:47 -0400
Subject: [PATCH 09/28] change init_param based on sklearn version

---
 sdgym/synthesizers/clavaddpm.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 2fed139d..2de7b967 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -26,6 +26,8 @@
 from sdgym.synthesizers.base import MultiTableBaselineSynthesizer
 from sdgym.synthesizers.tabddpm import CAT_MISSING_VALUE, TabDDPM, ohe_to_categories
 
+SKLEARN_VERSION = Version(sklearn.__version__)
+
 
 def get_group_data_dict(np_data, group_id_attrs=[0]):
     """Grouping dictionary from pipeline_utils.py."""
@@ -298,8 +300,7 @@ def pair_clustering_keep_id(
 
         cat_one_hot = np.empty((cat_converted.shape[0], 0))
         for col in range(cat_converted.shape[1]):
-            sklearn_version = Version(sklearn.__version__).release
-            if sklearn_version[:2] >= (1, 2):
+            if SKLEARN_VERSION.release[:2] >= (1, 2):
                 encoder = OneHotEncoder(sparse_output=False)
             else:
                 encoder = OneHotEncoder(sparse=False)
@@ -330,15 +331,19 @@ def pair_clustering_keep_id(
     child_group_lengths = np.array([len(group) for group in child_group_data], dtype=int)
     num_clusters = min(num_clusters, len(cluster_data))
 
+    init_param = 'k-means++'
+    if SKLEARN_VERSION.release[:2] < (1, 1):
+        init_param = 'k-means'
+
     if clustering_method == 'kmeans':
-        kmeans = KMeans(n_clusters=num_clusters, n_init='auto', init='k-means++', random_state=seed)
+        kmeans = KMeans(n_clusters=num_clusters, n_init='auto', init=init_param, random_state=seed)
         kmeans.fit(cluster_data)
         cluster_labels = kmeans.labels_
     elif clustering_method == 'both':
         gmm = GaussianMixture(
             n_components=num_clusters,
             covariance_type='diag',
-            init_params='k-means++',
+            init_params=init_param,
             tol=0.0001,
             random_state=seed,
         )
@@ -348,7 +353,7 @@ def pair_clustering_keep_id(
         gmm = BayesianGaussianMixture(
             n_components=num_clusters,
             covariance_type='diag',
-            init_params='k-means++',
+            init_params=init_param,
             tol=0.0001,
             random_state=seed,
         )

From e3513c253bf4022502e55758f64f244a28e6cfe8 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Mon, 27 Jul 2026 12:15:43 -0400
Subject: [PATCH 10/28] kmeans

---
 sdgym/synthesizers/clavaddpm.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 2de7b967..2d7d4ab2 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -333,10 +333,10 @@ def pair_clustering_keep_id(
 
     init_param = 'k-means++'
     if SKLEARN_VERSION.release[:2] < (1, 1):
-        init_param = 'k-means'
+        init_param = 'kmeans'
 
     if clustering_method == 'kmeans':
-        kmeans = KMeans(n_clusters=num_clusters, n_init='auto', init=init_param, random_state=seed)
+        kmeans = KMeans(n_clusters=num_clusters, n_init='auto', init='k-means++', random_state=seed)
         kmeans.fit(cluster_data)
         cluster_labels = kmeans.labels_
     elif clustering_method == 'both':

From f31299309b604f07fe8ec1b9c286cc295eed4c78 Mon Sep 17 00:00:00 2001
From: Sarah Alnegheimish <40212131+sarahmish@users.noreply.github.com>
Date: Mon, 27 Jul 2026 13:24:06 -0400
Subject: [PATCH 11/28] add space

---
 tests/unit/synthesizers/test_clavaddpm.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/unit/synthesizers/test_clavaddpm.py b/tests/unit/synthesizers/test_clavaddpm.py
index dfd55b23..2eca09ff 100644
--- a/tests/unit/synthesizers/test_clavaddpm.py
+++ b/tests/unit/synthesizers/test_clavaddpm.py
@@ -90,8 +90,10 @@ def test__get_trained_synthesizer(self, mock_clavaddpm):
         synthesizer = ClavaDDPMSynthesizer()
         data = {'table1': pd.DataFrame({'col1': [1, 2, 3]})}
         metadata = MagicMock()
+
         # Run
         trained_synthesizer = synthesizer._get_trained_synthesizer(data, metadata)
+
         # Assert
         assert isinstance(trained_synthesizer, ClavaDDPMSynthesizer)
         assert trained_synthesizer._internal_synthesizer is mock_clavaddpm.return_value

From 8ce08a5f63434ef40d65bcb0475402ae736310a0 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Thu, 30 Jul 2026 09:35:23 +0300
Subject: [PATCH 12/28] create primary key if none exist

---
 sdgym/synthesizers/clavaddpm.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 2d7d4ab2..df40149d 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -722,6 +722,17 @@ def _parse_metadata(self, metadata):
             ))
             self._children[parent].append(child)
 
+        # Resolve a primary key for every table.
+        self._primary_key = {}
+        self._synthetic_pk = set()
+        for name, table_meta in meta['tables'].items():
+            pk = table_meta.get('primary_key')
+            if pk is None:
+                pk = f'__{name}_pk__'
+                self._synthetic_pk.add(name)
+
+            self._primary_key[name] = pk
+
         # Preprocessing
         self._discrete_cols = {}
         self._datetime_cols = {}  # name -> {col: datetime_format}
@@ -773,6 +784,9 @@ def fit(self, data):
         for name in self._table_names:
             output_cols = list(data[name].columns)
             df = data[name].copy().reset_index(drop=True)
+            # tables without a primary key get an internal row-id
+            if name in self._synthetic_pk:
+                df[self._primary_key[name]] = np.arange(len(df))
 
             date_info = {}
             for col, date_format in self._datetime_cols[name].items():

From bc076e2181740a68436975a884c2f413e3c579fe Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Fri, 31 Jul 2026 19:47:52 +0300
Subject: [PATCH 13/28] add id column to discrete

---
 sdgym/synthesizers/clavaddpm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index df40149d..a3a6ad6d 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -741,7 +741,7 @@ def _parse_metadata(self, metadata):
             discrete, datetimes = [], {}
             for column, spec in table_meta['columns'].items():
                 sdtype = spec.get('sdtype', 'categorical')
-                if column in id_cols or sdtype == 'id':
+                if column in id_cols:
                     continue
                 if sdtype == 'datetime':
                     datetimes[column] = spec.get('datetime_format')

From 288b95bdc49628e6dfb93c1831cf1ed8ad22df40 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Mon, 3 Aug 2026 15:55:34 +0300
Subject: [PATCH 14/28] address comments

---
 sdgym/synthesizers/clavaddpm.py | 31 +++++++++++++++++--------------
 sdgym/synthesizers/tabddpm.py   |  3 ++-
 2 files changed, 19 insertions(+), 15 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index a3a6ad6d..3bd01111 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -18,6 +18,8 @@
 import sklearn
 import torch
 from packaging.version import Version
+from pandas.core.tools.datetimes import _guess_datetime_format_for_array
+from sdv.metadata import Metadata
 from sklearn.cluster import KMeans
 from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
 from sklearn.neighbors import NearestNeighbors
@@ -789,13 +791,14 @@ def fit(self, data):
                 df[self._primary_key[name]] = np.arange(len(df))
 
             date_info = {}
-            for col, date_format in self._datetime_cols[name].items():
-                if date_format is None:
-                    df[col] = pd.to_datetime(df[col], errors='coerce').dt.strftime('%Y-%m-%d')
-                    date_format = '%Y-%m-%d'
-                days_since, earliest = calculate_days_since_earliest_date(df[col], date_format)
+            for col, datetime_format in self._datetime_cols[name].items():
+                if datetime_format is None:
+                    datetime_array = df[df.notna()].astype(str).to_numpy()
+                    datetime_format = _guess_datetime_format_for_array(datetime_array)
+
+                days_since, earliest = calculate_days_since_earliest_date(df[col], datetime_format)
                 df[col] = np.asarray(days_since, dtype=float)
-                date_info[col] = (earliest, date_format)
+                date_info[col] = (earliest, datetime_format)
 
             df, label_encoders = table_label_encode(df, self._discrete_cols[name])
             domain = get_domain(df, self._id_cols(name), self._discrete_cols[name])
@@ -813,9 +816,7 @@ def fit(self, data):
 
         self._clustering()
         self._training()
-
         self._fitted = True
-        return self
 
     def _clustering(self):
         """clava_clustering: cluster every parent->child relationship."""
@@ -826,7 +827,7 @@ def _clustering(self):
             if parent is None:
                 continue
             if self.verbose:
-                sys.stdout.write(f'Clustering {parent} -> {child}')
+                sys.stdout.write(f'Clustering {parent} -> {child}\n')
 
             fk = next(f for p, f, _ppk in self._tables[child]['parents'] if p == parent)
 
@@ -855,7 +856,8 @@ def _training(self):
         self._models = {}
         for parent, child in self._relation_order:
             if self.verbose:
-                sys.stdout.write(f'Training {parent} -> {child}')
+                sys.stdout.write(f'Training {parent} -> {child}\n')
+
             df_with_cluster = self._tables[child]['df']
             df_without_id = get_df_without_id(df_with_cluster, self._id_cols(child))
             self._models[(parent, child)] = self._child_training(
@@ -877,10 +879,10 @@ def _child_training(self, df_without_id, domain, parent, child):
                 continue
             sdtype = 'categorical' if domain[column]['type'] == 'discrete' else 'numerical'
             columns_meta[column] = {'sdtype': sdtype}
-        table_metadata = {
-            'METADATA_SPEC_VERSION': 'V1',
+
+        table_metadata = Metadata.load_from_dict({
             'tables': {child: {'columns': columns_meta}},
-        }
+        })
 
         model = TabDDPM(
             table_metadata,
@@ -920,7 +922,8 @@ def sample(self, scale=1.0):
 
         for parent, child in self._relation_order:
             if self.verbose:
-                sys.stdout.write(f'Generating {parent} -> {child}')
+                sys.stdout.write(f'Generating {parent} -> {child}\n')
+
             result = self._models[(parent, child)]
             df_with_cluster = self._tables[child]['df']
             df_without_id = get_df_without_id(df_with_cluster, self._id_cols(child))
diff --git a/sdgym/synthesizers/tabddpm.py b/sdgym/synthesizers/tabddpm.py
index 33d0e0c7..f074a0e9 100644
--- a/sdgym/synthesizers/tabddpm.py
+++ b/sdgym/synthesizers/tabddpm.py
@@ -888,6 +888,7 @@ def __init__(self, diffusion, train_iter, lr, weight_decay, steps, device, verbo
         self.log_every = 100
         self.print_every = 500
         self.ema_every = 1000
+        self.verbose = verbose
 
     def _anneal_lr(self, step):
         frac_done = step / self.steps
@@ -927,7 +928,7 @@ def run_loop(self):
             if (step + 1) % self.log_every == 0:
                 mloss = np.around(curr_loss_multi / curr_count, 4)
                 gloss = np.around(curr_loss_gauss / curr_count, 4)
-                if (step + 1) % self.print_every == 0:
+                if self.verbose and (step + 1) % self.print_every == 0:
                     sys.stdout.write(
                         f'Step {(step + 1)}/{self.steps} '
                         f'MLoss: {mloss} GLoss: {gloss} Sum: {mloss + gloss}\n'

From 0f05bf78c36351da60e852c5eb0bb4e2bcc109ac Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Wed, 5 Aug 2026 01:10:30 +0300
Subject: [PATCH 15/28] change to in-house datetime guesser

---
 sdgym/synthesizers/clavaddpm.py           |  37 ++++-
 tests/unit/synthesizers/test_clavaddpm.py | 179 +++++++++++++++++++++-
 2 files changed, 211 insertions(+), 5 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 3bd01111..d74c7a31 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -10,7 +10,7 @@
 import logging
 import random
 import sys
-from collections import defaultdict
+from collections import Counter, defaultdict
 from datetime import datetime, timedelta
 
 import numpy as np
@@ -18,7 +18,7 @@
 import sklearn
 import torch
 from packaging.version import Version
-from pandas.core.tools.datetimes import _guess_datetime_format_for_array
+from pandas.tseries.api import guess_datetime_format
 from sdv.metadata import Metadata
 from sklearn.cluster import KMeans
 from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
@@ -31,6 +31,35 @@
 SKLEARN_VERSION = Version(sklearn.__version__)
 
 
+def guess_array_datetime_format(values, sample_size=100, dayfirst=False):
+    """Guess the most likely datetime format via majority vote over a sample.
+
+    Args:
+        values (list, numpy.array, pd.Series):
+            List of datetime values to inspect.
+        sample_size (int, optional):
+            Number of samples to use for guessing. Default 100.
+        dayfirst (bool, optional):
+            If True parses dates with the day first. Default False.
+
+    Returns:
+        str: Datetime format string, or None if it can not be guessed.
+    """
+    series = pd.Series(values)
+    series = series.dropna()
+    if series.empty:
+        return None
+
+    sample_size = min(sample_size, len(series))
+    sample = series.sample(sample_size, random_state=0)
+    sample = sample.astype(str).str.strip()
+    sample = sample[sample != '']
+
+    guesses = sample.apply(guess_datetime_format, dayfirst=dayfirst).dropna()
+    counts = Counter(guesses)
+    return top[0][0] if (top := counts.most_common(1)) else None
+
+
 def get_group_data_dict(np_data, group_id_attrs=[0]):
     """Grouping dictionary from pipeline_utils.py."""
     group_data_dict = {}
@@ -794,7 +823,7 @@ def fit(self, data):
             for col, datetime_format in self._datetime_cols[name].items():
                 if datetime_format is None:
                     datetime_array = df[df.notna()].astype(str).to_numpy()
-                    datetime_format = _guess_datetime_format_for_array(datetime_array)
+                    datetime_format = guess_array_datetime_format(datetime_array)
 
                 days_since, earliest = calculate_days_since_earliest_date(df[col], datetime_format)
                 df[col] = np.asarray(days_since, dtype=float)
@@ -1046,6 +1075,6 @@ def _sample_from_synthesizer(self, synthesizer, scale):
                 Defaults to 1.0.
 
         Returns:
-            dict:  A dict mapping table name to the sampled data.
+            dict: A dict mapping table name to the sampled data.
         """
         return synthesizer._internal_synthesizer.sample(scale)
diff --git a/tests/unit/synthesizers/test_clavaddpm.py b/tests/unit/synthesizers/test_clavaddpm.py
index 2eca09ff..decbf40d 100644
--- a/tests/unit/synthesizers/test_clavaddpm.py
+++ b/tests/unit/synthesizers/test_clavaddpm.py
@@ -3,9 +3,186 @@
 from unittest.mock import MagicMock, patch
 
 import pandas as pd
+import pytest
 from sdv.metadata import Metadata
 
-from sdgym.synthesizers.clavaddpm import ClavaDDPM, ClavaDDPMSynthesizer
+from sdgym.synthesizers.clavaddpm import (
+    ClavaDDPM,
+    ClavaDDPMSynthesizer,
+    guess_array_datetime_format,
+)
+
+
+def test_guess_array_datetime_format_empty_input_returns_none():
+    """Test guess array with an empty input returns None."""
+    # Run
+    result = guess_array_datetime_format([])
+
+    # Assert
+    assert result is None
+
+
+def test_guess_array_datetime_format_all_null_returns_none():
+    """Test guess array with all null inputs returns None."""
+    # Run
+    result = guess_array_datetime_format([None, None, float('nan')])
+
+    # Assert
+    assert result is None
+
+
+def test_guess_array_datetime_format_with_datetime_column():
+    """Test guess array works as expected with datetime values."""
+    # Setup
+    values = pd.to_datetime(['2020-01-01', '2020-01-02'])
+
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result == '%Y-%m-%d'
+
+
+def test_guess_array_datetime_format_non_date_numeric_values_return_none():
+    """Test guess array doesn't parse small integers as datetime format."""
+    # Run
+    result = guess_array_datetime_format([1, 2, 3, 4, 5])
+
+    # Assert
+    assert result is None
+
+
+def test_guess_array_datetime_format_consistent_format_detected():
+    """Test that guess array datetime works as expected."""
+    # Setup
+    values = ['2020-01-05', '2020-02-14', '2020-03-30']
+
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result == '%Y-%m-%d'
+
+
+def test_majority_format_wins_over_minority():
+    """Test that guess array datetime works finds the majority format."""
+    # Setup
+    values = ['2020-01-05'] * 7 + ['01/05/2020'] * 3
+
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result == '%Y-%m-%d'
+
+
+def test_guess_array_datetime_format_with_mixed_values():
+    """Test guess array returns the majority of valid dates."""
+    # Setup
+    values = ['2020-01-05', '2020-02-14', '2020-03-30'] + [''] * 5 + ['n/a'] * 5
+
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result == '%Y-%m-%d'
+
+
+def test_guess_array_datetime_format_with_whitespace_only_strings_are_ignored():
+    """Test guess array returns ignores whitespace."""
+    # Setup
+    values = ['2020-01-05', '2020-02-14', '', '   ', '2020-03-30', '\t']
+
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result == '%Y-%m-%d'
+
+
+def test_guess_array_datetime_format_with_all_unparseable_returns_none():
+    """Test guess array returns None when all values are not dates."""
+    # Setup
+    values = ['n/a', 'not a date', 'xx', '']
+
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result is None
+
+
+@pytest.mark.parametrize(
+    'values, expected',
+    [
+        (['2020-01-05', '2020-02-14', '2020-03-30'] + [None] * 10, '%Y-%m-%d'),
+        (['2020/01/05', '2020/02/14', '2020/03/30', '2020-02-14', '2020-03-30'], '%Y/%m/%d'),
+        (['2020-01-05 00:10:05', '2020-02-14 12:12:05', '2020-03-30'], '%Y-%m-%d %H:%M:%S'),
+    ],
+)
+def test_guess_array_datetime_format_with_expected_format(values, expected):
+    """Test guess array returns expected formats."""
+    # Run
+    result = guess_array_datetime_format(values)
+
+    # Assert
+    assert result == expected
+
+
+@pytest.mark.parametrize(
+    'dayfirst, expected',
+    [
+        (False, '%m-%d-%Y'),
+        (True, '%d-%m-%Y'),
+    ],
+)
+def test_guess_array_datetime_format_dayfirst_changes_ambiguous_guess(dayfirst, expected):
+    """Test guess array with dayfirst parameter clears up ambiguity."""
+    # Setup
+    values = ['01-02-2020'] * 5
+
+    # Run
+    result = guess_array_datetime_format(values, dayfirst=dayfirst)
+
+    # Assert
+    assert result == expected
+
+
+def test_sample_size_larger_than_population_does_not_error():
+    """Test guess array with larger sample size works."""
+    # Setup
+    values = ['2020-01-01', '2020-01-02']
+
+    # Run
+    result = guess_array_datetime_format(values, sample_size=1000)
+
+    # Assert
+    assert result == '%Y-%m-%d'
+
+
+def test_sample_size_zero_returns_none():
+    """Test guess array with zero sample size returns None."""
+    # Setup
+    values = ['2020-01-01', '2020-01-02']
+
+    # Run
+    result = guess_array_datetime_format(values, sample_size=0)
+
+    # Assert
+    assert result is None
+
+
+def test_reproducible_with_fixed_random_state():
+    """Test guess array retuns the same result."""
+    # Setup
+    values = ['2020-01-05'] * 150 + ['01/05/2020'] * 150
+
+    # Run
+    result_a = guess_array_datetime_format(values, sample_size=10)
+    result_b = guess_array_datetime_format(values, sample_size=10)
+
+    # Assert
+    assert result_a == result_b
 
 
 class TestClavaDDPMSynthesizer:

From 72ecee05f6de65a82ad7dd07b84bf899f38a01d8 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Wed, 5 Aug 2026 11:35:07 +0300
Subject: [PATCH 16/28] import guess_datetime_format depending on version

---
 sdgym/synthesizers/clavaddpm.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index d74c7a31..b97916bd 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -18,7 +18,6 @@
 import sklearn
 import torch
 from packaging.version import Version
-from pandas.tseries.api import guess_datetime_format
 from sdv.metadata import Metadata
 from sklearn.cluster import KMeans
 from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
@@ -45,6 +44,12 @@ def guess_array_datetime_format(values, sample_size=100, dayfirst=False):
     Returns:
         str: Datetime format string, or None if it can not be guessed.
     """
+    pandas_version = Version(pd.__version__).release[:2]
+    if pandas_version >= (2, 2):
+        from pandas.tseries.api import guess_datetime_format  # pandas >= 2.2.0
+    else:
+        from pandas._libs.tslibs.parsing import guess_datetime_format  # pandas < 2.2.0
+
     series = pd.Series(values)
     series = series.dropna()
     if series.empty:

From 29fee25352b0b281805d0cc5c15bf2b796b7e8fa Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Thu, 6 Aug 2026 11:32:15 +0300
Subject: [PATCH 17/28] cap categorical columns to 100

---
 sdgym/synthesizers/clavaddpm.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index b97916bd..fdd4f939 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -678,6 +678,10 @@ class ClavaDDPM:
             Multi-parent matching controls for unique matching.
         no_matching (bool):
             Multi-parent matching controls for no matching.
+        max_categories (int):
+            Maximum number of distinct values kept for a categorical column.
+            Columns with more are reduced to this many values, randomly sampled
+            from the data.
         device (str or None):
             Whether to use ``'cuda'`` or ``'cpu'``. If None, auto-select is used.
         seed (int):
@@ -707,6 +711,7 @@ def __init__(
         matching_batch_size=1000,
         unique_matching=True,
         no_matching=False,
+        max_categories=100,
         device=None,
         seed=0,
         verbose=False,
@@ -729,6 +734,7 @@ def __init__(
         self.matching_batch_size = matching_batch_size
         self.unique_matching = unique_matching
         self.no_matching = no_matching
+        self.max_categories = max_categories
         self.device = device
         self.seed = seed
         self.verbose = verbose
@@ -824,11 +830,22 @@ def fit(self, data):
             if name in self._synthetic_pk:
                 df[self._primary_key[name]] = np.arange(len(df))
 
+            # cap the cardinality of the discrete columns.
+            for col in self._discrete_cols[name]:
+                uniques = df[col].dropna().unique()
+                if len(uniques) > self.max_categories:
+                    kept = np.random.choice(uniques, size=self.max_categories, replace=False)
+                    outside = ~df[col].isin(kept)
+                    df.loc[outside, col] = np.random.choice(kept, size=int(outside.sum()))
+
             date_info = {}
             for col, datetime_format in self._datetime_cols[name].items():
                 if datetime_format is None:
                     datetime_array = df[df.notna()].astype(str).to_numpy()
                     datetime_format = guess_array_datetime_format(datetime_array)
+                    df[col] = pd.to_datetime(
+                        df[col], format=datetime_format, errors='coerce'
+                    ).dt.strftime()
 
                 days_since, earliest = calculate_days_since_earliest_date(df[col], datetime_format)
                 df[col] = np.asarray(days_since, dtype=float)

From b201719f94e9f6c4e5c5e6e599fdbc0a01c3bac0 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Thu, 6 Aug 2026 11:53:25 +0300
Subject: [PATCH 18/28] cast mixed formats to be uniform

---
 sdgym/synthesizers/clavaddpm.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index fdd4f939..1b39a6ac 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -843,9 +843,7 @@ def fit(self, data):
                 if datetime_format is None:
                     datetime_array = df[df.notna()].astype(str).to_numpy()
                     datetime_format = guess_array_datetime_format(datetime_array)
-                    df[col] = pd.to_datetime(
-                        df[col], format=datetime_format, errors='coerce'
-                    ).dt.strftime()
+                    df[col] = pd.to_datetime(df[col], format='mixed').dt.strftime(datetime_format)
 
                 days_since, earliest = calculate_days_since_earliest_date(df[col], datetime_format)
                 df[col] = np.asarray(days_since, dtype=float)

From 80482657830f3ce39ad78fb90dadcbcb3cfc92f5 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Mon, 10 Aug 2026 20:08:41 +0300
Subject: [PATCH 19/28] move preprocessing to fit only

---
 sdgym/synthesizers/clavaddpm.py | 243 ++++++++++++++++++++++----------
 1 file changed, 172 insertions(+), 71 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 1b39a6ac..06c7f6db 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -25,7 +25,12 @@
 from sklearn.preprocessing import LabelEncoder, MinMaxScaler, OneHotEncoder
 
 from sdgym.synthesizers.base import MultiTableBaselineSynthesizer
-from sdgym.synthesizers.tabddpm import CAT_MISSING_VALUE, TabDDPM, ohe_to_categories
+from sdgym.synthesizers.tabddpm import (
+    CAT_MISSING_VALUE,
+    TabDDPM,
+    _DataTransformer,
+    ohe_to_categories,
+)
 
 SKLEARN_VERSION = Version(sklearn.__version__)
 
@@ -201,6 +206,32 @@ def get_domain(df, id_cols, discrete_cols):
     return domain
 
 
+def decode_id_values(values, encoder, column):
+    """Inverse-transform label-encoded id codes back to the original id values.
+
+    Sampling can create more rows than the original table, so codes past the
+    encoder's fitted range get fresh values: numeric ids continue after the
+    largest original id, other ids become new ``'{column}_{code}'`` strings.
+    The extension is deterministic per code, so a foreign key decoded with its
+    parent's encoder still matches the parent's decoded primary key.
+    """
+    codes = np.asarray(values).astype('int64')
+    classes = encoder.classes_
+    decoded = np.empty(len(codes), dtype=object)
+    in_range = (codes >= 0) & (codes < len(classes))
+    decoded[in_range] = encoder.inverse_transform(codes[in_range])
+    if (~in_range).any():
+        if pd.api.types.is_numeric_dtype(classes):
+            # continue past the largest original id (code n_classes -> max + 1, ...)
+            decoded[~in_range] = codes[~in_range] + (classes.max() + 1 - len(classes))
+        else:
+            decoded[~in_range] = [f'{column}_{code}' for code in codes[~in_range]]
+
+    if pd.api.types.is_numeric_dtype(classes):
+        return decoded.astype(classes.dtype)
+    return decoded
+
+
 def topological_sort(graph):
     """Order tables into ``[parent, child]`` relations from preprocess_utils.py."""
     in_degree = {node: 0 for node in graph}
@@ -271,7 +302,6 @@ def pair_clustering_keep_id(
             else:
                 parent_num_cols.append((col_index, col))
 
-    child_primary_key_index = original_child_cols.index(child_primary_key)
     parent_primary_key_index = original_parent_cols.index(parent_primary_key)
     foreing_key_index = original_child_cols.index(foreign_key)
 
@@ -309,21 +339,9 @@ def pair_clustering_keep_id(
     joint_num_matrix = np.concatenate([sorted_child_num_data, sorted_parent_num_data], axis=1)
     joint_cat_matrix = np.concatenate([sorted_child_cat_data, sorted_parent_cat_data], axis=1)
 
-    # Impute missing numerical values with their column mean before clustering.
-    # The reference assumed pre-imputed data; this mirrors the mean-fill that
-    # TabDDPM's own transformer applies, so clustering (which cannot take NaNs)
-    # is robust to real-world tables with missing values.
-    if joint_num_matrix.shape[1] > 0:
-        joint_num_matrix = joint_num_matrix.astype(float)
-        col_means = np.nanmean(joint_num_matrix, axis=0)
-        col_means = np.where(np.isnan(col_means), 0.0, col_means)
-        nan_positions = np.where(np.isnan(joint_num_matrix))
-        joint_num_matrix[nan_positions] = np.take(col_means, nan_positions[1])
-
-    joint_num_matrix_p_index = sorted_child_num_data.shape[1]
-    cat_one_hot = None
     if joint_cat_matrix.shape[1] > 0:
         joint_cat_matrix_p_index = sorted_child_cat_data.shape[1]
+        joint_num_matrix_p_index = sorted_child_num_data.shape[1]
 
         cat_converted = []
         for i in range(joint_cat_matrix.shape[1]):
@@ -334,7 +352,10 @@ def pair_clustering_keep_id(
             cat_converted.append(label_encoder.fit_transform(joint_cat_matrix[:, i]).astype(float))
         cat_converted = np.vstack(cat_converted).T
 
+        # Initialize an empty array to store the encoded values
         cat_one_hot = np.empty((cat_converted.shape[0], 0))
+
+        # Loop through each column in the data and encode it
         for col in range(cat_converted.shape[1]):
             if SKLEARN_VERSION.release[:2] >= (1, 2):
                 encoder = OneHotEncoder(sparse_output=False)
@@ -342,15 +363,16 @@ def pair_clustering_keep_id(
                 encoder = OneHotEncoder(sparse=False)
 
             column = cat_converted[:, col].reshape(-1, 1)
-            cat_one_hot = np.concatenate((cat_one_hot, encoder.fit_transform(column)), axis=1)
+            encoded_column = encoder.fit_transform(column)
+            cat_one_hot = np.concatenate((cat_one_hot, encoded_column), axis=1)
 
         cat_one_hot[:, joint_cat_matrix_p_index:] = (
             parent_scale * cat_one_hot[:, joint_cat_matrix_p_index:]
         )
 
+    # Perform quantile normalization using QuantileTransformer
     num_min_max = min_max_normalize_sklearn(joint_num_matrix)
 
-    # key channel: parent identity, factorized so arbitrary key types normalize
     key_factorized = pd.factorize(sorted_parent_data_repeated[:, parent_primary_key_index])[0]
     key_min_max = min_max_normalize_sklearn(key_factorized.astype(float).reshape(-1, 1))
     key_scaled = key_scale * key_min_max
@@ -410,6 +432,7 @@ def pair_clustering_keep_id(
         curr_index = 0
         agree_rates = []
         for group_length in child_group_lengths:
+            # First, determine the most common label in the current group
             most_common_label_count = np.max(
                 np.bincount(cluster_labels[curr_index : curr_index + group_length])
             )
@@ -417,13 +440,17 @@ def pair_clustering_keep_id(
                 np.bincount(cluster_labels[curr_index : curr_index + group_length])
             )
             group_cluster_labels.append(group_cluster_label)
-            agree_rates.append(most_common_label_count / group_length)
+
+            # Compute agree rate using the most common label count
+            agree_rate = most_common_label_count / group_length
+            agree_rates.append(agree_rate)
+
+            # Then, update the curr_index for the next iteration
             curr_index += group_length
 
     group_assignment = np.repeat(group_cluster_labels, child_group_lengths, axis=0).reshape((-1, 1))
     sorted_child_data_with_cluster = np.concatenate([sorted_child_data, group_assignment], axis=1)
 
-    # per-cluster distribution of "how many children a parent has"
     group_labels_list = group_cluster_labels
     group_lengths_list = child_group_lengths.tolist()
     group_lengths_dict = {}
@@ -433,38 +460,49 @@ def pair_clustering_keep_id(
             group_lengths_dict[group_label] = defaultdict(int)
         group_lengths_dict[group_label][group_lengths_list[i]] += 1
 
-    group_lengths_prob_dicts = {
-        group_label: freq_to_prob(freq_dict)
-        for group_label, freq_dict in group_lengths_dict.items()
-    }
-
-    # attach cluster label to the child in its original order
-    sorted_child_ids = sorted_child_data[:, child_primary_key_index]
-    child_id_to_cluster = dict(zip(sorted_child_ids, group_assignment.flatten()))
-    child_df_with_cluster = child_df.copy()
-    child_df_with_cluster[relation_cluster_name] = (
-        child_df[child_primary_key].map(child_id_to_cluster).astype(int)
+    group_lengths_prob_dicts = {}
+    for group_label, freq_dict in group_lengths_dict.items():
+        group_lengths_prob_dicts[group_label] = freq_to_prob(freq_dict)
+
+    # recover the preprocessed data back to dataframe
+    child_df_with_cluster = pd.DataFrame(
+        sorted_child_data_with_cluster, columns=original_child_cols + [relation_cluster_name]
+    )
+
+    # recover child df order
+    child_df_with_cluster = pd.merge(  # noqa: PD015
+        child_df[[child_primary_key]],
+        child_df_with_cluster,
+        on=child_primary_key,
+        how='left',
     )
 
-    # a parent inherits its children's cluster; childless parents get a fresh id
     parent_id_to_cluster = {}
     for i in range(len(sorted_child_data)):
         parent_id = sorted_child_data[i, foreing_key_index]
         if parent_id in parent_id_to_cluster:
+            assert parent_id_to_cluster[parent_id] == sorted_child_data_with_cluster[i, -1]
             continue
         parent_id_to_cluster[parent_id] = sorted_child_data_with_cluster[i, -1]
 
     max_cluster_label = max(parent_id_to_cluster.values())
-    parent_df_with_cluster = parent_df.copy()
-    parent_clusters = parent_df[parent_primary_key].map(parent_id_to_cluster)
-    parent_df_with_cluster[relation_cluster_name] = parent_clusters.fillna(
-        max_cluster_label + 1
-    ).astype(int)
-
-    new_col_entry = {
-        'type': 'discrete',
-        'size': len(set(parent_df_with_cluster[relation_cluster_name])),
-    }
+
+    parent_data_clusters = []
+    for i in range(len(parent_data)):
+        if parent_data[i, parent_primary_key_index] in parent_id_to_cluster:
+            parent_data_clusters.append(
+                parent_id_to_cluster[parent_data[i, parent_primary_key_index]]
+            )
+        else:
+            parent_data_clusters.append(max_cluster_label + 1)
+
+    parent_data_clusters = np.array(parent_data_clusters).reshape(-1, 1)
+    parent_data_with_cluster = np.concatenate([parent_data, parent_data_clusters], axis=1)
+    parent_df_with_cluster = pd.DataFrame(
+        parent_data_with_cluster, columns=original_parent_cols + [relation_cluster_name]
+    )
+
+    new_col_entry = {'type': 'discrete', 'size': len(set(parent_data_clusters.flatten()))}
     parent_domain_dict[relation_cluster_name] = new_col_entry.copy()
     child_domain_dict[relation_cluster_name] = new_col_entry.copy()
 
@@ -740,6 +778,7 @@ def __init__(
         self.verbose = verbose
 
         self._parse_metadata(metadata)
+        self._metadata = metadata
         self._fitted = False
 
     def _parse_metadata(self, metadata):
@@ -751,9 +790,6 @@ def _parse_metadata(self, metadata):
         self._parents = {name: [] for name in self._table_names}
         self._children = {name: [] for name in self._table_names}
 
-        for name, table_meta in meta['tables'].items():
-            self._primary_key[name] = table_meta.get('primary_key')
-
         for relationship in meta.get('relationships', []):
             parent = relationship['parent_table_name']
             child = relationship['child_table_name']
@@ -769,7 +805,7 @@ def _parse_metadata(self, metadata):
         self._synthetic_pk = set()
         for name, table_meta in meta['tables'].items():
             pk = table_meta.get('primary_key')
-            if pk is None:
+            if pk is None or isinstance(pk, list):  # override composite key
                 pk = f'__{name}_pk__'
                 self._synthetic_pk.add(name)
 
@@ -778,19 +814,23 @@ def _parse_metadata(self, metadata):
         # Preprocessing
         self._discrete_cols = {}
         self._datetime_cols = {}  # name -> {col: datetime_format}
+        self._id_value_cols = {}
         for name, table_meta in meta['tables'].items():
             id_cols = self._id_cols(name)
-            discrete, datetimes = [], {}
+            discrete, datetimes, id_values = [], {}, []
             for column, spec in table_meta['columns'].items():
                 sdtype = spec.get('sdtype', 'categorical')
                 if column in id_cols:
                     continue
-                if sdtype == 'datetime':
+                if sdtype == 'id':
+                    id_values.append(column)
+                elif sdtype == 'datetime':
                     datetimes[column] = spec.get('datetime_format')
                 elif sdtype != 'numerical':
                     discrete.append(column)
             self._discrete_cols[name] = discrete
             self._datetime_cols[name] = datetimes
+            self._id_value_cols[name] = id_values
 
         graph = {name: {'children': self._children[name]} for name in self._table_names}
         self._relation_order = [tuple(edge) for edge in topological_sort(graph)]
@@ -803,6 +843,16 @@ def _num_clusters_for(self, child):
     def _id_cols(self, table):
         return [self._primary_key[table]] + [fk for _p, fk, _ppk in self._parents[table]]
 
+    def _get_values(self, data, table_name, key):
+        values = [data[table_name][key].astype(str).dropna().to_numpy().flatten()]
+        for child in self._table_names:
+            for parent, fk, _ppk in self._parents[child]:
+                if parent == table_name and fk in data[child].columns:
+                    keys = data[child][fk].astype(str).dropna().to_numpy().flatten()
+                    values.append(keys)
+
+        return np.concatenate(values)
+
     def fit(self, data):
         """Fit this model to the original data.
 
@@ -810,47 +860,82 @@ def fit(self, data):
             data (dict):
                 Dictionary mapping each table name to a ``pandas.DataFrame``.
         """
-        missing = set(self._table_names) - set(data)
-        if missing:
-            raise ValueError(f'Missing tables in data: {sorted(missing)}')
+        self._metadata.validate_data(data)
 
         random.seed(self.seed)
         np.random.seed(self.seed)
         torch.manual_seed(self.seed)
 
-        # Preprocess every table up front: turn datetimes into numeric
-        # day-counts, label-encode the discrete columns, and derive the
-        # ``domain``. Encoders / date anchors are kept so ``sample`` can invert
-        # them.
+        self._id_encoded_cols = defaultdict(dict)
+        for name in self._table_names:
+            pk = self._primary_key[name]
+            if not (
+                pk is None
+                or name in self._synthetic_pk
+                or pd.api.types.is_numeric_dtype(data[name][pk])
+            ):
+                key_values = self._get_values(data, name, pk)
+                encoder = LabelEncoder()
+                encoder.fit(key_values)
+                self._id_encoded_cols[name][pk] = encoder
+
+            for id_col in self._id_value_cols[name]:
+                if pd.api.types.is_numeric_dtype(data[name][id_col]):
+                    continue
+
+                key_values = self._get_values(data, name, id_col)
+                encoder = LabelEncoder()
+                encoder.fit(key_values)
+                self._id_encoded_cols[name][id_col] = encoder
+
+        for name in self._table_names:
+            for parent, fk, _ppk in self._parents[name]:
+                if parent in self._id_encoded_cols and _ppk in self._id_encoded_cols[parent]:
+                    self._id_encoded_cols[name][fk] = self._id_encoded_cols[parent][_ppk]
+
+        # preprocessing
         self._tables = {}
         for name in self._table_names:
             output_cols = list(data[name].columns)
             df = data[name].copy().reset_index(drop=True)
             # tables without a primary key get an internal row-id
             if name in self._synthetic_pk:
-                df[self._primary_key[name]] = np.arange(len(df))
+                df[self._primary_key[name]] = np.arange(len(df)) + 1
 
-            # cap the cardinality of the discrete columns.
-            for col in self._discrete_cols[name]:
-                uniques = df[col].dropna().unique()
-                if len(uniques) > self.max_categories:
-                    kept = np.random.choice(uniques, size=self.max_categories, replace=False)
-                    outside = ~df[col].isin(kept)
-                    df.loc[outside, col] = np.random.choice(kept, size=int(outside.sum()))
+            # label-encode id columns
+            for col, encoder in self._id_encoded_cols[name].items():
+                notna = df[col].notna()
+                df.loc[notna, col] = encoder.transform(df.loc[notna, col].astype(str))
+                df[col] = pd.to_numeric(df[col])
 
+            # preprocess_utils: encode datetimes as days since the earliest date
             date_info = {}
             for col, datetime_format in self._datetime_cols[name].items():
                 if datetime_format is None:
-                    datetime_array = df[df.notna()].astype(str).to_numpy()
-                    datetime_format = guess_array_datetime_format(datetime_array)
+                    datetime_format = guess_array_datetime_format(df[col])
                     df[col] = pd.to_datetime(df[col], format='mixed').dt.strftime(datetime_format)
 
                 days_since, earliest = calculate_days_since_earliest_date(df[col], datetime_format)
                 df[col] = np.asarray(days_since, dtype=float)
                 date_info[col] = (earliest, datetime_format)
 
+            # preprocess_utils: label-encode the discrete columns, derive the domain
             df, label_encoders = table_label_encode(df, self._discrete_cols[name])
-            domain = get_domain(df, self._id_cols(name), self._discrete_cols[name])
+            domain = get_domain(
+                df, self._id_cols(name) + self._id_value_cols[name], self._discrete_cols[name]
+            )
+
+            # impute missing values
+            continuous_cols = [col for col, spec in domain.items() if spec['type'] == 'continuous']
+            if continuous_cols:
+                imputer = _DataTransformer(
+                    {col: {'sdtype': 'numerical'} for col in continuous_cols},
+                    normalization=None,
+                    seed=self.seed,
+                )
+                imputer.fit(df[continuous_cols])
+                x_num, _ = imputer.transform(df[continuous_cols])
+                df[continuous_cols] = x_num.astype(float)
 
             self._tables[name] = {
                 'df': df,
@@ -861,6 +946,7 @@ def fit(self, data):
                 'parents': list(self._parents[name]),
                 'label_encoders': label_encoders,
                 'date_info': date_info,
+                'cluster_cols': [],
             }
 
         self._clustering()
@@ -896,6 +982,15 @@ def _clustering(self):
                 clustering_method=self.clustering_method,
                 seed=self.seed,
             )
+
+            cluster_col = f'{parent}_{child}_cluster'
+            self._tables[parent]['cluster_cols'].append(cluster_col)
+            self._tables[child]['cluster_cols'].append(cluster_col)
+            for col in self._tables[parent]['cluster_cols']:
+                parent_df[col] = parent_df[col].astype(int)
+            for col in self._tables[child]['cluster_cols']:
+                child_df[col] = child_df[col].astype(int)
+
             self._tables[parent]['df'] = parent_df
             self._tables[child]['df'] = child_df
             self._all_group_lengths_prob_dicts[(parent, child)] = group_lengths_prob_dicts
@@ -924,10 +1019,11 @@ def _child_training(self, df_without_id, domain, parent, child):
 
         columns_meta = {}
         for column in df_without_id.columns:
-            if column not in domain:
-                continue
-            sdtype = 'categorical' if domain[column]['type'] == 'discrete' else 'numerical'
-            columns_meta[column] = {'sdtype': sdtype}
+            if column in domain:
+                sdtype = 'categorical' if domain[column]['type'] == 'discrete' else 'numerical'
+                columns_meta[column] = {'sdtype': sdtype}
+            elif column in self._id_value_cols[child]:
+                columns_meta[column] = {'sdtype': 'id'}
 
         table_metadata = Metadata.load_from_dict({
             'tables': {child: {'columns': columns_meta}},
@@ -1032,7 +1128,7 @@ def sample(self, scale=1.0):
                     child,
                     self._tables[child]['parents'],
                     synthetic_tables,
-                    self._id_cols(child),
+                    self._id_cols(child) + self._id_value_cols[child],
                     n_clusters=self.num_matching_clusters,
                     unique_matching=self.unique_matching,
                     batch_size=self.matching_batch_size,
@@ -1044,7 +1140,7 @@ def sample(self, scale=1.0):
         return {child: self._postprocess(child, table) for child, table in final_tables.items()}
 
     def _postprocess(self, table, df):
-        """Invert fit-time preprocessing: decode discretes, rebuild dates, order cols."""
+        """Invert fit-time preprocessing: decode discretes, rebuild dates, decode ids."""
         df = df.copy()
         state = self._tables[table]
 
@@ -1059,6 +1155,11 @@ def _postprocess(self, table, df):
             if col in df.columns:
                 df[col] = reconstruct_dates(df[col].to_numpy(), earliest, date_format)
 
+        # decode the primary/foreign key columns back to the original id values
+        for col, encoder in self._id_encoded_cols[table].items():
+            if col in df.columns:
+                df[col] = decode_id_values(df[col].to_numpy(), encoder, col)
+
         return df[state['output_cols']].reset_index(drop=True)
 
 

From d1340b26d66afb8b225ef6393c38bd3242ab4a6c Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Thu, 13 Aug 2026 14:07:06 +0300
Subject: [PATCH 20/28] skip validation on tabddpm if run from clavaddpm

---
 sdgym/synthesizers/clavaddpm.py | 1 +
 sdgym/synthesizers/tabddpm.py   | 4 +++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 06c7f6db..5590aedb 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -1046,6 +1046,7 @@ def _child_training(self, df_without_id, domain, parent, child):
             seed=self.seed,
             verbose=self.verbose,
         )
+        model._skip_validation = True
         model.fit(df_without_id[list(columns_meta)])
         return model
 
diff --git a/sdgym/synthesizers/tabddpm.py b/sdgym/synthesizers/tabddpm.py
index 0af5fec6..f042f0d3 100644
--- a/sdgym/synthesizers/tabddpm.py
+++ b/sdgym/synthesizers/tabddpm.py
@@ -1030,6 +1030,7 @@ def __init__(
         if device is None:
             device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self._device = torch.device(device)
+        self._skip_validation = False
         self._fitted = False
 
     def fit(self, data):
@@ -1043,7 +1044,8 @@ def fit(self, data):
         if isinstance(data, pd.DataFrame):
             data_dict = {self._table_name: data}
 
-        self._metadata.validate_data(data_dict)
+        if not self._skip_validation:
+            self._metadata.validate_data(data_dict)
 
         torch.manual_seed(self.seed)
         np.random.seed(self.seed)

From 33ca6908d96bddbd929ad558cfcad48a10feaa60 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 09:58:21 +0300
Subject: [PATCH 21/28] reorganize code

---
 pyproject.toml                  |   1 +
 sdgym/synthesizers/clavaddpm.py | 551 ++++++++++++++++++++------------
 2 files changed, 348 insertions(+), 204 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 42aec112..b17cc9ee 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,6 +28,7 @@ dependencies = [
     "cloudpickle>=2.1.0;python_version<'3.14'",
     "cloudpickle>=3.1.1;python_version>='3.14'",
     'compress-pickle>=1.2.0',
+    'faiss-cpu>=1.10.0',
     'google-cloud-compute>=1.30.0',
     'google-auth>=2.14.1',
     'humanfriendly>=10.0',
diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 5590aedb..20e3dc45 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -13,6 +13,7 @@
 from collections import Counter, defaultdict
 from datetime import datetime, timedelta
 
+import faiss
 import numpy as np
 import pandas as pd
 import sklearn
@@ -21,8 +22,8 @@
 from sdv.metadata import Metadata
 from sklearn.cluster import KMeans
 from sklearn.mixture import BayesianGaussianMixture, GaussianMixture
-from sklearn.neighbors import NearestNeighbors
 from sklearn.preprocessing import LabelEncoder, MinMaxScaler, OneHotEncoder
+from tqdm import tqdm
 
 from sdgym.synthesizers.base import MultiTableBaselineSynthesizer
 from sdgym.synthesizers.tabddpm import (
@@ -35,42 +36,113 @@
 SKLEARN_VERSION = Version(sklearn.__version__)
 
 
-def guess_array_datetime_format(values, sample_size=100, dayfirst=False):
-    """Guess the most likely datetime format via majority vote over a sample.
+# --------------------------------------------------------------------------
+# ######################## From preprocess_utils.py ########################
+# --------------------------------------------------------------------------
 
-    Args:
-        values (list, numpy.array, pd.Series):
-            List of datetime values to inspect.
-        sample_size (int, optional):
-            Number of samples to use for guessing. Default 100.
-        dayfirst (bool, optional):
-            If True parses dates with the day first. Default False.
 
-    Returns:
-        str: Datetime format string, or None if it can not be guessed.
+def calculate_days_since_earliest_date(dates, date_format='%y%m%d'):
+    """Encode a date column as integer days since its earliest date from preprocess_utils.py.
+
+    Modified by SDGym to take different `date_format` and handle Nans.
     """
-    pandas_version = Version(pd.__version__).release[:2]
-    if pandas_version >= (2, 2):
-        from pandas.tseries.api import guess_datetime_format  # pandas >= 2.2.0
-    else:
-        from pandas._libs.tslibs.parsing import guess_datetime_format  # pandas < 2.2.0
+    parsed = [
+        datetime.strptime(str(date), date_format) if pd.notna(date) else None for date in dates
+    ]
+    earliest_date = min(date for date in parsed if date is not None)
+    days_since = [(date - earliest_date).days if date is not None else np.nan for date in parsed]
+    return days_since, earliest_date.strftime(date_format)
 
-    series = pd.Series(values)
-    series = series.dropna()
-    if series.empty:
-        return None
 
-    sample_size = min(sample_size, len(series))
-    sample = series.sample(sample_size, random_state=0)
-    sample = sample.astype(str).str.strip()
-    sample = sample[sample != '']
+def reconstruct_dates(days_since, earliest_date_str, date_format='%y%m%d'):
+    """Inverse of `calculate_days_since_earliest_date` from preprocess_utils.py.
 
-    guesses = sample.apply(guess_datetime_format, dayfirst=dayfirst).dropna()
-    counts = Counter(guesses)
-    return top[0][0] if (top := counts.most_common(1)) else None
+    Modified by SDGym to take different `date_format` and handle Nans.
+    """
+    earliest_date = datetime.strptime(earliest_date_str, date_format)
+    original_dates = [
+        (earliest_date + timedelta(days=int(round(float(days))))).strftime(date_format)
+        if pd.notna(days)
+        else None
+        for days in days_since
+    ]
+    return original_dates
+
+
+def table_label_encode(df, discrete_cols):
+    """Label-encode the discrete columns of a table from preprocess_utils.py."""
+    df = df.copy()
+    label_encoders = {}
+    for col in discrete_cols:
+        le = LabelEncoder()
+        df[col] = le.fit_transform(df[col])
+        label_encoders[col] = le
+    return df, label_encoders
+
+
+def table_label_decode(df, label_encoders):
+    """Inverse of `table_label_encode` from preprocess_utils.py."""
+    df = df.copy()
+    for col, le in label_encoders.items():
+        df[col] = le.inverse_transform(df[col])
+    return df
+
+
+def get_domain(df, id_cols, discrete_cols):
+    """Build the ``{col: {'size', 'type'}}`` domain of a table from preprocess_utils.py."""
+    domain = {}
+    for col in df.columns:
+        if col in discrete_cols:
+            domain[col] = {'size': len(df[col].unique()), 'type': 'discrete'}
+        elif col not in id_cols:
+            domain[col] = {'size': len(df[col].unique()), 'type': 'continuous'}
+    return domain
+
+
+def topological_sort(graph):
+    """Order tables into ``[parent, child]`` relations from preprocess_utils.py."""
+    # Initialize the indegree map and output
+    in_degree = {node: 0 for node in graph}
+    for node in graph:
+        for child in graph[node]['children']:
+            in_degree[child] += 1
+
+    # Queue for nodes with no incoming edges
+    zero_in_degree = [node for node, degree in in_degree.items() if degree == 0]
+
+    # Output list for storing the order
+    sorted_order = []
+
+    # Start with root nodes and format them with None as parent
+    for node in zero_in_degree:
+        sorted_order.append([None, node])
+
+    # Using a queue to maintain nodes to process
+    queue = zero_in_degree[:]
+
+    while queue:
+        current = queue.pop(0)
+        for child in graph[current]['children']:
+            in_degree[child] -= 1
+            if in_degree[child] == 0:
+                queue.append(child)
+            # Add each parent-child relationship as we process them
+            sorted_order.append([current, child])
 
+    return sorted_order
+
+
+# --------------------------------------------------------------------------
+# ######################### From pipeline_utils.py #########################
+# --------------------------------------------------------------------------
 
-def get_group_data_dict(np_data, group_id_attrs=[0]):
+
+def get_group_data_dict(
+    np_data,
+    group_id_attrs=[
+        0,
+    ],
+):
     """Grouping dictionary from pipeline_utils.py."""
     group_data_dict = {}
     data_len = len(np_data)
@@ -79,10 +151,16 @@ def get_group_data_dict(np_data, group_id_attrs=[0]):
         if row_id not in group_data_dict:
             group_data_dict[row_id] = []
         group_data_dict[row_id].append(np_data[i])
+
     return group_data_dict
 
 
-def get_group_data(np_data, group_id_attrs=[0]):
+def get_group_data(
+    np_data,
+    group_id_attrs=[
+        0,
+    ],
+):
     """Grouping list from pipeline_utils.py."""
     group_data_list = []
     data_len = len(np_data)
@@ -90,19 +168,26 @@ def get_group_data(np_data, group_id_attrs=[0]):
     while i < data_len:
         group = []
         row_id = np_data[i, group_id_attrs]
+
         while (np_data[i, group_id_attrs] == row_id).all():
             group.append(np_data[i])
             i += 1
             if i >= data_len:
                 break
-        group_data_list.append(np.array(group))
-    return np.array(group_data_list, dtype=object)
+        group = np.array(group)
+        group_data_list.append(group)
+    group_data_list = np.array(group_data_list, dtype=object)
+
+    return group_data_list
 
 
 def min_max_normalize_sklearn(matrix):
     """Apply MinMaxScaler to each column from pipeline_utils.py."""
     scaler = MinMaxScaler(feature_range=(-1, 1))
+
     normalized_data = np.empty((matrix.shape[0], 0))
+
+    # Apply MinMaxScaler to each column and concatenate the results
     for col in range(matrix.shape[1]):
         column = matrix[:, col].reshape(-1, 1)
         transformed_column = scaler.fit_transform(column)
@@ -111,24 +196,6 @@ def min_max_normalize_sklearn(matrix):
     return normalized_data
 
 
-def aggregate_and_sample(cluster_probabilities, child_group_lengths):
-    """Aggregate the distribution and sample from pipeline_modules.py."""
-    group_cluster_labels = []
-    curr_index = 0
-    agree_rates = []
-    for group_length in child_group_lengths:
-        group_probability_distribution = np.mean(
-            cluster_probabilities[curr_index : curr_index + group_length], axis=0
-        )
-        group_cluster_label = np.random.choice(
-            range(len(group_probability_distribution)), p=group_probability_distribution
-        )
-        group_cluster_labels.append(group_cluster_label)
-        agree_rates.append(np.max(group_probability_distribution))
-        curr_index += group_length
-    return group_cluster_labels, agree_rates
-
-
 def freq_to_prob(freq_dict):
     """Converts a dict of frequencies to a dict of probabilities from pipeline_utils.py."""
     prob_dict = {}
@@ -139,122 +206,185 @@ def freq_to_prob(freq_dict):
 
 def sample_from_dict(probabilities):
     """Sample using a dict of probabilities from pipeline_utils.py."""
+    # Generate a random number between 0 and 1
     random_number = random.random()
+
+    # Initialize cumulative sum and the selected key
     cumulative_sum = 0
     selected_key = None
+
+    # Iterate through the dictionary
     for key, probability in probabilities.items():
         cumulative_sum += probability
         if cumulative_sum >= random_number:
             selected_key = key
             break
+
     return selected_key
 
 
 def get_df_without_id(df, id_cols):
     """Drop id columns based on `id_cols` from pipeline_utils.py."""
-    return df.drop(columns=[col for col in id_cols if col in df.columns])
+    id_cols = [col for col in df.columns if '_id' in col]
+    return df.drop(columns=id_cols)
+
+
+def convert_to_unique_indices(indices):
+    """Convert indices to unique values from pipline_utils.py."""
+    occurrence = set()
+    max_index = len(indices)  # Assuming the range is the length of the list
+    replacement_candidates = set(range(max_index)) - set(indices)
+
+    for i, num in enumerate(tqdm(indices)):
+        if num in occurrence:
+            # Find the smallest number not in the list
+            replacement = min(replacement_candidates)
+            indices[i] = replacement
+            replacement_candidates.remove(replacement)
+        else:
+            occurrence.add(num)
 
+    return indices
 
-def calculate_days_since_earliest_date(dates, date_format='%y%m%d'):
-    """Encode a date column as integer days since its earliest date from preprocess_utils.py."""
-    parsed = [
-        datetime.strptime(str(date), date_format) if pd.notna(date) else None for date in dates
-    ]
-    earliest_date = min(date for date in parsed if date is not None)
-    days_since = [(date - earliest_date).days if date is not None else np.nan for date in parsed]
-    return days_since, earliest_date.strftime(date_format)
 
+def match_tables(A, B, n_clusters=25, unique_matching=True, batch_size=100):
+    """Nearest-neighbour match of every row of ``A`` to a row of ``B`` from pipline_utils.py."""
+    A = np.ascontiguousarray(A, dtype=np.float32)
+    B = np.ascontiguousarray(B, dtype=np.float32)
 
-def reconstruct_dates(days_since, earliest_date_str, date_format='%y%m%d'):
-    """Inverse of `calculate_days_since_earliest_date` from preprocess_utils.py."""
-    earliest_date = datetime.strptime(earliest_date_str, date_format)
-    return [
-        (earliest_date + timedelta(days=int(round(float(days))))).strftime(date_format)
-        if pd.notna(days)
-        else None
-        for days in days_since
-    ]
+    # Dimension of vectors
+    d = B.shape[1]
 
+    if unique_matching:
+        quantiser = faiss.IndexFlatL2(d)
+        index = faiss.IndexIVFFlat(quantiser, d, n_clusters, faiss.METRIC_L2)
+    else:
+        res = faiss.StandardGpuResources()
+        quantiser = faiss.IndexFlatL2(d)
+        index_cpu = faiss.IndexIVFFlat(quantiser, d, n_clusters, faiss.METRIC_L2)
+        index = faiss.index_cpu_to_gpu(res, 0, index_cpu)
+
+    index.train(B)
+    index.add(B)
+
+    # Initialize lists to store the results
+    all_indices = []
+    all_distances = []
+
+    if unique_matching:
+        batch_size = 1
+        n_batches = (A.shape[0] + batch_size - 1) // batch_size
+
+        for i in tqdm(range(n_batches)):
+            start = i * batch_size
+            end = min((i + 1) * batch_size, A.shape[0])
+            D, I = index.search(A[start:end], k=1)  # noqa: E741
+            index.remove_ids(I.flatten())
+            all_distances.append(D)
+            all_indices.append(I)
+
+        # Concatenate the results from all batches
+        all_distances = np.vstack(all_distances)
+        all_indices = np.vstack(all_indices)
+        distances = all_distances.flatten().tolist()
+        indices = all_indices.flatten().tolist()
+    else:
+        n_batches = (A.shape[0] + batch_size - 1) // batch_size
 
-def table_label_encode(df, discrete_cols):
-    """Label-encode the discrete columns of a table from preprocess_utils.py."""
-    df = df.copy()
-    label_encoders = {}
-    for col in discrete_cols:
-        le = LabelEncoder()
-        df[col] = le.fit_transform(df[col])
-        label_encoders[col] = le
-    return df, label_encoders
+        for i in tqdm(range(n_batches)):
+            start = i * batch_size
+            end = min((i + 1) * batch_size, A.shape[0])
+            D, I = index.search(A[start:end], k=1)  # noqa: E741
+            all_distances.append(D)
+            all_indices.append(I)
 
+        # Concatenate the results from all batches
+        all_distances = np.vstack(all_distances)
+        all_indices = np.vstack(all_indices)
+        distances = all_distances.flatten().tolist()
+        indices = all_indices.flatten().tolist()
+        indices = convert_to_unique_indices(indices)
+        assert len(indices) == len(set(indices))
 
-def table_label_decode(df, label_encoders):
-    """Inverse of :func:`table_label_encode` from preprocess_utils.py."""
-    df = df.copy()
-    for col, le in label_encoders.items():
-        df[col] = le.inverse_transform(df[col])
-    return df
+    return indices, distances
 
 
-def get_domain(df, id_cols, discrete_cols):
-    """Build the ``{col: {'size', 'type'}}`` domain of a table from preprocess_utils.py."""
-    domain = {}
-    for col in df.columns:
-        if col in discrete_cols:
-            domain[col] = {'size': len(df[col].unique()), 'type': 'discrete'}
-        elif col not in id_cols:
-            domain[col] = {'size': len(df[col].unique()), 'type': 'continuous'}
-    return domain
+def handle_multi_parent(
+    child,
+    parents,
+    synthetic_tables,
+    id_cols,
+    n_clusters=1,
+    unique_matching=True,
+    batch_size=100,
+    no_matching=False,
+):
+    """Reconcile a child generated once per parent into a single table from pipeline_utils.py.
+
+    Modified code to use 'fk' instead of assuming '{parent}_id'.
+    """
+    synthetic_child_dfs = [
+        (synthetic_tables[(parent, child)]['df'].copy(), fk) for parent, fk, _ppk in parents
+    ]
+    anchor_index = int(np.argmin([len(df) for df, _ in synthetic_child_dfs]))
+    anchor = synthetic_child_dfs[anchor_index]
+    synthetic_child_dfs.pop(anchor_index)
+    for df, fk in synthetic_child_dfs:
+        df_without_ids = get_df_without_id(df, id_cols)
+        anchor_df_without_ids = get_df_without_id(anchor[0], id_cols)
 
+        df_val = df_without_ids.to_numpy().astype(float)
+        anchor_val = anchor_df_without_ids.to_numpy().astype(float)
+        if len(df_val.shape) == 1:
+            df_val = df_val.reshape(-1, 1)
+            anchor_val = anchor_val.reshape(-1, 1)
 
-def decode_id_values(values, encoder, column):
-    """Inverse-transform label-encoded id codes back to the original id values.
+        indices, _ = match_tables(
+            anchor_val,
+            df_val,
+            n_clusters=n_clusters,
+            unique_matching=unique_matching,
+            batch_size=batch_size,
+        )
+        if no_matching:
+            # randomly shuffle the array
+            indices = np.random.permutation(indices)
+        df = df.iloc[indices]
+        anchor[0][fk] = df[fk].to_numpy()
+    return anchor[0]
 
-    Sampling can create more rows than the original table, so codes past the
-    encoder's fitted range get fresh values: numeric ids continue after the
-    largest original id, other ids become new ``'{column}_{code}'`` strings.
-    The extension is deterministic per code, so a foreign key decoded with its
-    parent's encoder still matches the parent's decoded primary key.
-    """
-    codes = np.asarray(values).astype('int64')
-    classes = encoder.classes_
-    decoded = np.empty(len(codes), dtype=object)
-    in_range = (codes >= 0) & (codes < len(classes))
-    decoded[in_range] = encoder.inverse_transform(codes[in_range])
-    if (~in_range).any():
-        if pd.api.types.is_numeric_dtype(classes):
-            # continue past the largest original id (code n_classes -> max + 1, ...)
-            decoded[~in_range] = codes[~in_range] + (classes.max() + 1 - len(classes))
-        else:
-            decoded[~in_range] = [f'{column}_{code}' for code in codes[~in_range]]
 
-    if pd.api.types.is_numeric_dtype(classes):
-        return decoded.astype(classes.dtype)
-    return decoded
+# --------------------------------------------------------------------------
+# ######################## From pipeline_modules.py ########################
+# --------------------------------------------------------------------------
 
 
-def topological_sort(graph):
-    """Order tables into ``[parent, child]`` relations from preprocess_utils.py."""
-    in_degree = {node: 0 for node in graph}
-    for node in graph:
-        for child in graph[node]['children']:
-            in_degree[child] += 1
+def aggregate_and_sample(cluster_probabilities, child_group_lengths):
+    """Aggregate the distribution and sample from pipeline_modules.py."""
+    group_cluster_labels = []
+    curr_index = 0
+    agree_rates = []
 
-    zero_in_degree = [node for node, degree in in_degree.items() if degree == 0]
+    for group_length in child_group_lengths:
+        # Aggregate the probability distributions by taking the mean
+        group_probability_distribution = np.mean(
+            cluster_probabilities[curr_index : curr_index + group_length], axis=0
+        )
 
-    sorted_order = []
-    for node in zero_in_degree:
-        sorted_order.append([None, node])
+        # Sample the label from the aggregated distribution
+        group_cluster_label = np.random.choice(
+            range(len(group_probability_distribution)), p=group_probability_distribution
+        )
+        group_cluster_labels.append(group_cluster_label)
 
-    queue = zero_in_degree[:]
-    while queue:
-        current = queue.pop(0)
-        for child in graph[current]['children']:
-            in_degree[child] -= 1
-            if in_degree[child] == 0:
-                queue.append(child)
-            sorted_order.append([current, child])
+        # Compute the max probability as the agree rate
+        max_probability = np.max(group_probability_distribution)
+        agree_rates.append(max_probability)
 
-    return sorted_order
+        # Update the curr_index for the next iteration
+        curr_index += group_length
+
+    return group_cluster_labels, agree_rates
 
 
 def pair_clustering_keep_id(
@@ -273,7 +403,12 @@ def pair_clustering_keep_id(
     clustering_method='kmeans',
     seed=0,
 ):
-    """Cluster child rows augmented with their parent's features from pipeline_modules.py."""
+    """Cluster child rows augmented with their parent's features from pipeline_modules.py.
+
+    Modified by SDGym to:
+    * remove output statements
+    * handle different versions of sklearn
+    """
     original_child_cols = list(child_df.columns)
     original_parent_cols = list(parent_df.columns)
 
@@ -509,77 +644,90 @@ def pair_clustering_keep_id(
     return parent_df_with_cluster, child_df_with_cluster, group_lengths_prob_dicts
 
 
-def match_tables(A, B, n_clusters=25, unique_matching=True, batch_size=100):
-    """Nearest-neighbour match of every row of ``A`` to a row of ``B``.
+# --------------------------------------------------------------------------
+# ########################### SDGym extra logic ############################
+# --------------------------------------------------------------------------
+def _check_faiss_installed():
+    try:
+        import faiss  # noqa: F401
+    except ImportError as e:
+        raise ImportError(
+            "\nThe 'faiss' package is not installed.\n"
+            'Please install it via conda before running this synthesizer., e.g.:\n\n'
+            '    conda install -c pytorch -c nvidia faiss-gpu\n'
+        ) from e
+
+
+def guess_array_datetime_format(values, sample_size=100, dayfirst=False):
+    """Guess the most likely datetime format via majority vote over a sample.
+
+    Args:
+        values (list, numpy.array, pd.Series):
+            List of datetime values to inspect.
+        sample_size (int, optional):
+            Number of samples to use for guessing. Default 100.
+        dayfirst (bool, optional):
+            If True parses dates with the day first. Default False.
 
-    n_clusters is used by the original implementation for faiss.IndexIVFFlat.
+    Returns:
+        str: Datetime format string, or None if it can not be guessed.
     """
-    A = np.ascontiguousarray(A, dtype=np.float32)
-    B = np.ascontiguousarray(B, dtype=np.float32)
+    pandas_version = Version(pd.__version__).release[:2]
+    if pandas_version >= (2, 2):
+        from pandas.tseries.api import guess_datetime_format  # pandas >= 2.2.0
+    else:
+        from pandas._libs.tslibs.parsing import guess_datetime_format  # pandas < 2.2.0
 
-    if not unique_matching:
-        distances, indices = NearestNeighbors(n_neighbors=1).fit(B).kneighbors(A)
-        return indices.flatten().tolist(), distances.flatten().tolist()
-
-    k = min(len(B), 50)
-    distances, indices = NearestNeighbors(n_neighbors=k).fit(B).kneighbors(A)
-    used = set()
-    matched_indices = []
-    matched_distances = []
-    for row_indices, row_distances in zip(indices, distances):
-        chosen, chosen_distance = None, 0.0
-        for candidate, distance in zip(row_indices, row_distances):
-            if int(candidate) not in used:
-                chosen, chosen_distance = int(candidate), float(distance)
-                break
-        if chosen is None:
-            chosen = next(j for j in range(len(B)) if j not in used)
-        used.add(chosen)
-        matched_indices.append(chosen)
-        matched_distances.append(chosen_distance)
+    series = pd.Series(values)
+    series = series.dropna()
+    if series.empty:
+        return None
 
-    return matched_indices, matched_distances
+    sample_size = min(sample_size, len(series))
+    sample = series.sample(sample_size, random_state=0)
+    sample = sample.astype(str).str.strip()
+    sample = sample[sample != '']
 
+    guesses = sample.apply(guess_datetime_format, dayfirst=dayfirst).dropna()
+    counts = Counter(guesses)
+    return top[0][0] if (top := counts.most_common(1)) else None
 
-def handle_multi_parent(
-    child,
-    parents,
-    synthetic_tables,
-    id_cols,
-    n_clusters=1,
-    unique_matching=True,
-    batch_size=100,
-    no_matching=False,
-):
-    """Reconcile a child generated once per parent into a single table."""
-    synthetic_child_dfs = [
-        (synthetic_tables[(parent, child)]['df'].copy(), fk) for parent, fk, _ppk in parents
-    ]
-    anchor_index = int(np.argmin([len(df) for df, _ in synthetic_child_dfs]))
-    anchor = synthetic_child_dfs[anchor_index]
-    synthetic_child_dfs.pop(anchor_index)
-    for df, fk in synthetic_child_dfs:
-        df_without_ids = get_df_without_id(df, id_cols)
-        anchor_df_without_ids = get_df_without_id(anchor[0], id_cols)
 
-        df_val = df_without_ids.to_numpy().astype(float)
-        anchor_val = anchor_df_without_ids.to_numpy().astype(float)
-        if len(df_val.shape) == 1:
-            df_val = df_val.reshape(-1, 1)
-            anchor_val = anchor_val.reshape(-1, 1)
+def decode_id_values(values, encoder, column):
+    """Inverse-transform label-encoded id codes back to the original id values.
 
-        indices, _ = match_tables(
-            anchor_val,
-            df_val,
-            n_clusters=n_clusters,
-            unique_matching=unique_matching,
-            batch_size=batch_size,
-        )
-        if no_matching:
-            indices = np.random.permutation(indices)
-        df = df.iloc[indices]
-        anchor[0][fk] = df[fk].to_numpy()
-    return anchor[0]
+    Sampling can create more rows than the original table, so codes past the
+    encoder's fitted range get fresh values: numeric ids continue after the
+    largest original id, other ids become new ``'{column}_{code}'`` strings.
+
+    Args:
+        values (list, numpy.array, pd.Series):
+            ID column values.
+        encoder (sklearn.preprocessing.LabelEncoder):
+            Fitted encoder.
+        column (str):
+            Column name.
+
+    Returns:
+        np.array:
+            List of decoded id values.
+    """
+    codes = np.asarray(values).astype('int64')
+    classes = encoder.classes_
+    decoded = np.empty(len(codes), dtype=object)
+    in_range = (codes >= 0) & (codes < len(classes))
+    decoded[in_range] = encoder.inverse_transform(codes[in_range])
+    if (~in_range).any():
+        if pd.api.types.is_numeric_dtype(classes):
+            # continue past the largest original id (code n_classes -> max + 1, ...)
+            decoded[~in_range] = codes[~in_range] + (classes.max() + 1 - len(classes))
+        else:
+            decoded[~in_range] = [f'{column}_{code}' for code in codes[~in_range]]
+
+    if pd.api.types.is_numeric_dtype(classes):
+        return decoded.astype(classes.dtype)
+
+    return decoded
 
 
 @torch.no_grad()
@@ -588,7 +736,6 @@ def _sample_step(diffusion, y):
 
     Mirrors ``GaussianMultinomialDiffusion._sample`` but takes explicit
     per-row labels instead of drawing them from the empirical distribution.
-    Similar to the original classifier-guided ``conditional_sample`` implemtation.
     """
     device = diffusion.log_alpha.device
     b = y.shape[0]
@@ -716,10 +863,6 @@ class ClavaDDPM:
             Multi-parent matching controls for unique matching.
         no_matching (bool):
             Multi-parent matching controls for no matching.
-        max_categories (int):
-            Maximum number of distinct values kept for a categorical column.
-            Columns with more are reduced to this many values, randomly sampled
-            from the data.
         device (str or None):
             Whether to use ``'cuda'`` or ``'cpu'``. If None, auto-select is used.
         seed (int):
@@ -749,11 +892,12 @@ def __init__(
         matching_batch_size=1000,
         unique_matching=True,
         no_matching=False,
-        max_categories=100,
         device=None,
         seed=0,
         verbose=False,
     ):
+        _check_faiss_installed()
+
         self.num_clusters = num_clusters
         self.parent_scale = parent_scale
         self.key_scale = key_scale
@@ -772,7 +916,6 @@ def __init__(
         self.matching_batch_size = matching_batch_size
         self.unique_matching = unique_matching
         self.no_matching = no_matching
-        self.max_categories = max_categories
         self.device = device
         self.seed = seed
         self.verbose = verbose

From 5fa3b40ee534d9afdfb18c5ebba7a2df6c7ea457 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 10:15:59 +0300
Subject: [PATCH 22/28] address comments

---
 sdgym/synthesizers/clavaddpm.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 20e3dc45..514fd91b 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -438,11 +438,11 @@ def pair_clustering_keep_id(
                 parent_num_cols.append((col_index, col))
 
     parent_primary_key_index = original_parent_cols.index(parent_primary_key)
-    foreing_key_index = original_child_cols.index(foreign_key)
+    foreing_key_index = original_child_cols.index(parent_primary_key)
 
     # sort child data by foreign key
     sorted_child_data = child_data[np.argsort(child_data[:, foreing_key_index])]
-    child_group_data_dict = get_group_data_dict(sorted_child_data, [foreing_key_index])
+    child_group_data_dict = get_group_data_dict(sorted_child_data, [foreing_key_index,])
 
     # sort parent data by primary key
     sorted_parent_data = parent_data[np.argsort(parent_data[:, parent_primary_key_index])]
@@ -505,7 +505,6 @@ def pair_clustering_keep_id(
             parent_scale * cat_one_hot[:, joint_cat_matrix_p_index:]
         )
 
-    # Perform quantile normalization using QuantileTransformer
     num_min_max = min_max_normalize_sklearn(joint_num_matrix)
 
     key_factorized = pd.factorize(sorted_parent_data_repeated[:, parent_primary_key_index])[0]
@@ -528,8 +527,12 @@ def pair_clustering_keep_id(
     if SKLEARN_VERSION.release[:2] < (1, 1):
         init_param = 'kmeans'
 
+    n_init = 'auto'
+    if SKLEARN_VERSION.release[:2] < (1, 2):
+        n_init = 10
+
     if clustering_method == 'kmeans':
-        kmeans = KMeans(n_clusters=num_clusters, n_init='auto', init='k-means++', random_state=seed)
+        kmeans = KMeans(n_clusters=num_clusters, n_init=n_init, init='k-means++', random_state=seed)
         kmeans.fit(cluster_data)
         cluster_labels = kmeans.labels_
     elif clustering_method == 'both':
@@ -928,7 +931,6 @@ def _parse_metadata(self, metadata):
         meta = metadata if isinstance(metadata, dict) else metadata.to_dict()
 
         self._table_names = list(meta['tables'])
-        self._primary_key = {}
         # child -> list of (parent_table, foreign_key_col, parent_primary_key)
         self._parents = {name: [] for name in self._table_names}
         self._children = {name: [] for name in self._table_names}
@@ -1312,7 +1314,6 @@ class ClavaDDPMSynthesizer(MultiTableBaselineSynthesizer):
 
     LOGGER = logging.getLogger(__name__)
     _MODEL_KWARGS = None
-    _MODALITY_FLAG = 'multi_table'
 
     def _fit(self, data, metadata):
         """Fit the synthesizer to the multi-table data.

From ec384ac0b278f61741d8c21d1d2cc8a2895134ed Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 10:25:33 +0300
Subject: [PATCH 23/28] fix lint

---
 sdgym/synthesizers/clavaddpm.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 514fd91b..6902003a 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -442,7 +442,12 @@ def pair_clustering_keep_id(
 
     # sort child data by foreign key
     sorted_child_data = child_data[np.argsort(child_data[:, foreing_key_index])]
-    child_group_data_dict = get_group_data_dict(sorted_child_data, [foreing_key_index,])
+    child_group_data_dict = get_group_data_dict(
+        sorted_child_data,
+        [
+            foreing_key_index,
+        ],
+    )
 
     # sort parent data by primary key
     sorted_parent_data = parent_data[np.argsort(parent_data[:, parent_primary_key_index])]

From 64c84c7b86dd32c3d820be4d09cfb7bbd1a6c966 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 21:19:24 +0300
Subject: [PATCH 24/28] fix index

---
 sdgym/synthesizers/clavaddpm.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 6902003a..2e7a933c 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -438,7 +438,7 @@ def pair_clustering_keep_id(
                 parent_num_cols.append((col_index, col))
 
     parent_primary_key_index = original_parent_cols.index(parent_primary_key)
-    foreing_key_index = original_child_cols.index(parent_primary_key)
+    foreing_key_index = original_child_cols.index(foreign_key)
 
     # sort child data by foreign key
     sorted_child_data = child_data[np.argsort(child_data[:, foreing_key_index])]

From dfdae004106fcce88bbcad8b7bc13f7f4d124b0a Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 21:31:14 +0300
Subject: [PATCH 25/28] seperate faiss req

---
 pyproject.toml | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index b17cc9ee..200a2bb0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,6 @@ dependencies = [
     "cloudpickle>=2.1.0;python_version<'3.14'",
     "cloudpickle>=3.1.1;python_version>='3.14'",
     'compress-pickle>=1.2.0',
-    'faiss-cpu>=1.10.0',
     'google-cloud-compute>=1.30.0',
     'google-auth>=2.14.1',
     'humanfriendly>=10.0',
@@ -86,8 +85,12 @@ realtabformer = [
     "torch>=2.6.0",
     'transformers<4.51',
 ]
+clavaddpm = [
+    'faiss-cpu>=1.10.0',
+    'numpy>=1.25.0',
+]
 test = [
-    'sdgym[realtabformer]',
+    'sdgym[realtabformer, clavaddpm]',
     'pytest>=6.2.5',
     'pytest-cov>=2.6.0',
     'jupyter>=1.0.0,<2',

From 5d8c746cd228b56f58b04a46f5f52be355caedf5 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 21:51:06 +0300
Subject: [PATCH 26/28] fix

---
 pyproject.toml | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 200a2bb0..eaf7b3d4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -86,11 +86,13 @@ realtabformer = [
     'transformers<4.51',
 ]
 clavaddpm = [
-    'faiss-cpu>=1.10.0',
+    "faiss-cpu>=1.10.0;python_version<'3.14'",
+    "faiss-cpu>=1.12.0;python_version>='3.14'",
     'numpy>=1.25.0',
 ]
 test = [
-    'sdgym[realtabformer, clavaddpm]',
+    'sdgym[realtabformer]',
+    'sdgym[clavaddpm]',
     'pytest>=6.2.5',
     'pytest-cov>=2.6.0',
     'jupyter>=1.0.0,<2',

From 09c06355b2d2e3383abf6edbf551355986f49648 Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Tue, 18 Aug 2026 22:14:51 +0300
Subject: [PATCH 27/28] add scipy dep

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index eaf7b3d4..726d459c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -88,6 +88,7 @@ realtabformer = [
 clavaddpm = [
     "faiss-cpu>=1.10.0;python_version<'3.14'",
     "faiss-cpu>=1.12.0;python_version>='3.14'",
+    'scipy>=1.9.2',
     'numpy>=1.25.0',
 ]
 test = [

From dbba3aa445859eabac22aa9dcf34c0423b9d93da Mon Sep 17 00:00:00 2001
From: sarahmish <sarahalnegheimish@gmail.com>
Date: Wed, 19 Aug 2026 09:30:30 +0300
Subject: [PATCH 28/28] import faiss after call

---
 sdgym/synthesizers/clavaddpm.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/sdgym/synthesizers/clavaddpm.py b/sdgym/synthesizers/clavaddpm.py
index 2e7a933c..66aaa5f2 100644
--- a/sdgym/synthesizers/clavaddpm.py
+++ b/sdgym/synthesizers/clavaddpm.py
@@ -13,7 +13,6 @@
 from collections import Counter, defaultdict
 from datetime import datetime, timedelta
 
-import faiss
 import numpy as np
 import pandas as pd
 import sklearn
@@ -660,9 +659,10 @@ def _check_faiss_installed():
         import faiss  # noqa: F401
     except ImportError as e:
         raise ImportError(
-            "\nThe 'faiss' package is not installed.\n"
-            'Please install it via conda before running this synthesizer., e.g.:\n\n'
-            '    conda install -c pytorch -c nvidia faiss-gpu\n'
+            "In order to use 'ClavaDDPMSynthesizer' you have to install the extra "
+            "dependencies by first installing 'faiss' library and other dependencies: \n\n"
+            "    conda install -c pytorch -c nvidia faiss-gpu\n"
+            "    pip install sdgym['clavaddpm']\n"
         ) from e