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

  36class DistroSimulator:
  37    def __init__(
  38        self,
  39        kernel="rbf",
  40        backend="numpy",
  41        n_clusters=5,
  42        clustering_method="kmeans",
  43        kde_kernel="gaussian",
  44        random_state=None,
  45        conformalize=False,
  46        residual_sampling="bootstrap",
  47        block_size=None,
  48        gmm_components=3,
  49        use_rff="auto",
  50        rff_components="auto",
  51        rff_gamma=None,
  52        kernel_approximation="rff",
  53        force_rff_threshold=1000,
  54    ):
  55        """
  56        Initialize the multivariate data generator.
  57
  58        Parameters:
  59        -----------
  60        kernel : str, default='rbf'
  61            Kernel type for KernelRidge regression
  62        backend : str, default='numpy'
  63            Backend for distance calculations ('numpy', 'gpu', 'tpu')
  64        n_clusters : int, default=5
  65            Number of clusters for stratified splitting
  66        clustering_method : str, default='kmeans'
  67            Clustering method for stratification ('kmeans' or 'gmm')
  68        random_state : int, default=None
  69            Random seed for reproducibility
  70        conformalize : bool
  71            Use split conformal prediction or not
  72        residual_sampling : str, default='bootstrap'
  73            Method for sampling residuals ('bootstrap', 'kde', 'gmm', 'block-bootstrap', 'me-bootstrap').
  74            Where 'me-bootstrap' refers to Maximum Entropy Bootstrap.
  75        block_size : int, default=None
  76            Block size for block bootstrap (if applicable)
  77        gmm_components : int, default=3
  78            Number of components for GMM sampling
  79        use_rff : bool or 'auto', default='auto'
  80            Whether to use kernel approximation. 'auto' enables for large datasets
  81        rff_components : int or 'auto', default='auto'
  82            Number of approximation components. 'auto' chooses based on data size
  83        rff_gamma : float, default=None
  84            Gamma parameter for approximation. If None, will be tuned.
  85        kernel_approximation : str, default='rff'
  86            Approximation method ('rff' or 'nystroem')
  87        force_rff_threshold : int, default=1000
  88            Auto-enable RFF when n_samples exceeds this threshold
  89        """
  90        self.kernel = kernel
  91        self.backend = backend
  92        self.n_clusters = n_clusters
  93        self.clustering_method = clustering_method
  94        self.random_state = random_state
  95        self.conformalize = conformalize
  96        self.residual_sampling = residual_sampling
  97        self.block_size = block_size
  98        self.gmm_components = gmm_components
  99        self.use_rff = use_rff
 100        self.rff_components = rff_components
 101        self.rff_gamma = rff_gamma
 102        self.kernel_approximation = kernel_approximation
 103        self.force_rff_threshold = force_rff_threshold
 104        self.kde_kernel = kde_kernel
 105
 106        if random_state is not None:
 107            np.random.seed(random_state)
 108            if JAX_AVAILABLE:
 109                key = jax.random.PRNGKey(random_state)
 110        # Validate sampling method
 111        valid_sampling_methods = [
 112            "bootstrap",
 113            "kde",
 114            "gmm",
 115            "block-bootstrap",
 116            "me-bootstrap",
 117        ]
 118        if residual_sampling not in valid_sampling_methods:
 119            raise ValueError(
 120                f"residual_sampling must be one of {valid_sampling_methods}"
 121            )
 122        # Validate approximation method
 123        valid_approximations = ["rff", "nystroem"]
 124        if kernel_approximation not in valid_approximations:
 125            raise ValueError(
 126                f"kernel_approximation must be one of {valid_approximations}"
 127            )
 128        # Initialize JAX if using GPU/TPU backend
 129        if backend in ["gpu", "tpu"] and JAX_AVAILABLE:
 130            self._setup_jax_backend()
 131        elif backend in ["gpu", "tpu"] and not JAX_AVAILABLE:
 132            print("JAX not available. Falling back to NumPy backend.")
 133            self.backend = "numpy"
 134        # Initialize attributes that will be set during fitting
 135        self.model = None
 136        self.residuals_ = None
 137        self.X_dist = None
 138        self.is_fitted = False
 139        self.best_params_ = None
 140        self.best_score_ = None
 141        self.cluster_labels_ = None
 142        self.cluster_model_ = None
 143        self.kde_model_ = None
 144        self.gmm_model_ = None
 145        self.scaler_ = None
 146        self.actual_rff_components_ = None
 147        self.actual_use_rff_ = None
 148
 149    def _setup_jax_backend(self):
 150        """Setup JAX backend for GPU/TPU acceleration."""
 151        if not JAX_AVAILABLE:
 152            raise ImportError("JAX is required for GPU/TPU backend")
 153
 154        # JIT compiled distance functions
 155        @jit
 156        def pairwise_sq_dists_jax(X1, X2):
 157            X1_sq = jnp.sum(X1**2, axis=1)[:, jnp.newaxis]
 158            X2_sq = jnp.sum(X2**2, axis=1)[jnp.newaxis, :]
 159            return X1_sq + X2_sq - 2 * X1 @ X2.T
 160
 161        @jit
 162        def cdist_jax(X1, X2):
 163            return vmap(
 164                lambda x: vmap(lambda y: jnp.sqrt(jnp.sum((x - y) ** 2)))(X2)
 165            )(X1)
 166
 167        self._pairwise_sq_dists_jax = pairwise_sq_dists_jax
 168        self._cdist_jax = cdist_jax
 169
 170    def _determine_components(self, n_samples):
 171        """Automatically determine optimal number of components."""
 172        if self.rff_components == "auto":
 173            # Optimized heuristic based on performance results
 174            if n_samples < 500:
 175                return min(50, n_samples)
 176            elif n_samples < 2000:
 177                return min(100, n_samples // 2)
 178            elif n_samples < 5000:
 179                return min(150, n_samples // 3)
 180            elif n_samples < 10000:
 181                return min(200, n_samples // 4)
 182            else:
 183                return min(300, n_samples // 5)
 184        else:
 185            return self.rff_components
 186
 187    def _create_model(self, gamma, alpha, use_rff=None):
 188        """Create the appropriate model based on RFF setting."""
 189        if use_rff is None:
 190            use_rff = self.actual_use_rff_
 191
 192        if use_rff:
 193            # Use kernel approximation with Ridge regression
 194            if self.rff_gamma is not None:
 195                effective_gamma = self.rff_gamma
 196            else:
 197                effective_gamma = gamma
 198            # Determine number of components
 199            n_components = self.actual_rff_components_
 200
 201            if self.kernel_approximation == "rff":
 202                approximator = RBFSampler(
 203                    gamma=effective_gamma,
 204                    n_components=n_components,
 205                    random_state=self.random_state,
 206                )
 207            else:  # nystroem
 208                approximator = Nystroem(
 209                    kernel="rbf",
 210                    gamma=effective_gamma,
 211                    n_components=n_components,
 212                    random_state=self.random_state,
 213                )
 214            # Create pipeline with scaling, approximation, and Ridge
 215            return Pipeline(
 216                [
 217                    ("scaler", StandardScaler()),
 218                    ("approx", approximator),
 219                    ("ridge", Ridge(alpha=alpha)),
 220                ]
 221            )
 222        # Standard KernelRidge
 223        return KernelRidge(kernel=self.kernel, gamma=gamma, alpha=alpha)
 224
 225    def _fit_residual_sampler(self, **kwargs):
 226        """Fit the chosen residual sampling model."""
 227        if self.residuals_ is None or len(self.residuals_) == 0:
 228            raise ValueError("No residuals available for fitting sampler")
 229
 230        if self.residual_sampling == "kde":
 231            kernel_bandwidths = {"bandwidth": np.logspace(-6, 6, 150)}
 232            grid = GridSearchCV(
 233                KernelDensity(kernel=self.kde_kernel, **kwargs),
 234                param_grid=kernel_bandwidths,
 235            )
 236            grid.fit(self.residuals_)
 237            self.kde_model_ = grid.best_estimator_
 238            self.kde_model_.fit(self.residuals_)
 239
 240        elif self.residual_sampling == "gmm":
 241            self.gmm_model_ = GaussianMixture(
 242                n_components=min(self.gmm_components, len(self.residuals_)),
 243                random_state=self.random_state,
 244                covariance_type="full",
 245            )
 246            self.gmm_model_.fit(self.residuals_)
 247
 248    def _sample_residuals(self, num_samples):
 249        """Sample residuals using the chosen method."""
 250        if self.residuals_ is None:
 251            raise ValueError("No residuals available for sampling")
 252
 253        if self.residual_sampling == "bootstrap":
 254            # Original bootstrap method
 255            n = self.residuals_.shape[0]
 256            idx = np.random.choice(n, num_samples, replace=True)
 257            return self.residuals_[idx]
 258
 259        elif self.residual_sampling == "kde":
 260            # Kernel Density Estimation sampling
 261            if self.kde_model_ is None:
 262                raise ValueError(
 263                    "KDE model not fitted. Call _fit_residual_sampler first."
 264                )
 265            # Sample from KDE
 266            return self.kde_model_.sample(num_samples)
 267
 268        elif self.residual_sampling == "gmm":
 269            # Gaussian Mixture Model sampling
 270            if self.gmm_model_ is None:
 271                raise ValueError(
 272                    "GMM model not fitted. Call _fit_residual_sampler first."
 273                )
 274            # Sample from GMM
 275            return self.gmm_model_.sample(num_samples)[0]
 276
 277        elif self.residual_sampling == "me-bootstrap":
 278            meb = MaximumEntropyBootstrap(random_state=self.random_state)
 279            # If residuals are shorter than num_samples, repeat or tile them
 280            residuals = self.residuals_.flatten()
 281            if residuals.shape[0] < num_samples:
 282                # Repeat residuals to reach num_samples
 283                repeats = int(np.ceil(num_samples / residuals.shape[0]))
 284                residuals = np.tile(residuals, repeats)[:num_samples]
 285            else:
 286                residuals = residuals[:num_samples]
 287            meb.fit(residuals)
 288            return meb.sample(1)[:, 0].reshape(-1, 1)
 289
 290        elif self.residual_sampling == "block-bootstrap":
 291            # Block Bootstrap sampling
 292            return bootstrap(
 293                self.residuals_, num_samples, block_size=self.block_size
 294            )
 295
 296        else:
 297            # Should not reach here due to validation in __init__
 298            raise ValueError(
 299                f"Unknown sampling method: {self.residual_sampling}"
 300            )
 301
 302    def _pairwise_sq_dists(self, X1, X2):
 303        """Compute pairwise squared Euclidean distances."""
 304        if self.backend in ["gpu", "tpu"] and JAX_AVAILABLE:
 305            X1_jax = jnp.array(X1)
 306            X2_jax = jnp.array(X2)
 307            result = self._pairwise_sq_dists_jax(X1_jax, X2_jax)
 308            return np.array(result)
 309        else:
 310            X1 = np.atleast_2d(X1)
 311            X2 = np.atleast_2d(X2)
 312            return (
 313                np.sum(X1**2, axis=1)[:, np.newaxis]
 314                + np.sum(X2**2, axis=1)[np.newaxis, :]
 315                - 2 * X1 @ X2.T
 316            )
 317
 318    def _compute_clusters(self, Y):
 319        """Compute cluster labels for stratified splitting."""
 320        if self.clustering_method == "kmeans":
 321            self.cluster_model_ = KMeans(
 322                n_clusters=self.n_clusters,
 323                random_state=self.random_state,
 324                n_init=10,
 325            )
 326        elif self.clustering_method == "gmm":
 327            self.cluster_model_ = GaussianMixture(
 328                n_components=self.n_clusters, random_state=self.random_state
 329            )
 330        else:
 331            raise ValueError("clustering_method must be 'kmeans' or 'gmm'")
 332        self.cluster_model_.fit(Y)
 333        return self.cluster_model_.predict(Y)
 334
 335    def _train_test_split(self, Y, n_train, sequential: bool = False):
 336        """Create train-test split. Stratified by clusters or sequential if specified."""
 337        try:
 338            n_samples = len(Y)
 339        except Exception:
 340            n_samples = Y.shape[0]
 341
 342        if sequential:
 343            # --- Sequential split (no shuffling, preserves temporal order)
 344            train_idx = np.arange(n_train)
 345            test_idx = np.arange(n_train, n_samples)
 346            return train_idx, test_idx
 347
 348        # --- Stratified split (default)
 349        self.cluster_labels_ = self._compute_clusters(Y)
 350        return train_test_split(
 351            np.arange(n_samples),
 352            train_size=n_train,
 353            stratify=self.cluster_labels_,
 354            random_state=self.random_state,
 355        )
 356
 357    def _mmd(self, u, v, kernel_sigma=1):
 358        """Maximum Mean Discrepancy between two distributions."""
 359        if u.ndim == 1:
 360            u = u.reshape(-1, 1)
 361        if v.ndim == 1:
 362            v = v.reshape(-1, 1)
 363
 364        def kmat(A, B):
 365            return np.exp(
 366                -self._pairwise_sq_dists(A, B) / (2 * kernel_sigma**2)
 367            )
 368
 369        return (
 370            np.mean(kmat(u, u)) + np.mean(kmat(v, v)) - 2 * np.mean(kmat(u, v))
 371        )
 372
 373    def _custom_energy_distance(self, u, v):
 374        """Energy distance between two distributions."""
 375        if u.ndim == 1:
 376            u = u.reshape(-1, 1)
 377        if v.ndim == 1:
 378            v = v.reshape(-1, 1)
 379
 380        n, d = u.shape
 381        m = v.shape[0]
 382
 383        if self.backend in ["gpu", "tpu"] and JAX_AVAILABLE:
 384            # JAX implementation
 385            u_jax = jnp.array(u)
 386            v_jax = jnp.array(v)
 387            dist_xx = self._cdist_jax(u_jax, u_jax)
 388            dist_yy = self._cdist_jax(v_jax, v_jax)
 389            dist_xy = self._cdist_jax(u_jax, v_jax)
 390            term1 = 2 * jnp.sum(dist_xy) / (n * m)
 391            term2 = jnp.sum(dist_xx) / (n * n)
 392            term3 = jnp.sum(dist_yy) / (m * m)
 393            return float(term1 - term2 - term3)
 394        else:
 395            # NumPy implementation
 396            dist_xx = cdist(u, u, metric="euclidean")
 397            dist_yy = cdist(v, v, metric="euclidean")
 398            dist_xy = cdist(u, v, metric="euclidean")
 399            term1 = 2 * np.sum(dist_xy) / (n * m)
 400            term2 = np.sum(dist_xx) / (n * n)
 401            term3 = np.sum(dist_yy) / (m * m)
 402            return term1 - term2 - term3
 403
 404    def _generate_pseudo(self, num_samples):
 405        """Generate synthetic data using the fitted model and residuals."""
 406        if not self.is_fitted:
 407            raise ValueError("Model not fitted. Call fit() first.")
 408        X_new = self.X_dist[:num_samples]
 409        # Handle prediction based on model type
 410        if self.actual_use_rff_:
 411            # For RFF pipeline
 412            preds = self.model.predict(X_new)
 413        else:
 414            # For standard KernelRidge
 415            preds = self.model.predict(X_new)
 416        if preds.ndim == 1:
 417            preds = preds.reshape(-1, 1)
 418        # Sample residuals using the chosen method
 419        return preds + self._sample_residuals(preds.shape[0])
 420
 421    def fit(self, Y, n_train=None, metric="energy", n_trials=50, **kwargs):
 422        """
 423        Fit the data generator to match the distribution of Y.
 424
 425        Parameters:
 426        -----------
 427        Y : array-like, shape (n_samples, n_features)
 428            Target multivariate data to emulate
 429        n_train : int, default=None
 430            Number of training samples (default: n_samples // 2)
 431        metric : str, default='energy'
 432            Distance metric for optimization ('energy', 'mmd', or 'wasserstein')
 433        n_trials : int, default=50
 434            Number of Optuna optimization trials
 435        **kwargs : dict
 436            Additional arguments for Optuna optimization
 437
 438        Returns:
 439        --------
 440        self : object
 441            Returns self
 442        """
 443        if Y.ndim == 1:
 444            Y = Y.reshape(-1, 1)
 445
 446        n, d = Y.shape
 447        self.n_features_ = d
 448        # Determine whether to use RFF
 449        if self.use_rff == "auto":
 450            self.actual_use_rff_ = n >= self.force_rff_threshold
 451        else:
 452            self.actual_use_rff_ = self.use_rff
 453        # Auto-enable RFF for large datasets with component determination
 454        if self.actual_use_rff_:
 455            self.actual_rff_components_ = self._determine_components(n)
 456            if self.use_rff == "auto":
 457                print(
 458                    f"Large dataset detected (n={n}). Auto-enabling {self.kernel_approximation.upper()} for scalability."
 459                )
 460
 461        if n_train is None:
 462            n_train = n // 2
 463        # Store the input distribution function
 464        self.X_dist = np.random.normal(0, 1, (n, d))
 465        # Create stratified train-test split
 466        if self.residual_sampling in ("block-bootstrap", "me-bootstrap"):
 467            train_idx, test_idx = self._train_test_split(
 468                Y, n_train, sequential=True
 469            )
 470        else:
 471            train_idx, test_idx = self._train_test_split(
 472                Y, n_train, sequential=False
 473            )
 474        Y_train = Y[train_idx]
 475        Y_test = Y[test_idx]
 476        X_train = self.X_dist[:n_train]
 477
 478        if self.conformalize:
 479
 480            def objective(trial):
 481                sigma = trial.suggest_float("sigma", 0.01, 10, log=True)
 482                lambd = trial.suggest_float("lambd", 1e-5, 1, log=True)
 483                gamma = 1 / (2 * sigma**2)
 484                # Determine proper training set size (50% of training data)
 485                n_proper_train = int(0.5 * len(Y_train))
 486                # Use stratified split for proper training and calibration sets
 487                proper_train_idx, calib_idx = self._train_test_split(
 488                    Y_train, n_proper_train
 489                )
 490                # Split the data
 491                X_proper_train = X_train[proper_train_idx]
 492                Y_proper_train = Y_train[proper_train_idx]
 493                X_calib = X_train[calib_idx]
 494                Y_calib = Y_train[calib_idx]
 495                # Standardize the response (Y) using proper training set statistics
 496                if not hasattr(self, "y_scaler_"):
 497                    self.y_scaler_ = StandardScaler()
 498                    Y_proper_train_scaled = self.y_scaler_.fit_transform(
 499                        Y_proper_train
 500                    )
 501                else:
 502                    Y_proper_train_scaled = self.y_scaler_.transform(
 503                        Y_proper_train
 504                    )
 505                # Create model with current parameters and fit on standardized proper training set
 506                model = self._create_model(gamma, lambd)
 507                model.fit(X_proper_train, Y_proper_train_scaled)
 508                # Get predictions on calibration set and transform back to original scale
 509                preds_calib_scaled = model.predict(X_calib)
 510                if preds_calib_scaled.ndim == 1:
 511                    preds_calib_scaled = preds_calib_scaled.reshape(-1, 1)
 512                # Transform predictions back to original scale
 513                preds_calib = self.y_scaler_.inverse_transform(
 514                    preds_calib_scaled
 515                )
 516                # Calculate residuals on calibration set in original scale
 517                res_calib = Y_calib - preds_calib
 518                # Standardize residuals using calibration set statistics
 519                if res_calib.ndim == 1:
 520                    res_mean = np.mean(res_calib)
 521                    res_std = np.std(res_calib, ddof=1)
 522                    # Avoid division by zero
 523                    res_std = res_std if res_std > 1e-10 else 1.0
 524                    res_calib_standardized = (res_calib - res_mean) / res_std
 525                else:
 526                    res_mean = np.mean(res_calib, axis=0)
 527                    res_std = np.std(res_calib, axis=0, ddof=1)
 528                    # Avoid division by zero
 529                    res_std = np.where(res_std > 1e-10, res_std, 1.0)
 530                    res_calib_standardized = (res_calib - res_mean) / res_std
 531
 532                # Store the calibrated standardized residuals for use in the generation method
 533                calibrated_residuals = (
 534                    res_calib_standardized * res_std + res_mean
 535                )
 536
 537                # Use the existing method to generate pseudo samples with conformal prediction
 538                Y_sim = self._generate_pseudo_with_model(
 539                    model, calibrated_residuals, len(Y_test)
 540                )
 541
 542                # Calculate distance metric
 543                if metric == "energy":
 544                    dist_val = self._custom_energy_distance(Y_test, Y_sim)
 545                elif metric == "mmd":
 546                    dist_val = self._mmd(Y_test, Y_sim)
 547                elif metric == "wasserstein" and d == 1:
 548                    dist_val = stats.wasserstein_distance(
 549                        Y_test.flatten(), Y_sim.flatten()
 550                    )
 551                else:
 552                    raise ValueError("Invalid metric for dimension")
 553
 554                return dist_val
 555
 556        else:
 557
 558            def objective(trial):
 559                sigma = trial.suggest_float("sigma", 0.01, 10, log=True)
 560                lambd = trial.suggest_float("lambd", 1e-5, 1, log=True)
 561                gamma = 1 / (2 * sigma**2)
 562                # Create model with current parameters
 563                model = self._create_model(gamma, lambd)
 564                model.fit(X_train, Y_train)
 565                preds_train = model.predict(X_train)
 566                if preds_train.ndim == 1:
 567                    preds_train = preds_train.reshape(-1, 1)
 568                res = Y_train - preds_train
 569                Y_sim = self._generate_pseudo_with_model(
 570                    model, res, len(Y_test)
 571                )
 572                if metric == "energy":
 573                    dist_val = self._custom_energy_distance(Y_test, Y_sim)
 574                elif metric == "mmd":
 575                    dist_val = self._mmd(Y_test, Y_sim)
 576                elif metric == "wasserstein" and d == 1:
 577                    dist_val = stats.wasserstein_distance(
 578                        Y_test.flatten(), Y_sim.flatten()
 579                    )
 580                else:
 581                    raise ValueError("Invalid metric for dimension")
 582                return dist_val
 583
 584        # Optimize hyperparameters
 585        study = optuna.create_study(direction="minimize")
 586        study.optimize(objective, n_trials=n_trials, **kwargs)
 587        # Store best parameters and fit final model
 588        self.best_params_ = study.best_params
 589        self.best_score_ = study.best_value
 590        sigma = self.best_params_["sigma"]
 591        lambd = self.best_params_["lambd"]
 592        gamma = 1 / (2 * sigma**2)
 593        # Fit final model with best parameters
 594        self.model = self._create_model(gamma, lambd)
 595        self.model.fit(X_train, Y_train)
 596        # Compute residuals
 597        preds_train = self.model.predict(X_train)
 598        if preds_train.ndim == 1:
 599            preds_train = preds_train.reshape(-1, 1)
 600        self.residuals_ = Y_train - preds_train
 601        # Fit the residual sampler
 602        self._fit_residual_sampler()
 603        self.is_fitted = True
 604        # Print final configuration
 605        if self.actual_use_rff_:
 606            print(
 607                f"  Using {self.kernel_approximation.upper()} with {self.actual_rff_components_} components"
 608            )
 609        else:
 610            print(f"  Using standard kernel method")
 611        return self
 612
 613    def _generate_pseudo_with_model(self, model, residuals, num_samples):
 614        """Helper method to generate data with a specific model."""
 615        X_new = self.X_dist[:num_samples]
 616
 617        # Handle prediction based on model type
 618        if hasattr(model, "named_steps"):
 619            # Pipeline (RFF or Nystroem)
 620            preds = model.predict(X_new)
 621        else:
 622            # Standard model
 623            preds = model.predict(X_new)
 624
 625        if preds.ndim == 1:
 626            preds = preds.reshape(-1, 1)
 627
 628        # Temporarily store original state
 629        original_residuals = self.residuals_
 630        original_kde = self.kde_model_
 631        original_gmm = self.gmm_model_
 632
 633        # Set residuals for this model
 634        self.residuals_ = residuals
 635
 636        # Fit sampler with the new residuals
 637        self._fit_residual_sampler()
 638
 639        # Sample residuals
 640        sampled_residuals = self._sample_residuals(num_samples)
 641
 642        # Restore original state
 643        self.residuals_ = original_residuals
 644        self.kde_model_ = original_kde
 645        self.gmm_model_ = original_gmm
 646
 647        return preds + sampled_residuals
 648
 649    def sample(self, n_samples=1):
 650        """
 651        Generate synthetic samples.
 652
 653        Parameters:
 654        -----------
 655        n_samples : int, default=1
 656            Number of samples to generate
 657
 658        Returns:
 659        --------
 660        Y_sim : array, shape (n_samples, n_features)
 661            Generated synthetic data
 662        """
 663        if not self.is_fitted:
 664            raise ValueError("Model not fitted. Call fit() first.")
 665        return self._generate_pseudo(n_samples)
 666
 667    def compare_approximation_methods(self, Y, n_train=None, n_trials=20):
 668        """
 669        Compare different kernel approximation methods.
 670
 671        Parameters:
 672        -----------
 673        Y : array-like
 674            Target data
 675        n_train : int, default=None
 676            Number of training samples
 677        n_trials : int, default=20
 678            Number of optimization trials
 679
 680        Returns:
 681        --------
 682        comparison_results : dict
 683            Comparison results
 684        """
 685        if Y.ndim == 1:
 686            Y = Y.reshape(-1, 1)
 687
 688        print("Comparing Kernel Approximation Methods...")
 689
 690        # Store original settings
 691        original_use_rff = self.use_rff
 692        original_approximation = self.kernel_approximation
 693        original_is_fitted = self.is_fitted
 694
 695        methods = ["rff", "nystroem"]
 696        results = {}
 697
 698        for method in methods:
 699            print(f"\nTesting {method.upper()}...")
 700            self.use_rff = True
 701            self.kernel_approximation = method
 702
 703            start_time = time()
 704            self.fit(Y, n_train=n_train, n_trials=n_trials)
 705            method_time = time() - start_time
 706            method_score = self.best_score_
 707            method_params = self.best_params_
 708
 709            results[method] = {
 710                "time": method_time,
 711                "score": method_score,
 712                "params": method_params,
 713                "components": self.actual_rff_components_,
 714            }
 715
 716        # Test standard method for comparison
 717        print(f"\nTesting Standard Kernel...")
 718        self.use_rff = False
 719        start_time = time()
 720        self.fit(Y, n_train=n_train, n_trials=n_trials)
 721        standard_time = time() - start_time
 722        standard_score = self.best_score_
 723        standard_params = self.best_params_
 724
 725        results["standard"] = {
 726            "time": standard_time,
 727            "score": standard_score,
 728            "params": standard_params,
 729            "components": "N/A",
 730        }
 731
 732        # Restore original settings
 733        self.use_rff = original_use_rff
 734        self.kernel_approximation = original_approximation
 735        self.is_fitted = original_is_fitted
 736
 737        # Print comparison
 738        print("\n" + "=" * 60)
 739        print("KERNEL APPROXIMATION COMPARISON RESULTS")
 740        print("=" * 60)
 741
 742        for method in ["standard"] + methods:
 743            data = results[method]
 744            print(f"\n{method.upper()}:")
 745            print(f"  Time: {data['time']:.2f}s")
 746            print(f"  Score: {data['score']:.6f}")
 747            print(f"  Components: {data['components']}")
 748            if method != "standard":
 749                speedup = standard_time / data["time"]
 750                score_ratio = data["score"] / standard_score
 751                print(f"  Speedup: {speedup:.2f}x")
 752                print(f"  Score Ratio: {score_ratio:.4f}")
 753
 754        return results
 755
 756    def compare_residual_sampling(self, n_samples=1000):
 757        """
 758        Compare different residual sampling methods visually.
 759
 760        Parameters:
 761        -----------
 762        n_samples : int, default=1000
 763            Number of samples to generate for comparison
 764        """
 765        if not self.is_fitted:
 766            raise ValueError("Model not fitted. Call fit() first.")
 767        # Store original sampling method
 768        original_sampling = self.residual_sampling
 769        # Generate samples with different methods
 770        sampling_methods = ["bootstrap", "kde", "gmm"]
 771        samples = {}
 772
 773        for method in sampling_methods:
 774            self.residual_sampling = method
 775            if method == "kde":
 776                self._fit_residual_sampler()
 777            elif method == "gmm":
 778                self._fit_residual_sampler()
 779            samples[method] = self._sample_residuals(n_samples)
 780        # Restore original method
 781        self.residual_sampling = original_sampling
 782        self._fit_residual_sampler()
 783        # Plot comparison
 784        n_dims = self.residuals_.shape[1]
 785        fig, axes = plt.subplots(
 786            n_dims,
 787            len(sampling_methods) + 1,
 788            figsize=(5 * (len(sampling_methods) + 1), 4 * n_dims),
 789        )
 790
 791        if n_dims == 1:
 792            axes = axes.reshape(1, -1)
 793
 794        for dim in range(n_dims):
 795            # Original residuals
 796            axes[dim, 0].hist(
 797                self.residuals_[:, dim], bins=30, alpha=0.7, density=True
 798            )
 799            axes[dim, 0].set_title(f"Original Residuals\nDim {dim+1}")
 800            axes[dim, 0].set_xlabel("Residual Value")
 801            axes[dim, 0].set_ylabel("Density")
 802
 803            # Sampled residuals
 804            for j, method in enumerate(sampling_methods):
 805                col = j + 1
 806                axes[dim, col].hist(
 807                    samples[method][:, dim], bins=30, alpha=0.7, density=True
 808                )
 809                axes[dim, col].set_title(
 810                    f"{method.upper()} Sampling\nDim {dim+1}"
 811                )
 812                axes[dim, col].set_xlabel("Residual Value")
 813                axes[dim, col].set_ylabel("Density")
 814
 815        plt.tight_layout()
 816        plt.show()
 817
 818        return samples
 819
 820    def _perm_test(self, Y_orig, Y_sim, stat_func, n_perm=1000):
 821        """Permutation test for distribution comparison."""
 822        if Y_orig.ndim == 1:
 823            Y_orig = Y_orig.reshape(-1, 1)
 824        if Y_sim.ndim == 1:
 825            Y_sim = Y_sim.reshape(-1, 1)
 826
 827        obs = stat_func(Y_orig, Y_sim)
 828        combined = np.vstack((Y_orig, Y_sim))
 829        n1 = Y_orig.shape[0]
 830        perms = np.zeros(n_perm)
 831
 832        for i in range(n_perm):
 833            idx = np.random.permutation(combined.shape[0])
 834            p1 = combined[idx[:n1]]
 835            p2 = combined[idx[n1:]]
 836            perms[i] = stat_func(p1, p2)
 837
 838        pval = (np.sum(perms >= obs) + 1) / (n_perm + 1)
 839        return obs, pval
 840
 841    def _fisher_z_test(self, r1, r2, n1, n2):
 842        """Fisher z-test for comparing correlation coefficients."""
 843        z1 = np.arctanh(r1)
 844        z2 = np.arctanh(r2)
 845        z = (z1 - z2) / np.sqrt(1 / (n1 - 3) + 1 / (n2 - 3))
 846        p = 2 * (1 - stats.norm.cdf(np.abs(z)))
 847        return z, p
 848
 849    def test_similarity(self, Y_orig, Y_sim, n_perm=1000):
 850        """
 851        Test statistical similarity between original and synthetic data.
 852
 853        Parameters:
 854        -----------
 855        Y_orig : array-like
 856            Original data
 857        Y_sim : array-like
 858            Synthetic data
 859        n_perm : int, default=1000
 860            Number of permutations for permutation tests
 861
 862        Returns:
 863        --------
 864        results : dict
 865            Dictionary containing test results
 866        """
 867        if Y_orig.ndim == 1:
 868            Y_orig = Y_orig.reshape(-1, 1)
 869        if Y_sim.ndim == 1:
 870            Y_sim = Y_sim.reshape(-1, 1)
 871
 872        d = Y_orig.shape[1]
 873        results = {}
 874        # Test 1: Perm with energy
 875        results["energy_perm"] = self._perm_test(
 876            Y_orig, Y_sim, self._custom_energy_distance, n_perm
 877        )
 878        # Test 2: Perm with MMD
 879        results["mmd_perm"] = self._perm_test(
 880            Y_orig, Y_sim, lambda u, v: self._mmd(u, v), n_perm
 881        )
 882
 883        # Test 3: Perm with avg Wasserstein on margins
 884        def avg_wass(u, v):
 885            return np.mean(
 886                [stats.wasserstein_distance(u[:, i], v[:, i]) for i in range(d)]
 887            )
 888
 889        results["avg_wass_perm"] = self._perm_test(
 890            Y_orig, Y_sim, avg_wass, n_perm
 891        )
 892        # Test 4: Min p-value from marginal KS tests
 893        ps_ks = [
 894            stats.ks_2samp(Y_orig[:, i], Y_sim[:, i]).pvalue for i in range(d)
 895        ]
 896        results["min_marginal_ks_p"] = min(ps_ks)
 897        # Test 5: Min p-value from marginal Anderson-Darling tests
 898        ps_ad = [
 899            stats.anderson_ksamp([Y_orig[:, i], Y_sim[:, i]]).significance_level
 900            for i in range(d)
 901        ]
 902        results["min_marginal_ad_p"] = min(ps_ad)
 903        # Test 6: Min p-value from marginal Cramer-von Mises tests
 904        ps_cvm = [
 905            stats.cramervonmises_2samp(Y_orig[:, i], Y_sim[:, i]).pvalue
 906            for i in range(d)
 907        ]
 908        results["min_marginal_cvm_p"] = min(ps_cvm)
 909        # Correlation test: Compare all pairwise correlations
 910        corr_results = {}
 911        pairs = [(i, j) for i in range(d) for j in range(i + 1, d)]
 912        for i, j in pairs:
 913            r_orig = stats.pearsonr(Y_orig[:, i], Y_orig[:, j])[0]
 914            r_sim = stats.pearsonr(Y_sim[:, i], Y_sim[:, j])[0]
 915            z, p = self._fisher_z_test(r_orig, r_sim, len(Y_orig), len(Y_sim))
 916            corr_results[f"corr_dim{i+1}_dim{j+1}"] = (r_orig, r_sim, z, p)
 917        results["corr_tests"] = corr_results
 918        return results
 919
 920    def compare_distributions(self, Y_orig, Y_sim, save_prefix=""):
 921        """
 922        Visual comparison of original and synthetic distributions.
 923
 924        Parameters:
 925        -----------
 926        Y_orig : array-like
 927            Original data
 928        Y_sim : array-like
 929            Synthetic data
 930        save_prefix : str, default=''
 931            Prefix for saving plots
 932        """
 933        if Y_orig.ndim == 1:
 934            Y_orig = Y_orig.reshape(-1, 1)
 935        if Y_sim.ndim == 1:
 936            Y_sim = Y_sim.reshape(-1, 1)
 937
 938        n, d = Y_orig.shape
 939
 940        # Create a figure with subplots for statistical tests
 941        fig, axes = plt.subplots(2, d, figsize=(6 * d, 10))
 942        if d == 1:
 943            axes = axes.reshape(2, 1)
 944
 945        # Statistical test results storage
 946        ks_results = []
 947        ad_results = []
 948
 949        for i in range(d):
 950            # Top row: Histograms with statistical test annotations
 951            ax_hist = axes[0, i]
 952
 953            # Plot histograms
 954            ax_hist.hist(
 955                Y_orig[:, i],
 956                alpha=0.5,
 957                label="Original",
 958                density=True,
 959                bins=20,
 960                color="blue",
 961            )
 962            ax_hist.hist(
 963                Y_sim[:, i],
 964                alpha=0.5,
 965                label="Simulated",
 966                density=True,
 967                bins=20,
 968                color="red",
 969            )
 970
 971            # Perform statistical tests
 972            # Kolmogorov-Smirnov test
 973            ks_stat, ks_pvalue = stats.ks_2samp(Y_orig[:, i], Y_sim[:, i])
 974            ks_results.append((ks_stat, ks_pvalue))
 975
 976            # Anderson-Darling test
 977            ad_result = stats.anderson_ksamp([Y_orig[:, i], Y_sim[:, i]])
 978            ad_stat = ad_result.statistic
 979            ad_critical = ad_result.critical_values
 980            ad_significance = ad_result.significance_level
 981            ad_results.append((ad_stat, ad_significance))
 982
 983            # Add test results to histogram plot
 984            textstr = "\n".join(
 985                (
 986                    f"KS test: p = {ks_pvalue:.4f}",
 987                    f"AD test: p < {ad_significance:.3f}",
 988                    f"AD stat: {ad_stat:.4f}",
 989                )
 990            )
 991            props = dict(boxstyle="round", facecolor="wheat", alpha=0.8)
 992            ax_hist.text(
 993                0.05,
 994                0.95,
 995                textstr,
 996                transform=ax_hist.transAxes,
 997                fontsize=10,
 998                verticalalignment="top",
 999                bbox=props,
1000            )
1001
1002            ax_hist.legend()
1003            ax_hist.set_title(
1004                f"Dimension {i+1} - Histograms with Statistical Tests"
1005            )
1006            ax_hist.set_xlabel("Value")
1007            ax_hist.set_ylabel("Density")
1008
1009            # Bottom row: ECDFs with KS test visualization
1010            ax_ecdf = axes[1, i]
1011
1012            # Compute ECDFs
1013            sorted_orig = np.sort(Y_orig[:, i])
1014            ecdf_orig = np.arange(1, len(sorted_orig) + 1) / len(sorted_orig)
1015            sorted_sim = np.sort(Y_sim[:, i])
1016            ecdf_sim = np.arange(1, len(sorted_sim) + 1) / len(sorted_sim)
1017
1018            # Plot ECDFs
1019            ax_ecdf.step(
1020                sorted_orig,
1021                ecdf_orig,
1022                label="Original",
1023                color="blue",
1024                linewidth=2,
1025            )
1026            ax_ecdf.step(
1027                sorted_sim,
1028                ecdf_sim,
1029                label="Simulated",
1030                color="red",
1031                linewidth=2,
1032            )
1033
1034            # Find the point of maximum difference for KS test
1035            # Combine and sort all values
1036            all_values = np.sort(np.concatenate([sorted_orig, sorted_sim]))
1037            # Compute ECDFs at all points
1038            ecdf_orig_all = np.searchsorted(
1039                sorted_orig, all_values, side="right"
1040            ) / len(sorted_orig)
1041            ecdf_sim_all = np.searchsorted(
1042                sorted_sim, all_values, side="right"
1043            ) / len(sorted_sim)
1044            # Find maximum difference
1045            diff = np.abs(ecdf_orig_all - ecdf_sim_all)
1046            max_idx = np.argmax(diff)
1047            max_x = all_values[max_idx]
1048            max_y1 = ecdf_orig_all[max_idx]
1049            max_y2 = ecdf_sim_all[max_idx]
1050
1051            # Mark the maximum difference point
1052            ax_ecdf.plot(
1053                [max_x, max_x],
1054                [max_y1, max_y2],
1055                "k-",
1056                linewidth=3,
1057                label=f"KS stat: {ks_stat:.4f}",
1058            )
1059            ax_ecdf.plot(max_x, max_y1, "ko", markersize=8)
1060            ax_ecdf.plot(max_x, max_y2, "ko", markersize=8)
1061
1062            ax_ecdf.legend()
1063            ax_ecdf.set_title(f"Dimension {i+1} - ECDFs with KS Statistic")
1064            ax_ecdf.set_xlabel("Value")
1065            ax_ecdf.set_ylabel("ECDF")
1066
1067        plt.tight_layout()
1068        if save_prefix:
1069            plt.savefig(
1070                f"{save_prefix}_statistical_comparison.png",
1071                dpi=300,
1072                bbox_inches="tight",
1073            )
1074        plt.show()
1075
1076        # Print comprehensive test results
1077        print("\n" + "=" * 60)
1078        print("COMPREHENSIVE STATISTICAL TEST RESULTS")
1079        print("=" * 60)
1080
1081        for i in range(d):
1082            ks_stat, ks_pvalue = ks_results[i]
1083            ad_stat, ad_significance = ad_results[i]
1084
1085            print(f"\nDimension {i+1}:")
1086            print(f"  Kolmogorov-Smirnov Test:")
1087            print(f"    Statistic: {ks_stat:.6f}")
1088            print(f"    p-value: {ks_pvalue:.6f}")
1089            print(
1090                f"    Significance: {'Not Significant' if ks_pvalue > 0.05 else 'SIGNIFICANT'}"
1091            )
1092
1093            print(f"  Anderson-Darling Test:")
1094            print(f"    Statistic: {ad_stat:.6f}")
1095            print(f"    Significance level: {ad_significance:.3f}")
1096            print(
1097                f"    Interpretation: {'Distributions differ' if ad_stat > ad_result.critical_values[2] else 'Distributions similar'}"
1098            )
1099
1100        # Create summary plot for all dimensions
1101        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
1102
1103        # KS test p-values across dimensions
1104        ks_pvalues = [result[1] for result in ks_results]
1105        dimensions = list(range(1, d + 1))
1106
1107        bars = ax1.bar(
1108            dimensions,
1109            ks_pvalues,
1110            color=["red" if p < 0.05 else "green" for p in ks_pvalues],
1111        )
1112        ax1.axhline(
1113            y=0.05, color="black", linestyle="--", alpha=0.7, label="α = 0.05"
1114        )
1115        ax1.set_xlabel("Dimension")
1116        ax1.set_ylabel("KS Test p-value")
1117        ax1.set_title("Kolmogorov-Smirnov Test Results\nby Dimension")
1118        ax1.set_xticks(dimensions)
1119        ax1.legend()
1120
1121        # Add value labels on bars
1122        for bar, pvalue in zip(bars, ks_pvalues):
1123            height = bar.get_height()
1124            ax1.text(
1125                bar.get_x() + bar.get_width() / 2.0,
1126                height,
1127                f"{pvalue:.3f}",
1128                ha="center",
1129                va="bottom",
1130            )
1131
1132        # AD test statistics across dimensions
1133        ad_stats = [result[0] for result in ad_results]
1134
1135        bars = ax2.bar(dimensions, ad_stats, color="skyblue")
1136        ax2.set_xlabel("Dimension")
1137        ax2.set_ylabel("AD Test Statistic")
1138        ax2.set_title("Anderson-Darling Test Statistics\nby Dimension")
1139        ax2.set_xticks(dimensions)
1140
1141        # Add value labels on bars
1142        for bar, stat in zip(bars, ad_stats):
1143            height = bar.get_height()
1144            ax2.text(
1145                bar.get_x() + bar.get_width() / 2.0,
1146                height,
1147                f"{stat:.3f}",
1148                ha="center",
1149                va="bottom",
1150            )
1151
1152        plt.tight_layout()
1153        if save_prefix:
1154            plt.savefig(
1155                f"{save_prefix}_test_summary.png", dpi=300, bbox_inches="tight"
1156            )
1157        plt.show()
1158
1159        # Additional: Q-Q plots for each dimension
1160        fig, axes = plt.subplots(1, d, figsize=(5 * d, 5))
1161        if d == 1:
1162            axes = [axes]
1163
1164        for i in range(d):
1165            # Sort data for Q-Q plot
1166            orig_sorted = np.sort(Y_orig[:, i])
1167            sim_sorted = np.sort(Y_sim[:, i])
1168
1169            # Generate theoretical quantiles
1170            n_orig = len(orig_sorted)
1171            n_sim = len(sim_sorted)
1172
1173            # Use smaller set for quantiles to avoid interpolation issues
1174            n_points = min(n_orig, n_sim, 1000)
1175            quantiles = np.linspace(0, 1, n_points)
1176
1177            orig_quantiles = np.quantile(orig_sorted, quantiles)
1178            sim_quantiles = np.quantile(sim_sorted, quantiles)
1179
1180            axes[i].plot(
1181                orig_quantiles, sim_quantiles, "o", alpha=0.6, markersize=3
1182            )
1183            min_val = min(orig_quantiles.min(), sim_quantiles.min())
1184            max_val = max(orig_quantiles.max(), sim_quantiles.max())
1185            axes[i].plot(
1186                [min_val, max_val],
1187                [min_val, max_val],
1188                "r--",
1189                alpha=0.8,
1190                linewidth=2,
1191            )
1192            axes[i].set_xlabel("Original Data Quantiles")
1193            axes[i].set_ylabel("Simulated Data Quantiles")
1194            axes[i].set_title(f"Dimension {i+1} - Q-Q Plot")
1195
1196            # Add correlation coefficient
1197            corr = np.corrcoef(orig_quantiles, sim_quantiles)[0, 1]
1198            axes[i].text(
1199                0.05,
1200                0.95,
1201                f"Corr: {corr:.4f}",
1202                transform=axes[i].transAxes,
1203                bbox=dict(
1204                    boxstyle="round,pad=0.3", facecolor="white", alpha=0.8
1205                ),
1206                verticalalignment="top",
1207            )
1208
1209        plt.tight_layout()
1210        if save_prefix:
1211            plt.savefig(
1212                f"{save_prefix}_qq_plots.png", dpi=300, bbox_inches="tight"
1213            )
1214        plt.show()
1215
1216        return {
1217            "ks_results": ks_results,
1218            "ad_results": ad_results,
1219            "dimensions": d,
1220        }