Metadata-Version: 2.4
Name: bstabdiff
Version: 0.2.1
Summary: BSTabDiff: Block-Subunit Diffusion Priors for High-Dimensional Tabular Data Generation (NeurIPS 2026)
Home-page: https://www.zadidhabib.com/bstabdiff.html
Author: Al Zadid Sultan Bin Habib, Md Younus Ahamed, Prashnna Kumar Gyawali, Gianfranco Doretto, Donald A. Adjeroh
License: MIT
Project-URL: Homepage, https://www.zadidhabib.com/bstabdiff.html
Project-URL: Repository, https://github.com/zadid6pretam/BSTabDiff
Project-URL: Issues, https://github.com/zadid6pretam/BSTabDiff/issues
Project-URL: Paper, https://openreview.net/forum?id=ciPCMXX1aj
Keywords: tabular-data,high-dimensional-data,HDLSS,synthetic-data,generative-modeling,diffusion,normalizing-flow,feature-ordering,GO-BS,missing-data,mixed-type-data,neurips
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy<3.0,>=2.2.6
Requires-Dist: pandas<3.0,>=2.2.3
Requires-Dist: scikit-learn<2.0,>=1.7.2
Requires-Dist: scipy<2.0,>=1.15.3
Requires-Dist: matplotlib<4.0,>=3.10.8
Requires-Dist: tqdm<5.0,>=4.67.3
Requires-Dist: torch<3.0,>=2.5.1
Requires-Dist: kmeans_gpu==0.0.5
Requires-Dist: optuna<5.0,>=4.8.0
Requires-Dist: psutil<8.0,>=7.0.0
Requires-Dist: codecarbon<4.0,>=3.2.6
Dynamic: license-file

# BSTabDiff: Block-Subunit Diffusion Priors for High-Dimensional Tabular Data Generation (NeurIPS 2026)


<!-- ========================================================= -->
<!-- BSTabDiff: Framework and Implementation -->
<!-- ========================================================= -->

