synthe

  7class AdaptiveHistogramSampler:
  8    def __init__(self, n_bins=10, method="quantile", seed=123):
  9        self.n_bins = n_bins
 10        self.method = method
 11        self.rng = np.random.default_rng(seed)
 12        self.bin_edges = None
 13        self.bin_indices = None
 14        self.unique_bins = None
 15        self.bin_probs = None
 16        self.X = None
 17        self.n = None
 18        self.d = None
 19
 20    def fit(self, X):
 21        self.X = np.asarray(X)
 22        self.n, self.d = self.X.shape
 23
 24        self.bin_edges = []
 25        for j in range(self.d):
 26            xj = self.X[:, j]
 27            if self.method == "quantile":
 28                edges_j = np.quantile(xj, np.linspace(0, 1, self.n_bins + 1))
 29            else:
 30                edges_j = np.linspace(xj.min(), xj.max(), self.n_bins + 1)
 31            self.bin_edges.append(edges_j)
 32
 33        # Assign points to bins
 34        bin_idx = np.zeros((self.n, self.d), dtype=int)
 35        for j in range(self.d):
 36            bin_idx[:, j] = np.digitize(self.X[:, j], self.bin_edges[j]) - 1
 37            bin_idx[:, j] = np.clip(bin_idx[:, j], 0, self.n_bins - 1)
 38        self.bin_indices = bin_idx
 39
 40        bin_ids = np.ravel_multi_index(
 41            self.bin_indices.T, (self.n_bins,) * self.d
 42        )
 43        unique_bins, counts = np.unique(bin_ids, return_counts=True)
 44        self.unique_bins = unique_bins
 45        self.bin_probs = counts / counts.sum()
 46
 47    def sample(
 48        self,
 49        n_samples,
 50        oversample=False,
 51        oversample_method="bootstrap",
 52        jitter_scale=0.05,
 53    ):
 54        if self.bin_probs is None:
 55            raise RuntimeError("You must call `fit` before `sample`.")
 56
 57        chosen_bins = self.rng.choice(
 58            self.unique_bins, size=n_samples, p=self.bin_probs
 59        )
 60
 61        if not oversample:
 62            return self._subsample_existing(chosen_bins)
 63
 64        if oversample_method == "uniform":
 65            return self._oversample_uniform(chosen_bins)
 66        elif oversample_method == "bootstrap":
 67            return self._oversample_bootstrap(chosen_bins)
 68        elif oversample_method == "jitter":
 69            return self._oversample_jitter(chosen_bins, jitter_scale)
 70        else:
 71            raise ValueError(f"Unknown oversample_method: {oversample_method}")
 72
 73    # --- Internal helpers ------------------------------------------------
 74    def _subsample_existing(self, chosen_bins):
 75        bin_ids = np.ravel_multi_index(
 76            self.bin_indices.T, (self.n_bins,) * self.d
 77        )
 78        X_sampled = []
 79        for b in chosen_bins:
 80            idx_in_bin = np.where(bin_ids == b)[0]
 81            i = self.rng.choice(idx_in_bin)
 82            X_sampled.append(self.X[i])
 83        return np.array(X_sampled)
 84
 85    def _oversample_uniform(self, chosen_bins):
 86        X_sampled = []
 87        for b in chosen_bins:
 88            multi_idx = np.unravel_index(b, (self.n_bins,) * self.d)
 89            coords = []
 90            for j, bi in enumerate(multi_idx):
 91                left = self.bin_edges[j][bi]
 92                right = self.bin_edges[j][bi + 1]
 93                coords.append(self.rng.uniform(left, right))
 94            X_sampled.append(coords)
 95        return np.array(X_sampled)
 96
 97    def _oversample_bootstrap(self, chosen_bins):
 98        return self._subsample_existing(chosen_bins)
 99
