An AI-powered Bayesian generative modeling approach for causal inference in observational studies

CausalBGM is an AI-powered Bayesian generative modeling approach that captures the causal relationship among covariates, treatment, and outcome variables.

CausalBGM adopts a Bayesian iterative approach to update the model parameters and the posterior distribution of latent features until convergence. This framework leverages the power of AI to capture complex dependencies among variables while adhering to the Bayesian principles.

CausalBGM was developed with Python3.9, TensorFlow2.10, and TensorFlow Probability. Now both Python and R packages for CausalBGM are available! Besides, we provide a console program to run CausalBGM directly. For more information, checkout the Document.

CausalBGM Main Applications

• Point estimate of ATE, ITE, ADRF, CATE.

• Posterior interval estimate of ATE, ITE, ADRF, CATE with user-specific significant level α (alpha).

Installation

See detailed installation instructions in our website. Briefly, CausalBGM can be installed via

pip install bayesgm

where bayesgm is a Python package providing AI-powered Bayesian generative modeling approaches and CausalBGM is one of the methods focusing on causal inference.

Usage

A detailed Python tutorial can be found at our website. The source Python notebook for the detailed tutorial is provided at here.

Example Usage of CausalBGM

import yaml
import numpy as np
import bayesgm
from bayesgm.models import CausalBGM
from bayesgm.datasets import Sim_Hirano_Imbens_sampler

params = yaml.safe_load(open('src/configs/Sim_Hirano_Imbens.yaml', 'r'))
x, y, v = Sim_Hirano_Imbens_sampler(N=20000, v_dim=200).load_all()

# Instantiate a CausalBGM model
model = CausalBGM(params=params, random_seed=None)

# Perform Encoding Generative Modeling (EGM) initialization (optional but recommended)
model.egm_init(data=(x, y, v), n_iter=30000, batches_per_eval=500, verbose=1)

# Train the CausalBGM model with an iterative updating algorithm
model.fit(data=(x, y, v), epochs=100, epochs_per_eval=10, verbose=1)

# Make predictions using the trained CausalBGM model
causal_pre, pos_intervals = model.predict(
  data=(x, y, v), alpha=0.01, n_mcmc=3000, x_values=np.linspace(0, 3, 20), q_sd=1.0
)

Datasets

bayesgm package provides several built-in simulation datasets from bayesgm.datasets.

For semi-synthetic dataset, users need to create a CausalBGM/data folder and uncompress the dataset in the CausalBGM/data folder.

• Twin dataset. Google Drive download link.

• ACIC 2018 datasets. Google Drive download link.

Main References

• Qiao Liu and Wing Hung Wong. An AI-powered Bayesian generative modeling approach for causal inference in observational studies [J]. arXiv preprint arXiv:2501.00755, 2025.

• Qiao Liu, Zhongren Chen, and Wing Hung Wong. An encoding generative modeling approach to dimension reduction and covariate adjustment in causal inference with observational studies [J]. PNAS, 121 (23) e2322376121, 2024.

Support

Found a bug or would like to see a feature implemented? Feel free to submit an issue.

Have a question or would like to start a new discussion? You can also always send us an e-mail.

Your help to improve CausalBGM is highly appreciated! For further information visit website.