flowMC

Normalizing-flow enhanced sampling package for probabilistic inference

Warning

flowMC has not yet reached v1.0.0, and the API may change. Higher-level APIs are more stable, while intermediate-level APIs (such as the resource strategy interface) may undergo major revisions for performance improvements.

flowMC is a JAX-based Python package for normalizing-flow enhanced Markov chain Monte Carlo (MCMC) sampling. The code is open source under the MIT license and is under active development.

flowMC implements methods described in Gabrié et al. (2021) and Gabrié et al. (2022). See the accompanying paper, Wong, Gabrié, Foreman-Mackey (2023), for more details.

• Just-in-time compilation support
• Native GPU acceleration
• Effective for multi-modal problems
• Minimal hyperparameter tuning required

Installation

The simplest way to install flowMC is through pip:

pip install flowMC

This will install the latest stable release and its dependencies. flowMC is built on JAX and Equinox. By default, this installs the CPU version of JAX from PyPI. If you have a GPU and want to leverage hardware acceleration, install the CUDA-enabled version:

pip install flowMC[cuda]

If you want to install the latest version of flowMC, you can clone this repo and install it locally:

git clone https://github.com/GW-JAX-Team/flowMC.git
cd flowMC
pip install -e .

Additional optional dependencies are available:

• flowMC[docs]: Documentation dependencies
• flowMC[codeqa]: Code quality tools
• flowMC[visualize]: Visualization dependencies

We recommend using uv to manage your Python environment. After cloning the repository, run uv sync to create a virtual environment with all dependencies installed.

Attribution

If you use flowMC in your research, please cite the following papers:

@article{Wong:2022xvh,
    author = "Wong, Kaze W. k. and Gabri\'e, Marylou and Foreman-Mackey, Daniel",
    title = "{flowMC: Normalizing flow enhanced sampling package for probabilistic inference in JAX}",
    eprint = "2211.06397",
    archivePrefix = "arXiv",
    primaryClass = "astro-ph.IM",
    doi = "10.21105/joss.05021",
    journal = "J. Open Source Softw.",
    volume = "8",
    number = "83",
    pages = "5021",
    year = "2023"
}

@article{Gabrie:2021tlu,
    author = "Gabri\'e, Marylou and Rotskoff, Grant M. and Vanden-Eijnden, Eric",
    title = "{Adaptive Monte Carlo augmented with normalizing flows}",
    eprint = "2105.12603",
    archivePrefix = "arXiv",
    primaryClass = "physics.data-an",
    doi = "10.1073/pnas.2109420119",
    journal = "Proc. Nat. Acad. Sci.",
    volume = "119",
    number = "10",
    pages = "e2109420119",
    year = "2022"
}