![Python](https://img.shields.io/badge/Python-3.8%2B-blue)
![License](https://img.shields.io/badge/License-MIT-green)
![Task](https://img.shields.io/badge/Task-Tabular%20Data%20Synthesis-orange)
![Model](https://img.shields.io/badge/Model-BSTabDiff-blueviolet)
![Architecture](https://img.shields.io/badge/Architecture-Block--Subunit%20Generator-informational)
![Latent Prior](https://img.shields.io/badge/Latent%20Prior-Diffusion%20%2F%20Flow-purple)
![Regime](https://img.shields.io/badge/Regime-HDLSS%20%26%20Beyond-critical)
![Evaluation](https://img.shields.io/badge/Evaluation-18%20Datasets%20%7C%2010%20Baselines-teal)

<!-- ========================================================= -->
<!-- NeurIPS 2026: Main Conference -->
<!-- ========================================================= -->

[![NeurIPS](https://img.shields.io/badge/NeurIPS-2026%20Main%20Conference-8A2BE2)](https://neurips.cc/virtual/2026/loc/atlanta/poster/151140)
[![Status](https://img.shields.io/badge/Status-Accepted-brightgreen)](https://neurips.cc/virtual/2026/loc/atlanta/poster/151140)
[![NeurIPS Poster](https://img.shields.io/badge/NeurIPS-Poster%20Page-6A5ACD)](https://neurips.cc/virtual/2026/loc/atlanta/poster/151140)
[![OpenReview](https://img.shields.io/badge/OpenReview-Main%20Paper-B31B1B)](https://openreview.net/forum?id=ciPCMXX1aj)

<!-- ========================================================= -->
<!-- ICLR 2026: Earlier DeLTa Workshop Version -->
<!-- ========================================================= -->

[![Earlier Version](https://img.shields.io/badge/Earlier%20Version-ICLR%202026%20DeLTa-blue)](https://iclr.cc/virtual/2026/10017199)
[![ICLR Workshop Paper](https://img.shields.io/badge/ICLR-Workshop%20Paper-4169E1)](https://iclr.cc/virtual/2026/10017199)
[![Workshop OpenReview](https://img.shields.io/badge/OpenReview-Workshop%20Version-red)](https://openreview.net/forum?id=RKNDy0KhGT)
[![DeLTa Workshop](https://img.shields.io/badge/DeLTa%202026-Workshop%20Website-1f6feb)](https://delta-workshop.github.io/DeLTa2026/)



<p align="center">
  <img src="./BSTabDiffArchi.png" alt="BSTabDiff Architecture" width="900">
</p>

BSTabDiff is a block-subunit generative framework designed primarily for **High-Dimensional Low-Sample Size (HDLSS) tabular data synthesis**. Instead of learning dense dependence directly in the original `m`-dimensional feature space, BSTabDiff organizes features into **`M` dependency-aware latent blocks, where `M << m`**, optionally using **Graph-guided Ordering with Block Segmentation (GO-BS)** or its feature-clustered variant **GO-BS-FC** to construct coherent feature neighborhoods. Global dependence is learned in the compact block-latent space using diffusion or normalizing-flow priors, while the full-dimensional table is reconstructed through block-conditioned copula decoding, feature-wise marginal transformations, and explicit missingness modeling. This design makes BSTabDiff particularly suitable for omics-style and other HDLSS problems, where direct high-dimensional density estimation is often ill-conditioned. Across HDLSS benchmarks, BSTabDiff produces high-fidelity synthetic data with strong downstream utility, often approaching real-data performance, while additional experiments indicate that the block-subunit design can also generalize beyond the strict HDLSS regime.

## Citation

Al Zadid Sultan Bin Habib, Md Younus Ahamed, Prashnna Kumar Gyawali, Gianfranco Doretto, and Donald A. Adjeroh.  
**“BSTabDiff: Block-Subunit Diffusion Priors for High-Dimensional Tabular Data Generation.”**  
In *Advances in Neural Information Processing Systems (NeurIPS)*, 2026.


BibTeX:
```bibtex
@inproceedings{habib2026bstabdiff,
  title     = {BSTabDiff: Block-Subunit Diffusion Priors for High-Dimensional Tabular Data Generation},
  author    = {Habib, Al Zadid Sultan Bin and Ahamed, Md Younus and Gyawali, Prashnna Kumar and Doretto, Gianfranco and Adjeroh, Donald A.},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}
```

### Publication Venues

**NeurIPS 2026 - Main Conference (Poster)**  
[NeurIPS Poster Page](https://neurips.cc/virtual/2026/loc/atlanta/poster/151140) | [OpenReview](https://openreview.net/forum?id=ciPCMXX1aj)

**Earlier Version - ICLR 2026 DeLTa Workshop**  
[ICLR Workshop Paper](https://iclr.cc/virtual/2026/10017199) | [OpenReview](https://openreview.net/forum?id=RKNDy0KhGT) | [Workshop Website](https://delta-workshop.github.io/DeLTa2026/)


## Files and Repository Structure

The repository is organized into three main components:

```text
BSTabDiff/
├── bstabdiff/
├── Experimental Notebooks/
├── Smoke Tests/
├── BSTabDiffArchi.png
├── README.md
├── requirements.txt
├── pyproject.toml
├── setup.cfg
└── LICENSE
```

### 1. Model Package: `bstabdiff/`

The `bstabdiff/` folder contains the core BSTabDiff implementation and the public Python API.

- **`__init__.py`**  
  Package initializer exposing the main BSTabDiff classes, utilities, and high-level API.

- **`bstabdiff_gobs.py`**  
  Extended BSTabDiff implementation with **Graph-guided Ordering with Block Segmentation (GO-BS)** and **GO-BS-FC**, including dependency-aware feature ordering, block construction, boundary refinement, and integration with the BSTabDiff generator.

- **`estimator.py`**  
  High-level estimator interface providing the user-facing `BSTabDiff` API for model initialization, `fit(...)`, synthetic sampling, and configuration of the underlying generator and ordering variants.

For most users, the recommended interface is:

```python
from bstabdiff import BSTabDiff
```

### 2. `Experimental Notebooks/`

The **`Experimental Notebooks/`** folder contains cleaned versions of most of the notebooks used during the experiments reported in the previous ICLR Workshop paper, NeurIPS paper and NeurIPS rebuttal. These notebooks are provided primarily for **verification of the reported results and displayed outputs**, including dataset-specific experiments, downstream evaluation, ablation studies, fidelity analysis, GO-BS/GO-BS-FC experiments, resource analysis, and additional rebuttal experiments. Because the codebase evolved during development, some notebooks reference earlier or intermediate model/package names that were used at different stages of the project. They should therefore be viewed mainly as experimental records for reproducing and inspecting the reported results, while the current `bstabdiff/` package should be used for new experiments and applications. As stated during the NeurIPS rebuttal, we provide the **current model package together with experimental notebooks containing displayed outputs/results** to make the reported findings easier to inspect and verify. We have cleaned and organized most of the notebooks before public release. A small number of auxiliary development notebooks are not yet included because they require additional cleanup and documentation before they can be released in a reasonably readable form. See the folder here: **[Experimental Notebooks](https://github.com/zadid6pretam/BSTabDiff/tree/main/Experimental%20Notebooks)**

### 3. `Smoke Tests/`

The **`Smoke Tests/`** folder contains lightweight notebooks intended to verify installation, API behavior, training, sampling, tuning, and different supported data configurations without requiring the full experimental pipeline.

- **`BSTabDiff_SmokeTest.ipynb`**  
  Basic end-to-end smoke test for fitting BSTabDiff and generating synthetic samples.

- **`BSTabDiff_Tuning_Smoke_Test.ipynb`**  
  Lightweight demonstration of hyperparameter tuning and subsequent BSTabDiff training using selected parameters.

- **`BSTabDiff_MixedData_SmokeTest.ipynb`**  
  Demonstrates BSTabDiff on mixed-type data containing both continuous and categorical variables, including explicit missingness.

- **`BSTabDiff_PIP_Install_Check.ipynb`**  
  Google Colab smoke test demonstrating installation of BSTabDiff through `pip`, importing the installed package, fitting the model, and generating synthetic data in a clean notebook environment.

- **`BSTabDiff_CarbonTracking.ipynb`**  
  Demonstrates BSTabDiff energy and CO₂e tracking with CodeCarbon on the Colon dataset, together with a separate 5×5 cross-validated TSTR evaluation using Logistic Regression.

See the folder here: **[Smoke Tests](https://github.com/zadid6pretam/BSTabDiff/tree/main/Smoke%20Tests)**

### Architecture Figure

- **`BSTabDiffArchi.png`**  
  High-level architecture figure illustrating the BSTabDiff framework, including block-latent generation, diffusion/flow priors, feature decoding, missingness modeling, and reconstruction in the observed feature space.

### Other Top-Level Files

- **`requirements.txt`** - Main Python dependencies used by the repository.
- **`pyproject.toml`** - Build-system and package metadata.
- **`setup.cfg`** - Package configuration and installation metadata.
- **`LICENSE`** - Repository license.
- **`README.md`** - Project overview, installation instructions, usage examples, reproducibility information, and links to experimental resources.

### Tested Environment

- numpy==2.2.6
- pandas==2.3.3
- scikit-learn==1.7.2
- scipy==1.15.3
- matplotlib==3.10.8
- tqdm==4.67.3
- torch==2.5.1
- catboost==1.2.8
- tabpfn==6.3.1
- kmeans_gpu==0.0.5
- optuna==4.8.0
- psutil==7.0.0
- codecarbon==3.2.6
- jupyter==1.1.1
- ipykernel==7.2.0

## Installation

You can install **BSTabDiff** in several ways depending on your workflow.

---

### Option 1: Clone the Repository (Recommended for Development)

```bash
git clone https://github.com/zadid6pretam/BSTabDiff.git
cd BSTabDiff
pip install -r requirements.txt
pip install -e .
```

### Option 2: Install Directly from GitHub (No Cloning Needed)

```bash
pip install "git+https://github.com/zadid6pretam/BSTabDiff.git"
```

### Option 3: Use a Virtual Environment

```bash
python -m venv bstabdiff-env
source bstabdiff-env/bin/activate  # On Windows: bstabdiff-env\Scripts\activate

git clone https://github.com/zadid6pretam/BSTabDiff.git
cd BSTabDiff
pip install -r requirements.txt
pip install -e .
```

### Option 4: Local Install Without Editable Mode

```bash
git clone https://github.com/zadid6pretam/BSTabDiff.git
cd BSTabDiff
pip install -r requirements.txt
pip install .
```

### Option 5: Install from PyPI (Planned)

```bash
pip install bstabdiff
```

## Practical Configuration Rules

BSTabDiff has two main dataset-dependent structural choices:

1. the number of latent blocks, `M`;
2. whether to use **GO-BS** or **GO-BS-FC** for structure-aware feature ordering and block construction.

The following rules provide practical defaults for new datasets. They are intended as initialization guidelines rather than rigid prescriptions; users performing dataset-specific hyperparameter optimization may tune these choices directly.

### 1. Selecting the Number of Blocks `M`

Let

- `n` = number of samples,
- `m` = number of features,
- `ρ = m / n` = feature-to-sample ratio,
- `κ = n × m` = sample-feature product (total number of data cells).

For **Optuna-tuned configurations**, `M` is treated as a dataset-specific hyperparameter and can be optimized directly.

For the **default/no-tuning configuration**, BSTabDiff initializes `M` using both the HDLSS severity `ρ` and dataset scale `κ`:

```math
M=
\begin{cases}
m, & m \le 64\\
64, & m>64 \text{ and } \rho<10\\
\min\left(\widetilde{M}(\kappa),m\right), & \rho\ge10
\end{cases}
```

where

```math
\widetilde M(\kappa)=
\begin{cases}
32, & \kappa < 2.5\times10^{5}\\
64, & 2.5\times10^{5}\leq\kappa<10^{6}\\
128, & 10^{6}\leq\kappa<3\times10^{6}\\
192, & \kappa\geq3\times10^{6}
\end{cases}
```

In practical terms:

| Condition | Default `M` |
|---|---:|
| `m <= 64` | `m` |
| `m > 64` and `ρ < 10` | `64` |
| `ρ >= 10` and `κ < 2.5e5` | `32` |
| `ρ >= 10` and `2.5e5 <= κ < 1e6` | `64` |
| `ρ >= 10` and `1e6 <= κ < 3e6` | `128` |
| `ρ >= 10` and `κ >= 3e6` | `192` |

The final value is always capped by the number of available features:

```python
M = min(M, m)
```

A compact implementation is:

```python
def select_default_M(n, m):
    rho = m / n
    kappa = n * m

    if m <= 64:
        return m

    if rho < 10:
        return 64

    if kappa < 2.5e5:
        M = 32
    elif kappa < 1e6:
        M = 64
    elif kappa < 3e6:
        M = 128
    else:
        M = 192

    return min(M, m)
```

This rule is designed as a practical initialization. If block coherence or held-out validation utility stabilizes at a smaller value, `M` can be reduced accordingly. For tuned experiments, the user may instead include `M` directly in the Optuna search space.

### 2. Selecting GO-BS vs. GO-BS-FC

BSTabDiff provides two structure-aware ordering variants:

- **GO-BS**: the standard/sample clustered Graph-guided Ordering with Block Segmentation variant;
- **GO-BS-FC**: the feature clustered variant, designed to reduce ordering cost on sufficiently large high-dimensional datasets.

The choice is based on both:

- the dataset **shape**, characterized by `ρ = m / n`;
- the dataset **computational scale**, characterized by `κ = n × m`.

Let $\mathcal{R}_{\rho}(D)$ denote the dataset regime assigned by the BSTabDiff/DynaTab/GOTabPFN-style taxonomy, e.g. `HDLSS`, `HDHSS`, `LDHSS`, `LDLSS`, or `MixedRegime`. For each regime $r$, define a regime-specific scale threshold:

```math
T_r=\mathrm{median}\left\{\kappa_i:\mathcal{R}_{\rho}(D_i)=r\right\}
```

A new dataset is then classified as:

```math
\mathrm{Scale}(D)=\mathrm{Micro}\quad\mathrm{if}\quad \kappa<T_{\mathcal{R}_{\rho}(D)}
```

```math
\mathrm{Scale}(D)=\mathrm{Macro}\quad\mathrm{if}\quad \kappa\ge T_{\mathcal{R}_{\rho}(D)}
```


The default ordering variant is:

```math
\mathcal{V}(D)=
\begin{cases}
\mathrm{GO\!-\!BS\!-\!FC}, & \mathrm{Scale}(D)=\mathrm{Macro}\ \land\ \mathcal{R}_{\rho}(D)\in\{\mathrm{HDLSS},\mathrm{HDHSS}\}\\
\mathrm{GO\!-\!BS}, & \mathrm{otherwise}
\end{cases}
```
In practical terms:

| Dataset condition | Recommended variant |
|---|---|
| Macro-HDLSS | **GO-BS-FC** |
| Macro-HDHSS | **GO-BS-FC** |
| Micro-HDLSS | **GO-BS** |
| Micro-HDHSS | **GO-BS** |
| LDHSS | **GO-BS** |
| LDLSS | **GO-BS** |
| MixedRegime | **GO-BS** |

Therefore:

> **Use GO-BS-FC for Macro-HDLSS or Macro-HDHSS datasets; use GO-BS otherwise.**

The motivation is computational. For very large high-dimensional datasets, full ordering can become the primary bottleneck. GO-BS-FC first partitions the feature space into smaller feature clusters and performs ordering within these smaller subproblems, making it preferable for large `m` or large `κ`.

GO-BS remains the default for Micro-scale and otherwise computationally manageable datasets, where the standard sample clustered dependency construction can be performed affordably.

#### Near-boundary tie-breaker

If a dataset lies close to its regime-specific threshold $T_r$, bootstrap cluster stability can be used as a secondary diagnostic:

- more stable **sample clusters** → prefer **GO-BS**;
- more stable **feature clusters** → prefer **GO-BS-FC**.

This is a secondary tie breaker; the Micro/Macro rule above remains the primary default selection rule.

The examples below are complete runnable scripts. Each example uses the Colon dataset for illustration; for another dataset, update the dataset path, target column, and any dataset-specific pre-selected parameters.

## Example Usage

The seven examples below are complete, standalone scripts. Each can be copied directly and adapted to a new dataset by changing the dataset path, target column, and where applicable dataset-specific pre-selected parameters.

### Example 1 - Default, No Tuning, No Permutation
```python
import os, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.simplefilter("ignore")

import numpy as np
import pandas as pd
import torch
from bstabdiff import BSTabDiff

SEED = 42
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
N_SYN = 200
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def select_default_M(n, m):
    rho = m / n
    kappa = n * m
    if m <= 64:
        return m
    if rho < 10:
        return 64
    if kappa < 2.5e5:
        M = 32
    elif kappa < 1e6:
        M = 64
    elif kappa < 3e6:
        M = 128
    else:
        M = 192
    return min(M, m)

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X_df = df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number])
    X = np.nan_to_num(X_df.to_numpy(dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

set_seed(SEED)
X, y = load_data()
n, m = X.shape
M = select_default_M(n, m)

print(f"Dataset: n={n}, m={m}, rho={m/n:.4f}, kappa={n*m}, M={M}")
print(f"Device: {DEVICE}")

model = BSTabDiff(
    feature_specs=None,
    M=M,
    blocks=None,
    permute_features=False,
    prior_type="diffusion",
    device=DEVICE,
    random_state=SEED,
    prior_epochs=10000,
    prior_batch=128,
    prior_lr=1e-3,
    verbose_every=0,
    save_dir=None,
    save_name="bstabdiff_default",
    save_best=True,
    use_ema=True,
    ema_decay=0.999,
    use_gobs=False,
    use_gobs_fc=False
)

model.fit(X, y)
X_syn, R_syn, y_syn = model.sample(n_samples=N_SYN, return_mask=True)

print("X_syn:", X_syn.shape)
print("R_syn:", R_syn.shape)
print("y_syn:", y_syn.shape if y_syn is not None else None)
```

### Example 2 - No Permutation with Optuna Tuning
```python
import os, gc, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import torch
import optuna
from sklearn.model_selection import StratifiedShuffleSplit, RepeatedStratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.utils import check_random_state
from optuna.samplers import TPESampler
from bstabdiff import BSTabDiff

GLOBAL_SEED = 42
CV_SEED = 0
N_TRIALS = 5  # Increase for full tuning
INNER_VAL_SIZE = 0.2
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def cleanup():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X_df = df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number])
    X = np.nan_to_num(X_df.to_numpy(dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

def make_lr():
    return Pipeline([
        ("imp", SimpleImputer(strategy="median")),
        ("sc", StandardScaler()),
        ("clf", LogisticRegression(max_iter=10000, class_weight="balanced", solver="lbfgs"))
    ])

def safe_auc_or_acc(clf, X, y):
    try:
        return float(roc_auc_score(y, clf.predict_proba(X)[:, 1]))
    except Exception:
        return float(accuracy_score(y, clf.predict(X)))

set_seed(GLOBAL_SEED)
X, y = load_data()
print(f"Dataset: n={X.shape[0]}, m={X.shape[1]}")
print(f"Device: {DEVICE}")

splitter = StratifiedShuffleSplit(n_splits=1, test_size=INNER_VAL_SIZE, random_state=GLOBAL_SEED)
(itr, iva), = splitter.split(X, y)
X_inner_tr, y_inner_tr = X[itr], y[itr]
X_inner_va, y_inner_va = X[iva], y[iva]

def objective(trial):
    params = {
        "M": trial.suggest_categorical("M", [8, 12, 16, 24, 32, 48, 64]),
        "prior_type": trial.suggest_categorical("prior_type", ["diffusion", "flow"]),
        "prior_epochs": trial.suggest_categorical("prior_epochs", [300, 500, 800, 1200, 1600]),
        "prior_batch": trial.suggest_categorical("prior_batch", [16, 32, 64, 128]),
        "prior_lr": trial.suggest_float("prior_lr", 5e-5, 5e-3, log=True),
        "ema_decay": trial.suggest_float("ema_decay", 0.99, 0.9999, log=True),
        "n_syn": trial.suggest_categorical("n_syn", [64, 100, 128, 200, 256])
    }
    trial_seed = GLOBAL_SEED + 1000 + trial.number
    try:
        model = BSTabDiff(
            feature_specs=None,
            M=params["M"],
            blocks=None,
            permute_features=False,
            prior_type=params["prior_type"],
            device=DEVICE,
            random_state=trial_seed,
            prior_epochs=params["prior_epochs"],
            prior_batch=params["prior_batch"],
            prior_lr=params["prior_lr"],
            verbose_every=0,
            save_dir=None,
            save_name=f"bstabdiff_noperm_trial_{trial.number}",
            save_best=True,
            use_ema=True,
            ema_decay=params["ema_decay"],
            use_gobs=False,
            use_gobs_fc=False
        )
        model.fit(X_inner_tr, y_inner_tr)
        X_syn, _, y_syn = model.sample(n_samples=params["n_syn"], return_mask=True)
        X_syn = np.nan_to_num(np.asarray(X_syn, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        y_syn = np.asarray(y_syn).reshape(-1).astype(int)
        if np.unique(y_syn).size < 2:
            return 0.0
        clf = make_lr()
        clf.fit(X_syn, y_syn)
        score = safe_auc_or_acc(clf, X_inner_va, y_inner_va)
        trial.set_user_attr("raw_validation_score", score)
        return score - 1e-4 * params["M"]
    except Exception as e:
        trial.set_user_attr("error", repr(e))
        return 0.0
    finally:
        cleanup()

study = optuna.create_study(
    direction="maximize",
    sampler=TPESampler(seed=GLOBAL_SEED, multivariate=True, warn_independent_sampling=False)
)
study.optimize(objective, n_trials=N_TRIALS, show_progress_bar=True)

best = dict(study.best_params)
print("\nBest objective:", study.best_value)
print("Best parameters:", best)

def final_eval(best):
    rng = check_random_state(CV_SEED)
    rkf = RepeatedStratifiedKFold(n_splits=5, n_repeats=5, random_state=CV_SEED)
    tstr_acc, tstr_auc, trtr_acc, trtr_auc = [], [], [], []
    for fold, (tr, te) in enumerate(rkf.split(X, y), 1):
        Xtr, ytr = X[tr], y[tr]
        Xte, yte = X[te], y[te]
        fold_seed = int(rng.randint(0, 2**31 - 1))
        model = BSTabDiff(
            feature_specs=None,
            M=best["M"],
            blocks=None,
            permute_features=False,
            prior_type=best["prior_type"],
            device=DEVICE,
            random_state=fold_seed,
            prior_epochs=best["prior_epochs"],
            prior_batch=best["prior_batch"],
            prior_lr=best["prior_lr"],
            verbose_every=0,
            save_dir=None,
            save_name="bstabdiff_noperm_best",
            save_best=True,
            use_ema=True,
            ema_decay=best["ema_decay"],
            use_gobs=False,
            use_gobs_fc=False
        )
        model.fit(Xtr, ytr)
        X_syn, _, y_syn = model.sample(n_samples=best["n_syn"], return_mask=True)
        X_syn = np.nan_to_num(np.asarray(X_syn, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        y_syn = np.asarray(y_syn).reshape(-1).astype(int)

        clf = make_lr()
        clf.fit(X_syn, y_syn)
        p = clf.predict_proba(Xte)[:, 1]
        tstr_acc.append(accuracy_score(yte, (p >= 0.5).astype(int)))
        tstr_auc.append(roc_auc_score(yte, p))

        clf_real = make_lr()
        clf_real.fit(Xtr, ytr)
        p_real = clf_real.predict_proba(Xte)[:, 1]
        trtr_acc.append(accuracy_score(yte, (p_real >= 0.5).astype(int)))
        trtr_auc.append(roc_auc_score(yte, p_real))

        print(f"Fold {fold:02d}/25 | TSTR ACC={tstr_acc[-1]:.4f} AUC={tstr_auc[-1]:.4f}")
        del model, clf, clf_real
        cleanup()

    print("\nTUNED NO PERMUTATION - STRICT 5x5")
    print(f"TSTR ACC = {np.mean(tstr_acc):.4f} ± {np.std(tstr_acc):.4f}")
    print(f"TSTR AUC = {np.mean(tstr_auc):.4f} ± {np.std(tstr_auc):.4f}")
    print(f"TRTR ACC = {np.mean(trtr_acc):.4f} ± {np.std(trtr_acc):.4f}")
    print(f"TRTR AUC = {np.mean(trtr_auc):.4f} ± {np.std(trtr_auc):.4f}")

final_eval(best)
study.trials_dataframe().to_csv("bstabdiff_optuna_noperm_trials.csv", index=False)
```

### Example 3 - GO-BS with Optuna Tuning
```python
import os, gc, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import torch
import optuna
from sklearn.model_selection import StratifiedShuffleSplit, RepeatedStratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.utils import check_random_state
from optuna.samplers import TPESampler
from bstabdiff import BSTabDiff

GLOBAL_SEED = 42
CV_SEED = 0
N_TRIALS = 5  # Increase for full tuning
INNER_VAL_SIZE = 0.2
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

FIXED = {
    "M": 8,
    "prior_type": "flow",
    "prior_epochs": 300,
    "prior_batch": 64,
    "prior_lr": 3.2931e-3,
    "ema_decay": 0.9926,
    "n_syn": 64
}

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def cleanup():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X_df = df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number])
    X = np.nan_to_num(X_df.to_numpy(dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

def make_lr():
    return Pipeline([
        ("imp", SimpleImputer(strategy="median")),
        ("sc", StandardScaler()),
        ("clf", LogisticRegression(max_iter=10000, class_weight="balanced", solver="lbfgs"))
    ])

def safe_auc_or_acc(clf, X, y):
    try:
        return float(roc_auc_score(y, clf.predict_proba(X)[:, 1]))
    except Exception:
        return float(accuracy_score(y, clf.predict(X)))

set_seed(GLOBAL_SEED)
X, y = load_data()
print(f"Dataset: n={X.shape[0]}, m={X.shape[1]}")
print(f"Device: {DEVICE}")
print("Fixed generator parameters:", FIXED)

splitter = StratifiedShuffleSplit(n_splits=1, test_size=INNER_VAL_SIZE, random_state=GLOBAL_SEED)
(itr, iva), = splitter.split(X, y)
X_inner_tr, y_inner_tr = X[itr], y[itr]
X_inner_va, y_inner_va = X[iva], y[iva]

def objective(trial):
    params = {
        "gobs_num_clusters": trial.suggest_categorical("gobs_num_clusters", [3, 4, 5, 6, 7, 8, 10, 12]),
        "gobs_metric": trial.suggest_categorical("gobs_metric", ["correlation", "cosine", "manhattan", "euclidean", "kl_divergence"]),
        "gobs_bins": trial.suggest_categorical("gobs_bins", [16, 24, 32, 48, 64]),
        "gobs_top_k": trial.suggest_categorical("gobs_top_k", [None, 16, 32, 64, 128, 256]),
        "gobs_refine_order": trial.suggest_categorical("gobs_refine_order", [True, False]),
        "gobs_direction_select": trial.suggest_categorical("gobs_direction_select", [True, False]),
        "gobs_refine_passes": trial.suggest_categorical("gobs_refine_passes", [0, 1, 2, 3]),
        "gobs_boundary_refine_passes": trial.suggest_categorical("gobs_boundary_refine_passes", [0, 1, 3, 5, 8]),
        "gobs_boundary_window": trial.suggest_categorical("gobs_boundary_window", [4, 8, 12, 16, 24]),
        "gobs_lambda_cross": trial.suggest_float("gobs_lambda_cross", 0.1, 5.0, log=True),
        "gobs_gamma_balance": trial.suggest_float("gobs_gamma_balance", 1e-8, 1e-2, log=True)
    }
    trial_seed = GLOBAL_SEED + 1000 + trial.number
    try:
        model = BSTabDiff(
            feature_specs=None,
            M=FIXED["M"],
            blocks=None,
            permute_features=False,
            prior_type=FIXED["prior_type"],
            device=DEVICE,
            random_state=trial_seed,
            prior_epochs=FIXED["prior_epochs"],
            prior_batch=FIXED["prior_batch"],
            prior_lr=FIXED["prior_lr"],
            verbose_every=0,
            save_dir=None,
            save_name=f"bstabdiff_gobs_trial_{trial.number}",
            save_best=True,
            use_ema=True,
            ema_decay=FIXED["ema_decay"],
            use_gobs=True,
            use_gobs_fc=False,
            gobs_num_clusters=params["gobs_num_clusters"],
            gobs_metric=params["gobs_metric"],
            gobs_bins=params["gobs_bins"],
            gobs_top_k=params["gobs_top_k"],
            gobs_refine_order=params["gobs_refine_order"],
            gobs_direction_select=params["gobs_direction_select"],
            gobs_refine_passes=params["gobs_refine_passes"],
            gobs_boundary_refine_passes=params["gobs_boundary_refine_passes"],
            gobs_boundary_window=params["gobs_boundary_window"],
            gobs_lambda_cross=params["gobs_lambda_cross"],
            gobs_gamma_balance=params["gobs_gamma_balance"],
            gobs_use_cpu_kmeans=False
        )
        model.fit(X_inner_tr, y_inner_tr)
        X_syn, _, y_syn = model.sample(n_samples=FIXED["n_syn"], return_mask=True)
        X_syn = np.nan_to_num(np.asarray(X_syn, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        y_syn = np.asarray(y_syn).reshape(-1).astype(int)
        if np.unique(y_syn).size < 2:
            return 0.0
        clf = make_lr()
        clf.fit(X_syn, y_syn)
        score = safe_auc_or_acc(clf, X_inner_va, y_inner_va)
        trial.set_user_attr("raw_validation_score", score)
        penalty = 1e-4 * params["gobs_num_clusters"]
        if params["gobs_top_k"] is not None:
            penalty += 1e-7 * params["gobs_top_k"]
        return score - penalty
    except Exception as e:
        trial.set_user_attr("error", repr(e))
        return 0.0
    finally:
        cleanup()

study = optuna.create_study(
    direction="maximize",
    sampler=TPESampler(seed=GLOBAL_SEED, multivariate=True, warn_independent_sampling=False)
)
study.optimize(objective, n_trials=N_TRIALS, show_progress_bar=True)

best = dict(study.best_params)
print("\nBest objective:", study.best_value)
print("Best GO-BS parameters:", best)

def final_eval(best):
    rng = check_random_state(CV_SEED)
    rkf = RepeatedStratifiedKFold(n_splits=5, n_repeats=5, random_state=CV_SEED)
    tstr_acc, tstr_auc, trtr_acc, trtr_auc = [], [], [], []
    for fold, (tr, te) in enumerate(rkf.split(X, y), 1):
        Xtr, ytr = X[tr], y[tr]
        Xte, yte = X[te], y[te]
        fold_seed = int(rng.randint(0, 2**31 - 1))
        model = BSTabDiff(
            feature_specs=None,
            M=FIXED["M"],
            blocks=None,
            permute_features=False,
            prior_type=FIXED["prior_type"],
            device=DEVICE,
            random_state=fold_seed,
            prior_epochs=FIXED["prior_epochs"],
            prior_batch=FIXED["prior_batch"],
            prior_lr=FIXED["prior_lr"],
            verbose_every=0,
            save_dir=None,
            save_name="bstabdiff_gobs_best",
            save_best=True,
            use_ema=True,
            ema_decay=FIXED["ema_decay"],
            use_gobs=True,
            use_gobs_fc=False,
            gobs_num_clusters=best["gobs_num_clusters"],
            gobs_metric=best["gobs_metric"],
            gobs_bins=best["gobs_bins"],
            gobs_top_k=best["gobs_top_k"],
            gobs_refine_order=best["gobs_refine_order"],
            gobs_direction_select=best["gobs_direction_select"],
            gobs_refine_passes=best["gobs_refine_passes"],
            gobs_boundary_refine_passes=best["gobs_boundary_refine_passes"],
            gobs_boundary_window=best["gobs_boundary_window"],
            gobs_lambda_cross=best["gobs_lambda_cross"],
            gobs_gamma_balance=best["gobs_gamma_balance"],
            gobs_use_cpu_kmeans=False
        )
        model.fit(Xtr, ytr)
        X_syn, _, y_syn = model.sample(n_samples=FIXED["n_syn"], return_mask=True)
        X_syn = np.nan_to_num(np.asarray(X_syn, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        y_syn = np.asarray(y_syn).reshape(-1).astype(int)

        clf = make_lr()
        clf.fit(X_syn, y_syn)
        p = clf.predict_proba(Xte)[:, 1]
        tstr_acc.append(accuracy_score(yte, (p >= 0.5).astype(int)))
        tstr_auc.append(roc_auc_score(yte, p))

        clf_real = make_lr()
        clf_real.fit(Xtr, ytr)
        p_real = clf_real.predict_proba(Xte)[:, 1]
        trtr_acc.append(accuracy_score(yte, (p_real >= 0.5).astype(int)))
        trtr_auc.append(roc_auc_score(yte, p_real))

        print(f"Fold {fold:02d}/25 | TSTR ACC={tstr_acc[-1]:.4f} AUC={tstr_auc[-1]:.4f}")
        del model, clf, clf_real
        cleanup()

    print("\nTUNED GO-BS - STRICT 5x5")
    print(f"TSTR ACC = {np.mean(tstr_acc):.4f} ± {np.std(tstr_acc):.4f}")
    print(f"TSTR AUC = {np.mean(tstr_auc):.4f} ± {np.std(tstr_auc):.4f}")
    print(f"TRTR ACC = {np.mean(trtr_acc):.4f} ± {np.std(trtr_acc):.4f}")
    print(f"TRTR AUC = {np.mean(trtr_auc):.4f} ± {np.std(trtr_auc):.4f}")

final_eval(best)
study.trials_dataframe().to_csv("bstabdiff_optuna_gobs_trials.csv", index=False)
```

### Example 4 - GO-BS-FC with Optuna Tuning
```python
import os, gc, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import torch
import optuna
from sklearn.model_selection import StratifiedShuffleSplit, RepeatedStratifiedKFold
from sklearn.pipeline import Pipeline
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, roc_auc_score
from sklearn.utils import check_random_state
from optuna.samplers import TPESampler
from bstabdiff import BSTabDiff

GLOBAL_SEED = 42
CV_SEED = 0
N_TRIALS = 5  # Increase for full tuning
INNER_VAL_SIZE = 0.2
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

FIXED = {
    "M": 8,
    "prior_type": "flow",
    "prior_epochs": 300,
    "prior_batch": 64,
    "prior_lr": 3.2931e-3,
    "ema_decay": 0.9926,
    "n_syn": 64
}

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def cleanup():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X = np.nan_to_num(
        df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number]).to_numpy(dtype=np.float32),
        nan=0.0, posinf=0.0, neginf=0.0
    )
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

def make_lr():
    return Pipeline([
        ("imp", SimpleImputer(strategy="median")),
        ("sc", StandardScaler()),
        ("clf", LogisticRegression(max_iter=10000, class_weight="balanced", solver="lbfgs"))
    ])

def safe_auc_or_acc(clf, X, y):
    try:
        return float(roc_auc_score(y, clf.predict_proba(X)[:, 1]))
    except Exception:
        return float(accuracy_score(y, clf.predict(X)))

def sample_gobs_fc(trial):
    return {
        "gobs_num_clusters": trial.suggest_int("gobs_num_clusters", 3, 10),
        "gobs_metric": trial.suggest_categorical("gobs_metric", ["correlation", "cosine", "euclidean", "manhattan"]),
        "gobs_bins": trial.suggest_categorical("gobs_bins", [16, 32, 64]),
        "gobs_top_k": trial.suggest_categorical("gobs_top_k", [None, 8, 16, 32, 64]),
        "gobs_refine_order": trial.suggest_categorical("gobs_refine_order", [True, False]),
        "gobs_direction_select": trial.suggest_categorical("gobs_direction_select", [True, False]),
        "gobs_refine_passes": trial.suggest_int("gobs_refine_passes", 0, 3),
        "gobs_boundary_refine_passes": trial.suggest_int("gobs_boundary_refine_passes", 1, 8),
        "gobs_boundary_window": trial.suggest_categorical("gobs_boundary_window", [4, 8, 12, 16]),
        "gobs_lambda_cross": trial.suggest_float("gobs_lambda_cross", 0.1, 5.0, log=True),
        "gobs_gamma_balance": trial.suggest_float("gobs_gamma_balance", 1e-6, 1e-1, log=True)
    }

def build_model(gobs, seed, name):
    return BSTabDiff(
        feature_specs=None,
        M=FIXED["M"],
        blocks=None,
        permute_features=False,
        prior_type=FIXED["prior_type"],
        device=DEVICE,
        random_state=seed,
        prior_epochs=FIXED["prior_epochs"],
        prior_batch=FIXED["prior_batch"],
        prior_lr=FIXED["prior_lr"],
        verbose_every=0,
        save_dir=None,
        save_name=name,
        save_best=True,
        use_ema=True,
        ema_decay=FIXED["ema_decay"],
        use_gobs=True,
        use_gobs_fc=True,
        gobs_use_cpu_kmeans=False,
        **gobs
    )

set_seed(GLOBAL_SEED)
X, y = load_data()
print(f"Dataset: n={X.shape[0]}, m={X.shape[1]}")
print(f"Device: {DEVICE}")
print("Fixed generator parameters:", FIXED)

splitter = StratifiedShuffleSplit(n_splits=1, test_size=INNER_VAL_SIZE, random_state=GLOBAL_SEED)
(itr, iva), = splitter.split(X, y)
X_inner_tr, y_inner_tr = X[itr], y[itr]
X_inner_va, y_inner_va = X[iva], y[iva]

def objective(trial):
    gobs = sample_gobs_fc(trial)
    trial_seed = GLOBAL_SEED + 1000 + trial.number
    try:
        model = build_model(gobs, trial_seed, f"bstabdiff_gobs_fc_trial_{trial.number}")
        model.fit(X_inner_tr, y_inner_tr)
        X_syn, _, y_syn = model.sample(n_samples=FIXED["n_syn"], return_mask=True)
        X_syn = np.nan_to_num(np.asarray(X_syn, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        y_syn = np.asarray(y_syn).reshape(-1).astype(int)
        if np.unique(y_syn).size < 2:
            return 0.0
        clf = make_lr()
        clf.fit(X_syn, y_syn)
        score = safe_auc_or_acc(clf, X_inner_va, y_inner_va)
        trial.set_user_attr("raw_validation_score", score)
        return score - 1e-4 * FIXED["M"]
    except Exception as e:
        trial.set_user_attr("error", repr(e))
        return 0.0
    finally:
        cleanup()

study = optuna.create_study(
    direction="maximize",
    sampler=TPESampler(seed=GLOBAL_SEED, multivariate=True, warn_independent_sampling=False)
)
study.optimize(objective, n_trials=N_TRIALS, show_progress_bar=True)

best = dict(study.best_params)
print("\nBest objective:", study.best_value)
print("Best GO-BS-FC parameters:", best)

def final_eval(best):
    rng = check_random_state(CV_SEED)
    rkf = RepeatedStratifiedKFold(n_splits=5, n_repeats=5, random_state=CV_SEED)
    tstr_acc, tstr_auc, trtr_acc, trtr_auc = [], [], [], []
    for fold, (tr, te) in enumerate(rkf.split(X, y), 1):
        Xtr, ytr = X[tr], y[tr]
        Xte, yte = X[te], y[te]
        fold_seed = int(rng.randint(0, 2**31 - 1))
        model = build_model(best, fold_seed, "bstabdiff_gobs_fc_best")
        model.fit(Xtr, ytr)
        X_syn, _, y_syn = model.sample(n_samples=FIXED["n_syn"], return_mask=True)
        X_syn = np.nan_to_num(np.asarray(X_syn, dtype=np.float32), nan=0.0, posinf=0.0, neginf=0.0)
        y_syn = np.asarray(y_syn).reshape(-1).astype(int)

        clf = make_lr()
        clf.fit(X_syn, y_syn)
        p = clf.predict_proba(Xte)[:, 1]
        tstr_acc.append(accuracy_score(yte, (p >= 0.5).astype(int)))
        tstr_auc.append(roc_auc_score(yte, p))

        clf_real = make_lr()
        clf_real.fit(Xtr, ytr)
        p_real = clf_real.predict_proba(Xte)[:, 1]
        trtr_acc.append(accuracy_score(yte, (p_real >= 0.5).astype(int)))
        trtr_auc.append(roc_auc_score(yte, p_real))

        print(f"Fold {fold:02d}/25 | TSTR ACC={tstr_acc[-1]:.4f} AUC={tstr_auc[-1]:.4f}")
        del model, clf, clf_real
        cleanup()

    print("\nTUNED GO-BS-FC - STRICT 5x5")
    print(f"TSTR ACC = {np.mean(tstr_acc):.4f} ± {np.std(tstr_acc):.4f}")
    print(f"TSTR AUC = {np.mean(tstr_auc):.4f} ± {np.std(tstr_auc):.4f}")
    print(f"TRTR ACC = {np.mean(trtr_acc):.4f} ± {np.std(trtr_acc):.4f}")
    print(f"TRTR AUC = {np.mean(trtr_auc):.4f} ± {np.std(trtr_auc):.4f}")

final_eval(best)
study.trials_dataframe().to_csv("bstabdiff_optuna_gobs_fc_trials.csv", index=False)
```

### Example 5 - No Permutation with Pre-selected Parameters
```python
import os, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import torch
from bstabdiff import BSTabDiff

SEED = 42
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
N_SYN = 64
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

PARAMS = {
    "M": 8,
    "prior_type": "flow",
    "prior_epochs": 300,
    "prior_batch": 64,
    "prior_lr": 3.2931e-3,
    "ema_decay": 0.9926
}

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X = np.nan_to_num(
        df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number]).to_numpy(dtype=np.float32),
        nan=0.0, posinf=0.0, neginf=0.0
    )
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

set_seed(SEED)
X, y = load_data()

print(f"Dataset: n={X.shape[0]}, m={X.shape[1]}")
print(f"Device: {DEVICE}")
print("Pre-selected parameters:", PARAMS)

model = BSTabDiff(
    feature_specs=None,
    M=PARAMS["M"],
    blocks=None,
    permute_features=False,
    prior_type=PARAMS["prior_type"],
    device=DEVICE,
    random_state=SEED,
    prior_epochs=PARAMS["prior_epochs"],
    prior_batch=PARAMS["prior_batch"],
    prior_lr=PARAMS["prior_lr"],
    verbose_every=0,
    save_dir=None,
    save_name="bstabdiff_noperm_preselected",
    save_best=True,
    use_ema=True,
    ema_decay=PARAMS["ema_decay"],
    use_gobs=False,
    use_gobs_fc=False
)

model.fit(X, y)
X_syn, R_syn, y_syn = model.sample(n_samples=N_SYN, return_mask=True)

X_syn = np.asarray(X_syn, dtype=np.float32)
R_syn = np.asarray(R_syn)
y_syn = None if y_syn is None else np.asarray(y_syn).astype(int)

print("X_syn:", X_syn.shape)
print("R_syn:", R_syn.shape)
print("y_syn:", y_syn.shape if y_syn is not None else None)

pd.DataFrame(X_syn).to_csv("bstabdiff_noperm_synthetic_X.csv", index=False)
if y_syn is not None:
    pd.DataFrame({"label": y_syn}).to_csv("bstabdiff_noperm_synthetic_y.csv", index=False)

print("Saved synthetic data.")
```

### Example 6 - GO-BS with Pre-selected Parameters
```python
import os, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import torch
from bstabdiff import BSTabDiff

SEED = 42
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
N_SYN = 64
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

PARAMS = {
    "M": 8,
    "prior_type": "flow",
    "prior_epochs": 300,
    "prior_batch": 64,
    "prior_lr": 3.2931e-3,
    "ema_decay": 0.9926,
    "gobs_num_clusters": 3,
    "gobs_metric": "correlation",
    "gobs_bins": 24,
    "gobs_top_k": 16,
    "gobs_refine_order": False,
    "gobs_direction_select": False,
    "gobs_refine_passes": 3,
    "gobs_boundary_refine_passes": 5,
    "gobs_boundary_window": 16,
    "gobs_lambda_cross": 2.3989,
    "gobs_gamma_balance": 5.19e-6
}

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X = np.nan_to_num(
        df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number]).to_numpy(dtype=np.float32),
        nan=0.0, posinf=0.0, neginf=0.0
    )
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

set_seed(SEED)
X, y = load_data()

print(f"Dataset: n={X.shape[0]}, m={X.shape[1]}")
print(f"Device: {DEVICE}")
print("Pre-selected parameters:", PARAMS)

model = BSTabDiff(
    feature_specs=None,
    M=PARAMS["M"],
    blocks=None,
    permute_features=False,
    prior_type=PARAMS["prior_type"],
    device=DEVICE,
    random_state=SEED,
    prior_epochs=PARAMS["prior_epochs"],
    prior_batch=PARAMS["prior_batch"],
    prior_lr=PARAMS["prior_lr"],
    verbose_every=0,
    save_dir=None,
    save_name="bstabdiff_gobs_preselected",
    save_best=True,
    use_ema=True,
    ema_decay=PARAMS["ema_decay"],
    use_gobs=True,
    use_gobs_fc=False,
    gobs_num_clusters=PARAMS["gobs_num_clusters"],
    gobs_metric=PARAMS["gobs_metric"],
    gobs_bins=PARAMS["gobs_bins"],
    gobs_top_k=PARAMS["gobs_top_k"],
    gobs_refine_order=PARAMS["gobs_refine_order"],
    gobs_direction_select=PARAMS["gobs_direction_select"],
    gobs_refine_passes=PARAMS["gobs_refine_passes"],
    gobs_boundary_refine_passes=PARAMS["gobs_boundary_refine_passes"],
    gobs_boundary_window=PARAMS["gobs_boundary_window"],
    gobs_lambda_cross=PARAMS["gobs_lambda_cross"],
    gobs_gamma_balance=PARAMS["gobs_gamma_balance"],
    gobs_use_cpu_kmeans=False
)

model.fit(X, y)
X_syn, R_syn, y_syn = model.sample(n_samples=N_SYN, return_mask=True)

X_syn = np.asarray(X_syn, dtype=np.float32)
R_syn = np.asarray(R_syn)
y_syn = None if y_syn is None else np.asarray(y_syn).astype(int)

print("X_syn:", X_syn.shape)
print("R_syn:", R_syn.shape)
print("y_syn:", y_syn.shape if y_syn is not None else None)

pd.DataFrame(X_syn).to_csv("bstabdiff_gobs_synthetic_X.csv", index=False)
if y_syn is not None:
    pd.DataFrame({"label": y_syn}).to_csv("bstabdiff_gobs_synthetic_y.csv", index=False)

print("Saved synthetic data.")
```

### Example 7 - GO-BS-FC with Pre-selected Parameters
```python
import os, random, warnings
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd
import torch
from bstabdiff import BSTabDiff

SEED = 42
CSV_PATH = "coloncancer_encoded.csv" #change your dataset file here
LABEL_COL = "label" #change your target column here
N_SYN = 64
DEVICE = "cuda:0" if torch.cuda.is_available() else "cpu"

PARAMS = {
    "M": 8,
    "prior_type": "flow",
    "prior_epochs": 300,
    "prior_batch": 64,
    "prior_lr": 3.2931e-3,
    "ema_decay": 0.9926,
    "gobs_num_clusters": 3,
    "gobs_metric": "correlation",
    "gobs_bins": 24,
    "gobs_top_k": 16,
    "gobs_refine_order": False,
    "gobs_direction_select": False,
    "gobs_refine_passes": 3,
    "gobs_boundary_refine_passes": 5,
    "gobs_boundary_window": 16,
    "gobs_lambda_cross": 2.3989,
    "gobs_gamma_balance": 5.19e-6
}

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

def load_data():
    df = pd.read_csv(CSV_PATH)
    if LABEL_COL not in df.columns:
        raise ValueError(f"Label column '{LABEL_COL}' not found.")
    y_raw = df[LABEL_COL].values
    X = np.nan_to_num(
        df.drop(columns=[LABEL_COL]).select_dtypes(include=[np.number]).to_numpy(dtype=np.float32),
        nan=0.0, posinf=0.0, neginf=0.0
    )
    _, y = np.unique(y_raw, return_inverse=True)
    return X, y.astype(int)

set_seed(SEED)
X, y = load_data()

print(f"Dataset: n={X.shape[0]}, m={X.shape[1]}")
print(f"Device: {DEVICE}")
print("Pre-selected parameters:", PARAMS)

model = BSTabDiff(
    feature_specs=None,
    M=PARAMS["M"],
    blocks=None,
    permute_features=False,
    prior_type=PARAMS["prior_type"],
    device=DEVICE,
    random_state=SEED,
    prior_epochs=PARAMS["prior_epochs"],
    prior_batch=PARAMS["prior_batch"],
    prior_lr=PARAMS["prior_lr"],
    verbose_every=0,
    save_dir=None,
    save_name="bstabdiff_gobs_fc_preselected",
    save_best=True,
    use_ema=True,
    ema_decay=PARAMS["ema_decay"],
    use_gobs=True,
    use_gobs_fc=True,
    gobs_num_clusters=PARAMS["gobs_num_clusters"],
    gobs_metric=PARAMS["gobs_metric"],
    gobs_bins=PARAMS["gobs_bins"],
    gobs_top_k=PARAMS["gobs_top_k"],
    gobs_refine_order=PARAMS["gobs_refine_order"],
    gobs_direction_select=PARAMS["gobs_direction_select"],
    gobs_refine_passes=PARAMS["gobs_refine_passes"],
    gobs_boundary_refine_passes=PARAMS["gobs_boundary_refine_passes"],
    gobs_boundary_window=PARAMS["gobs_boundary_window"],
    gobs_lambda_cross=PARAMS["gobs_lambda_cross"],
    gobs_gamma_balance=PARAMS["gobs_gamma_balance"],
    gobs_use_cpu_kmeans=False
)

model.fit(X, y)
X_syn, R_syn, y_syn = model.sample(n_samples=N_SYN, return_mask=True)

X_syn = np.asarray(X_syn, dtype=np.float32)
R_syn = np.asarray(R_syn)
y_syn = None if y_syn is None else np.asarray(y_syn).astype(int)

print("X_syn:", X_syn.shape)
print("R_syn:", R_syn.shape)
print("y_syn:", y_syn.shape if y_syn is not None else None)

pd.DataFrame(X_syn).to_csv("bstabdiff_gobs_fc_synthetic_X.csv", index=False)
if y_syn is not None:
    pd.DataFrame({"label": y_syn}).to_csv("bstabdiff_gobs_fc_synthetic_y.csv", index=False)

print("Saved synthetic data.")
```

### Mixed-Type Data with Missingness

BSTabDiff also supports **mixed-type tabular data** containing both continuous and categorical features, including datasets with missing values. Categorical variables are specified explicitly through `FeatureSpec`, while missingness is modeled jointly during generation.

For a complete runnable example using a dummy mixed-type dataset with categorical values and structured missingness, see: **[BSTabDiff Mixed-Data Smoke Test](https://github.com/zadid6pretam/BSTabDiff/blob/main/Smoke%20Tests/BSTabDiff_MixedData_SmokeTest.ipynb)**

## Our Previous Related Work Involving Tabular Data

BSTabDiff is part of our broader line of work on tabular deep learning, feature ordering, multimodal learning, and high-dimensional tabular modeling.

### GOTabPFN
- **Title:** GOTabPFN: From Feature Ordering to Compact Tokenization for Tabular Foundation Models on High-Dimensional Data
- **Venue:** ICML 2026
- **GitHub:** https://github.com/zadid6pretam/GOTabPFN

### iSyncTab
- **Title:** iSyncTab: Learning Cross-Modal Feature Sequencing for Image-Tabular Data via Neural Synchrony
- **Venue:** ECCV 2026
- **GitHub:** https://github.com/zadid6pretam/iSyncTab

### iStructTab
- **Title:** iStructTab: Structured Feature Sequencing for Multimodal Learning of Image and Tabular Data
- **Venue:** ICPR 2026
- **GitHub:** https://github.com/zadid6pretam/iStructTab

### DynaTab
- **Title:** DynaTab: Dynamic Feature Ordering as Neural Rewiring for High-Dimensional Tabular Data
- **Venue:** AAAI 2026 NeuroAI Workshop
- **GitHub:** https://github.com/zadid6pretam/DynaTab

### TabSeq
- **Title:** TabSeq: A Framework for Deep Learning on Tabular Data via Sequential Ordering
- **Venue:** ICPR 2024
- **GitHub:** https://github.com/zadid6pretam/TabSeq

### Side Projects

#### ZAYAN
- **Title:** ZAYAN: Disentangled Contrastive Transformer for Tabular Remote Sensing Data
- **Venue:** ICPR 2026
- **GitHub:** https://github.com/zadid6pretam/ZAYAN

#### AugTab
- **Title:** AugTab: Learnable Feature Augmentation for Low-Dimensional Tabular Data
- **Venue:** ECML-PKDD 2026 Research Track
- **GitHub:** https://github.com/zadid6pretam/AugTab

## Contact

For any questions, issues, or suggestions related to this repository, please feel free to contact us or open an issue on GitHub.
