CSlearn

A Python library for context-specific causal graphical models.
Docs »

Overview

cslearn is a Python package for CStree models—a family of graphical causal models for multivariate discrete data that encode context-specific independence (CSI). CStrees generalize DAG models while remaining tractable.

The package implements CSlearn, a three-phase structure-learning algorithm:

1. DAG pre-screening via PC or GRaSP to restrict candidate parent sets
2. Order MCMC (Gibbs sampler) over topological orderings
3. Exact staging search under a sparsity bound

Installation

cslearn requires graphviz to be installed on your system.

On Debian/Ubuntu/Linux:

sudo apt install graphviz libgraphviz-dev pkg-config

On macOS (Homebrew):

brew install graphviz

On Windows, install graphviz from https://graphviz.org/download/ and ensure it is on your PATH.

Then install the package:

pip install cslearn

See the full installation instructions for development setup.

Quick start

from cslearn import CStree, sample_cstree

# Sample a random CStree and simulate data
tree = sample_cstree([2, 2, 3, 2], max_cvars=2, prob_cvar=0.5)
tree.sample_stage_parameters(alpha=2.0)
data = tree.sample(500)

# Learn a CStree from data
learned = CStree().fit(data)

# Predict the last variable for five held-out observations
# (row 0 of `data` holds the cardinalities, so skip it)
predictions = learned.predict(data.iloc[1:6, :-1])

See the example notebooks for walkthroughs covering CStree construction and visualization, structure learning with exact and Gibbs-sampler search, LDAG representations on the alarm and Sachs datasets, and prediction.

Paper experiments

The simulation experiments and figures from the accompanying paper are in src/expt/. Aggregated results (CSVs) are committed to the repository; precomputed intermediates (~5.9 GB) are on Zenodo at https://doi.org/10.5281/zenodo.21198084. See src/expt/README.md for reproduction instructions.

Reference

If you use this package, please cite the accompanying paper:

Rios, F. L., Markham, A. & Solus, L. (2024). Scalable Structure Learning for Sparse Context-Specific Systems. arXiv:2402.07762

@misc{rios2024scalablestructurelearningsparse,
      title={Scalable Structure Learning for Sparse Context-Specific Systems},
      author={Felix Leopoldo Rios and Alex Markham and Liam Solus},
      year={2024},
      eprint={2402.07762},
      archivePrefix={arXiv},
      primaryClass={stat.ML},
      url={https://arxiv.org/abs/2402.07762},
}

Contributing

Contributions are welcome. Please open an issue or pull request on GitHub.