100    def _oversample_jitter(self, chosen_bins, jitter_scale):
101        base_points = self._subsample_existing(chosen_bins)
102        noise = self.rng.normal(scale=jitter_scale, size=base_points.shape)
103        return base_points + noise
104
105    # --- Visualization ----------------------------------------------------
106    def plot_comparison(self, X_sampled, bins=30):
107        """Plot joint 2D histogram and marginal distributions (for d=2)."""
108        if self.d != 2:
109            raise ValueError("plot_comparison only supports 2D currently.")
110
111        fig = plt.figure(figsize=(10, 10))
112        grid = plt.GridSpec(4, 4, hspace=0.3, wspace=0.3)
113
114        main_ax = fig.add_subplot(grid[1:, :-1])
115        y_hist = fig.add_subplot(grid[0, :-1], sharex=main_ax)
116        x_hist = fig.add_subplot(grid[1:, -1], sharey=main_ax)
117
118        # 2D histogram
119        main_ax.hist2d(
120            self.X[:, 0], self.X[:, 1], bins=bins, alpha=0.5, cmap="Blues"
121        )
122        main_ax.hist2d(
123            X_sampled[:, 0], X_sampled[:, 1], bins=bins, alpha=0.5, cmap="Reds"
124        )
125        main_ax.set_xlabel("X1")
126        main_ax.set_ylabel("X2")
127        main_ax.set_title("Joint distribution")
128
129        # Marginals
130        y_hist.hist(
131            self.X[:, 0], bins=bins, color="blue", alpha=0.5, density=True
132        )
133        y_hist.hist(
134            X_sampled[:, 0], bins=bins, color="red", alpha=0.5, density=True
135        )
136        x_hist.hist(
137            self.X[:, 1],
138            bins=bins,
139            orientation="horizontal",
140            color="blue",
141            alpha=0.5,
142            density=True,
143        )
144        x_hist.hist(
145            X_sampled[:, 1],
146            bins=bins,
147            orientation="horizontal",
148            color="red",
149            alpha=0.5,
150            density=True,
151        )
152
153        y_hist.axis("off")
154        x_hist.axis("off")
155
156        plt.show()
157
158    # --- Goodness of fit tests -------------------------------------------
159    def goodness_of_fit(self, X_sampled):
160        """
161        Compare marginals with Kolmogorov–Smirnov and Anderson–Darling tests.
162        Returns dict of test results for each dimension.
163        """
164        results = {}
165        for j in range(self.d):
166            x_orig = self.X[:, j]
167            x_samp = X_sampled[:, j]
168
169            # KS test
170            ks_stat, ks_p = stats.ks_2samp(x_orig, x_samp)
171
172            # Anderson-Darling test
173            ad_result = stats.anderson_ksamp([x_orig, x_samp])
174
175            results[f"dim_{j}"] = {
176                "ks_statistic": ks_stat,
177                "ks_pvalue": ks_p,
178                "ad_statistic": ad_result.statistic,
179                "ad_significance_level": ad_result.significance_level,
180            }
181        return results

  11class DiffusionModel(BaseEstimator):
  12    """
  13    Sklearn-compatible diffusion model with MMD-based and noise-prediction training.
  14
  15    Implements both traditional DDPM (noise prediction with MSE) and novel MMD-based
  16    training that directly minimizes distribution mismatch between true posterior
  17    and learned transitions using Maximum Mean Discrepancy.
  18
  19    Parameters
  20    ----------
  21    timesteps : int, default=1000
  22        Number of diffusion timesteps
  23    beta_start : float, default=0.0001
  24        Initial noise variance
  25    beta_end : float, default=0.02
  26        Final noise variance
  27    model : sklearn estimator, optional
  28        Base model for reverse process (default: Ridge with alpha=1.0)
  29    schedule : {'linear', 'cosine'}, default='linear'
  30        Noise schedule type
  31    use_pca : bool, default=False
  32        Apply PCA for dimensionality reduction (recommended for >100 dims)
  33    pca_components : int, default=50
  34        Number of PCA components if use_pca=True
  35    variance_type : {'fixed_small', 'fixed_large', 'learned'}, default='fixed_small'
  36        Variance schedule for reverse process
  37    random_state : int, optional
  38        Random seed for reproducibility
  39    batch_size : int, default=32
  40        Batch size for training data generation
  41    training_objective : {'noise', 'mmd', 'hybrid'}, default='noise'
  42        Training objective:
  43        - 'noise': Traditional DDPM noise prediction with MSE loss
  44        - 'mmd': Direct MMD minimization between true and learned posteriors
  45        - 'hybrid': Combine both objectives
  46    mmd_samples_per_step : int, default=10
  47        Number of samples to draw per timestep for MMD estimation
  48    mmd_kernel : {'rbf', 'imq', 'linear'}, default='rbf'
  49        Kernel for MMD computation
  50    mmd_bandwidth : float or 'auto', default='auto'
  51        Kernel bandwidth (gamma for RBF)
  52
  53    Examples
  54    --------
  55    Traditional noise-prediction training:
  56    >>> model = DiffusionModel(timesteps=100, training_objective='noise')
  57    >>> model.fit(X, n_steps=1000)
  58
  59    MMD-based training (distribution matching):
  60    >>> model = DiffusionModel(timesteps=100, training_objective='mmd',
  61    ...                        mmd_samples_per_step=20)
  62    >>> model.fit(X, n_steps=1000)
  63
  64    Hybrid approach:
  65    >>> model = DiffusionModel(timesteps=100, training_objective='hybrid')
  66    >>> model.fit(X, n_steps=1000)
  67    """
  68
  69    def __init__(
  70        self,
  71        timesteps: int = 1000,
  72        beta_start: float = 0.0001,
  73        beta_end: float = 0.02,
  74        schedule: Literal["linear", "cosine"] = "linear",
  75        model: Optional[BaseEstimator] = None,
  76        use_pca: bool = False,
  77        pca_components: int = 50,
  78        variance_type: Literal[
  79            "fixed_small", "fixed_large", "learned"
  80        ] = "fixed_small",
  81        random_state: Optional[int] = None,
  82        batch_size: int = 32,
  83        training_objective: Literal["noise", "mmd", "hybrid"] = "noise",
  84        mmd_samples_per_step: int = 10,
  85        mmd_kernel: Literal["rbf", "imq", "linear"] = "rbf",
  86        mmd_bandwidth: Union[float, str] = "auto",
  87    ):
  88        self.timesteps = timesteps
  89        self.beta_start = beta_start
  90        self.beta_end = beta_end
  91        self.schedule = schedule
  92        self.model = model
  93        self.use_pca = use_pca
  94        self.pca_components = pca_components
  95        self.variance_type = variance_type
  96        self.random_state = random_state
  97        self.batch_size = batch_size
  98        self.training_objective = training_objective
  99        self.mmd_samples_per_step = mmd_samples_per_step
 100        self.mmd_kernel = mmd_kernel
 101        self.mmd_bandwidth = mmd_bandwidth
 102
 103        # Initialize random state generator
 104        self._rng = np.random.RandomState(random_state)
 105
 106        # Input validation
 107        self._validate_parameters()
 108        self._init_noise_schedule()
 109
 110    def _validate_parameters(self) -> None:
 111        """Validate input parameters with comprehensive checks"""
 112        if self.beta_start >= self.beta_end:
 113            raise ValueError("beta_start must be less than beta_end")
 114        if self.timesteps <= 0:
 115            raise ValueError("timesteps must be positive")
 116        if self.batch_size <= 0:
 117            raise ValueError("batch_size must be positive")
 118        if self.schedule not in ["linear", "cosine"]:
 119            raise ValueError("schedule must be 'linear' or 'cosine'")
 120        if self.variance_type not in ["fixed_small", "fixed_large", "learned"]:
 121            raise ValueError(
 122                "variance_type must be 'fixed_small', 'fixed_large', or 'learned'"
 123            )
 124        if self.training_objective not in ["noise", "mmd", "hybrid"]:
 125            raise ValueError(
 126                "training_objective must be 'noise', 'mmd', or 'hybrid'"
 127            )
 128        if self.mmd_kernel not in ["rbf", "imq", "linear"]:
 129            raise ValueError("mmd_kernel must be 'rbf', 'imq', or 'linear'")
 130        if self.mmd_samples_per_step <= 0:
 131            raise ValueError("mmd_samples_per_step must be positive")
 132
 133    def _init_noise_schedule(self) -> None:
 134        """Initialize forward diffusion noise schedule with numerical stability"""
 135        if self.schedule == "linear":
 136            self.betas = np.linspace(
 137                self.beta_start, self.beta_end, self.timesteps
 138            )
 139        elif self.schedule == "cosine":
 140            s = 0.008
 141            steps = np.arange(self.timesteps + 1, dtype=np.float64)
 142            alphas_bar = (
 143                np.cos(((steps / self.timesteps) + s) / (1 + s) * np.pi * 0.5)
 144                ** 2
 145            )
 146            alphas_bar = alphas_bar / alphas_bar[0]
 147            self.betas = np.clip(
 148                1 - (alphas_bar[1:] / alphas_bar[:-1]), 0, 0.999
 149            )
 150
 151        self.alphas = 1.0 - self.betas
 152        # Numerical stability: clamp away from 0
 153        self.alphas_cumprod = np.clip(np.cumprod(self.alphas), 1e-8, 1.0)
 154        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
 155        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
 156
 157        # Compute posterior variance (beta_tilde) for true posterior
 158        alphas_cumprod_prev = np.concatenate([[1.0], self.alphas_cumprod[:-1]])
 159        self.posterior_variance = (
 160            self.betas
 161            * (1.0 - alphas_cumprod_prev)
 162            / (1.0 - self.alphas_cumprod)
 163        )
 164
 165        # Compute posterior mean coefficients
 166        self.posterior_mean_coef1 = (
 167            self.betas
 168            * np.sqrt(alphas_cumprod_prev)
 169            / (1.0 - self.alphas_cumprod)
 170        )
 171        self.posterior_mean_coef2 = (
 172            (1.0 - alphas_cumprod_prev)
 173            * np.sqrt(self.alphas)
 174            / (1.0 - self.alphas_cumprod)
 175        )
 176
 177    def _validate_data(self, X: np.ndarray) -> np.ndarray:
 178        """Validate input data with comprehensive checks"""
 179        X = np.asarray(X)
 180        if X.ndim != 2:
 181            raise ValueError(f"Expected 2D array, got {X.ndim}D array instead")
 182        if np.any(np.isnan(X)) or np.any(np.isinf(X)):
 183            raise ValueError("Input contains NaN or infinite values")
 184        if X.std(axis=0).min() == 0:
 185            warnings.warn("Some features have zero variance")
 186        return X
 187
 188    def _validate_timestep(self, t: Union[int, np.ndarray]) -> None:
 189        """Validate timestep bounds"""
 190        t_array = np.atleast_1d(t)
 191        if np.any(t_array < 0) or np.any(t_array >= self.timesteps):
 192            raise ValueError(
 193                f"Timestep values out of range [0, {self.timesteps})"
 194            )
 195
 196    def forward_diffusion(
 197        self, x0: np.ndarray, t: np.ndarray, noise: Optional[np.ndarray] = None
 198    ) -> Tuple[np.ndarray, np.ndarray]:
 199        """
 200        Forward diffusion: q(x_t | x_0) = N(sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I)
 201
 202        Parameters
 203        ----------
 204        x0 : ndarray of shape (n_samples, n_features)
 205            Original data samples
 206        t : ndarray of shape (n_samples,)
 207            Timesteps (0 to timesteps-1)
 208        noise : ndarray, optional
 209            Noise to add (generated if None)
 210
 211        Returns
 212        -------
 213        xt : ndarray
 214            Noised data
 215        noise : ndarray
 216            Noise that was added
 217        """
 218        x0 = self._validate_data(x0)
 219        self._validate_timestep(t)
 220
 221        if noise is None:
 222            noise = self._rng.randn(*x0.shape)
 223
 224        # Vectorized computation for batch processing
 225        sqrt_alpha = self.sqrt_alphas_cumprod[t][:, np.newaxis]
 226        sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t][
 227            :, np.newaxis
 228        ]
 229
 230        xt = sqrt_alpha * x0 + sqrt_one_minus_alpha * noise
 231        return xt, noise
 232
 233    def _sample_true_posterior(
 234        self, x0: np.ndarray, xt: np.ndarray, t: int, n_samples: int = 1
 235    ) -> np.ndarray:
 236        """
 237        Sample from true posterior q(x_{t-1} | x_t, x_0)
 238
 239        This is the ground truth denoising distribution that we want to match.
 240
 241        Parameters
 242        ----------
 243        x0 : ndarray of shape (batch_size, n_features)
 244            Clean data
 245        xt : ndarray of shape (batch_size, n_features)
 246            Noisy data at timestep t
 247        t : int
 248            Current timestep
 249        n_samples : int
 250            Number of samples to draw per data point
 251
 252        Returns
 253        -------
 254        samples : ndarray of shape (batch_size * n_samples, n_features)
 255            Samples from true posterior
 256        """
 257        if t == 0:
 258            # At t=0, posterior is deterministic: x_{-1} doesn't exist, return x_0
 259            return np.repeat(x0, n_samples, axis=0)
 260
 261        batch_size = x0.shape[0]
 262
 263        # Posterior mean: μ̃_t = coef1 * x_0 + coef2 * x_t
 264        coef1 = self.posterior_mean_coef1[t]
 265        coef2 = self.posterior_mean_coef2[t]
 266        posterior_mean = coef1 * x0 + coef2 * xt
 267
 268        # Posterior variance: σ̃_t²
 269        posterior_var = self.posterior_variance[t]
 270
 271        # Sample from N(μ̃_t, σ̃_t² I)
 272        samples = []
 273        for _ in range(n_samples):
 274            noise = self._rng.randn(*x0.shape) * np.sqrt(posterior_var)
 275            sample = posterior_mean + noise
 276            samples.append(sample)
 277
 278        return np.vstack(samples)
 279
 280    def _sample_learned_transition(
 281        self, xt: np.ndarray, t: int, n_samples: int = 1
 282    ) -> np.ndarray:
 283        """
 284        Sample from learned transition p_θ(x_{t-1} | x_t)
 285
 286        Parameters
 287        ----------
 288        xt : ndarray of shape (batch_size, n_features)
 289            Noisy data at timestep t
 290        t : int
 291            Current timestep
 292        n_samples : int
 293            Number of samples to draw per data point
 294
 295        Returns
 296        -------
 297        samples : ndarray of shape (batch_size * n_samples, n_features)
 298            Samples from learned transition
 299        """
 300        batch_size = xt.shape[0]
 301        t_array = np.full(batch_size, t)
 302
 303        # Predict noise using learned model
 304        features = self._create_features(xt, t_array)
 305        pred_noise = self.model_.predict(features).reshape(xt.shape)
 306
 307        # Compute predicted mean
 308        alpha = self.alphas[t]
 309        alpha_bar = self.alphas_cumprod[t]
 310        beta = self.betas[t]
 311
 312        coef1 = 1.0 / np.sqrt(alpha)
 313        coef2 = beta / np.sqrt(1.0 - alpha_bar)
 314        predicted_mean = coef1 * (xt - coef2 * pred_noise)
 315
 316        # Get variance
 317        if self.variance_type == "fixed_small":
 318            variance = self.posterior_variance[t]
 319        elif self.variance_type == "fixed_large":
 320            variance = beta
 321        else:  # learned - for now use fixed_small
 322            variance = self.posterior_variance[t]
 323
 324        # Sample from N(μ_θ, σ² I)
 325        if t == 0:
 326            # Deterministic at final step
 327            return np.repeat(predicted_mean, n_samples, axis=0)
 328
 329        samples = []
 330        for _ in range(n_samples):
 331            noise = self._rng.randn(*xt.shape) * np.sqrt(variance)
 332            sample = predicted_mean + noise
 333            samples.append(sample)
 334
 335        return np.vstack(samples)
 336
 337    def _compute_kernel(
 338        self, X: np.ndarray, Y: np.ndarray, gamma: Optional[float] = None
 339    ) -> np.ndarray:
 340        """
 341        Compute kernel matrix K(X, Y) using specified kernel
 342
 343        Parameters
 344        ----------
 345        X : ndarray of shape (n, d)
 346        Y : ndarray of shape (m, d)
 347        gamma : float, optional
 348            Kernel bandwidth
 349
 350        Returns
 351        -------
 352        K : ndarray of shape (n, m)
 353            Kernel matrix
 354        """
 355        if self.mmd_kernel == "linear":
 356            return X @ Y.T
 357
 358        # Compute squared distances
 359        try:
 360            from scipy.spatial.distance import cdist
 361
 362            sq_dist = cdist(X, Y, metric="sqeuclidean")
 363        except ImportError:
 364            sq_dist = np.sum((X[:, None, :] - Y[None, :, :]) ** 2, axis=2)
 365
 366        if gamma is None:
 367            # Use median heuristic
 368            if len(X) > 100:
 369                sample_idx = self._rng.choice(len(X), 100, replace=False)
 370                sq_dist_sample = sq_dist[sample_idx][:, :100]
 371            else:
 372                sq_dist_sample = sq_dist
 373            gamma = 1.0 / np.median(sq_dist_sample[sq_dist_sample > 0])
 374
 375        if self.mmd_kernel == "rbf":
 376            return np.exp(-gamma * sq_dist)
 377        elif self.mmd_kernel == "imq":
 378            # Inverse multiquadric: (||x-y||² + c²)^(-β)
 379            c = 1.0
 380            beta = 0.5
 381            return (sq_dist + c**2) ** (-beta)
 382        else:
 383            raise ValueError(f"Unknown kernel: {self.mmd_kernel}")
 384
 385    def _compute_mmd_unbiased(
 386        self, X: np.ndarray, Y: np.ndarray, gamma: Optional[float] = None
 387    ) -> float:
 388        """
 389        Compute unbiased MMD² estimator
 390
 391        MMD²(P, Q) = E[k(u,u')] + E[k(v,v')] - 2E[k(u,v)]
 392
 393        Uses unbiased estimator that excludes diagonal terms.
 394        """
 395        n, m = len(X), len(Y)
 396
 397        # Compute kernel matrices
 398        K_XX = self._compute_kernel(X, X, gamma)
 399        K_YY = self._compute_kernel(Y, Y, gamma)
 400        K_XY = self._compute_kernel(X, Y, gamma)
 401
 402        # Unbiased estimator: exclude diagonal
 403        if n > 1:
 404            K_XX_unbiased = (K_XX.sum() - np.trace(K_XX)) / (n * (n - 1))
 405        else:
 406            K_XX_unbiased = 0
 407
 408        if m > 1:
 409            K_YY_unbiased = (K_YY.sum() - np.trace(K_YY)) / (m * (m - 1))
 410        else:
 411            K_YY_unbiased = 0
 412
 413        K_XY_mean = K_XY.mean()
 414
 415        mmd_sq = K_XX_unbiased + K_YY_unbiased - 2 * K_XY_mean
 416
 417        return max(0.0, mmd_sq)  # Ensure non-negative
 418
 419    def _create_features(self, x: np.ndarray, t: np.ndarray) -> np.ndarray:
 420        """
 421        Create features for the regressor with multi-frequency temporal encoding
 422        """
 423        # Normalize timestep
 424        t_norm = t / self.timesteps
 425
 426        # Multi-frequency positional encoding
 427        frequencies = [1.0, 2.0, 4.0]
 428        t_encoded_parts = []
 429
 430        for freq in frequencies:
 431            t_encoded_parts.extend(
 432                [
 433                    np.sin(2 * np.pi * freq * t_norm),
 434                    np.cos(2 * np.pi * freq * t_norm),
 435                ]
 436            )
 437
 438        t_encoded_parts.append(t_norm)
 439        t_encoded = np.column_stack(t_encoded_parts)
 440
 441        # Concatenate spatial and temporal features
 442        features = np.hstack([x, t_encoded])
 443        return features
 444
 445    def fit(
 446        self,
 447        X: np.ndarray,
 448        y: Optional[np.ndarray] = None,
 449        n_steps: int = 1000,
 450    ) -> "DiffusionModel":
 451        """
 452        Train the reverse diffusion model using specified objective
 453
 454        Parameters
 455        ----------
 456        X : ndarray of shape (n_samples, n_features)
 457            Training data (will be normalized internally)
 458        y : ignored
 459            Not used, present for API consistency
 460        n_steps : int, default=1000
 461            Number of training iterations
 462
 463        Returns
 464        -------
 465        self : DiffusionModel
 466            Fitted estimator
 467        """
 468        X = self._validate_data(X)
 469        n_samples, n_features = X.shape
 470
 471        # Store normalization params
 472        self.X_mean_ = X.mean(axis=0)
 473        self.X_std_ = X.std(axis=0) + 1e-8
 474        X_norm = (X - self.X_mean_) / self.X_std_
 475
 476        # Optional PCA
 477        if self.use_pca and n_features > 100:
 478            self.pca_ = PCA(
 479                n_components=min(self.pca_components, n_features),
 480                random_state=self.random_state,
 481            )
 482            X_norm = self.pca_.fit_transform(X_norm)
 483            n_features = X_norm.shape[1]
 484        else:
 485            self.pca_ = None
 486
 487        self.n_features_ = n_features
 488
 489        # Initialize model
 490        if self.model is None:
 491            self.model_ = Ridge(alpha=1.0, random_state=self.random_state)
 492        else:
 493            self.model_ = clone(self.model)
 494
 495        # Train based on objective
 496        if self.training_objective == "noise":
 497            self._fit_noise_prediction(X_norm, n_samples, n_steps)
 498        elif self.training_objective == "mmd":
 499            self._fit_mmd_matching(X_norm, n_samples, n_steps)
 500        elif self.training_objective == "hybrid":
 501            self._fit_hybrid(X_norm, n_samples, n_steps)
 502
 503        self.is_fitted_ = True
 504        return self
 505
 506    def _fit_noise_prediction(
 507        self, X_norm: np.ndarray, n_samples: int, n_steps: int
 508    ) -> None:
 509        """Traditional DDPM training: predict noise with MSE loss"""
 510        n_batches = max(1, n_steps // self.batch_size)
 511
 512        X_train_batches: List[np.ndarray] = []
 513        y_train_batches: List[np.ndarray] = []
 514
 515        for batch_idx in range(n_batches):
 516            indices = self._rng.randint(n_samples, size=self.batch_size)
 517            t_batch = self._rng.randint(self.timesteps, size=self.batch_size)
 518
 519            x0_batch = X_norm[indices]
 520            noise_batch = self._rng.randn(self.batch_size, self.n_features_)
 521
 522            xt_batch, _ = self.forward_diffusion(x0_batch, t_batch, noise_batch)
 523            features = self._create_features(xt_batch, t_batch)
 524
 525            X_train_batches.append(features)
 526            y_train_batches.append(noise_batch)
 527
 528        X_train = np.vstack(X_train_batches)
 529        y_train = np.vstack(y_train_batches)
 530
 531        self.model_.fit(X_train, y_train)
 532
 533    def _fit_mmd_matching(
 534        self, X_norm: np.ndarray, n_samples: int, n_steps: int
 535    ) -> None:
 536        """
 537        MMD-based training: directly minimize distribution mismatch
 538
 539        This implements the algorithm from the theoretical document:
 540        1. Sample (x_0, t, x_t)
 541        2. Draw samples from true posterior q(x_{t-1}|x_t,x_0)
 542        3. Draw samples from learned transition p_θ(x_{t-1}|x_t)
 543        4. Compute MMD² between these sample sets
 544        5. Update θ to minimize MMD
 545        """
 546        print(
 547            f"Training with MMD objective (kernel={self.mmd_kernel}, samples={self.mmd_samples_per_step})"
 548        )
 549
 550        # Determine gamma for kernel
 551        if self.mmd_bandwidth == "auto":
 552            gamma = None  # Will be computed adaptively
 553        else:
 554            gamma = float(self.mmd_bandwidth)
 555
 556        # For sklearn models without iterative updates, we need to:
 557        # 1. Generate training data with MMD-weighted importance
 558        # 2. Fit the model on this data
 559
 560        # We'll use an iterative refinement approach:
 561        # Start with noise prediction, then refine with MMD
 562        print("Phase 1: Initial noise prediction training...")
 563        self._fit_noise_prediction(X_norm, n_samples, n_steps // 2)
 564
 565        print("Phase 2: MMD refinement...")
 566        # Generate MMD-weighted training data
 567        X_train_batches: List[np.ndarray] = []
 568        y_train_batches: List[np.ndarray] = []
 569        weights_batches: List[np.ndarray] = []
 570
 571        n_batches = max(1, (n_steps // 2) // self.batch_size)
 572
 573        for batch_idx in range(n_batches):
 574            if batch_idx % 10 == 0:
 575                print(f"  Batch {batch_idx}/{n_batches}")
 576
 577            indices = self._rng.randint(n_samples, size=self.batch_size)
 578            t_batch = self._rng.randint(1, self.timesteps, size=self.batch_size)
 579
 580            x0_batch = X_norm[indices]
 581            noise_batch = self._rng.randn(self.batch_size, self.n_features_)
 582            xt_batch, _ = self.forward_diffusion(x0_batch, t_batch, noise_batch)
 583
 584            # For each timestep, compute MMD between true and learned transitions
 585            batch_weights = []
 586            for i in range(self.batch_size):
 587                t = t_batch[i]
 588                x0_i = x0_batch[i: i + 1]
 589                xt_i = xt_batch[i: i + 1]
 590
 591                # Sample from both distributions
 592                true_samples = self._sample_true_posterior(
 593                    x0_i, xt_i, t, n_samples=self.mmd_samples_per_step
 594                )
 595                learned_samples = self._sample_learned_transition(
 596                    xt_i, t, n_samples=self.mmd_samples_per_step
 597                )
 598
 599                # Compute MMD
 600                mmd = self._compute_mmd_unbiased(
 601                    true_samples, learned_samples, gamma
 602                )
 603
 604                # Weight by MMD (higher mismatch = higher weight)
 605                weight = 1.0 + mmd
 606                batch_weights.append(weight)
 607
 608            features = self._create_features(xt_batch, t_batch)
 609            X_train_batches.append(features)
 610            y_train_batches.append(noise_batch)
 611            weights_batches.append(np.array(batch_weights))
 612
 613        # Fit with weighted samples
 614        X_train = np.vstack(X_train_batches)
 615        y_train = np.vstack(y_train_batches)
 616        weights = np.concatenate(weights_batches)
 617
 618        # Normalize weights
 619        weights = weights / weights.sum() * len(weights)
 620
 621        print("  Final weighted fit...")
 622        if (
 623            hasattr(self.model_, "fit")
 624            and "sample_weight" in self.model_.fit.__code__.co_varnames
 625        ):
 626            self.model_.fit(X_train, y_train, sample_weight=weights)
 627        else:
 628            # Fallback: repeat samples according to weights
 629            weights_int = np.round(weights).astype(int)
 630            indices_weighted = np.repeat(np.arange(len(X_train)), weights_int)
 631            self.model_.fit(
 632                X_train[indices_weighted], y_train[indices_weighted]
 633            )
 634
 635        print("MMD training complete!")
 636
 637    def _fit_hybrid(
 638        self, X_norm: np.ndarray, n_samples: int, n_steps: int
 639    ) -> None:
 640        """Hybrid training: combine noise prediction and MMD objectives"""
 641        # Split training: 60% noise prediction, 40% MMD refinement
 642        n_noise_steps = int(0.6 * n_steps)
 643        n_mmd_steps = n_steps - n_noise_steps
 644
 645        print(
 646            f"Hybrid training: {n_noise_steps} noise steps + {n_mmd_steps} MMD steps"
 647        )
 648        self._fit_noise_prediction(X_norm, n_samples, n_noise_steps)
 649
 650        # Continue with MMD refinement
 651        self.training_objective = "mmd"  # Temporarily switch
 652        self._fit_mmd_matching(X_norm, n_samples, n_mmd_steps)
 653        self.training_objective = "hybrid"  # Restore
 654
 655    def sample(
 656        self,
 657        n_samples: int = 1,
 658        return_trajectory: bool = False,
 659        ddim: bool = False,
 660        ddim_steps: int = 50,
 661    ) -> np.ndarray:
 662        """
 663        Generate samples via reverse diffusion
 664
 665        Parameters
 666        ----------
 667        n_samples : int, default=1
 668            Number of samples to generate
 669        return_trajectory : bool, default=False
 670            If True, return all intermediate denoising steps
 671        ddim : bool, default=False
 672            Use DDIM for faster deterministic sampling
 673        ddim_steps : int, default=50
 674            Number of DDIM steps (only used if ddim=True)
 675
 676        Returns
 677        -------
 678        samples : ndarray
 679            Generated samples (denormalized to original scale)
 680        """
 681        check_is_fitted(self, ["model_", "X_mean_", "X_std_", "n_features_"])
 682
 683        if ddim:
 684            return self._sample_ddim(n_samples, ddim_steps, return_trajectory)
 685
 686        # Standard DDPM sampling
 687        x = self._rng.randn(n_samples, self.n_features_)
 688        trajectory = [x.copy()] if return_trajectory else None
 689
 690        for t in reversed(range(self.timesteps)):
 691            t_array = np.full(n_samples, t)
 692            features = self._create_features(x, t_array)
 693
 694            pred_noise = self.model_.predict(features).reshape(x.shape)
 695
 696            alpha = self.alphas[t]
 697            alpha_bar = self.alphas_cumprod[t]
 698            beta = self.betas[t]
 699
 700            coef1 = 1.0 / np.sqrt(alpha)
 701            coef2 = beta / np.sqrt(1.0 - alpha_bar)
 702            mean = coef1 * (x - coef2 * pred_noise)
 703
 704            if t > 0:
 705                if self.variance_type == "fixed_small":
 706                    variance = self.posterior_variance[t]
 707                else:
 708                    variance = beta
 709                noise = np.sqrt(variance) * self._rng.randn(*x.shape)
 710                x = mean + noise
 711            else:
 712                x = mean
 713
 714            if return_trajectory:
 715                trajectory.append(x.copy())
 716
 717        return self._postprocess_samples(x, trajectory, return_trajectory)
 718
 719    def _sample_ddim(
 720        self, n_samples: int, steps: int, return_trajectory: bool
 721    ) -> np.ndarray:
 722        """DDIM sampling for faster generation"""
 723        skip = max(1, self.timesteps // steps)
 724        timesteps = np.arange(0, self.timesteps, skip)[::-1]
 725
 726        x = self._rng.randn(n_samples, self.n_features_)
 727        trajectory = [x.copy()] if return_trajectory else None
 728
 729        for i, t in enumerate(timesteps):
 730            t_array = np.full(n_samples, t)
 731            features = self._create_features(x, t_array)
 732            pred_noise = self.model_.predict(features).reshape(x.shape)
 733
 734            alpha_bar = self.alphas_cumprod[t]
 735            alpha_bar_prev = (
 736                self.alphas_cumprod[timesteps[i + 1]]
 737                if i < len(timesteps) - 1
 738                else 1.0
 739            )
 740
 741            pred_x0 = (x - np.sqrt(1 - alpha_bar) * pred_noise) / np.sqrt(
 742                alpha_bar
 743            )
 744            dir_xt = np.sqrt(1 - alpha_bar_prev) * pred_noise
 745            x = np.sqrt(alpha_bar_prev) * pred_x0 + dir_xt
 746
 747            if return_trajectory:
 748                trajectory.append(x.copy())
 749
 750        return self._postprocess_samples(x, trajectory, return_trajectory)
 751
 752    def _postprocess_samples(
 753        self,
 754        x: np.ndarray,
 755        trajectory: Optional[List[np.ndarray]],
 756        return_trajectory: bool,
 757    ) -> np.ndarray:
 758        """Apply inverse transforms to generated samples"""
 759        if return_trajectory and trajectory is not None:
 760            processed_trajectory = []
 761            for step in trajectory:
 762                if self.pca_ is not None:
 763                    step = self.pca_.inverse_transform(step)
 764                step = step * self.X_std_ + self.X_mean_
 765                processed_trajectory.append(step)
 766            return np.array(processed_trajectory)
 767
 768        if self.pca_ is not None:
 769            x = self.pca_.inverse_transform(x)
 770
 771        x = x * self.X_std_ + self.X_mean_
 772        return x
 773
 774    def predict(self, n_samples: int = 1, **kwargs) -> np.ndarray:
 775        """Sklearn-style predict method (alias for sample)"""
 776        return self.sample(n_samples=n_samples, **kwargs)
 777
 778    def score(self, X: np.ndarray, metric: str = "mmd") -> float:
 779        """Compute negative reconstruction error (higher is better)"""
 780        return -self.reconstruction_error(X, metric=metric)
 781
 782    def reconstruction_error(
 783        self,
 784        X: np.ndarray,
 785        metric: str = "mmd",
 786        n_samples: Optional[int] = None,
 787        gamma: Union[str, float] = "auto",
 788    ) -> float:
 789        """
 790        Compute distributional reconstruction error
 791
 792        Parameters
 793        ----------
 794        X : ndarray
 795            Original data distribution
 796        metric : str, default='mmd'
 797            Error metric: 'mmd' or 'energy'
 798        n_samples : int, optional
 799            Number of samples to generate
 800        gamma : str or float, default='auto'
 801            RBF bandwidth for MMD
 802
 803        Returns
 804        -------
 805        error : float
 806            Reconstruction error (lower is better)
 807        """
 808        check_is_fitted(self, ["model_", "X_mean_", "X_std_"])
 809
 810        if n_samples is None:
 811            n_samples = min(1000, len(X))
 812
 813        X_reconstructed = self.sample(n_samples)
 814
 815        if metric == "mmd":
 816            return self._compute_mmd(X, X_reconstructed, gamma)
 817        elif metric == "energy":
 818            return self._compute_energy_distance(X, X_reconstructed)
 819        else:
 820            raise ValueError(f"Metric must be 'mmd' or 'energy', got: {metric}")
 821
 822    def _compute_mmd(
 823        self, X: np.ndarray, Y: np.ndarray, gamma: Union[str, float] = "auto"
 824    ) -> float:
 825        """Compute MMD with adaptive bandwidth (for evaluation)"""
 826        max_samples = 1000
 827        if len(X) > max_samples:
 828            indices = self._rng.choice(len(X), max_samples, replace=False)
 829            X = X[indices]
 830        if len(Y) > max_samples:
 831            indices = self._rng.choice(len(Y), max_samples, replace=False)
 832            Y = Y[indices]
 833
 834        if gamma == "auto":
 835            n_bandwidth = min(100, len(X), len(Y))
 836            X_bw = X[:n_bandwidth]
 837            Y_bw = Y[:n_bandwidth]
 838            combined = np.vstack([X_bw, Y_bw])
 839
 840            from scipy.spatial.distance import pdist
 841
 842            squared_dists = pdist(combined, metric="sqeuclidean")
 843            gamma = 1.0 / np.median(squared_dists)
 844
 845        try:
 846            from scipy.spatial.distance import cdist
 847
 848            XX = cdist(X, X, metric="sqeuclidean")
 849            YY = cdist(Y, Y, metric="sqeuclidean")
 850            XY = cdist(X, Y, metric="sqeuclidean")
 851
 852            K_XX = np.exp(-gamma * XX).mean()
 853            K_YY = np.exp(-gamma * YY).mean()
 854            K_XY = np.exp(-gamma * XY).mean()
 855        except ImportError:
 856            chunk_size = 100
 857
 858            def chunked_kernel_mean(A, B):
 859                total = 0.0
 860                count = 0
 861                for i in range(0, len(A), chunk_size):
 862                    A_chunk = A[i: i + chunk_size]
 863                    for j in range(0, len(B), chunk_size):
 864                        B_chunk = B[j: j + chunk_size]
 865                        sq_dist = np.sum(
 866                            (A_chunk[:, None, :] - B_chunk[None, :, :]) ** 2,
 867                            axis=2,
 868                        )
 869                        total += np.exp(-gamma * sq_dist).sum()
 870                        count += len(A_chunk) * len(B_chunk)
 871                return total / count
 872
 873            K_XX = chunked_kernel_mean(X, X)
 874            K_YY = chunked_kernel_mean(Y, Y)
 875            K_XY = chunked_kernel_mean(X, Y)
 876
 877        mmd = K_XX + K_YY - 2 * K_XY
 878        return max(0, mmd)
 879
 880    def _compute_energy_distance(self, X: np.ndarray, Y: np.ndarray) -> float:
 881        """Compute energy distance"""
 882        try:
 883            from scipy.spatial.distance import pdist, cdist
 884
 885            max_samples = 500
 886            if len(X) > max_samples:
 887                indices = self._rng.choice(len(X), max_samples, replace=False)
 888                X = X[indices]
 889            if len(Y) > max_samples:
 890                indices = self._rng.choice(len(Y), max_samples, replace=False)
 891                Y = Y[indices]
 892
 893            xy_dist = cdist(X, Y, metric="euclidean").mean()
 894            xx_dist = pdist(X, metric="euclidean").mean() if len(X) > 1 else 0
 895            yy_dist = pdist(Y, metric="euclidean").mean() if len(Y) > 1 else 0
 896
 897            energy = 2 * xy_dist - xx_dist - yy_dist
 898            return max(0, energy)
 899        except ImportError:
 900            max_samples = 500
 901            if len(X) > max_samples:
 902                indices = self._rng.choice(len(X), max_samples, replace=False)
 903                X = X[indices]
 904            if len(Y) > max_samples:
 905                indices = self._rng.choice(len(Y), max_samples, replace=False)
 906                Y = Y[indices]
 907
 908            def chunked_pairwise_dist(A, B, chunk_size=100):
 909                total = 0.0
 910                count = 0
 911                for i in range(0, len(A), chunk_size):
 912                    A_chunk = A[i: i + chunk_size]
 913                    for j in range(0, len(B), chunk_size):
 914                        B_chunk = B[j: j + chunk_size]
 915                        dist = np.sqrt(
 916                            np.sum(
 917                                (A_chunk[:, None, :] - B_chunk[None, :, :])
 918                                ** 2,
 919                                axis=2,
 920                            )
 921                        )
 922                        total += dist.sum()
 923                        count += len(A_chunk) * len(B_chunk)
 924                return total / count
 925
 926            xy_dist = chunked_pairwise_dist(X, Y)
 927
 928            n, m = len(X), len(Y)
 929            if n > 1:
 930                xx_dist = chunked_pairwise_dist(X, X)
 931                xx_dist = (xx_dist * n * n) / (n * (n - 1))
 932            else:
 933                xx_dist = 0
 934
 935            if m > 1:
 936                yy_dist = chunked_pairwise_dist(Y, Y)
 937                yy_dist = (yy_dist * m * m) / (m * (m - 1))
 938            else:
 939                yy_dist = 0
 940
 941            energy = 2 * xy_dist - xx_dist - yy_dist
 942            return max(0, energy)
 943
 944    def optimize_hyperparameters(
 945        self,
 946        X: np.ndarray,
 947        n_calls: int = 20,
 948        metric: str = "mmd",
 949        cv_splits: int = 3,
 950    ) -> dict:
 951        """Bayesian optimization of hyperparameters with cross-validation"""
 952        try:
 953            from skopt import gp_minimize
 954            from skopt.space import Real, Integer, Categorical
 955        except ImportError:
 956            raise ImportError(
 957                "scikit-optimize not installed. Install with: pip install scikit-optimize"
 958            )
 959
 960        space = [
 961            Integer(100, 2000, name="timesteps"),
 962            Real(1e-5, 1e-3, name="beta_start", prior="log-uniform"),
 963            Real(0.01, 0.05, name="beta_end"),
 964            Categorical(["linear", "cosine"], name="schedule"),
 965            Real(0.1, 10.0, name="ridge_alpha", prior="log-uniform"),
 966        ]
 967
 968        def objective(params):
 969            try:
 970                model = DiffusionModel(
 971                    timesteps=params[0],
 972                    beta_start=params[1],
 973                    beta_end=params[2],
 974                    schedule=params[3],
 975                    model=Ridge(
 976                        alpha=params[4], random_state=self.random_state
 977                    ),
 978                    random_state=self.random_state,
 979                    batch_size=self.batch_size,
 980                    use_pca=self.use_pca,
 981                    pca_components=self.pca_components,
 982                    training_objective=self.training_objective,
 983                )
 984
 985                errors = []
 986                for _ in range(cv_splits):
 987                    model.fit(X, n_steps=500)
 988                    X_eval = X[: min(100, len(X))]
 989                    error = model.reconstruction_error(
 990                        X_eval, metric=metric, n_samples=100
 991                    )
 992                    errors.append(error)
 993
 994                return np.mean(errors)
 995            except Exception as e:
 996                warnings.warn(f"Optimization iteration failed: {e}")
 997                return 1e6
 998
 999        result = gp_minimize(
1000            objective,
1001            space,
1002            n_calls=n_calls,
1003            random_state=self.random_state,
1004            verbose=False,
1005        )
1006
1007        best_params = {
1008            "timesteps": result.x[0],
1009            "beta_start": result.x[1],
1010            "beta_end": result.x[2],
1011            "schedule": result.x[3],
1012            "model": Ridge(alpha=result.x[4], random_state=self.random_state),
1013        }
1014
1015        self.set_params(**best_params)
1016        self._init_noise_schedule()
1017        self.fit(X, n_steps=2000)
1018
1019        print(f"\n✅ Optimization complete!")
1020        print(f"Best {metric.upper()} error: {result.fun:.6f}")
1021        print(f"Best parameters:")
1022        for key, value in best_params.items():
1023            if key != "model":
1024                print(f"  {key}: {value}")
1025            else:
1026                print(f"  model: Ridge(alpha={result.x[4]:.4f})")
1027
1028        return best_params